Directly load https://hub.tensorflow.google.cn/* from GCS

PiperOrigin-RevId: 547499610

tensorflow_hub/compressed_module_resolver.py

@@ -15,14 +15,18 @@
 """Functions to resolve TF-Hub Module stored in compressed TGZ format."""
 
 import hashlib
+import logging
 import urllib
 
 import tensorflow as tf
 from tensorflow_hub import resolver
 
 
 LOCK_FILE_TIMEOUT_SEC = 10 * 60  # 10 minutes
-
+_HUB_TF_GOOGLE_CN = "https://hub.tensorflow.google.cn/"
+_GCS_GOOGLE_CN_TEMPLATE = (
+    "https://gcs.tensorflow.google.cn/tfhub-modules/%s.tar.gz"
+)
 _COMPRESSED_FORMAT_QUERY = ("tf-hub-format", "compressed")
 
 
@@ -58,6 +62,16 @@ def __call__(self, handle):
 
     def download(handle, tmp_dir):
       """Fetch a module via HTTP(S), handling redirect and download headers."""
+      # Directly load the model from GCS.
+      if handle.startswith(_HUB_TF_GOOGLE_CN):
+        full_model_name = handle[len(_HUB_TF_GOOGLE_CN):]
+        gcs_cn_url = _GCS_GOOGLE_CN_TEMPLATE % full_model_name
+        logging.info("Directly downloading %s", gcs_cn_url)
+        response = self._call_urlopen(gcs_cn_url)
+        return resolver.DownloadManager(handle).download_and_uncompress(
+            response, tmp_dir
+        )
+
       request = urllib.request.Request(
           self._append_compressed_format_query(handle))
       response = self._call_urlopen(request)

tensorflow_hub/compressed_module_resolver_test.py

@@ -21,23 +21,26 @@
 import tarfile
 import tempfile
 import unittest
+from unittest import mock
+import urllib.request
 import uuid
 
 from absl import flags
 from absl.testing import parameterized
 import tensorflow as tf
-
 from tensorflow_hub import compressed_module_resolver
 from tensorflow_hub import resolver
 from tensorflow_hub import test_utils
 from tensorflow_hub import tf_utils
 
+
 FLAGS = flags.FLAGS
 
 
 class HttpCompressedFileResolverTest(tf.test.TestCase, parameterized.TestCase):
 
   def setUp(self):
+    super().setUp()
     # Set current directory to test temp directory where we can create
     # files and serve them through the HTTP server.
     os.chdir(self.get_temp_dir())
@@ -235,6 +238,27 @@ def testCorruptedArchive(self):
           "http://localhost:%d/bad_archive.tar.gz does not appear "
           "to be a valid module." % self.redirect_server_port, str(e))
 
+  def testLoadFromCn(self):
+    http_resolver = compressed_module_resolver.HttpCompressedFileResolver()
+
+    with mock.patch.object(
+        urllib.request,
+        "urlopen",
+        autospec=True,
+        return_value=urllib.request.urlopen(
+            "http://localhost:%d/mock_module.tar.gz" % self.server_port
+        ),
+    ) as mock_urlopen:
+      path = http_resolver(
+          "https://hub.tensorflow.google.cn/google/bit/s-r50x1/1"
+      )
+
+    mock_urlopen.assert_called_once_with(
+        "https://gcs.tensorflow.google.cn/tfhub-modules/google/bit/s-r50x1/1.tar.gz",
+        context=mock.ANY,
+    )
+    self.assertCountEqual(os.listdir(path), ["file1", "file2", "file3"])
+
 
 if __name__ == "__main__":
   tf.test.main()