/* 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
|
