(beam) branch master updated: Pass location for Vertex endpoint retrieval (#31171)

[jr...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

jrmccluskey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new 7dbdbdfbad4 Pass location for Vertex endpoint retrieval (#31171)
7dbdbdfbad4 is described below

commit 7dbdbdfbad45f50eab1b1de7f5589b3f28903934
Author: Jack McCluskey <34...@users.noreply.github.com>
AuthorDate: Fri May 3 12:50:07 2024 -0400

    Pass location for Vertex endpoint retrieval (#31171)
---
 .../apache_beam/ml/inference/vertex_ai_inference.py       | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

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 95660441a84..4c4163accfb 100644
--- a/sdks/python/apache_beam/ml/inference/vertex_ai_inference.py
+++ b/sdks/python/apache_beam/ml/inference/vertex_ai_inference.py
@@ -133,9 +133,11 @@ class VertexAIModelHandlerJSON(ModelHandler[Any,
     # Check for liveness here but don't try to actually store the endpoint
     # in the class yet
     self.endpoint_name = endpoint_id
+    self.location = location
     self.is_private = private
 
-    _ = self._retrieve_endpoint(self.endpoint_name, self.is_private)
+    _ = self._retrieve_endpoint(
+        self.endpoint_name, self.location, self.is_private)
 
     # Configure AdaptiveThrottler and throttling metrics for client-side
     # throttling behavior.
@@ -147,7 +149,8 @@ class VertexAIModelHandlerJSON(ModelHandler[Any,
         window_ms=1, bucket_ms=1, overload_ratio=2)
 
   def _retrieve_endpoint(
-      self, endpoint_id: str, is_private: bool) -> aiplatform.Endpoint:
+      self, endpoint_id: str, location: str,
+      is_private: bool) -> aiplatform.Endpoint:
     """Retrieves an AI Platform endpoint and queries it for liveness/deployed
     models.
 
@@ -162,10 +165,11 @@ class VertexAIModelHandlerJSON(ModelHandler[Any,
     """
     if is_private:
       endpoint: aiplatform.Endpoint = aiplatform.PrivateEndpoint(
-          endpoint_name=endpoint_id)
+          endpoint_name=endpoint_id, location=location)
       LOGGER.debug("Treating endpoint %s as private", endpoint_id)
     else:
-      endpoint = aiplatform.Endpoint(endpoint_name=endpoint_id)
+      endpoint = aiplatform.Endpoint(
+          endpoint_name=endpoint_id, location=location)
       LOGGER.debug("Treating endpoint %s as public", endpoint_id)
 
     try:
@@ -185,7 +189,8 @@ class VertexAIModelHandlerJSON(ModelHandler[Any,
     """
     # Check to make sure the endpoint is still active since pipeline
     # construction time
-    ep = self._retrieve_endpoint(self.endpoint_name, self.is_private)
+    ep = self._retrieve_endpoint(
+        self.endpoint_name, self.location, self.is_private)
     return ep
 
   @retry.with_exponential_backoff(