Make GCS ops work in open source

tensorflow/contrib/cloud/__init__.py

@@ -27,8 +27,9 @@
 
 _allowed_symbols = [
     'BigQueryReader',
-    'ConfigureColabSession',
-    'ConfigureGcs',
+    'BlockCacheParams',
+    'configure_colab_session',
+    'configure_gcs',
     'ConfigureGcsHook',
 ]
 remove_undocumented(__name__, _allowed_symbols)

tensorflow/contrib/cloud/kernels/BUILD

@@ -79,6 +79,7 @@ tf_kernel_library(
     srcs = ["gcs_config_ops.cc"],
     visibility = ["//tensorflow:internal"],
     deps = [
+        "//tensorflow/contrib/cloud:gcs_config_ops_op_lib",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core/platform/cloud:curl_http_request",

tensorflow/contrib/cloud/ops/gcs_config_ops.cc

@@ -21,50 +21,12 @@ namespace tensorflow {
 
 REGISTER_OP("GcsConfigureCredentials")
     .Input("json: string")
-    .SetShapeFn(shape_inference::NoOutputs)
-    .Doc(R"doc(
-Configures the credentials used by the GCS client of the local TF runtime.
-
-The json input can be of the format:
-
-1. Refresh Token:
-{
-  "client_id": "<redacted>",
-  "client_secret": "<redacted>",
-  "refresh_token: "<redacted>",
-  "type": "authorized_user",
-}
-
-2. Service Account:
-{
-  "type": "service_account",
-  "project_id": "<redacted>",
-  "private_key_id": "<redacted>",
-  "private_key": "------BEGIN PRIVATE KEY-----\n<REDACTED>\n-----END PRIVATE KEY------\n",
-  "client_email": "<REDACTED>@<REDACTED>.iam.gserviceaccount.com",
-  "client_id": "<REDACTED>",
-  # Some additional fields elided
-}
-
-Note the credentials established through this method are shared across all
-sessions run on this runtime.
-
-Note be sure to feed the inputs to this op to ensure the credentials are not
-stored in a constant op within the graph that might accidentally be checkpointed
-or in other ways be persisted or exfiltrated.
-)doc");
+    .SetShapeFn(shape_inference::NoOutputs);
 
 REGISTER_OP("GcsConfigureBlockCache")
     .Input("max_cache_size: uint64")
     .Input("block_size: uint64")
     .Input("max_staleness: uint64")
-    .SetShapeFn(shape_inference::NoOutputs)
-    .Doc(R"doc(
-Re-configures the GCS block cache with the new configuration values.
-
-If the values are the same as already configured values, this op is a no-op. If
-they are different, the current contents of the block cache is dropped, and a
-new block cache is created fresh.
-)doc");
+    .SetShapeFn(shape_inference::NoOutputs);
 
 }  // namespace tensorflow

tensorflow/core/api_def/base_api/api_def_GcsConfigureBlockCache.pbtxt

@@ -0,0 +1,9 @@
+op {
+  graph_op_name: "GcsConfigureBlockCache"
+  summary: "Re-configures the GCS block cache with the new configuration values."
+  description: <<END
+If the values are the same as already configured values, this op is a no-op. If
+they are different, the current contents of the block cache is dropped, and a
+new block cache is created fresh.
+END
+}

tensorflow/core/api_def/base_api/api_def_GcsConfigureCredentials.pbtxt

@@ -0,0 +1,33 @@
+op {
+  graph_op_name: "GcsConfigureCredentials"
+  summary: "Configures the credentials used by the GCS client of the local TF runtime."
+  description: <<END0
+The json input can be of the format:
+
+1. Refresh Token:
+{
+  "client_id": "<redacted>",
+  "client_secret": "<redacted>",
+  "refresh_token: "<redacted>",
+  "type": "authorized_user",
+}
+
+2. Service Account:
+{
+  "type": "service_account",
+  "project_id": "<redacted>",
+  "private_key_id": "<redacted>",
+  "private_key": "------BEGIN PRIVATE KEY-----\n<REDACTED>\n-----END PRIVATE KEY------\n",
+  "client_email": "<REDACTED>@<REDACTED>.iam.gserviceaccount.com",
+  "client_id": "<REDACTED>",
+  # Some additional fields elided
+}
+
+Note the credentials established through this method are shared across all
+sessions run on this runtime.
+
+Note be sure to feed the inputs to this op to ensure the credentials are not
+stored in a constant op within the graph that might accidentally be checkpointed
+or in other ways be persisted or exfiltrated.
+END0
+}

tensorflow/core/platform/cloud/gcs_file_system.cc

@@ -804,7 +804,9 @@ void GcsFileSystem::ResetFileBlockCache(size_t block_size_bytes,
   mutex_lock l(block_cache_lock_);
   file_block_cache_ =
       MakeFileBlockCache(block_size_bytes, max_bytes, max_staleness_secs);
-  stats_->Configure(this, &throttle_, file_block_cache_.get());
+  if (stats_) {
+    stats_->Configure(this, &throttle_, file_block_cache_.get());
+  }
 }
 
 // A helper function to build a FileBlockCache for GcsFileSystem.

tensorflow/core/platform/default/build_config.bzl

@@ -633,6 +633,7 @@ def tf_additional_cloud_op_deps():
       "//tensorflow:with_gcp_support_ios_override": [],
       "//tensorflow:with_gcp_support": [
         "//tensorflow/contrib/cloud:bigquery_reader_ops_op_lib",
+        "//tensorflow/contrib/cloud:gcs_config_ops_op_lib",
       ],
       "//conditions:default": [],
   })
@@ -645,6 +646,7 @@ def tf_additional_cloud_kernel_deps():
       "//tensorflow:with_gcp_support_ios_override": [],
       "//tensorflow:with_gcp_support": [
         "//tensorflow/contrib/cloud/kernels:bigquery_reader_ops",
+        "//tensorflow/contrib/cloud/kernels:gcs_config_ops",
       ],
       "//conditions:default": [],
   })