[beam] branch master updated: Allow use of private endpoints for Vertex AI requests (#27458)

[tv...@apache.org]

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

tvalentyn 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 30b1b8c1d72 Allow use of private endpoints for Vertex AI requests (#27458)
30b1b8c1d72 is described below

commit 30b1b8c1d72243b51104deff435539c7dee3fd3c
Author: Jack McCluskey <34...@users.noreply.github.com>
AuthorDate: Wed Jul 12 15:28:12 2023 -0400

    Allow use of private endpoints for Vertex AI requests (#27458)
---
 .../apache_beam/ml/inference/vertex_ai_inference.py  | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 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 34fd6d27448..a0e0d9d3f8f 100644
--- a/sdks/python/apache_beam/ml/inference/vertex_ai_inference.py
+++ b/sdks/python/apache_beam/ml/inference/vertex_ai_inference.py
@@ -66,6 +66,7 @@ class VertexAIModelHandlerJSON(ModelHandler[Any,
       project: str,
       location: str,
       experiment: Optional[str] = None,
+      network: Optional[str] = None,
       **kwargs):
     """Implementation of the ModelHandler interface for Vertex AI.
     **NOTE:** This API and its implementation are under development and
@@ -73,19 +74,30 @@ class VertexAIModelHandlerJSON(ModelHandler[Any,
     Unlike other ModelHandler implementations, this does not load the model
     being used onto the worker and instead makes remote queries to a
     Vertex AI endpoint. In that way it functions more like a mid-pipeline
-    IO. At present this implementation only supports public endpoints with
-    a maximum request size of 1.5 MB.
+    IO. Public Vertex AI endpoints have a maximum request size of 1.5 MB.
+    If you wish to make larger requests and use a private endpoint, provide
+    the Compute Engine network you wish to use.
+
     Args:
       endpoint_id: the numerical ID of the Vertex AI endpoint to query
       project: the GCP project name where the endpoint is deployed
       location: the GCP location where the endpoint is deployed
-      experiment (Optional): experiment label to apply to the queries
+      experiment: optional. experiment label to apply to the
+        queries
+      network: optional. the full name of the Compute Engine
+        network the endpoint is deployed on; used for private
+        endpoints only.
+        Ex: "projects/12345/global/networks/myVPC"
     """
 
     self._env_vars = kwargs.get('env_vars', {})
     # TODO: support the full list of options for aiplatform.init()
     # See https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform#google_cloud_aiplatform_init
-    aiplatform.init(project=project, location=location, experiment=experiment)
+    aiplatform.init(
+        project=project,
+        location=location,
+        experiment=experiment,
+        network=network)
 
     # Check for liveness here but don't try to actually store the endpoint
     # in the class yet