• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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/core/platform/env.h"
17 
18 #include <sys/stat.h>
19 
20 #include <deque>
21 #include <utility>
22 #include <vector>
23 
24 #include "tensorflow/core/platform/env_time.h"
25 #include "tensorflow/core/platform/errors.h"
26 #include "tensorflow/core/platform/host_info.h"
27 #include "tensorflow/core/platform/path.h"
28 #include "tensorflow/core/platform/platform.h"
29 #include "tensorflow/core/platform/protobuf.h"
30 #include "tensorflow/core/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 tensorflow {
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 
Register(const std::string & scheme,FileSystemRegistry::Factory factory)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 
Register(const std::string & scheme,std::unique_ptr<FileSystem> filesystem)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 
Lookup(const std::string & scheme)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 
GetRegisteredFileSystemSchemes(std::vector<std::string> * schemes)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 
Env()110 Env::Env() : file_system_registry_(new FileSystemRegistryImpl) {}
111 
GetFileSystemForFile(const std::string & fname,FileSystem ** result)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 
GetRegisteredFileSystemSchemes(std::vector<std::string> * schemes)129 Status Env::GetRegisteredFileSystemSchemes(std::vector<std::string>* schemes) {
130   return file_system_registry_->GetRegisteredFileSystemSchemes(schemes);
131 }
132 
RegisterFileSystem(const std::string & scheme,FileSystemRegistry::Factory factory)133 Status Env::RegisterFileSystem(const std::string& scheme,
134                                FileSystemRegistry::Factory factory) {
135   return file_system_registry_->Register(scheme, std::move(factory));
136 }
137 
RegisterFileSystem(const std::string & scheme,std::unique_ptr<FileSystem> filesystem)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 
SetOption(const std::string & scheme,const std::string & key,const std::string & value)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 
SetOption(const std::string & scheme,const std::string & key,const std::vector<string> & values)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 
SetOption(const std::string & scheme,const std::string & key,const std::vector<int64_t> & values)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 
SetOption(const std::string & scheme,const std::string & key,const std::vector<double> & values)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 
FlushFileSystemCaches()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 
NewRandomAccessFile(const string & fname,std::unique_ptr<RandomAccessFile> * result)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 
NewReadOnlyMemoryRegionFromFile(const string & fname,std::unique_ptr<ReadOnlyMemoryRegion> * result)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 
NewWritableFile(const string & fname,std::unique_ptr<WritableFile> * result)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 
NewAppendableFile(const string & fname,std::unique_ptr<WritableFile> * result)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 
FileExists(const string & fname)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 
FilesExist(const std::vector<string> & files,std::vector<Status> * status)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 
GetChildren(const string & dir,std::vector<string> * result)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 
GetMatchingPaths(const string & pattern,std::vector<string> * results)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 
DeleteFile(const string & fname)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 
RecursivelyCreateDir(const string & dirname)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 
CreateDir(const string & dirname)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 
DeleteDir(const string & dirname)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 
Stat(const string & fname,FileStatistics * stat)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 
IsDirectory(const string & fname)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 
HasAtomicMove(const string & path,bool * has_atomic_move)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 
DeleteRecursively(const string & dirname,int64_t * undeleted_files,int64_t * undeleted_dirs)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 
GetFileSize(const string & fname,uint64 * file_size)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 
RenameFile(const string & src,const string & target)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 
CopyFile(const string & src,const string & target)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 
GetExecutablePath()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 
LocalTempFilename(string * filename)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 
CreateUniqueFileName(string * prefix,const string & suffix)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 
GetProcessId()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 
~Thread()463 Thread::~Thread() {}
464 
~EnvWrapper()465 EnvWrapper::~EnvWrapper() {}
466 
ReadFileToString(Env * env,const string & fname,string * data)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 
WriteStringToFile(Env * env,const string & fname,const StringPiece & data)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 
FileSystemCopyFile(FileSystem * src_fs,const string & src,FileSystem * target_fs,const string & target)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:
FileStream(RandomAccessFile * file)545   explicit FileStream(RandomAccessFile* file) : file_(file), pos_(0) {}
546 
BackUp(int count)547   void BackUp(int count) override { pos_ -= count; }
Skip(int count)548   bool Skip(int count) override {
549     pos_ += count;
550     return true;
551   }
ByteCount() const552   int64_t ByteCount() const override { return pos_; }
status() const553   Status status() const { return status_; }
554 
Next(const void ** data,int * size)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 
WriteBinaryProto(Env * env,const string & fname,const protobuf::MessageLite & proto)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 
ReadBinaryProto(Env * env,const string & fname,protobuf::MessageLite * proto)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 
WriteTextProto(Env * env,const string & fname,const protobuf::Message & proto)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 
ReadTextProto(Env * env,const string & fname,protobuf::Message * proto)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 
ReadTextOrBinaryProto(Env * env,const string & fname,protobuf::Message * proto)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 
ReadTextOrBinaryProto(Env * env,const string & fname,protobuf::MessageLite * proto)630 Status ReadTextOrBinaryProto(Env* env, const string& fname,
631                              protobuf::MessageLite* proto) {
632   return ReadBinaryProto(env, fname, proto);
633 }
634 
635 }  // namespace tensorflow
636