1/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations 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
51namespace tsl {
52
53// 128KB copy buffer
54constexpr size_t kCopyFileBufferSize = 128 * 1024;
55
56class 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
71Status 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
82Status 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
92FileSystem* 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
101Status 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
110Env::Env() : file_system_registry_(new FileSystemRegistryImpl) {}
111
112Status 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
129Status Env::GetRegisteredFileSystemSchemes(std::vector<std::string>* schemes) {
130 return file_system_registry_->GetRegisteredFileSystemSchemes(schemes);
131}
132
133Status Env::RegisterFileSystem(const std::string& scheme,
134 FileSystemRegistry::Factory factory) {
135 return file_system_registry_->Register(scheme, std::move(factory));
136}
137
138Status Env::RegisterFileSystem(const std::string& scheme,
139 std::unique_ptr<FileSystem> filesystem) {
140 return file_system_registry_->Register(scheme, std::move(filesystem));
141}
142
143Status 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
153Status 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
163Status 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
173Status 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
183Status 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
195Status 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
202Status 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
209Status 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
216Status 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
223Status Env::FileExists(const string& fname) {
224 FileSystem* fs;
225 TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
226 return fs->FileExists(fname);
227}
228
229bool 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
275Status 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
281Status 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
288Status Env::DeleteFile(const string& fname) {
289 FileSystem* fs;
290 TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
291 return fs->DeleteFile(fname);
292}
293
294Status Env::RecursivelyCreateDir(const string& dirname) {
295 FileSystem* fs;
296 TF_RETURN_IF_ERROR(GetFileSystemForFile(dirname, &fs));
297 return fs->RecursivelyCreateDir(dirname);
298}
299
300Status Env::CreateDir(const string& dirname) {
301 FileSystem* fs;
302 TF_RETURN_IF_ERROR(GetFileSystemForFile(dirname, &fs));
303 return fs->CreateDir(dirname);
304}
305
306Status Env::DeleteDir(const string& dirname) {
307 FileSystem* fs;
308 TF_RETURN_IF_ERROR(GetFileSystemForFile(dirname, &fs));
309 return fs->DeleteDir(dirname);
310}
311
312Status 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
318Status Env::IsDirectory(const string& fname) {
319 FileSystem* fs;
320 TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
321 return fs->IsDirectory(fname);
322}
323
324Status 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
330Status 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
337Status 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
343Status 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
355Status 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
366string 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
421bool 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
436bool 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
455int32 Env::GetProcessId() {
456#ifdef PLATFORM_WINDOWS
457 return static_cast<int32>(GetCurrentProcessId());
458#else
459 return static_cast<int32>(getpid());
460#endif
461}
462
463Thread::~Thread() {}
464
465EnvWrapper::~EnvWrapper() {}
466
467Status 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
496Status 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
510Status 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.
542namespace {
543class 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
579Status 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
586Status 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
601Status 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
610Status 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
622Status 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
630Status ReadTextOrBinaryProto(Env* env, const string& fname,
631 protobuf::MessageLite* proto) {
632 return ReadBinaryProto(env, fname, proto);
633}
634
635} // namespace tsl
636