1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
|
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/util/memmapped_file_system.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/util/memmapped_file_system_writer.h"
namespace tensorflow {
namespace {
// Names of files in memmapped environment.
constexpr char kTensor1FileName[] = "memmapped_package://t1";
constexpr char kTensor2FileName[] = "memmapped_package://t2";
constexpr char kProtoFileName[] = "memmapped_package://b";
constexpr int kTestGraphDefVersion = 666;
Status CreateMemmappedFileSystemFile(const string& filename, bool corrupted,
Tensor* test_tensor) {
Env* env = Env::Default();
MemmappedFileSystemWriter writer;
TF_RETURN_IF_ERROR(writer.InitializeToFile(env, filename));
// Try to write a tensor and proto.
test::FillFn<float>(test_tensor,
[](int i) { return static_cast<float>(i * i); });
TF_RETURN_IF_ERROR(writer.SaveTensor(*test_tensor, kTensor1FileName));
// Create a proto with some fields.
GraphDef graph_def;
graph_def.mutable_versions()->set_producer(kTestGraphDefVersion);
graph_def.mutable_versions()->set_min_consumer(kTestGraphDefVersion);
TF_RETURN_IF_ERROR(writer.SaveProtobuf(graph_def, kProtoFileName));
// Save a tensor after the proto to check that alignment works.
test::FillFn<float>(test_tensor,
[](int i) { return static_cast<float>(i) * i * i; });
TF_RETURN_IF_ERROR(writer.SaveTensor(*test_tensor, kTensor2FileName));
if (!corrupted) {
// Flush and close the file.
TF_RETURN_IF_ERROR(writer.FlushAndClose());
}
return Status::OK();
}
TEST(MemmappedFileSystemTest, SimpleTest) {
const TensorShape test_tensor_shape = {10, 200};
Tensor test_tensor(DT_FLOAT, test_tensor_shape);
const string dir = testing::TmpDir();
const string filename = io::JoinPath(dir, "memmapped_env_test");
TF_ASSERT_OK(CreateMemmappedFileSystemFile(filename, false, &test_tensor));
// Check that we can memmap the created file.
MemmappedEnv memmapped_env(Env::Default());
TF_ASSERT_OK(memmapped_env.InitializeFromFile(filename));
// Try to load a proto from the file.
GraphDef test_graph_def;
TF_EXPECT_OK(
ReadBinaryProto(&memmapped_env, kProtoFileName, &test_graph_def));
EXPECT_EQ(kTestGraphDefVersion, test_graph_def.versions().producer());
EXPECT_EQ(kTestGraphDefVersion, test_graph_def.versions().min_consumer());
// Check that we can correctly get a tensor memory.
std::unique_ptr<ReadOnlyMemoryRegion> memory_region;
TF_ASSERT_OK(memmapped_env.NewReadOnlyMemoryRegionFromFile(kTensor2FileName,
&memory_region));
// The memory region can be bigger but not less than Tensor size.
ASSERT_GE(memory_region->length(), test_tensor.TotalBytes());
EXPECT_EQ(test_tensor.tensor_data(),
StringPiece(static_cast<const char*>(memory_region->data()),
test_tensor.TotalBytes()));
// Check that GetFileSize works.
uint64 file_size = 0;
TF_ASSERT_OK(memmapped_env.GetFileSize(kTensor2FileName, &file_size));
EXPECT_EQ(test_tensor.TotalBytes(), file_size);
// Check that Stat works.
FileStatistics stat;
TF_ASSERT_OK(memmapped_env.Stat(kTensor2FileName, &stat));
EXPECT_EQ(test_tensor.TotalBytes(), stat.length);
// Check that if file not found correct error message returned.
EXPECT_EQ(
error::NOT_FOUND,
memmapped_env.NewReadOnlyMemoryRegionFromFile("bla-bla", &memory_region)
.code());
// Check FileExists.
TF_EXPECT_OK(memmapped_env.FileExists(kTensor2FileName));
EXPECT_EQ(error::Code::NOT_FOUND,
memmapped_env.FileExists("bla-bla-bla").code());
}
TEST(MemmappedFileSystemTest, NotInitialized) {
MemmappedEnv memmapped_env(Env::Default());
std::unique_ptr<ReadOnlyMemoryRegion> memory_region;
EXPECT_EQ(
error::FAILED_PRECONDITION,
memmapped_env
.NewReadOnlyMemoryRegionFromFile(kTensor1FileName, &memory_region)
.code());
std::unique_ptr<RandomAccessFile> file;
EXPECT_EQ(error::FAILED_PRECONDITION,
memmapped_env.NewRandomAccessFile(kProtoFileName, &file).code());
}
TEST(MemmappedFileSystemTest, Corrupted) {
// Create a corrupted file (it is not closed it properly).
const TensorShape test_tensor_shape = {100, 200};
Tensor test_tensor(DT_FLOAT, test_tensor_shape);
const string dir = testing::TmpDir();
const string filename = io::JoinPath(dir, "memmapped_env_corrupted_test");
TF_ASSERT_OK(CreateMemmappedFileSystemFile(filename, true, &test_tensor));
MemmappedFileSystem memmapped_env;
ASSERT_NE(memmapped_env.InitializeFromFile(Env::Default(), filename),
Status::OK());
}
TEST(MemmappedFileSystemTest, ProxyToDefault) {
MemmappedEnv memmapped_env(Env::Default());
const string dir = testing::TmpDir();
const string filename = io::JoinPath(dir, "test_file");
// Check that we can create write and read ordinary file.
std::unique_ptr<WritableFile> writable_file_temp;
TF_ASSERT_OK(memmapped_env.NewAppendableFile(filename, &writable_file_temp));
// Making sure to clean up after the test finishes.
const auto adh = [&memmapped_env, &filename](WritableFile* f) {
delete f;
TF_CHECK_OK(memmapped_env.DeleteFile(filename));
};
std::unique_ptr<WritableFile, decltype(adh)> writable_file(
writable_file_temp.release(), adh);
const string test_string = "bla-bla-bla";
TF_ASSERT_OK(writable_file->Append(test_string));
TF_ASSERT_OK(writable_file->Close());
uint64 file_length = 0;
TF_EXPECT_OK(memmapped_env.GetFileSize(filename, &file_length));
EXPECT_EQ(test_string.length(), file_length);
FileStatistics stat;
TF_EXPECT_OK(memmapped_env.Stat(filename, &stat));
EXPECT_EQ(test_string.length(), stat.length);
std::unique_ptr<RandomAccessFile> random_access_file;
TF_ASSERT_OK(
memmapped_env.NewRandomAccessFile(filename, &random_access_file));
}
} // namespace
} // namespace tensorflow
|
