feat: Add private async REST support for transport override

matthew29tang · copybara-github · commit c0b31e267b15 · 2024-10-29T10:08:49.000-07:00
PiperOrigin-RevId: 691063218
diff --git a/google/cloud/aiplatform/initializer.py b/google/cloud/aiplatform/initializer.py
@@ -23,7 +23,7 @@
 import logging
 import os
 import types
-from typing import Iterator, List, Optional, Sequence, Tuple, Type, TypeVar, Union
+from typing import Any, Iterator, List, Optional, Sequence, Tuple, Type, TypeVar, Union
 
 from google.api_core import client_options
 from google.api_core import gapic_v1
@@ -46,6 +46,15 @@
     encryption_spec_v1beta1 as gca_encryption_spec_v1beta1,
 )
 
+try:
+    import google.auth.aio
+
+    AsyncCredentials = google.auth.aio.credentials.Credentials
+    _HAS_ASYNC_CRED_DEPS = True
+except ImportError:
+    AsyncCredentials = Any
+    _HAS_ASYNC_CRED_DEPS = False
+
 _TVertexAiServiceClientWithOverride = TypeVar(
     "_TVertexAiServiceClientWithOverride",
     bound=utils.VertexAiServiceClientWithOverride,
@@ -121,6 +130,7 @@ def __init__(self):
         self._api_transport = None
         self._request_metadata = None
         self._resource_type = None
+        self._async_rest_credentials = None
 
     def init(
         self,
@@ -590,15 +600,24 @@ def create_client(
         }
 
         # Do not pass "grpc", rely on gapic defaults unless "rest" is specified
-        if self._api_transport == "rest":
-            if "Async" in client_class.__name__:
-                # Warn user that "rest" is not supported and use grpc instead
+        if self._api_transport == "rest" and "Async" in client_class.__name__:
+            # User requests async rest
+            if self._async_rest_credentials:
+                # Rest async recieves credentials from _async_rest_credentials
+                kwargs["credentials"] = self._async_rest_credentials
+                kwargs["transport"] = "rest_asyncio"
+            else:
+                # Rest async was specified, but no async credentials were set.
+                # Fallback to gRPC instead.
                 logging.warning(
-                    "REST is not supported for async clients, "
-                    + "falling back to grpc."
+                    "REST async clients requires async credentials set using "
+                    + "aiplatform.initializer._set_async_rest_credentials().\n"
+                    + "Falling back to grpc since no async rest credentials "
+                    + "were detected."
                 )
-            else:
-                kwargs["transport"] = self._api_transport
+        elif self._api_transport == "rest":
+            # User requests sync REST
+            kwargs["transport"] = self._api_transport
 
         client = client_class(**kwargs)
         # We only wrap the client if the request_metadata is set at the creation time.
@@ -672,6 +691,29 @@ def __call__(self, *args, **kwargs):
 )
 
 
+def _set_async_rest_credentials(credentials: AsyncCredentials):
+    """Private method to set async REST credentials."""
+    if global_config._api_transport != "rest":
+        raise ValueError(
+            "Async REST credentials can only be set when using REST transport."
+        )
+    elif not _HAS_ASYNC_CRED_DEPS or not isinstance(credentials, AsyncCredentials):
+        raise ValueError(
+            "Async REST transport requires async credentials of type"
+            + f"{AsyncCredentials} which is only supported in "
+            + "google-auth >= 2.35.0.\n\n"
+            + "Install the following dependencies:\n"
+            + "pip install google-api-core[grpc, async_rest] >= 2.21.0\n"
+            + "pip install google-auth[aiohttp] >= 2.35.0\n\n"
+            + "Example usage:\n"
+            + "from google.auth.aio.credentials import StaticCredentials\n"
+            + "async_credentials = StaticCredentials(token=YOUR_TOKEN_HERE)\n"
+            + "aiplatform.initializer._set_async_rest_credentials("
+            + "credentials=async_credentials)"
+        )
+    global_config._async_rest_credentials = credentials
+
+
 def _get_function_name_from_stack_frame(frame) -> str:
     """Gates fully qualified function or method name.
 
diff --git a/setup.py b/setup.py
@@ -195,6 +195,9 @@
     + profiler_extra_require
     + tokenization_testing_extra_require
     + [
+        # aiohttp is required for async rest tests (need google-auth[aiohttp],
+        # but can't specify extras in constraints files)
+        "aiohttp",
         "bigframes; python_version>='3.10'",
         # google-api-core 2.x is required since kfp requires protobuf > 4
         "google-api-core >= 2.11, < 3.0.0",
diff --git a/testing/constraints-3.10.txt b/testing/constraints-3.10.txt
@@ -1,7 +1,8 @@
 # -*- coding: utf-8 -*-
 # This constraints file is required for unit tests.
 # List all library dependencies and extras in this file.
-google-api-core
+google-api-core==2.21.0 # Tests google-api-core with rest async support
+google-auth==2.35.0 # Tests google-auth with rest async support
 proto-plus==1.22.3
 protobuf
 mock==4.0.2
diff --git a/testing/constraints-3.11.txt b/testing/constraints-3.11.txt
@@ -1,7 +1,8 @@
 # -*- coding: utf-8 -*-
 # This constraints file is required for unit tests.
 # List all library dependencies and extras in this file.
-google-api-core
+google-api-core==2.21.0 # Tests google-api-core with rest async support
+google-auth==2.35.0 # Tests google-auth with rest async support
 proto-plus
 protobuf
 mock==4.0.2
diff --git a/testing/constraints-3.12.txt b/testing/constraints-3.12.txt
@@ -1,7 +1,8 @@
 # -*- coding: utf-8 -*-
 # This constraints file is required for unit tests.
 # List all library dependencies and extras in this file.
-google-api-core
+google-api-core==2.21.0 # Tests google-api-core with rest async support
+google-auth==2.35.0 # Tests google-auth with rest async support
 proto-plus
 protobuf
 mock==4.0.2
diff --git a/testing/constraints-3.8.txt b/testing/constraints-3.8.txt
@@ -3,7 +3,7 @@
 # are correct in setup.py
 # List *all* library dependencies and extras in this file.
 google-api-core==2.17.1 # Increased for gapic owlbot presubmit tests
-google-auth==2.14.1
+google-auth==2.14.1 # Tests google-auth without rest async support
 proto-plus==1.22.3
 protobuf
 mock==4.0.2
diff --git a/testing/constraints-3.9.txt b/testing/constraints-3.9.txt
@@ -1,7 +1,8 @@
 # -*- coding: utf-8 -*-
 # This constraints file is required for unit tests.
 # List all library dependencies and extras in this file.
-google-api-core
+google-api-core==2.21.0 # Tests google-api-core with rest async support
+google-auth==2.35.0 # Tests google-auth with rest async support
 proto-plus==1.22.3
 protobuf
 mock==4.0.2
diff --git a/tests/system/aiplatform/test_initializer.py b/tests/system/aiplatform/test_initializer.py
@@ -15,9 +15,12 @@
 # limitations under the License.
 #
 
+import pytest
+
 from google.auth import credentials as auth_credentials
 
 from google.cloud import aiplatform
+from google.cloud.aiplatform import initializer as aiplatform_initializer
 from tests.system.aiplatform import e2e_base
 
 
@@ -39,3 +42,18 @@ def test_init_calls_set_google_auth_default(self):
         # init() with only project shouldn't overwrite creds
         aiplatform.init(project=e2e_base._PROJECT)
         assert aiplatform.initializer.global_config.credentials == creds
+
+    def test_init_rest_async_incorrect_credentials(self):
+        # Async REST credentials must be explicitly set using
+        # _set_async_rest_credentials() for async REST transport.
+        creds = auth_credentials.AnonymousCredentials()
+        aiplatform.init(
+            project=e2e_base._PROJECT,
+            location=e2e_base._LOCATION,
+            api_transport="rest",
+        )
+
+        # System tests are run on Python 3.10 which has async deps.
+        with pytest.raises(ValueError):
+            # Expect a ValueError for passing in sync credentials.
+            aiplatform_initializer._set_async_rest_credentials(credentials=creds)
diff --git a/tests/system/aiplatform/test_language_models.py b/tests/system/aiplatform/test_language_models.py
@@ -86,12 +86,10 @@ def test_text_generation_preview_count_tokens(self, api_transport):
         assert response.total_billable_characters
 
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("api_transport", ["grpc", "rest"])
-    async def test_text_generation_model_predict_async(self, api_transport):
+    async def test_text_generation_model_predict_async(self):
         aiplatform.init(
             project=e2e_base._PROJECT,
             location=e2e_base._LOCATION,
-            api_transport=api_transport,
         )
 
         model = TextGenerationModel.from_pretrained("google/text-bison@001")
@@ -227,12 +225,10 @@ def test_chat_model_preview_count_tokens(self, api_transport):
         )
 
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("api_transport", ["grpc", "rest"])
-    async def test_chat_model_async(self, api_transport):
+    async def test_chat_model_async(self):
         aiplatform.init(
             project=e2e_base._PROJECT,
             location=e2e_base._LOCATION,
-            api_transport=api_transport,
         )
 
         chat_model = ChatModel.from_pretrained("google/chat-bison@001")
@@ -343,12 +339,10 @@ def test_text_embedding(self, api_transport):
         assert embeddings[1].statistics.truncated
 
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("api_transport", ["grpc", "rest"])
-    async def test_text_embedding_async(self, api_transport):
+    async def test_text_embedding_async(self):
         aiplatform.init(
             project=e2e_base._PROJECT,
             location=e2e_base._LOCATION,
-            api_transport=api_transport,
         )
 
         model = TextEmbeddingModel.from_pretrained("google/textembedding-gecko@001")
diff --git a/tests/system/vertexai/test_generative_models.py b/tests/system/vertexai/test_generative_models.py
diff --git a/tests/unit/aiplatform/test_language_models.py b/tests/unit/aiplatform/test_language_models.py
diff --git a/tests/unit/vertexai/test_generative_models.py b/tests/unit/vertexai/test_generative_models.py
diff --git a/vertexai/generative_models/_generative_models.py b/vertexai/generative_models/_generative_models.py
