fix: Fix dedicated endpoint DNS is empty

vertex-sdk-bot · copybara-github · commit fc6d893b2191 · 2025-05-19T17:14:29.000-07:00
PiperOrigin-RevId: 760826747
diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py
@@ -791,8 +791,15 @@ def dedicated_endpoint_dns(self) -> Optional[str]:
         If dedicated endpoint is not enabled, this property returns None.
         """
         if re.match(r"^projects/.*/endpoints/.*$", self._gca_resource.name):
-            self._assert_gca_resource_is_available()
-            return getattr(self._gca_resource, "dedicated_endpoint_dns", None)
+            dedicated_endpoint_dns = getattr(
+                self._gca_resource, "dedicated_endpoint_dns", None
+            )
+            if self.dedicated_endpoint_enabled and not dedicated_endpoint_dns:
+                self._sync_gca_resource()
+                dedicated_endpoint_dns = getattr(
+                    self._gca_resource, "dedicated_endpoint_dns", None
+                )
+            return dedicated_endpoint_dns
         return None
 
     @property
@@ -2452,7 +2459,7 @@ def predict(
                 model_resource_name=prediction_response.model,
             )
 
-        if self.dedicated_endpoint_dns is None:
+        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."
@@ -2610,7 +2617,7 @@ def raw_predict(
             )
 
         if self.dedicated_endpoint_enabled:
-            if self.dedicated_endpoint_dns is None:
+            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."
@@ -2682,7 +2689,7 @@ def stream_raw_predict(
             )
 
         if self.dedicated_endpoint_enabled:
-            if self.dedicated_endpoint_dns is None:
+            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."
diff --git a/tests/unit/aiplatform/test_endpoints.py b/tests/unit/aiplatform/test_endpoints.py
@@ -364,6 +364,20 @@ def get_dedicated_endpoint_mock():
         yield get_endpoint_mock
 
 
+@pytest.fixture
+def get_dedicated_endpoint_no_dns_mock():
+    with mock.patch.object(
+        endpoint_service_client.EndpointServiceClient, "get_endpoint"
+    ) as get_endpoint_mock:
+        get_endpoint_mock.return_value = gca_endpoint.Endpoint(
+            display_name=_TEST_DISPLAY_NAME,
+            name=_TEST_ENDPOINT_NAME,
+            encryption_spec=_TEST_ENCRYPTION_SPEC,
+            dedicated_endpoint_enabled=True,
+        )
+        yield get_endpoint_mock
+
+
 @pytest.fixture
 def get_model_mock():
     with mock.patch.object(
@@ -2662,6 +2676,21 @@ def test_predict_dedicated_endpoint(self, predict_endpoint_http_mock):
             timeout=None,
         )
 
+    @pytest.mark.usefixtures("get_dedicated_endpoint_no_dns_mock")
+    def test_predict_dedicated_endpoint_without_dns(self, predict_endpoint_http_mock):
+        test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME)
+
+        with pytest.raises(ValueError) as err:
+            test_endpoint.predict(
+                instances=_TEST_INSTANCES,
+                parameters={"param": 3.0},
+                use_dedicated_endpoint=True,
+            )
+        assert err.match(
+            regexp=r"Dedicated endpoint DNS is empty. Please make sure endpoint"
+            "and model are ready before making a prediction."
+        )
+
     @pytest.mark.usefixtures("get_dedicated_endpoint_mock")
     def test_predict_dedicated_endpoint_with_timeout(self, predict_endpoint_http_mock):
         test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME)
