feat: Add support for Prediction dedicated endpoint. predict/rawPredict/streamRawPredict can use dedicated DNS to access the dedicated endpoint.

PiperOrigin-RevId: 667018843

Author: vertex-sdk-bot

google/cloud/aiplatform/models.py

@@ -782,6 +782,7 @@ def create(
  enable_request_response_logging=False,
  request_response_logging_sampling_rate: Optional[float] = None,
  request_response_logging_bq_destination_table: Optional[str] = None,
+ dedicated_endpoint_enabled=False,
  ) -> "Endpoint":
  """Creates a new endpoint.
 
@@ -849,6 +850,10 @@ def create(
  request_response_logging_bq_destination_table (str):
  Optional. The request response logging bigquery destination. If not set, will create a table with name:
  ``bq://{project_id}.logging_{endpoint_display_name}_{endpoint_id}.request_response_logging``.
+ dedicated_endpoint_enabled (bool):
+ Optional. If enabled, a dedicated dns will be created and your
+ traffic will be fully isolated from other customers' traffic and
+ latency will be reduced.
 
  Returns:
  endpoint (aiplatform.Endpoint):
@@ -893,6 +898,7 @@ def create(
  create_request_timeout=create_request_timeout,
  endpoint_id=endpoint_id,
  predict_request_response_logging_config=predict_request_response_logging_config,
+ dedicated_endpoint_enabled=dedicated_endpoint_enabled,
  )
 
  @classmethod
@@ -918,6 +924,7 @@ def _create(
  private_service_connect_config: Optional[
  gca_service_networking.PrivateServiceConnectConfig
  ] = None,
+ dedicated_endpoint_enabled=False,
  ) -> "Endpoint":
  """Creates a new endpoint by calling the API client.
 
