You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by jr...@apache.org on 2025/04/08 14:31:24 UTC
(beam) branch vertexRefactor created (now 9e6bb258f12)
This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a change to branch vertexRefactor
in repository https://gitbox.apache.org/repos/asf/beam.git
at 9e6bb258f12 Migrate Vertex AI Model Handler to new base class
This branch includes the following new commits:
new 9e6bb258f12 Migrate Vertex AI Model Handler to new base class
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(beam) 01/01: Migrate Vertex AI Model Handler to new base class
Posted by jr...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch vertexRefactor
in repository https://gitbox.apache.org/repos/asf/beam.git
commit 9e6bb258f12176a1c1b0d64149e9c7a6da1116ea
Author: Jack McCluskey <th...@gmail.com>
AuthorDate: Tue Apr 8 10:31:09 2025 -0400
Migrate Vertex AI Model Handler to new base class
---
sdks/python/apache_beam/io/requestresponse.py | 3 +-
.../ml/inference/vertex_ai_inference.py | 76 +++++-----------------
2 files changed, 18 insertions(+), 61 deletions(-)
diff --git a/sdks/python/apache_beam/io/requestresponse.py b/sdks/python/apache_beam/io/requestresponse.py
index d7011e5a8ff..213a1a3fc7e 100644
--- a/sdks/python/apache_beam/io/requestresponse.py
+++ b/sdks/python/apache_beam/io/requestresponse.py
@@ -43,7 +43,6 @@ from apache_beam import pvalue
from apache_beam.coders import coders
from apache_beam.io.components.adaptive_throttler import AdaptiveThrottler
from apache_beam.metrics import Metrics
-from apache_beam.ml.inference.vertex_ai_inference import MSEC_TO_SEC
from apache_beam.transforms.util import BatchElements
from apache_beam.utils import retry
@@ -58,6 +57,8 @@ DEFAULT_TIMEOUT_SECS = 30
# for cache record.
DEFAULT_CACHE_ENTRY_TTL_SEC = 24 * 60 * 60
+MSEC_TO_SEC = 1000
+
_LOGGER = logging.getLogger(__name__)
__all__ = [
diff --git a/sdks/python/apache_beam/ml/inference/vertex_ai_inference.py b/sdks/python/apache_beam/ml/inference/vertex_ai_inference.py
index 4c4163accfb..471f2379cfb 100644
--- a/sdks/python/apache_beam/ml/inference/vertex_ai_inference.py
+++ b/sdks/python/apache_beam/ml/inference/vertex_ai_inference.py
@@ -16,26 +16,19 @@
#
import logging
-import time
+from collections.abc import Iterable
+from collections.abc import Mapping
+from collections.abc import Sequence
from typing import Any
-from typing import Dict
-from typing import Iterable
-from typing import Mapping
from typing import Optional
-from typing import Sequence
from google.api_core.exceptions import ServerError
from google.api_core.exceptions import TooManyRequests
from google.cloud import aiplatform
-from apache_beam.io.components.adaptive_throttler import AdaptiveThrottler
-from apache_beam.metrics.metric import Metrics
from apache_beam.ml.inference import utils
-from apache_beam.ml.inference.base import ModelHandler
from apache_beam.ml.inference.base import PredictionResult
-from apache_beam.utils import retry
-
-MSEC_TO_SEC = 1000
+from apache_beam.ml.inference.base import RemoteModelHandler
LOGGER = logging.getLogger("VertexAIModelHandlerJSON")
@@ -59,9 +52,9 @@ def _retry_on_appropriate_gcp_error(exception):
return isinstance(exception, (TooManyRequests, ServerError))
-class VertexAIModelHandlerJSON(ModelHandler[Any,
- PredictionResult,
- aiplatform.Endpoint]):
+class VertexAIModelHandlerJSON(RemoteModelHandler[Any,
+ PredictionResult,
+ aiplatform.Endpoint]):
def __init__(
self,
endpoint_id: str,
@@ -139,14 +132,10 @@ class VertexAIModelHandlerJSON(ModelHandler[Any,
_ = self._retrieve_endpoint(
self.endpoint_name, self.location, self.is_private)
- # Configure AdaptiveThrottler and throttling metrics for client-side
- # throttling behavior.
- # See https://docs.google.com/document/d/1ePorJGZnLbNCmLD9mR7iFYOdPsyDA1rDnTpYnbdrzSU/edit?usp=sharing
- # for more details.
- self.throttled_secs = Metrics.counter(
- VertexAIModelHandlerJSON, "cumulativeThrottlingSeconds")
- self.throttler = AdaptiveThrottler(
- window_ms=1, bucket_ms=1, overload_ratio=2)
+ super().__init__(
+ namespace='VertexAIModelHandlerJSON',
+ retry_filter=_retry_on_appropriate_gcp_error,
+ **kwargs)
def _retrieve_endpoint(
self, endpoint_id: str, location: str,
@@ -183,7 +172,7 @@ class VertexAIModelHandlerJSON(ModelHandler[Any,
return endpoint
- def load_model(self) -> aiplatform.Endpoint:
+ def create_client(self) -> aiplatform.Endpoint:
"""Loads the Endpoint object used to build and send prediction request to
Vertex AI.
"""
@@ -193,39 +182,11 @@ class VertexAIModelHandlerJSON(ModelHandler[Any,
self.endpoint_name, self.location, self.is_private)
return ep
- @retry.with_exponential_backoff(
- num_retries=5, retry_filter=_retry_on_appropriate_gcp_error)
- def get_request(
- self,
- batch: Sequence[Any],
- model: aiplatform.Endpoint,
- throttle_delay_secs: int,
- inference_args: Optional[Dict[str, Any]]):
- while self.throttler.throttle_request(time.time() * MSEC_TO_SEC):
- LOGGER.info(
- "Delaying request for %d seconds due to previous failures",
- throttle_delay_secs)
- time.sleep(throttle_delay_secs)
- self.throttled_secs.inc(throttle_delay_secs)
-
- try:
- req_time = time.time()
- prediction = model.predict(
- instances=list(batch), parameters=inference_args)
- self.throttler.successful_request(req_time * MSEC_TO_SEC)
- return prediction
- except TooManyRequests as e:
- LOGGER.warning("request was limited by the service with code %i", e.code)
- raise
- except Exception as e:
- LOGGER.error("unexpected exception raised as part of request, got %s", e)
- raise
-
- def run_inference(
+ def request(
self,
batch: Sequence[Any],
model: aiplatform.Endpoint,
- inference_args: Optional[Dict[str, Any]] = None
+ inference_args: Optional[dict[str, Any]] = None
) -> Iterable[PredictionResult]:
""" Sends a prediction request to a Vertex AI endpoint containing batch
of inputs and matches that input with the prediction response from
@@ -242,16 +203,11 @@ class VertexAIModelHandlerJSON(ModelHandler[Any,
Returns:
An iterable of Predictions.
"""
-
- # Endpoint.predict returns a Prediction type with the prediction values
- # along with model metadata
- prediction = self.get_request(
- batch, model, throttle_delay_secs=5, inference_args=inference_args)
-
+ prediction = model.predict(instances=list(batch), parameters=inference_args)
return utils._convert_to_result(
batch, prediction.predictions, prediction.deployed_model_id)
- def validate_inference_args(self, inference_args: Optional[Dict[str, Any]]):
+ def validate_inference_args(self, inference_args: Optional[dict[str, Any]]):
pass
def batch_elements_kwargs(self) -> Mapping[str, Any]: