airflow.providers.google.cloud.hooks.vertex_ai.prediction_service

Module Contents

Classes

PredictionServiceHook

Hook for Google Cloud Vertex AI Prediction API.
class airflow.providers.google.cloud.hooks.vertex_ai.prediction_service.PredictionServiceHook(gcp_conn_id='google_cloud_default', impersonation_chain=None, **kwargs)[source]

Bases: airflow.providers.google.common.hooks.base_google.GoogleBaseHook

Hook for Google Cloud Vertex AI Prediction API.

get_prediction_service_client(region=None)[source]

Return PredictionServiceClient object.

Parameters

region (str | None) – The ID of the Google Cloud region that the service belongs to. Default is None.

Returns

google.cloud.aiplatform_v1.services.prediction_service.client.PredictionServiceClient instance.

Return type

google.cloud.aiplatform_v1.PredictionServiceClient

predict(endpoint_id, instances, location, project_id=PROVIDE_PROJECT_ID, parameters=None, retry=DEFAULT, timeout=None, metadata=())[source]

Perform an online prediction and returns the prediction result in the response.

Parameters

  • endpoint_id (str) – Name of the endpoint_id requested to serve the prediction.

  • instances (list[str]) – Required. The instances that are the input to the prediction call. A DeployedModel may have an upper limit on the number of instances it supports per request, and when it is exceeded the prediction call errors in case of AutoML Models, or, in case of customer created Models, the behaviour is as documented by that Model.

  • parameters (dict[str, str] | None) – Additional domain-specific parameters, any string must be up to 25000 characters long.

  • project_id (str) – ID of the Google Cloud project where model is located if None then default project_id is used.

  • location (str) – The location of the project.

  • retry (google.api_core.retry.Retry | google.api_core.gapic_v1.method._MethodDefault) – A retry object used to retry requests. If None is specified, requests will not be retried.

  • timeout (float | None) – The amount of time, in seconds, to wait for the request to complete. Note that if retry is specified, the timeout applies to each individual attempt.

  • metadata (collections.abc.Sequence[tuple[str, str]]) – Additional metadata that is provided to the method.

Was this entry helpful?