@@ -984,6 +991,10 @@ def _create(
  private_service_connect_config (aiplatform.service_network.PrivateServiceConnectConfig):
  If enabled, the endpoint can be accessible via [Private Service Connect](https://cloud.google.com/vpc/docs/private-service-connect).
  Cannot be enabled when network is specified.
+ dedicated_endpoint_enabled (bool):
+ Optional. If enabled, a dedicated dns will be created and your
+ traffic will be fully isolated from other customers' traffic and
+ latency will be reduced.
 
  Returns:
  endpoint (aiplatform.Endpoint):
@@ -1002,6 +1013,7 @@ def _create(
  network=network,
  predict_request_response_logging_config=predict_request_response_logging_config,
  private_service_connect_config=private_service_connect_config,
+ dedicated_endpoint_enabled=dedicated_endpoint_enabled,
  )
 
  operation_future = api_client.create_endpoint(
@@ -2167,9 +2179,18 @@ def predict(
  parameters: Optional[Dict] = None,
  timeout: Optional[float] = None,
  use_raw_predict: Optional[bool] = False,
+ *,
+ use_dedicated_endpoint: Optional[bool] = False,
  ) -> Prediction:
  """Make a prediction against this Endpoint.
 
+ For dedicated endpoint, set use_dedicated_endpoint = True:
+ ```
+ response = my_endpoint.predict(instances=[...],
+ use_dedicated_endpoint=True)
+ my_predictions = response.predictions
+ ```
+
  Args:
  instances (List):
  Required. The instances that are the input to the
@@ -2194,6 +2215,9 @@ def predict(
  use_raw_predict (bool):
  Optional. Default value is False. If set to True, the underlying prediction call will be made
  against Endpoint.raw_predict().
+ use_dedicated_endpoint (bool):
+ Optional. Default value is False. If set to True, the underlying prediction call will be made
+ using the dedicated endpoint dns.
 
  Returns:
  prediction (aiplatform.Prediction):
@@ -2204,6 +2228,7 @@ def predict(
  raw_predict_response = self.raw_predict(
  body=json.dumps({"instances": instances, "parameters": parameters}),
  headers={"Content-Type": "application/json"},
+ use_dedicated_endpoint=use_dedicated_endpoint,
  )
  json_response = raw_predict_response.json()
  return Prediction(
@@ -2219,6 +2244,51 @@ def predict(
  _RAW_PREDICT_MODEL_VERSION_ID_KEY, None
  ),
  )
+
+ if use_dedicated_endpoint:
+ self._sync_gca_resource_if_skipped()
+ if (
+ not self._gca_resource.dedicated_endpoint_enabled
+ or self._gca_resource.dedicated_endpoint_dns is None
+ ):
+ raise ValueError(
+ "Dedicated endpoint is not enabled or DNS is empty."
+ "Please make sure endpoint has dedicated endpoint enabled"
+ "and model are ready before making a prediction."
+ )
+
+ if not self.authorized_session:
+ self.credentials._scopes = constants.base.DEFAULT_AUTHED_SCOPES
+ self.authorized_session = google_auth_requests.AuthorizedSession(
+ self.credentials
+ )
+
+ headers = {
+ "Content-Type": "application/json",
+ }
+
+ url = f"https://{self._gca_resource.dedicated_endpoint_dns}/v1/{self.resource_name}:predict"
+ response = self.authorized_session.post(
+ url=url,
+ data=json.dumps(
+ {
+ "instances": instances,
+ "parameters": parameters,
+ }
+ ),
+ headers=headers,
+ )
+
+ prediction_response = json.loads(response.text)
+
+ return Prediction(
+ predictions=prediction_response.get("predictions"),
+ metadata=prediction_response.get("metadata"),
+ deployed_model_id=prediction_response.get("deployedModelId"),
+ model_resource_name=prediction_response.get("model"),
+ model_version_id=prediction_response.get("modelVersionId"),
+ )
+
  else:
  prediction_response = self._prediction_client.predict(
  endpoint=self._gca_resource.name,
@@ -2307,7 +2377,11 @@ async def predict_async(
  )
 
  def raw_predict(
- self, body: bytes, headers: Dict[str, str]
+ self,
+ body: bytes,
+ headers: Dict[str, str],
+ *,
+ use_dedicated_endpoint: Optional[bool] = False,
  ) -> requests.models.Response:
  """Makes a prediction request using arbitrary headers.
 
@@ -2317,6 +2391,12 @@ def raw_predict(
  body = b'{"instances":[{"feat_1":val_1, "feat_2":val_2}]}'
  headers = {'Content-Type':'application/json'}
  )
+ # For dedicated endpoint:
+ response = my_endpoint.raw_predict(
+ body = b'{"instances":[{"feat_1":val_1, "feat_2":val_2}]}',
+ headers = {'Content-Type':'application/json'},
+ dedicated_endpoint=True,
+ )
  status_code = response.status_code
  results = json.dumps(response.text)
 
@@ -2325,6 +2405,9 @@ def raw_predict(
  The body of the prediction request in bytes. This must not exceed 1.5 mb per request.
  headers (Dict[str, str]):
  The header of the request as a dictionary. There are no restrictions on the header.
+ use_dedicated_endpoint (bool):
+ Optional. Default value is False. If set to True, the underlying prediction call will be made
+ using the dedicated endpoint dns.
 
  Returns:
  A requests.models.Response object containing the status code and prediction results.
@@ -2338,12 +2421,29 @@ def raw_predict(
  if self.raw_predict_request_url is None:
  self.raw_predict_request_url = f"https://{self.location}-{constants.base.API_BASE_PATH}/v1/projects/{self.project}/locations/{self.location}/endpoints/{self.name}:rawPredict"
 
- return self.authorized_session.post(
- url=self.raw_predict_request_url, data=body, headers=headers
- )
+ url = self.raw_predict_request_url
+
+ if use_dedicated_endpoint:
+ self._sync_gca_resource_if_skipped()
+ if (
+ not self._gca_resource.dedicated_endpoint_enabled
+ or self._gca_resource.dedicated_endpoint_dns is None
+ ):
+ raise ValueError(
+ "Dedicated endpoint is not enabled or DNS is empty."
+ "Please make sure endpoint has dedicated endpoint enabled"
+ "and model are ready before making a prediction."
+ )
+ url = f"https://{self._gca_resource.dedicated_endpoint_dns}/v1/{self.resource_name}:rawPredict"
+
+ return self.authorized_session.post(url=url, data=body, headers=headers)
 
  def stream_raw_predict(
- self, body: bytes, headers: Dict[str, str]
+ self,
+ body: bytes,
+ headers: Dict[str, str],
+ *,
+ use_dedicated_endpoint: Optional[bool] = False,
  ) -> Iterator[requests.models.Response]:
  """Makes a streaming prediction request using arbitrary headers.
 
@@ -2358,13 +2458,28 @@ def stream_raw_predict(
  stream_result = json.dumps(response.text)
  ```
 
+ For dedicated endpoint:
+ ```
+ my_endpoint = aiplatform.Endpoint(ENDPOINT_ID)
+ for stream_response in my_endpoint.stream_raw_predict(
+ body = b'{"instances":[{"feat_1":val_1, "feat_2":val_2}]}',
+ headers = {'Content-Type':'application/json'},
+ use_dedicated_endpoint=True,
+ ):
+ status_code = response.status_code
+ stream_result = json.dumps(response.text)
+ ```
+
  Args:
  body (bytes):
  The body of the prediction request in bytes. This must not
  exceed 10 mb per request.
  headers (Dict[str, str]):
  The header of the request as a dictionary. There are no
  restrictions on the header.
+ use_dedicated_endpoint (bool):
+ Optional. Default value is False. If set to True, the underlying prediction call will be made
+ using the dedicated endpoint dns.
 
  Yields:
  predictions (Iterator[requests.models.Response]):
@@ -2379,8 +2494,23 @@ def stream_raw_predict(
  if self.stream_raw_predict_request_url is None:
  self.stream_raw_predict_request_url = f"https://{self.location}-{constants.base.API_BASE_PATH}/v1/projects/{self.project}/locations/{self.location}/endpoints/{self.name}:streamRawPredict"
 
+ url = self.raw_predict_request_url
+
+ if use_dedicated_endpoint:
+ self._sync_gca_resource_if_skipped()
+ if (
+ not self._gca_resource.dedicated_endpoint_enabled
+ or self._gca_resource.dedicated_endpoint_dns is None
+ ):
+ raise ValueError(
+ "Dedicated endpoint is not enabled or DNS is empty."
+ "Please make sure endpoint has dedicated endpoint enabled"
+ "and model are ready before making a prediction."
+ )
+ url = f"https://{self._gca_resource.dedicated_endpoint_dns}/v1/{self.resource_name}:streamRawPredict"
+
  with self.authorized_session.post(
- url=self.stream_raw_predict_request_url,
+ url=url,
  data=body,
  headers=headers,
  stream=True,