feat: Support PSC-I ingress in Ray Client

yinghsienwu · copybara-github · commit da76253f3b53 · 2024-10-22T15:55:53.000-07:00
PiperOrigin-RevId: 688719900
diff --git a/google/cloud/aiplatform/vertex_ray/client_builder.py b/google/cloud/aiplatform/vertex_ray/client_builder.py
@@ -177,7 +177,10 @@ def connect(self) -> _VertexRayClientContext:
         public_address = self.response.resource_runtime.access_uris.get(
             "RAY_CLIENT_ENDPOINT"
         )
-        if public_address:
+        private_address = self.response.resource_runtime.access_uris.get(
+            "RAY_HEAD_NODE_INTERNAL_IP"
+        )
+        if public_address and not private_address:
             self._credentials = grpc.ssl_channel_credentials()
             bearer_token = _validation_utils.get_bearer_token()
             self._metadata = [
