Add memory management to plugin interface.
Since on Windows memory allocated by one shared object must be freed from the same object, we need to provide a mechanism for plugins to provide their own memory management routines. A previous change made core TensorFlow provide an allocator for plugins to use, however it is more cumbersome to get this allocator to be used in every function in the implementation provided by the plugin. Plus, we still needed to also provide a deleter. This approach simplifies the interface and also allow one to alter these routines to debug issues with plugin's memory allocation. Tested manually, both on Windows and in Linux. Part of the work for modular filesystem plugins. For more details, consult the RFC at https://github.com/tensorflow/community/blob/master/rfcs/20190506-filesystem-plugin-modular-tensorflow.md PiperOrigin-RevId: 291408627 Change-Id: I4fc919701e2bb39928453a02dd6019b559edbbb1
This commit is contained in:
parent
0da59c63fa
commit
12d516b9cd
@ -56,7 +56,7 @@ extern "C" {
|
||||
/// Lifetime: The wrapper data structures are owned by core TensorFlow. The data
|
||||
/// pointed to by the `void*` members is always owned by the plugin. The plugin
|
||||
/// will provide functions to call to allocate and deallocate this data (see
|
||||
/// next section) and core TensorFlow ensures to call these at the proper time.
|
||||
/// next sections) and core TensorFlow ensures to call these at the proper time.
|
||||
///
|
||||
/// Plugins will never receive a `TF_*` pointer that is `nullptr`. Core
|
||||
/// TensorFlow will never touch the `void*` wrapped by these structures, except
|
||||
@ -601,6 +601,10 @@ typedef struct TF_FilesystemOps {
|
||||
///
|
||||
/// Plugins must not return `nullptr`. Returning empty strings is allowed.
|
||||
///
|
||||
/// The allocation and freeing of memory must happen via the functions sent to
|
||||
/// core TensorFlow upon registration (see the `TF_FilesystemPluginInfo`
|
||||
/// structure in Section 4).
|
||||
///
|
||||
/// This function will be called by core TensorFlow to clean up all path
|
||||
/// arguments for all other methods in the filesystem API.
|
||||
///
|
||||
@ -618,6 +622,10 @@ typedef struct TF_FilesystemOps {
|
||||
/// In case of error, plugins must set `status` to a value different than
|
||||
/// `TF_OK`, free memory allocated for `entries` and return -1.
|
||||
///
|
||||
/// The allocation and freeing of memory must happen via the functions sent to
|
||||
/// core TensorFlow upon registration (see the `TF_FilesystemPluginInfo`
|
||||
/// structure in Section 4).
|
||||
///
|
||||
/// Plugins:
|
||||
/// * Must set `status` to `TF_OK` if all children were returned.
|
||||
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't point to a
|
||||
@ -654,6 +662,10 @@ typedef struct TF_FilesystemOps {
|
||||
/// different than `TF_OK`, free any memory that might have been allocated for
|
||||
/// `entries` and return -1.
|
||||
///
|
||||
/// The allocation and freeing of memory must happen via the functions sent to
|
||||
/// core TensorFlow upon registration (see the `TF_FilesystemPluginInfo`
|
||||
/// structure in Section 4).
|
||||
///
|
||||
/// Plugins:
|
||||
/// * Must set `status` to `TF_OK` if all matches were returned.
|
||||
/// * Might use any other error value for `status` to signal other errors.
|
||||
@ -741,8 +753,11 @@ constexpr size_t TF_FILESYSTEM_OPS_SIZE = sizeof(TF_FilesystemOps);
|
||||
/// * `TF_InitPlugin` function: must be present in the plugin shared object as
|
||||
/// it will be called by core TensorFlow when the filesystem plugin is
|
||||
/// loaded;
|
||||
/// * `TF_FilesystemPluginInfo` struct: used to transfer information between
|
||||
/// * `TF_FilesystemPluginOps` struct: used to transfer information between
|
||||
/// plugins and core TensorFlow about the operations provided and metadata;
|
||||
/// * `TF_FilesystemPluginInfo` struct: similar to the above structure, but
|
||||
/// collects information about all the file schemes that the plugin provides
|
||||
/// support for, as well as about the plugin's memory handling routines;
|
||||
/// * `TF_SetFilesystemVersionMetadata` function: must be called by plugins in
|
||||
/// their `TF_InitPlugin` to record the versioning information the plugins
|
||||
/// are compiled against.
|
||||
@ -774,7 +789,7 @@ constexpr size_t TF_FILESYSTEM_OPS_SIZE = sizeof(TF_FilesystemOps);
|
||||
/// IMPORTANT: To maintain binary compatibility, the layout of this structure
|
||||
/// must not change! In the unlikely case that a new type of file needs to be
|
||||
/// supported, add the new ops and metadata at the end of the structure.
|
||||
typedef struct TF_FilesystemPluginInfo {
|
||||
typedef struct TF_FilesystemPluginOps {
|
||||
char* scheme;
|
||||
int filesystem_ops_abi;
|
||||
int filesystem_ops_api;
|
||||
@ -792,6 +807,29 @@ typedef struct TF_FilesystemPluginInfo {
|
||||
int read_only_memory_region_ops_api;
|
||||
size_t read_only_memory_region_ops_size;
|
||||
TF_ReadOnlyMemoryRegionOps* read_only_memory_region_ops;
|
||||
} TF_FilesystemPluginOps;
|
||||
|
||||
/// This structure gathers together all the operations provided by the plugin.
|
||||
///
|
||||
/// Plugins must provide exactly `num_schemes` elements in the `ops` array.
|
||||
///
|
||||
/// Since memory that is allocated by the DSO gets transferred to core
|
||||
/// TensorFlow, we need to provide a way for the allocation and deallocation to
|
||||
/// match. This is why this structure also defines `plugin_memory_allocate` and
|
||||
/// `plugin_memory_free` members.
|
||||
///
|
||||
/// All memory allocated by the plugin that will be owned by core TensorFlow
|
||||
/// must be allocated using the allocator in this structure. Core TensorFlow
|
||||
/// will use the deallocator to free this memory once it no longer needs it.
|
||||
///
|
||||
/// IMPORTANT: To maintain binary compatibility, the layout of this structure
|
||||
/// must not change! In the unlikely case that new global operations must be
|
||||
/// provided, add them at the end of the structure.
|
||||
typedef struct TF_FilesystemPluginInfo {
|
||||
size_t num_schemes;
|
||||
TF_FilesystemPluginOps* ops;
|
||||
void* (*plugin_memory_allocate)(size_t size);
|
||||
void (*plugin_memory_free)(void* ptr);
|
||||
} TF_FilesystemPluginInfo;
|
||||
|
||||
/// Convenience function for setting the versioning metadata.
|
||||
@ -801,19 +839,19 @@ typedef struct TF_FilesystemPluginInfo {
|
||||
/// We want this to be defined in the plugin's memory space and we guarantee
|
||||
/// that core TensorFlow will never call this.
|
||||
static inline void TF_SetFilesystemVersionMetadata(
|
||||
TF_FilesystemPluginInfo* info) {
|
||||
info->filesystem_ops_abi = TF_FILESYSTEM_OPS_ABI;
|
||||
info->filesystem_ops_api = TF_FILESYSTEM_OPS_API;
|
||||
info->filesystem_ops_size = TF_FILESYSTEM_OPS_SIZE;
|
||||
info->random_access_file_ops_abi = TF_RANDOM_ACCESS_FILE_OPS_ABI;
|
||||
info->random_access_file_ops_api = TF_RANDOM_ACCESS_FILE_OPS_API;
|
||||
info->random_access_file_ops_size = TF_RANDOM_ACCESS_FILE_OPS_SIZE;
|
||||
info->writable_file_ops_abi = TF_WRITABLE_FILE_OPS_ABI;
|
||||
info->writable_file_ops_api = TF_WRITABLE_FILE_OPS_API;
|
||||
info->writable_file_ops_size = TF_WRITABLE_FILE_OPS_SIZE;
|
||||
info->read_only_memory_region_ops_abi = TF_READ_ONLY_MEMORY_REGION_OPS_ABI;
|
||||
info->read_only_memory_region_ops_api = TF_READ_ONLY_MEMORY_REGION_OPS_API;
|
||||
info->read_only_memory_region_ops_size = TF_READ_ONLY_MEMORY_REGION_OPS_SIZE;
|
||||
TF_FilesystemPluginOps* ops) {
|
||||
ops->filesystem_ops_abi = TF_FILESYSTEM_OPS_ABI;
|
||||
ops->filesystem_ops_api = TF_FILESYSTEM_OPS_API;
|
||||
ops->filesystem_ops_size = TF_FILESYSTEM_OPS_SIZE;
|
||||
ops->random_access_file_ops_abi = TF_RANDOM_ACCESS_FILE_OPS_ABI;
|
||||
ops->random_access_file_ops_api = TF_RANDOM_ACCESS_FILE_OPS_API;
|
||||
ops->random_access_file_ops_size = TF_RANDOM_ACCESS_FILE_OPS_SIZE;
|
||||
ops->writable_file_ops_abi = TF_WRITABLE_FILE_OPS_ABI;
|
||||
ops->writable_file_ops_api = TF_WRITABLE_FILE_OPS_API;
|
||||
ops->writable_file_ops_size = TF_WRITABLE_FILE_OPS_SIZE;
|
||||
ops->read_only_memory_region_ops_abi = TF_READ_ONLY_MEMORY_REGION_OPS_ABI;
|
||||
ops->read_only_memory_region_ops_api = TF_READ_ONLY_MEMORY_REGION_OPS_API;
|
||||
ops->read_only_memory_region_ops_size = TF_READ_ONLY_MEMORY_REGION_OPS_SIZE;
|
||||
}
|
||||
|
||||
/// Initializes a TensorFlow plugin.
|
||||
@ -828,16 +866,14 @@ static inline void TF_SetFilesystemVersionMetadata(
|
||||
/// manage themselves). In both of these cases, core TensorFlow looks for
|
||||
/// the `TF_InitPlugin` symbol and calls this function.
|
||||
///
|
||||
/// All memory allocated by this function must be allocated via the `allocator`
|
||||
/// argument.
|
||||
///
|
||||
/// For every filesystem URI scheme that this plugin supports, the plugin must
|
||||
/// add one `TF_FilesystemPluginInfo` entry in `plugin_info`.
|
||||
/// add one `TF_FilesystemPluginInfo` entry in `plugin_info->ops` and call
|
||||
/// `TF_SetFilesystemVersionMetadata` for that entry.
|
||||
///
|
||||
/// Returns number of entries in `plugin_info` (i.e., number of URI schemes
|
||||
/// supported).
|
||||
TF_CAPI_EXPORT extern int TF_InitPlugin(void* (*allocator)(size_t size),
|
||||
TF_FilesystemPluginInfo** plugin_info);
|
||||
/// Plugins must also initialize `plugin_info->plugin_memory_allocate` and
|
||||
/// `plugin_info->plugin_memory_free` to ensure memory allocated by plugin is
|
||||
/// freed in a compatible way.
|
||||
TF_CAPI_EXPORT extern void TF_InitPlugin(TF_FilesystemPluginInfo* plugin_info);
|
||||
|
||||
#ifdef __cplusplus
|
||||
} // end extern "C"
|
||||
|
@ -164,16 +164,18 @@ Status ModularFileSystem::GetChildren(const std::string& dir,
|
||||
|
||||
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
|
||||
std::string translated_name = TranslateName(dir);
|
||||
char** children;
|
||||
// Note that `children` is allocated by the plugin and freed by core
|
||||
// TensorFlow, so we need to use `plugin_memory_free_` here.
|
||||
char** children = nullptr;
|
||||
const int num_children =
|
||||
ops_->get_children(filesystem_.get(), translated_name.c_str(), &children,
|
||||
plugin_status.get());
|
||||
if (num_children >= 0) {
|
||||
for (int i = 0; i < num_children; i++) {
|
||||
result->push_back(std::string(children[i]));
|
||||
free(children[i]);
|
||||
plugin_memory_free_(children[i]);
|
||||
}
|
||||
free(children);
|
||||
plugin_memory_free_(children);
|
||||
}
|
||||
|
||||
return StatusFromTF_Status(plugin_status.get());
|
||||
@ -185,15 +187,17 @@ Status ModularFileSystem::GetMatchingPaths(const std::string& pattern,
|
||||
return internal::GetMatchingPaths(this, Env::Default(), pattern, result);
|
||||
|
||||
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
|
||||
char** matches;
|
||||
// Note that `matches` is allocated by the plugin and freed by core
|
||||
// TensorFlow, so we need to use `plugin_memory_free_` here.
|
||||
char** matches = nullptr;
|
||||
const int num_matches = ops_->get_matching_paths(
|
||||
filesystem_.get(), pattern.c_str(), &matches, plugin_status.get());
|
||||
if (num_matches >= 0) {
|
||||
for (int i = 0; i < num_matches; i++) {
|
||||
result->push_back(std::string(matches[i]));
|
||||
free(matches[i]);
|
||||
plugin_memory_free_(matches[i]);
|
||||
}
|
||||
free(matches);
|
||||
plugin_memory_free_(matches);
|
||||
}
|
||||
|
||||
return StatusFromTF_Status(plugin_status.get());
|
||||
@ -357,7 +361,8 @@ std::string ModularFileSystem::TranslateName(const std::string& name) const {
|
||||
CHECK(p != nullptr) << "TranslateName(" << name << ") returned nullptr";
|
||||
|
||||
std::string ret(p);
|
||||
free(p);
|
||||
// Since `p` is allocated by plugin, free it using plugin's method.
|
||||
plugin_memory_free_(p);
|
||||
return ret;
|
||||
}
|
||||
|
||||
|
@ -46,12 +46,16 @@ class ModularFileSystem final : public FileSystem {
|
||||
std::unique_ptr<const TF_RandomAccessFileOps> random_access_file_ops,
|
||||
std::unique_ptr<const TF_WritableFileOps> writable_file_ops,
|
||||
std::unique_ptr<const TF_ReadOnlyMemoryRegionOps>
|
||||
read_only_memory_region_ops)
|
||||
read_only_memory_region_ops,
|
||||
std::function<void*(size_t)> plugin_memory_allocate,
|
||||
std::function<void(void*)> plugin_memory_free)
|
||||
: filesystem_(std::move(filesystem)),
|
||||
ops_(std::move(filesystem_ops)),
|
||||
random_access_file_ops_(std::move(random_access_file_ops)),
|
||||
writable_file_ops_(std::move(writable_file_ops)),
|
||||
read_only_memory_region_ops_(std::move(read_only_memory_region_ops)) {}
|
||||
read_only_memory_region_ops_(std::move(read_only_memory_region_ops)),
|
||||
plugin_memory_allocate_(std::move(plugin_memory_allocate)),
|
||||
plugin_memory_free_(std::move(plugin_memory_free)) {}
|
||||
|
||||
~ModularFileSystem() override { ops_->cleanup(filesystem_.get()); }
|
||||
|
||||
@ -93,6 +97,8 @@ class ModularFileSystem final : public FileSystem {
|
||||
std::unique_ptr<const TF_WritableFileOps> writable_file_ops_;
|
||||
std::unique_ptr<const TF_ReadOnlyMemoryRegionOps>
|
||||
read_only_memory_region_ops_;
|
||||
std::function<void*(size_t)> plugin_memory_allocate_;
|
||||
std::function<void(void*)> plugin_memory_free_;
|
||||
TF_DISALLOW_COPY_AND_ASSIGN(ModularFileSystem);
|
||||
};
|
||||
|
||||
|
@ -50,21 +50,21 @@ static Status CheckABI(int pluginABI, int coreABI, StringPiece where) {
|
||||
// If the numbers don't match, plugin cannot be loaded.
|
||||
//
|
||||
// Uses the simpler `CheckABI(int, int, StringPiece)`.
|
||||
static Status ValidateABI(const TF_FilesystemPluginInfo* info) {
|
||||
static Status ValidateABI(const TF_FilesystemPluginOps* ops) {
|
||||
TF_RETURN_IF_ERROR(
|
||||
CheckABI(info->filesystem_ops_abi, TF_FILESYSTEM_OPS_ABI, "filesystem"));
|
||||
CheckABI(ops->filesystem_ops_abi, TF_FILESYSTEM_OPS_ABI, "filesystem"));
|
||||
|
||||
if (info->random_access_file_ops != nullptr)
|
||||
TF_RETURN_IF_ERROR(CheckABI(info->random_access_file_ops_abi,
|
||||
if (ops->random_access_file_ops != nullptr)
|
||||
TF_RETURN_IF_ERROR(CheckABI(ops->random_access_file_ops_abi,
|
||||
TF_RANDOM_ACCESS_FILE_OPS_ABI,
|
||||
"random access file"));
|
||||
|
||||
if (info->writable_file_ops != nullptr)
|
||||
TF_RETURN_IF_ERROR(CheckABI(info->writable_file_ops_abi,
|
||||
if (ops->writable_file_ops != nullptr)
|
||||
TF_RETURN_IF_ERROR(CheckABI(ops->writable_file_ops_abi,
|
||||
TF_WRITABLE_FILE_OPS_ABI, "writable file"));
|
||||
|
||||
if (info->read_only_memory_region_ops != nullptr)
|
||||
TF_RETURN_IF_ERROR(CheckABI(info->read_only_memory_region_ops_abi,
|
||||
if (ops->read_only_memory_region_ops != nullptr)
|
||||
TF_RETURN_IF_ERROR(CheckABI(ops->read_only_memory_region_ops_abi,
|
||||
TF_READ_ONLY_MEMORY_REGION_OPS_ABI,
|
||||
"read only memory region"));
|
||||
|
||||
@ -83,19 +83,19 @@ static void CheckAPI(int plugin_API, int core_API, StringPiece where) {
|
||||
// Checks if the plugin and core API numbers match, for all operations.
|
||||
//
|
||||
// Uses the simpler `CheckAPIHelper(int, int, StringPiece)`.
|
||||
static void ValidateAPI(const TF_FilesystemPluginInfo* info) {
|
||||
CheckAPI(info->filesystem_ops_api, TF_FILESYSTEM_OPS_API, "filesystem");
|
||||
static void ValidateAPI(const TF_FilesystemPluginOps* ops) {
|
||||
CheckAPI(ops->filesystem_ops_api, TF_FILESYSTEM_OPS_API, "filesystem");
|
||||
|
||||
if (info->random_access_file_ops != nullptr)
|
||||
CheckAPI(info->random_access_file_ops_api, TF_RANDOM_ACCESS_FILE_OPS_API,
|
||||
if (ops->random_access_file_ops != nullptr)
|
||||
CheckAPI(ops->random_access_file_ops_api, TF_RANDOM_ACCESS_FILE_OPS_API,
|
||||
"random access file");
|
||||
|
||||
if (info->writable_file_ops != nullptr)
|
||||
CheckAPI(info->writable_file_ops_api, TF_WRITABLE_FILE_OPS_API,
|
||||
if (ops->writable_file_ops != nullptr)
|
||||
CheckAPI(ops->writable_file_ops_api, TF_WRITABLE_FILE_OPS_API,
|
||||
"writable file");
|
||||
|
||||
if (info->read_only_memory_region_ops != nullptr)
|
||||
CheckAPI(info->read_only_memory_region_ops_api,
|
||||
if (ops->read_only_memory_region_ops != nullptr)
|
||||
CheckAPI(ops->read_only_memory_region_ops_api,
|
||||
TF_READ_ONLY_MEMORY_REGION_OPS_API, "read only memory region");
|
||||
}
|
||||
|
||||
@ -177,27 +177,27 @@ static Status ValidateHelper(const TF_ReadOnlyMemoryRegionOps* ops) {
|
||||
// individual function table and then checks that the function table for a
|
||||
// specific file type exists if the plugin offers support for creating that
|
||||
// type of files.
|
||||
static Status ValidateOperations(const TF_FilesystemPluginInfo* info) {
|
||||
TF_RETURN_IF_ERROR(ValidateHelper(info->filesystem_ops));
|
||||
TF_RETURN_IF_ERROR(ValidateHelper(info->random_access_file_ops));
|
||||
TF_RETURN_IF_ERROR(ValidateHelper(info->writable_file_ops));
|
||||
TF_RETURN_IF_ERROR(ValidateHelper(info->read_only_memory_region_ops));
|
||||
static Status ValidateOperations(const TF_FilesystemPluginOps* ops) {
|
||||
TF_RETURN_IF_ERROR(ValidateHelper(ops->filesystem_ops));
|
||||
TF_RETURN_IF_ERROR(ValidateHelper(ops->random_access_file_ops));
|
||||
TF_RETURN_IF_ERROR(ValidateHelper(ops->writable_file_ops));
|
||||
TF_RETURN_IF_ERROR(ValidateHelper(ops->read_only_memory_region_ops));
|
||||
|
||||
if (info->filesystem_ops->new_random_access_file != nullptr &&
|
||||
info->random_access_file_ops == nullptr)
|
||||
if (ops->filesystem_ops->new_random_access_file != nullptr &&
|
||||
ops->random_access_file_ops == nullptr)
|
||||
return errors::FailedPrecondition(
|
||||
"Filesystem allows creation of random access files but no "
|
||||
"operations on them have been supplied.");
|
||||
|
||||
if ((info->filesystem_ops->new_writable_file != nullptr ||
|
||||
info->filesystem_ops->new_appendable_file != nullptr) &&
|
||||
info->writable_file_ops == nullptr)
|
||||
if ((ops->filesystem_ops->new_writable_file != nullptr ||
|
||||
ops->filesystem_ops->new_appendable_file != nullptr) &&
|
||||
ops->writable_file_ops == nullptr)
|
||||
return errors::FailedPrecondition(
|
||||
"Filesystem allows creation of writable files but no "
|
||||
"operations on them have been supplied.");
|
||||
|
||||
if (info->filesystem_ops->new_read_only_memory_region_from_file != nullptr &&
|
||||
info->read_only_memory_region_ops == nullptr)
|
||||
if (ops->filesystem_ops->new_read_only_memory_region_from_file != nullptr &&
|
||||
ops->read_only_memory_region_ops == nullptr)
|
||||
return errors::FailedPrecondition(
|
||||
"Filesystem allows creation of readonly memory regions but no "
|
||||
"operations on them have been supplied.");
|
||||
@ -232,18 +232,23 @@ static std::unique_ptr<const T> CopyToCore(const T* plugin_ops,
|
||||
}
|
||||
|
||||
// Registers one filesystem from the plugin.
|
||||
static Status RegisterFileSystem(const TF_FilesystemPluginInfo* info) {
|
||||
//
|
||||
// Must be called only with `index` a valid index in `info->ops`.
|
||||
static Status RegisterFileSystem(const TF_FilesystemPluginInfo* info,
|
||||
int index) {
|
||||
// Step 1: Copy all the function tables to core TensorFlow memory space
|
||||
auto core_filesystem_ops = CopyToCore<TF_FilesystemOps>(
|
||||
info->filesystem_ops, info->filesystem_ops_size);
|
||||
info->ops[index].filesystem_ops, info->ops[index].filesystem_ops_size);
|
||||
auto core_random_access_file_ops = CopyToCore<TF_RandomAccessFileOps>(
|
||||
info->random_access_file_ops, info->random_access_file_ops_size);
|
||||
auto core_writable_file_ops = CopyToCore<TF_WritableFileOps>(
|
||||
info->writable_file_ops, info->writable_file_ops_size);
|
||||
info->ops[index].random_access_file_ops,
|
||||
info->ops[index].random_access_file_ops_size);
|
||||
auto core_writable_file_ops =
|
||||
CopyToCore<TF_WritableFileOps>(info->ops[index].writable_file_ops,
|
||||
info->ops[index].writable_file_ops_size);
|
||||
auto core_read_only_memory_region_ops =
|
||||
CopyToCore<TF_ReadOnlyMemoryRegionOps>(
|
||||
info->read_only_memory_region_ops,
|
||||
info->read_only_memory_region_ops_size);
|
||||
info->ops[index].read_only_memory_region_ops,
|
||||
info->ops[index].read_only_memory_region_ops_size);
|
||||
|
||||
// Step 2: Initialize the opaque filesystem structure
|
||||
auto filesystem = tensorflow::MakeUnique<TF_Filesystem>();
|
||||
@ -256,32 +261,46 @@ static Status RegisterFileSystem(const TF_FilesystemPluginInfo* info) {
|
||||
|
||||
// Step 3: Actual registration
|
||||
return Env::Default()->RegisterFileSystem(
|
||||
info->scheme, tensorflow::MakeUnique<tensorflow::ModularFileSystem>(
|
||||
std::move(filesystem), std::move(core_filesystem_ops),
|
||||
std::move(core_random_access_file_ops),
|
||||
std::move(core_writable_file_ops),
|
||||
std::move(core_read_only_memory_region_ops)));
|
||||
info->ops[index].scheme,
|
||||
tensorflow::MakeUnique<tensorflow::ModularFileSystem>(
|
||||
std::move(filesystem), std::move(core_filesystem_ops),
|
||||
std::move(core_random_access_file_ops),
|
||||
std::move(core_writable_file_ops),
|
||||
std::move(core_read_only_memory_region_ops),
|
||||
info->plugin_memory_allocate, info->plugin_memory_free));
|
||||
}
|
||||
|
||||
// Registers all filesystems, if plugin is providing valid information.
|
||||
// Registers filesystem at `index`, if plugin is providing valid information.
|
||||
//
|
||||
// Extracted to a separate function so that pointers inside `info` are freed
|
||||
// by the caller regardless of whether validation/registration failed or not.
|
||||
//
|
||||
// Must be called only with `index` a valid index in `info->ops`.
|
||||
static Status ValidateAndRegisterFilesystems(
|
||||
const TF_FilesystemPluginInfo* info) {
|
||||
TF_RETURN_IF_ERROR(ValidateScheme(info->scheme));
|
||||
TF_RETURN_IF_ERROR(ValidateABI(info));
|
||||
ValidateAPI(info); // we just warn on API number mismatch
|
||||
TF_RETURN_IF_ERROR(ValidateOperations(info));
|
||||
TF_RETURN_IF_ERROR(RegisterFileSystem(info));
|
||||
const TF_FilesystemPluginInfo* info, int index) {
|
||||
TF_RETURN_IF_ERROR(ValidateScheme(info->ops[index].scheme));
|
||||
TF_RETURN_IF_ERROR(ValidateABI(&info->ops[index]));
|
||||
ValidateAPI(&info->ops[index]); // we just warn on API number mismatch
|
||||
TF_RETURN_IF_ERROR(ValidateOperations(&info->ops[index]));
|
||||
TF_RETURN_IF_ERROR(RegisterFileSystem(info, index));
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
// Alocates memory in plugin DSO.
|
||||
//
|
||||
// Provided by core TensorFlow so that it can free this memory after DSO is
|
||||
// loaded and filesystem information has been used to register the filesystem.
|
||||
static void* basic_allocator(size_t size) { return calloc(1, size); }
|
||||
// Ensures that the plugin provides the required memory management operations.
|
||||
static Status ValidatePluginMemoryRoutines(
|
||||
const TF_FilesystemPluginInfo* info) {
|
||||
if (info->plugin_memory_allocate == nullptr)
|
||||
return errors::FailedPrecondition(
|
||||
"Cannot load filesystem plugin which does not provide "
|
||||
"`plugin_memory_allocate`");
|
||||
|
||||
if (info->plugin_memory_free == nullptr)
|
||||
return errors::FailedPrecondition(
|
||||
"Cannot load filesystem plugin which does not provide "
|
||||
"`plugin_memory_free`");
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
namespace filesystem_registration {
|
||||
|
||||
@ -297,26 +316,28 @@ Status RegisterFilesystemPluginImpl(const std::string& dso_path) {
|
||||
env->GetSymbolFromLibrary(dso_handle, "TF_InitPlugin", &dso_symbol));
|
||||
|
||||
// Step 3: Call `TF_InitPlugin`
|
||||
TF_FilesystemPluginInfo* info = nullptr;
|
||||
auto TF_InitPlugin = reinterpret_cast<int (*)(
|
||||
decltype(&basic_allocator), TF_FilesystemPluginInfo**)>(dso_symbol);
|
||||
int num_schemes = TF_InitPlugin(&basic_allocator, &info);
|
||||
if (num_schemes < 0 || info == nullptr)
|
||||
return errors::InvalidArgument("DSO returned invalid filesystem data");
|
||||
TF_FilesystemPluginInfo info;
|
||||
memset(&info, 0, sizeof(info));
|
||||
auto TF_InitPlugin =
|
||||
reinterpret_cast<int (*)(TF_FilesystemPluginInfo*)>(dso_symbol);
|
||||
TF_InitPlugin(&info);
|
||||
|
||||
// Step 4: Validate and register all filesystems
|
||||
// Step 4: Ensure plugin provides the memory management functions.
|
||||
TF_RETURN_IF_ERROR(ValidatePluginMemoryRoutines(&info));
|
||||
|
||||
// Step 5: Validate and register all filesystems
|
||||
// Try to register as many filesystems as possible.
|
||||
// Free memory once we no longer need it
|
||||
Status status;
|
||||
for (int i = 0; i < num_schemes; i++) {
|
||||
status.Update(ValidateAndRegisterFilesystems(&info[i]));
|
||||
free(info[i].scheme);
|
||||
free(info[i].filesystem_ops);
|
||||
free(info[i].random_access_file_ops);
|
||||
free(info[i].writable_file_ops);
|
||||
free(info[i].read_only_memory_region_ops);
|
||||
for (int i = 0; i < info.num_schemes; i++) {
|
||||
status.Update(ValidateAndRegisterFilesystems(&info, i));
|
||||
info.plugin_memory_free(info.ops[i].scheme);
|
||||
info.plugin_memory_free(info.ops[i].filesystem_ops);
|
||||
info.plugin_memory_free(info.ops[i].random_access_file_ops);
|
||||
info.plugin_memory_free(info.ops[i].writable_file_ops);
|
||||
info.plugin_memory_free(info.ops[i].read_only_memory_region_ops);
|
||||
}
|
||||
free(info);
|
||||
info.plugin_memory_free(info.ops);
|
||||
return status;
|
||||
}
|
||||
|
||||
|
@ -31,6 +31,9 @@ limitations under the License.
|
||||
// Implementation of a filesystem for POSIX environments.
|
||||
// This filesystem will support `file://` and empty (local) URI schemes.
|
||||
|
||||
static void* plugin_memory_allocate(size_t size) { return calloc(1, size); }
|
||||
static void plugin_memory_free(void* ptr) { free(ptr); }
|
||||
|
||||
// SECTION 1. Implementation for `TF_RandomAccessFile`
|
||||
// ----------------------------------------------------------------------------
|
||||
namespace tf_random_access_file {
|
||||
@ -43,7 +46,9 @@ typedef struct PosixFile {
|
||||
static void Cleanup(TF_RandomAccessFile* file) {
|
||||
auto posix_file = static_cast<PosixFile*>(file->plugin_file);
|
||||
close(posix_file->fd);
|
||||
free(const_cast<char*>(posix_file->filename));
|
||||
// This would be safe to free using `free` directly as it is only opaque.
|
||||
// However, it is better to be consistent everywhere.
|
||||
plugin_memory_free(const_cast<char*>(posix_file->filename));
|
||||
delete posix_file;
|
||||
}
|
||||
|
||||
@ -98,7 +103,7 @@ typedef struct PosixFile {
|
||||
|
||||
static void Cleanup(TF_WritableFile* file) {
|
||||
auto posix_file = static_cast<PosixFile*>(file->plugin_file);
|
||||
free(const_cast<char*>(posix_file->filename));
|
||||
plugin_memory_free(const_cast<char*>(posix_file->filename));
|
||||
delete posix_file;
|
||||
}
|
||||
|
||||
@ -381,12 +386,13 @@ static int GetChildren(const TF_Filesystem* filesystem, const char* path,
|
||||
if (num_entries < 0) {
|
||||
TF_SetStatusFromIOError(status, errno, path);
|
||||
} else {
|
||||
*entries = static_cast<char**>(calloc(num_entries, sizeof((*entries)[0])));
|
||||
*entries = static_cast<char**>(
|
||||
plugin_memory_allocate(num_entries * sizeof((*entries)[0])));
|
||||
for (int i = 0; i < num_entries; i++) {
|
||||
(*entries)[i] = strdup(dir_entries[i]->d_name);
|
||||
free(dir_entries[i]);
|
||||
plugin_memory_free(dir_entries[i]);
|
||||
}
|
||||
free(dir_entries);
|
||||
plugin_memory_free(dir_entries);
|
||||
}
|
||||
|
||||
return num_entries;
|
||||
@ -394,65 +400,59 @@ static int GetChildren(const TF_Filesystem* filesystem, const char* path,
|
||||
|
||||
} // namespace tf_posix_filesystem
|
||||
|
||||
int TF_InitPlugin(void* (*allocator)(size_t), TF_FilesystemPluginInfo** info) {
|
||||
const int num_schemes = 2;
|
||||
*info = static_cast<TF_FilesystemPluginInfo*>(
|
||||
allocator(num_schemes * sizeof((*info)[0])));
|
||||
static void ProvideFilesystemSupportFor(TF_FilesystemPluginOps* ops,
|
||||
const char* uri) {
|
||||
TF_SetFilesystemVersionMetadata(ops);
|
||||
ops->scheme = strdup(uri);
|
||||
|
||||
for (int i = 0; i < num_schemes; i++) {
|
||||
TF_FilesystemPluginInfo* current_info = &((*info)[i]);
|
||||
TF_SetFilesystemVersionMetadata(current_info);
|
||||
ops->random_access_file_ops = static_cast<TF_RandomAccessFileOps*>(
|
||||
plugin_memory_allocate(TF_RANDOM_ACCESS_FILE_OPS_SIZE));
|
||||
ops->random_access_file_ops->cleanup = tf_random_access_file::Cleanup;
|
||||
ops->random_access_file_ops->read = tf_random_access_file::Read;
|
||||
|
||||
current_info->random_access_file_ops = static_cast<TF_RandomAccessFileOps*>(
|
||||
allocator(TF_RANDOM_ACCESS_FILE_OPS_SIZE));
|
||||
current_info->random_access_file_ops->cleanup =
|
||||
tf_random_access_file::Cleanup;
|
||||
current_info->random_access_file_ops->read = tf_random_access_file::Read;
|
||||
ops->writable_file_ops = static_cast<TF_WritableFileOps*>(
|
||||
plugin_memory_allocate(TF_WRITABLE_FILE_OPS_SIZE));
|
||||
ops->writable_file_ops->cleanup = tf_writable_file::Cleanup;
|
||||
ops->writable_file_ops->append = tf_writable_file::Append;
|
||||
ops->writable_file_ops->tell = tf_writable_file::Tell;
|
||||
ops->writable_file_ops->flush = tf_writable_file::Flush;
|
||||
ops->writable_file_ops->sync = tf_writable_file::Sync;
|
||||
ops->writable_file_ops->close = tf_writable_file::Close;
|
||||
|
||||
current_info->writable_file_ops =
|
||||
static_cast<TF_WritableFileOps*>(allocator(TF_WRITABLE_FILE_OPS_SIZE));
|
||||
current_info->writable_file_ops->cleanup = tf_writable_file::Cleanup;
|
||||
current_info->writable_file_ops->append = tf_writable_file::Append;
|
||||
current_info->writable_file_ops->tell = tf_writable_file::Tell;
|
||||
current_info->writable_file_ops->flush = tf_writable_file::Flush;
|
||||
current_info->writable_file_ops->sync = tf_writable_file::Sync;
|
||||
current_info->writable_file_ops->close = tf_writable_file::Close;
|
||||
ops->read_only_memory_region_ops = static_cast<TF_ReadOnlyMemoryRegionOps*>(
|
||||
plugin_memory_allocate(TF_READ_ONLY_MEMORY_REGION_OPS_SIZE));
|
||||
ops->read_only_memory_region_ops->cleanup =
|
||||
tf_read_only_memory_region::Cleanup;
|
||||
ops->read_only_memory_region_ops->data = tf_read_only_memory_region::Data;
|
||||
ops->read_only_memory_region_ops->length = tf_read_only_memory_region::Length;
|
||||
|
||||
current_info->read_only_memory_region_ops =
|
||||
static_cast<TF_ReadOnlyMemoryRegionOps*>(
|
||||
allocator(TF_READ_ONLY_MEMORY_REGION_OPS_SIZE));
|
||||
current_info->read_only_memory_region_ops->cleanup =
|
||||
tf_read_only_memory_region::Cleanup;
|
||||
current_info->read_only_memory_region_ops->data =
|
||||
tf_read_only_memory_region::Data;
|
||||
current_info->read_only_memory_region_ops->length =
|
||||
tf_read_only_memory_region::Length;
|
||||
|
||||
current_info->filesystem_ops =
|
||||
static_cast<TF_FilesystemOps*>(allocator(TF_FILESYSTEM_OPS_SIZE));
|
||||
current_info->filesystem_ops->init = tf_posix_filesystem::Init;
|
||||
current_info->filesystem_ops->cleanup = tf_posix_filesystem::Cleanup;
|
||||
current_info->filesystem_ops->new_random_access_file =
|
||||
tf_posix_filesystem::NewRandomAccessFile;
|
||||
current_info->filesystem_ops->new_writable_file =
|
||||
tf_posix_filesystem::NewWritableFile;
|
||||
current_info->filesystem_ops->new_appendable_file =
|
||||
tf_posix_filesystem::NewAppendableFile;
|
||||
current_info->filesystem_ops->new_read_only_memory_region_from_file =
|
||||
tf_posix_filesystem::NewReadOnlyMemoryRegionFromFile;
|
||||
current_info->filesystem_ops->create_dir = tf_posix_filesystem::CreateDir;
|
||||
current_info->filesystem_ops->delete_file = tf_posix_filesystem::DeleteFile;
|
||||
current_info->filesystem_ops->delete_dir = tf_posix_filesystem::DeleteDir;
|
||||
current_info->filesystem_ops->rename_file = tf_posix_filesystem::RenameFile;
|
||||
current_info->filesystem_ops->copy_file = tf_posix_filesystem::CopyFile;
|
||||
current_info->filesystem_ops->path_exists = tf_posix_filesystem::PathExists;
|
||||
current_info->filesystem_ops->stat = tf_posix_filesystem::Stat;
|
||||
current_info->filesystem_ops->get_children =
|
||||
tf_posix_filesystem::GetChildren;
|
||||
}
|
||||
|
||||
(*info)[0].scheme = strdup("");
|
||||
(*info)[1].scheme = strdup("file");
|
||||
|
||||
return num_schemes;
|
||||
ops->filesystem_ops = static_cast<TF_FilesystemOps*>(
|
||||
plugin_memory_allocate(TF_FILESYSTEM_OPS_SIZE));
|
||||
ops->filesystem_ops->init = tf_posix_filesystem::Init;
|
||||
ops->filesystem_ops->cleanup = tf_posix_filesystem::Cleanup;
|
||||
ops->filesystem_ops->new_random_access_file =
|
||||
tf_posix_filesystem::NewRandomAccessFile;
|
||||
ops->filesystem_ops->new_writable_file = tf_posix_filesystem::NewWritableFile;
|
||||
ops->filesystem_ops->new_appendable_file =
|
||||
tf_posix_filesystem::NewAppendableFile;
|
||||
ops->filesystem_ops->new_read_only_memory_region_from_file =
|
||||
tf_posix_filesystem::NewReadOnlyMemoryRegionFromFile;
|
||||
ops->filesystem_ops->create_dir = tf_posix_filesystem::CreateDir;
|
||||
ops->filesystem_ops->delete_file = tf_posix_filesystem::DeleteFile;
|
||||
ops->filesystem_ops->delete_dir = tf_posix_filesystem::DeleteDir;
|
||||
ops->filesystem_ops->rename_file = tf_posix_filesystem::RenameFile;
|
||||
ops->filesystem_ops->copy_file = tf_posix_filesystem::CopyFile;
|
||||
ops->filesystem_ops->path_exists = tf_posix_filesystem::PathExists;
|
||||
ops->filesystem_ops->stat = tf_posix_filesystem::Stat;
|
||||
ops->filesystem_ops->get_children = tf_posix_filesystem::GetChildren;
|
||||
}
|
||||
|
||||
void TF_InitPlugin(TF_FilesystemPluginInfo* info) {
|
||||
info->plugin_memory_allocate = plugin_memory_allocate;
|
||||
info->plugin_memory_free = plugin_memory_free;
|
||||
info->num_schemes = 2;
|
||||
info->ops = static_cast<TF_FilesystemPluginOps*>(
|
||||
plugin_memory_allocate(info->num_schemes * sizeof(info->ops[0])));
|
||||
ProvideFilesystemSupportFor(&info->ops[0], "");
|
||||
ProvideFilesystemSupportFor(&info->ops[1], "file");
|
||||
}
|
||||
|
@ -21,6 +21,9 @@ limitations under the License.
|
||||
// Implementation of a filesystem for POSIX environments.
|
||||
// This filesystem will support `file://` and empty (local) URI schemes.
|
||||
|
||||
static void* plugin_memory_allocate(size_t size) { return calloc(1, size); }
|
||||
static void plugin_memory_free(void* ptr) { free(ptr); }
|
||||
|
||||
// SECTION 1. Implementation for `TF_RandomAccessFile`
|
||||
// ----------------------------------------------------------------------------
|
||||
namespace tf_random_access_file {
|
||||
@ -53,18 +56,18 @@ namespace tf_windows_filesystem {
|
||||
|
||||
} // namespace tf_windows_filesystem
|
||||
|
||||
int TF_InitPlugin(void* (*allocator)(size_t), TF_FilesystemPluginInfo** info) {
|
||||
const int num_schemes = 2;
|
||||
*info = static_cast<TF_FilesystemPluginInfo*>(
|
||||
allocator(num_schemes * sizeof((*info)[0])));
|
||||
|
||||
for (int i = 0; i < num_schemes; i++) {
|
||||
TF_FilesystemPluginInfo* current_info = &((*info)[i]);
|
||||
TF_SetFilesystemVersionMetadata(current_info);
|
||||
}
|
||||
|
||||
(*info)[0].scheme = strdup("");
|
||||
(*info)[1].scheme = strdup("file");
|
||||
|
||||
return num_schemes;
|
||||
static void ProvideFilesystemSupportFor(TF_FilesystemPluginOps* ops,
|
||||
const char* uri) {
|
||||
TF_SetFilesystemVersionMetadata(ops);
|
||||
ops->scheme = strdup(uri);
|
||||
}
|
||||
|
||||
void TF_InitPlugin(TF_FilesystemPluginInfo* info) {
|
||||
info->plugin_memory_allocate = plugin_memory_allocate;
|
||||
info->plugin_memory_free = plugin_memory_free;
|
||||
info->num_schemes = 2;
|
||||
info->ops = static_cast<TF_FilesystemPluginOps*>(
|
||||
plugin_memory_allocate(info->num_schemes * sizeof(info->ops[0])));
|
||||
ProvideFilesystemSupportFor(&info->ops[0], "");
|
||||
ProvideFilesystemSupportFor(&info->ops[1], "file");
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user