apache-airflow-providers-google 10.10.0__py3-none-any.whl → 10.10.1__py3-none-any.whl

airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py

@@ -15,17 +15,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""This module contains Google Vertex AI operators.
 
-.. spelling:word-list::
+"""This module contains Google Vertex AI operators."""
 
-    undeployed
-    undeploy
-    Undeploys
-    aiplatform
-    FieldMask
-    unassigns
-"""
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py

@@ -15,16 +15,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""This module contains Google Vertex AI operators.
 
-.. spelling:word-list::
+"""This module contains Google Vertex AI operators."""
 
-    irreproducible
-    codepoints
-    Tensorboard
-    aiplatform
-    myVPC
-"""
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence
@@ -67,7 +60,7 @@ class CreateHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
     :param max_trial_count: Required. The desired total number of Trials.
     :param parallel_trial_count: Required. The desired number of Trials to run in parallel.
     :param worker_pool_specs: Required. The spec of the worker pools including machine type and Docker
-        image. Can provided as a list of dictionaries or list of WorkerPoolSpec proto messages.
+        image. Can be provided as a list of dictionaries or list of WorkerPoolSpec proto messages.
     :param base_output_dir: Optional. GCS output directory of job. If not provided a timestamped
         directory in the staging directory will be used.
     :param custom_job_labels: Optional. The labels with user-defined metadata to organize CustomJobs.

airflow/providers/google/cloud/operators/vertex_ai/model_service.py

@@ -15,13 +15,8 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""This module contains Google Vertex AI operators.
+"""This module contains Google Vertex AI operators."""
 
