Add integrations for ucaip

Allow users to use all the aiplatform services

Author: bovard

Dockerfile

@@ -273,6 +273,8 @@ RUN pip install --upgrade cython && \
     pip install google-cloud-language==2.* && \
     pip install google-cloud-videointelligence==2.* && \
     pip install google-cloud-vision==2.* && \
+    # After launch this should be installed from pip
+    pip install git+https://github.com/googleapis/python-aiplatform.git@mb-release && \ 
     pip install ortools && \
     pip install scattertext && \
     # Pandas data reader

patches/kaggle_gcp.py

@@ -187,6 +187,23 @@ def monkeypatch_bq(bq_client, *args, **kwargs):
             bq_client, *args, **kwargs)
     return bigquery
 
+# Monkey patch classes that use the init method
+# eg
+# from google.cloud import aiplatform
+# aiplatform.init(args)
+def monkeypatch_init(client_klass, kaggle_kernel_credentials):
+    client_init = client_klass.init
+    def patched_init(self, *args, **kwargs):
+        specified_credentials = kwargs.get('credentials')
+        if specified_credentials is None:
+            Log.info("No credentials specified, using KaggleKernelCredentials.")
+            kwargs['credentials'] = kaggle_kernel_credentials
+            return client_init(self, *args, **kwargs)
+
+    if (not has_been_monkeypatched(client_klass.init)):
+        client_klass.init = patched_init
+        Log.info(f"Client patched: {client_klass}")
+
 def monkeypatch_client(client_klass, kaggle_kernel_credentials):
     client_init = client_klass.__init__
     def patched_init(self, *args, **kwargs):
@@ -309,6 +326,22 @@ def init_natural_language():
     monkeypatch_client(language.LanguageServiceAsyncClient, kernel_credentials)
     return language
 
+def init_ucaip():
+    from google.cloud import aiplatform
+    if not is_user_secrets_token_set():
+        return
+
+    from kaggle_gcp import get_integrations
+    if not get_integrations().has_cloudai():
+        return
+
+    from kaggle_secrets import GcpTarget
+    from kaggle_gcp import KaggleKernelCredentials
+    kaggle_kernel_credentials = KaggleKernelCredentials(target=GcpTarget.CLOUDAI)
+
+    # Patch the ucaip init method, this flows down to all ucaip services
+    monkeypatch_init(aiplatform.initializer.global_config, kaggle_kernel_credentials)
+
 def init_video_intelligence():
     from google.cloud import videointelligence
     if not is_user_secrets_token_set():
@@ -352,6 +385,7 @@ def init():
     init_natural_language()
     init_video_intelligence()
     init_vision()
+    init_ucaip()
 
 # We need to initialize the monkeypatching of the client libraries
 # here since there is a circular dependency between our import hook version

test

@@ -70,7 +70,6 @@ readonly PATTERN
 
 set -x
 docker run --rm --net=none -v /tmp/python-build:/tmp/python-build "$IMAGE_TAG" rm -rf /tmp/python-build/*
-
 docker rm jupyter_test || true
 mkdir -p /tmp/python-build/tmp
 mkdir -p /tmp/python-build/devshm

tests/test_ucaip.py

@@ -0,0 +1,24 @@
+import unittest
+
+from unittest.mock import Mock
+
+from kaggle_gcp import KaggleKernelCredentials, init_ucaip
+from test.support import EnvironmentVarGuard
+
+def _make_credentials():
+    import google.auth.credentials
+    return Mock(spec=google.auth.credentials.Credentials)
+
+class TestUcaip(unittest.TestCase):
+
+    def test_user_provided_credentials(self):
+        credentials = _make_credentials()
+        env = EnvironmentVarGuard()
+        env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar')
+        env.set('KAGGLE_KERNEL_INTEGRATIONS', 'CLOUDAI')
+        with env:
+            from google.cloud import aiplatform
+            init_ucaip()
+            aiplatform.init(credentials=credentials)
+            self.assertNotIsInstance(aiplatform.initializer.global_config.credentials, KaggleKernelCredentials)
+            self.assertIsNotNone(aiplatform.initializer.global_config.credentials)