Skip to content

Commit fe5c6ae

Browse files
kevemantensorflower-gardener
authored andcommitted
Implement a file factory mechanism to handle network file systems.
- Env dispatches to a FileSystem interface - FileSystemFactory is used to look up the correct FileSystem implementation based on the prefix of the filename - Provide a registration mechanism to register different factories Change: 119268846
1 parent 0bc7958 commit fe5c6ae

16 files changed

Lines changed: 1155 additions & 483 deletions

File tree

tensorflow/core/BUILD

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,7 @@ cc_library(
172172
"lib/strings/strcat.h",
173173
"lib/strings/stringprintf.h",
174174
"platform/env.h",
175+
"platform/file_system.h",
175176
"platform/host_info.h", # TODO(josh11b): make internal
176177
"platform/init_main.h",
177178
"platform/logging.h",
@@ -615,6 +616,7 @@ ANDROID_TF_LITE_HDRS = [
615616
"framework/tensor.h",
616617
"platform/default/integral_types.h",
617618
"platform/env.h",
619+
"platform/file_system.h",
618620
"platform/logging.h",
619621
"platform/platform.h",
620622
"platform/types.h",
@@ -882,6 +884,7 @@ filegroup(
882884
"platform/default/protobuf.h",
883885
"platform/default/thread_annotations.h",
884886
"platform/env.h",
887+
"platform/file_system.h",
885888
"platform/host_info.h",
886889
"platform/logging.h",
887890
"platform/macros.h",

tensorflow/core/kernels/immutable_constant_op_test.cc

Lines changed: 17 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -54,18 +54,17 @@ class TestReadOnlyMemoryRegion : public ReadOnlyMemoryRegion {
5454
uint64 length_;
5555
};
5656

57-
// A mock environment class that creates ReadOnlyMemoryRegion from allocated
58-
// memory.
59-
class TestEnvironment : public EnvWrapper {
57+
// A mock file system and environment class that creates ReadOnlyMemoryRegion
58+
// from allocated memory.
59+
class TestFileSystem : public NullFileSystem {
6060
public:
61-
explicit TestEnvironment(Env* env) : EnvWrapper(env) {}
62-
~TestEnvironment() override = default;
61+
~TestFileSystem() override = default;
6362
Status NewReadOnlyMemoryRegionFromFile(
6463
const string& fname, ReadOnlyMemoryRegion** result) override {
6564
float val = 0;
6665
// For the tests create in-memory regions with float values equal to the
6766
// first letter of the region name.
68-
switch (fname.front()) {
67+
switch (GetNameFromURI(fname).front()) {
6968
case '2':
7069
val = 2.0f;
7170
break;
@@ -84,20 +83,23 @@ class TestEnvironment : public EnvWrapper {
8483
}
8584
};
8685

86+
REGISTER_FILE_SYSTEM("test", TestFileSystem);
87+
8788
struct ImmutableConstantOpTest {};
8889

8990
TEST(ImmutableConstantOpTest, Simple) {
9091
const TensorShape kTestTensorShape({4, 1});
9192
const TensorShape kTestTensorShapeT({1, 4});
9293
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
93-
Node* node1 = ops::ImmutableConst(DT_FLOAT, kTestTensorShape, "2", b.opts());
94-
Node* node2 = ops::ImmutableConst(DT_FLOAT, kTestTensorShapeT, "3", b.opts());
94+
Node* node1 =
95+
ops::ImmutableConst(DT_FLOAT, kTestTensorShape, "test://2", b.opts());
96+
Node* node2 =
97+
ops::ImmutableConst(DT_FLOAT, kTestTensorShapeT, "test://3", b.opts());
9598
Node* result = ops::MatMul(node1, node2, b.opts());
9699
GraphDef graph_def;
97100
TF_ASSERT_OK(b.ToGraphDef(&graph_def));
98-
std::unique_ptr<Env> env_ptr(new TestEnvironment(Env::Default()));
99101
SessionOptions session_options;
100-
session_options.env = env_ptr.get();
102+
session_options.env = Env::Default();
101103
session_options.config.mutable_graph_options()
102104
->mutable_optimizer_options()
103105
->set_opt_level(OptimizerOptions_Level_L0);
@@ -120,14 +122,15 @@ TEST(ImmutableConstantOpTest, ExecutionError) {
120122
const TensorShape kBadTensorShape({40, 100});
121123
const TensorShape kTestTensorShapeT({1, 4});
122124
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
123-
Node* node1 = ops::ImmutableConst(DT_FLOAT, kBadTensorShape, "2", b.opts());
124-
Node* node2 = ops::ImmutableConst(DT_FLOAT, kTestTensorShapeT, "3", b.opts());
125+
Node* node1 =
126+
ops::ImmutableConst(DT_FLOAT, kBadTensorShape, "test://2", b.opts());
127+
Node* node2 =
128+
ops::ImmutableConst(DT_FLOAT, kTestTensorShapeT, "test://3", b.opts());
125129
Node* result = ops::MatMul(node1, node2, b.opts());
126130
GraphDef graph_def;
127131
TF_ASSERT_OK(b.ToGraphDef(&graph_def));
128-
std::unique_ptr<Env> env_ptr(new TestEnvironment(Env::Default()));
129132
SessionOptions session_options;
130-
session_options.env = env_ptr.get();
133+
session_options.env = Env::Default();
131134
std::unique_ptr<Session> session(NewSession(session_options));
132135
ASSERT_TRUE(session != nullptr) << "Failed to create session";
133136
TF_ASSERT_OK(session->Create(graph_def)) << "Can't create test graph";

tensorflow/core/platform/env.cc

Lines changed: 86 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,16 +15,100 @@ limitations under the License.
1515

1616
#include "tensorflow/core/platform/env.h"
1717
#include "tensorflow/core/lib/core/errors.h"
18+
#include "tensorflow/core/lib/gtl/map_util.h"
1819
#include "tensorflow/core/lib/gtl/stl_util.h"
1920
#include "tensorflow/core/platform/protobuf.h"
2021

2122
namespace tensorflow {
2223

2324
Env::~Env() {}
2425

25-
RandomAccessFile::~RandomAccessFile() {}
26+
Status Env::GetFileSystemForFile(const string& fname, FileSystem** result) {
27+
string scheme = GetSchemeFromURI(fname);
28+
FileSystem* file_system = GlobalFileSystemRegistry()->Lookup(scheme);
29+
if (!file_system) {
30+
return errors::Unimplemented("File system scheme ", scheme,
31+
" not implemented");
32+
}
33+
*result = file_system;
34+
return Status::OK();
35+
}
36+
37+
Status Env::NewRandomAccessFile(const string& fname,
38+
RandomAccessFile** result) {
39+
FileSystem* fs;
40+
TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
41+
return fs->NewRandomAccessFile(fname, result);
42+
}
43+
44+
Status Env::NewWritableFile(const string& fname, WritableFile** result) {
45+
FileSystem* fs;
46+
TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
47+
return fs->NewWritableFile(fname, result);
48+
}
49+
50+
Status Env::NewAppendableFile(const string& fname, WritableFile** result) {
51+
FileSystem* fs;
52+
TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
53+
return fs->NewAppendableFile(fname, result);
54+
}
55+
56+
Status Env::NewReadOnlyMemoryRegionFromFile(const string& fname,
57+
ReadOnlyMemoryRegion** result) {
58+
FileSystem* fs;
59+
TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
60+
return fs->NewReadOnlyMemoryRegionFromFile(fname, result);
61+
}
62+
63+
bool Env::FileExists(const string& fname) {
64+
FileSystem* fs;
65+
if (!GetFileSystemForFile(fname, &fs).ok()) {
66+
return false;
67+
}
68+
return fs->FileExists(fname);
69+
}
2670

27-
WritableFile::~WritableFile() {}
71+
Status Env::GetChildren(const string& dir, std::vector<string>* result) {
72+
FileSystem* fs;
73+
TF_RETURN_IF_ERROR(GetFileSystemForFile(dir, &fs));
74+
return fs->GetChildren(dir, result);
75+
}
76+
77+
Status Env::DeleteFile(const string& fname) {
78+
FileSystem* fs;
79+
TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
80+
return fs->DeleteFile(fname);
81+
}
82+
83+
Status Env::CreateDir(const string& dirname) {
84+
FileSystem* fs;
85+
TF_RETURN_IF_ERROR(GetFileSystemForFile(dirname, &fs));
86+
return fs->CreateDir(dirname);
87+
}
88+
89+
Status Env::DeleteDir(const string& dirname) {
90+
FileSystem* fs;
91+
TF_RETURN_IF_ERROR(GetFileSystemForFile(dirname, &fs));
92+
return fs->DeleteDir(dirname);
93+
}
94+
95+
Status Env::GetFileSize(const string& fname, uint64* file_size) {
96+
FileSystem* fs;
97+
TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
98+
return fs->GetFileSize(fname, file_size);
99+
}
100+
101+
Status Env::RenameFile(const string& src, const string& target) {
102+
FileSystem* src_fs;
103+
FileSystem* target_fs;
104+
TF_RETURN_IF_ERROR(GetFileSystemForFile(src, &src_fs));
105+
TF_RETURN_IF_ERROR(GetFileSystemForFile(target, &target_fs));
106+
if (src_fs != target_fs) {
107+
return errors::Unimplemented("Renaming ", src, " to ", target,
108+
" not implemented");
109+
}
110+
return src_fs->RenameFile(src, target);
111+
}
28112

29113
Thread::~Thread() {}
30114

0 commit comments

Comments
 (0)