Merge pull request #37734 from tensorflow:mm-static-filesystem-registration

PiperOrigin-RevId: 302044106
Change-Id: I7c4aca9fe90accb26d9579425bb0c8d99eb5d359

tensorflow/c/experimental/filesystem/BUILD

@@ -24,9 +24,14 @@ cc_library(
     srcs = [
         "modular_filesystem.cc",
         "modular_filesystem_registration.cc",
+    ],
+    hdrs = [
+        "modular_filesystem.h",
         "modular_filesystem_registration.h",
     ],
-    hdrs = ["modular_filesystem.h"],
+    # TODO(mihaimaruseac): Visibility should be more restrictive once we
+    # convert to modular filesystems everywhere
+    visibility = ["//visibility:public"],
     deps = [
         ":filesystem_interface",
         "//tensorflow/c:tf_status_helper",

tensorflow/c/experimental/filesystem/modular_filesystem.cc

@@ -440,7 +440,25 @@ Status ModularWritableFile::Tell(int64* position) {
 }
 
 Status RegisterFilesystemPlugin(const std::string& dso_path) {
-  return filesystem_registration::RegisterFilesystemPluginImpl(dso_path);
+  // Step 1: Load plugin
+  Env* env = Env::Default();
+  void* dso_handle;
+  TF_RETURN_IF_ERROR(env->LoadLibrary(dso_path.c_str(), &dso_handle));
+
+  // Step 2: Load symbol for `TF_InitPlugin`
+  void* dso_symbol;
+  TF_RETURN_IF_ERROR(
+      env->GetSymbolFromLibrary(dso_handle, "TF_InitPlugin", &dso_symbol));
+
+  // 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_registration.cc

@@ -14,7 +14,6 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/c/experimental/filesystem/modular_filesystem_registration.h"
 
-#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
 #include "tensorflow/c/experimental/filesystem/modular_filesystem.h"
 #include "tensorflow/c/tf_status_internal.h"
 #include "tensorflow/core/platform/env.h"
@@ -304,40 +303,22 @@ static Status ValidatePluginMemoryRoutines(
 
 namespace filesystem_registration {
 
-Status RegisterFilesystemPluginImpl(const std::string& dso_path) {
-  // Step 1: Load plugin
-  Env* env = Env::Default();
-  void* dso_handle;
-  TF_RETURN_IF_ERROR(env->LoadLibrary(dso_path.c_str(), &dso_handle));
+Status RegisterFilesystemPluginImpl(const TF_FilesystemPluginInfo* info) {
+  TF_RETURN_IF_ERROR(ValidatePluginMemoryRoutines(info));
 
-  // Step 2: Load symbol for `TF_InitPlugin`
-  void* dso_symbol;
-  TF_RETURN_IF_ERROR(
-      env->GetSymbolFromLibrary(dso_handle, "TF_InitPlugin", &dso_symbol));
-
-  // 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: Ensure plugin provides the memory management functions.
-  TF_RETURN_IF_ERROR(ValidatePluginMemoryRoutines(&info));
-
-  // Step 5: Validate and register all filesystems
+  // 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);
+  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);
+  info->plugin_memory_free(info->ops);
   return status;
 }
 

tensorflow/c/experimental/filesystem/modular_filesystem_registration.h

@@ -15,12 +15,17 @@ limitations under the License.
 #ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_MODULAR_FILESYSTEM_REGISTRATION_H_
 #define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_MODULAR_FILESYSTEM_REGISTRATION_H_
 
+#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
 #include "tensorflow/core/platform/status.h"
 
 namespace tensorflow {
 namespace filesystem_registration {
 
-Status RegisterFilesystemPluginImpl(const std::string& dso_path);
+// Implementation for filesystem registration
+//
+// Don't call this directly. Instead call `RegisterFilesystemPlugin`.
+// Exposed only for static registration of local filesystems.
+Status RegisterFilesystemPluginImpl(const TF_FilesystemPluginInfo* info);
 
 }  // namespace filesystem_registration
 }  // namespace tensorflow

tensorflow/c/experimental/filesystem/plugins/posix/BUILD

@@ -19,6 +19,7 @@ tf_cc_shared_object(
 cc_library(
     name = "posix_filesystem_impl",
     srcs = ["posix_filesystem.cc"],
+    hdrs = ["posix_filesystem.h"],
     deps = [
         ":posix_filesystem_helper",
         "//tensorflow/c:tf_status",
@@ -26,6 +27,20 @@ cc_library(
     ],
 )
 
+# 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",
+    ],
+    alwayslink = 1,
+)
+
 # Library implementing helper functionality, so that the above only contains
 # the API implementation for modular filesystems.
 cc_library(

tensorflow/c/experimental/filesystem/plugins/posix/posix_filesystem.cc

@@ -12,6 +12,8 @@ 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>
@@ -24,7 +26,6 @@ limitations under the License.
 #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_status.h"
 

tensorflow/c/experimental/filesystem/plugins/posix/posix_filesystem.h

@@ -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_static.cc

@@ -0,0 +1,38 @@
+/* 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/filesystem_interface.h"
+#include "tensorflow/c/experimental/filesystem/modular_filesystem_registration.h"
+#include "tensorflow/c/experimental/filesystem/plugins/posix/posix_filesystem.h"
+
+namespace tensorflow {
+
+// Register the POSIX filesystems statically.
+// Return value will be unused
+bool StaticallyRegisterLocalFilesystems() {
+  TF_FilesystemPluginInfo info;
+  TF_InitPlugin(&info);
+  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