1 | /* Copyright 2015 The TensorFlow Authors. All Rights Reserved. |
2 | |
3 | Licensed under the Apache License, Version 2.0 (the "License"); |
4 | you may not use this file except in compliance with the License. |
5 | You may obtain a copy of the License at |
6 | |
7 | http://www.apache.org/licenses/LICENSE-2.0 |
8 | |
9 | Unless required by applicable law or agreed to in writing, software |
10 | distributed under the License is distributed on an "AS IS" BASIS, |
11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | See the License for the specific language governing permissions and |
13 | limitations under the License. |
14 | ==============================================================================*/ |
15 | |
16 | #include "tensorflow/tsl/platform/env.h" |
17 | |
18 | #include <sys/stat.h> |
19 | |
20 | #include <deque> |
21 | #include <utility> |
22 | #include <vector> |
23 | |
24 | #include "tensorflow/tsl/platform/env_time.h" |
25 | #include "tensorflow/tsl/platform/errors.h" |
26 | #include "tensorflow/tsl/platform/host_info.h" |
27 | #include "tensorflow/tsl/platform/path.h" |
28 | #include "tensorflow/tsl/platform/platform.h" |
29 | #include "tensorflow/tsl/platform/protobuf.h" |
30 | #include "tensorflow/tsl/platform/stringprintf.h" |
31 | |
32 | #if defined(__APPLE__) |
33 | #include <mach-o/dyld.h> |
34 | #endif |
35 | #if defined(__FreeBSD__) |
36 | #include <sys/sysctl.h> |
37 | #endif |
38 | #if defined(PLATFORM_WINDOWS) |
39 | #include <windows.h> |
40 | #undef DeleteFile |
41 | #undef CopyFile |
42 | #include "tensorflow/tsl/platform/windows/wide_char.h" |
43 | #define PATH_MAX MAX_PATH |
44 | #else |
45 | #include <fcntl.h> |
46 | #include <string.h> |
47 | #include <sys/types.h> |
48 | #include <unistd.h> |
49 | #endif |
50 | |
51 | namespace tsl { |
52 | |
53 | // 128KB copy buffer |
54 | constexpr size_t kCopyFileBufferSize = 128 * 1024; |
55 | |
56 | class FileSystemRegistryImpl : public FileSystemRegistry { |
57 | public: |
58 | Status Register(const std::string& scheme, Factory factory) override; |
59 | Status Register(const std::string& scheme, |
60 | std::unique_ptr<FileSystem> filesystem) override; |
61 | FileSystem* Lookup(const std::string& scheme) override; |
62 | Status GetRegisteredFileSystemSchemes( |
63 | std::vector<std::string>* schemes) override; |
64 | |
65 | private: |
66 | mutable mutex mu_; |
67 | mutable std::unordered_map<std::string, std::unique_ptr<FileSystem>> registry_ |
68 | TF_GUARDED_BY(mu_); |
69 | }; |
70 | |
71 | Status FileSystemRegistryImpl::Register(const std::string& scheme, |
72 | FileSystemRegistry::Factory factory) { |
73 | mutex_lock lock(mu_); |
74 | if (!registry_.emplace(scheme, std::unique_ptr<FileSystem>(factory())) |
75 | .second) { |
76 | return errors::AlreadyExists("File factory for " , scheme, |
77 | " already registered" ); |
78 | } |
79 | return OkStatus(); |
80 | } |
81 | |
82 | Status FileSystemRegistryImpl::Register( |
83 | const std::string& scheme, std::unique_ptr<FileSystem> filesystem) { |
84 | mutex_lock lock(mu_); |
85 | if (!registry_.emplace(scheme, std::move(filesystem)).second) { |
86 | return errors::AlreadyExists("File system for " , scheme, |
87 | " already registered" ); |
88 | } |
89 | return OkStatus(); |
90 | } |
91 | |
92 | FileSystem* FileSystemRegistryImpl::Lookup(const std::string& scheme) { |
93 | mutex_lock lock(mu_); |
94 | const auto found = registry_.find(scheme); |
95 | if (found == registry_.end()) { |
96 | return nullptr; |
97 | } |
98 | return found->second.get(); |
99 | } |
100 | |
101 | Status FileSystemRegistryImpl::GetRegisteredFileSystemSchemes( |
102 | std::vector<std::string>* schemes) { |
103 | mutex_lock lock(mu_); |
104 | for (const auto& e : registry_) { |
105 | schemes->push_back(e.first); |
106 | } |
107 | return OkStatus(); |
108 | } |
109 | |
110 | Env::Env() : file_system_registry_(new FileSystemRegistryImpl) {} |
111 | |
112 | Status Env::GetFileSystemForFile(const std::string& fname, |
113 | FileSystem** result) { |
114 | StringPiece scheme, host, path; |
115 | io::ParseURI(fname, &scheme, &host, &path); |
116 | FileSystem* file_system = file_system_registry_->Lookup(std::string(scheme)); |
117 | if (!file_system) { |
118 | if (scheme.empty()) { |
119 | scheme = "[local]" ; |
120 | } |
121 | |
122 | return errors::Unimplemented("File system scheme '" , scheme, |
123 | "' not implemented (file: '" , fname, "')" ); |
124 | } |
125 | *result = file_system; |
126 | return OkStatus(); |
127 | } |
128 | |
129 | Status Env::GetRegisteredFileSystemSchemes(std::vector<std::string>* schemes) { |
130 | return file_system_registry_->GetRegisteredFileSystemSchemes(schemes); |
131 | } |
132 | |
133 | Status Env::RegisterFileSystem(const std::string& scheme, |
134 | FileSystemRegistry::Factory factory) { |
135 | return file_system_registry_->Register(scheme, std::move(factory)); |
136 | } |
137 | |
138 | Status Env::RegisterFileSystem(const std::string& scheme, |
139 | std::unique_ptr<FileSystem> filesystem) { |
140 | return file_system_registry_->Register(scheme, std::move(filesystem)); |
141 | } |
142 | |
143 | Status Env::SetOption(const std::string& scheme, const std::string& key, |
144 | const std::string& value) { |
145 | FileSystem* file_system = file_system_registry_->Lookup(scheme); |
146 | if (!file_system) { |
147 | return errors::Unimplemented("File system scheme '" , scheme, |
148 | "' not found to set configuration" ); |
149 | } |
150 | return file_system->SetOption(key, value); |
151 | } |
152 | |
153 | Status Env::SetOption(const std::string& scheme, const std::string& key, |
154 | const std::vector<string>& values) { |
155 | FileSystem* file_system = file_system_registry_->Lookup(scheme); |
156 | if (!file_system) { |
157 | return errors::Unimplemented("File system scheme '" , scheme, |
158 | "' not found to set configuration" ); |
159 | } |
160 | return file_system->SetOption(key, values); |
161 | } |
162 | |
163 | Status Env::SetOption(const std::string& scheme, const std::string& key, |
164 | const std::vector<int64_t>& values) { |
165 | FileSystem* file_system = file_system_registry_->Lookup(scheme); |
166 | if (!file_system) { |
167 | return errors::Unimplemented("File system scheme '" , scheme, |
168 | "' not found to set configuration" ); |
169 | } |
170 | return file_system->SetOption(key, values); |
171 | } |
172 | |
173 | Status Env::SetOption(const std::string& scheme, const std::string& key, |
174 | const std::vector<double>& values) { |
175 | FileSystem* file_system = file_system_registry_->Lookup(scheme); |
176 | if (!file_system) { |
177 | return errors::Unimplemented("File system scheme '" , scheme, |
178 | "' not found to set configuration" ); |
179 | } |
180 | return file_system->SetOption(key, values); |
181 | } |
182 | |
183 | Status Env::FlushFileSystemCaches() { |
184 | std::vector<string> schemes; |
185 | TF_RETURN_IF_ERROR(GetRegisteredFileSystemSchemes(&schemes)); |
186 | for (const string& scheme : schemes) { |
187 | FileSystem* fs = nullptr; |
188 | TF_RETURN_IF_ERROR( |
189 | GetFileSystemForFile(io::CreateURI(scheme, "" , "" ), &fs)); |
190 | fs->FlushCaches(); |
191 | } |
192 | return OkStatus(); |
193 | } |
194 | |
195 | Status Env::NewRandomAccessFile(const string& fname, |
196 | std::unique_ptr<RandomAccessFile>* result) { |
197 | FileSystem* fs; |
198 | TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs)); |
199 | return fs->NewRandomAccessFile(fname, result); |
200 | } |
201 | |
202 | Status Env::NewReadOnlyMemoryRegionFromFile( |
203 | const string& fname, std::unique_ptr<ReadOnlyMemoryRegion>* result) { |
204 | FileSystem* fs; |
205 | TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs)); |
206 | return fs->NewReadOnlyMemoryRegionFromFile(fname, result); |
207 | } |
208 | |
209 | Status Env::NewWritableFile(const string& fname, |
210 | std::unique_ptr<WritableFile>* result) { |
211 | FileSystem* fs; |
212 | TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs)); |
213 | return fs->NewWritableFile(fname, result); |
214 | } |
215 | |
216 | Status Env::NewAppendableFile(const string& fname, |
217 | std::unique_ptr<WritableFile>* result) { |
218 | FileSystem* fs; |
219 | TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs)); |
220 | return fs->NewAppendableFile(fname, result); |
221 | } |
222 | |
223 | Status Env::FileExists(const string& fname) { |
224 | FileSystem* fs; |
225 | TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs)); |
226 | return fs->FileExists(fname); |
227 | } |
228 | |
229 | bool Env::FilesExist(const std::vector<string>& files, |
230 | std::vector<Status>* status) { |
231 | std::unordered_map<string, std::vector<string>> files_per_fs; |
232 | for (const auto& file : files) { |
233 | StringPiece scheme, host, path; |
234 | io::ParseURI(file, &scheme, &host, &path); |
235 | files_per_fs[string(scheme)].push_back(file); |
236 | } |
237 | |
238 | std::unordered_map<string, Status> per_file_status; |
239 | bool result = true; |
240 | for (auto itr : files_per_fs) { |
241 | FileSystem* file_system = file_system_registry_->Lookup(itr.first); |
242 | bool fs_result; |
243 | std::vector<Status> local_status; |
244 | std::vector<Status>* fs_status = status ? &local_status : nullptr; |
245 | if (!file_system) { |
246 | fs_result = false; |
247 | if (fs_status) { |
248 | Status s = errors::Unimplemented("File system scheme '" , itr.first, |
249 | "' not implemented" ); |
250 | local_status.resize(itr.second.size(), s); |
251 | } |
252 | } else { |
253 | fs_result = file_system->FilesExist(itr.second, fs_status); |
254 | } |
255 | if (fs_status) { |
256 | result &= fs_result; |
257 | for (size_t i = 0; i < itr.second.size(); ++i) { |
258 | per_file_status[itr.second[i]] = fs_status->at(i); |
259 | } |
260 | } else if (!fs_result) { |
261 | // Return early |
262 | return false; |
263 | } |
264 | } |
265 | |
266 | if (status) { |
267 | for (const auto& file : files) { |
268 | status->push_back(per_file_status[file]); |
269 | } |
270 | } |
271 | |
272 | return result; |
273 | } |
274 | |
275 | Status Env::GetChildren(const string& dir, std::vector<string>* result) { |
276 | FileSystem* fs; |
277 | TF_RETURN_IF_ERROR(GetFileSystemForFile(dir, &fs)); |
278 | return fs->GetChildren(dir, result); |
279 | } |
280 | |
281 | Status Env::GetMatchingPaths(const string& pattern, |
282 | std::vector<string>* results) { |
283 | FileSystem* fs; |
284 | TF_RETURN_IF_ERROR(GetFileSystemForFile(pattern, &fs)); |
285 | return fs->GetMatchingPaths(pattern, results); |
286 | } |
287 | |
288 | Status Env::DeleteFile(const string& fname) { |
289 | FileSystem* fs; |
290 | TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs)); |
291 | return fs->DeleteFile(fname); |
292 | } |
293 | |
294 | Status Env::RecursivelyCreateDir(const string& dirname) { |
295 | FileSystem* fs; |
296 | TF_RETURN_IF_ERROR(GetFileSystemForFile(dirname, &fs)); |
297 | return fs->RecursivelyCreateDir(dirname); |
298 | } |
299 | |
300 | Status Env::CreateDir(const string& dirname) { |
301 | FileSystem* fs; |
302 | TF_RETURN_IF_ERROR(GetFileSystemForFile(dirname, &fs)); |
303 | return fs->CreateDir(dirname); |
304 | } |
305 | |
306 | Status Env::DeleteDir(const string& dirname) { |
307 | FileSystem* fs; |
308 | TF_RETURN_IF_ERROR(GetFileSystemForFile(dirname, &fs)); |
309 | return fs->DeleteDir(dirname); |
310 | } |
311 | |
312 | Status Env::Stat(const string& fname, FileStatistics* stat) { |
313 | FileSystem* fs; |
314 | TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs)); |
315 | return fs->Stat(fname, stat); |
316 | } |
317 | |
318 | Status Env::IsDirectory(const string& fname) { |
319 | FileSystem* fs; |
320 | TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs)); |
321 | return fs->IsDirectory(fname); |
322 | } |
323 | |
324 | Status Env::HasAtomicMove(const string& path, bool* has_atomic_move) { |
325 | FileSystem* fs; |
326 | TF_RETURN_IF_ERROR(GetFileSystemForFile(path, &fs)); |
327 | return fs->HasAtomicMove(path, has_atomic_move); |
328 | } |
329 | |
330 | Status Env::DeleteRecursively(const string& dirname, int64_t* undeleted_files, |
331 | int64_t* undeleted_dirs) { |
332 | FileSystem* fs; |
333 | TF_RETURN_IF_ERROR(GetFileSystemForFile(dirname, &fs)); |
334 | return fs->DeleteRecursively(dirname, undeleted_files, undeleted_dirs); |
335 | } |
336 | |
337 | Status Env::GetFileSize(const string& fname, uint64* file_size) { |
338 | FileSystem* fs; |
339 | TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs)); |
340 | return fs->GetFileSize(fname, file_size); |
341 | } |
342 | |
343 | Status Env::RenameFile(const string& src, const string& target) { |
344 | FileSystem* src_fs; |
345 | FileSystem* target_fs; |
346 | TF_RETURN_IF_ERROR(GetFileSystemForFile(src, &src_fs)); |
347 | TF_RETURN_IF_ERROR(GetFileSystemForFile(target, &target_fs)); |
348 | if (src_fs != target_fs) { |
349 | return errors::Unimplemented("Renaming " , src, " to " , target, |
350 | " not implemented" ); |
351 | } |
352 | return src_fs->RenameFile(src, target); |
353 | } |
354 | |
355 | Status Env::CopyFile(const string& src, const string& target) { |
356 | FileSystem* src_fs; |
357 | FileSystem* target_fs; |
358 | TF_RETURN_IF_ERROR(GetFileSystemForFile(src, &src_fs)); |
359 | TF_RETURN_IF_ERROR(GetFileSystemForFile(target, &target_fs)); |
360 | if (src_fs == target_fs) { |
361 | return src_fs->CopyFile(src, target); |
362 | } |
363 | return FileSystemCopyFile(src_fs, src, target_fs, target); |
364 | } |
365 | |
366 | string Env::GetExecutablePath() { |
367 | char exe_path[PATH_MAX] = {0}; |
368 | #ifdef __APPLE__ |
369 | uint32_t buffer_size(0U); |
370 | _NSGetExecutablePath(nullptr, &buffer_size); |
371 | std::vector<char> unresolved_path(buffer_size); |
372 | _NSGetExecutablePath(unresolved_path.data(), &buffer_size); |
373 | CHECK(realpath(unresolved_path.data(), exe_path)); |
374 | #elif defined(__FreeBSD__) |
375 | int mib[4] = {CTL_KERN, KERN_PROC, KERN_PROC_PATHNAME, -1}; |
376 | size_t exe_path_size = PATH_MAX; |
377 | |
378 | if (sysctl(mib, 4, exe_path, &exe_path_size, NULL, 0) != 0) { |
379 | // Resolution of path failed |
380 | return "" ; |
381 | } |
382 | #elif defined(PLATFORM_WINDOWS) |
383 | HMODULE hModule = GetModuleHandleW(NULL); |
384 | WCHAR wc_file_path[MAX_PATH] = {0}; |
385 | GetModuleFileNameW(hModule, wc_file_path, MAX_PATH); |
386 | string file_path = WideCharToUtf8(wc_file_path); |
387 | std::copy(file_path.begin(), file_path.end(), exe_path); |
388 | #else |
389 | char buf[PATH_MAX] = {0}; |
390 | int path_length = readlink("/proc/self/exe" , buf, sizeof(buf) - 1); |
391 | CHECK_NE(-1, path_length); |
392 | |
393 | if (strstr(buf, "python" ) != nullptr) { |
394 | // Discard the path of the python binary, and any flags. |
395 | int fd = open("/proc/self/cmdline" , O_RDONLY); |
396 | int cmd_length = read(fd, buf, PATH_MAX - 1); |
397 | CHECK_NE(-1, cmd_length); |
398 | int token_pos = 0; |
399 | for (bool token_is_first_or_flag = true; token_is_first_or_flag;) { |
400 | // Get token length, including null |
401 | int token_len = strlen(&buf[token_pos]) + 1; |
402 | token_is_first_or_flag = false; |
403 | // Check if we can skip without overshooting |
404 | if (token_pos + token_len < cmd_length) { |
405 | token_pos += token_len; |
406 | token_is_first_or_flag = (buf[token_pos] == '-'); // token is a flag |
407 | } |
408 | } |
409 | snprintf(exe_path, sizeof(exe_path), "%s" , &buf[token_pos]); |
410 | } else { |
411 | snprintf(exe_path, sizeof(exe_path), "%s" , buf); |
412 | } |
413 | |
414 | #endif |
415 | // Make sure it's null-terminated: |
416 | exe_path[sizeof(exe_path) - 1] = 0; |
417 | |
418 | return exe_path; |
419 | } |
420 | |
421 | bool Env::LocalTempFilename(string* filename) { |
422 | std::vector<string> dirs; |
423 | GetLocalTempDirectories(&dirs); |
424 | |
425 | // Try each directory, as they might be full, have inappropriate |
426 | // permissions or have different problems at times. |
427 | for (const string& dir : dirs) { |
428 | *filename = io::JoinPath(dir, "tempfile-" ); |
429 | if (CreateUniqueFileName(filename, "" )) { |
430 | return true; |
431 | } |
432 | } |
433 | return false; |
434 | } |
435 | |
436 | bool Env::CreateUniqueFileName(string* prefix, const string& suffix) { |
437 | int32_t tid = GetCurrentThreadId(); |
438 | int32_t pid = GetProcessId(); |
439 | long long now_microsec = NowMicros(); // NOLINT |
440 | |
441 | *prefix += strings::Printf("%s-%x-%d-%llx" , port::Hostname().c_str(), tid, |
442 | pid, now_microsec); |
443 | |
444 | if (!suffix.empty()) { |
445 | *prefix += suffix; |
446 | } |
447 | if (FileExists(*prefix).ok()) { |
448 | prefix->clear(); |
449 | return false; |
450 | } else { |
451 | return true; |
452 | } |
453 | } |
454 | |
455 | int32 Env::GetProcessId() { |
456 | #ifdef PLATFORM_WINDOWS |
457 | return static_cast<int32>(GetCurrentProcessId()); |
458 | #else |
459 | return static_cast<int32>(getpid()); |
460 | #endif |
461 | } |
462 | |
463 | Thread::~Thread() {} |
464 | |
465 | EnvWrapper::~EnvWrapper() {} |
466 | |
467 | Status ReadFileToString(Env* env, const string& fname, string* data) { |
468 | uint64 file_size; |
469 | Status s = env->GetFileSize(fname, &file_size); |
470 | if (!s.ok()) { |
471 | return s; |
472 | } |
473 | std::unique_ptr<RandomAccessFile> file; |
474 | s = env->NewRandomAccessFile(fname, &file); |
475 | if (!s.ok()) { |
476 | return s; |
477 | } |
478 | data->resize(file_size); |
479 | char* p = &*data->begin(); |
480 | StringPiece result; |
481 | s = file->Read(0, file_size, &result, p); |
482 | if (!s.ok()) { |
483 | data->clear(); |
484 | } else if (result.size() != file_size) { |
485 | s = errors::Aborted("File " , fname, " changed while reading: " , file_size, |
486 | " vs. " , result.size()); |
487 | data->clear(); |
488 | } else if (result.data() == p) { |
489 | // Data is already in the correct location |
490 | } else { |
491 | memmove(p, result.data(), result.size()); |
492 | } |
493 | return s; |
494 | } |
495 | |
496 | Status WriteStringToFile(Env* env, const string& fname, |
497 | const StringPiece& data) { |
498 | std::unique_ptr<WritableFile> file; |
499 | Status s = env->NewWritableFile(fname, &file); |
500 | if (!s.ok()) { |
501 | return s; |
502 | } |
503 | s = file->Append(data); |
504 | if (s.ok()) { |
505 | s = file->Close(); |
506 | } |
507 | return s; |
508 | } |
509 | |
510 | Status FileSystemCopyFile(FileSystem* src_fs, const string& src, |
511 | FileSystem* target_fs, const string& target) { |
512 | std::unique_ptr<RandomAccessFile> src_file; |
513 | TF_RETURN_IF_ERROR(src_fs->NewRandomAccessFile(src, &src_file)); |
514 | |
515 | // When `target` points to a directory, we need to create a file within. |
516 | string target_name; |
517 | if (target_fs->IsDirectory(target).ok()) { |
518 | target_name = io::JoinPath(target, io::Basename(src)); |
519 | } else { |
520 | target_name = target; |
521 | } |
522 | |
523 | std::unique_ptr<WritableFile> target_file; |
524 | TF_RETURN_IF_ERROR(target_fs->NewWritableFile(target_name, &target_file)); |
525 | |
526 | uint64 offset = 0; |
527 | std::unique_ptr<char[]> scratch(new char[kCopyFileBufferSize]); |
528 | Status s = OkStatus(); |
529 | while (s.ok()) { |
530 | StringPiece result; |
531 | s = src_file->Read(offset, kCopyFileBufferSize, &result, scratch.get()); |
532 | if (!(s.ok() || s.code() == error::OUT_OF_RANGE)) { |
533 | return s; |
534 | } |
535 | TF_RETURN_IF_ERROR(target_file->Append(result)); |
536 | offset += result.size(); |
537 | } |
538 | return target_file->Close(); |
539 | } |
540 | |
541 | // A ZeroCopyInputStream on a RandomAccessFile. |
542 | namespace { |
543 | class FileStream : public protobuf::io::ZeroCopyInputStream { |
544 | public: |
545 | explicit FileStream(RandomAccessFile* file) : file_(file), pos_(0) {} |
546 | |
547 | void BackUp(int count) override { pos_ -= count; } |
548 | bool Skip(int count) override { |
549 | pos_ += count; |
550 | return true; |
551 | } |
552 | int64_t ByteCount() const override { return pos_; } |
553 | Status status() const { return status_; } |
554 | |
555 | bool Next(const void** data, int* size) override { |
556 | StringPiece result; |
557 | Status s = file_->Read(pos_, kBufSize, &result, scratch_); |
558 | if (result.empty()) { |
559 | status_ = s; |
560 | return false; |
561 | } |
562 | pos_ += result.size(); |
563 | *data = result.data(); |
564 | *size = result.size(); |
565 | return true; |
566 | } |
567 | |
568 | private: |
569 | static constexpr int kBufSize = 512 << 10; |
570 | |
571 | RandomAccessFile* file_; |
572 | int64_t pos_; |
573 | Status status_; |
574 | char scratch_[kBufSize]; |
575 | }; |
576 | |
577 | } // namespace |
578 | |
579 | Status WriteBinaryProto(Env* env, const string& fname, |
580 | const protobuf::MessageLite& proto) { |
581 | string serialized; |
582 | proto.AppendToString(&serialized); |
583 | return WriteStringToFile(env, fname, serialized); |
584 | } |
585 | |
586 | Status ReadBinaryProto(Env* env, const string& fname, |
587 | protobuf::MessageLite* proto) { |
588 | std::unique_ptr<RandomAccessFile> file; |
589 | TF_RETURN_IF_ERROR(env->NewRandomAccessFile(fname, &file)); |
590 | std::unique_ptr<FileStream> stream(new FileStream(file.get())); |
591 | protobuf::io::CodedInputStream coded_stream(stream.get()); |
592 | |
593 | if (!proto->ParseFromCodedStream(&coded_stream) || |
594 | !coded_stream.ConsumedEntireMessage()) { |
595 | TF_RETURN_IF_ERROR(stream->status()); |
596 | return errors::DataLoss("Can't parse " , fname, " as binary proto" ); |
597 | } |
598 | return OkStatus(); |
599 | } |
600 | |
601 | Status WriteTextProto(Env* env, const string& fname, |
602 | const protobuf::Message& proto) { |
603 | string serialized; |
604 | if (!protobuf::TextFormat::PrintToString(proto, &serialized)) { |
605 | return errors::FailedPrecondition("Unable to convert proto to text." ); |
606 | } |
607 | return WriteStringToFile(env, fname, serialized); |
608 | } |
609 | |
610 | Status ReadTextProto(Env* env, const string& fname, protobuf::Message* proto) { |
611 | std::unique_ptr<RandomAccessFile> file; |
612 | TF_RETURN_IF_ERROR(env->NewRandomAccessFile(fname, &file)); |
613 | std::unique_ptr<FileStream> stream(new FileStream(file.get())); |
614 | |
615 | if (!protobuf::TextFormat::Parse(stream.get(), proto)) { |
616 | TF_RETURN_IF_ERROR(stream->status()); |
617 | return errors::DataLoss("Can't parse " , fname, " as text proto" ); |
618 | } |
619 | return OkStatus(); |
620 | } |
621 | |
622 | Status ReadTextOrBinaryProto(Env* env, const string& fname, |
623 | protobuf::Message* proto) { |
624 | if (ReadTextProto(env, fname, proto).ok()) { |
625 | return OkStatus(); |
626 | } |
627 | return ReadBinaryProto(env, fname, proto); |
628 | } |
629 | |
630 | Status ReadTextOrBinaryProto(Env* env, const string& fname, |
631 | protobuf::MessageLite* proto) { |
632 | return ReadBinaryProto(env, fname, proto); |
633 | } |
634 | |
635 | } // namespace tsl |
636 | |