/* 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/lib/core/errors.h"
|
#include "tensorflow/core/lib/strings/str_util.h"
|
#include "tensorflow/core/platform/protobuf.h"
|
#include "tensorflow/core/util/memmapped_file_system.pb.h"
|
|
namespace tensorflow {
|
|
namespace {
|
|
uint64 DecodeUint64LittleEndian(const uint8* buffer) {
|
uint64 result = 0;
|
for (int i = 0; i < static_cast<int>(sizeof(uint64)); ++i) {
|
result |= static_cast<uint64>(buffer[i]) << (8 * i);
|
}
|
return result;
|
}
|
|
} // namespace
|
|
namespace {
|
|
class ReadOnlyMemoryRegionFromMemmapped : public ReadOnlyMemoryRegion {
|
public:
|
ReadOnlyMemoryRegionFromMemmapped(const void* data, uint64 length)
|
: data_(data), length_(length) {}
|
~ReadOnlyMemoryRegionFromMemmapped() override = default;
|
const void* data() override { return data_; }
|
uint64 length() override { return length_; }
|
|
private:
|
const void* const data_;
|
const uint64 length_;
|
// intentionally copyable
|
};
|
|
class RandomAccessFileFromMemmapped : public RandomAccessFile {
|
public:
|
RandomAccessFileFromMemmapped(const void* data, uint64 length)
|
: data_(data), length_(length) {}
|
|
~RandomAccessFileFromMemmapped() override = default;
|
|
Status Name(StringPiece* result) const override {
|
return errors::Unimplemented(
|
"RandomAccessFileFromMemmapped does not support Name()");
|
}
|
|
Status Read(uint64 offset, size_t to_read, StringPiece* result,
|
char* scratch) const override {
|
if (offset >= length_) {
|
*result = StringPiece(scratch, 0);
|
return Status(error::OUT_OF_RANGE, "Read after file end");
|
}
|
const uint64 region_left =
|
std::min(length_ - offset, static_cast<uint64>(to_read));
|
*result =
|
StringPiece(reinterpret_cast<const char*>(data_) + offset, region_left);
|
return (region_left == to_read)
|
? Status::OK()
|
: Status(error::OUT_OF_RANGE, "Read less bytes than requested");
|
}
|
|
private:
|
const void* const data_;
|
const uint64 length_;
|
// intentionally copyable
|
};
|
|
} // namespace
|
|
MemmappedFileSystem::MemmappedFileSystem() {}
|
|
Status MemmappedFileSystem::FileExists(const string& fname) {
|
if (!mapped_memory_) {
|
return errors::FailedPrecondition("MemmappedEnv is not initialized");
|
}
|
const auto dir_element = directory_.find(fname);
|
if (dir_element != directory_.end()) {
|
return Status::OK();
|
}
|
return errors::NotFound(fname, " not found");
|
}
|
|
Status MemmappedFileSystem::NewRandomAccessFile(
|
const string& filename, std::unique_ptr<RandomAccessFile>* result) {
|
if (!mapped_memory_) {
|
return errors::FailedPrecondition("MemmappedEnv is not initialized");
|
}
|
const auto dir_element = directory_.find(filename);
|
if (dir_element == directory_.end()) {
|
return errors::NotFound("Region ", filename, " is not found");
|
}
|
result->reset(new RandomAccessFileFromMemmapped(
|
GetMemoryWithOffset(dir_element->second.offset),
|
dir_element->second.length));
|
return Status::OK();
|
}
|
|
Status MemmappedFileSystem::NewReadOnlyMemoryRegionFromFile(
|
const string& filename, std::unique_ptr<ReadOnlyMemoryRegion>* result) {
|
if (!mapped_memory_) {
|
return errors::FailedPrecondition("MemmappedEnv is not initialized");
|
}
|
const auto dir_element = directory_.find(filename);
|
if (dir_element == directory_.end()) {
|
return errors::NotFound("Region ", filename, " is not found");
|
}
|
result->reset(new ReadOnlyMemoryRegionFromMemmapped(
|
GetMemoryWithOffset(dir_element->second.offset),
|
dir_element->second.length));
|
return Status::OK();
|
}
|
|
Status MemmappedFileSystem::GetFileSize(const string& filename, uint64* size) {
|
if (!mapped_memory_) {
|
return errors::FailedPrecondition("MemmappedEnv is not initialized");
|
}
|
const auto dir_element = directory_.find(filename);
|
if (dir_element == directory_.end()) {
|
return errors::NotFound("Region ", filename, " is not found");
|
}
|
*size = dir_element->second.length;
|
return Status::OK();
|
}
|
|
Status MemmappedFileSystem::Stat(const string& fname, FileStatistics* stat) {
|
uint64 size;
|
auto status = GetFileSize(fname, &size);
|
if (status.ok()) {
|
stat->length = size;
|
}
|
return status;
|
}
|
|
Status MemmappedFileSystem::NewWritableFile(const string& filename,
|
std::unique_ptr<WritableFile>* wf) {
|
return errors::Unimplemented("memmapped format doesn't support writing");
|
}
|
|
Status MemmappedFileSystem::NewAppendableFile(
|
const string& filename, std::unique_ptr<WritableFile>* result) {
|
return errors::Unimplemented("memmapped format doesn't support writing");
|
}
|
|
Status MemmappedFileSystem::GetChildren(const string& filename,
|
std::vector<string>* strings) {
|
return errors::Unimplemented("memmapped format doesn't support GetChildren");
|
}
|
|
Status MemmappedFileSystem::GetMatchingPaths(const string& pattern,
|
std::vector<string>* results) {
|
return errors::Unimplemented(
|
"memmapped format doesn't support GetMatchingPaths");
|
}
|
|
Status MemmappedFileSystem::DeleteFile(const string& filename) {
|
return errors::Unimplemented("memmapped format doesn't support DeleteFile");
|
}
|
|
Status MemmappedFileSystem::CreateDir(const string& dirname) {
|
return errors::Unimplemented("memmapped format doesn't support CreateDir");
|
}
|
|
Status MemmappedFileSystem::DeleteDir(const string& dirname) {
|
return errors::Unimplemented("memmapped format doesn't support DeleteDir");
|
}
|
|
Status MemmappedFileSystem::RenameFile(const string& filename_from,
|
const string& filename_to) {
|
return errors::Unimplemented("memmapped format doesn't support RenameFile");
|
}
|
|
const void* MemmappedFileSystem::GetMemoryWithOffset(uint64 offset) const {
|
return reinterpret_cast<const uint8*>(mapped_memory_->data()) + offset;
|
}
|
|
#if defined(_MSC_VER)
|
constexpr char* MemmappedFileSystem::kMemmappedPackagePrefix;
|
constexpr char* MemmappedFileSystem::kMemmappedPackageDefaultGraphDef;
|
#else
|
constexpr char MemmappedFileSystem::kMemmappedPackagePrefix[];
|
constexpr char MemmappedFileSystem::kMemmappedPackageDefaultGraphDef[];
|
#endif
|
|
Status MemmappedFileSystem::InitializeFromFile(Env* env,
|
const string& filename) {
|
TF_RETURN_IF_ERROR(
|
env->NewReadOnlyMemoryRegionFromFile(filename, &mapped_memory_));
|
directory_.clear();
|
if (mapped_memory_->length() <= sizeof(uint64)) {
|
return errors::DataLoss("Corrupted memmapped model file: ", filename,
|
" Invalid package size");
|
}
|
const auto memory_start =
|
reinterpret_cast<const uint8*>(mapped_memory_->data());
|
const uint64 directory_offset = DecodeUint64LittleEndian(
|
memory_start + mapped_memory_->length() - sizeof(uint64));
|
if (directory_offset > mapped_memory_->length() - sizeof(uint64)) {
|
return errors::DataLoss("Corrupted memmapped model file: ", filename,
|
" Invalid directory offset");
|
}
|
MemmappedFileSystemDirectory proto_directory;
|
if (!ParseProtoUnlimited(
|
&proto_directory, memory_start + directory_offset,
|
mapped_memory_->length() - directory_offset - sizeof(uint64))) {
|
return errors::DataLoss("Corrupted memmapped model file: ", filename,
|
" Can't parse its internal directory");
|
}
|
|
// Iterating in reverse order to get lengths of elements;
|
uint64 prev_element_offset = directory_offset;
|
for (auto element_iter = proto_directory.element().rbegin();
|
element_iter != proto_directory.element().rend(); ++element_iter) {
|
// Check that the element offset is in the right range.
|
if (element_iter->offset() >= prev_element_offset) {
|
return errors::DataLoss("Corrupted memmapped model file: ", filename,
|
" Invalid offset of internal component");
|
}
|
if (!directory_
|
.insert(std::make_pair(
|
element_iter->name(),
|
FileRegion(element_iter->offset(),
|
prev_element_offset - element_iter->offset())))
|
.second) {
|
return errors::DataLoss("Corrupted memmapped model file: ", filename,
|
" Duplicate name of internal component ",
|
element_iter->name());
|
}
|
prev_element_offset = element_iter->offset();
|
}
|
return Status::OK();
|
}
|
|
bool MemmappedFileSystem::IsMemmappedPackageFilename(const string& filename) {
|
return str_util::StartsWith(filename, kMemmappedPackagePrefix);
|
}
|
|
namespace {
|
bool IsValidRegionChar(char c) {
|
return (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') ||
|
(c >= '0' && c <= '9') || c == '_' || c == '.';
|
}
|
} // namespace
|
|
bool MemmappedFileSystem::IsWellFormedMemmappedPackageFilename(
|
const string& filename) {
|
if (!IsMemmappedPackageFilename(filename)) {
|
return false;
|
}
|
for (char c :
|
filename.substr(strlen(kMemmappedPackagePrefix),
|
filename.length() - strlen(kMemmappedPackagePrefix))) {
|
if (!IsValidRegionChar(c)) {
|
return false;
|
}
|
}
|
return true;
|
}
|
|
MemmappedEnv::MemmappedEnv(Env* env) : EnvWrapper(env) {}
|
|
Status MemmappedEnv::GetFileSystemForFile(const string& fname,
|
FileSystem** result) {
|
if (MemmappedFileSystem::IsMemmappedPackageFilename(fname)) {
|
if (!memmapped_file_system_) {
|
return errors::FailedPrecondition(
|
"MemmappedEnv is not initialized from a file.");
|
}
|
*result = memmapped_file_system_.get();
|
return Status::OK();
|
}
|
return EnvWrapper::GetFileSystemForFile(fname, result);
|
}
|
|
Status MemmappedEnv::GetRegisteredFileSystemSchemes(
|
std::vector<string>* schemes) {
|
const auto status = EnvWrapper::GetRegisteredFileSystemSchemes(schemes);
|
if (status.ok()) {
|
schemes->emplace_back(MemmappedFileSystem::kMemmappedPackagePrefix);
|
}
|
return status;
|
}
|
|
Status MemmappedEnv::InitializeFromFile(const string& package_filename) {
|
std::unique_ptr<MemmappedFileSystem> file_system_ptr(new MemmappedFileSystem);
|
const auto status =
|
file_system_ptr->InitializeFromFile(target(), package_filename);
|
if (status.ok()) {
|
memmapped_file_system_ = std::move(file_system_ptr);
|
}
|
return status;
|
}
|
|
} // namespace tensorflow
|
