fix: route Endpoint.predict_async via dedicated DNS when the endpoint is dedicated

Endpoint.predict_async was hardcoded to the public regional GAPIC async
client, which rejects dedicated endpoint traffic with FailedPrecondition.
The synchronous predict() already detects dedicated_endpoint_enabled and
POSTs via AuthorizedSession to https://{dedicated_endpoint_dns}/v1/{name}:predict;
this change mirrors that decision for the async path using aiohttp.

When the endpoint is not dedicated, predict_async behavior is unchanged
(still goes through PredictionServiceAsyncClient), so non-dedicated
callers see no diff.

The dedicated branch:
  - Raises ValueError when dedicated_endpoint_dns is empty, matching the
    sync path's error message.
  - Lazily builds and caches an aiohttp.ClientSession on the Endpoint
    instance (self._authorized_session_async).
  - Refreshes self.credentials per call (AuthorizedSession does this
    transparently in sync; async path must do it explicitly) and sends
    the bearer token in the Authorization header.
  - Returns a Prediction with the same field mapping the sync dedicated
    path uses.

PiperOrigin-RevId: 922598716

Author: IronPan

google/cloud/aiplatform/models.py

@@ -20,6 +20,7 @@
 import re
 import shutil
 import tempfile
+import aiohttp
 import requests
 from typing import (
     Any,
@@ -696,6 +697,8 @@ def __init__(
         self._gca_resource = gca_endpoint_compat.Endpoint(name=endpoint_name)
 
         self.authorized_session = None
+        # Lazy aiohttp ClientSession for async predict on dedicated endpoints.
+        self._authorized_session_async = None
 
     @property
     def _prediction_client(self) -> utils.PredictionClientWithOverride:
@@ -2715,6 +2718,11 @@ async def predict_async(
             my_predictions = response.predictions
             ```
 
+        For dedicated endpoints (``dedicated_endpoint_enabled=True``), the call
+        is routed via async HTTPS POST to the endpoint's dedicated DNS,
+        mirroring the synchronous ``predict()`` dedicated-endpoint path.
+        Otherwise the GAPIC async prediction client is used.
+
         Args:
             instances (List):
                 Required. The instances that are the input to the
@@ -2740,29 +2748,85 @@ async def predict_async(
         Returns:
             prediction (aiplatform.Prediction):
                 Prediction with returned predictions and Model ID.
+
+        Raises:
+            ValueError: If the dedicated endpoint DNS is empty for dedicated
+                endpoints, or if the prediction request to a dedicated endpoint
+                returns a non-200 status.
         """
         self.wait()
 
-        prediction_response = await self._prediction_async_client.predict(
-            endpoint=self._gca_resource.name,
-            instances=instances,
-            parameters=parameters,
-            timeout=timeout,
-        )
-        if prediction_response._pb.metadata:
-            metadata = json_format.MessageToDict(prediction_response._pb.metadata)
+        if not self.dedicated_endpoint_enabled:
+            prediction_response = await self._prediction_async_client.predict(
+                endpoint=self._gca_resource.name,
+                instances=instances,
+                parameters=parameters,
+                timeout=timeout,
+            )
+            if prediction_response._pb.metadata:
+                metadata = json_format.MessageToDict(prediction_response._pb.metadata)
+            else:
+                metadata = None
+
+            return Prediction(
+                predictions=[
+                    json_format.MessageToDict(item)
+                    for item in prediction_response.predictions.pb
+                ],
+                metadata=metadata,
+                deployed_model_id=prediction_response.deployed_model_id,
+                model_version_id=prediction_response.model_version_id,
+                model_resource_name=prediction_response.model,
+            )
+
+        # Dedicated endpoint: REST POST to the dedicated DNS via aiohttp.
+        if not self.dedicated_endpoint_dns:
+            raise ValueError(
+                "Dedicated endpoint DNS is empty. Please make sure endpoint"
+                "and model are ready before making a prediction."
+            )
+
+        if parameters is not None:
+            data = json.dumps({"instances": instances, "parameters": parameters})
         else:
-            metadata = None
+            data = json.dumps({"instances": instances})
+
+        if self._authorized_session_async is None:
+            self.credentials._scopes = constants.base.DEFAULT_AUTHED_SCOPES
+            self._authorized_session_async = aiohttp.ClientSession()
+
+        # Refresh the bearer token per call. ``AuthorizedSession`` (sync)
+        # handles refresh internally; in the async path we do it explicitly.
+        self.credentials.refresh(google_auth_requests.Request())
+        headers = {
+            "Authorization": f"Bearer {self.credentials.token}",
+            "Content-Type": "application/json",
+        }
+        url = f"https://{self.dedicated_endpoint_dns}/v1/{self.resource_name}:predict"
+        aiohttp_timeout = (
+            aiohttp.ClientTimeout(total=timeout) if timeout is not None else None
+        )
+
+        async with self._authorized_session_async.post(
+            url=url,
+            data=data,
+            headers=headers,
+            timeout=aiohttp_timeout,
+        ) as response:
+            if response.status != 200:
+                text = await response.text()
+                raise ValueError(
+                    f"Failed to make prediction request. Status code:"
+                    f"{response.status}, response: {text}."
+                )
+            prediction_response = await response.json()
 
         return Prediction(
-            predictions=[
-                json_format.MessageToDict(item)
-                for item in prediction_response.predictions.pb
-            ],
-            metadata=metadata,
-            deployed_model_id=prediction_response.deployed_model_id,
-            model_version_id=prediction_response.model_version_id,
-            model_resource_name=prediction_response.model,
+            predictions=prediction_response.get("predictions"),
+            metadata=prediction_response.get("metadata"),
+            deployed_model_id=prediction_response.get("deployedModelId"),
+            model_resource_name=prediction_response.get("model"),
+            model_version_id=prediction_response.get("modelVersionId"),
         )
 
     def raw_predict(

tests/unit/aiplatform/test_endpoints.py

@@ -19,6 +19,8 @@
 from datetime import datetime, timedelta
 from importlib import reload
 import json
+import aiohttp
+from aiohttp import web as aiohttp_web
 import requests
 from unittest import mock
 from google.protobuf import duration_pb2
@@ -662,6 +664,38 @@ def predict_async_client_predict_mock():
         yield predict_mock
 
 
+@pytest.fixture
+def predict_endpoint_aiohttp_mock():
+    """Mocks aiohttp.ClientSession.post for dedicated-endpoint async predict.
+
+    ClientSession.post returns a _RequestContextManager (async ctx mgr), so the
+    mock must implement __aenter__/__aexit__. Also patches Credentials.refresh
+    to a no-op so we don't hit real auth during tests.
+    """
+    payload = {
+        "predictions": _TEST_PREDICTION,
+        "metadata": _TEST_METADATA,
+        "deployedModelId": _TEST_DEPLOYED_MODELS[0].id,
+        "model": _TEST_MODEL_NAME,
+        "modelVersionId": "1",
+    }
+    response = mock.AsyncMock()
+    response.status = 200
+    response.json = mock.AsyncMock(return_value=payload)
+    response.text = mock.AsyncMock(return_value=json.dumps(payload))
+
+    post_cm = mock.MagicMock()
+    post_cm.__aenter__ = mock.AsyncMock(return_value=response)
+    post_cm.__aexit__ = mock.AsyncMock(return_value=None)
+
+    with mock.patch.object(
+        aiohttp.ClientSession, "post", return_value=post_cm
+    ) as post_mock, mock.patch.object(
+        auth_credentials.AnonymousCredentials, "refresh"
+    ):
+        yield post_mock
+
+
 @pytest.fixture
 def predict_client_direct_predict_mock():
     with mock.patch.object(
@@ -3687,6 +3721,175 @@ async def test_predict_async(self, predict_async_client_predict_mock):
             timeout=None,
         )
 
+    @pytest.mark.asyncio
+    @pytest.mark.usefixtures("get_dedicated_endpoint_mock")
+    async def test_predict_async_dedicated_endpoint(
+        self, predict_endpoint_aiohttp_mock
+    ):
+        """Async predict on a dedicated endpoint routes via aiohttp REST."""
+        test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME)
+
+        test_prediction = await test_endpoint.predict_async(
+            instances=_TEST_INSTANCES, parameters={"param": 3.0}
+        )
+
+        true_prediction = models.Prediction(
+            predictions=_TEST_PREDICTION,
+            deployed_model_id=_TEST_ID,
+            metadata=_TEST_METADATA,
+            model_version_id=_TEST_VERSION_ID,
+            model_resource_name=_TEST_MODEL_NAME,
+        )
+        assert true_prediction == test_prediction
+
+        predict_endpoint_aiohttp_mock.assert_called_once()
+        call_kwargs = predict_endpoint_aiohttp_mock.call_args.kwargs
+        assert call_kwargs["url"] == (
+            f"https://{_TEST_DEDICATED_ENDPOINT_DNS}/v1/projects/"
+            f"{_TEST_PROJECT}/locations/{_TEST_LOCATION}/endpoints/{_TEST_ID}:predict"
+        )
+        assert call_kwargs["data"] == (
+            '{"instances": [[1.0, 2.0, 3.0], [1.0, 3.0, 4.0]],'
+            ' "parameters": {"param": 3.0}}'
+        )
+        assert call_kwargs["headers"]["Content-Type"] == "application/json"
+        assert call_kwargs["headers"]["Authorization"].startswith("Bearer ")
+        assert call_kwargs["timeout"] is None
+
+    @pytest.mark.asyncio
+    @pytest.mark.usefixtures("get_dedicated_endpoint_no_dns_mock")
+    async def test_predict_async_dedicated_endpoint_without_dns(
+        self, predict_endpoint_aiohttp_mock
+    ):
+        """Async predict on a dedicated endpoint with no DNS raises ValueError."""
+        test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME)
+
+        with pytest.raises(ValueError) as err:
+            await test_endpoint.predict_async(
+                instances=_TEST_INSTANCES, parameters={"param": 3.0}
+            )
+        assert err.match(
+            regexp=r"Dedicated endpoint DNS is empty. Please make sure endpoint"
+            "and model are ready before making a prediction."
+        )
+
+    @pytest.mark.asyncio
+    @pytest.mark.usefixtures("get_dedicated_endpoint_mock")
+    async def test_predict_async_dedicated_endpoint_with_timeout(
+        self, predict_endpoint_aiohttp_mock
+    ):
+        """A non-None timeout is forwarded as an aiohttp.ClientTimeout."""
+        test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME)
+
+        await test_endpoint.predict_async(
+            instances=_TEST_INSTANCES,
+            parameters={"param": 3.0},
+            timeout=_TEST_PREDICT_TIMEOUT,
+        )
+
+        predict_endpoint_aiohttp_mock.assert_called_once()
+        call_timeout = predict_endpoint_aiohttp_mock.call_args.kwargs["timeout"]
+        assert isinstance(call_timeout, aiohttp.ClientTimeout)
+        assert call_timeout.total == _TEST_PREDICT_TIMEOUT
+
+    @pytest.mark.asyncio
+    @pytest.mark.usefixtures("get_dedicated_endpoint_mock")
+    async def test_predict_async_dedicated_endpoint_integration(self):
+        """Hermetic integration test: real aiohttp request against a local server.
+
+        Unlike the mock-based tests above, this exercises the real
+        aiohttp.ClientSession code path (connection, request encoding,
+        response decoding, JSON parsing, async context manager teardown) by
+        pointing the dedicated DNS at an in-process aiohttp.web server.
+
+        The server validates the URL path, headers, and body shape the SDK
+        sends, and returns a Vertex-shaped JSON response. The SDK builds the
+        URL as ``https://{dns}/...``; we wrap the session to downgrade the
+        scheme to ``http://`` so the local server can serve plain HTTP without
+        needing self-signed TLS plumbing in the test.
+        """
+        captured = {}
+
+        async def predict_handler(request: aiohttp_web.Request):
+            captured["path"] = request.path
+            captured["headers"] = dict(request.headers)
+            captured["body"] = await request.json()
+            return aiohttp_web.json_response(
+                {
+                    "predictions": _TEST_PREDICTION,
+                    "metadata": _TEST_METADATA,
+                    "deployedModelId": _TEST_DEPLOYED_MODELS[0].id,
+                    "model": _TEST_MODEL_NAME,
+                    "modelVersionId": "1",
+                }
+            )
+
+        app = aiohttp_web.Application()
+        app.router.add_post(
+            f"/v1/{_TEST_ENDPOINT_NAME}:predict", predict_handler
+        )
+        runner = aiohttp_web.AppRunner(app)
+        await runner.setup()
+        site = aiohttp_web.TCPSite(runner, "127.0.0.1", 0)
+        await site.start()
+        port = site._server.sockets[0].getsockname()[1]
+
+        try:
+            test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME)
+
+            # Force the dedicated DNS to point at the local server.
+            with mock.patch.object(
+                models.Endpoint,
+                "dedicated_endpoint_dns",
+                new_callable=mock.PropertyMock,
+                return_value=f"127.0.0.1:{port}",
+            ), mock.patch.object(
+                auth_credentials.AnonymousCredentials, "refresh"
+            ):
+                # Inject a session that downgrades https:// -> http:// so the
+                # local test server can serve plain HTTP. Real aiohttp does
+                # everything else (real socket, real request/response cycle).
+                real_session = aiohttp.ClientSession()
+                original_post = real_session.post
+
+                def http_post(url, **kwargs):
+                    return original_post(
+                        url.replace("https://", "http://", 1), **kwargs
+                    )
+
+                real_session.post = http_post
+                test_endpoint._authorized_session_async = real_session
+                # Give the mocked credentials a fake token.
+                test_endpoint.credentials.token = "test-token"
+
+                try:
+                    prediction = await test_endpoint.predict_async(
+                        instances=_TEST_INSTANCES,
+                        parameters={"param": 3.0},
+                    )
+                finally:
+                    await real_session.close()
+
+            # Verify the SDK sent the right wire request.
+            assert captured["path"] == f"/v1/{_TEST_ENDPOINT_NAME}:predict"
+            assert captured["headers"]["Content-Type"] == "application/json"
+            assert captured["headers"]["Authorization"] == "Bearer test-token"
+            assert captured["body"] == {
+                "instances": _TEST_INSTANCES,
+                "parameters": {"param": 3.0},
+            }
+
+            # Verify the response was parsed correctly into a Prediction.
+            assert prediction == models.Prediction(
+                predictions=_TEST_PREDICTION,
+                deployed_model_id=_TEST_ID,
+                metadata=_TEST_METADATA,
+                model_version_id=_TEST_VERSION_ID,
+                model_resource_name=_TEST_MODEL_NAME,
+            )
+        finally:
+            await runner.cleanup()
+
     @pytest.mark.usefixtures("get_endpoint_mock")
     def test_explain(self, predict_client_explain_mock):
         test_endpoint = models.Endpoint(_TEST_ID)