wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 5 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/c/experimental/filesystem/BUILD
+77
View File
@@ -0,0 +1,77 @@
# Experimental filesystem C APIs for TensorFlow.
# Will be moved in proper place once all filesystems are converted to the
# modular framework.
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
# This is only for plugins
cc_library(
name = "filesystem_interface",
hdrs = ["filesystem_interface.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/c:tf_file_statistics",
"//tensorflow/c:tf_status",
],
)
# Core TensorFlow depends on this, will be included in main library
cc_library(
name = "modular_filesystem",
srcs = [
"modular_filesystem.cc",
"modular_filesystem_registration.cc",
],
hdrs = [
"modular_filesystem.h",
"modular_filesystem_registration.h",
],
# TODO(b/139060984): Visibility should be more restrictive once we
# convert to modular filesystems everywhere
visibility = ["//visibility:public"],
deps = [
":filesystem_interface",
"//tensorflow/c:tf_file_statistics",
"//tensorflow/c:tf_status_helper",
"//tensorflow/c:tf_status_internal",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core/platform:env",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:file_statistics",
"//tensorflow/core/platform:status",
"//tensorflow/core/platform:strcat",
"//tensorflow/core/platform:stringpiece",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
],
)
# Compliance test for modules and for interface
tf_cc_test(
name = "modular_filesystem_test",
size = "small",
srcs = ["modular_filesystem_test.cc"],
linkopts = ["-ldl"],
tags = [
"manual", # Requires DSOs as arguments, eventual setup
"notap", # b/139060984, requires implementing modular support for Google filesystem
],
deps = [
":modular_filesystem",
"//tensorflow/core:framework_internal",
"//tensorflow/core/lib/io:path",
"//tensorflow/core/platform:env",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:stacktrace_handler",
"//tensorflow/core/platform:test",
],
)
tensorflow/c/experimental/filesystem/filesystem_interface.h
+998
View File
@@ -0,0 +1,998 @@
/* Copyright 2019 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_FILESYSTEM_INTERFACE_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_FILESYSTEM_INTERFACE_H_
#include <stddef.h>
#include <stdint.h>
#include "tensorflow/c/tf_file_statistics.h"
#include "tensorflow/c/tf_status.h"
/// This is the interop header between core TensorFlow and modular filesystem
/// plugins (see initial RFC https://github.com/tensorflow/community/pull/101).
///
/// Both core TensorFlow and every plugin will use this header. The associated
/// `.cc` file is only used by core TensorFlow to implement checking needed for
/// plugin registration and ensuring API and ABI compatibility. Plugin authors
/// don't need to read the `.cc` file but they should consult every section of
/// this file to ensure a compliant plugin can be built and that the plugin can
/// be used without recompilation in the widest range of TensorFlow versions.
///
/// The header is divided into sections, as follows:
/// 1. Opaque plugin private data structures and wrappers for type safety;
/// 2. Function tables for plugin functionality;
/// 3. Versioning metadata;
/// 4. Plugin registration API and the DSO entry point.
#ifdef __cplusplus
extern "C" {
#endif // __cplusplus
/// SECTION 1. Opaque data structures to hold plugin specific data
/// ----------------------------------------------------------------------------
///
/// The following data structures incorporate a `void*` that is opaque to
/// TensorFlow but can be used by each filesystem plugin to represent internal
/// data.
///
/// We prefer to have these structures instead of passing `void*` into
/// method signatures to have some type of type safety: for example, operations
/// that are only valid on random access files have a `TF_RandomAccessFile`
/// argument.
///
/// 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 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
/// to initialize it as `nullptr`.
typedef struct TF_RandomAccessFile {
void* plugin_file;
} TF_RandomAccessFile;
typedef struct TF_WritableFile {
void* plugin_file;
} TF_WritableFile;
typedef struct TF_ReadOnlyMemoryRegion {
void* plugin_memory_region;
} TF_ReadOnlyMemoryRegion;
typedef struct TF_Filesystem {
void* plugin_filesystem;
} TF_Filesystem;
// The named union is needed here (as opposed to
// inside the `TF_Filesystem_Option_Value` struct)
// as MSVC does not recognize `typeof`.
typedef union TF_Filesystem_Option_Value_Union {
int64_t int_val;
double real_val;
struct {
char* buf;
int buf_length;
} buffer_val;
} TF_Filesystem_Option_Value_Union;
typedef struct TF_Filesystem_Option_Value {
int type_tag; // type of values in the values union
int num_values; // number of values
TF_Filesystem_Option_Value_Union*
values; // owned (plugins must make a copy if storing this)
} TF_Filesystem_Option_Value;
typedef enum TF_Filesystem_Option_Type {
TF_Filesystem_Option_Type_Int = 0,
TF_Filesystem_Option_Type_Real,
TF_Filesystem_Option_Type_Buffer,
TF_Filesystem_Num_Option_Types, // must always be the last item
} TF_Filesystem_Option_Type;
typedef struct TF_Filesystem_Option {
char* name; // null terminated, owned
char* description; // null terminated, owned
int per_file; // bool actually, but bool is not a C type
TF_Filesystem_Option_Value* value; // owned
} TF_Filesystem_Option;
/// SECTION 2. Function tables for functionality provided by plugins
/// ----------------------------------------------------------------------------
///
/// The following data structures represent the function tables for operations
/// that plugins provide (some are mandatory, some are optional, with or without
/// a default implementation).
///
/// Each plugin implements the operations that are supported and TensorFlow will
/// properly handle the cases when an operation is not supported (i.e., return
/// the corresponding `Status` value).
///
/// REQUIRED OPERATIONS: All required operations are marked as such, including
/// operations which are conditionally required. If the presence of an operation
/// `foo` requires operation `bar` to be present, this is specified in `foo`. If
/// the entire set of operations in a table is not provided, use `nullptr` for
/// the struct pointer (e.g., when a file type is not supported).
///
/// DEFAULT IMPLEMENTATIONS: Some operations have default implementations that
/// TensorFlow uses in case the plugin doesn't supply its own version. An
/// operation `foo` might have a default implementation which uses `bar` and
/// `foobar`. If the plugin supplies `bar` and `foobar`, TensorFlow can use the
/// default implementation of `foo`.
///
/// During plugin loading, plugins will call the registration function provided
/// by this interface, supplying values for each of these structures. Core
/// TensorFlow checks that the plugin supplies all mandatory operations and
/// then copies these tables to a different memory location, marking the new
/// operation tables as read-only. Once a plugin is loaded, none of these
/// operation pointers may change.
///
/// There are 4 function tables: one for each of the 3 file objects in
/// TensorFlow (i.e., `RandomAccessFile`, `WritableFile`,
/// `ReadOnlyMemoryRegion`) and one for all the operations a `Filesystem`
/// implements. Each of them is in a 1-to-1 correspondence with the wrapper
/// structures from the first section: these tables only contain function
/// pointers that operate on the corresponding data. Thus, the first argument of
/// each of these functions is a pointer to the paired struct and this argument
/// can be used to track state in between calls (from an object oriented point
/// of view, this can be viewed as a "vtable" for a "class" -- that is the
/// corresponding struct above --; the first argument is in place of `this`).
///
/// Except where noted otherwise, all pointer arguments are owned by core
/// TensorFlow and are guaranteed to not be `nullptr`.
///
/// All path-like arguments are null terminated `char*` strings. Plugins can
/// assume that before any function using path arguments is invoked, the path is
/// made canonical by calling the function provided by `translate_name` or a
/// default implementation of that (supplied by core TensorFlow).
///
/// The only time the pointer to the `TF_*` structures from section 1 is not
/// marked `const` in these functions is when these function are either
/// allocating or deallocating the plugin specific data. That is, in the 4
/// `cleanup` functions (one for each data structure), the `init` function for
/// `TF_Filesystem` and the `new_*` methods of `TF_FilesystemOps` to initialize
/// the 3 types of files. In all other cases, there is no need to modify the
/// address of the opaque data pointer, hence the wrapper pointer is marked
/// `const`.
///
/// For consistency, the arguments on all these functions follow the same
/// pattern: first we have the opaque pointer argument ("this" above), then the
/// input arguments, then the in-out arguments (if any) and we finish the
/// argument list with the out arguments. We only use the return type for an out
/// parameter if that is a plain C type, as this ensures ABI compatibility
/// (returning structures has issues in case compiler options affect
/// optimizations such as RVO). If a status needs to be returned from these
/// methods, the last argument is always a `TF_Status *` (or an array of such
/// pointers) owned by core TensorFlow and guaranteed to not be `nullptr`.
///
/// To ensure ABI and API compatibility, we have out-of-bounds data that is used
/// by both core TensorFlow and the plugin at load time. We don't include this
/// data in the structures here to prevent cases when padding/packing enabled by
/// different compiler options breaks compatibility. For more details about how
/// this is used, please consult next sections. Here we just wrap these tables
/// in lint warnings so that changes here cause changes to the versioning data
/// as well. Here is a short summary of what changes are allowed:
/// * adding a new method at the end of a table is allowed at any time;
/// * any other change to these tables is only allowed on a major TensorFlow
/// version change (e.g., from 2.x to 3.0). This is provided as an escape
/// hatch to allow cleaning up these tables. Since any of these changes
/// break ABI compatibility and cause all plugins to be recompiled, these
/// type of changes should be extremely rare.
///
/// Next section will detail this as well as some corner cases that are out of
/// scope for now.
// LINT.IfChange
typedef struct TF_RandomAccessFileOps {
/// Releases resources associated with `*file`.
///
/// Requires that `*file` is not used in any concurrent or subsequent
/// operations.
///
/// This operation must be provided. See "REQUIRED OPERATIONS" above.
void (*cleanup)(TF_RandomAccessFile* file);
/// Reads up to `n` bytes from `*file` starting at `offset`.
///
/// The output is in `buffer`, core TensorFlow owns the buffer and guarantees
/// that at least `n` bytes are available.
///
/// Returns number of bytes read or -1 in case of error. Because of this
/// constraint and the fact that `ssize_t` is not defined in `stdint.h`/C++
/// standard, the return type is `int64_t`.
///
/// This is thread safe.
///
/// Note: the `buffer` argument is NOT a null terminated string!
///
/// Plugins:
/// * Must set `status` to `TF_OK` if exactly `n` bytes have been read.
/// * Must set `status` to `TF_OUT_OF_RANGE` if fewer than `n` bytes have
/// been read due to EOF.
/// * Must return -1 for any other error and must set `status` to any
/// other value to provide more information about the error.
int64_t (*read)(const TF_RandomAccessFile* file, uint64_t offset, size_t n,
char* buffer, TF_Status* status);
} TF_RandomAccessFileOps;
// LINT.ThenChange(:random_access_file_ops_version)
// LINT.IfChange
typedef struct TF_WritableFileOps {
/// Releases resources associated with `*file`.
///
/// Requires that `*file` is not used in any concurrent or subsequent
/// operations.
///
/// This operation must be provided. See "REQUIRED OPERATIONS" above.
void (*cleanup)(TF_WritableFile* file);
/// Appends `buffer` of size `n` to `*file`.
///
/// Core TensorFlow owns `buffer` and guarantees at least `n` bytes of storage
/// that can be used to write data.
///
/// Note: the `buffer` argument is NOT a null terminated string!
///
/// Plugins:
/// * Must set `status` to `TF_OK` if exactly `n` bytes have been written.
/// * Must set `status` to `TF_RESOURCE_EXHAUSTED` if fewer than `n` bytes
/// have been written, potentially due to quota/disk space.
/// * Might use any other error value for `status` to signal other errors.
void (*append)(const TF_WritableFile* file, const char* buffer, size_t n,
TF_Status* status);
/// Returns the current write position in `*file`.
///
/// Plugins should ensure that the implementation is idempotent, 2 identical
/// calls result in the same answer.
///
/// Plugins:
/// * Must set `status` to `TF_OK` and return current position if no error.
/// * Must set `status` to any other value and return -1 in case of error.
int64_t (*tell)(const TF_WritableFile* file, TF_Status* status);
/// Flushes `*file` and syncs contents to filesystem.
///
/// This call might not block, and when it returns the contents might not have
/// been fully persisted.
///
/// DEFAULT IMPLEMENTATION: No op.
void (*flush)(const TF_WritableFile* file, TF_Status* status);
/// Syncs contents of `*file` with the filesystem.
///
/// This call should block until filesystem confirms that all buffers have
/// been flushed and persisted.
///
/// DEFAULT IMPLEMENTATION: No op.
void (*sync)(const TF_WritableFile* file, TF_Status* status);
/// Closes `*file`.
///
/// Flushes all buffers and deallocates all resources.
///
/// Calling `close` must not result in calling `cleanup`.
///
/// Core TensorFlow will never call `close` twice.
void (*close)(const TF_WritableFile* file, TF_Status* status);
} TF_WritableFileOps;
// LINT.ThenChange(:writable_file_ops_version)
// LINT.IfChange
typedef struct TF_ReadOnlyMemoryRegionOps {
/// Releases resources associated with `*region`.
///
/// Requires that `*region` is not used in any concurrent or subsequent
/// operations.
///
/// This operation must be provided. See "REQUIRED OPERATIONS" above.
void (*cleanup)(TF_ReadOnlyMemoryRegion* region);
/// Returns a pointer to the memory region.
///
/// This operation must be provided. See "REQUIRED OPERATIONS" above.
const void* (*data)(const TF_ReadOnlyMemoryRegion* region);
/// Returns the length of the memory region in bytes.
///
/// This operation must be provided. See "REQUIRED OPERATIONS" above.
uint64_t (*length)(const TF_ReadOnlyMemoryRegion* region);
} TF_ReadOnlyMemoryRegionOps;
// LINT.ThenChange(:read_only_memory_region_ops_version)
// LINT.IfChange
typedef struct TF_FilesystemOps {
/// Acquires all resources used by the filesystem.
///
/// This operation must be provided. See "REQUIRED OPERATIONS" above.
void (*init)(TF_Filesystem* filesystem, TF_Status* status);
/// Releases all resources used by the filesystem
///
/// NOTE: TensorFlow does not unload DSOs. Thus, the only way a filesystem
/// won't be registered anymore is if this function gets called by core
/// TensorFlow and the `TF_Filesystem*` object is destroyed. However, due to
/// registration being done in a static instance of `Env`, the destructor of
/// `FileSystem` is never called (see
/// https://github.com/tensorflow/tensorflow/issues/27535). In turn, this
/// function will never be called. There are plans to refactor registration
/// and fix this.
///
/// TODO(b/139060984): After all filesystems are converted, revisit note.
///
/// This operation must be provided. See "REQUIRED OPERATIONS" above.
void (*cleanup)(TF_Filesystem* filesystem);
/// Creates a new random access read-only file from given `path`.
///
/// After this call `file` may be concurrently accessed by multiple threads.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `file` was updated.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't point to an
/// existing file or one of the parent entries in `path` doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` points to a
/// directory or if it is invalid (e.g., malformed, or has a parent entry
/// which is a file).
/// * Might use any other error value for `status` to signal other errors.
///
/// REQUIREMENTS: If plugins implement this, they must also provide a filled
/// `TF_RandomAccessFileOps` table. See "REQUIRED OPERATIONS" above.
void (*new_random_access_file)(const TF_Filesystem* filesystem,
const char* path, TF_RandomAccessFile* file,
TF_Status* status);
/// Creates an object to write to a file with the specified `path`.
///
/// If the file already exists, it is deleted and recreated. The `file` object
/// must only be accessed by one thread at a time.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `file` was updated.
/// * Must set `status` to `TF_NOT_FOUND` if one of the parents entries in
/// `path` doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` points to a
/// directory or if it is invalid.
/// * Might use any other error value for `status` to signal other errors.
///
/// REQUIREMENTS: If plugins implement this, they must also provide a filled
/// `TF_WritableFileOps` table. See "REQUIRED OPERATIONS" above.
void (*new_writable_file)(const TF_Filesystem* filesystem, const char* path,
TF_WritableFile* file, TF_Status* status);
/// Creates an object to append to a file with the specified `path`.
///
/// If the file doesn't exists, it is first created with empty contents.
/// The `file` object must only be accessed by one thread at a time.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `file` was updated.
/// * Must set `status` to `TF_NOT_FOUND` if one of the parents entries in
/// `path` doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` points to a
/// directory or if it is invalid.
/// * Might use any other error value for `status` to signal other errors.
///
/// REQUIREMENTS: If plugins implement this, they must also provide a filled
/// `TF_WritableFileOps` table. See "REQUIRED OPERATIONS" above.
void (*new_appendable_file)(const TF_Filesystem* filesystem, const char* path,
TF_WritableFile* file, TF_Status* status);
/// Creates a read-only region of memory from contents of `path`.
///
/// After this call `region` may be concurrently accessed by multiple threads.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `region` was updated.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't point to an
/// existing file or one of the parent entries in `path` doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` points to a
/// directory or if it is invalid.
/// * Must set `status` to `TF_INVALID_ARGUMENT` if `path` points to an
/// empty file.
/// * Might use any other error value for `status` to signal other errors.
///
/// REQUIREMENTS: If plugins implement this, they must also provide a filled
/// `TF_ReadOnlyMemoryRegionOps` table. See "REQUIRED OPERATIONS" above.
void (*new_read_only_memory_region_from_file)(const TF_Filesystem* filesystem,
const char* path,
TF_ReadOnlyMemoryRegion* region,
TF_Status* status);
/// Creates the directory specified by `path`, assuming parent exists.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if directory was created.
/// * Must set `status` to `TF_NOT_FOUND` if one of the parents entries in
/// `path` doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` is invalid.
/// * Must set `status` to `TF_ALREADY_EXISTS` if `path` already exists.
/// * Might use any other error value for `status` to signal other errors.
void (*create_dir)(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
/// Creates the directory specified by `path` and all needed ancestors.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if directory was created.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` is invalid or
/// if it exists but is not a directory.
/// * Might use any other error value for `status` to signal other errors.
///
/// NOTE: The requirements specify that `TF_ALREADY_EXISTS` is not returned if
/// directory exists. Similarly, `TF_NOT_FOUND` is not be returned, as the
/// missing directory entry and all its descendants will be created by the
/// plugin.
///
/// DEFAULT IMPLEMENTATION: Creates directories one by one. Needs
/// `path_exists`, `is_directory`, and `create_dir`.
void (*recursively_create_dir)(const TF_Filesystem* filesystem,
const char* path, TF_Status* status);
/// Deletes the file specified by `path`.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if file was deleted.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` points to a
/// directory or if it is invalid.
/// * Might use any other error value for `status` to signal other errors.
void (*delete_file)(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
/// Deletes the empty directory specified by `path`.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if directory was deleted.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` does not point
/// to a directory, if `path` is invalid, or if directory is not empty.
/// * Might use any other error value for `status` to signal other errors.
void (*delete_dir)(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
/// Deletes the directory specified by `path` and all its contents.
///
/// This is accomplished by traversing directory tree rooted at `path` and
/// deleting entries as they are encountered, from leaves to root. Each plugin
/// is free to choose a different approach which obtains similar results.
///
/// On successful deletion, `status` must be `TF_OK` and `*undeleted_files`
/// and `*undeleted_dirs` must be 0. On unsuccessful deletion, `status` must
/// be set to the reason why one entry couldn't be removed and the proper
/// count must be updated. If the deletion is unsuccessful because the
/// traversal couldn't start, `*undeleted_files` must be set to 0 and
/// `*undeleted_dirs` must be set to 1.
///
/// TODO(b/139060984): After all filesystems are converted, consider
/// invariant about `*undeleted_files` and `*undeleted_dirs`.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if directory was deleted.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` is invalid.
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: Does a BFS traversal of tree rooted at `path`,
/// deleting entries as needed. Needs `path_exists`, `get_children`,
/// `is_directory`, `delete_file`, and `delete_dir`.
void (*delete_recursively)(const TF_Filesystem* filesystem, const char* path,
uint64_t* undeleted_files,
uint64_t* undeleted_dirs, TF_Status* status);
/// Renames the file given by `src` to that in `dst`.
///
/// Replaces `dst` if it exists. In case of error, both `src` and `dst` keep
/// the same state as before the call.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if rename was completed.
/// * Must set `status` to `TF_NOT_FOUND` if one of the parents entries in
/// either `src` or `dst` doesn't exist or if the specified `src` path
/// doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if either `src` or
/// `dst` is a directory or if either of them is invalid.
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: Copies file and deletes original. Needs
/// `copy_file`. and `delete_file`.
void (*rename_file)(const TF_Filesystem* filesystem, const char* src,
const char* dst, TF_Status* status);
/// Copies the file given by `src` to that in `dst`.
///
/// Similar to `rename_file`, but both `src` and `dst` exist after this call
/// with the same contents. In case of error, both `src` and `dst` keep the
/// same state as before the call.
///
/// If `dst` is a directory, creates a file with the same name as the source
/// inside the target directory.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if rename was completed.
/// * Must set `status` to `TF_NOT_FOUND` if one of the parents entries in
/// either `src` or `dst` doesn't exist or if the specified `src` path
/// doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if either `src` or
/// `dst` is a directory or if either of them is invalid.
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: Reads from `src` and writes to `dst`. Needs
/// `new_random_access_file` and `new_writable_file`.
void (*copy_file)(const TF_Filesystem* filesystem, const char* src,
const char* dst, TF_Status* status);
/// Checks if `path` exists.
///
/// Note that this doesn't differentiate between files and directories.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `path` exists.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't point to a
/// filesystem entry.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` is invalid.
/// * Might use any other error value for `status` to signal other errors.
void (*path_exists)(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
/// Checks if all values in `paths` exist in the filesystem.
///
/// Returns `true` if and only if calling `path_exists` on each entry in
/// `paths` would set `status` to `TF_OK`.
///
/// Caller guarantees that:
/// * `paths` has exactly `num_files` entries.
/// * `statuses` is either null or an array of `num_files` non-null elements
/// of type `TF_Status*`.
///
/// If `statuses` is not null, plugins must fill each element with detailed
/// status for each file, as if calling `path_exists` on each one. Core
/// TensorFlow initializes the `statuses` array and plugins must use
/// `TF_SetStatus` to set each element instead of directly assigning.
///
/// DEFAULT IMPLEMENTATION: Checks existence of every file. Needs
/// `path_exists`.
bool (*paths_exist)(const TF_Filesystem* filesystem, char** paths,
int num_files, TF_Status** statuses);
/// Obtains statistics for the given `path`.
///
/// Updates `stats` only if `status` is set to `TF_OK`.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `path` exists.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't point to a
/// filesystem entry.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` is invalid.
/// * Might use any other error value for `status` to signal other errors.
void (*stat)(const TF_Filesystem* filesystem, const char* path,
TF_FileStatistics* stats, TF_Status* status);
/// Checks whether the given `path` is a directory or not.
///
/// If `status` is not `TF_OK`, returns `false`, otherwise returns the same
/// as the `is_directory` member of a `TF_FileStatistics` that would be used
/// on the equivalent call of `stat`.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `path` exists.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't point to a
/// filesystem entry.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` is invalid.
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: Gets statistics about `path`. Needs `stat`.
bool (*is_directory)(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
/// Returns the size of the file given by `path`.
///
/// If `status` is not `TF_OK`, return value is undefined. Otherwise, returns
/// the same as `length` member of a `TF_FileStatistics` that would be used on
/// the equivalent call of `stat`.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `path` exists.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't point to a
/// filesystem entry.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` is invalid or
/// points to a directory.
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: Gets statistics about `path`. Needs `stat`.
int64_t (*get_file_size)(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
/// Translates `uri` to a filename for the filesystem
///
/// A filesystem is registered for a specific scheme and all of the methods
/// should work with URIs. Hence, each filesystem needs to be able to
/// translate from an URI to a path on the filesystem. For example, this
/// function could translate `fs:///path/to/a/file` into `/path/to/a/file`, if
/// implemented by a filesystem registered to handle the `fs://` scheme.
///
/// A new `char*` buffer must be allocated by this method. Core TensorFlow
/// manages the lifetime of the buffer after the call. Thus, all callers of
/// this method must take ownership of the returned pointer.
///
/// The implementation should clean up paths, including but not limited to,
/// removing duplicate `/`s, and resolving `..` and `.`.
///
/// 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.
///
/// DEFAULT IMPLEMENTATION: Uses `io::CleanPath` and `io::ParseURI`.
char* (*translate_name)(const TF_Filesystem* filesystem, const char* uri);
/// Finds all entries in the directory given by `path`.
///
/// The returned entries are paths relative to `path`.
///
/// Plugins must allocate `entries` to hold all names that need to be returned
/// and return the size of `entries`. Caller takes ownership of `entries`
/// after the call.
///
/// 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
/// filesystem entry or if one of the parents entries in `path` doesn't
/// exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if one of the parent
/// entries in `path` is not a directory, or if `path` is a file.
/// * Might use any other error value for `status` to signal other errors.
int (*get_children)(const TF_Filesystem* filesystem, const char* path,
char*** entries, TF_Status* status);
/// Finds all entries matching the regular expression given by `glob`.
///
/// Pattern must match the entire entry name, not just a substring.
///
/// pattern: { term }
/// term:
/// '*': matches any sequence of non-'/' characters
/// '?': matches a single non-'/' character
/// '[' [ '^' ] { match-list } ']':
/// matches any single character (not) on the list
/// c: matches character c (c != '*', '?', '\\', '[')
/// '\\' c: matches character c
/// character-range:
/// c: matches character c (c != '\\', '-', ']')
/// '\\' c: matches character c
/// lo '-' hi: matches character c for lo <= c <= hi
///
/// Implementations must allocate `entries` to hold all names that need to be
/// returned and return the size of `entries`. Caller takes ownership of
/// `entries` after the call.
///
/// In case of error, the implementations must set `status` to a value
/// 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.
///
/// DEFAULT IMPLEMENTATION: Scans the directory tree (in parallel if possible)
/// and fills `*entries`. Needs `get_children` and `is_directory`.
int (*get_matching_paths)(const TF_Filesystem* filesystem, const char* glob,
char*** entries, TF_Status* status);
/// Flushes any filesystem cache currently in memory
///
/// DEFAULT IMPLEMENTATION: No op.
void (*flush_caches)(const TF_Filesystem* filesystem);
/// Returns pointer to an array of available configuration options and their
/// current/default values in `options` and number of options in array in
/// `num_options`. Ownership of the array is transferred to caller and the
/// caller is responsible of freeing the buffers using respective file systems
/// allocation API.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `options` and `num_options` set.
/// If there is no configurable option, `num_options` should be 0.
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: return 0 options and `TF_OK`.
void (*get_filesystem_configuration)(const TF_Filesystem* filesystem,
TF_Filesystem_Option** options,
int* num_options, TF_Status* status);
/// Updates filesystem configuration with options passed in `options`. It can
/// contain full set of options supported by the filesystem or just a subset
/// of them. Ownership of options and buffers therein belongs to the caller
/// and any buffers need to be allocated through filesystem allocation API.
/// Filesystems may choose to ignore configuration errors but should at least
/// display a warning or error message to warn the users.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if options are updated.
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: return `TF_NOT_FOUND`.
void (*set_filesystem_configuration)(const TF_Filesystem* filesystem,
const TF_Filesystem_Option* options,
int num_options, TF_Status* status);
/// Returns the value of the filesystem option given in `key` in `option`.
/// Valid values of the `key` are returned by
/// `get_file_system_configuration_keys` call. Ownership of the
/// `option` is transferred to caller. Buffers therein should be allocated and
/// freed by the relevant filesystems allocation API.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `option` is set
/// * Must set `status` to `TF_NOT_FOUND` if the key is invalid
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: return `TF_NOT_FOUND`.
void (*get_filesystem_configuration_option)(const TF_Filesystem* filesystem,
const char* key,
TF_Filesystem_Option** option,
TF_Status* status);
/// Sets the value of the filesystem option given in `key` to value in
/// `option`. Valid values of the `key` are returned by
/// `get_file_system_configuration_keys` call. Ownership of the `option` and
/// the `key` belogs to the caller. Buffers therein should be allocated and
/// freed by the filesystems allocation API.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `option` is set/updated
/// * Must set `status` to `TF_NOT_FOUND` if the key is invalid
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: return `TF_NOT_FOUND`.
void (*set_filesystem_configuration_option)(
const TF_Filesystem* filesystem, const TF_Filesystem_Option* option,
TF_Status* status);
/// Returns a list of valid configuration keys in `keys` array and number of
/// keys in `num_keys`. Ownership of the buffers in `keys` are transferred to
/// caller and needs to be freed using relevant filesystem allocation API.
///
/// Plugins:
/// * Must set `status` to `TF_OK` on success. If there are no configurable
/// keys, `num_keys` should be set to 0
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: return `TF_OK` and `num_keys`=0.
void (*get_filesystem_configuration_keys)(const TF_Filesystem* filesystem,
char** keys, int* num_keys,
TF_Status* status);
} TF_FilesystemOps;
// LINT.ThenChange(:filesystem_ops_version)
/// SECTION 3. ABI and API compatibility
/// ----------------------------------------------------------------------------
///
/// In this section we define constants and macros to record versioning
/// information for each of the structures in section 2: ABI and API versions
/// and the number of functions in each of the function tables (which is
/// automatically determined, so ignored for the rest of this comment).
///
/// Since filesystem plugins are outside of TensorFlow's code tree, they are not
/// tied with TensorFlow releases and should have their own versioning metadata
/// in addition with the data discussed in this section. Each plugin author can
/// use a custom scheme, but it should only relate to changes in plugin code.
/// This section only touches metadata related to the versioning of this
/// interface that is shared by all possible plugins.
///
/// The API number increases whenever we break API compatibility while still
/// maintaining ABI compatibility. This happens only in the following cases:
/// 1. A new method is added _at the end_ of the function table.
/// 2. Preconditions or postconditions for one operation in these function
/// table change. Note that only core TensorFlow is able to impose these
/// invariants (i.e., guarantee the preconditions before calling the operation
/// and check the postconditions after the operation returns). If plugins need
/// additional invariants, they should be checked on the plugin side and the
/// `status` out variable should be updated accordingly (e.g., to include
/// plugin version information that relates to the condition change).
///
/// All other changes to the data structures (e.g., method removal, method
/// reordering, argument reordering, adding or removing arguments, changing the
/// type or the constness of a parameter, etc.) results in an ABI breakage.
/// Thus, we should not do any of these types of changes, except, potentially,
/// when we are releasing a new major version of TensorFlow. This is an escape
/// hatch, to be used rarely, preferably only to cleanup these structures.
/// Whenever we do these changes, the ABI number must be increased.
///
/// Next section will detail how this metadata is used at plugin registration to
/// only load compatible plugins and discard all others.
// LINT.IfChange(random_access_file_ops_version)
constexpr int TF_RANDOM_ACCESS_FILE_OPS_API = 0;
constexpr int TF_RANDOM_ACCESS_FILE_OPS_ABI = 0;
constexpr size_t TF_RANDOM_ACCESS_FILE_OPS_SIZE =
sizeof(TF_RandomAccessFileOps);
// LINT.ThenChange()
// LINT.IfChange(writable_file_ops_version)
constexpr int TF_WRITABLE_FILE_OPS_API = 0;
constexpr int TF_WRITABLE_FILE_OPS_ABI = 0;
constexpr size_t TF_WRITABLE_FILE_OPS_SIZE = sizeof(TF_WritableFileOps);
// LINT.ThenChange()
// LINT.IfChange(read_only_memory_region_ops_version)
constexpr int TF_READ_ONLY_MEMORY_REGION_OPS_API = 0;
constexpr int TF_READ_ONLY_MEMORY_REGION_OPS_ABI = 0;
constexpr size_t TF_READ_ONLY_MEMORY_REGION_OPS_SIZE =
sizeof(TF_ReadOnlyMemoryRegionOps);
// LINT.ThenChange()
// LINT.IfChange(filesystem_ops_version)
constexpr int TF_FILESYSTEM_OPS_API = 0;
constexpr int TF_FILESYSTEM_OPS_ABI = 0;
constexpr size_t TF_FILESYSTEM_OPS_SIZE = sizeof(TF_FilesystemOps);
// LINT.ThenChange()
/// SECTION 4. Plugin registration and initialization
/// ----------------------------------------------------------------------------
///
/// In this section we define the API used by core TensorFlow to initialize a
/// filesystem provided by a plugin. That is, we define the following:
/// * `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_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.
///
/// The `TF_InitPlugin` function is used by plugins to set up the data
/// structures that implement this interface, as presented in Section 2. In
/// order to not have plugin shared objects call back symbols defined in core
/// TensorFlow, `TF_InitPlugin` has a `TF_FilesystemPluginInfo` argument which
/// the plugin must fill (using the `TF_SetFilesystemVersionMetadata` for the
/// metadata and setting up all the supported operations and the URI schemes
/// that are supported).
/// This structure incorporates the operations defined in Section 2 and the
/// metadata defined in section 3, allowing plugins to define different ops
/// for different URI schemes.
///
/// Every URI scheme is of the form "fs" for URIs of form "fs:///path/to/file".
/// For local filesystems (i.e., when the URI is "/path/to/file"), the scheme
/// must be "". The scheme must never be `nullptr`.
///
/// Every plugin fills this in `TF_InitPlugin`, using the alocator passed as
/// argument to allocate memory. After `TF_InitPlugin` finishes, core
/// TensorFlow uses the information present in this to initialize filesystems
/// for the URI schemes that the plugin requests.
///
/// All pointers defined in this structure point to memory allocated by the DSO
/// using an allocator provided by core TensorFlow when calling `TF_InitPlugin`.
///
/// 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_FilesystemPluginOps {
char* scheme;
int filesystem_ops_abi;
int filesystem_ops_api;
size_t filesystem_ops_size;
TF_FilesystemOps* filesystem_ops;
int random_access_file_ops_abi;
int random_access_file_ops_api;
size_t random_access_file_ops_size;
TF_RandomAccessFileOps* random_access_file_ops;
int writable_file_ops_abi;
int writable_file_ops_api;
size_t writable_file_ops_size;
TF_WritableFileOps* writable_file_ops;
int read_only_memory_region_ops_abi;
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.
///
/// The argument is guaranteed to not be `nullptr`.
///
/// 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_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.
///
/// Must be implemented by the plugin DSO. It is called by TensorFlow runtime.
///
/// Filesystem plugins can be loaded on demand by users via
/// `Env::LoadLibrary` or during TensorFlow's startup if they are on certain
/// paths (although this has a security risk if two plugins register for the
/// same filesystem and the malicious one loads before the legimitate one -
/// but we consider this to be something that users should care about and
/// manage themselves). In both of these cases, core TensorFlow looks for
/// the `TF_InitPlugin` symbol and calls this function.
///
/// For every filesystem URI scheme that this plugin supports, the plugin must
/// add one `TF_FilesystemPluginInfo` entry in `plugin_info->ops` and call
/// `TF_SetFilesystemVersionMetadata` for that entry.
///
/// 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"
#endif // __cplusplus
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_FILESYSTEM_INTERFACE_H_
tensorflow/c/experimental/filesystem/modular_filesystem.cc
+577
View File
@@ -0,0 +1,577 @@
/* Copyright 2019 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/c/experimental/filesystem/modular_filesystem.h"
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/log/check.h"
#include "absl/strings/string_view.h"
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
#include "tensorflow/c/experimental/filesystem/modular_filesystem_registration.h"
#include "tensorflow/c/tf_file_statistics.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_status_helper.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/file_system.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/file_statistics.h"
#include "tensorflow/core/platform/file_system.h"
#include "tensorflow/core/platform/file_system_helper.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/strcat.h"
#include "tensorflow/core/platform/stringpiece.h"
#include "tensorflow/core/platform/types.h"
// TODO(b/139060984): After all filesystems are converted, all calls to
// methods from `FileSystem` will have to be replaced to calls to private
// methods here, as part of making this class a singleton and the only way to
// register/use filesystems.
namespace tensorflow {
using UniquePtrTo_TF_Status =
::std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)>;
Status ModularFileSystem::NewRandomAccessFile(
const std::string& fname, std::unique_ptr<RandomAccessFile>* result) {
if (ops_->new_random_access_file == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", fname, " does not support NewRandomAccessFile()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
auto file = std::make_unique<TF_RandomAccessFile>();
std::string translated_name = TranslateName(fname);
ops_->new_random_access_file(filesystem_.get(), translated_name.c_str(),
file.get(), plugin_status.get());
if (TF_GetCode(plugin_status.get()) == TF_OK)
*result = std::make_unique<ModularRandomAccessFile>(
translated_name, std::move(file), random_access_file_ops_.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::NewWritableFile(
const std::string& fname, std::unique_ptr<WritableFile>* result) {
if (ops_->new_writable_file == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", fname, " does not support NewWritableFile()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
auto file = std::make_unique<TF_WritableFile>();
std::string translated_name = TranslateName(fname);
ops_->new_writable_file(filesystem_.get(), translated_name.c_str(),
file.get(), plugin_status.get());
if (TF_GetCode(plugin_status.get()) == TF_OK)
*result = std::make_unique<ModularWritableFile>(
translated_name, std::move(file), writable_file_ops_.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::NewAppendableFile(
const std::string& fname, std::unique_ptr<WritableFile>* result) {
if (ops_->new_appendable_file == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", fname, " does not support NewAppendableFile()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
auto file = std::make_unique<TF_WritableFile>();
std::string translated_name = TranslateName(fname);
ops_->new_appendable_file(filesystem_.get(), translated_name.c_str(),
file.get(), plugin_status.get());
if (TF_GetCode(plugin_status.get()) == TF_OK)
*result = std::make_unique<ModularWritableFile>(
translated_name, std::move(file), writable_file_ops_.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::NewReadOnlyMemoryRegionFromFile(
const std::string& fname, std::unique_ptr<ReadOnlyMemoryRegion>* result) {
if (ops_->new_read_only_memory_region_from_file == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", fname,
" does not support NewReadOnlyMemoryRegionFromFile()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
auto region = std::make_unique<TF_ReadOnlyMemoryRegion>();
std::string translated_name = TranslateName(fname);
ops_->new_read_only_memory_region_from_file(
filesystem_.get(), translated_name.c_str(), region.get(),
plugin_status.get());
if (TF_GetCode(plugin_status.get()) == TF_OK)
*result = std::make_unique<ModularReadOnlyMemoryRegion>(
std::move(region), read_only_memory_region_ops_.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::FileExists(absl::string_view fname) {
if (ops_->path_exists == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", fname, " does not support FileExists()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
const std::string translated_name = TranslateName(fname);
ops_->path_exists(filesystem_.get(), translated_name.c_str(),
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
bool ModularFileSystem::FilesExist(const std::vector<std::string>& files,
std::vector<Status>* status) {
if (ops_->paths_exist == nullptr)
return FileSystem::FilesExist(files, status);
std::vector<char*> translated_names;
translated_names.reserve(files.size());
for (int i = 0; i < files.size(); i++)
translated_names.push_back(strdup(TranslateName(files[i]).c_str()));
bool result;
if (status == nullptr) {
result = ops_->paths_exist(filesystem_.get(), translated_names.data(),
files.size(), nullptr);
} else {
std::vector<TF_Status*> plugin_status;
plugin_status.reserve(files.size());
for (int i = 0; i < files.size(); i++)
plugin_status.push_back(TF_NewStatus());
result = ops_->paths_exist(filesystem_.get(), translated_names.data(),
files.size(), plugin_status.data());
for (int i = 0; i < files.size(); i++) {
status->push_back(StatusFromTF_Status(plugin_status[i]));
TF_DeleteStatus(plugin_status[i]);
}
}
for (int i = 0; i < files.size(); i++) free(translated_names[i]);
return result;
}
Status ModularFileSystem::GetChildren(const std::string& dir,
std::vector<std::string>* result) {
if (ops_->get_children == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", dir, " does not support GetChildren()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(dir);
// 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]));
plugin_memory_free_(children[i]);
}
plugin_memory_free_(children);
}
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::GetMatchingPaths(const std::string& pattern,
std::vector<std::string>* result) {
if (ops_->get_matching_paths == nullptr)
return internal::GetMatchingPaths(this, Env::Default(), pattern, result);
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
// 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]));
plugin_memory_free_(matches[i]);
}
plugin_memory_free_(matches);
}
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::DeleteFile(const std::string& fname) {
if (ops_->delete_file == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", fname, " does not support DeleteFile()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(fname);
ops_->delete_file(filesystem_.get(), translated_name.c_str(),
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::DeleteRecursively(const std::string& dirname,
int64_t* undeleted_files,
int64_t* undeleted_dirs) {
if (undeleted_files == nullptr || undeleted_dirs == nullptr)
return errors::FailedPrecondition(
"DeleteRecursively must not be called with `undeleted_files` or "
"`undeleted_dirs` set to NULL");
if (ops_->delete_recursively == nullptr)
return FileSystem::DeleteRecursively(dirname, undeleted_files,
undeleted_dirs);
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(dirname);
uint64_t plugin_undeleted_files, plugin_undeleted_dirs;
ops_->delete_recursively(filesystem_.get(), translated_name.c_str(),
&plugin_undeleted_files, &plugin_undeleted_dirs,
plugin_status.get());
*undeleted_files = plugin_undeleted_files;
*undeleted_dirs = plugin_undeleted_dirs;
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::DeleteDir(const std::string& dirname) {
if (ops_->delete_dir == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", dirname, " does not support DeleteDir()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(dirname);
ops_->delete_dir(filesystem_.get(), translated_name.c_str(),
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::RecursivelyCreateDir(const std::string& dirname) {
if (ops_->recursively_create_dir == nullptr)
return FileSystem::RecursivelyCreateDir(dirname);
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(dirname);
ops_->recursively_create_dir(filesystem_.get(), translated_name.c_str(),
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::CreateDir(const std::string& dirname) {
if (ops_->create_dir == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", dirname, " does not support CreateDir()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(dirname);
ops_->create_dir(filesystem_.get(), translated_name.c_str(),
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::Stat(const std::string& fname, FileStatistics* stat) {
if (ops_->stat == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", fname, " does not support Stat()"));
if (stat == nullptr)
return errors::InvalidArgument("FileStatistics pointer must not be NULL");
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(fname);
TF_FileStatistics stats;
ops_->stat(filesystem_.get(), translated_name.c_str(), &stats,
plugin_status.get());
if (TF_GetCode(plugin_status.get()) == TF_OK) {
stat->length = stats.length;
stat->mtime_nsec = stats.mtime_nsec;
stat->is_directory = stats.is_directory;
}
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::IsDirectory(const std::string& name) {
if (ops_->is_directory == nullptr) return FileSystem::IsDirectory(name);
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(name);
ops_->is_directory(filesystem_.get(), translated_name.c_str(),
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::GetFileSize(const std::string& fname,
uint64* file_size) {
if (ops_->get_file_size == nullptr) {
FileStatistics stat;
Status status = Stat(fname, &stat);
if (!status.ok()) return status;
if (stat.is_directory)
return errors::FailedPrecondition("Called GetFileSize on a directory");
*file_size = stat.length;
return status;
}
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(fname);
*file_size = ops_->get_file_size(filesystem_.get(), translated_name.c_str(),
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::RenameFile(const std::string& src,
const std::string& target) {
if (ops_->rename_file == nullptr) {
Status status = CopyFile(src, target);
if (status.ok()) status = DeleteFile(src);
return status;
}
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_src = TranslateName(src);
std::string translated_target = TranslateName(target);
ops_->rename_file(filesystem_.get(), translated_src.c_str(),
translated_target.c_str(), plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::CopyFile(const std::string& src,
const std::string& target) {
if (ops_->copy_file == nullptr) return FileSystem::CopyFile(src, target);
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_src = TranslateName(src);
std::string translated_target = TranslateName(target);
ops_->copy_file(filesystem_.get(), translated_src.c_str(),
translated_target.c_str(), plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
std::string ModularFileSystem::TranslateName(absl::string_view name) const {
if (ops_->translate_name == nullptr) return FileSystem::TranslateName(name);
char* p = ops_->translate_name(filesystem_.get(), std::string(name).c_str());
CHECK(p != nullptr) << "TranslateName(" << name << ") returned nullptr";
std::string ret(p);
// Since `p` is allocated by plugin, free it using plugin's method.
plugin_memory_free_(p);
return ret;
}
void ModularFileSystem::FlushCaches() {
if (ops_->flush_caches != nullptr) ops_->flush_caches(filesystem_.get());
}
Status ModularFileSystem::SetOption(const std::string& name,
const std::vector<string>& values) {
if (ops_->set_filesystem_configuration == nullptr) {
return errors::Unimplemented(
"Filesystem does not support SetConfiguration()");
}
if (values.empty()) {
return errors::InvalidArgument(
"SetConfiguration() needs number of values > 0");
}
TF_Filesystem_Option option;
memset(&option, 0, sizeof(option));
option.name = const_cast<char*>(name.c_str());
TF_Filesystem_Option_Value option_value;
memset(&option_value, 0, sizeof(option_value));
option_value.type_tag = TF_Filesystem_Option_Type_Buffer;
option_value.num_values = values.size();
std::vector<TF_Filesystem_Option_Value_Union> option_values(values.size());
for (size_t i = 0; i < values.size(); i++) {
memset(&option_values[i], 0, sizeof(option_values[i]));
option_values[i].buffer_val.buf = const_cast<char*>(values[i].c_str());
option_values[i].buffer_val.buf_length = values[i].size();
}
option_value.values = &option_values[0];
option.value = &option_value;
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
ops_->set_filesystem_configuration(filesystem_.get(), &option, 1,
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::SetOption(const std::string& name,
const std::vector<int64_t>& values) {
if (ops_->set_filesystem_configuration == nullptr) {
return errors::Unimplemented(
"Filesystem does not support SetConfiguration()");
}
if (values.empty()) {
return errors::InvalidArgument(
"SetConfiguration() needs number of values > 0");
}
TF_Filesystem_Option option;
memset(&option, 0, sizeof(option));
option.name = const_cast<char*>(name.c_str());
TF_Filesystem_Option_Value option_value;
memset(&option_value, 0, sizeof(option_value));
option_value.type_tag = TF_Filesystem_Option_Type_Int;
option_value.num_values = values.size();
std::vector<TF_Filesystem_Option_Value_Union> option_values(values.size());
for (size_t i = 0; i < values.size(); i++) {
memset(&option_values[i], 0, sizeof(option_values[i]));
option_values[i].int_val = values[i];
}
option_value.values = &option_values[0];
option.value = &option_value;
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
ops_->set_filesystem_configuration(filesystem_.get(), &option, 1,
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::SetOption(const std::string& name,
const std::vector<double>& values) {
if (ops_->set_filesystem_configuration == nullptr) {
return errors::Unimplemented(
"Filesystem does not support SetConfiguration()");
}
if (values.empty()) {
return errors::InvalidArgument(
"SetConfiguration() needs number of values > 0");
}
TF_Filesystem_Option option;
memset(&option, 0, sizeof(option));
option.name = const_cast<char*>(name.c_str());
TF_Filesystem_Option_Value option_value;
memset(&option_value, 0, sizeof(option_value));
option_value.type_tag = TF_Filesystem_Option_Type_Real;
option_value.num_values = values.size();
std::vector<TF_Filesystem_Option_Value_Union> option_values(values.size());
for (size_t i = 0; i < values.size(); i++) {
memset(&option_values[i], 0, sizeof(option_values[i]));
option_values[i].real_val = values[i];
}
option_value.values = &option_values[0];
option.value = &option_value;
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
ops_->set_filesystem_configuration(filesystem_.get(), &option, 1,
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularRandomAccessFile::Read(uint64 offset, size_t n,
StringPiece* result, char* scratch) const {
if (ops_->read == nullptr)
return errors::Unimplemented(
tensorflow::strings::StrCat("Read() not implemented for ", filename_));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
int64_t read =
ops_->read(file_.get(), offset, n, scratch, plugin_status.get());
if (read > 0) *result = StringPiece(scratch, read);
return StatusFromTF_Status(plugin_status.get());
}
Status ModularRandomAccessFile::Name(StringPiece* result) const {
*result = filename_;
return OkStatus();
}
Status ModularWritableFile::Append(StringPiece data) {
if (ops_->append == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Append() not implemented for ", filename_));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
ops_->append(file_.get(), data.data(), data.size(), plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularWritableFile::Close() {
if (ops_->close == nullptr)
return errors::Unimplemented(
tensorflow::strings::StrCat("Close() not implemented for ", filename_));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
ops_->close(file_.get(), plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularWritableFile::Flush() {
if (ops_->flush == nullptr) return OkStatus();
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
ops_->flush(file_.get(), plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularWritableFile::Sync() {
if (ops_->sync == nullptr) return Flush();
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
ops_->sync(file_.get(), plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularWritableFile::Name(StringPiece* result) const {
*result = filename_;
return OkStatus();
}
Status ModularWritableFile::Tell(int64_t* position) {
if (ops_->tell == nullptr)
return errors::Unimplemented(
tensorflow::strings::StrCat("Tell() not implemented for ", filename_));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
*position = ops_->tell(file_.get(), plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status RegisterFilesystemPlugin(const std::string& dso_path) {
// Step 1: Load plugin
Env* env = Env::Default();
void* dso_handle;
TF_RETURN_IF_ERROR(env->LoadDynamicLibrary(dso_path.c_str(), &dso_handle));
// Step 2: Load symbol for `TF_InitPlugin`
void* dso_symbol;
TF_RETURN_WITH_CONTEXT_IF_ERROR(
env->GetSymbolFromLibrary(dso_handle, "TF_InitPlugin", &dso_symbol),
"Failed to load TF_InitPlugin symbol for DSO: ", dso_path);
// Step 3: Call `TF_InitPlugin`
TF_FilesystemPluginInfo info;
memset(&info, 0, sizeof(info));
auto TF_InitPlugin =
reinterpret_cast<int (*)(TF_FilesystemPluginInfo*)>(dso_symbol);
TF_InitPlugin(&info);
// Step 4: Do the actual registration
return filesystem_registration::RegisterFilesystemPluginImpl(&info);
}
} // namespace tensorflow
tensorflow/c/experimental/filesystem/modular_filesystem.h
+197
View File
@@ -0,0 +1,197 @@
/* Copyright 2019 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_MODULAR_FILESYSTEM_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_MODULAR_FILESYSTEM_H_
#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
#include "xla/tsl/platform/file_system.h"
#include "tensorflow/core/platform/file_statistics.h"
#include "tensorflow/core/platform/file_system.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/stringpiece.h"
#include "tensorflow/core/platform/types.h"
/// This file builds classes needed to hold a filesystem implementation in the
/// modular world. Once all TensorFlow filesystems are converted to use the
/// plugin based approach, this file will replace the one in core/platform and
/// the names will lose the `Modular` part. Until that point, the `Modular*`
/// classes here are experimental and subject to breaking changes.
/// For documentation on these methods, consult `core/platform/filesystem.h`.
namespace tensorflow {
// TODO(b/143949615): After all filesystems are converted, this file will be
// moved to core/platform, and this class can become a singleton and replace the
// need for `Env::Default()`. At that time, we might decide to remove the need
// for `Env::Default()` altogether, but that's a different project, not in
// scope for now. I'm just mentioning this here as that transition will mean
// removal of the registration part from `Env` and adding it here instead: we
// will need tables to hold for each scheme the function tables that implement
// the needed functionality instead of the current `FileSystemRegistry` code in
// `core/platform/env.cc`.
class ModularFileSystem final : public FileSystem {
public:
ModularFileSystem(
std::unique_ptr<TF_Filesystem> filesystem,
std::unique_ptr<const TF_FilesystemOps> filesystem_ops,
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,
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)),
plugin_memory_allocate_(std::move(plugin_memory_allocate)),
plugin_memory_free_(std::move(plugin_memory_free)) {}
~ModularFileSystem() override { ops_->cleanup(filesystem_.get()); }
absl::Status NewRandomAccessFile(
const std::string& fname,
std::unique_ptr<RandomAccessFile>* result) override;
absl::Status NewWritableFile(const std::string& fname,
std::unique_ptr<WritableFile>* result) override;
absl::Status NewAppendableFile(
const std::string& fname, std::unique_ptr<WritableFile>* result) override;
absl::Status NewReadOnlyMemoryRegionFromFile(
const std::string& fname,
std::unique_ptr<ReadOnlyMemoryRegion>* result) override;
absl::Status FileExists(absl::string_view fname) override;
bool FilesExist(const std::vector<std::string>& files,
std::vector<absl::Status>* status) override;
absl::Status GetChildren(const std::string& dir,
std::vector<std::string>* result) override;
absl::Status GetMatchingPaths(const std::string& pattern,
std::vector<std::string>* results) override;
absl::Status DeleteFile(const std::string& fname) override;
absl::Status DeleteRecursively(const std::string& dirname,
int64_t* undeleted_files,
int64_t* undeleted_dirs) override;
absl::Status DeleteDir(const std::string& dirname) override;
absl::Status RecursivelyCreateDir(const std::string& dirname) override;
absl::Status CreateDir(const std::string& dirname) override;
absl::Status Stat(const std::string& fname, FileStatistics* stat) override;
absl::Status IsDirectory(const std::string& fname) override;
absl::Status GetFileSize(const std::string& fname,
uint64_t* file_size) override;
absl::Status RenameFile(const std::string& src,
const std::string& target) override;
absl::Status CopyFile(const std::string& src,
const std::string& target) override;
std::string TranslateName(absl::string_view name) const override;
void FlushCaches() override;
absl::Status SetOption(const std::string& name,
const std::vector<std::string>& values) override;
absl::Status SetOption(const std::string& name,
const std::vector<int64_t>& values) override;
absl::Status SetOption(const std::string& name,
const std::vector<double>& values) override;
private:
std::unique_ptr<TF_Filesystem> filesystem_;
std::unique_ptr<const TF_FilesystemOps> ops_;
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_;
std::function<void*(size_t)> plugin_memory_allocate_;
std::function<void(void*)> plugin_memory_free_;
ModularFileSystem(const ModularFileSystem&) = delete;
void operator=(const ModularFileSystem&) = delete;
};
class ModularRandomAccessFile final : public RandomAccessFile {
public:
ModularRandomAccessFile(const std::string& filename,
std::unique_ptr<TF_RandomAccessFile> file,
const TF_RandomAccessFileOps* ops)
: filename_(filename), file_(std::move(file)), ops_(ops) {}
~ModularRandomAccessFile() override { ops_->cleanup(file_.get()); }
absl::Status Read(uint64_t offset, size_t n, absl::string_view* result,
char* scratch) const override;
absl::Status Name(absl::string_view* result) const override;
private:
std::string filename_;
std::unique_ptr<TF_RandomAccessFile> file_;
const TF_RandomAccessFileOps* ops_; // not owned
ModularRandomAccessFile(const ModularRandomAccessFile&) = delete;
void operator=(const ModularRandomAccessFile&) = delete;
};
class ModularWritableFile final : public WritableFile {
public:
ModularWritableFile(const std::string& filename,
std::unique_ptr<TF_WritableFile> file,
const TF_WritableFileOps* ops)
: filename_(filename), file_(std::move(file)), ops_(ops) {}
~ModularWritableFile() override { ops_->cleanup(file_.get()); }
absl::Status Append(absl::string_view data) override;
absl::Status Close() override;
absl::Status Flush() override;
absl::Status Sync() override;
absl::Status Name(absl::string_view* result) const override;
absl::Status Tell(int64_t* position) override;
private:
std::string filename_;
std::unique_ptr<TF_WritableFile> file_;
const TF_WritableFileOps* ops_; // not owned
ModularWritableFile(const ModularWritableFile&) = delete;
void operator=(const ModularWritableFile&) = delete;
};
class ModularReadOnlyMemoryRegion final : public ReadOnlyMemoryRegion {
public:
ModularReadOnlyMemoryRegion(std::unique_ptr<TF_ReadOnlyMemoryRegion> region,
const TF_ReadOnlyMemoryRegionOps* ops)
: region_(std::move(region)), ops_(ops) {}
~ModularReadOnlyMemoryRegion() override { ops_->cleanup(region_.get()); };
const void* data() override { return ops_->data(region_.get()); }
uint64_t length() override { return ops_->length(region_.get()); }
private:
std::unique_ptr<TF_ReadOnlyMemoryRegion> region_;
const TF_ReadOnlyMemoryRegionOps* ops_; // not owned
ModularReadOnlyMemoryRegion(const ModularReadOnlyMemoryRegion&) = delete;
void operator=(const ModularReadOnlyMemoryRegion&) = delete;
};
// Registers a filesystem plugin so that core TensorFlow can use it.
absl::Status RegisterFilesystemPlugin(const std::string& dso_path);
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_MODULAR_FILESYSTEM_H_
tensorflow/c/experimental/filesystem/modular_filesystem_registration.cc
+338
View File
@@ -0,0 +1,338 @@
/* Copyright 2020 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/c/experimental/filesystem/modular_filesystem_registration.h"
#include <algorithm>
#include <cstddef>
#include <memory>
#include <utility>
#include "absl/log/log.h"
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
#include "tensorflow/c/experimental/filesystem/modular_filesystem.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_status_internal.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/strcat.h"
#include "tensorflow/core/platform/stringpiece.h"
namespace tensorflow {
// Checks that all schemes provided by a plugin are valid.
// TODO(b/139060984): More validation could be done here, based on supported
// charset, maximum length, etc. Punting it for later.
static Status ValidateScheme(const char* scheme) {
if (scheme == nullptr)
return errors::InvalidArgument(
"Attempted to register filesystem with `nullptr` URI scheme");
return OkStatus();
}
// Checks if the plugin and core ABI numbers match.
//
// If the numbers don't match, plugin cannot be loaded.
static Status CheckABI(int pluginABI, int coreABI, StringPiece where) {
if (pluginABI != coreABI)
return errors::FailedPrecondition(
strings::StrCat("Plugin ABI (", pluginABI, ") for ", where,
" operations doesn't match expected core ABI (",
coreABI, "). Plugin cannot be loaded."));
return OkStatus();
}
// Checks if the plugin and core ABI numbers match, for all operations.
//
// If the numbers don't match, plugin cannot be loaded.
//
// Uses the simpler `CheckABI(int, int, StringPiece)`.
static Status ValidateABI(const TF_FilesystemPluginOps* ops) {
TF_RETURN_IF_ERROR(
CheckABI(ops->filesystem_ops_abi, TF_FILESYSTEM_OPS_ABI, "filesystem"));
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 (ops->writable_file_ops != nullptr)
TF_RETURN_IF_ERROR(CheckABI(ops->writable_file_ops_abi,
TF_WRITABLE_FILE_OPS_ABI, "writable file"));
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"));
return OkStatus();
}
// Checks if the plugin and core API numbers match, logging mismatches.
static void CheckAPI(int plugin_API, int core_API, StringPiece where) {
if (plugin_API != core_API) {
VLOG(0) << "Plugin API (" << plugin_API << ") for " << where
<< " operations doesn't match expected core API (" << core_API
<< "). Plugin will be loaded but functionality might be missing.";
}
}
// Checks if the plugin and core API numbers match, for all operations.
//
// Uses the simpler `CheckAPIHelper(int, int, StringPiece)`.
static void ValidateAPI(const TF_FilesystemPluginOps* ops) {
CheckAPI(ops->filesystem_ops_api, TF_FILESYSTEM_OPS_API, "filesystem");
if (ops->random_access_file_ops != nullptr)
CheckAPI(ops->random_access_file_ops_api, TF_RANDOM_ACCESS_FILE_OPS_API,
"random access file");
if (ops->writable_file_ops != nullptr)
CheckAPI(ops->writable_file_ops_api, TF_WRITABLE_FILE_OPS_API,
"writable file");
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");
}
// Validates the filesystem operations supplied by the plugin.
static Status ValidateHelper(const TF_FilesystemOps* ops) {
if (ops == nullptr)
return errors::FailedPrecondition(
"Trying to register filesystem without operations");
if (ops->init == nullptr)
return errors::FailedPrecondition(
"Trying to register filesystem without `init` operation");
if (ops->cleanup == nullptr)
return errors::FailedPrecondition(
"Trying to register filesystem without `cleanup` operation");
return OkStatus();
}
// Validates the random access file operations supplied by the plugin.
static Status ValidateHelper(const TF_RandomAccessFileOps* ops) {
if (ops == nullptr) {
// We allow filesystems where files can only be written to (from TF code)
return OkStatus();
}
if (ops->cleanup == nullptr)
return errors::FailedPrecondition(
"Trying to register filesystem without `cleanup` operation on random "
"access files");
return OkStatus();
}
// Validates the writable file operations supplied by the plugin.
static Status ValidateHelper(const TF_WritableFileOps* ops) {
if (ops == nullptr) {
// We allow read-only filesystems
return OkStatus();
}
if (ops->cleanup == nullptr)
return errors::FailedPrecondition(
"Trying to register filesystem without `cleanup` operation on writable "
"files");
return OkStatus();
}
// Validates the read only memory region operations given by the plugin.
static Status ValidateHelper(const TF_ReadOnlyMemoryRegionOps* ops) {
if (ops == nullptr) {
// read only memory region support is always optional
return OkStatus();
}
if (ops->cleanup == nullptr)
return errors::FailedPrecondition(
"Trying to register filesystem without `cleanup` operation on read "
"only memory regions");
if (ops->data == nullptr)
return errors::FailedPrecondition(
"Trying to register filesystem without `data` operation on read only "
"memory regions");
if (ops->length == nullptr)
return errors::FailedPrecondition(
"Trying to register filesystem without `length` operation on read only "
"memory regions");
return OkStatus();
}
// Validates the operations supplied by the plugin.
//
// Uses the 4 simpler `ValidateHelper(const TF_...*)` to validate each
// 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_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 (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 ((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 (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.");
return OkStatus();
}
// Copies a function table from plugin memory space to core memory space.
//
// This has three benefits:
// * allows having newer plugins than the current core TensorFlow: the
// additional entries in the plugin's table are just discarded;
// * allows having older plugins than the current core TensorFlow (though
// we are still warning users): the entries that core TensorFlow expects
// but plugins didn't provide will be set to `nullptr` values and core
// TensorFlow will know to not call these on behalf of users;
// * increased security as plugins will not be able to alter function table
// after loading up. Thus, malicious plugins can't alter functionality to
// probe for gadgets inside core TensorFlow. We can even protect the area
// of memory where the copies reside to not allow any more writes to it
// after all copies are created.
template <typename T>
static std::unique_ptr<const T> CopyToCore(const T* plugin_ops,
size_t plugin_size) {
if (plugin_ops == nullptr) return nullptr;
size_t copy_size = std::min(plugin_size, sizeof(T));
auto core_ops = std::make_unique<T>();
memset(core_ops.get(), 0, sizeof(T));
memcpy(core_ops.get(), plugin_ops, copy_size);
return core_ops;
}
// Registers one filesystem from the plugin.
//
// 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->ops[index].filesystem_ops, info->ops[index].filesystem_ops_size);
auto core_random_access_file_ops = CopyToCore<TF_RandomAccessFileOps>(
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->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 = std::make_unique<TF_Filesystem>();
TF_Status* c_status = TF_NewStatus();
Status status = OkStatus();
core_filesystem_ops->init(filesystem.get(), c_status);
status = Status(c_status->status);
TF_DeleteStatus(c_status);
if (!status.ok()) return status;
// Step 3: Actual registration
return Env::Default()->RegisterFileSystem(
info->ops[index].scheme,
std::make_unique<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 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, 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 OkStatus();
}
// 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 OkStatus();
}
namespace filesystem_registration {
Status RegisterFilesystemPluginImpl(const TF_FilesystemPluginInfo* info) {
TF_RETURN_IF_ERROR(ValidatePluginMemoryRoutines(info));
// 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 < 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);
}
info->plugin_memory_free(info->ops);
return status;
}
} // namespace filesystem_registration
} // namespace tensorflow
tensorflow/c/experimental/filesystem/modular_filesystem_registration.h
+34
View File
@@ -0,0 +1,34 @@
/* Copyright 2020 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_MODULAR_FILESYSTEM_REGISTRATION_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_MODULAR_FILESYSTEM_REGISTRATION_H_
#include "absl/status/status.h"
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
namespace filesystem_registration {
// Implementation for filesystem registration
//
// Don't call this directly. Instead call `RegisterFilesystemPlugin`.
// Exposed only for static registration of local filesystems.
absl::Status RegisterFilesystemPluginImpl(const TF_FilesystemPluginInfo* info);
} // namespace filesystem_registration
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_MODULAR_FILESYSTEM_REGISTRATION_H_
tensorflow/c/experimental/filesystem/modular_filesystem_test.cc
+1803
View File
File diff suppressed because it is too large Load Diff
tensorflow/c/experimental/filesystem/plugins/gcs/BUILD
+139
View File
@@ -0,0 +1,139 @@
# Experimental gcs filesystem plugin.
load("//tensorflow:tensorflow.bzl", "get_win_copts", "tf_cc_shared_object", "tf_cc_test")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
# Filesystem implementation for GCS environments
tf_cc_shared_object(
name = "gcs_filesystem",
framework_so = [],
linkstatic = False,
per_os_targets = 1,
visibility = ["//visibility:public"],
deps = [":gcs_filesystem_impl"],
)
# The real implementation of the filesystem.
cc_library(
name = "gcs_filesystem_impl",
srcs = ["gcs_filesystem.cc"],
hdrs = ["gcs_filesystem.h"],
copts = select({
"//conditions:default": [],
"//tensorflow:windows": get_win_copts(),
}),
deps = [
":expiring_lru_cache",
":gcs_helper",
":ram_file_block_cache",
"//tensorflow/c:env",
"//tensorflow/c:tf_status",
"//tensorflow/c/experimental/filesystem:filesystem_interface",
"@com_github_googlecloudplatform_google_cloud_cpp//:storage_client",
"@com_github_googlecloudplatform_google_cloud_cpp//google/cloud:google_cloud_cpp_common",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/synchronization",
],
)
cc_library(
name = "gcs_helper",
srcs = ["gcs_helper.cc"],
hdrs = ["gcs_helper.h"],
linkstatic = 1,
deps = [
"//tensorflow/c:env",
],
)
cc_library(
name = "cleanup",
hdrs = ["cleanup.h"],
)
cc_library(
name = "ram_file_block_cache",
srcs = ["ram_file_block_cache.cc"],
hdrs = ["ram_file_block_cache.h"],
deps = [
":cleanup",
"//tensorflow/c:env",
"//tensorflow/c:tf_status",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
],
)
tf_cc_test(
name = "ram_file_block_cache_test",
size = "small",
srcs = ["ram_file_block_cache_test.cc"],
deps = [
":ram_file_block_cache",
"//tensorflow/c:tf_status_internal",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform/cloud:now_seconds_env",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
],
)
tf_cc_test(
name = "gcs_filesystem_test",
srcs = [
"gcs_filesystem_test.cc",
],
tags = [
"manual",
"notap",
],
deps = [
":gcs_filesystem_impl",
"//tensorflow/core/platform:path",
"//tensorflow/core/platform:stacktrace_handler",
"//tensorflow/core/platform:test",
"@com_github_googlecloudplatform_google_cloud_cpp//google/cloud:google_cloud_cpp_common",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "expiring_lru_cache",
hdrs = ["expiring_lru_cache.h"],
deps = [
"//tensorflow/c:env",
"//tensorflow/c:tf_status",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/synchronization",
],
)
tf_cc_test(
name = "expiring_lru_cache_test",
size = "small",
srcs = ["expiring_lru_cache_test.cc"],
deps = [
":expiring_lru_cache",
"//tensorflow/c:tf_status_helper",
"//tensorflow/c:tf_status_internal",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform/cloud:now_seconds_env",
],
)
tensorflow/c/experimental/filesystem/plugins/gcs/cleanup.h
+109
View File
@@ -0,0 +1,109 @@
/* Copyright 2020 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.
==============================================================================*/
// MakeCleanup(f) returns an RAII cleanup object that calls 'f' in its
// destructor. The easiest way to use MakeCleanup is with a lambda argument,
// capturing the return value in an 'auto' local variable. Most users will not
// need more sophisticated syntax than that.
//
// Example:
// void func() {
// FILE* fp = fopen("data.txt", "r");
// if (fp == nullptr) return;
// auto fp_cleaner = gtl::MakeCleanup([fp] { fclose(fp); });
// // No matter what, fclose(fp) will happen.
// DataObject d;
// while (ReadDataObject(fp, &d)) {
// if (d.IsBad()) {
// LOG(ERROR) << "Bad Data";
// return;
// }
// PushGoodData(d);
// }
// }
//
// You can use Cleanup<F> directly, instead of using MakeCleanup and auto,
// but there's rarely a reason to do that.
//
// You can call 'release()' on a Cleanup object to cancel the cleanup.
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_CLEANUP_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_CLEANUP_H_
#include <type_traits>
#include <utility>
namespace tf_gcs_filesystem {
// A move-only RAII object that calls a stored cleanup functor when
// destroyed. Cleanup<F> is the return type of gtl::MakeCleanup(F).
template <typename F>
class Cleanup {
public:
Cleanup() : released_(true), f_() {}
template <typename G>
explicit Cleanup(G&& f) // NOLINT
: f_(std::forward<G>(f)) {} // NOLINT(build/c++11)
Cleanup(Cleanup&& src) // NOLINT
: released_(src.is_released()), f_(src.release()) {}
// Implicitly move-constructible from any compatible Cleanup<G>.
// The source will be released as if src.release() were called.
// A moved-from Cleanup can be safely destroyed or reassigned.
template <typename G>
Cleanup(Cleanup<G>&& src) // NOLINT
: released_(src.is_released()), f_(src.release()) {}
// Assignment to a Cleanup object behaves like destroying it
// and making a new one in its place, analogous to unique_ptr
// semantics.
Cleanup& operator=(Cleanup&& src) { // NOLINT
if (!released_) f_();
released_ = src.released_;
f_ = src.release();
return *this;
}
~Cleanup() {
if (!released_) f_();
}
// Releases the cleanup function instead of running it.
// Hint: use c.release()() to run early.
F release() {
released_ = true;
return std::move(f_);
}
bool is_released() const { return released_; }
private:
static_assert(!std::is_reference<F>::value, "F must not be a reference");
bool released_ = false;
F f_;
};
template <int&... ExplicitParameterBarrier, typename F,
typename DecayF = typename std::decay<F>::type>
Cleanup<DecayF> MakeCleanup(F&& f) {
return Cleanup<DecayF>(std::forward<F>(f));
}
} // namespace tf_gcs_filesystem
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_CLEANUP_H_
tensorflow/c/experimental/filesystem/plugins/gcs/expiring_lru_cache.h
+191
View File
@@ -0,0 +1,191 @@
/* Copyright 2020 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_EXPIRING_LRU_CACHE_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_EXPIRING_LRU_CACHE_H_
#include <functional>
#include <list>
#include <map>
#include <memory>
#include <string>
#include "absl/base/thread_annotations.h"
#include "absl/synchronization/mutex.h"
#include "tensorflow/c/env.h"
#include "tensorflow/c/tf_status.h"
namespace tf_gcs_filesystem {
/// \brief An LRU cache of string keys and arbitrary values, with configurable
/// max item age (in seconds) and max entries.
///
/// This class is thread safe.
template <typename T>
class ExpiringLRUCache {
public:
/// A `max_age` of 0 means that nothing is cached. A `max_entries` of 0 means
/// that there is no limit on the number of entries in the cache (however, if
/// `max_age` is also 0, the cache will not be populated).
ExpiringLRUCache(uint64_t max_age, size_t max_entries,
std::function<uint64_t()> timer_seconds = TF_NowSeconds)
: max_age_(max_age),
max_entries_(max_entries),
timer_seconds_(timer_seconds) {}
/// Insert `value` with key `key`. This will replace any previous entry with
/// the same key.
void Insert(const std::string& key, const T& value) {
if (max_age_ == 0) {
return;
}
absl::MutexLock lock(mu_);
InsertLocked(key, value);
}
// Delete the entry with key `key`. Return true if the entry was found for
// `key`, false if the entry was not found. In both cases, there is no entry
// with key `key` existed after the call.
bool Delete(const std::string& key) {
absl::MutexLock lock(mu_);
return DeleteLocked(key);
}
/// Look up the entry with key `key` and copy it to `value` if found. Returns
/// true if an entry was found for `key`, and its timestamp is not more than
/// max_age_ seconds in the past.
bool Lookup(const std::string& key, T* value) {
if (max_age_ == 0) {
return false;
}
absl::MutexLock lock(mu_);
return LookupLocked(key, value);
}
typedef std::function<void(const std::string&, T*, TF_Status*)> ComputeFunc;
/// Look up the entry with key `key` and copy it to `value` if found. If not
/// found, call `compute_func`. If `compute_func` set `status` to `TF_OK`,
/// store a copy of the output parameter in the cache, and another copy in
/// `value`.
void LookupOrCompute(const std::string& key, T* value,
const ComputeFunc& compute_func, TF_Status* status) {
if (max_age_ == 0) {
return compute_func(key, value, status);
}
// Note: we hold onto mu_ for the rest of this function. In practice, this
// is okay, as stat requests are typically fast, and concurrent requests are
// often for the same file. Future work can split this up into one lock per
// key if this proves to be a significant performance bottleneck.
absl::MutexLock lock(mu_);
if (LookupLocked(key, value)) {
return TF_SetStatus(status, TF_OK, "");
}
compute_func(key, value, status);
if (TF_GetCode(status) == TF_OK) {
InsertLocked(key, *value);
}
}
/// Clear the cache.
void Clear() {
absl::MutexLock lock(mu_);
cache_.clear();
lru_list_.clear();
}
/// Accessors for cache parameters.
uint64_t max_age() const { return max_age_; }
size_t max_entries() const { return max_entries_; }
private:
struct Entry {
/// The timestamp (seconds) at which the entry was added to the cache.
uint64_t timestamp;
/// The entry's value.
T value;
/// A list iterator pointing to the entry's position in the LRU list.
std::list<std::string>::iterator lru_iterator;
};
bool LookupLocked(const std::string& key, T* value)
ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
auto it = cache_.find(key);
if (it == cache_.end()) {
return false;
}
lru_list_.erase(it->second.lru_iterator);
if (timer_seconds_() - it->second.timestamp > max_age_) {
cache_.erase(it);
return false;
}
*value = it->second.value;
lru_list_.push_front(it->first);
it->second.lru_iterator = lru_list_.begin();
return true;
}
void InsertLocked(const std::string& key, const T& value)
ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
lru_list_.push_front(key);
Entry entry{timer_seconds_(), value, lru_list_.begin()};
auto insert = cache_.insert(std::make_pair(key, entry));
if (!insert.second) {
lru_list_.erase(insert.first->second.lru_iterator);
insert.first->second = entry;
} else if (max_entries_ > 0 && cache_.size() > max_entries_) {
cache_.erase(lru_list_.back());
lru_list_.pop_back();
}
}
bool DeleteLocked(const std::string& key) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
auto it = cache_.find(key);
if (it == cache_.end()) {
return false;
}
lru_list_.erase(it->second.lru_iterator);
cache_.erase(it);
return true;
}
/// The maximum age of entries in the cache, in seconds. A value of 0 means
/// that no entry is ever placed in the cache.
const uint64_t max_age_;
/// The maximum number of entries in the cache. A value of 0 means there is no
/// limit on entry count.
const size_t max_entries_;
/// The callback to read timestamps.
std::function<uint64_t()> timer_seconds_;
/// Guards access to the cache and the LRU list.
absl::Mutex mu_;
/// The cache (a map from string key to Entry).
std::map<std::string, Entry> cache_ ABSL_GUARDED_BY(mu_);
/// The LRU list of entries. The front of the list identifies the most
/// recently accessed entry.
std::list<std::string> lru_list_ ABSL_GUARDED_BY(mu_);
};
} // namespace tf_gcs_filesystem
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_EXPIRING_LRU_CACHE_H_
tensorflow/c/experimental/filesystem/plugins/gcs/expiring_lru_cache_test.cc
+216
View File
@@ -0,0 +1,216 @@
/* Copyright 2020 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/c/experimental/filesystem/plugins/gcs/expiring_lru_cache.h"
#include <cstdint>
#include <memory>
#include <string>
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_status_internal.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/cloud/now_seconds_env.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
TEST(ExpiringLRUCacheTest, MaxAge) {
const std::string key = "a";
std::unique_ptr<NowSecondsEnv> env(new NowSecondsEnv);
tf_gcs_filesystem::ExpiringLRUCache<int> cache(
1, 0, [&env]() { return env->NowSeconds(); });
env->SetNowSeconds(1);
// Verify that replacement of an existing element works, and updates the
// timestamp of the entry.
cache.Insert(key, 41);
env->SetNowSeconds(2);
cache.Insert(key, 42);
// 1 second after the most recent insertion, the entry is still valid.
env->SetNowSeconds(3);
int value = 0;
EXPECT_TRUE(cache.Lookup(key, &value));
EXPECT_EQ(value, 42);
// 2 seconds after the most recent insertion, the entry is no longer valid.
env->SetNowSeconds(4);
EXPECT_FALSE(cache.Lookup(key, &value));
// Re-insert the entry.
cache.Insert(key, 43);
EXPECT_TRUE(cache.Lookup(key, &value));
EXPECT_EQ(value, 43);
// The entry is valid 1 second after the insertion...
env->SetNowSeconds(5);
value = 0;
EXPECT_TRUE(cache.Lookup(key, &value));
EXPECT_EQ(value, 43);
// ...but is no longer valid 2 seconds after the insertion.
env->SetNowSeconds(6);
EXPECT_FALSE(cache.Lookup(key, &value));
}
TEST(ExpiringLRUCacheTest, MaxEntries) {
// max_age of 0 means nothing will be cached.
tf_gcs_filesystem::ExpiringLRUCache<int> cache1(0, 4);
cache1.Insert("a", 1);
int value = 0;
EXPECT_FALSE(cache1.Lookup("a", &value));
// Now set max_age = 1 and verify the LRU eviction logic.
tf_gcs_filesystem::ExpiringLRUCache<int> cache2(1, 4);
cache2.Insert("a", 1);
cache2.Insert("b", 2);
cache2.Insert("c", 3);
cache2.Insert("d", 4);
EXPECT_TRUE(cache2.Lookup("a", &value));
EXPECT_EQ(value, 1);
EXPECT_TRUE(cache2.Lookup("b", &value));
EXPECT_EQ(value, 2);
EXPECT_TRUE(cache2.Lookup("c", &value));
EXPECT_EQ(value, 3);
EXPECT_TRUE(cache2.Lookup("d", &value));
EXPECT_EQ(value, 4);
// Insertion of "e" causes "a" to be evicted, but the other entries are still
// there.
cache2.Insert("e", 5);
EXPECT_FALSE(cache2.Lookup("a", &value));
EXPECT_TRUE(cache2.Lookup("b", &value));
EXPECT_EQ(value, 2);
EXPECT_TRUE(cache2.Lookup("c", &value));
EXPECT_EQ(value, 3);
EXPECT_TRUE(cache2.Lookup("d", &value));
EXPECT_EQ(value, 4);
EXPECT_TRUE(cache2.Lookup("e", &value));
EXPECT_EQ(value, 5);
}
TEST(ExpiringLRUCacheTest, LookupOrCompute) {
// max_age of 0 means we should always compute.
uint64_t num_compute_calls = 0;
tf_gcs_filesystem::ExpiringLRUCache<int>::ComputeFunc compute_func =
[&num_compute_calls](const std::string& key, int* value,
TF_Status* status) {
*value = num_compute_calls;
num_compute_calls++;
return TF_SetStatus(status, TF_OK, "");
};
tf_gcs_filesystem::ExpiringLRUCache<int> cache1(0, 4);
int value = -1;
TF_Status status;
cache1.LookupOrCompute("a", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 0);
EXPECT_EQ(num_compute_calls, 1);
// re-read the same value, expect another lookup
cache1.LookupOrCompute("a", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 1);
EXPECT_EQ(num_compute_calls, 2);
// Define a new cache with max_age > 0 and verify correct behavior.
tf_gcs_filesystem::ExpiringLRUCache<int> cache2(2, 4);
num_compute_calls = 0;
value = -1;
// Read our first value
cache2.LookupOrCompute("a", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 0);
EXPECT_EQ(num_compute_calls, 1);
// Re-read, exepct no additional function compute_func calls.
cache2.LookupOrCompute("a", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 0);
EXPECT_EQ(num_compute_calls, 1);
// Read a sequence of additional values, eventually evicting "a".
cache2.LookupOrCompute("b", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 1);
EXPECT_EQ(num_compute_calls, 2);
cache2.LookupOrCompute("c", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 2);
EXPECT_EQ(num_compute_calls, 3);
cache2.LookupOrCompute("d", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 3);
EXPECT_EQ(num_compute_calls, 4);
cache2.LookupOrCompute("e", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 4);
EXPECT_EQ(num_compute_calls, 5);
// Verify the other values remain in the cache.
cache2.LookupOrCompute("b", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 1);
EXPECT_EQ(num_compute_calls, 5);
cache2.LookupOrCompute("c", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 2);
EXPECT_EQ(num_compute_calls, 5);
cache2.LookupOrCompute("d", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 3);
EXPECT_EQ(num_compute_calls, 5);
// Re-read "a", ensure it is re-computed.
cache2.LookupOrCompute("a", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 5);
EXPECT_EQ(num_compute_calls, 6);
}
TEST(ExpiringLRUCacheTest, Clear) {
tf_gcs_filesystem::ExpiringLRUCache<int> cache(1, 4);
cache.Insert("a", 1);
cache.Insert("b", 2);
cache.Insert("c", 3);
cache.Insert("d", 4);
int value = 0;
EXPECT_TRUE(cache.Lookup("a", &value));
EXPECT_EQ(value, 1);
EXPECT_TRUE(cache.Lookup("b", &value));
EXPECT_EQ(value, 2);
EXPECT_TRUE(cache.Lookup("c", &value));
EXPECT_EQ(value, 3);
EXPECT_TRUE(cache.Lookup("d", &value));
EXPECT_EQ(value, 4);
cache.Clear();
EXPECT_FALSE(cache.Lookup("a", &value));
EXPECT_FALSE(cache.Lookup("b", &value));
EXPECT_FALSE(cache.Lookup("c", &value));
EXPECT_FALSE(cache.Lookup("d", &value));
}
TEST(ExpiringLRUCacheTest, Delete) {
// Insert an entry.
tf_gcs_filesystem::ExpiringLRUCache<int> cache(1, 4);
cache.Insert("a", 1);
int value = 0;
EXPECT_TRUE(cache.Lookup("a", &value));
EXPECT_EQ(value, 1);
// Delete the entry.
EXPECT_TRUE(cache.Delete("a"));
EXPECT_FALSE(cache.Lookup("a", &value));
// Try deleting the entry again.
EXPECT_FALSE(cache.Delete("a"));
EXPECT_FALSE(cache.Lookup("a", &value));
}
} // namespace
} // namespace tensorflow
tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc
+1240
View File
File diff suppressed because it is too large Load Diff
tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.h
+117
View File
@@ -0,0 +1,117 @@
/* Copyright 2020 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_FILESYSTEM_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_FILESYSTEM_H_
#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include "absl/base/thread_annotations.h"
#include "absl/synchronization/mutex.h"
#include "google/cloud/storage/client.h"
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
#include "tensorflow/c/experimental/filesystem/plugins/gcs/expiring_lru_cache.h"
#include "tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h"
#include "tensorflow/c/tf_status.h"
void ParseGCSPath(const std::string& fname, bool object_empty_ok,
std::string* bucket, std::string* object, TF_Status* status);
namespace tf_random_access_file {
void Cleanup(TF_RandomAccessFile* file);
int64_t Read(const TF_RandomAccessFile* file, uint64_t offset, size_t n,
char* buffer, TF_Status* status);
} // namespace tf_random_access_file
namespace tf_writable_file {
void Cleanup(TF_WritableFile* file);
void Append(const TF_WritableFile* file, const char* buffer, size_t n,
TF_Status* status);
int64_t Tell(const TF_WritableFile* file, TF_Status* status);
void Flush(const TF_WritableFile* file, TF_Status* status);
void Sync(const TF_WritableFile* file, TF_Status* status);
void Close(const TF_WritableFile* file, TF_Status* status);
} // namespace tf_writable_file
namespace tf_read_only_memory_region {
void Cleanup(TF_ReadOnlyMemoryRegion* region);
const void* Data(const TF_ReadOnlyMemoryRegion* region);
uint64_t Length(const TF_ReadOnlyMemoryRegion* region);
} // namespace tf_read_only_memory_region
namespace tf_gcs_filesystem {
typedef struct GcsFileStat {
TF_FileStatistics base;
int64_t generation_number;
} GcsFileStat;
typedef struct GCSFile {
google::cloud::storage::Client gcs_client; // owned
bool compose;
absl::Mutex block_cache_lock;
std::shared_ptr<RamFileBlockCache> file_block_cache
ABSL_GUARDED_BY(block_cache_lock);
uint64_t block_size; // Reads smaller than block_size will trigger a read
// of block_size.
std::unique_ptr<ExpiringLRUCache<GcsFileStat>> stat_cache;
GCSFile(google::cloud::storage::Client&& gcs_client);
// This constructor is used for testing purpose only.
GCSFile(google::cloud::storage::Client&& gcs_client, bool compose,
uint64_t block_size, size_t max_bytes, uint64_t max_staleness,
uint64_t stat_cache_max_age, size_t stat_cache_max_entries);
} GCSFile;
// This function is used to initialize a filesystem without the need of setting
// manually environement variables.
void InitTest(TF_Filesystem* filesystem, bool compose, uint64_t block_size,
size_t max_bytes, uint64_t max_staleness,
uint64_t stat_cache_max_age, size_t stat_cache_max_entries,
TF_Status* status);
void Init(TF_Filesystem* filesystem, TF_Status* status);
void Cleanup(TF_Filesystem* filesystem);
void NewRandomAccessFile(const TF_Filesystem* filesystem, const char* path,
TF_RandomAccessFile* file, TF_Status* status);
void NewWritableFile(const TF_Filesystem* filesystem, const char* path,
TF_WritableFile* file, TF_Status* status);
void NewAppendableFile(const TF_Filesystem* filesystem, const char* path,
TF_WritableFile* file, TF_Status* status);
void NewReadOnlyMemoryRegionFromFile(const TF_Filesystem* filesystem,
const char* path,
TF_ReadOnlyMemoryRegion* region,
TF_Status* status);
int64_t GetFileSize(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
void PathExists(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
void CreateDir(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
int GetChildren(const TF_Filesystem* filesystem, const char* path,
char*** entries, TF_Status* status);
void DeleteFile(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
void Stat(const TF_Filesystem* filesystem, const char* path,
TF_FileStatistics* stats, TF_Status* status);
void DeleteDir(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
void CopyFile(const TF_Filesystem* filesystem, const char* src, const char* dst,
TF_Status* status);
void RenameFile(const TF_Filesystem* filesystem, const char* src,
const char* dst, TF_Status* status);
} // namespace tf_gcs_filesystem
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_FILESYSTEM_H_
tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem_test.cc
+705
View File
@@ -0,0 +1,705 @@
/* Copyright 2020 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/c/experimental/filesystem/plugins/gcs/gcs_filesystem.h"
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <iterator>
#include <memory>
#include <random>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include "absl/strings/match.h"
#include "absl/strings/string_view.h"
#include "third_party/cloud_cpp/google/cloud/common_options.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/stacktrace_handler.h"
#include "tensorflow/core/platform/test.h"
#define ASSERT_TF_OK(x) ASSERT_EQ(TF_OK, TF_GetCode(x)) << TF_Message(x)
#define EXPECT_TF_OK(x) EXPECT_EQ(TF_OK, TF_GetCode(x)) << TF_Message(x)
static const char* content = "abcdefghijklmnopqrstuvwxyz1234567890";
// We will work with content_view instead of content.
static const absl::string_view content_view = content;
namespace gcs = google::cloud::storage;
static std::string InitializeTmpDir() {
// This env should be something like `gs://bucket/path`
const char* test_dir = getenv("GCS_TEST_TMPDIR");
if (test_dir != nullptr) {
std::string bucket, object;
TF_Status* status = TF_NewStatus();
ParseGCSPath(test_dir, true, &bucket, &object, status);
if (TF_GetCode(status) != TF_OK) {
TF_DeleteStatus(status);
return "";
}
TF_DeleteStatus(status);
// We add a random value into `test_dir` to ensures that two consecutive
// runs are unlikely to clash.
std::random_device rd;
std::mt19937 gen(rd());
std::uniform_int_distribution<> distribution;
std::string rng_val = std::to_string(distribution(gen));
return tensorflow::io::JoinPath(std::string(test_dir), rng_val);
} else {
return "";
}
}
static std::string* GetTmpDir() {
static std::string tmp_dir = InitializeTmpDir();
if (tmp_dir == "")
return nullptr;
else
return &tmp_dir;
}
namespace tensorflow {
namespace {
// TODO(vnvo2409): Refactor `gcs_filesystem_test` to remove unnecessary tests
// after porting all tests from
// `//tensorflow/core/platform/cloud:gcs_file_system_test`.
class GCSFilesystemTest : public ::testing::Test {
public:
void SetUp() override {
root_dir_ = io::JoinPath(
*GetTmpDir(),
::testing::UnitTest::GetInstance()->current_test_info()->name());
status_ = TF_NewStatus();
filesystem_ = new TF_Filesystem;
filesystem_->plugin_filesystem = nullptr;
// Because different tests requires different setup for filesystem. We
// initialize filesystem in each testcase.
}
void TearDown() override {
TF_DeleteStatus(status_);
if (filesystem_->plugin_filesystem != nullptr)
tf_gcs_filesystem::Cleanup(filesystem_);
delete filesystem_;
}
std::string GetURIForPath(absl::string_view path) {
const std::string translated_name =
tensorflow::io::JoinPath(root_dir_, path);
return translated_name;
}
std::unique_ptr<TF_WritableFile, void (*)(TF_WritableFile* file)>
GetWriter() {
std::unique_ptr<TF_WritableFile, void (*)(TF_WritableFile * file)> writer(
new TF_WritableFile, [](TF_WritableFile* file) {
if (file != nullptr) {
if (file->plugin_file != nullptr) tf_writable_file::Cleanup(file);
delete file;
}
});
writer->plugin_file = nullptr;
return writer;
}
std::unique_ptr<TF_RandomAccessFile, void (*)(TF_RandomAccessFile* file)>
GetReader() {
std::unique_ptr<TF_RandomAccessFile, void (*)(TF_RandomAccessFile * file)>
reader(new TF_RandomAccessFile, [](TF_RandomAccessFile* file) {
if (file != nullptr) {
if (file->plugin_file != nullptr)
tf_random_access_file::Cleanup(file);
delete file;
}
});
reader->plugin_file = nullptr;
return reader;
}
void WriteString(const std::string& path, const std::string& content) {
auto writer = GetWriter();
tf_gcs_filesystem::NewWritableFile(filesystem_, path.c_str(), writer.get(),
status_);
if (TF_GetCode(status_) != TF_OK) return;
tf_writable_file::Append(writer.get(), content.c_str(), content.length(),
status_);
if (TF_GetCode(status_) != TF_OK) return;
tf_writable_file::Close(writer.get(), status_);
if (TF_GetCode(status_) != TF_OK) return;
}
std::string ReadAll(const std::string& path) {
auto reader = GetReader();
tf_gcs_filesystem::NewRandomAccessFile(filesystem_, path.c_str(),
reader.get(), status_);
if (TF_GetCode(status_) != TF_OK) return "";
auto file_size =
tf_gcs_filesystem::GetFileSize(filesystem_, path.c_str(), status_);
if (TF_GetCode(status_) != TF_OK) return "";
std::string content;
content.resize(file_size);
auto read = tf_random_access_file::Read(reader.get(), 0, file_size,
&content[0], status_);
if (TF_GetCode(status_) != TF_OK) return "";
if (read >= 0) content.resize(read);
if (file_size != content.size())
TF_SetStatus(
status_, TF_DATA_LOSS,
std::string("expected " + std::to_string(file_size) + " got " +
std::to_string(content.size()) + " bytes")
.c_str());
return content;
}
protected:
TF_Filesystem* filesystem_;
TF_Status* status_;
private:
std::string root_dir_;
};
::testing::AssertionResult WriteToServer(const std::string& path, size_t offset,
size_t length, gcs::Client* gcs_client,
TF_Status* status) {
std::string bucket, object;
ParseGCSPath(path, false, &bucket, &object, status);
if (TF_GetCode(status) != TF_OK)
return ::testing::AssertionFailure() << TF_Message(status);
auto writer = gcs_client->WriteObject(bucket, object);
writer.write(content + offset, length);
writer.Close();
if (writer.metadata()) {
return ::testing::AssertionSuccess();
} else {
return ::testing::AssertionFailure()
<< writer.metadata().status().message();
}
}
::testing::AssertionResult InsertObject(const std::string& path,
const std::string& content,
gcs::Client* gcs_client,
TF_Status* status) {
std::string bucket, object;
ParseGCSPath(path, false, &bucket, &object, status);
if (TF_GetCode(status) != TF_OK)
return ::testing::AssertionFailure() << TF_Message(status);
auto metadata = gcs_client->InsertObject(bucket, object, content);
if (metadata)
return ::testing::AssertionSuccess();
else
return ::testing::AssertionFailure() << metadata.status().message();
}
::testing::AssertionResult CompareSubString(int64_t offset, size_t length,
absl::string_view result,
size_t read) {
// Result isn't a null-terminated string so we have to wrap it inside a
// `string_view`
if (length == read && content_view.substr(offset, length) ==
absl::string_view(result).substr(0, read))
return ::testing::AssertionSuccess();
else
return ::testing::AssertionFailure()
<< "Result: " << absl::string_view(result).substr(0, read)
<< " Read: " << read;
}
::testing::AssertionResult CompareWithServer(const std::string& path,
size_t offset, size_t length,
gcs::Client* gcs_client,
TF_Status* status) {
std::string bucket, object;
ParseGCSPath(path, false, &bucket, &object, status);
if (TF_GetCode(status) != TF_OK)
return ::testing::AssertionFailure() << TF_Message(status);
auto reader = gcs_client->ReadObject(bucket, object);
if (!reader) {
return ::testing::AssertionFailure() << reader.status().message();
} else {
std::string content{std::istreambuf_iterator<char>{reader}, {}};
return CompareSubString(offset, length, content, content.length());
}
}
TEST_F(GCSFilesystemTest, ParseGCSPath) {
std::string bucket, object;
ParseGCSPath("gs://bucket/path/to/object", false, &bucket, &object, status_);
ASSERT_TF_OK(status_);
ASSERT_EQ(bucket, "bucket");
ASSERT_EQ(object, "path/to/object");
ParseGCSPath("gs://bucket/", true, &bucket, &object, status_);
ASSERT_TF_OK(status_);
ASSERT_EQ(bucket, "bucket");
ParseGCSPath("bucket/path/to/object", false, &bucket, &object, status_);
ASSERT_EQ(TF_GetCode(status_), TF_INVALID_ARGUMENT);
// bucket name must end with "/"
ParseGCSPath("gs://bucket", true, &bucket, &object, status_);
ASSERT_EQ(TF_GetCode(status_), TF_INVALID_ARGUMENT);
ParseGCSPath("gs://bucket/", false, &bucket, &object, status_);
ASSERT_EQ(TF_GetCode(status_), TF_INVALID_ARGUMENT);
}
TEST_F(GCSFilesystemTest, RandomAccessFile) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
<< TF_Message(status_);
std::string filepath = GetURIForPath("a_file");
TF_RandomAccessFile* file = new TF_RandomAccessFile;
tf_gcs_filesystem::NewRandomAccessFile(filesystem_, filepath.c_str(), file,
status_);
ASSERT_TF_OK(status_);
char* result = new char[content_view.length()];
int64_t read = tf_random_access_file::Read(file, 0, 1, result, status_);
ASSERT_EQ(read, -1) << "Read: " << read;
ASSERT_EQ(TF_GetCode(status_), TF_NOT_FOUND) << TF_Message(status_);
TF_SetStatus(status_, TF_OK, "");
auto gcs_file =
static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
ASSERT_TRUE(WriteToServer(filepath, 0, content_view.length(),
&gcs_file->gcs_client, status_));
read = tf_random_access_file::Read(file, 0, content_view.length(), result,
status_);
ASSERT_TF_OK(status_);
ASSERT_TRUE(CompareSubString(0, content_view.length(), result, read));
read = tf_random_access_file::Read(file, 0, 4, result, status_);
ASSERT_TF_OK(status_);
ASSERT_TRUE(CompareSubString(0, 4, result, read));
read = tf_random_access_file::Read(file, content_view.length() - 2, 4, result,
status_);
ASSERT_EQ(TF_GetCode(status_), TF_OUT_OF_RANGE) << TF_Message(status_);
ASSERT_TRUE(CompareSubString(content_view.length() - 2, 2, result, read));
delete[] result;
tf_random_access_file::Cleanup(file);
delete file;
}
TEST_F(GCSFilesystemTest, WritableFile) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
<< TF_Message(status_);
std::string filepath = GetURIForPath("a_file");
TF_WritableFile* file = new TF_WritableFile;
tf_gcs_filesystem::NewWritableFile(filesystem_, filepath.c_str(), file,
status_);
ASSERT_TF_OK(status_);
tf_writable_file::Append(file, content, 4, status_);
ASSERT_TF_OK(status_);
auto length = tf_writable_file::Tell(file, status_);
ASSERT_EQ(length, 4);
ASSERT_TF_OK(status_);
tf_writable_file::Flush(file, status_);
ASSERT_TF_OK(status_);
auto gcs_file =
static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
ASSERT_TRUE(
CompareWithServer(filepath, 0, 4, &gcs_file->gcs_client, status_));
tf_writable_file::Append(file, content + 4, 4, status_);
ASSERT_TF_OK(status_);
length = tf_writable_file::Tell(file, status_);
ASSERT_EQ(length, 8);
ASSERT_TF_OK(status_);
tf_writable_file::Flush(file, status_);
ASSERT_TF_OK(status_);
ASSERT_TRUE(
CompareWithServer(filepath, 0, 8, &gcs_file->gcs_client, status_));
tf_writable_file::Close(file, status_);
ASSERT_TF_OK(status_);
tf_writable_file::Cleanup(file);
// Testing for compose objects
gcs_file->compose = true;
filepath = GetURIForPath("b_file");
tf_gcs_filesystem::NewWritableFile(filesystem_, filepath.c_str(), file,
status_);
ASSERT_TF_OK(status_);
tf_writable_file::Append(file, content, 4, status_);
ASSERT_TF_OK(status_);
length = tf_writable_file::Tell(file, status_);
ASSERT_EQ(length, 4);
ASSERT_TF_OK(status_);
tf_writable_file::Flush(file, status_);
ASSERT_TF_OK(status_);
ASSERT_TRUE(
CompareWithServer(filepath, 0, 4, &gcs_file->gcs_client, status_));
tf_writable_file::Append(file, content + 4, 4, status_);
ASSERT_TF_OK(status_);
length = tf_writable_file::Tell(file, status_);
ASSERT_EQ(length, 8);
ASSERT_TF_OK(status_);
tf_writable_file::Flush(file, status_);
ASSERT_TF_OK(status_);
ASSERT_TRUE(
CompareWithServer(filepath, 0, 8, &gcs_file->gcs_client, status_));
tf_writable_file::Close(file, status_);
ASSERT_TF_OK(status_);
tf_writable_file::Cleanup(file);
delete file;
}
TEST_F(GCSFilesystemTest, ReadOnlyMemoryRegion) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
<< TF_Message(status_);
std::string path = GetURIForPath("a_file");
auto gcs_file =
static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
ASSERT_TRUE(WriteToServer(path, 0, 0, &gcs_file->gcs_client, status_));
TF_ReadOnlyMemoryRegion* region = new TF_ReadOnlyMemoryRegion;
tf_gcs_filesystem::NewReadOnlyMemoryRegionFromFile(filesystem_, path.c_str(),
region, status_);
ASSERT_EQ(TF_GetCode(status_), TF_INVALID_ARGUMENT) << TF_Message(status_);
TF_SetStatus(status_, TF_OK, "");
ASSERT_TRUE(WriteToServer(path, 0, content_view.length(),
&gcs_file->gcs_client, status_));
tf_gcs_filesystem::NewReadOnlyMemoryRegionFromFile(filesystem_, path.c_str(),
region, status_);
ASSERT_TF_OK(status_);
auto length = tf_read_only_memory_region::Length(region);
ASSERT_EQ(length, content_view.length());
auto data =
static_cast<const char*>(tf_read_only_memory_region::Data(region));
ASSERT_TRUE(CompareSubString(0, content_view.length(), data, length));
tf_read_only_memory_region::Cleanup(region);
delete region;
}
TEST_F(GCSFilesystemTest, PathExists) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_);
const std::string path = GetURIForPath("PathExists");
tf_gcs_filesystem::PathExists(filesystem_, path.c_str(), status_);
EXPECT_EQ(TF_NOT_FOUND, TF_GetCode(status_)) << TF_Message(status_);
TF_SetStatus(status_, TF_OK, "");
WriteString(path, "test");
ASSERT_TF_OK(status_);
tf_gcs_filesystem::PathExists(filesystem_, path.c_str(), status_);
EXPECT_TF_OK(status_);
}
TEST_F(GCSFilesystemTest, GetChildren) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_);
const std::string base = GetURIForPath("GetChildren");
tf_gcs_filesystem::CreateDir(filesystem_, base.c_str(), status_);
EXPECT_TF_OK(status_);
const std::string file = io::JoinPath(base, "TestFile.csv");
WriteString(file, "test");
EXPECT_TF_OK(status_);
const std::string subdir = io::JoinPath(base, "SubDir");
tf_gcs_filesystem::CreateDir(filesystem_, subdir.c_str(), status_);
EXPECT_TF_OK(status_);
const std::string subfile = io::JoinPath(subdir, "TestSubFile.csv");
WriteString(subfile, "test");
EXPECT_TF_OK(status_);
char** entries;
auto num_entries = tf_gcs_filesystem::GetChildren(filesystem_, base.c_str(),
&entries, status_);
EXPECT_TF_OK(status_);
std::vector<std::string> childrens;
for (int i = 0; i < num_entries; ++i) {
childrens.push_back(entries[i]);
}
std::sort(childrens.begin(), childrens.end());
EXPECT_EQ(std::vector<std::string>({"SubDir/", "TestFile.csv"}), childrens);
}
TEST_F(GCSFilesystemTest, DeleteFile) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_);
const std::string path = GetURIForPath("DeleteFile");
WriteString(path, "test");
ASSERT_TF_OK(status_);
tf_gcs_filesystem::DeleteFile(filesystem_, path.c_str(), status_);
EXPECT_TF_OK(status_);
tf_gcs_filesystem::PathExists(filesystem_, path.c_str(), status_);
EXPECT_EQ(TF_GetCode(status_), TF_NOT_FOUND);
}
TEST_F(GCSFilesystemTest, CreateDir) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_);
const std::string dir = GetURIForPath("CreateDir");
tf_gcs_filesystem::CreateDir(filesystem_, dir.c_str(), status_);
EXPECT_TF_OK(status_);
TF_FileStatistics stat;
tf_gcs_filesystem::Stat(filesystem_, dir.c_str(), &stat, status_);
EXPECT_TF_OK(status_);
EXPECT_TRUE(stat.is_directory);
}
TEST_F(GCSFilesystemTest, DeleteDir) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_);
const std::string dir = GetURIForPath("DeleteDir");
const std::string file = io::JoinPath(dir, "DeleteDirFile.csv");
WriteString(file, "test");
ASSERT_TF_OK(status_);
tf_gcs_filesystem::DeleteDir(filesystem_, dir.c_str(), status_);
EXPECT_EQ(TF_GetCode(status_), TF_FAILED_PRECONDITION);
TF_SetStatus(status_, TF_OK, "");
tf_gcs_filesystem::DeleteFile(filesystem_, file.c_str(), status_);
EXPECT_TF_OK(status_);
tf_gcs_filesystem::DeleteDir(filesystem_, dir.c_str(), status_);
EXPECT_TF_OK(status_);
TF_FileStatistics stat;
tf_gcs_filesystem::Stat(filesystem_, dir.c_str(), &stat, status_);
EXPECT_EQ(TF_GetCode(status_), TF_NOT_FOUND) << TF_Message(status_);
}
TEST_F(GCSFilesystemTest, StatFile) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_);
const std::string path = GetURIForPath("StatFile");
WriteString(path, "test");
ASSERT_TF_OK(status_);
TF_FileStatistics stat;
tf_gcs_filesystem::Stat(filesystem_, path.c_str(), &stat, status_);
EXPECT_TF_OK(status_);
EXPECT_EQ(4, stat.length);
EXPECT_FALSE(stat.is_directory);
}
TEST_F(GCSFilesystemTest, RenameFile) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_);
const std::string src = GetURIForPath("RenameFileSrc");
const std::string dst = GetURIForPath("RenameFileDst");
WriteString(src, "test");
ASSERT_TF_OK(status_);
tf_gcs_filesystem::RenameFile(filesystem_, src.c_str(), dst.c_str(), status_);
EXPECT_TF_OK(status_);
auto result = ReadAll(dst);
EXPECT_TF_OK(status_);
EXPECT_EQ("test", result);
}
TEST_F(GCSFilesystemTest, RenameFileOverwrite) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_);
const std::string src = GetURIForPath("RenameFileOverwriteSrc");
const std::string dst = GetURIForPath("RenameFileOverwriteDst");
WriteString(src, "test_old");
ASSERT_TF_OK(status_);
WriteString(dst, "test_new");
ASSERT_TF_OK(status_);
tf_gcs_filesystem::PathExists(filesystem_, dst.c_str(), status_);
EXPECT_TF_OK(status_);
tf_gcs_filesystem::RenameFile(filesystem_, src.c_str(), dst.c_str(), status_);
EXPECT_TF_OK(status_);
auto result = ReadAll(dst);
EXPECT_TF_OK(status_);
EXPECT_EQ("test_old", result);
}
// These tests below are ported from
// `//tensorflow/core/platform/cloud:gcs_file_system_test`
TEST_F(GCSFilesystemTest, NewRandomAccessFile_NoBlockCache) {
tf_gcs_filesystem::InitTest(filesystem_, false, 0, 0, 0, 0, 0, status_);
ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
<< TF_Message(status_);
std::string path = GetURIForPath("a_file");
auto gcs_file =
static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
ASSERT_TRUE(InsertObject(path, "0123456789", &gcs_file->gcs_client, status_));
TF_RandomAccessFile* file = new TF_RandomAccessFile;
tf_gcs_filesystem::NewRandomAccessFile(filesystem_, path.c_str(), file,
status_);
ASSERT_TF_OK(status_);
std::string result;
result.resize(6);
int64_t read = tf_random_access_file::Read(file, 0, 6, &result[0], status_);
ASSERT_EQ(read, 6) << "Read: " << read << "\n";
ASSERT_TF_OK(status_);
ASSERT_EQ(result, "012345") << "Result: " << result << "\n";
read = tf_random_access_file::Read(file, 6, 6, &result[0], status_);
ASSERT_EQ(read, 4) << "Read: " << read << "\n";
ASSERT_EQ(TF_GetCode(status_), TF_OUT_OF_RANGE) << TF_Message(status_);
result.resize(read);
ASSERT_EQ(result, "6789") << "Result: " << result << "\n";
}
TEST_F(GCSFilesystemTest, NewRandomAccessFile_Buffered) {
tf_gcs_filesystem::InitTest(filesystem_, false, 10, 0, 0, 0, 0, status_);
ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
<< TF_Message(status_);
std::string path = GetURIForPath("a_file");
auto gcs_file =
static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
ASSERT_TRUE(InsertObject(path, "0123456789", &gcs_file->gcs_client, status_));
TF_RandomAccessFile* file = new TF_RandomAccessFile;
tf_gcs_filesystem::NewRandomAccessFile(filesystem_, path.c_str(), file,
status_);
ASSERT_TF_OK(status_);
std::string result;
result.resize(6);
int64_t read = tf_random_access_file::Read(file, 0, 6, &result[0], status_);
ASSERT_EQ(read, 6) << "Read: " << read << "\n";
ASSERT_TF_OK(status_);
ASSERT_EQ(result, "012345") << "Result: " << result << "\n";
read = tf_random_access_file::Read(file, 6, 6, &result[0], status_);
ASSERT_EQ(read, 4) << "Read: " << read << "\n";
ASSERT_EQ(TF_GetCode(status_), TF_OUT_OF_RANGE) << TF_Message(status_);
result.resize(read);
ASSERT_EQ(result, "6789") << "Result: " << result << "\n";
}
TEST_F(GCSFilesystemTest, NewRandomAccessFile_Buffered_ReadAtEOF) {
tf_gcs_filesystem::InitTest(filesystem_, false, 10, 0, 0, 0, 0, status_);
ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
<< TF_Message(status_);
std::string path = GetURIForPath("a_file");
auto gcs_file =
static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
ASSERT_TRUE(InsertObject(path, "0123456789", &gcs_file->gcs_client, status_));
TF_RandomAccessFile* file = new TF_RandomAccessFile;
tf_gcs_filesystem::NewRandomAccessFile(filesystem_, path.c_str(), file,
status_);
ASSERT_TF_OK(status_);
std::string result;
result.resize(10);
int64_t read = tf_random_access_file::Read(file, 0, result.length(),
&result[0], status_);
ASSERT_EQ(read, 10) << "Read: " << read << "\n";
ASSERT_TF_OK(status_);
ASSERT_EQ(result, "0123456789") << "Result: " << result << "\n";
read = tf_random_access_file::Read(file, result.length(), result.length(),
&result[0], status_);
ASSERT_EQ(read, 0) << "Read: " << read << "\n";
ASSERT_EQ(TF_GetCode(status_), TF_OUT_OF_RANGE) << TF_Message(status_);
result.resize(read);
ASSERT_EQ(result, "") << "Result: " << result << "\n";
}
TEST_F(GCSFilesystemTest, NewRandomAccessFile_Buffered_CachedOutOfRange) {
tf_gcs_filesystem::InitTest(filesystem_, false, 10, 0, 0, 0, 0, status_);
ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
<< TF_Message(status_);
std::string path = GetURIForPath("a_file");
auto gcs_file =
static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
ASSERT_TRUE(InsertObject(path, "012345678", &gcs_file->gcs_client, status_));
TF_RandomAccessFile* file = new TF_RandomAccessFile;
tf_gcs_filesystem::NewRandomAccessFile(filesystem_, path.c_str(), file,
status_);
ASSERT_TF_OK(status_);
std::string result;
result.resize(5);
int64_t read = tf_random_access_file::Read(file, 0, result.length(),
&result[0], status_);
ASSERT_EQ(read, 5) << "Read: " << read << "\n";
ASSERT_TF_OK(status_);
ASSERT_EQ(result, "01234") << "Result: " << result << "\n";
read = tf_random_access_file::Read(file, 4, result.length(), &result[0],
status_);
ASSERT_EQ(read, 5) << "Read: " << read << "\n";
ASSERT_TF_OK(status_);
result.resize(read);
ASSERT_EQ(result, "45678") << "Result: " << result << "\n";
read = tf_random_access_file::Read(file, 5, result.length(), &result[0],
status_);
ASSERT_EQ(read, 4) << "Read: " << read << "\n";
ASSERT_EQ(TF_GetCode(status_), TF_OUT_OF_RANGE) << TF_Message(status_);
result.resize(read);
ASSERT_EQ(result, "5678") << "Result: " << result << "\n";
}
TEST_F(GCSFilesystemTest, TestGetStorageClientOptions) {
tf_gcs_filesystem::InitTest(filesystem_, false, 10, 0, 0, 0, 0, status_);
ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
<< TF_Message(status_);
auto gcs_file =
static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
// Get the raw connection
auto connection = gcs_file->gcs_client.raw_client();
// Get the options from the connection
auto options = connection->options();
// Access the UserAgentProductsOptions which contain information about
// the UserAgent
auto user_agent_products =
options.get<google::cloud::UserAgentProductsOption>();
bool found_x_goog_api_client = false;
for (const auto& product : user_agent_products) {
if (absl::StrContains(product, "TensorFlow-C-API")) {
found_x_goog_api_client = true;
break;
}
}
EXPECT_TRUE(found_x_goog_api_client);
}
} // namespace
} // namespace tensorflow
GTEST_API_ int main(int argc, char** argv) {
tensorflow::testing::InstallStacktraceHandler();
if (!GetTmpDir()) {
std::cerr << "Could not read GCS_TEST_TMPDIR env";
return -1;
}
::testing::InitGoogleTest(&argc, argv);
return RUN_ALL_TESTS();
}
tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.cc
+42
View File
@@ -0,0 +1,42 @@
/* Copyright 2020 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/c/experimental/filesystem/plugins/gcs/gcs_helper.h"
#include <stdio.h>
#include <cstdio>
#include <fstream>
#include <ios>
#include <string>
#include <utility>
TempFile::TempFile(const std::string& temp_file_name, std::ios::openmode mode)
: std::fstream(temp_file_name, mode), name_(temp_file_name) {}
TempFile::TempFile(TempFile&& rhs)
: std::fstream(std::move(rhs)), name_(std::move(rhs.name_)) {}
TempFile::~TempFile() {
std::fstream::close();
std::remove(name_.c_str());
}
const std::string TempFile::getName() const { return name_; }
bool TempFile::truncate() {
std::fstream::close();
std::fstream::open(name_, std::ios::binary | std::ios::out);
return std::fstream::is_open();
}
tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.h
+35
View File
@@ -0,0 +1,35 @@
/* Copyright 2020 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_HELPER_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_HELPER_H_
#include <fstream>
#include <ios>
#include <string>
class TempFile : public std::fstream {
public:
// We should specify openmode each time we call TempFile.
TempFile(const std::string& temp_file_name, std::ios::openmode mode);
TempFile(TempFile&& rhs);
~TempFile() override;
const std::string getName() const;
bool truncate();
private:
const std::string name_;
};
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_HELPER_H_
tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.cc
+316
View File
@@ -0,0 +1,316 @@
/* Copyright 2020 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/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h"
#include <cstdint>
#include <cstring>
#include <iostream>
#include <iterator>
#include <limits>
#include <memory>
#include <sstream>
#include <string>
#include <utility>
#include <vector>
#include "absl/synchronization/mutex.h"
#include "absl/time/time.h"
#include "tensorflow/c/experimental/filesystem/plugins/gcs/cleanup.h"
namespace tf_gcs_filesystem {
bool RamFileBlockCache::BlockNotStale(const std::shared_ptr<Block>& block) {
absl::MutexLock l(block->mu);
if (block->state != FetchState::FINISHED) {
return true; // No need to check for staleness.
}
if (max_staleness_ == 0) return true; // Not enforcing staleness.
return timer_seconds_() - block->timestamp <= max_staleness_;
}
std::shared_ptr<RamFileBlockCache::Block> RamFileBlockCache::Lookup(
const Key& key) {
absl::MutexLock lock(mu_);
auto entry = block_map_.find(key);
if (entry != block_map_.end()) {
if (BlockNotStale(entry->second)) {
return entry->second;
} else {
// Remove the stale block and continue.
RemoveFile_Locked(key.first);
}
}
// Insert a new empty block, setting the bookkeeping to sentinel values
// in order to update them as appropriate.
auto new_entry = std::make_shared<Block>();
lru_list_.push_front(key);
lra_list_.push_front(key);
new_entry->lru_iterator = lru_list_.begin();
new_entry->lra_iterator = lra_list_.begin();
new_entry->timestamp = timer_seconds_();
block_map_.emplace(std::make_pair(key, new_entry));
return new_entry;
}
// Remove blocks from the cache until we do not exceed our maximum size.
void RamFileBlockCache::Trim() {
while (!lru_list_.empty() && cache_size_ > max_bytes_) {
RemoveBlock(block_map_.find(lru_list_.back()));
}
}
/// Move the block to the front of the LRU list if it isn't already there.
void RamFileBlockCache::UpdateLRU(const Key& key,
const std::shared_ptr<Block>& block,
TF_Status* status) {
absl::MutexLock lock(mu_);
if (block->timestamp == 0) {
// The block was evicted from another thread. Allow it to remain evicted.
return TF_SetStatus(status, TF_OK, "");
}
if (block->lru_iterator != lru_list_.begin()) {
lru_list_.erase(block->lru_iterator);
lru_list_.push_front(key);
block->lru_iterator = lru_list_.begin();
}
// Check for inconsistent state. If there is a block later in the same file
// in the cache, and our current block is not block size, this likely means
// we have inconsistent state within the cache. Note: it's possible some
// incomplete reads may still go undetected.
if (block->data.size() < block_size_) {
Key fmax = std::make_pair(key.first, std::numeric_limits<size_t>::max());
auto fcmp = block_map_.upper_bound(fmax);
if (fcmp != block_map_.begin() && key < (--fcmp)->first) {
return TF_SetStatus(status, TF_INTERNAL,
"Block cache contents are inconsistent.");
}
}
Trim();
return TF_SetStatus(status, TF_OK, "");
}
void RamFileBlockCache::MaybeFetch(const Key& key,
const std::shared_ptr<Block>& block,
TF_Status* status) {
bool downloaded_block = false;
auto reconcile_state = MakeCleanup([this, &downloaded_block, &key, &block] {
// Perform this action in a cleanup callback to avoid locking mu_ after
// locking block->mu.
if (downloaded_block) {
absl::MutexLock l(mu_);
// Do not update state if the block is already to be evicted.
if (block->timestamp != 0) {
// Use capacity() instead of size() to account for all memory
// used by the cache.
cache_size_ += block->data.capacity();
// Put to beginning of LRA list.
lra_list_.erase(block->lra_iterator);
lra_list_.push_front(key);
block->lra_iterator = lra_list_.begin();
block->timestamp = timer_seconds_();
}
}
});
// Loop until either block content is successfully fetched, or our request
// encounters an error.
absl::MutexLock l(block->mu);
TF_SetStatus(status, TF_OK, "");
while (true) {
switch (block->state) {
case FetchState::ERROR:
// TF_FALLTHROUGH_INTENDED
case FetchState::CREATED:
block->state = FetchState::FETCHING;
block->mu.unlock(); // Release the lock while making the API call.
block->data.clear();
block->data.resize(block_size_, 0);
int64_t bytes_transferred;
bytes_transferred = block_fetcher_(key.first, key.second, block_size_,
block->data.data(), status);
block->mu.lock(); // Reacquire the lock immediately afterwards
if (TF_GetCode(status) == TF_OK) {
block->data.resize(bytes_transferred, 0);
// Shrink the data capacity to the actual size used.
// NOLINTNEXTLINE: shrink_to_fit() may not shrink the capacity.
std::vector<char>(block->data).swap(block->data);
downloaded_block = true;
block->state = FetchState::FINISHED;
} else {
block->state = FetchState::ERROR;
}
block->cond_var.SignalAll();
return;
case FetchState::FETCHING:
block->cond_var.WaitWithTimeout(&block->mu, absl::Minutes(1));
if (block->state == FetchState::FINISHED) {
return TF_SetStatus(status, TF_OK, "");
}
// Re-loop in case of errors.
break;
case FetchState::FINISHED:
return TF_SetStatus(status, TF_OK, "");
}
}
return TF_SetStatus(
status, TF_INTERNAL,
"Control flow should never reach the end of RamFileBlockCache::Fetch.");
}
int64_t RamFileBlockCache::Read(const std::string& filename, size_t offset,
size_t n, char* buffer, TF_Status* status) {
if (n == 0) {
TF_SetStatus(status, TF_OK, "");
return 0;
}
if (!IsCacheEnabled() || (n > max_bytes_)) {
// The cache is effectively disabled, so we pass the read through to the
// fetcher without breaking it up into blocks.
return block_fetcher_(filename, offset, n, buffer, status);
}
// Calculate the block-aligned start and end of the read.
size_t start = block_size_ * (offset / block_size_);
size_t finish = block_size_ * ((offset + n) / block_size_);
if (finish < offset + n) {
finish += block_size_;
}
size_t total_bytes_transferred = 0;
// Now iterate through the blocks, reading them one at a time.
for (size_t pos = start; pos < finish; pos += block_size_) {
Key key = std::make_pair(filename, pos);
// Look up the block, fetching and inserting it if necessary, and update the
// LRU iterator for the key and block.
std::shared_ptr<Block> block = Lookup(key);
if (!block) {
std::cerr << "No block for key " << key.first << "@" << key.second;
abort();
}
MaybeFetch(key, block, status);
if (TF_GetCode(status) != TF_OK) return -1;
UpdateLRU(key, block, status);
if (TF_GetCode(status) != TF_OK) return -1;
// Copy the relevant portion of the block into the result buffer.
const auto& data = block->data;
if (offset >= pos + data.size()) {
// The requested offset is at or beyond the end of the file. This can
// happen if `offset` is not block-aligned, and the read returns the last
// block in the file, which does not extend all the way out to `offset`.
std::stringstream os;
os << "EOF at offset " << offset << " in file " << filename
<< " at position " << pos << " with data size " << data.size();
TF_SetStatus(status, TF_OUT_OF_RANGE, std::move(os).str().c_str());
return total_bytes_transferred;
}
auto begin = data.begin();
if (offset > pos) {
// The block begins before the slice we're reading.
begin += offset - pos;
}
auto end = data.end();
if (pos + data.size() > offset + n) {
// The block extends past the end of the slice we're reading.
end -= (pos + data.size()) - (offset + n);
}
if (begin < end) {
size_t bytes_to_copy = end - begin;
memcpy(&buffer[total_bytes_transferred], &*begin, bytes_to_copy);
total_bytes_transferred += bytes_to_copy;
}
if (data.size() < block_size_) {
// The block was a partial block and thus signals EOF at its upper bound.
break;
}
}
TF_SetStatus(status, TF_OK, "");
return total_bytes_transferred;
}
bool RamFileBlockCache::ValidateAndUpdateFileSignature(
const std::string& filename, int64_t file_signature) {
absl::MutexLock lock(mu_);
auto it = file_signature_map_.find(filename);
if (it != file_signature_map_.end()) {
if (it->second == file_signature) {
return true;
}
// Remove the file from cache if the signatures don't match.
RemoveFile_Locked(filename);
it->second = file_signature;
return false;
}
file_signature_map_[filename] = file_signature;
return true;
}
size_t RamFileBlockCache::CacheSize() const {
absl::MutexLock lock(mu_);
return cache_size_;
}
void RamFileBlockCache::Prune() {
while (!stop_pruning_thread_.WaitForNotificationWithTimeout(
absl::Microseconds(1000000))) {
absl::MutexLock lock(mu_);
uint64_t now = timer_seconds_();
while (!lra_list_.empty()) {
auto it = block_map_.find(lra_list_.back());
if (now - it->second->timestamp <= max_staleness_) {
// The oldest block is not yet expired. Come back later.
break;
}
// We need to make a copy of the filename here, since it could otherwise
// be used within RemoveFile_Locked after `it` is deleted.
RemoveFile_Locked(std::string(it->first.first));
}
}
}
void RamFileBlockCache::Flush() {
absl::MutexLock lock(mu_);
block_map_.clear();
lru_list_.clear();
lra_list_.clear();
cache_size_ = 0;
}
void RamFileBlockCache::RemoveFile(const std::string& filename) {
absl::MutexLock lock(mu_);
RemoveFile_Locked(filename);
}
void RamFileBlockCache::RemoveFile_Locked(const std::string& filename) {
Key begin = std::make_pair(filename, 0);
auto it = block_map_.lower_bound(begin);
while (it != block_map_.end() && it->first.first == filename) {
auto next = std::next(it);
RemoveBlock(it);
it = next;
}
}
void RamFileBlockCache::RemoveBlock(BlockMap::iterator entry) {
// This signals that the block is removed, and should not be inadvertently
// reinserted into the cache in UpdateLRU.
entry->second->timestamp = 0;
lru_list_.erase(entry->second->lru_iterator);
lra_list_.erase(entry->second->lra_iterator);
cache_size_ -= entry->second->data.capacity();
block_map_.erase(entry);
}
} // namespace tf_gcs_filesystem
tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h
+270
View File
@@ -0,0 +1,270 @@
/* Copyright 2020 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_RAM_FILE_BLOCK_CACHE_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_RAM_FILE_BLOCK_CACHE_H_
#include <cstddef>
#include <cstdint>
#include <functional>
#include <iostream>
#include <list>
#include <map>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/base/thread_annotations.h"
#include "absl/log/log.h"
#include "absl/strings/str_format.h"
#include "absl/synchronization/mutex.h"
#include "absl/synchronization/notification.h"
#include "tensorflow/c/env.h"
#include "tensorflow/c/tf_status.h"
namespace tf_gcs_filesystem {
/// \brief An LRU block cache of file contents, keyed by {filename, offset}.
///
/// This class should be shared by read-only random access files on a remote
/// filesystem (e.g. GCS).
class RamFileBlockCache {
public:
/// The callback executed when a block is not found in the cache, and needs to
/// be fetched from the backing filesystem. This callback is provided when the
/// cache is constructed. It returns total bytes read ( -1 in case of errors
/// ). The `status` should be `TF_OK` as long as the read from the remote
/// filesystem succeeded (similar to the semantics of the read(2) system
/// call).
typedef std::function<int64_t(const std::string& filename, size_t offset,
size_t buffer_size, char* buffer,
TF_Status* status)>
BlockFetcher;
RamFileBlockCache(size_t block_size, size_t max_bytes, uint64_t max_staleness,
BlockFetcher block_fetcher,
std::function<uint64_t()> timer_seconds = TF_NowSeconds)
: block_size_(block_size),
max_bytes_(max_bytes),
max_staleness_(max_staleness),
block_fetcher_(block_fetcher),
timer_seconds_(timer_seconds),
pruning_thread_(nullptr,
[](TF_Thread* thread) { TF_JoinThread(thread); }) {
if (max_staleness_ > 0) {
TF_ThreadOptions thread_options;
TF_DefaultThreadOptions(&thread_options);
pruning_thread_.reset(
TF_StartThread(&thread_options, "TF_prune_FBC", PruneThread, this));
}
VLOG(1) << absl::StrFormat("GCS file block cache is %s.\n",
(IsCacheEnabled() ? "enabled" : "disabled"));
}
~RamFileBlockCache() {
if (pruning_thread_) {
stop_pruning_thread_.Notify();
// Destroying pruning_thread_ will block until Prune() receives the above
// notification and returns.
pruning_thread_.reset();
}
}
/// Read `n` bytes from `filename` starting at `offset` into `buffer`. It
/// returns total bytes read ( -1 in case of errors ). This method will set
/// `status` to:
///
/// 1) The error from the remote filesystem, if the read from the remote
/// filesystem failed.
/// 2) `TF_FAILED_PRECONDITION` if the read from the remote filesystem
/// succeeded,
/// but the read returned a partial block, and the LRU cache contained a
/// block at a higher offset (indicating that the partial block should have
/// been a full block).
/// 3) `TF_OUT_OF_RANGE` if the read from the remote filesystem succeeded, but
/// the file contents do not extend past `offset` and thus nothing was
/// placed in `out`.
/// 4) `TF_OK` otherwise (i.e. the read succeeded, and at least one byte was
/// placed
/// in `buffer`).
///
/// Caller is responsible for allocating memory for `buffer`.
/// `buffer` will be left unchanged in case of errors.
int64_t Read(const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status);
// Validate the given file signature with the existing file signature in the
// cache. Returns true if the signature doesn't change or the file doesn't
// exist before. If the signature changes, update the existing signature with
// the new one and remove the file from cache.
bool ValidateAndUpdateFileSignature(const std::string& filename,
int64_t file_signature)
ABSL_LOCKS_EXCLUDED(mu_);
/// Remove all cached blocks for `filename`.
void RemoveFile(const std::string& filename) ABSL_LOCKS_EXCLUDED(mu_);
/// Remove all cached data.
void Flush() ABSL_LOCKS_EXCLUDED(mu_);
/// Accessors for cache parameters.
size_t block_size() const { return block_size_; }
size_t max_bytes() const { return max_bytes_; }
uint64_t max_staleness() const { return max_staleness_; }
/// The current size (in bytes) of the cache.
size_t CacheSize() const ABSL_LOCKS_EXCLUDED(mu_);
// Returns true if the cache is enabled. If false, the BlockFetcher callback
// is always executed during Read.
bool IsCacheEnabled() const { return block_size_ > 0 && max_bytes_ > 0; }
// We can not pass a lambda with capture as a function pointer to
// `TF_StartThread`, so we have to wrap `Prune` inside a static function.
static void PruneThread(void* param) {
auto ram_file_block_cache = static_cast<RamFileBlockCache*>(param);
ram_file_block_cache->Prune();
}
private:
/// The size of the blocks stored in the LRU cache, as well as the size of the
/// reads from the underlying filesystem.
const size_t block_size_;
/// The maximum number of bytes (sum of block sizes) allowed in the LRU cache.
const size_t max_bytes_;
/// The maximum staleness of any block in the LRU cache, in seconds.
const uint64_t max_staleness_;
/// The callback to read a block from the underlying filesystem.
const BlockFetcher block_fetcher_;
/// The callback to read timestamps.
const std::function<uint64_t()> timer_seconds_;
/// \brief The key type for the file block cache.
///
/// The file block cache key is a {filename, offset} pair.
typedef std::pair<std::string, size_t> Key;
/// \brief The state of a block.
///
/// A block begins in the CREATED stage. The first thread will attempt to read
/// the block from the filesystem, transitioning the state of the block to
/// FETCHING. After completing, if the read was successful the state should
/// be FINISHED. Otherwise the state should be ERROR. A subsequent read can
/// re-fetch the block if the state is ERROR.
enum class FetchState {
CREATED,
FETCHING,
FINISHED,
ERROR,
};
/// \brief A block of a file.
///
/// A file block consists of the block data, the block's current position in
/// the LRU cache, the timestamp (seconds since epoch) at which the block
/// was cached, a coordination lock, and state & condition variables.
///
/// Thread safety:
/// The iterator and timestamp fields should only be accessed while holding
/// the block-cache-wide mu_ instance variable. The state variable should only
/// be accessed while holding the Block's mu lock. The data vector should only
/// be accessed after state == FINISHED, and it should never be modified.
///
/// In order to prevent deadlocks, never grab the block-cache-wide mu_ lock
/// AFTER grabbing any block's mu lock. It is safe to grab mu without locking
/// mu_.
struct Block {
/// The block data.
std::vector<char> data;
/// A list iterator pointing to the block's position in the LRU list.
std::list<Key>::iterator lru_iterator;
/// A list iterator pointing to the block's position in the LRA list.
std::list<Key>::iterator lra_iterator;
/// The timestamp (seconds since epoch) at which the block was cached.
uint64_t timestamp;
/// Mutex to guard state variable
absl::Mutex mu;
/// The state of the block.
FetchState state ABSL_GUARDED_BY(mu) = FetchState::CREATED;
/// Wait on cond_var if state is FETCHING.
absl::CondVar cond_var;
};
/// \brief The block map type for the file block cache.
///
/// The block map is an ordered map from Key to Block.
typedef std::map<Key, std::shared_ptr<Block>> BlockMap;
/// Prune the cache by removing files with expired blocks.
void Prune() ABSL_LOCKS_EXCLUDED(mu_);
bool BlockNotStale(const std::shared_ptr<Block>& block)
ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
/// Look up a Key in the block cache.
std::shared_ptr<Block> Lookup(const Key& key) ABSL_LOCKS_EXCLUDED(mu_);
void MaybeFetch(const Key& key, const std::shared_ptr<Block>& block,
TF_Status* status) ABSL_LOCKS_EXCLUDED(mu_);
/// Trim the block cache to make room for another entry.
void Trim() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
/// Update the LRU iterator for the block at `key`.
void UpdateLRU(const Key& key, const std::shared_ptr<Block>& block,
TF_Status* status) ABSL_LOCKS_EXCLUDED(mu_);
/// Remove all blocks of a file, with mu_ already held.
void RemoveFile_Locked(const std::string& filename)
ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
/// Remove the block `entry` from the block map and LRU list, and update the
/// cache size accordingly.
void RemoveBlock(BlockMap::iterator entry) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
/// The cache pruning thread that removes files with expired blocks.
std::unique_ptr<TF_Thread, std::function<void(TF_Thread*)>> pruning_thread_;
/// Notification for stopping the cache pruning thread.
absl::Notification stop_pruning_thread_;
/// Guards access to the block map, LRU list, and cached byte count.
mutable absl::Mutex mu_;
/// The block map (map from Key to Block).
BlockMap block_map_ ABSL_GUARDED_BY(mu_);
/// The LRU list of block keys. The front of the list identifies the most
/// recently accessed block.
std::list<Key> lru_list_ ABSL_GUARDED_BY(mu_);
/// The LRA (least recently added) list of block keys. The front of the list
/// identifies the most recently added block.
///
/// Note: blocks are added to lra_list_ only after they have successfully been
/// fetched from the underlying block store.
std::list<Key> lra_list_ ABSL_GUARDED_BY(mu_);
/// The combined number of bytes in all of the cached blocks.
size_t cache_size_ ABSL_GUARDED_BY(mu_) = 0;
// A filename->file_signature map.
std::map<std::string, int64_t> file_signature_map_ ABSL_GUARDED_BY(mu_);
};
} // namespace tf_gcs_filesystem
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_RAM_FILE_BLOCK_CACHE_H_
tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache_test.cc
+614
View File
@@ -0,0 +1,614 @@
/* Copyright 2020 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/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h"
#include <algorithm>
#include <cstdint>
#include <cstring>
#include <list>
#include <memory>
#include <set>
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/ascii.h"
#include "absl/synchronization/blocking_counter.h"
#include "absl/synchronization/notification.h"
#include "absl/time/time.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_status_internal.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/cloud/now_seconds_env.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
absl::Status ReadCache(tf_gcs_filesystem::RamFileBlockCache* cache,
const std::string& filename, size_t offset, size_t n,
std::vector<char>* out) {
out->clear();
out->resize(n, 0);
TF_Status status;
auto bytes_transferred =
cache->Read(filename, offset, n, out->data(), &status);
if (bytes_transferred >= 0) {
EXPECT_LE(bytes_transferred, n);
out->resize(bytes_transferred, n);
}
return status.status;
}
TEST(RamFileBlockCacheTest, IsCacheEnabled) {
auto fetcher = [](const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
// Do nothing.
TF_SetStatus(status, TF_OK, "");
return 0;
};
tf_gcs_filesystem::RamFileBlockCache cache1(0, 0, 0, fetcher);
tf_gcs_filesystem::RamFileBlockCache cache2(16, 0, 0, fetcher);
tf_gcs_filesystem::RamFileBlockCache cache3(0, 32, 0, fetcher);
tf_gcs_filesystem::RamFileBlockCache cache4(16, 32, 0, fetcher);
EXPECT_FALSE(cache1.IsCacheEnabled());
EXPECT_FALSE(cache2.IsCacheEnabled());
EXPECT_FALSE(cache3.IsCacheEnabled());
EXPECT_TRUE(cache4.IsCacheEnabled());
}
TEST(RamFileBlockCacheTest, ValidateAndUpdateFileSignature) {
int calls = 0;
auto fetcher = [&calls](const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
calls++;
memset(buffer, 'x', n);
TF_SetStatus(status, TF_OK, "");
return n;
};
std::string filename = "file";
tf_gcs_filesystem::RamFileBlockCache cache(16, 32, 0, fetcher);
std::vector<char> out;
// First read.
EXPECT_TRUE(cache.ValidateAndUpdateFileSignature(filename, 123));
TF_EXPECT_OK(ReadCache(&cache, filename, 0, 16, &out));
EXPECT_EQ(calls, 1);
// Second read. Hit cache.
EXPECT_TRUE(cache.ValidateAndUpdateFileSignature(filename, 123));
TF_EXPECT_OK(ReadCache(&cache, filename, 0, 16, &out));
EXPECT_EQ(calls, 1);
// Third read. File signatures are different.
EXPECT_FALSE(cache.ValidateAndUpdateFileSignature(filename, 321));
TF_EXPECT_OK(ReadCache(&cache, filename, 0, 16, &out));
EXPECT_EQ(calls, 2);
}
TEST(RamFileBlockCacheTest, PassThrough) {
const std::string want_filename = "foo/bar";
const size_t want_offset = 42;
const size_t want_n = 1024;
int calls = 0;
auto fetcher = [&calls, want_filename, want_offset, want_n](
const std::string& got_filename, size_t got_offset,
size_t got_n, char* buffer, TF_Status* status) -> int64_t {
EXPECT_EQ(got_filename, want_filename);
EXPECT_EQ(got_offset, want_offset);
EXPECT_EQ(got_n, want_n);
calls++;
memset(buffer, 'x', got_n);
TF_SetStatus(status, TF_OK, "");
return got_n;
};
// If block_size, max_bytes, or both are zero, or want_n is larger than
// max_bytes the cache is a pass-through.
tf_gcs_filesystem::RamFileBlockCache cache1(1, 0, 0, fetcher);
tf_gcs_filesystem::RamFileBlockCache cache2(0, 1, 0, fetcher);
tf_gcs_filesystem::RamFileBlockCache cache3(0, 0, 0, fetcher);
tf_gcs_filesystem::RamFileBlockCache cache4(1000, 1000, 0, fetcher);
std::vector<char> out;
TF_EXPECT_OK(ReadCache(&cache1, want_filename, want_offset, want_n, &out));
EXPECT_EQ(calls, 1);
TF_EXPECT_OK(ReadCache(&cache2, want_filename, want_offset, want_n, &out));
EXPECT_EQ(calls, 2);
TF_EXPECT_OK(ReadCache(&cache3, want_filename, want_offset, want_n, &out));
EXPECT_EQ(calls, 3);
TF_EXPECT_OK(ReadCache(&cache4, want_filename, want_offset, want_n, &out));
EXPECT_EQ(calls, 4);
}
TEST(RamFileBlockCacheTest, BlockAlignment) {
// Initialize a 256-byte buffer. This is the file underlying the reads we'll
// do in this test.
const size_t size = 256;
std::vector<char> buf;
buf.reserve(size);
for (int i = 0; i < size; i++) {
buf.push_back(i);
}
// The fetcher just fetches slices of the buffer.
auto fetcher = [&buf](const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
int64_t bytes_transferred;
if (offset < buf.size()) {
size_t bytes_to_copy = std::min<size_t>(buf.size() - offset, n);
memcpy(buffer, buf.data() + offset, bytes_to_copy);
bytes_transferred = bytes_to_copy;
} else {
bytes_transferred = 0;
}
TF_SetStatus(status, TF_OK, "");
return bytes_transferred;
};
for (size_t block_size = 2; block_size <= 4; block_size++) {
// Make a cache of N-byte block size (1 block) and verify that reads of
// varying offsets and lengths return correct data.
tf_gcs_filesystem::RamFileBlockCache cache(block_size, block_size, 0,
fetcher);
for (size_t offset = 0; offset < 10; offset++) {
for (size_t n = block_size - 2; n <= block_size + 2; n++) {
std::vector<char> got;
TF_EXPECT_OK(ReadCache(&cache, "", offset, n, &got));
// Verify the size of the read.
if (offset + n <= size) {
// Expect a full read.
EXPECT_EQ(got.size(), n) << "block size = " << block_size
<< ", offset = " << offset << ", n = " << n;
} else {
// Expect a partial read.
EXPECT_EQ(got.size(), size - offset)
<< "block size = " << block_size << ", offset = " << offset
<< ", n = " << n;
}
// Verify the contents of the read.
std::vector<char>::const_iterator begin = buf.begin() + offset;
std::vector<char>::const_iterator end =
offset + n > buf.size() ? buf.end() : begin + n;
std::vector<char> want(begin, end);
EXPECT_EQ(got, want) << "block size = " << block_size
<< ", offset = " << offset << ", n = " << n;
}
}
}
}
TEST(RamFileBlockCacheTest, CacheHits) {
const size_t block_size = 16;
std::set<size_t> calls;
auto fetcher = [&calls, block_size](const std::string& filename,
size_t offset, size_t n, char* buffer,
TF_Status* status) -> int64_t {
EXPECT_EQ(n, block_size);
EXPECT_EQ(offset % block_size, 0);
EXPECT_EQ(calls.find(offset), calls.end()) << "at offset " << offset;
calls.insert(offset);
memset(buffer, 'x', n);
TF_SetStatus(status, TF_OK, "");
return n;
};
const uint32_t block_count = 256;
tf_gcs_filesystem::RamFileBlockCache cache(
block_size, block_count * block_size, 0, fetcher);
std::vector<char> out;
out.resize(block_count, 0);
// The cache has space for `block_count` blocks. The loop with i = 0 should
// fill the cache, and the loop with i = 1 should be all cache hits. The
// fetcher checks that it is called once and only once for each offset (to
// fetch the corresponding block).
for (int i = 0; i < 2; i++) {
for (int j = 0; j < block_count; j++) {
TF_EXPECT_OK(ReadCache(&cache, "", block_size * j, block_size, &out));
}
}
}
TEST(RamFileBlockCacheTest, OutOfRange) {
// Tests reads of a 24-byte file with block size 16.
const size_t block_size = 16;
const size_t file_size = 24;
bool first_block = false;
bool second_block = false;
auto fetcher = [block_size, file_size, &first_block, &second_block](
const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
EXPECT_EQ(n, block_size);
EXPECT_EQ(offset % block_size, 0);
size_t bytes_to_copy = 0;
if (offset == 0) {
// The first block (16 bytes) of the file.
memset(buffer, 'x', n);
bytes_to_copy = n;
first_block = true;
} else if (offset == block_size) {
// The second block (8 bytes) of the file.
bytes_to_copy = file_size - block_size;
memset(buffer, 'x', bytes_to_copy);
second_block = true;
}
TF_SetStatus(status, TF_OK, "");
return bytes_to_copy;
};
tf_gcs_filesystem::RamFileBlockCache cache(block_size, block_size, 0,
fetcher);
std::vector<char> out;
// Reading the first 16 bytes should be fine.
TF_EXPECT_OK(ReadCache(&cache, "", 0, block_size, &out));
EXPECT_TRUE(first_block);
EXPECT_EQ(out.size(), block_size);
// Reading at offset file_size + 4 will read the second block (since the read
// at file_size + 4 = 28 will be aligned to an offset of 16) but will return
// OutOfRange because the offset is past the end of the 24-byte file.
absl::Status status = ReadCache(&cache, "", file_size + 4, 4, &out);
EXPECT_EQ(status.code(), error::OUT_OF_RANGE);
EXPECT_TRUE(second_block);
// Reading the second full block will return 8 bytes, from a cache hit.
second_block = false;
TF_EXPECT_OK(ReadCache(&cache, "", block_size, block_size, &out));
EXPECT_FALSE(second_block);
EXPECT_EQ(out.size(), file_size - block_size);
}
TEST(RamFileBlockCacheTest, Inconsistent) {
// Tests the detection of interrupted reads leading to partially filled blocks
// where we expected complete blocks.
const size_t block_size = 16;
// This fetcher returns OK but only fills in one byte for any offset.
auto fetcher = [block_size](const std::string& filename, size_t offset,
size_t n, char* buffer,
TF_Status* status) -> int64_t {
EXPECT_EQ(n, block_size);
EXPECT_EQ(offset % block_size, 0);
EXPECT_GE(n, 1);
memset(buffer, 'x', 1);
TF_SetStatus(status, TF_OK, "");
return 1;
};
tf_gcs_filesystem::RamFileBlockCache cache(block_size, 2 * block_size, 0,
fetcher);
std::vector<char> out;
// Read the second block; this should yield an OK status and a single byte.
TF_EXPECT_OK(ReadCache(&cache, "", block_size, block_size, &out));
EXPECT_EQ(out.size(), 1);
// Now read the first block; this should yield an INTERNAL error because we
// had already cached a partial block at a later position.
absl::Status status = ReadCache(&cache, "", 0, block_size, &out);
EXPECT_EQ(status.code(), error::INTERNAL);
}
TEST(RamFileBlockCacheTest, LRU) {
const size_t block_size = 16;
std::list<size_t> calls;
auto fetcher = [&calls, block_size](const std::string& filename,
size_t offset, size_t n, char* buffer,
TF_Status* status) -> int64_t {
EXPECT_EQ(n, block_size);
EXPECT_FALSE(calls.empty()) << "at offset = " << offset;
if (!calls.empty()) {
EXPECT_EQ(offset, calls.front());
calls.pop_front();
}
memset(buffer, 'x', n);
TF_SetStatus(status, TF_OK, "");
return n;
};
const uint32_t block_count = 2;
tf_gcs_filesystem::RamFileBlockCache cache(
block_size, block_count * block_size, 0, fetcher);
std::vector<char> out;
// Read blocks from the cache, and verify the LRU behavior based on the
// fetcher calls that the cache makes.
calls.push_back(0);
// Cache miss - drains an element from `calls`.
TF_EXPECT_OK(ReadCache(&cache, "", 0, 1, &out));
// Cache hit - does not drain an element from `calls`.
TF_EXPECT_OK(ReadCache(&cache, "", 0, 1, &out));
calls.push_back(block_size);
// Cache miss followed by cache hit.
TF_EXPECT_OK(ReadCache(&cache, "", block_size, 1, &out));
TF_EXPECT_OK(ReadCache(&cache, "", block_size, 1, &out));
calls.push_back(2 * block_size);
// Cache miss followed by cache hit. Causes eviction of LRU element.
TF_EXPECT_OK(ReadCache(&cache, "", 2 * block_size, 1, &out));
TF_EXPECT_OK(ReadCache(&cache, "", 2 * block_size, 1, &out));
// LRU element was at offset 0. Cache miss.
calls.push_back(0);
TF_EXPECT_OK(ReadCache(&cache, "", 0, 1, &out));
// Element at 2 * block_size is still in cache, and this read should update
// its position in the LRU list so it doesn't get evicted by the next read.
TF_EXPECT_OK(ReadCache(&cache, "", 2 * block_size, 1, &out));
// Element at block_size was evicted. Reading this element will also cause
// the LRU element (at 0) to be evicted.
calls.push_back(block_size);
TF_EXPECT_OK(ReadCache(&cache, "", block_size, 1, &out));
// Element at 0 was evicted again.
calls.push_back(0);
TF_EXPECT_OK(ReadCache(&cache, "", 0, 1, &out));
}
TEST(RamFileBlockCacheTest, MaxStaleness) {
int calls = 0;
auto fetcher = [&calls](const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
calls++;
memset(buffer, 'x', n);
TF_SetStatus(status, TF_OK, "");
return n;
};
std::vector<char> out;
std::unique_ptr<NowSecondsEnv> env(new NowSecondsEnv);
// Create a cache with max staleness of 2 seconds, and verify that it works as
// expected.
tf_gcs_filesystem::RamFileBlockCache cache1(
8, 16, 2 /* max staleness */, fetcher,
[&env]() { return env->NowSeconds(); });
// Execute the first read to load the block.
TF_EXPECT_OK(ReadCache(&cache1, "", 0, 1, &out));
EXPECT_EQ(calls, 1);
// Now advance the clock one second at a time and redo the read. The call
// count should advance every 3 seconds (i.e. every time the staleness is
// greater than 2).
for (int i = 1; i <= 10; i++) {
env->SetNowSeconds(i + 1);
TF_EXPECT_OK(ReadCache(&cache1, "", 0, 1, &out));
EXPECT_EQ(calls, 1 + i / 3);
}
// Now create a cache with max staleness of 0, and verify that it also works
// as expected.
calls = 0;
env->SetNowSeconds(0);
tf_gcs_filesystem::RamFileBlockCache cache2(
8, 16, 0 /* max staleness */, fetcher,
[&env]() { return env->NowSeconds(); });
// Execute the first read to load the block.
TF_EXPECT_OK(ReadCache(&cache2, "", 0, 1, &out));
EXPECT_EQ(calls, 1);
// Advance the clock by a huge amount and verify that the cached block is
// used to satisfy the read.
env->SetNowSeconds(365 * 24 * 60 * 60); // ~1 year, just for fun.
TF_EXPECT_OK(ReadCache(&cache2, "", 0, 1, &out));
EXPECT_EQ(calls, 1);
}
TEST(RamFileBlockCacheTest, RemoveFile) {
int calls = 0;
auto fetcher = [&calls](const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
calls++;
char c = (filename == "a") ? 'a' : (filename == "b") ? 'b' : 'x';
if (offset > 0) {
// The first block is lower case and all subsequent blocks are upper case.
c = absl::ascii_toupper(c);
}
memset(buffer, c, n);
TF_SetStatus(status, TF_OK, "");
return n;
};
// This cache has space for 4 blocks; we'll read from two files.
const size_t n = 3;
tf_gcs_filesystem::RamFileBlockCache cache(8, 32, 0, fetcher);
std::vector<char> out;
std::vector<char> a(n, 'a');
std::vector<char> b(n, 'b');
std::vector<char> A(n, 'A');
std::vector<char> B(n, 'B');
// Fill the cache.
TF_EXPECT_OK(ReadCache(&cache, "a", 0, n, &out));
EXPECT_EQ(out, a);
EXPECT_EQ(calls, 1);
TF_EXPECT_OK(ReadCache(&cache, "a", 8, n, &out));
EXPECT_EQ(out, A);
EXPECT_EQ(calls, 2);
TF_EXPECT_OK(ReadCache(&cache, "b", 0, n, &out));
EXPECT_EQ(out, b);
EXPECT_EQ(calls, 3);
TF_EXPECT_OK(ReadCache(&cache, "b", 8, n, &out));
EXPECT_EQ(out, B);
EXPECT_EQ(calls, 4);
// All four blocks should be in the cache now.
TF_EXPECT_OK(ReadCache(&cache, "a", 0, n, &out));
EXPECT_EQ(out, a);
TF_EXPECT_OK(ReadCache(&cache, "a", 8, n, &out));
EXPECT_EQ(out, A);
TF_EXPECT_OK(ReadCache(&cache, "b", 0, n, &out));
EXPECT_EQ(out, b);
TF_EXPECT_OK(ReadCache(&cache, "b", 8, n, &out));
EXPECT_EQ(out, B);
EXPECT_EQ(calls, 4);
// Remove the blocks from "a".
cache.RemoveFile("a");
// Both blocks from "b" should still be there.
TF_EXPECT_OK(ReadCache(&cache, "b", 0, n, &out));
EXPECT_EQ(out, b);
TF_EXPECT_OK(ReadCache(&cache, "b", 8, n, &out));
EXPECT_EQ(out, B);
EXPECT_EQ(calls, 4);
// The blocks from "a" should not be there.
TF_EXPECT_OK(ReadCache(&cache, "a", 0, n, &out));
EXPECT_EQ(out, a);
EXPECT_EQ(calls, 5);
TF_EXPECT_OK(ReadCache(&cache, "a", 8, n, &out));
EXPECT_EQ(out, A);
EXPECT_EQ(calls, 6);
}
TEST(RamFileBlockCacheTest, Prune) {
int calls = 0;
auto fetcher = [&calls](const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
calls++;
memset(buffer, 'x', n);
TF_SetStatus(status, TF_OK, "");
return n;
};
std::vector<char> out;
// Our fake environment is initialized with the current timestamp.
std::unique_ptr<NowSecondsEnv> env(new NowSecondsEnv);
uint64_t now = Env::Default()->NowSeconds();
env->SetNowSeconds(now);
tf_gcs_filesystem::RamFileBlockCache cache(
8, 32, 1 /* max staleness */, fetcher,
[&env]() { return env->NowSeconds(); });
// Read three blocks into the cache, and advance the timestamp by one second
// with each read. Start with a block of "a" at the current timestamp `now`.
TF_EXPECT_OK(ReadCache(&cache, "a", 0, 1, &out));
// Now load a block of a different file "b" at timestamp `now` + 1
env->SetNowSeconds(now + 1);
TF_EXPECT_OK(ReadCache(&cache, "b", 0, 1, &out));
// Now load a different block of file "a" at timestamp `now` + 1. When the
// first block of "a" expires, this block should also be removed because it
// also belongs to file "a".
TF_EXPECT_OK(ReadCache(&cache, "a", 8, 1, &out));
// Ensure that all blocks are in the cache (i.e. reads are cache hits).
EXPECT_EQ(cache.CacheSize(), 24);
EXPECT_EQ(calls, 3);
TF_EXPECT_OK(ReadCache(&cache, "a", 0, 1, &out));
TF_EXPECT_OK(ReadCache(&cache, "b", 0, 1, &out));
TF_EXPECT_OK(ReadCache(&cache, "a", 8, 1, &out));
EXPECT_EQ(calls, 3);
// Advance the fake timestamp so that "a" becomes stale via its first block.
env->SetNowSeconds(now + 2);
// The pruning thread periodically compares env->NowSeconds() with the oldest
// block's timestamp to see if it should evict any files. At the current fake
// timestamp of `now` + 2, file "a" is stale because its first block is stale,
// but file "b" is not stale yet. Thus, once the pruning thread wakes up (in
// one second of wall time), it should remove "a" and leave "b" alone.
uint64_t start = Env::Default()->NowSeconds();
do {
Env::Default()->SleepForMicroseconds(100000);
} while (cache.CacheSize() == 24 && Env::Default()->NowSeconds() - start < 3);
// There should be one block left in the cache, and it should be the first
// block of "b".
EXPECT_EQ(cache.CacheSize(), 8);
TF_EXPECT_OK(ReadCache(&cache, "b", 0, 1, &out));
EXPECT_EQ(calls, 3);
// Advance the fake time to `now` + 3, at which point "b" becomes stale.
env->SetNowSeconds(now + 3);
// Wait for the pruner to remove "b".
start = Env::Default()->NowSeconds();
do {
Env::Default()->SleepForMicroseconds(100000);
} while (cache.CacheSize() == 8 && Env::Default()->NowSeconds() - start < 3);
// The cache should now be empty.
EXPECT_EQ(cache.CacheSize(), 0);
}
TEST(RamFileBlockCacheTest, ParallelReads) {
// This fetcher won't respond until either `callers` threads are calling it
// concurrently (at which point it will respond with success to all callers),
// or 10 seconds have elapsed (at which point it will respond with an error).
const int callers = 4;
absl::BlockingCounter counter(callers);
absl::Notification notification;
auto fetcher = [&counter, &notification](
const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
if (counter.DecrementCount()) {
notification.Notify();
// This call to `Wait()` is not expected to block. Calling `Wait()` here
// allows us to satisfy `BlockingCounter`'s requirement: "When `Wait()`
// returns, it is legal to destroy the `BlockingCounter`.".
counter.Wait();
}
if (!notification.WaitForNotificationWithTimeout(absl::Seconds(10))) {
// This avoids having the test time out, which is harder to debug.
TF_SetStatus(status, TF_FAILED_PRECONDITION,
"desired concurrency not reached");
return -1;
}
memset(buffer, 'x', n);
TF_SetStatus(status, TF_OK, "");
return n;
};
const int block_size = 8;
tf_gcs_filesystem::RamFileBlockCache cache(
block_size, 2 * callers * block_size, 0, fetcher);
std::vector<std::unique_ptr<Thread>> threads;
threads.reserve(callers);
for (int i = 0; i < callers; i++) {
threads.emplace_back(
Env::Default()->StartThread({}, "caller", [block_size, &cache, i]() {
std::vector<char> out;
TF_EXPECT_OK(
ReadCache(&cache, "a", i * block_size, block_size, &out));
std::vector<char> x(block_size, 'x');
EXPECT_EQ(out, x);
}));
}
// The `threads` destructor blocks until the threads can be joined, once their
// respective reads finish (which happens once they are all concurrently being
// executed, or 10 seconds have passed).
}
TEST(RamFileBlockCacheTest, CoalesceConcurrentReads) {
// Concurrent reads to the same file blocks should be de-duplicated.
const size_t block_size = 16;
int num_requests = 0;
absl::Notification notification;
auto fetcher = [&num_requests, &notification, block_size](
const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
EXPECT_EQ(n, block_size);
EXPECT_EQ(offset, 0);
num_requests++;
memset(buffer, 'x', n);
notification.Notify();
// Wait for other thread to issue read.
Env::Default()->SleepForMicroseconds(100000); // 0.1 secs
TF_SetStatus(status, TF_OK, "");
return n;
};
tf_gcs_filesystem::RamFileBlockCache cache(block_size, block_size, 0,
fetcher);
// Fork off thread for parallel read.
std::unique_ptr<Thread> concurrent(
Env::Default()->StartThread({}, "concurrent", [block_size, &cache] {
std::vector<char> out;
TF_EXPECT_OK(ReadCache(&cache, "", 0, block_size / 2, &out));
EXPECT_EQ(out.size(), block_size / 2);
}));
notification.WaitForNotification();
std::vector<char> out;
TF_EXPECT_OK(ReadCache(&cache, "", block_size / 2, block_size / 2, &out));
EXPECT_EQ(out.size(), block_size / 2);
EXPECT_EQ(1, num_requests);
}
TEST(RamFileBlockCacheTest, Flush) {
int calls = 0;
auto fetcher = [&calls](const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
calls++;
memset(buffer, 'x', n);
TF_SetStatus(status, TF_OK, "");
return n;
};
tf_gcs_filesystem::RamFileBlockCache cache(16, 32, 0, fetcher);
std::vector<char> out;
TF_EXPECT_OK(ReadCache(&cache, "", 0, 16, &out));
TF_EXPECT_OK(ReadCache(&cache, "", 0, 16, &out));
EXPECT_EQ(calls, 1);
cache.Flush();
TF_EXPECT_OK(ReadCache(&cache, "", 0, 16, &out));
EXPECT_EQ(calls, 2);
}
} // namespace
} // namespace tensorflow
tensorflow/c/experimental/filesystem/plugins/posix/BUILD
+67
View File
@@ -0,0 +1,67 @@
# Experimental posix filesystem plugin.
load("//tensorflow:tensorflow.bzl", "tf_cc_shared_object")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:private"],
licenses = ["notice"],
)
# Filesystem implementation for POSIX environments: Linux, MacOS, Android, etc.
tf_cc_shared_object(
name = "libposix_filesystem.so",
framework_so = [],
linkstatic = False,
visibility = ["//visibility:public"],
deps = [":posix_filesystem_impl"],
)
# The real implementation of the filesystem.
cc_library(
name = "posix_filesystem_impl",
srcs = ["posix_filesystem.cc"],
hdrs = ["posix_filesystem.h"],
deps = [
":posix_filesystem_helper",
"//tensorflow/c:tf_file_statistics",
"//tensorflow/c:tf_status",
"//tensorflow/c/experimental/filesystem:filesystem_interface",
],
)
# Since building pip package and API tests require a filesystem, we provide a
# static registration target that they should link against.
cc_library(
name = "posix_filesystem_static",
srcs = ["posix_filesystem_static.cc"],
visibility = ["//visibility:public"],
deps = [
":posix_filesystem_impl",
"//tensorflow/c/experimental/filesystem:filesystem_interface",
"//tensorflow/c/experimental/filesystem:modular_filesystem",
"//tensorflow/core/platform:status",
"@com_google_absl//absl/log",
],
alwayslink = 1,
)
# Library implementing helper functionality, so that the above only contains
# the API implementation for modular filesystems.
cc_library(
name = "posix_filesystem_helper",
srcs = ["posix_filesystem_helper.cc"],
hdrs = ["posix_filesystem_helper.h"],
deps = [":copy_file"],
)
# On Linux, we can copy files faster using `sendfile`. But not elsewhere.
# Hence, this private library to select which implementation to use.
cc_library(
name = "copy_file",
srcs = select({
"//tensorflow:linux_x86_64": ["copy_file_linux.cc"],
"//conditions:default": ["copy_file_portable.cc"],
}),
hdrs = ["copy_file.h"],
)
tensorflow/c/experimental/filesystem/plugins/posix/copy_file.h
+32
View File
@@ -0,0 +1,32 @@
/* Copyright 2019 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_COPY_FILE_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_COPY_FILE_H_
#include <sys/stat.h>
namespace tf_posix_filesystem {
// Transfers up to `size` bytes from `dst_fd` to `src_fd`.
//
// This method uses `sendfile` if available (i.e., linux 2.6.33 or later) or an
// intermediate buffer if not.
//
// Returns number of bytes transferred or -1 on failure.
int CopyFileContents(int dst_fd, int src_fd, off_t size);
} // namespace tf_posix_filesystem
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_COPY_FILE_H_
tensorflow/c/experimental/filesystem/plugins/posix/copy_file_linux.cc
+46
View File
@@ -0,0 +1,46 @@
/* Copyright 2019 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 <limits.h>
#include <stdint.h>
#include <sys/sendfile.h>
#include <cstddef>
#include "tensorflow/c/experimental/filesystem/plugins/posix/copy_file.h"
namespace tf_posix_filesystem {
// Transfers up to `size` bytes from `dst_fd` to `src_fd`.
//
// This method uses `sendfile` specific to linux after 2.6.33.
int CopyFileContents(int dst_fd, int src_fd, off_t size) {
off_t offset = 0;
int bytes_transferred = 0;
int rc = 1;
// When `sendfile` returns 0 we stop copying and let callers handle this.
while (offset < size && rc > 0) {
// Use uint64 for safe compare SSIZE_MAX
uint64_t chunk = size - offset;
if (chunk > SSIZE_MAX) chunk = SSIZE_MAX;
rc = sendfile(dst_fd, src_fd, &offset, static_cast<size_t>(chunk));
if (rc < 0) return -1;
bytes_transferred += rc;
}
return bytes_transferred;
}
} // namespace tf_posix_filesystem
tensorflow/c/experimental/filesystem/plugins/posix/copy_file_portable.cc
+58
View File
@@ -0,0 +1,58 @@
/* Copyright 2019 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 <stdint.h>
#include <unistd.h>
#include <memory>
#include "tensorflow/c/experimental/filesystem/plugins/posix/copy_file.h"
namespace tf_posix_filesystem {
// Transfers up to `size` bytes from `dst_fd` to `src_fd`.
//
// This method uses a temporary buffer to hold contents.
int CopyFileContents(int dst_fd, int src_fd, off_t size) {
// Use a copy buffer of 128KB but don't store it on the stack.
constexpr static size_t kPosixCopyFileBufferSize = 128 * 1024;
std::unique_ptr<char[]> buffer(new char[kPosixCopyFileBufferSize]);
off_t offset = 0;
int bytes_transferred = 0;
int rc = 1;
// When `sendfile` returns 0 we stop copying and let callers handle this.
while (offset < size && rc > 0) {
size_t chunk = size - offset;
if (chunk > kPosixCopyFileBufferSize) chunk = kPosixCopyFileBufferSize;
rc = read(src_fd, buffer.get(), chunk);
if (rc < 0) return -1;
int total_write = 0;
int total_read = rc;
while (total_write < total_read && rc > 0) {
rc = write(dst_fd, buffer.get() + total_write, total_read - total_write);
if (rc < 0) return -1;
total_write += rc;
bytes_transferred += rc;
offset += rc;
}
}
return bytes_transferred;
}
} // namespace tf_posix_filesystem
tensorflow/c/experimental/filesystem/plugins/posix/posix_filesystem.cc
+461
View File
@@ -0,0 +1,461 @@
/* Copyright 2019 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/c/experimental/filesystem/plugins/posix/posix_filesystem.h"
#include <dirent.h>
#include <errno.h>
#include <fcntl.h>
#include <limits.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
#include "tensorflow/c/experimental/filesystem/plugins/posix/posix_filesystem_helper.h"
#include "tensorflow/c/tf_file_statistics.h"
#include "tensorflow/c/tf_status.h"
// 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 {
typedef struct PosixFile {
const char* filename;
int fd;
} PosixFile;
static void Cleanup(TF_RandomAccessFile* file) {
auto posix_file = static_cast<PosixFile*>(file->plugin_file);
close(posix_file->fd);
// 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;
}
static int64_t Read(const TF_RandomAccessFile* file, uint64_t offset, size_t n,
char* buffer, TF_Status* status) {
auto posix_file = static_cast<PosixFile*>(file->plugin_file);
char* dst = buffer;
int64_t read = 0;
while (n > 0) {
// Some platforms, notably macs, throw `EINVAL` if `pread` is asked to read
// more than fits in a 32-bit integer.
size_t requested_read_length;
if (n > INT32_MAX)
requested_read_length = INT32_MAX;
else
requested_read_length = n;
// `pread` returns a `ssize_t` on POSIX, but due to interface being
// cross-platform, return type of `Read` is `int64_t`.
int64_t r = int64_t{pread(posix_file->fd, dst, requested_read_length,
static_cast<off_t>(offset))};
if (r > 0) {
dst += r;
offset += static_cast<uint64_t>(r);
n -= r; // safe as 0 < r <= n so n will never underflow
read += r;
} else if (r == 0) {
TF_SetStatus(status, TF_OUT_OF_RANGE, "Read fewer bytes than requested");
break;
} else if (errno == EINTR || errno == EAGAIN) {
// Retry
} else {
TF_SetStatusFromIOError(status, errno, posix_file->filename);
break;
}
}
return read;
}
} // namespace tf_random_access_file
// SECTION 2. Implementation for `TF_WritableFile`
// ----------------------------------------------------------------------------
namespace tf_writable_file {
typedef struct PosixFile {
const char* filename;
FILE* handle;
} PosixFile;
static void Cleanup(TF_WritableFile* file) {
auto posix_file = static_cast<PosixFile*>(file->plugin_file);
plugin_memory_free(const_cast<char*>(posix_file->filename));
delete posix_file;
}
static void Append(const TF_WritableFile* file, const char* buffer, size_t n,
TF_Status* status) {
auto posix_file = static_cast<PosixFile*>(file->plugin_file);
size_t r = fwrite(buffer, 1, n, posix_file->handle);
if (r != n)
TF_SetStatusFromIOError(status, errno, posix_file->filename);
else
TF_SetStatus(status, TF_OK, "");
}
static int64_t Tell(const TF_WritableFile* file, TF_Status* status) {
auto posix_file = static_cast<PosixFile*>(file->plugin_file);
// POSIX's `ftell` returns `long`, do a manual cast.
int64_t position = int64_t{ftell(posix_file->handle)};
if (position < 0)
TF_SetStatusFromIOError(status, errno, posix_file->filename);
else
TF_SetStatus(status, TF_OK, "");
return position;
}
static void Flush(const TF_WritableFile* file, TF_Status* status) {
auto posix_file = static_cast<PosixFile*>(file->plugin_file);
TF_SetStatus(status, TF_OK, "");
if (fflush(posix_file->handle) != 0)
TF_SetStatusFromIOError(status, errno, posix_file->filename);
}
static void Sync(const TF_WritableFile* file, TF_Status* status) {
// For historical reasons, this does the same as `Flush` at the moment.
// TODO(b/144055243): This should use `fsync`/`sync`.
Flush(file, status);
}
static void Close(const TF_WritableFile* file, TF_Status* status) {
auto posix_file = static_cast<PosixFile*>(file->plugin_file);
if (fclose(posix_file->handle) != 0)
TF_SetStatusFromIOError(status, errno, posix_file->filename);
else
TF_SetStatus(status, TF_OK, "");
}
} // namespace tf_writable_file
// SECTION 3. Implementation for `TF_ReadOnlyMemoryRegion`
// ----------------------------------------------------------------------------
namespace tf_read_only_memory_region {
typedef struct PosixMemoryRegion {
const void* const address;
const uint64_t length;
} PosixMemoryRegion;
static void Cleanup(TF_ReadOnlyMemoryRegion* region) {
auto r = static_cast<PosixMemoryRegion*>(region->plugin_memory_region);
munmap(const_cast<void*>(r->address), r->length);
delete r;
}
static const void* Data(const TF_ReadOnlyMemoryRegion* region) {
auto r = static_cast<PosixMemoryRegion*>(region->plugin_memory_region);
return r->address;
}
static uint64_t Length(const TF_ReadOnlyMemoryRegion* region) {
auto r = static_cast<PosixMemoryRegion*>(region->plugin_memory_region);
return r->length;
}
} // namespace tf_read_only_memory_region
// SECTION 4. Implementation for `TF_Filesystem`, the actual filesystem
// ----------------------------------------------------------------------------
namespace tf_posix_filesystem {
static void Init(TF_Filesystem* filesystem, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
}
static void Cleanup(TF_Filesystem* filesystem) {}
static void NewRandomAccessFile(const TF_Filesystem* filesystem,
const char* path, TF_RandomAccessFile* file,
TF_Status* status) {
int fd = open(path, O_RDONLY);
if (fd < 0) {
TF_SetStatusFromIOError(status, errno, path);
return;
}
struct stat st;
fstat(fd, &st);
if (S_ISDIR(st.st_mode)) {
TF_SetStatus(status, TF_FAILED_PRECONDITION, "path is a directory");
close(fd);
return;
}
file->plugin_file = new tf_random_access_file::PosixFile({strdup(path), fd});
TF_SetStatus(status, TF_OK, "");
}
static void NewWritableFile(const TF_Filesystem* filesystem, const char* path,
TF_WritableFile* file, TF_Status* status) {
FILE* f = fopen(path, "w");
if (f == nullptr) {
TF_SetStatusFromIOError(status, errno, path);
return;
}
file->plugin_file = new tf_writable_file::PosixFile({strdup(path), f});
TF_SetStatus(status, TF_OK, "");
}
static void NewAppendableFile(const TF_Filesystem* filesystem, const char* path,
TF_WritableFile* file, TF_Status* status) {
FILE* f = fopen(path, "a");
if (f == nullptr) {
TF_SetStatusFromIOError(status, errno, path);
return;
}
file->plugin_file = new tf_writable_file::PosixFile({strdup(path), f});
TF_SetStatus(status, TF_OK, "");
}
static void NewReadOnlyMemoryRegionFromFile(const TF_Filesystem* filesystem,
const char* path,
TF_ReadOnlyMemoryRegion* region,
TF_Status* status) {
int fd = open(path, O_RDONLY);
if (fd < 0) {
TF_SetStatusFromIOError(status, errno, path);
return;
}
struct stat st;
fstat(fd, &st);
if (S_ISDIR(st.st_mode)) {
TF_SetStatus(status, TF_FAILED_PRECONDITION, "path is a directory");
} else {
const void* address =
mmap(nullptr, st.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
if (address == MAP_FAILED) {
TF_SetStatusFromIOError(status, errno, path);
} else {
region->plugin_memory_region =
new tf_read_only_memory_region::PosixMemoryRegion{
address, static_cast<uint64_t>(st.st_size)};
TF_SetStatus(status, TF_OK, "");
}
}
close(fd);
}
static void CreateDir(const TF_Filesystem* filesystem, const char* path,
TF_Status* status) {
if (strlen(path) == 0)
TF_SetStatus(status, TF_ALREADY_EXISTS, "already exists");
else if (mkdir(path, /*mode=*/0755) != 0)
TF_SetStatusFromIOError(status, errno, path);
else
TF_SetStatus(status, TF_OK, "");
}
static void DeleteFile(const TF_Filesystem* filesystem, const char* path,
TF_Status* status) {
if (unlink(path) != 0)
TF_SetStatusFromIOError(status, errno, path);
else
TF_SetStatus(status, TF_OK, "");
}
static void DeleteDir(const TF_Filesystem* filesystem, const char* path,
TF_Status* status) {
if (rmdir(path) != 0)
TF_SetStatusFromIOError(status, errno, path);
else
TF_SetStatus(status, TF_OK, "");
}
static void RenameFile(const TF_Filesystem* filesystem, const char* src,
const char* dst, TF_Status* status) {
// If target is a directory return TF_FAILED_PRECONDITION.
// Target might be missing, so don't error in that case.
struct stat st;
if (stat(dst, &st) != 0) {
if (errno != ENOENT) {
TF_SetStatusFromIOError(status, errno, dst);
return;
}
} else if (S_ISDIR(st.st_mode)) {
TF_SetStatus(status, TF_FAILED_PRECONDITION, "target path is a directory");
return;
}
// We cannot rename directories yet, so prevent this.
if (stat(src, &st) != 0) {
TF_SetStatusFromIOError(status, errno, src);
return;
} else if (S_ISDIR(st.st_mode)) {
TF_SetStatus(status, TF_FAILED_PRECONDITION, "source path is a directory");
return;
}
// Do the actual rename. Here both arguments are filenames.
if (rename(src, dst) != 0)
TF_SetStatusFromIOError(status, errno, dst);
else
TF_SetStatus(status, TF_OK, "");
}
static void CopyFile(const TF_Filesystem* filesystem, const char* src,
const char* dst, TF_Status* status) {
// If target is a directory return TF_FAILED_PRECONDITION.
// Target might be missing, so don't error in that case.
struct stat st;
if (stat(dst, &st) != 0) {
if (errno != ENOENT) {
TF_SetStatusFromIOError(status, errno, dst);
return;
}
} else if (S_ISDIR(st.st_mode)) {
TF_SetStatus(status, TF_FAILED_PRECONDITION, "target path is a directory");
return;
}
// We cannot copy directories yet, so prevent this.
if (stat(src, &st) != 0) {
TF_SetStatusFromIOError(status, errno, src);
return;
} else if (S_ISDIR(st.st_mode)) {
TF_SetStatus(status, TF_FAILED_PRECONDITION, "source path is a directory");
return;
}
// Both `src` and `dst` point to files here. Delegate to helper.
if (TransferFileContents(src, dst, st.st_mode, st.st_size) < 0)
TF_SetStatusFromIOError(status, errno, dst);
else
TF_SetStatus(status, TF_OK, "");
}
static void PathExists(const TF_Filesystem* filesystem, const char* path,
TF_Status* status) {
if (access(path, F_OK) != 0)
TF_SetStatusFromIOError(status, errno, path);
else
TF_SetStatus(status, TF_OK, "");
}
static void Stat(const TF_Filesystem* filesystem, const char* path,
TF_FileStatistics* stats, TF_Status* status) {
struct stat sbuf;
if (stat(path, &sbuf) != 0) {
TF_SetStatusFromIOError(status, errno, path);
} else {
stats->length = sbuf.st_size;
stats->mtime_nsec = sbuf.st_mtime * (1000 * 1000 * 1000);
stats->is_directory = S_ISDIR(sbuf.st_mode);
TF_SetStatus(status, TF_OK, "");
}
}
static int GetChildren(const TF_Filesystem* filesystem, const char* path,
char*** entries, TF_Status* status) {
struct dirent** dir_entries = nullptr;
/* we don't promise entries would be sorted */
int num_entries =
scandir(path, &dir_entries, RemoveSpecialDirectoryEntries, nullptr);
if (num_entries < 0) {
TF_SetStatusFromIOError(status, errno, path);
} else {
*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);
plugin_memory_free(dir_entries[i]);
}
plugin_memory_free(dir_entries);
}
return num_entries;
}
} // namespace tf_posix_filesystem
static void ProvideFilesystemSupportFor(TF_FilesystemPluginOps* ops,
const char* uri) {
TF_SetFilesystemVersionMetadata(ops);
ops->scheme = strdup(uri);
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;
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;
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;
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");
}
tensorflow/c/experimental/filesystem/plugins/posix/posix_filesystem.h
+31
View File
@@ -0,0 +1,31 @@
/* Copyright 2020 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_POSIX_FILESYSTEM_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_POSIX_FILESYSTEM_H_
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
// Initialize the POSIX filesystem.
//
// In general, the `TF_InitPlugin` symbol doesn't need to be exposed in a header
// file, since the plugin registration will look for the symbol in the DSO file
// that provides the filesystem functionality. However, the POSIX filesystem
// needs to be statically registered in some tests and utilities for building
// the API files at the time of creating the pip package. Hence, we need to
// expose this function so that this filesystem can be statically registered
// when needed.
void TF_InitPlugin(TF_FilesystemPluginInfo* info);
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_POSIX_FILESYSTEM_H_
tensorflow/c/experimental/filesystem/plugins/posix/posix_filesystem_helper.cc
+65
View File
@@ -0,0 +1,65 @@
/* Copyright 2019 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/c/experimental/filesystem/plugins/posix/posix_filesystem_helper.h"
#include <errno.h>
#include <fcntl.h>
#include <limits.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include "tensorflow/c/experimental/filesystem/plugins/posix/copy_file.h"
namespace tf_posix_filesystem {
int TransferFileContents(const char* src, const char* dst, mode_t mode,
off_t size) {
int src_fd = open(src, O_RDONLY);
if (src_fd < 0) return -1;
// When creating file, use the same permissions as original
mode_t open_mode = mode & (S_IRWXU | S_IRWXG | S_IRWXO);
// O_WRONLY | O_CREAT | O_TRUNC:
// Open file for write and if file does not exist, create the file.
// If file exists, truncate its size to 0.
int dst_fd = open(dst, O_WRONLY | O_CREAT | O_TRUNC, open_mode);
if (dst_fd < 0) {
close(src_fd);
return -1;
}
// Both files have been opened, do the transfer.
// Since errno would be overridden by `close` below, save it here.
int error_code = 0;
if (CopyFileContents(dst_fd, src_fd, size) < 0) error_code = errno;
close(src_fd);
close(dst_fd);
if (error_code != 0) {
errno = error_code;
return -1;
} else {
return 0;
}
}
int RemoveSpecialDirectoryEntries(const struct dirent* entry) {
return strcmp(entry->d_name, ".") != 0 && strcmp(entry->d_name, "..") != 0;
}
} // namespace tf_posix_filesystem
tensorflow/c/experimental/filesystem/plugins/posix/posix_filesystem_helper.h
+37
View File
@@ -0,0 +1,37 @@
/* Copyright 2019 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_POSIX_FILESYSTEM_HELPER_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_POSIX_FILESYSTEM_HELPER_H_
#include <dirent.h>
#include <sys/stat.h>
namespace tf_posix_filesystem {
// Copies up to `size` of `src` to `dst`, creating destination if needed.
//
// Callers should pass size of `src` in `size` and the permissions of `src` in
// `mode`. The later is only used if `dst` needs to be created.
int TransferFileContents(const char* src, const char* dst, mode_t mode,
off_t size);
// Returns true only if `entry` points to an entry other than `.` or `..`.
//
// This is a filter for `scandir`.
int RemoveSpecialDirectoryEntries(const struct dirent* entry);
} // namespace tf_posix_filesystem
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_POSIX_FILESYSTEM_HELPER_H_
tensorflow/c/experimental/filesystem/plugins/posix/posix_filesystem_static.cc
+41
View File
@@ -0,0 +1,41 @@
/* Copyright 2020 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 "absl/log/log.h"
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
#include "tensorflow/c/experimental/filesystem/modular_filesystem_registration.h"
#include "tensorflow/c/experimental/filesystem/plugins/posix/posix_filesystem.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
// Register the POSIX filesystems statically.
// Return value will be unused
bool StaticallyRegisterLocalFilesystems() {
TF_FilesystemPluginInfo info;
TF_InitPlugin(&info);
absl::Status status =
filesystem_registration::RegisterFilesystemPluginImpl(&info);
if (!status.ok()) {
VLOG(0) << "Static POSIX filesystem could not be registered: " << status;
return false;
}
return true;
}
// Perform the actual registration
static bool unused = StaticallyRegisterLocalFilesystems();
} // namespace tensorflow
tensorflow/c/experimental/filesystem/plugins/windows/BUILD
+35
View File
@@ -0,0 +1,35 @@
# Experimental windows filesystem plugin.
load("//tensorflow:tensorflow.bzl", "get_win_copts", "tf_cc_shared_object")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
# Filesystem implementation for Windows environment
tf_cc_shared_object(
name = "windows_filesystem.dll",
framework_so = [],
linkstatic = False,
tags = [
"manual",
"nobuilder",
"notap",
],
visibility = ["//visibility:public"],
deps = [":windows_filesystem_impl"],
)
# The real implementation of the filesystem.
cc_library(
name = "windows_filesystem_impl",
srcs = ["windows_filesystem.cc"],
copts = get_win_copts(),
tags = [
"manual",
"nobuilder",
"notap",
],
deps = ["//tensorflow/c/experimental/filesystem:filesystem_interface"],
)
tensorflow/c/experimental/filesystem/plugins/windows/windows_filesystem.cc
+72
View File
@@ -0,0 +1,72 @@
/* Copyright 2020 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 <stdlib.h>
#include <string.h>
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
// 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 {
// TODO(b/139060984): Implement later
} // namespace tf_random_access_file
// SECTION 2. Implementation for `TF_WritableFile`
// ----------------------------------------------------------------------------
namespace tf_writable_file {
// TODO(b/139060984): Implement later
} // namespace tf_writable_file
// SECTION 3. Implementation for `TF_ReadOnlyMemoryRegion`
// ----------------------------------------------------------------------------
namespace tf_read_only_memory_region {
// TODO(b/139060984): Implement later
} // namespace tf_read_only_memory_region
// SECTION 4. Implementation for `TF_Filesystem`, the actual filesystem
// ----------------------------------------------------------------------------
namespace tf_windows_filesystem {
// TODO(b/139060984): Implement later
} // namespace tf_windows_filesystem
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");
}