-.. spelling:word-list::
-
-    aiplatform
-    camelCase
-"""
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence
@@ -50,7 +45,10 @@ class DeleteModelOperator(GoogleCloudBaseOperator):
 
     :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
     :param region: Required. The ID of the Google Cloud region that the service belongs to.
-    :param model_id: Required. The name of the Model resource to be deleted.
+    :param model_id: Required. The ID of the Model resource to be deleted.
+        Could be in format `projects/{project}/locations/{location}/models/{model_id}@{version_id}` or
+        `projects/{project}/locations/{location}/models/{model_id}@{version_alias}` if model
+        has several versions.
     :param retry: Designation of what errors, if any, should be retried.
     :param timeout: The timeout for this request.
     :param metadata: Strings which should be sent along with the request as metadata.
@@ -95,7 +93,7 @@ class DeleteModelOperator(GoogleCloudBaseOperator):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
-
+        self.model_id = self.model_id.rpartition("@")[0] if "@" in self.model_id else self.model_id
         try:
             self.log.info("Deleting model: %s", self.model_id)
             operation = hook.delete_model(
@@ -112,13 +110,91 @@ class DeleteModelOperator(GoogleCloudBaseOperator):
             self.log.info("The Model ID %s does not exist.", self.model_id)
 
 
+class GetModelOperator(GoogleCloudBaseOperator):
+    """
+    Retrieves a Model.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param region: Required. The ID of the Google Cloud region that the service belongs to.
+    :param model_id: Required. The ID of the Model resource to be retrieved.
+        Could be in format `projects/{project}/locations/{location}/models/{model_id}@{version_id}` or
+        `projects/{project}/locations/{location}/models/{model_id}@{version_alias}` if model has
+        several versions.
+    :param retry: Designation of what errors, if any, should be retried.
+    :param timeout: The timeout for this request.
+    :param metadata: Strings which should be sent along with the request as metadata.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    template_fields = ("region", "model_id", "project_id", "impersonation_chain")
+    operator_extra_links = (VertexAIModelLink(),)
+
+    def __init__(
+        self,
+        *,
+        region: str,
+        project_id: str,
+        model_id: str,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.region = region
+        self.project_id = project_id
+        self.model_id = model_id
+        self.retry = retry
+        self.timeout = timeout
+        self.metadata = metadata
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        hook = ModelServiceHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        self.model_id = self.model_id.rpartition("@")[0] if "@" in self.model_id else self.model_id
+        try:
+            self.log.info("Retrieving model: %s", self.model_id)
+            model = hook.get_model(
+                project_id=self.project_id,
+                region=self.region,
+                model_id=self.model_id,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            )
+            self.log.info("Model found. Model ID: %s", self.model_id)
+
+            self.xcom_push(context, key="model_id", value=self.model_id)
+            VertexAIModelLink.persist(context=context, task_instance=self, model_id=self.model_id)
+            return Model.to_dict(model)
+        except NotFound:
+            self.log.info("The Model ID %s does not exist.", self.model_id)
+
+
 class ExportModelOperator(GoogleCloudBaseOperator):
     """
     Exports a trained, exportable Model to a location specified by the user.
 
     :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
     :param region: Required. The ID of the Google Cloud region that the service belongs to.
-    :param model_id: Required. The resource name of the Model to export.
+    :param model_id: Required. The ID of the Model to export.
+        Could be in format `projects/{project}/locations/{location}/models/{model_id}@{version_id}` or
+        `projects/{project}/locations/{location}/models/{model_id}@{version_alias}` if model has
+        several versions.
     :param output_config:  Required. The desired output location and configuration.
     :param retry: Designation of what errors, if any, should be retried.
     :param timeout: The timeout for this request.
@@ -195,8 +271,8 @@ class ListModelsOperator(GoogleCloudBaseOperator):
     :param retry: Designation of what errors, if any, should be retried.
     :param filter: An expression for filtering the results of the request. For field names both
         snake_case and camelCase are supported.
-        -  ``model`` supports = and !=. ``model`` represents the Model ID, i.e. the last segment of the
-        Model's [resource name][google.cloud.aiplatform.v1.Model.name].
+        -  ``model`` supports = and !=. ``model`` represents the Model ID,  Could be in format the
+        last segment of the Model's [resource name][google.cloud.aiplatform.v1.Model.name].
         -  ``display_name`` supports = and !=
         -  ``labels`` supports general map functions that is:
         --  ``labels.key=value`` - key:value equality
@@ -285,7 +361,7 @@ class UploadModelOperator(GoogleCloudBaseOperator):
 
     :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
     :param region: Required. The ID of the Google Cloud region that the service belongs to.
-    :param model:  Required. The Model to create.
+    :param model: Required. The Model to create.
     :param retry: Designation of what errors, if any, should be retried.
     :param timeout: The timeout for this request.
     :param metadata: Strings which should be sent along with the request as metadata.
@@ -349,3 +425,366 @@ class UploadModelOperator(GoogleCloudBaseOperator):
         self.xcom_push(context, key="model_id", value=model_id)
         VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
         return model_resp
+
+
+class ListModelVersionsOperator(GoogleCloudBaseOperator):
+    """
+    Lists Model versions in a Location.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param region: Required. The ID of the Google Cloud region that the service belongs to.
+    :param model_id: Required. The ID of the model to list versions for.
+        Could be in format `projects/{project}/locations/{location}/models/{model_id}@{version_id}` or
+        `projects/{project}/locations/{location}/models/{model_id}@{version_alias}` if model has
+        several versions.
+    :param retry: Designation of what errors, if any, should be retried.
+    :param timeout: The timeout for this request.
+    :param metadata: Strings which should be sent along with the request as metadata.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    template_fields = ("model_id", "region", "project_id", "impersonation_chain")
+
+    def __init__(
+        self,
+        *,
+        region: str,
+        project_id: str,
+        model_id: str,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.region = region
+        self.project_id = project_id
+        self.model_id = model_id
+        self.retry = retry
+        self.timeout = timeout
+        self.metadata = metadata
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        hook = ModelServiceHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        self.log.info("Retrieving versions list from model: %s", self.model_id)
+        results = hook.list_model_versions(
+            project_id=self.project_id,
+            region=self.region,
+            model_id=self.model_id,
+            retry=self.retry,
+            timeout=self.timeout,
+            metadata=self.metadata,
+        )
+        for result in results:
+            model = Model.to_dict(result)
+            self.log.info("Model name: %s;", model["name"])
+            self.log.info("Model version: %s, model alias %s;", model["version_id"], model["version_aliases"])
+        return [Model.to_dict(result) for result in results]
+
+
+class SetDefaultVersionOnModelOperator(GoogleCloudBaseOperator):
+    """
+    Sets the desired Model version as Default.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param region: Required. The ID of the Google Cloud region that the service belongs to.
+    :param model_id: Required. The ID of the model to set as default.
+        Should be in format `projects/{project}/locations/{location}/models/{model_id}@{version_id}` or
+        `projects/{project}/locations/{location}/models/{model_id}@{version_alias}` if model
+        has several versions.
+    :param retry: Designation of what errors, if any, should be retried.
+    :param timeout: The timeout for this request.
+    :param metadata: Strings which should be sent along with the request as metadata.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    template_fields = ("model_id", "project_id", "impersonation_chain")
+    operator_extra_links = (VertexAIModelLink(),)
+
+    def __init__(
+        self,
+        *,
+        region: str,
+        project_id: str,
+        model_id: str,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.region = region
+        self.project_id = project_id
+        self.model_id = model_id
+        self.retry = retry
+        self.timeout = timeout
+        self.metadata = metadata
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        hook = ModelServiceHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        self.log.info(
+            "Setting version %s as default on model %s", self.model_id.rpartition("@")[0], self.model_id
+        )
+
+        updated_model = hook.set_version_as_default(
+            region=self.region,
+            model_id=self.model_id,
+            project_id=self.project_id,
+            retry=self.retry,
+            timeout=self.timeout,
+            metadata=self.metadata,
+        )
+        VertexAIModelLink.persist(context=context, task_instance=self, model_id=self.model_id)
+        return Model.to_dict(updated_model)
+
+
+class AddVersionAliasesOnModelOperator(GoogleCloudBaseOperator):
+    """
+    Adds version aliases for the Model.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param region: Required. The ID of the Google Cloud region that the service belongs to.
+    :param model_id: Required. The ID of the model to add version aliases for.
+        Should be in format `projects/{project}/locations/{location}/models/{model_id}@{version_id}` or
+        `projects/{project}/locations/{location}/models/{model_id}@{version_alias}`.
+    :param version_aliases: List of version aliases to be added to model version.
+    :param retry: Designation of what errors, if any, should be retried.
+    :param timeout: The timeout for this request.
+    :param metadata: Strings which should be sent along with the request as metadata.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    template_fields = ("model_id", "project_id", "impersonation_chain")
+    operator_extra_links = (VertexAIModelLink(),)
+
+    def __init__(
+        self,
+        *,
+        region: str,
+        project_id: str,
+        model_id: str,
+        version_aliases: Sequence[str],
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.region = region
+        self.project_id = project_id
+        self.model_id = model_id
+        self.version_aliases = version_aliases
+        self.retry = retry
+        self.timeout = timeout
+        self.metadata = metadata
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        hook = ModelServiceHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        self.log.info(
+            "Adding aliases %s to model version %s", self.version_aliases, self.model_id.rpartition("@")[0]
+        )
+
+        updated_model = hook.add_version_aliases(
+            region=self.region,
+            model_id=self.model_id,
+            version_aliases=self.version_aliases,
+            project_id=self.project_id,
+            retry=self.retry,
+            timeout=self.timeout,
+            metadata=self.metadata,
+        )
+        VertexAIModelLink.persist(context=context, task_instance=self, model_id=self.model_id)
+        return Model.to_dict(updated_model)
+
+
+class DeleteVersionAliasesOnModelOperator(GoogleCloudBaseOperator):
+    """
+    Deletes version aliases for the Model.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param region: Required. The ID of the Google Cloud region that the service belongs to.
+    :param model_id: Required. The ID of the model to delete version aliases from.
+        Should be in format `projects/{project}/locations/{location}/models/{model_id}@{version_id}` or
+        `projects/{project}/locations/{location}/models/{model_id}@{version_alias}`.
+    :param version_aliases: List of version aliases to be deleted from model version.
+    :param retry: Designation of what errors, if any, should be retried.
+    :param timeout: The timeout for this request.
+    :param metadata: Strings which should be sent along with the request as metadata.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    template_fields = ("model_id", "project_id", "impersonation_chain")
+    operator_extra_links = (VertexAIModelLink(),)
+
+    def __init__(
+        self,
+        *,
+        region: str,
+        project_id: str,
+        model_id: str,
+        version_aliases: Sequence[str],
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.region = region
+        self.project_id = project_id
+        self.model_id = model_id
+        self.version_aliases = version_aliases
+        self.retry = retry
+        self.timeout = timeout
+        self.metadata = metadata
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        hook = ModelServiceHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        self.log.info(
+            "Deleting aliases %s from model version %s",
+            self.version_aliases,
+            self.model_id.rpartition("@")[0],
+        )
+
+        updated_model = hook.delete_version_aliases(
+            region=self.region,
+            model_id=self.model_id,
+            version_aliases=self.version_aliases,
+            project_id=self.project_id,
+            retry=self.retry,
+            timeout=self.timeout,
+            metadata=self.metadata,
+        )
+        VertexAIModelLink.persist(context=context, task_instance=self, model_id=self.model_id)
+        return Model.to_dict(updated_model)
+
+
+class DeleteModelVersionOperator(GoogleCloudBaseOperator):
+    """
+    Delete Model version in a Location.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param region: Required. The ID of the Google Cloud region that the service belongs to.
+    :param model_id: Required. The ID of the Model in which to delete version.
+        Should be in format `projects/{project}/locations/{location}/models/{model_id}@{version_id}` or
+        `projects/{project}/locations/{location}/models/{model_id}@{version_alias}`
+        several versions.
+    :param retry: Designation of what errors, if any, should be retried.
+    :param timeout: The timeout for this request.
+    :param metadata: Strings which should be sent along with the request as metadata.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    template_fields = ("model_id", "project_id", "impersonation_chain")
+
+    def __init__(
+        self,
+        *,
+        region: str,
+        project_id: str,
+        model_id: str,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.region = region
+        self.project_id = project_id
+        self.model_id = model_id
+        self.retry = retry
+        self.timeout = timeout
+        self.metadata = metadata
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        hook = ModelServiceHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        try:
+            self.log.info("Deleting model version: %s", self.model_id)
+            operation = hook.delete_model_version(
+                project_id=self.project_id,
+                region=self.region,
+                model_id=self.model_id,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            )
+            hook.wait_for_operation(timeout=self.timeout, operation=operation)
+            self.log.info("Model version was deleted.")
+        except NotFound:
+            self.log.info("The Model ID %s does not exist.", self.model_id)