Use kagglehub inside Kaggle Notebook for tensorflow_hub (#1342)

Currently, if you try using tensorflow_hub with a kaggle.com URL inside
a Kaggle notebook but the model isn't attached you would get an error
message telling you to attach the model.

With this change, we are calling `kagglehub` instead and will attach the
model for the user if it is missing.

http://b/311184735

Author: rosbo

patches/kaggle_module_resolver.py

@@ -1,12 +1,13 @@
 import os
 import re
+import kagglehub
 
 from tensorflow_hub import resolver
 
-url_pattern = re.compile(r"https?://([a-z]+\.)?kaggle.com/models/[^\\/]+/(?P<model>[^\\/]+)/frameworks/(?P<framework>[^\\/]+)/variations/(?P<variation>[^\\/]+)/versions/(?P<version>[0-9]+)$")
+url_pattern = re.compile(r"https?://([a-z]+\.)?kaggle.com/models/(?P<owner>[^\\/]+)/(?P<model>[^\\/]+)/frameworks/(?P<framework>[^\\/]+)/variations/(?P<variation>[^\\/]+)/versions/(?P<version>[0-9]+)$")
 
 def _is_on_kaggle_notebook():
-    return os.getenv("KAGGLE_CONTAINER_NAME") != None
+    return os.getenv("KAGGLE_KERNEL_RUN_TYPE") != None
 
 def _is_kaggle_handle(handle):
     return url_pattern.match(handle) != None
@@ -17,9 +18,4 @@ def is_supported(self, handle):
 
     def __call__(self, handle):
         m = url_pattern.match(handle)
-        local_path = f"/kaggle/input/{m.group('model')}/{m.group('framework').lower()}/{m.group('variation')}/{m.group('version')}"
-        if not os.path.exists(local_path):
-            # TODO(b/268256777) Attach model & wait until ready instead.
-            raise RuntimeError(f"You have to attach the '{handle}' model to your Kaggle notebook.")
-        
-        return local_path
+        return kagglehub.model_download(f"{m.group('owner')}/{m.group('model')}/{m.group('framework').lower()}/{m.group('variation')}/{m.group('version')}")

tests/test_kaggle_module_resolver.py

@@ -2,62 +2,100 @@
 
 import os
 import threading
+import json
 
 import tensorflow as tf
 import tensorflow_hub as hub
 
 from http.server import BaseHTTPRequestHandler, HTTPServer
 from test.support.os_helper import EnvironmentVarGuard
+from contextlib import contextmanager
+from kagglehub.exceptions import BackendError
 
+MOUNT_PATH = "/kaggle/input"
 
-class TestKaggleModuleResolver(unittest.TestCase):
-    class HubHTTPHandler(BaseHTTPRequestHandler):
-        def do_GET(self):
-            self.send_response(200)
-            self.send_header('Content-Type', 'application/gzip')
-            self.end_headers()
+@contextmanager
+def create_test_server(handler_class):
+    hostname = 'localhost'
+    port = 8080
+    addr = f"http://{hostname}:{port}"
 
-            with open('/input/tests/data/model.tar.gz', 'rb') as model_archive:
-                self.wfile.write(model_archive.read())
+    # Simulates we are inside a Kaggle environment.
+    env = EnvironmentVarGuard()
+    env.set('KAGGLE_KERNEL_RUN_TYPE', 'Interactive')
+    env.set('KAGGLE_USER_SECRETS_TOKEN', 'foo jwt token')
+    env.set('KAGGLE_DATA_PROXY_TOKEN', 'foo proxy token')
+    env.set('KAGGLE_DATA_PROXY_URL', addr)
 
-    def _test_client(self, client_func, handler):
-        with HTTPServer(('localhost', 8080), handler) as test_server:
+    with env:
+        with HTTPServer((hostname, port), handler_class) as test_server:
             threading.Thread(target=test_server.serve_forever).start()
 
             try:
-                client_func()            
+                yield addr
             finally:
                 test_server.shutdown()
 
-    def test_kaggle_resolver_succeeds(self):
-        # Simulates we are inside a Kaggle environment.
-        env = EnvironmentVarGuard()
-        env.set('KAGGLE_CONTAINER_NAME', 'foo')
-        # Attach model to right directory.
-        os.makedirs('/kaggle/input/foomodule/tensorflow2/barvar')
-        os.symlink('/input/tests/data/saved_model/', '/kaggle/input/foomodule/tensorflow2/barvar/2', target_is_directory=True)
+class HubHTTPHandler(BaseHTTPRequestHandler):
+    def do_GET(self):
+        self.send_response(200)
+        self.send_header('Content-Type', 'application/gzip')
+        self.end_headers()
+
+        with open('/input/tests/data/model.tar.gz', 'rb') as model_archive:
+            self.wfile.write(model_archive.read())
 
-        with env:
+class KaggleJwtHandler(BaseHTTPRequestHandler):
+    def do_POST(self):
+        if self.path.endswith("AttachDatasourceUsingJwtRequest"):
+            content_length = int(self.headers["Content-Length"])
+            request = json.loads(self.rfile.read(content_length))
+            model_ref = request["modelRef"]
+
+            self.send_response(200)
+            self.send_header("Content-type", "application/json")
+            self.end_headers()
+
+            if model_ref['ModelSlug'] == 'unknown':
+                self.wfile.write(bytes(json.dumps({
+                    "wasSuccessful": False,
+                }), "utf-8"))
+                return
+
+            # Load the files
+            mount_slug = f"{model_ref['ModelSlug']}/{model_ref['Framework']}/{model_ref['InstanceSlug']}/{model_ref['VersionNumber']}"
+            os.makedirs(os.path.dirname(os.path.join(MOUNT_PATH, mount_slug)))
+            os.symlink('/input/tests/data/saved_model/', os.path.join(MOUNT_PATH, mount_slug), target_is_directory=True)
+
+            # Return the response
+            self.wfile.write(bytes(json.dumps({
+                "wasSuccessful": True,
+                "result": {
+                    "mountSlug": mount_slug,
+                },
+            }), "utf-8"))
+        else:
+            self.send_response(404)
+            self.wfile.write(bytes(f"Unhandled path: {self.path}", "utf-8"))
+
+class TestKaggleModuleResolver(unittest.TestCase):
+    def test_kaggle_resolver_succeeds(self):        
+        with create_test_server(KaggleJwtHandler) as addr:
             test_inputs = tf.ones([1,4])
             layer = hub.KerasLayer("https://kaggle.com/models/foo/foomodule/frameworks/TensorFlow2/variations/barvar/versions/2")
             self.assertEqual([1, 1], layer(test_inputs).shape)
 
     def test_kaggle_resolver_not_attached_throws(self):
-        # Simulates we are inside a Kaggle environment.
-        env = EnvironmentVarGuard()
-        env.set('KAGGLE_CONTAINER_NAME', 'foo')
-        with env:
-            with self.assertRaisesRegex(RuntimeError, '.*attach.*'):
-                hub.KerasLayer("https://kaggle.com/models/foo/foomodule/frameworks/TensorFlow2/variations/barvar/versions/2")
+        with create_test_server(KaggleJwtHandler) as addr:
+            with self.assertRaises(BackendError):
+                hub.KerasLayer("https://kaggle.com/models/foo/unknown/frameworks/TensorFlow2/variations/barvar/versions/2")
 
     def test_http_resolver_succeeds(self):
-        def call_hub():
+        with create_test_server(HubHTTPHandler) as addr:
             test_inputs = tf.ones([1,4])
-            layer = hub.KerasLayer('http://localhost:8080/model.tar.gz')
+            layer = hub.KerasLayer(f'{addr}/model.tar.gz')
             self.assertEqual([1, 1], layer(test_inputs).shape)
 
-        self._test_client(call_hub, TestKaggleModuleResolver.HubHTTPHandler)
-
     def test_local_path_resolver_succeeds(self):
         test_inputs = tf.ones([1,4])
         layer = hub.KerasLayer('/input/tests/data/saved_model')