From 92204b36b08044db4ed958dfaed7672071cd37d4 Mon Sep 17 00:00:00 2001 From: VladaZakharova <80038284+VladaZakharova@users.noreply.github.com> Date: Wed, 18 Oct 2023 10:45:22 +0200 Subject: [PATCH] Migrate legacy version of AI Platform Prediction to VertexAI (#34922) --- .../google/cloud/hooks/vertex_ai/auto_ml.py | 158 +++++- .../hooks/vertex_ai/batch_prediction_job.py | 10 +- .../cloud/hooks/vertex_ai/custom_job.py | 92 +++- .../cloud/hooks/vertex_ai/endpoint_service.py | 12 +- .../vertex_ai/hyperparameter_tuning_job.py | 11 +- .../cloud/hooks/vertex_ai/model_service.py | 226 +++++++- .../google/cloud/operators/mlengine.py | 116 +++++ .../cloud/operators/vertex_ai/auto_ml.py | 45 ++ .../vertex_ai/batch_prediction_job.py | 10 +- .../cloud/operators/vertex_ai/custom_job.py | 488 ++++++++++-------- .../operators/vertex_ai/endpoint_service.py | 10 +- .../vertex_ai/hyperparameter_tuning_job.py | 11 +- .../operators/vertex_ai/model_service.py | 463 ++++++++++++++++- .../cloud/utils/mlengine_operator_utils.py | 8 +- .../operators/cloud/mlengine.rst | 123 ++--- .../operators/cloud/vertex_ai.rst | 75 +++ docs/spelling_wordlist.txt | 32 ++ tests/always/test_project_structure.py | 9 + .../hooks/vertex_ai/test_model_service.py | 195 +++++++ .../google/cloud/operators/test_vertex_ai.py | 209 ++++++++ .../cloud/ml_engine/example_mlengine.py | 407 +++++++-------- .../cloud/ml_engine/example_mlengine_async.py | 329 ------------ ...xample_vertex_ai_auto_ml_video_training.py | 16 + .../vertex_ai/example_vertex_ai_custom_job.py | 22 + .../example_vertex_ai_model_service.py | 84 +++ 25 files changed, 2245 insertions(+), 916 deletions(-) delete mode 100644 tests/system/providers/google/cloud/ml_engine/example_mlengine_async.py diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py b/airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py index 16e6b1c87c335..28062b444351b 100644 --- a/airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py +++ b/airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py @@ -15,33 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -""" -This module contains a Google Cloud Vertex AI hook. - -.. spelling:word-list:: - - aiplatform - au - codepoints - milli - mae - quantile - quantiles - Quantiles - rmse - rmsle - rmspe - wape - prc - roc - Jetson - forecasted - Struct - sentimentMax - TrainingPipeline - targetColumn - optimizationObjective -""" +"""This module contains a Google Cloud Vertex AI hook.""" from __future__ import annotations import warnings @@ -314,6 +288,10 @@ def create_auto_ml_tabular_training_job( export_evaluated_data_items_bigquery_destination_uri: str | None = None, export_evaluated_data_items_override_destination: bool = False, sync: bool = True, + parent_model: str | None = None, + is_default_version: bool | None = None, + model_version_aliases: list[str] | None = None, + model_version_description: str | None = None, ) -> tuple[models.Model | None, str]: """ Create an AutoML Tabular Training Job. @@ -327,6 +305,24 @@ def create_auto_ml_tabular_training_job( [google.cloud.aiplatform.v1beta1.TrainingPipeline.training_task_definition]. For tabular Datasets, all their data is exported to training, to pick and choose from. :param target_column: Required. The name of the column values of which the Model is to predict. + :param parent_model: Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + Only set this field when training a new version of an existing model. + :param is_default_version: Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + New model uploads, i.e. version 1, will always be "default" aliased. + :param model_version_aliases: Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + :param model_version_description: Optional. The description of the model version + being uploaded by this job. :param optimization_prediction_type: The type of prediction the Model is to produce. "classification" - Predict one out of multiple target values is picked for each row. "regression" - Predict a value based on its relation to other values. This type is available only @@ -498,6 +494,10 @@ def create_auto_ml_tabular_training_job( ), export_evaluated_data_items_override_destination=export_evaluated_data_items_override_destination, sync=sync, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, ) training_id = self.extract_training_id(self._job.resource_name) if model: @@ -546,6 +546,10 @@ def create_auto_ml_forecasting_training_job( model_display_name: str | None = None, model_labels: dict[str, str] | None = None, sync: bool = True, + parent_model: str | None = None, + is_default_version: bool | None = None, + model_version_aliases: list[str] | None = None, + model_version_description: str | None = None, ) -> tuple[models.Model | None, str]: """ Create an AutoML Forecasting Training Job. @@ -560,6 +564,24 @@ def create_auto_ml_forecasting_training_job( Datasets, all their data is exported to training, to pick and choose from. :param target_column: Required. Name of the column that the Model is to predict values for. :param time_column: Required. Name of the column that identifies time order in the time series. + :param parent_model: Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + Only set this field when training a new version of an existing model. + :param is_default_version: Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + New model uploads, i.e. version 1, will always be "default" aliased. + :param model_version_aliases: Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + :param model_version_description: Optional. The description of the model version + being uploaded by this job. :param time_series_identifier_column: Required. Name of the column that identifies the time series. :param unavailable_at_forecast_columns: Required. Column names of columns that are unavailable at forecast. Each column contains information for the given entity (identified by the @@ -731,6 +753,10 @@ def create_auto_ml_forecasting_training_job( model_display_name=model_display_name, model_labels=model_labels, sync=sync, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, ) training_id = self.extract_training_id(self._job.resource_name) if model: @@ -767,6 +793,10 @@ def create_auto_ml_image_training_job( model_labels: dict[str, str] | None = None, disable_early_stopping: bool = False, sync: bool = True, + parent_model: str | None = None, + is_default_version: bool | None = None, + model_version_aliases: list[str] | None = None, + model_version_description: str | None = None, ) -> tuple[models.Model | None, str]: """ Create an AutoML Image Training Job. @@ -784,6 +814,24 @@ def create_auto_ml_image_training_job( "object_detection" - Predict a value based on its relation to other values. This type is available only to columns that contain semantically numeric values, i.e. integers or floating point number, even if stored as e.g. strings. + :param parent_model: Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + Only set this field when training a new version of an existing model. + :param is_default_version: Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + New model uploads, i.e. version 1, will always be "default" aliased. + :param model_version_aliases: Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + :param model_version_description: Optional. The description of the model version + being uploaded by this job. :param multi_label: Required. Default is False. If false, a single-label (multi-class) Model will be trained (i.e. assuming that for each image just up to one annotation may be applicable). If true, a multi-label Model will be trained (i.e. assuming that for each image multiple annotations may @@ -907,6 +955,10 @@ def create_auto_ml_image_training_job( model_labels=model_labels, disable_early_stopping=disable_early_stopping, sync=sync, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, ) training_id = self.extract_training_id(self._job.resource_name) if model: @@ -940,6 +992,10 @@ def create_auto_ml_text_training_job( model_display_name: str | None = None, model_labels: dict[str, str] | None = None, sync: bool = True, + parent_model: str | None = None, + is_default_version: bool | None = None, + model_version_aliases: list[str] | None = None, + model_version_description: str | None = None, ) -> tuple[models.Model | None, str]: """ Create an AutoML Text Training Job. @@ -960,6 +1016,24 @@ def create_auto_ml_text_training_job( "sentiment" - A sentiment analysis model inspects text data and identifies the prevailing emotional opinion within it, especially to determine a writer's attitude as positive, negative, or neutral. + :param parent_model: Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + Only set this field when training a new version of an existing model. + :param is_default_version: Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + New model uploads, i.e. version 1, will always be "default" aliased. + :param model_version_aliases: Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + :param model_version_description: Optional. The description of the model version + being uploaded by this job. :param multi_label: Required and only applicable for text classification task. If false, a single-label (multi-class) Model will be trained (i.e. assuming that for each text snippet just up to one annotation may be applicable). If true, a multi-label Model will be trained (i.e. @@ -1044,6 +1118,10 @@ def create_auto_ml_text_training_job( model_display_name=model_display_name, model_labels=model_labels, sync=sync, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, ) training_id = self.extract_training_id(self._job.resource_name) if model: @@ -1074,6 +1152,10 @@ def create_auto_ml_video_training_job( model_display_name: str | None = None, model_labels: dict[str, str] | None = None, sync: bool = True, + parent_model: str | None = None, + is_default_version: bool | None = None, + model_version_aliases: list[str] | None = None, + model_version_description: str | None = None, ) -> tuple[models.Model | None, str]: """ Create an AutoML Video Training Job. @@ -1094,6 +1176,24 @@ def create_auto_ml_video_training_job( pre-defined, custom labels. "action_recognition" - A video action recognition model pinpoints the location of actions with short temporal durations (~1 second). + :param parent_model: Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + Only set this field when training a new version of an existing model. + :param is_default_version: Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + New model uploads, i.e. version 1, will always be "default" aliased. + :param model_version_aliases: Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + :param model_version_description: Optional. The description of the model version + being uploaded by this job. :param model_type: Required. One of the following: "CLOUD" - available for "classification", "object_tracking" and "action_recognition" A Model best tailored to be used within Google Cloud, and which cannot be exported. @@ -1175,6 +1275,10 @@ def create_auto_ml_video_training_job( model_display_name=model_display_name, model_labels=model_labels, sync=sync, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, ) training_id = self.extract_training_id(self._job.resource_name) if model: diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py b/airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py index 37f623aa00710..f5681c9063dd3 100644 --- a/airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py +++ b/airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py @@ -15,15 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""This module contains a Google Cloud Vertex AI hook. - -.. spelling:word-list:: - - jsonl - codepoints - aiplatform - gapic -""" +"""This module contains a Google Cloud Vertex AI hook.""" from __future__ import annotations from typing import TYPE_CHECKING, Sequence diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py b/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py index 6acee36701dcf..ae8141e8749cf 100644 --- a/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py +++ b/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py @@ -303,6 +303,10 @@ def _run_job( timestamp_split_column_name: str | None = None, tensorboard: str | None = None, sync=True, + parent_model: str | None = None, + is_default_version: bool | None = None, + model_version_aliases: list[str] | None = None, + model_version_description: str | None = None, ) -> tuple[models.Model | None, str, str]: """Run Job for training pipeline.""" model = job.run( @@ -332,6 +336,10 @@ def _run_job( timestamp_split_column_name=timestamp_split_column_name, tensorboard=tensorboard, sync=sync, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, ) training_id = self.extract_training_id(job.resource_name) custom_job_id = self.extract_custom_job_id( @@ -361,7 +369,7 @@ def cancel_pipeline_job( """ Cancels a PipelineJob. - Starts asynchronous cancellation on the PipelineJob. The server makes a best + Starts asynchronous cancellation on the PipelineJob. The server makes the best effort to cancel the pipeline, but success is not guaranteed. Clients can use [PipelineService.GetPipelineJob][google.cloud.aiplatform.v1.PipelineService.GetPipelineJob] or other methods to check whether the cancellation succeeded or whether the pipeline completed despite @@ -403,7 +411,7 @@ def cancel_training_pipeline( Cancels a TrainingPipeline. Starts asynchronous cancellation on the TrainingPipeline. The server makes - a best effort to cancel the pipeline, but success is not guaranteed. Clients can use + the best effort to cancel the pipeline, but success is not guaranteed. Clients can use [PipelineService.GetTrainingPipeline][google.cloud.aiplatform.v1.PipelineService.GetTrainingPipeline] or other methods to check whether the cancellation succeeded or whether the pipeline completed despite cancellation. On successful cancellation, the TrainingPipeline is not deleted; instead it becomes a @@ -443,7 +451,7 @@ def cancel_custom_job( """ Cancels a CustomJob. - Starts asynchronous cancellation on the CustomJob. The server makes a best effort + Starts asynchronous cancellation on the CustomJob. The server makes the best effort to cancel the job, but success is not guaranteed. Clients can use [JobService.GetCustomJob][google.cloud.aiplatform.v1.JobService.GetCustomJob] or other methods to check whether the cancellation succeeded or whether the job completed despite cancellation. On @@ -599,6 +607,10 @@ def create_custom_container_training_job( model_instance_schema_uri: str | None = None, model_parameters_schema_uri: str | None = None, model_prediction_schema_uri: str | None = None, + parent_model: str | None = None, + is_default_version: bool | None = None, + model_version_aliases: list[str] | None = None, + model_version_description: str | None = None, labels: dict[str, str] | None = None, training_encryption_spec_key_name: str | None = None, model_encryption_spec_key_name: str | None = None, @@ -716,6 +728,24 @@ def create_custom_container_training_job( and probably different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access. + :param parent_model: Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + Only set this field when training a new version of an existing model. + :param is_default_version: Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + New model uploads, i.e. version 1, will always be "default" aliased. + :param model_version_aliases: Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + :param model_version_description: Optional. The description of the model version + being uploaded by this job. :param project_id: Project to run training in. :param region: Location to run training in. :param labels: Optional. The labels with user-defined metadata to @@ -933,6 +963,10 @@ def create_custom_container_training_job( timestamp_split_column_name=timestamp_split_column_name, tensorboard=tensorboard, sync=sync, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, ) return model, training_id, custom_job_id @@ -990,6 +1024,10 @@ def create_custom_python_package_training_job( predefined_split_column_name: str | None = None, timestamp_split_column_name: str | None = None, tensorboard: str | None = None, + parent_model: str | None = None, + is_default_version: bool | None = None, + model_version_aliases: list[str] | None = None, + model_version_description: str | None = None, sync=True, ) -> tuple[models.Model | None, str, str]: """ @@ -1074,6 +1112,24 @@ def create_custom_python_package_training_job( and probably different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access. + :param parent_model: Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + Only set this field when training a new version of an existing model. + :param is_default_version: Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + New model uploads, i.e. version 1, will always be "default" aliased. + :param model_version_aliases: Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + :param model_version_description: Optional. The description of the model version + being uploaded by this job. :param project_id: Project to run training in. :param region: Location to run training in. :param labels: Optional. The labels with user-defined metadata to @@ -1291,6 +1347,10 @@ def create_custom_python_package_training_job( timestamp_split_column_name=timestamp_split_column_name, tensorboard=tensorboard, sync=sync, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, ) return model, training_id, custom_job_id @@ -1315,6 +1375,10 @@ def create_custom_training_job( model_instance_schema_uri: str | None = None, model_parameters_schema_uri: str | None = None, model_prediction_schema_uri: str | None = None, + parent_model: str | None = None, + is_default_version: bool | None = None, + model_version_aliases: list[str] | None = None, + model_version_description: str | None = None, labels: dict[str, str] | None = None, training_encryption_spec_key_name: str | None = None, model_encryption_spec_key_name: str | None = None, @@ -1432,6 +1496,24 @@ def create_custom_training_job( and probably different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access. + :param parent_model: Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + Only set this field when training a new version of an existing model. + :param is_default_version: Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + New model uploads, i.e. version 1, will always be "default" aliased. + :param model_version_aliases: Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + :param model_version_description: Optional. The description of the model version + being uploaded by this job. :param project_id: Project to run training in. :param region: Location to run training in. :param labels: Optional. The labels with user-defined metadata to @@ -1649,6 +1731,10 @@ def create_custom_training_job( timestamp_split_column_name=timestamp_split_column_name, tensorboard=tensorboard, sync=sync, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, ) return model, training_id, custom_job_id diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py b/airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py index c9a12d6077087..11cd35943d413 100644 --- a/airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py +++ b/airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py @@ -15,17 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""This module contains a Google Cloud Vertex AI hook. - -.. spelling:word-list:: - - undeployed - undeploy - Undeploys - aiplatform - FieldMask - unassigns -""" +"""This module contains a Google Cloud Vertex AI hook.""" from __future__ import annotations from typing import TYPE_CHECKING, Sequence diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py b/airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py index 1592ab4551579..fc2bc3da75454 100644 --- a/airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py +++ b/airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py @@ -15,16 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""This module contains a Google Cloud Vertex AI hook. - -.. spelling:word-list:: - - irreproducible - codepoints - Tensorboard - aiplatform - myVPC -""" +"""This module contains a Google Cloud Vertex AI hook.""" from __future__ import annotations from typing import TYPE_CHECKING, Sequence diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py b/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py index 4565a9e55258e..f9dfb04314517 100644 --- a/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py +++ b/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py @@ -15,13 +15,9 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""This module contains a Google Cloud Vertex AI hook. -.. spelling:word-list:: +"""This module contains a Google Cloud Vertex AI hook.""" - aiplatform - camelCase -""" from __future__ import annotations from typing import TYPE_CHECKING, Sequence @@ -36,7 +32,10 @@ if TYPE_CHECKING: from google.api_core.operation import Operation from google.api_core.retry import Retry - from google.cloud.aiplatform_v1.services.model_service.pagers import ListModelsPager + from google.cloud.aiplatform_v1.services.model_service.pagers import ( + ListModelsPager, + ListModelVersionsPager, + ) from google.cloud.aiplatform_v1.types import Model, model_service @@ -236,3 +235,218 @@ def upload_model( metadata=metadata, ) return result + + @GoogleBaseHook.fallback_to_default_project_id + def list_model_versions( + self, + region: str, + project_id: str, + model_id: str, + retry: Retry | _MethodDefault = DEFAULT, + timeout: float | None = None, + metadata: Sequence[tuple[str, str]] = (), + ) -> ListModelVersionsPager: + """ + Lists all versions of the existing 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 output versions for. + :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. + """ + client = self.get_model_service_client(region) + name = client.model_path(project_id, region, model_id) + + result = client.list_model_versions( + request={ + "name": name, + }, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return result + + @GoogleBaseHook.fallback_to_default_project_id + def delete_model_version( + self, + region: str, + project_id: str, + model_id: str, + retry: Retry | _MethodDefault = DEFAULT, + timeout: float | None = None, + metadata: Sequence[tuple[str, str]] = (), + ) -> Operation: + """ + Deletes version of the Model. The version could not be deleted if this version is 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 in which to delete 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. + """ + client = self.get_model_service_client(region) + name = client.model_path(project_id, region, model_id) + + result = client.delete_model_version( + request={ + "name": name, + }, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return result + + @GoogleBaseHook.fallback_to_default_project_id + def get_model( + self, + region: str, + project_id: str, + model_id: str, + retry: Retry | _MethodDefault = DEFAULT, + timeout: float | None = None, + metadata: Sequence[tuple[str, str]] = (), + ) -> Model: + """ + Retrieves Model of specific name and version. If version is not specified, the default is retrieved. + + :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 retrieve. + :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. + """ + client = self.get_model_service_client(region) + name = client.model_path(project_id, region, model_id) + + result = client.get_model( + request={ + "name": name, + }, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return result + + @GoogleBaseHook.fallback_to_default_project_id + def set_version_as_default( + self, + region: str, + model_id: str, + project_id: str, + retry: Retry | _MethodDefault = DEFAULT, + timeout: float | None = None, + metadata: Sequence[tuple[str, str]] = (), + ) -> Model: + """ + Set current version of the Model 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. + :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. + """ + client = self.get_model_service_client(region) + name = client.model_path(project_id, region, model_id) + + result = client.merge_version_aliases( + request={ + "name": name, + "version_aliases": ["default"], + }, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return result + + @GoogleBaseHook.fallback_to_default_project_id + def add_version_aliases( + self, + region: str, + model_id: str, + project_id: str, + version_aliases: Sequence[str], + retry: Retry | _MethodDefault = DEFAULT, + timeout: float | None = None, + metadata: Sequence[tuple[str, str]] = (), + ) -> Model: + """ + Add list of version aliases to specific version of 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 aliases to. + :param version_aliases: Required. List of version aliases to be added for specific 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. + """ + client = self.get_model_service_client(region) + name = client.model_path(project_id, region, model_id) + + for alias in version_aliases: + if alias.startswith("-"): + raise AirflowException("Name of the alias can't start with '-'") + + result = client.merge_version_aliases( + request={ + "name": name, + "version_aliases": version_aliases, + }, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return result + + @GoogleBaseHook.fallback_to_default_project_id + def delete_version_aliases( + self, + region: str, + model_id: str, + project_id: str, + version_aliases: Sequence[str], + retry: Retry | _MethodDefault = DEFAULT, + timeout: float | None = None, + metadata: Sequence[tuple[str, str]] = (), + ) -> Model: + """ + Delete list of version aliases of specific version of 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 aliases from. + :param version_aliases: Required. List of version aliases to be deleted from specific 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. + """ + client = self.get_model_service_client(region) + name = client.model_path(project_id, region, model_id) + if "default" in version_aliases: + raise AirflowException( + "Default alias can't be deleted. " + "Make sure to assign this alias to another version before deletion" + ) + aliases_for_delete = ["-" + alias for alias in version_aliases] + + result = client.merge_version_aliases( + request={ + "name": name, + "version_aliases": aliases_for_delete, + }, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return result diff --git a/airflow/providers/google/cloud/operators/mlengine.py b/airflow/providers/google/cloud/operators/mlengine.py index 30d1225c0561e..9758d411bb446 100644 --- a/airflow/providers/google/cloud/operators/mlengine.py +++ b/airflow/providers/google/cloud/operators/mlengine.py @@ -82,6 +82,10 @@ class MLEngineStartBatchPredictionJobOperator(GoogleCloudBaseOperator): """ Start a Google Cloud ML Engine prediction job. + This operator is deprecated. Please use + :class:`airflow.providers.google.cloud.operators.vertex_ai.batch_prediction.CreateBatchPredictionJobOperator` + instead. + .. seealso:: For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:MLEngineStartBatchPredictionJobOperator` @@ -210,6 +214,14 @@ def __init__( self._labels = labels self._impersonation_chain = impersonation_chain + warnings.warn( + "This operator is deprecated. All the functionality of legacy " + "MLEngine and new features are available on the Vertex AI platform. " + "Please use `CreateBatchPredictionJobOperator`", + AirflowProviderDeprecationWarning, + stacklevel=3, + ) + if not self._project_id: raise AirflowException("Google Cloud project id is required.") if not self._job_id: @@ -364,6 +376,9 @@ class MLEngineCreateModelOperator(GoogleCloudBaseOperator): """ Creates a new model. + This operator is deprecated. Please use appropriate VertexAI operator from + :class:`airflow.providers.google.cloud.operators.vertex_ai` instead. + .. seealso:: For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:MLEngineCreateModelOperator` @@ -407,6 +422,14 @@ def __init__( self._gcp_conn_id = gcp_conn_id self._impersonation_chain = impersonation_chain + warnings.warn( + "This operator is deprecated. All the functionality of legacy " + "MLEngine and new features are available on the Vertex AI platform. " + "Please use appropriate VertexAI operator.", + AirflowProviderDeprecationWarning, + stacklevel=3, + ) + def execute(self, context: Context): hook = MLEngineHook( gcp_conn_id=self._gcp_conn_id, @@ -429,6 +452,9 @@ class MLEngineGetModelOperator(GoogleCloudBaseOperator): """ Gets a particular model. + This operator is deprecated. Please use + :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.GetModelOperator` instead. + .. seealso:: For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:MLEngineGetModelOperator` @@ -472,6 +498,14 @@ def __init__( self._gcp_conn_id = gcp_conn_id self._impersonation_chain = impersonation_chain + warnings.warn( + "This operator is deprecated. All the functionality of legacy " + "MLEngine and new features are available on the Vertex AI platform. " + "Please use `GetModelOperator`", + AirflowProviderDeprecationWarning, + stacklevel=3, + ) + def execute(self, context: Context): hook = MLEngineHook( gcp_conn_id=self._gcp_conn_id, @@ -493,6 +527,10 @@ class MLEngineDeleteModelOperator(GoogleCloudBaseOperator): """ Deletes a model. + This operator is deprecated. Please use + :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.DeleteModelOperator` instead. + + .. seealso:: For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:MLEngineDeleteModelOperator` @@ -541,6 +579,14 @@ def __init__( self._gcp_conn_id = gcp_conn_id self._impersonation_chain = impersonation_chain + warnings.warn( + "This operator is deprecated. All the functionality of legacy " + "MLEngine and new features are available on the Vertex AI platform. " + "Please use `DeleteModelOperator`", + AirflowProviderDeprecationWarning, + stacklevel=3, + ) + def execute(self, context: Context): hook = MLEngineHook( gcp_conn_id=self._gcp_conn_id, @@ -682,6 +728,8 @@ class MLEngineCreateVersionOperator(GoogleCloudBaseOperator): """ Creates a new version in the model. + This operator is deprecated. Please use parent_model parameter of VertexAI operators instead. + .. seealso:: For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:MLEngineCreateVersionOperator` @@ -731,6 +779,14 @@ def __init__( self._impersonation_chain = impersonation_chain self._validate_inputs() + warnings.warn( + "This operator is deprecated. All the functionality of legacy " + "MLEngine and new features are available on the Vertex AI platform. " + "Please use parent_model parameter for VertexAI operators instead.", + AirflowProviderDeprecationWarning, + stacklevel=3, + ) + def _validate_inputs(self): if not self._model_name: raise AirflowException("The model_name parameter could not be empty.") @@ -763,6 +819,10 @@ class MLEngineSetDefaultVersionOperator(GoogleCloudBaseOperator): """ Sets a version in the model. + This operator is deprecated. Please use + :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.SetDefaultVersionOnModelOperator` + instead. + .. seealso:: For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:MLEngineSetDefaultVersionOperator` @@ -812,6 +872,14 @@ def __init__( self._impersonation_chain = impersonation_chain self._validate_inputs() + warnings.warn( + "This operator is deprecated. All the functionality of legacy " + "MLEngine and new features are available on the Vertex AI platform. " + "Please use `SetDefaultVersionOnModelOperator` instead.", + AirflowProviderDeprecationWarning, + stacklevel=3, + ) + def _validate_inputs(self): if not self._model_name: raise AirflowException("The model_name parameter could not be empty.") @@ -844,6 +912,10 @@ class MLEngineListVersionsOperator(GoogleCloudBaseOperator): """ Lists all available versions of the model. + This operator is deprecated. Please use + :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.ListModelVersionsOperator` + instead. + .. seealso:: For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:MLEngineListVersionsOperator` @@ -889,6 +961,14 @@ def __init__( self._impersonation_chain = impersonation_chain self._validate_inputs() + warnings.warn( + "This operator is deprecated. All the functionality of legacy " + "MLEngine and new features are available on the Vertex AI platform. " + "Please use `ListModelVersionsOperator` instead.", + AirflowProviderDeprecationWarning, + stacklevel=3, + ) + def _validate_inputs(self): if not self._model_name: raise AirflowException("The model_name parameter could not be empty.") @@ -918,6 +998,10 @@ class MLEngineDeleteVersionOperator(GoogleCloudBaseOperator): """ Deletes the version from the model. + This operator is deprecated. Please use + :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.DeleteModelVersionOperator` + instead. + .. seealso:: For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:MLEngineDeleteVersionOperator` @@ -967,6 +1051,14 @@ def __init__( self._impersonation_chain = impersonation_chain self._validate_inputs() + warnings.warn( + "This operator is deprecated. All the functionality of legacy " + "MLEngine and new features are available on the Vertex AI platform. " + "Please use `DeleteModelVersionOperator` instead.", + AirflowProviderDeprecationWarning, + stacklevel=3, + ) + def _validate_inputs(self): if not self._model_name: raise AirflowException("The model_name parameter could not be empty.") @@ -998,6 +1090,10 @@ class MLEngineStartTrainingJobOperator(GoogleCloudBaseOperator): """ Operator for launching a MLEngine training job. + This operator is deprecated. Please use + :class:`airflow.providers.google.cloud.operators.vertex_ai.custom_job.CreateCustomPythonPackageTrainingJobOperator` + instead. + .. seealso:: For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:MLEngineStartTrainingJobOperator` @@ -1124,6 +1220,14 @@ def __init__( self.deferrable = deferrable self.cancel_on_kill = cancel_on_kill + warnings.warn( + "This operator is deprecated. All the functionality of legacy " + "MLEngine and new features are available on the Vertex AI platform. " + "Please use `CreateCustomPythonPackageTrainingJobOperator` instead.", + AirflowProviderDeprecationWarning, + stacklevel=3, + ) + custom = self._scale_tier is not None and self._scale_tier.upper() == "CUSTOM" custom_image = ( custom @@ -1328,6 +1432,10 @@ class MLEngineTrainingCancelJobOperator(GoogleCloudBaseOperator): """ Operator for cleaning up failed MLEngine training job. + This operator is deprecated. Please use + :class:`airflow.providers.google.cloud.operators.vertex_ai.custom_job.CancelCustomTrainingJobOperator` + instead. + :param job_id: A unique templated id for the submitted Google MLEngine training job. (templated) :param project_id: The Google Cloud project name within which MLEngine training job should run. @@ -1366,6 +1474,14 @@ def __init__( self._gcp_conn_id = gcp_conn_id self._impersonation_chain = impersonation_chain + warnings.warn( + "This operator is deprecated. All the functionality of legacy " + "MLEngine and new features are available on the Vertex AI platform. " + "Please use `CancelCustomTrainingJobOperator` instead.", + AirflowProviderDeprecationWarning, + stacklevel=3, + ) + if not self._project_id: raise AirflowException("Google Cloud project id is required.") diff --git a/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py b/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py index d9b0426695104..465677a81b3b1 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py +++ b/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py @@ -15,7 +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.""" + from __future__ import annotations from typing import TYPE_CHECKING, Sequence @@ -50,6 +52,10 @@ def __init__( region: str, display_name: str, labels: dict[str, str] | None = None, + parent_model: str | None = None, + is_default_version: bool | None = None, + model_version_aliases: list[str] | None = None, + model_version_description: str | None = None, training_encryption_spec_key_name: str | None = None, model_encryption_spec_key_name: str | None = None, # RUN @@ -67,6 +73,10 @@ def __init__( self.region = region self.display_name = display_name self.labels = labels + self.parent_model = parent_model + self.is_default_version = is_default_version + self.model_version_aliases = model_version_aliases + self.model_version_description = model_version_description self.training_encryption_spec_key_name = training_encryption_spec_key_name self.model_encryption_spec_key_name = model_encryption_spec_key_name # START Run param @@ -90,6 +100,7 @@ class CreateAutoMLForecastingTrainingJobOperator(AutoMLTrainingJobBaseOperator): """Create AutoML Forecasting Training job.""" template_fields = ( + "parent_model", "dataset_id", "region", "impersonation_chain", @@ -158,11 +169,16 @@ def execute(self, context: Context): gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, ) + self.parent_model = self.parent_model.rpartition("@")[0] if self.parent_model else None model, training_id = self.hook.create_auto_ml_forecasting_training_job( project_id=self.project_id, region=self.region, display_name=self.display_name, dataset=datasets.TimeSeriesDataset(dataset_name=self.dataset_id), + parent_model=self.parent_model, + is_default_version=self.is_default_version, + model_version_aliases=self.model_version_aliases, + model_version_description=self.model_version_description, target_column=self.target_column, time_column=self.time_column, time_series_identifier_column=self.time_series_identifier_column, @@ -202,6 +218,7 @@ def execute(self, context: Context): if model: result = Model.to_dict(model) model_id = self.hook.extract_model_id(result) + self.xcom_push(context, key="model_id", value=model_id) VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id) else: result = model # type: ignore @@ -214,6 +231,7 @@ class CreateAutoMLImageTrainingJobOperator(AutoMLTrainingJobBaseOperator): """Create Auto ML Image Training job.""" template_fields = ( + "parent_model", "dataset_id", "region", "impersonation_chain", @@ -254,11 +272,16 @@ def execute(self, context: Context): gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, ) + self.parent_model = self.parent_model.rpartition("@")[0] if self.parent_model else None model, training_id = self.hook.create_auto_ml_image_training_job( project_id=self.project_id, region=self.region, display_name=self.display_name, dataset=datasets.ImageDataset(dataset_name=self.dataset_id), + parent_model=self.parent_model, + is_default_version=self.is_default_version, + model_version_aliases=self.model_version_aliases, + model_version_description=self.model_version_description, prediction_type=self.prediction_type, multi_label=self.multi_label, model_type=self.model_type, @@ -282,6 +305,7 @@ def execute(self, context: Context): if model: result = Model.to_dict(model) model_id = self.hook.extract_model_id(result) + self.xcom_push(context, key="model_id", value=model_id) VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id) else: result = model # type: ignore @@ -294,6 +318,7 @@ class CreateAutoMLTabularTrainingJobOperator(AutoMLTrainingJobBaseOperator): """Create Auto ML Tabular Training job.""" template_fields = ( + "parent_model", "dataset_id", "region", "impersonation_chain", @@ -351,6 +376,7 @@ def execute(self, context: Context): impersonation_chain=self.impersonation_chain, ) credentials, _ = self.hook.get_credentials_and_project_id() + self.parent_model = self.parent_model.rpartition("@")[0] if self.parent_model else None model, training_id = self.hook.create_auto_ml_tabular_training_job( project_id=self.project_id, region=self.region, @@ -360,6 +386,10 @@ def execute(self, context: Context): project=self.project_id, credentials=credentials, ), + parent_model=self.parent_model, + is_default_version=self.is_default_version, + model_version_aliases=self.model_version_aliases, + model_version_description=self.model_version_description, target_column=self.target_column, optimization_prediction_type=self.optimization_prediction_type, optimization_objective=self.optimization_objective, @@ -393,6 +423,7 @@ def execute(self, context: Context): if model: result = Model.to_dict(model) model_id = self.hook.extract_model_id(result) + self.xcom_push(context, key="model_id", value=model_id) VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id) else: result = model # type: ignore @@ -405,6 +436,7 @@ class CreateAutoMLTextTrainingJobOperator(AutoMLTrainingJobBaseOperator): """Create Auto ML Text Training job.""" template_fields = [ + "parent_model", "dataset_id", "region", "impersonation_chain", @@ -439,6 +471,7 @@ def execute(self, context: Context): gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, ) + self.parent_model = self.parent_model.rpartition("@")[0] if self.parent_model else None model, training_id = self.hook.create_auto_ml_text_training_job( project_id=self.project_id, region=self.region, @@ -459,11 +492,16 @@ def execute(self, context: Context): model_display_name=self.model_display_name, model_labels=self.model_labels, sync=self.sync, + parent_model=self.parent_model, + is_default_version=self.is_default_version, + model_version_aliases=self.model_version_aliases, + model_version_description=self.model_version_description, ) if model: result = Model.to_dict(model) model_id = self.hook.extract_model_id(result) + self.xcom_push(context, key="model_id", value=model_id) VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id) else: result = model # type: ignore @@ -476,6 +514,7 @@ class CreateAutoMLVideoTrainingJobOperator(AutoMLTrainingJobBaseOperator): """Create Auto ML Video Training job.""" template_fields = ( + "parent_model", "dataset_id", "region", "impersonation_chain", @@ -504,6 +543,7 @@ def execute(self, context: Context): gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, ) + self.parent_model = self.parent_model.rpartition("@")[0] if self.parent_model else None model, training_id = self.hook.create_auto_ml_video_training_job( project_id=self.project_id, region=self.region, @@ -521,11 +561,16 @@ def execute(self, context: Context): model_display_name=self.model_display_name, model_labels=self.model_labels, sync=self.sync, + parent_model=self.parent_model, + is_default_version=self.is_default_version, + model_version_aliases=self.model_version_aliases, + model_version_description=self.model_version_description, ) if model: result = Model.to_dict(model) model_id = self.hook.extract_model_id(result) + self.xcom_push(context, key="model_id", value=model_id) VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id) else: result = model # type: ignore diff --git a/airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py b/airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py index f6d9229dfe8e3..7daf0ec5ab3cd 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py +++ b/airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py @@ -15,15 +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.""" - jsonl - codepoints - aiplatform - gapic -""" from __future__ import annotations from typing import TYPE_CHECKING, Sequence @@ -54,7 +48,7 @@ class CreateBatchPredictionJobOperator(GoogleCloudBaseOperator): :param region: Required. The ID of the Google Cloud region that the service belongs to. :param batch_prediction_job: Required. The BatchPredictionJob to create. :param job_display_name: Required. The user-defined name of the BatchPredictionJob. The name can be - up to 128 characters long and can be consist of any UTF-8 characters. + up to 128 characters long and can consist of any UTF-8 characters. :param model_name: Required. A fully-qualified model resource name or model ID. :param instances_format: Required. The format in which instances are provided. Must be one of the formats listed in `Model.supported_input_storage_formats`. Default is "jsonl" when using diff --git a/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py b/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py index 87611c10a6580..0b3d221675e07 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +++ b/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py @@ -61,6 +61,10 @@ def __init__( model_instance_schema_uri: str | None = None, model_parameters_schema_uri: str | None = None, model_prediction_schema_uri: str | None = None, + parent_model: str | None = None, + is_default_version: bool | None = None, + model_version_aliases: list[str] | None = None, + model_version_description: str | None = None, labels: dict[str, str] | None = None, training_encryption_spec_key_name: str | None = None, model_encryption_spec_key_name: str | None = None, @@ -114,6 +118,10 @@ def __init__( self.model_parameters_schema_uri = model_parameters_schema_uri self.model_prediction_schema_uri = model_prediction_schema_uri self.labels = labels + self.parent_model = parent_model + self.is_default_version = is_default_version + self.model_version_aliases = model_version_aliases + self.model_version_description = model_version_description self.training_encryption_spec_key_name = training_encryption_spec_key_name self.model_encryption_spec_key_name = model_encryption_spec_key_name self.staging_bucket = staging_bucket @@ -192,48 +200,66 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator): the network. :param model_description: The description of the Model. :param model_instance_schema_uri: Optional. Points to a YAML file stored on Google Cloud - Storage describing the format of a single instance, which - are used in - ``PredictRequest.instances``, - ``ExplainRequest.instances`` - and - ``BatchPredictionJob.input_config``. - The schema is defined as an OpenAPI 3.0.2 `Schema - Object `__. - AutoML Models always have this field populated by AI - Platform. Note: The URI given on output will be immutable - and probably different, including the URI scheme, than the - one given on input. The output URI will point to a location - where the user only has a read access. + Storage describing the format of a single instance, which + are used in + ``PredictRequest.instances``, + ``ExplainRequest.instances`` + and + ``BatchPredictionJob.input_config``. + The schema is defined as an OpenAPI 3.0.2 `Schema + Object `__. + AutoML Models always have this field populated by AI + Platform. Note: The URI given on output will be immutable + and probably different, including the URI scheme, than the + one given on input. The output URI will point to a location + where the user only has a read access. :param model_parameters_schema_uri: Optional. Points to a YAML file stored on Google Cloud - Storage describing the parameters of prediction and - explanation via - ``PredictRequest.parameters``, - ``ExplainRequest.parameters`` - and - ``BatchPredictionJob.model_parameters``. - The schema is defined as an OpenAPI 3.0.2 `Schema - Object `__. - AutoML Models always have this field populated by AI - Platform, if no parameters are supported it is set to an - empty string. Note: The URI given on output will be - immutable and probably different, including the URI scheme, - than the one given on input. The output URI will point to a - location where the user only has a read access. + Storage describing the parameters of prediction and + explanation via + ``PredictRequest.parameters``, + ``ExplainRequest.parameters`` + and + ``BatchPredictionJob.model_parameters``. + The schema is defined as an OpenAPI 3.0.2 `Schema + Object `__. + AutoML Models always have this field populated by AI + Platform, if no parameters are supported it is set to an + empty string. Note: The URI given on output will be + immutable and probably different, including the URI scheme, + than the one given on input. The output URI will point to a + location where the user only has a read access. :param model_prediction_schema_uri: Optional. Points to a YAML file stored on Google Cloud - Storage describing the format of a single prediction - produced by this Model, which are returned via - ``PredictResponse.predictions``, - ``ExplainResponse.explanations``, - and - ``BatchPredictionJob.output_config``. - The schema is defined as an OpenAPI 3.0.2 `Schema - Object `__. - AutoML Models always have this field populated by AI - Platform. Note: The URI given on output will be immutable - and probably different, including the URI scheme, than the - one given on input. The output URI will point to a location - where the user only has a read access. + Storage describing the format of a single prediction + produced by this Model, which are returned via + ``PredictResponse.predictions``, + ``ExplainResponse.explanations``, + and + ``BatchPredictionJob.output_config``. + The schema is defined as an OpenAPI 3.0.2 `Schema + Object `__. + AutoML Models always have this field populated by AI + Platform. Note: The URI given on output will be immutable + and probably different, including the URI scheme, than the + one given on input. The output URI will point to a location + where the user only has a read access. + :param parent_model: Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + Only set this field when training a new version of an existing model. + :param is_default_version: Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + New model uploads, i.e. version 1, will always be "default" aliased. + :param model_version_aliases: Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + :param model_version_description: Optional. The description of the model version + being uploaded by this job. :param project_id: Project to run training in. :param region: Location to run training in. :param labels: Optional. The labels with user-defined metadata to @@ -409,6 +435,7 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator): template_fields = ( "region", "command", + "parent_model", "dataset_id", "impersonation_chain", ) @@ -428,6 +455,8 @@ def execute(self, context: Context): gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, ) + self.parent_model = self.parent_model.rpartition("@")[0] if self.parent_model else None + model, training_id, custom_job_id = self.hook.create_custom_container_training_job( project_id=self.project_id, region=self.region, @@ -445,6 +474,10 @@ def execute(self, context: Context): model_instance_schema_uri=self.model_instance_schema_uri, model_parameters_schema_uri=self.model_parameters_schema_uri, model_prediction_schema_uri=self.model_prediction_schema_uri, + parent_model=self.parent_model, + is_default_version=self.is_default_version, + model_version_aliases=self.model_version_aliases, + model_version_description=self.model_version_description, labels=self.labels, training_encryption_spec_key_name=self.training_encryption_spec_key_name, model_encryption_spec_key_name=self.model_encryption_spec_key_name, @@ -481,6 +514,7 @@ def execute(self, context: Context): if model: result = Model.to_dict(model) model_id = self.hook.extract_model_id(result) + self.xcom_push(context, key="model_id", value=model_id) VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id) else: result = model # type: ignore @@ -537,78 +571,96 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator the network. :param model_description: The description of the Model. :param model_instance_schema_uri: Optional. Points to a YAML file stored on Google Cloud - Storage describing the format of a single instance, which - are used in - ``PredictRequest.instances``, - ``ExplainRequest.instances`` - and - ``BatchPredictionJob.input_config``. - The schema is defined as an OpenAPI 3.0.2 `Schema - Object `__. - AutoML Models always have this field populated by AI - Platform. Note: The URI given on output will be immutable - and probably different, including the URI scheme, than the - one given on input. The output URI will point to a location - where the user only has a read access. + Storage describing the format of a single instance, which + are used in + ``PredictRequest.instances``, + ``ExplainRequest.instances`` + and + ``BatchPredictionJob.input_config``. + The schema is defined as an OpenAPI 3.0.2 `Schema + Object `__. + AutoML Models always have this field populated by AI + Platform. Note: The URI given on output will be immutable + and probably different, including the URI scheme, than the + one given on input. The output URI will point to a location + where the user only has a read access. :param model_parameters_schema_uri: Optional. Points to a YAML file stored on Google Cloud - Storage describing the parameters of prediction and - explanation via - ``PredictRequest.parameters``, - ``ExplainRequest.parameters`` - and - ``BatchPredictionJob.model_parameters``. - The schema is defined as an OpenAPI 3.0.2 `Schema - Object `__. - AutoML Models always have this field populated by AI - Platform, if no parameters are supported it is set to an - empty string. Note: The URI given on output will be - immutable and probably different, including the URI scheme, - than the one given on input. The output URI will point to a - location where the user only has a read access. + Storage describing the parameters of prediction and + explanation via + ``PredictRequest.parameters``, + ``ExplainRequest.parameters`` + and + ``BatchPredictionJob.model_parameters``. + The schema is defined as an OpenAPI 3.0.2 `Schema + Object `__. + AutoML Models always have this field populated by AI + Platform, if no parameters are supported it is set to an + empty string. Note: The URI given on output will be + immutable and probably different, including the URI scheme, + than the one given on input. The output URI will point to a + location where the user only has a read access. :param model_prediction_schema_uri: Optional. Points to a YAML file stored on Google Cloud - Storage describing the format of a single prediction - produced by this Model, which are returned via - ``PredictResponse.predictions``, - ``ExplainResponse.explanations``, - and - ``BatchPredictionJob.output_config``. - The schema is defined as an OpenAPI 3.0.2 `Schema - Object `__. - AutoML Models always have this field populated by AI - Platform. Note: The URI given on output will be immutable - and probably different, including the URI scheme, than the - one given on input. The output URI will point to a location - where the user only has a read access. + Storage describing the format of a single prediction + produced by this Model, which are returned via + ``PredictResponse.predictions``, + ``ExplainResponse.explanations``, + and + ``BatchPredictionJob.output_config``. + The schema is defined as an OpenAPI 3.0.2 `Schema + Object `__. + AutoML Models always have this field populated by AI + Platform. Note: The URI given on output will be immutable + and probably different, including the URI scheme, than the + one given on input. The output URI will point to a location + where the user only has a read access. + :param parent_model: Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + Only set this field when training a new version of an existing model. + :param is_default_version: Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + New model uploads, i.e. version 1, will always be "default" aliased. + :param model_version_aliases: Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + :param model_version_description: Optional. The description of the model version + being uploaded by this job. :param project_id: Project to run training in. :param region: Location to run training in. :param labels: Optional. The labels with user-defined metadata to - organize TrainingPipelines. - Label keys and values can be no longer than 64 - characters, can only - contain lowercase letters, numeric characters, - underscores and dashes. International characters - are allowed. - See https://goo.gl/xmQnxf for more information - and examples of labels. + organize TrainingPipelines. + Label keys and values can be no longer than 64 + characters, can only + contain lowercase letters, numeric characters, + underscores and dashes. International characters + are allowed. + See https://goo.gl/xmQnxf for more information + and examples of labels. :param training_encryption_spec_key_name: Optional. The Cloud KMS resource identifier of the customer - managed encryption key used to protect the training pipeline. Has the - form: - ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``. - The key needs to be in the same region as where the compute - resource is created. + managed encryption key used to protect the training pipeline. Has the + form: + ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``. + The key needs to be in the same region as where the compute + resource is created. - If set, this TrainingPipeline will be secured by this key. + If set, this TrainingPipeline will be secured by this key. - Note: Model trained by this TrainingPipeline is also secured - by this key if ``model_to_upload`` is not set separately. + Note: Model trained by this TrainingPipeline is also secured + by this key if ``model_to_upload`` is not set separately. :param model_encryption_spec_key_name: Optional. The Cloud KMS resource identifier of the customer - managed encryption key used to protect the model. Has the - form: - ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``. - The key needs to be in the same region as where the compute - resource is created. + managed encryption key used to protect the model. Has the + form: + ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``. + The key needs to be in the same region as where the compute + resource is created. - If set, the trained Model will be secured by this key. + If set, the trained Model will be secured by this key. :param staging_bucket: Bucket used to stage source and training artifacts. :param dataset: Vertex AI to fit this training against. :param annotation_schema_uri: Google Cloud Storage URI points to a YAML file describing @@ -628,19 +680,19 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator and ``annotation_schema_uri``. :param model_display_name: If the script produces a managed Vertex AI Model. The display name of - the Model. The name can be up to 128 characters long and can be consist - of any UTF-8 characters. + the Model. The name can be up to 128 characters long and can be consist + of any UTF-8 characters. - If not provided upon creation, the job's display_name is used. + If not provided upon creation, the job's display_name is used. :param model_labels: Optional. The labels with user-defined metadata to - organize your Models. - Label keys and values can be no longer than 64 - characters, can only - contain lowercase letters, numeric characters, - underscores and dashes. International characters - are allowed. - See https://goo.gl/xmQnxf for more information - and examples of labels. + organize your Models. + Label keys and values can be no longer than 64 + characters, can only + contain lowercase letters, numeric characters, + underscores and dashes. International characters + are allowed. + See https://goo.gl/xmQnxf for more information + and examples of labels. :param base_output_dir: GCS output directory of job. If not provided a timestamped directory in the staging directory will be used. @@ -653,38 +705,38 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator - AIP_TENSORBOARD_LOG_DIR: a Cloud Storage URI of a directory intended for saving TensorBoard logs, i.e. /logs/ :param service_account: Specifies the service account for workload run-as account. - Users submitting jobs must have act-as permission on this run-as account. + Users submitting jobs must have act-as permission on this run-as account. :param network: The full name of the Compute Engine network to which the job - should be peered. - Private services access must already be configured for the network. - If left unspecified, the job is not peered with any network. + should be peered. + Private services access must already be configured for the network. + If left unspecified, the job is not peered with any network. :param bigquery_destination: Provide this field if `dataset` is a BiqQuery dataset. - The BigQuery project location where the training data is to - be written to. In the given project a new dataset is created - with name - ``dataset___`` - where timestamp is in YYYY_MM_DDThh_mm_ss_sssZ format. All - training input data will be written into that dataset. In - the dataset three tables will be created, ``training``, - ``validation`` and ``test``. - - - AIP_DATA_FORMAT = "bigquery". - - AIP_TRAINING_DATA_URI ="bigquery_destination.dataset_*.training" - - AIP_VALIDATION_DATA_URI = "bigquery_destination.dataset_*.validation" - - AIP_TEST_DATA_URI = "bigquery_destination.dataset_*.test" + The BigQuery project location where the training data is to + be written to. In the given project a new dataset is created + with name + ``dataset___`` + where timestamp is in YYYY_MM_DDThh_mm_ss_sssZ format. All + training input data will be written into that dataset. In + the dataset three tables will be created, ``training``, + ``validation`` and ``test``. + + - AIP_DATA_FORMAT = "bigquery". + - AIP_TRAINING_DATA_URI ="bigquery_destination.dataset_*.training" + - AIP_VALIDATION_DATA_URI = "bigquery_destination.dataset_*.validation" + - AIP_TEST_DATA_URI = "bigquery_destination.dataset_*.test" :param args: Command line arguments to be passed to the Python script. :param environment_variables: Environment variables to be passed to the container. - Should be a dictionary where keys are environment variable names - and values are environment variable values for those names. - At most 10 environment variables can be specified. - The Name of the environment variable must be unique. + Should be a dictionary where keys are environment variable names + and values are environment variable values for those names. + At most 10 environment variables can be specified. + The Name of the environment variable must be unique. :param replica_count: The number of worker replicas. If replica count = 1 then one chief - replica will be provisioned. If replica_count > 1 the remainder will be - provisioned as a worker replica pool. + replica will be provisioned. If replica_count > 1 the remainder will be + provisioned as a worker replica pool. :param machine_type: The type of machine to use for training. :param accelerator_type: Hardware accelerator type. One of ACCELERATOR_TYPE_UNSPECIFIED, - NVIDIA_TESLA_K80, NVIDIA_TESLA_P100, NVIDIA_TESLA_V100, NVIDIA_TESLA_P4, - NVIDIA_TESLA_T4 + NVIDIA_TESLA_K80, NVIDIA_TESLA_P100, NVIDIA_TESLA_V100, NVIDIA_TESLA_P4, + NVIDIA_TESLA_T4 :param accelerator_count: The number of accelerators to attach to a worker replica. :param boot_disk_type: Type of the boot disk, default is `pd-ssd`. Valid values: `pd-ssd` (Persistent Disk Solid State Drive) or @@ -752,6 +804,7 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator """ template_fields = ( + "parent_model", "region", "dataset_id", "impersonation_chain", @@ -774,6 +827,7 @@ def execute(self, context: Context): gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, ) + self.parent_model = self.parent_model.rpartition("@")[0] if self.parent_model else None model, training_id, custom_job_id = self.hook.create_custom_python_package_training_job( project_id=self.project_id, region=self.region, @@ -792,6 +846,10 @@ def execute(self, context: Context): model_instance_schema_uri=self.model_instance_schema_uri, model_parameters_schema_uri=self.model_parameters_schema_uri, model_prediction_schema_uri=self.model_prediction_schema_uri, + parent_model=self.parent_model, + is_default_version=self.is_default_version, + model_version_aliases=self.model_version_aliases, + model_version_description=self.model_version_description, labels=self.labels, training_encryption_spec_key_name=self.training_encryption_spec_key_name, model_encryption_spec_key_name=self.model_encryption_spec_key_name, @@ -828,6 +886,7 @@ def execute(self, context: Context): if model: result = Model.to_dict(model) model_id = self.hook.extract_model_id(result) + self.xcom_push(context, key="model_id", value=model_id) VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id) else: result = model # type: ignore @@ -884,78 +943,96 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator): the network. :param model_description: The description of the Model. :param model_instance_schema_uri: Optional. Points to a YAML file stored on Google Cloud - Storage describing the format of a single instance, which - are used in - ``PredictRequest.instances``, - ``ExplainRequest.instances`` - and - ``BatchPredictionJob.input_config``. - The schema is defined as an OpenAPI 3.0.2 `Schema - Object `__. - AutoML Models always have this field populated by AI - Platform. Note: The URI given on output will be immutable - and probably different, including the URI scheme, than the - one given on input. The output URI will point to a location - where the user only has a read access. + Storage describing the format of a single instance, which + are used in + ``PredictRequest.instances``, + ``ExplainRequest.instances`` + and + ``BatchPredictionJob.input_config``. + The schema is defined as an OpenAPI 3.0.2 `Schema + Object `__. + AutoML Models always have this field populated by AI + Platform. Note: The URI given on output will be immutable + and probably different, including the URI scheme, than the + one given on input. The output URI will point to a location + where the user only has a read access. :param model_parameters_schema_uri: Optional. Points to a YAML file stored on Google Cloud - Storage describing the parameters of prediction and - explanation via - ``PredictRequest.parameters``, - ``ExplainRequest.parameters`` - and - ``BatchPredictionJob.model_parameters``. - The schema is defined as an OpenAPI 3.0.2 `Schema - Object `__. - AutoML Models always have this field populated by AI - Platform, if no parameters are supported it is set to an - empty string. Note: The URI given on output will be - immutable and probably different, including the URI scheme, - than the one given on input. The output URI will point to a - location where the user only has a read access. + Storage describing the parameters of prediction and + explanation via + ``PredictRequest.parameters``, + ``ExplainRequest.parameters`` + and + ``BatchPredictionJob.model_parameters``. + The schema is defined as an OpenAPI 3.0.2 `Schema + Object `__. + AutoML Models always have this field populated by AI + Platform, if no parameters are supported it is set to an + empty string. Note: The URI given on output will be + immutable and probably different, including the URI scheme, + than the one given on input. The output URI will point to a + location where the user only has a read access. :param model_prediction_schema_uri: Optional. Points to a YAML file stored on Google Cloud - Storage describing the format of a single prediction - produced by this Model, which are returned via - ``PredictResponse.predictions``, - ``ExplainResponse.explanations``, - and - ``BatchPredictionJob.output_config``. - The schema is defined as an OpenAPI 3.0.2 `Schema - Object `__. - AutoML Models always have this field populated by AI - Platform. Note: The URI given on output will be immutable - and probably different, including the URI scheme, than the - one given on input. The output URI will point to a location - where the user only has a read access. + Storage describing the format of a single prediction + produced by this Model, which are returned via + ``PredictResponse.predictions``, + ``ExplainResponse.explanations``, + and + ``BatchPredictionJob.output_config``. + The schema is defined as an OpenAPI 3.0.2 `Schema + Object `__. + AutoML Models always have this field populated by AI + Platform. Note: The URI given on output will be immutable + and probably different, including the URI scheme, than the + one given on input. The output URI will point to a location + where the user only has a read access. + :param parent_model: Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + Only set this field when training a new version of an existing model. + :param is_default_version: Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + New model uploads, i.e. version 1, will always be "default" aliased. + :param model_version_aliases: Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + :param model_version_description: Optional. The description of the model version + being uploaded by this job. :param project_id: Project to run training in. :param region: Location to run training in. :param labels: Optional. The labels with user-defined metadata to - organize TrainingPipelines. - Label keys and values can be no longer than 64 - characters, can only - contain lowercase letters, numeric characters, - underscores and dashes. International characters - are allowed. - See https://goo.gl/xmQnxf for more information - and examples of labels. + organize TrainingPipelines. + Label keys and values can be no longer than 64 + characters, can only + contain lowercase letters, numeric characters, + underscores and dashes. International characters + are allowed. + See https://goo.gl/xmQnxf for more information + and examples of labels. :param training_encryption_spec_key_name: Optional. The Cloud KMS resource identifier of the customer - managed encryption key used to protect the training pipeline. Has the - form: - ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``. - The key needs to be in the same region as where the compute - resource is created. + managed encryption key used to protect the training pipeline. Has the + form: + ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``. + The key needs to be in the same region as where the compute + resource is created. - If set, this TrainingPipeline will be secured by this key. + If set, this TrainingPipeline will be secured by this key. - Note: Model trained by this TrainingPipeline is also secured - by this key if ``model_to_upload`` is not set separately. + Note: Model trained by this TrainingPipeline is also secured + by this key if ``model_to_upload`` is not set separately. :param model_encryption_spec_key_name: Optional. The Cloud KMS resource identifier of the customer - managed encryption key used to protect the model. Has the - form: - ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``. - The key needs to be in the same region as where the compute - resource is created. + managed encryption key used to protect the model. Has the + form: + ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``. + The key needs to be in the same region as where the compute + resource is created. - If set, the trained Model will be secured by this key. + If set, the trained Model will be secured by this key. :param staging_bucket: Bucket used to stage source and training artifacts. :param dataset: Vertex AI to fit this training against. :param annotation_schema_uri: Google Cloud Storage URI points to a YAML file describing @@ -1101,6 +1178,7 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator): template_fields = ( "region", "script_path", + "parent_model", "requirements", "dataset_id", "impersonation_chain", @@ -1123,6 +1201,8 @@ def execute(self, context: Context): gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, ) + self.parent_model = self.parent_model.rpartition("@")[0] if self.parent_model else None + model, training_id, custom_job_id = self.hook.create_custom_training_job( project_id=self.project_id, region=self.region, @@ -1141,6 +1221,10 @@ def execute(self, context: Context): model_instance_schema_uri=self.model_instance_schema_uri, model_parameters_schema_uri=self.model_parameters_schema_uri, model_prediction_schema_uri=self.model_prediction_schema_uri, + parent_model=self.parent_model, + is_default_version=self.is_default_version, + model_version_aliases=self.model_version_aliases, + model_version_description=self.model_version_description, labels=self.labels, training_encryption_spec_key_name=self.training_encryption_spec_key_name, model_encryption_spec_key_name=self.model_encryption_spec_key_name, @@ -1177,6 +1261,7 @@ def execute(self, context: Context): if model: result = Model.to_dict(model) model_id = self.hook.extract_model_id(result) + self.xcom_push(context, key="model_id", value=model_id) VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id) else: result = model # type: ignore @@ -1276,7 +1361,8 @@ def execute(self, context: Context): class ListCustomTrainingJobOperator(GoogleCloudBaseOperator): - """Lists CustomTrainingJob, CustomPythonTrainingJob, or CustomContainerTrainingJob in a Location. + """ + Lists CustomTrainingJob, CustomPythonTrainingJob, or CustomContainerTrainingJob 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. diff --git a/airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py b/airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py index 9f2ea4dac8101..3a53efbf4d7e9 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py +++ b/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 diff --git a/airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py b/airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py index 70cbddc701f53..060eeff249471 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py +++ b/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. diff --git a/airflow/providers/google/cloud/operators/vertex_ai/model_service.py b/airflow/providers/google/cloud/operators/vertex_ai/model_service.py index 379557fa8e6ff..c44da5ffd58c5 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/model_service.py +++ b/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 @@ def execute(self, context: Context): 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 @@ def execute(self, context: Context): 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 @@ def execute(self, context: Context): 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) diff --git a/airflow/providers/google/cloud/utils/mlengine_operator_utils.py b/airflow/providers/google/cloud/utils/mlengine_operator_utils.py index feffbd2f356f0..cc85705653e06 100644 --- a/airflow/providers/google/cloud/utils/mlengine_operator_utils.py +++ b/airflow/providers/google/cloud/utils/mlengine_operator_utils.py @@ -56,9 +56,15 @@ def create_evaluate_ops( dag: DAG | None = None, py_interpreter="python3", ) -> tuple[MLEngineStartBatchPredictionJobOperator, BeamRunPythonPipelineOperator, PythonOperator]: - """ + r""" Creates Operators needed for model evaluation and returns. + This function is deprecated. All the functionality of legacy MLEngine and new features are available + on the Vertex AI platform. + + To create and view Model Evaluation, please check the documentation: + https://cloud.google.com/vertex-ai/docs/evaluation/using-model-evaluation#create_an_evaluation. + It gets prediction over inputs via Cloud ML Engine BatchPrediction API by calling MLEngineBatchPredictionOperator, then summarize and validate the result via Cloud Dataflow using DataFlowPythonOperator. diff --git a/docs/apache-airflow-providers-google/operators/cloud/mlengine.rst b/docs/apache-airflow-providers-google/operators/cloud/mlengine.rst index 4e28fb54062ae..9c71e885d8d6a 100644 --- a/docs/apache-airflow-providers-google/operators/cloud/mlengine.rst +++ b/docs/apache-airflow-providers-google/operators/cloud/mlengine.rst @@ -26,6 +26,11 @@ in the cloud, and use models to make predictions for new data. AI Platform is a of tools for training, evaluating, and tuning machine learning models. AI Platform can also be used to deploy a trained model, make predictions, and manage various model versions. +The legacy versions of AI Platform Training, AI Platform Prediction, AI Platform Pipelines, +and AI Platform Data Labeling Service are deprecated and will no longer be available on +Google Cloud after their shutdown date. All the functionality of legacy AI Platform and new +features are available on the Vertex AI platform. + Prerequisite tasks ^^^^^^^^^^^^^^^^^^ @@ -40,19 +45,15 @@ This creates a virtual machine that can run code specified in the trainer file, contains the main application code. A job can be initiated with the :class:`~airflow.providers.google.cloud.operators.mlengine.MLEngineStartTrainingJobOperator`. -.. exampleinclude:: /../../tests/system/providers/google/cloud/ml_engine/example_mlengine.py - :language: python - :dedent: 4 - :start-after: [START howto_operator_gcp_mlengine_training] - :end-before: [END howto_operator_gcp_mlengine_training] - -Also for all this action you can use operator in the deferrable mode: +This operator is deprecated. Please, use +:class:`~airflow.providers.google.cloud.operators.vertex_ai.custom_job.CreateCustomPythonPackageTrainingJobOperator` +instead. -.. exampleinclude:: /../../tests/system/providers/google/cloud/ml_engine/example_mlengine_async.py +.. exampleinclude:: /../../tests/system/providers/google/cloud/ml_engine/example_mlengine.py :language: python :dedent: 4 - :start-after: [START howto_operator_gcp_mlengine_training_async] - :end-before: [END howto_operator_gcp_mlengine_training_async] + :start-after: [START howto_operator_create_custom_python_training_job_v1] + :end-before: [END howto_operator_create_custom_python_training_job_v1] .. _howto/operator:MLEngineCreateModelOperator: @@ -63,11 +64,16 @@ A model is a container that can hold multiple model versions. A new model can be The ``model`` field should be defined with a dictionary containing the information about the model. ``name`` is a required field in this dictionary. +This operator is deprecated. The model is created as a result of running Vertex AI operators that create training jobs +of any types. For example, you can use +:class:`~airflow.providers.google.cloud.operators.vertex_ai.custom_job.CreateCustomPythonPackageTrainingJobOperator`. +The result of running this operator will be ready-to-use model saved in Model Registry. + .. exampleinclude:: /../../tests/system/providers/google/cloud/ml_engine/example_mlengine.py :language: python :dedent: 4 - :start-after: [START howto_operator_gcp_mlengine_create_model] - :end-before: [END howto_operator_gcp_mlengine_create_model] + :start-after: [START howto_operator_create_custom_python_training_job_v1] + :end-before: [END howto_operator_create_custom_python_training_job_v1] .. _howto/operator:MLEngineGetModelOperator: @@ -77,6 +83,10 @@ The :class:`~airflow.providers.google.cloud.operators.mlengine.MLEngineGetModelO can be used to obtain a model previously created. To obtain the correct model, ``model_name`` must be defined in the operator. +This operator is deprecated. Please, use +:class:`~airflow.providers.google.cloud.operators.vertex_ai.model_service.GetModelOperator` +instead. + .. exampleinclude:: /../../tests/system/providers/google/cloud/ml_engine/example_mlengine.py :language: python :dedent: 4 @@ -104,20 +114,26 @@ The model must be specified by ``model_name``, and the ``version`` parameter sho all the information about the version. Within the ``version`` parameter's dictionary, the ``name`` field is required. +This operator is deprecated. Please, use +:class:`~airflow.providers.google.cloud.operators.vertex_ai.custom_job.CreateCustomPythonPackageTrainingJobOperator` +instead. In this case, the new version of specific model could be created by specifying existing model id in +``parent_model`` parameter when running Training Job. This will ensure that new version of model will be trained except +of creating new model. + .. exampleinclude:: /../../tests/system/providers/google/cloud/ml_engine/example_mlengine.py :language: python :dedent: 4 - :start-after: [START howto_operator_gcp_mlengine_create_version1] - :end-before: [END howto_operator_gcp_mlengine_create_version1] + :start-after: [START howto_operator_create_custom_python_training_job_v1] + :end-before: [END howto_operator_create_custom_python_training_job_v1] -The :class:`~airflow.providers.google.cloud.operators.mlengine.MLEngineCreateVersionOperator` +The :class:`~airflow.providers.google.cloud.operators.vertex_ai.custom_job.CreateCustomPythonPackageTrainingJobOperator` can also be used to create more versions with varying parameters. .. exampleinclude:: /../../tests/system/providers/google/cloud/ml_engine/example_mlengine.py :language: python :dedent: 4 - :start-after: [START howto_operator_gcp_mlengine_create_version2] - :end-before: [END howto_operator_gcp_mlengine_create_version2] + :start-after: [START howto_operator_create_custom_python_training_job_v2] + :end-before: [END howto_operator_create_custom_python_training_job_v2] .. _howto/operator:MLEngineSetDefaultVersionOperator: .. _howto/operator:MLEngineListVersionsOperator: @@ -128,6 +144,13 @@ By default, the model code will run using the default model version. You can set :class:`~airflow.providers.google.cloud.operators.mlengine.MLEngineSetDefaultVersionOperator` by specifying the ``model_name`` and ``version_name`` parameters. +This operator is deprecated. Please, use +:class:`~airflow.providers.google.cloud.operators.vertex_ai.model_service.SetDefaultVersionOnModelOperator` +instead. The desired model version to be set as default could be passed with the model ID in ``model_id`` parameter +in format ``projects/{project}/locations/{location}/models/{model_id}@{version_id}`` or +``projects/{project}/locations/{location}/models/{model_id}@{version_alias}``. By default, the first model version +created will be marked as default. + .. exampleinclude:: /../../tests/system/providers/google/cloud/ml_engine/example_mlengine.py :language: python :dedent: 4 @@ -138,23 +161,17 @@ To list the model versions available, use the :class:`~airflow.providers.google.cloud.operators.mlengine.MLEngineListVersionsOperator` while specifying the ``model_name`` parameter. +This operator is deprecated. Please, use +:class:`~airflow.providers.google.cloud.operators.vertex_ai.model_service.ListModelVersionsOperator` +instead. You can pass the name of the desired model in ``model_id`` parameter. If the model ID is passed +with version aliases, the operator will output all the versions available for this model. + .. exampleinclude:: /../../tests/system/providers/google/cloud/ml_engine/example_mlengine.py :language: python :dedent: 4 :start-after: [START howto_operator_gcp_mlengine_list_versions] :end-before: [END howto_operator_gcp_mlengine_list_versions] -You can use :ref:`Jinja templating ` with the ``project_id`` and ``model`` -fields to dynamically determine their values. The result are saved to :ref:`XCom `, -allowing them to be used by other operators. In this case, the -:class:`~airflow.operators.bash.BashOperator` is used to print the version information. - -.. exampleinclude:: /../../tests/system/providers/google/cloud/ml_engine/example_mlengine.py - :language: python - :dedent: 4 - :start-after: [START howto_operator_gcp_mlengine_print_versions] - :end-before: [END howto_operator_gcp_mlengine_print_versions] - .. _howto/operator:MLEngineStartBatchPredictionJobOperator: Making predictions @@ -164,11 +181,15 @@ A Google Cloud AI Platform prediction job can be started with the For specifying the model origin, you need to provide either the ``model_name``, ``uri``, or ``model_name`` and ``version_name``. If you do not provide the ``version_name``, the operator will use the default model version. +This operator is deprecated. Please, use +:class:`~airflow.providers.google.cloud.operators.vertex_ai.batch_prediction_job.CreateBatchPredictionJobOperator` +instead. + .. exampleinclude:: /../../tests/system/providers/google/cloud/ml_engine/example_mlengine.py :language: python :dedent: 4 - :start-after: [START howto_operator_gcp_mlengine_get_prediction] - :end-before: [END howto_operator_gcp_mlengine_get_prediction] + :start-after: [START howto_operator_start_batch_prediction] + :end-before: [END howto_operator_start_batch_prediction] .. _howto/operator:MLEngineDeleteVersionOperator: .. _howto/operator:MLEngineDeleteModelOperator: @@ -179,6 +200,10 @@ A model version can be deleted with the :class:`~airflow.providers.google.cloud.operators.mlengine.MLEngineDeleteVersionOperator` by the ``version_name`` and ``model_name`` parameters. +This operator is deprecated. Please, use +:class:`~airflow.providers.google.cloud.operators.vertex_ai.model_service.DeleteModelVersionOperator` +instead. The default version could not be deleted on the model. + .. exampleinclude:: /../../tests/system/providers/google/cloud/ml_engine/example_mlengine.py :language: python :dedent: 4 @@ -189,6 +214,10 @@ You can also delete a model with the :class:`~airflow.providers.google.cloud.operators.mlengine.MLEngineDeleteModelOperator` by providing the ``model_name`` parameter. +This operator is deprecated. Please, use +:class:`~airflow.providers.google.cloud.operators.vertex_ai.model_service.DeleteModelOperator` +instead. + .. exampleinclude:: /../../tests/system/providers/google/cloud/ml_engine/example_mlengine.py :language: python :dedent: 4 @@ -197,36 +226,12 @@ by providing the ``model_name`` parameter. Evaluating a model ^^^^^^^^^^^^^^^^^^ -To evaluate a prediction and model, specify a metric function to generate a summary and customize -the evaluation of the model. This function receives a dictionary derived from a json in the batch -prediction result, then returns a tuple of metrics. - -.. exampleinclude:: /../../tests/system/providers/google/cloud/ml_engine/example_mlengine.py - :language: python - :dedent: 4 - :start-after: [START howto_operator_gcp_mlengine_get_metric] - :end-before: [END howto_operator_gcp_mlengine_get_metric] - -To evaluate a prediction and model, it's useful to have a function to validate the summary result. -This function receives a dictionary of the averaged metrics the function above generated. It then -raises an exception if a task fails or should not proceed. -.. exampleinclude:: /../../tests/system/providers/google/cloud/ml_engine/example_mlengine.py - :language: python - :dedent: 4 - :start-after: [START howto_operator_gcp_mlengine_validate_error] - :end-before: [END howto_operator_gcp_mlengine_validate_error] - -Prediction results and a model summary can be generated through a function such as -:class:`~airflow.providers.google.cloud.utils.mlengine_operator_utils.create_evaluate_ops`. -It makes predictions using the specified inputs and then summarizes and validates the result. The -functions created above should be passed in through the ``metric_fn_and_keys`` and ``validate_fn`` fields. - -.. exampleinclude:: /../../tests/system/providers/google/cloud/ml_engine/example_mlengine.py - :language: python - :dedent: 4 - :start-after: [START howto_operator_gcp_mlengine_evaluate] - :end-before: [END howto_operator_gcp_mlengine_evaluate] +This function is deprecated. All the functionality of legacy MLEngine and new features are available on +the Vertex AI platform. +To create and view Model Evaluation, please check the documentation: +`Evaluate models using Vertex AI +`__ Reference ^^^^^^^^^ diff --git a/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst b/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst index 0ab36f0ab920c..f9e426a032cbb 100644 --- a/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst +++ b/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst @@ -146,6 +146,17 @@ For this Job you should put path to your local training script inside ``script_p :start-after: [START how_to_cloud_vertex_ai_create_custom_training_job_operator] :end-before: [END how_to_cloud_vertex_ai_create_custom_training_job_operator] +Additionally, you can create new version of existing Training Job instead. In this case, the result will be new +version of existing Model instead of new Model created in Model Registry. This can be done by specifying +``parent_model`` parameter when running Training Job. + +.. exampleinclude:: /../../tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_job.py + :language: python + :dedent: 4 + :start-after: [START how_to_cloud_vertex_ai_create_custom_training_job_v2_operator] + :end-before: [END how_to_cloud_vertex_ai_create_custom_training_job_v2_operator] + + You can get a list of Training Jobs using :class:`~airflow.providers.google.cloud.operators.vertex_ai.custom_job.ListCustomTrainingJobOperator`. @@ -236,6 +247,16 @@ put dataset id to ``dataset_id`` parameter in operator. :start-after: [START how_to_cloud_vertex_ai_create_auto_ml_video_training_job_operator] :end-before: [END how_to_cloud_vertex_ai_create_auto_ml_video_training_job_operator] +Additionally, you can create new version of existing AutoML Video Training Job. In this case, the result will be new +version of existing Model instead of new Model created in Model Registry. This can be done by specifying +``parent_model`` parameter when running AutoML Video Training Job. + +.. exampleinclude:: /../../tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_auto_ml_video_training.py + :language: python + :dedent: 4 + :start-after: [START how_to_cloud_vertex_ai_create_auto_ml_video_training_job_v2_operator] + :end-before: [END how_to_cloud_vertex_ai_create_auto_ml_video_training_job_v2_operator] + You can get a list of AutoML Training Jobs using :class:`~airflow.providers.google.cloud.operators.vertex_ai.auto_ml.ListAutoMLTrainingJobOperator`. @@ -414,6 +435,60 @@ To get a model list you can use :start-after: [START how_to_cloud_vertex_ai_list_models_operator] :end-before: [END how_to_cloud_vertex_ai_list_models_operator] +To retrieve model by its ID you can use +:class:`~airflow.providers.google.cloud.operators.vertex_ai.model_service.GetModelOperator`. + +.. exampleinclude:: /../../tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_model_service.py + :language: python + :dedent: 4 + :start-after: [START how_to_cloud_vertex_ai_get_model_operator] + :end-before: [END how_to_cloud_vertex_ai_get_model_operator] + +To list all model versions you can use +:class:`~airflow.providers.google.cloud.operators.vertex_ai.model_service.ListModelVersionsOperator`. + +.. exampleinclude:: /../../tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_model_service.py + :language: python + :dedent: 4 + :start-after: [START how_to_cloud_vertex_ai_list_model_versions_operator] + :end-before: [END how_to_cloud_vertex_ai_list_model_versions_operator] + +To set a specific version of model as a default one you can use +:class:`~airflow.providers.google.cloud.operators.vertex_ai.model_service.SetDefaultVersionOnModelOperator`. + +.. exampleinclude:: /../../tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_model_service.py + :language: python + :dedent: 4 + :start-after: [START how_to_cloud_vertex_ai_set_version_as_default_operator] + :end-before: [END how_to_cloud_vertex_ai_set_version_as_default_operator] + +To add aliases to specific version of model you can use +:class:`~airflow.providers.google.cloud.operators.vertex_ai.model_service.AddVersionAliasesOnModelOperator`. + +.. exampleinclude:: /../../tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_model_service.py + :language: python + :dedent: 4 + :start-after: [START how_to_cloud_vertex_ai_add_version_aliases_operator] + :end-before: [END how_to_cloud_vertex_ai_add_version_aliases_operator] + +To delete aliases from specific version of model you can use +:class:`~airflow.providers.google.cloud.operators.vertex_ai.model_service.DeleteVersionAliasesOnModelOperator`. + +.. exampleinclude:: /../../tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_model_service.py + :language: python + :dedent: 4 + :start-after: [START how_to_cloud_vertex_ai_delete_version_aliases_operator] + :end-before: [END how_to_cloud_vertex_ai_delete_version_aliases_operator] + +To delete specific version of model you can use +:class:`~airflow.providers.google.cloud.operators.vertex_ai.model_service.DeleteModelVersionOperator`. + +.. exampleinclude:: /../../tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_model_service.py + :language: python + :dedent: 4 + :start-after: [START how_to_cloud_vertex_ai_delete_version_operator] + :end-before: [END how_to_cloud_vertex_ai_delete_version_operator] + Reference ^^^^^^^^^ diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index b37bdf3c2277a..a71bd7b790c0b 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -22,6 +22,7 @@ AgentKey aio aiobotocore AioSession +aiplatform Airbnb airbnb Airbyte @@ -95,6 +96,7 @@ Atlassian atlassian attr attrs +au Auth auth authenticator @@ -195,6 +197,7 @@ burstable bytestring cacert callables +camelCase Cancelled cancelled carbonite @@ -263,6 +266,7 @@ cnt codebase Codecov codecov +codepoints Colour colour colours @@ -581,6 +585,7 @@ Fernet fernet fetchmany fetchone +FieldMask Filebeat filehandle fileloc @@ -611,6 +616,7 @@ fmt fn fo followsa +forecasted formatter formatters Formaturas @@ -626,6 +632,7 @@ Fundera ga Gantt gantt +gapic gbq gcc gcloud @@ -790,6 +797,7 @@ Investorise io ip iPython +irreproducible IRSA isfile ish @@ -811,6 +819,7 @@ jdbc jdk jenkins JenkinsRequest +Jetson Jiajie Jinja jinja @@ -834,6 +843,7 @@ JPype js Json json +jsonl jthomas Jupyter jupyter @@ -924,6 +934,7 @@ Lyft machineTypes macOS macosx +mae mailto makedirs makedsn @@ -959,6 +970,7 @@ microsoft middleware Midnights midnights +milli milton minicluster minikube @@ -991,6 +1003,7 @@ mysql mysqlclient mysqldb mytaxi +myVPC Namenode namenode Namespace @@ -1062,6 +1075,7 @@ Opsgenie opsgenie Optimise optimise +optimizationObjective optionality ora oracledb @@ -1142,6 +1156,7 @@ postMessage Potiuk powershell pql +prc Pre pre precheck @@ -1220,6 +1235,9 @@ pywinrm qds Qingping Qplum +quantile +Quantiles +quantiles Quantopian Qubole qubole @@ -1289,8 +1307,12 @@ retransmits rfc ricard rideable +rmse +rmsle +rmspe Roadmap Robinhood +roc RoleBinding romeoandjuliet rootcss @@ -1348,6 +1370,7 @@ seekable segmentGranularity Sendgrid sendgrid +sentimentMax serde serialise serializable @@ -1439,6 +1462,7 @@ str StrictUndefined Stringified stringified +Struct subchart subclassed Subclasses @@ -1501,6 +1525,7 @@ tableId Tabular tagKey tagValue +targetColumn task_group TaskDecorator Taskfail @@ -1529,6 +1554,7 @@ templating templatize templatized tenantId +Tensorboard tensorflow Teradata Terraform @@ -1562,6 +1588,7 @@ tooltip tooltips traceback tracebacks +TrainingPipeline travis triage triaging @@ -1591,10 +1618,14 @@ umask Un un unarchived +unassigns uncommenting Undead undead Undeads +undeploy +undeployed +Undeploys ungenerated unicode unicodecsv @@ -1670,6 +1701,7 @@ VolumeMount volumeMounts vpc WaiterModel +wape warmup Wasb wasb diff --git a/tests/always/test_project_structure.py b/tests/always/test_project_structure.py index f64a5077f23a6..696ce6dad5d1b 100644 --- a/tests/always/test_project_structure.py +++ b/tests/always/test_project_structure.py @@ -260,6 +260,15 @@ class TestGoogleProviderProjectStructure(ExampleCoverageTest, AssetsCoverageTest "airflow.providers.google.cloud.operators.dataproc.DataprocSubmitPySparkJobOperator", "airflow.providers.google.cloud.operators.mlengine.MLEngineManageModelOperator", "airflow.providers.google.cloud.operators.mlengine.MLEngineManageVersionOperator", + "airflow.providers.google.cloud.operators.mlengine.MLEngineCreateModelOperator", + "airflow.providers.google.cloud.operators.mlengine.MLEngineCreateVersionOperator", + "airflow.providers.google.cloud.operators.mlengine.MLEngineDeleteModelOperator", + "airflow.providers.google.cloud.operators.mlengine.MLEngineDeleteVersionOperator", + "airflow.providers.google.cloud.operators.mlengine.MLEngineGetModelOperator", + "airflow.providers.google.cloud.operators.mlengine.MLEngineListVersionsOperator", + "airflow.providers.google.cloud.operators.mlengine.MLEngineSetDefaultVersionOperator", + "airflow.providers.google.cloud.operators.mlengine.MLEngineStartBatchPredictionJobOperator", + "airflow.providers.google.cloud.operators.mlengine.MLEngineStartTrainingJobOperator", "airflow.providers.google.cloud.operators.dataflow.DataflowCreateJavaJobOperator", "airflow.providers.google.cloud.operators.bigquery.BigQueryPatchDatasetOperator", "airflow.providers.google.cloud.operators.dataflow.DataflowCreatePythonJobOperator", diff --git a/tests/providers/google/cloud/hooks/vertex_ai/test_model_service.py b/tests/providers/google/cloud/hooks/vertex_ai/test_model_service.py index 1b684fb0bd479..8798f09dc5aa6 100644 --- a/tests/providers/google/cloud/hooks/vertex_ai/test_model_service.py +++ b/tests/providers/google/cloud/hooks/vertex_ai/test_model_service.py @@ -37,6 +37,7 @@ BASE_STRING = "airflow.providers.google.common.hooks.base_google.{}" MODEL_SERVICE_STRING = "airflow.providers.google.cloud.hooks.vertex_ai.model_service.{}" +TEST_VERSION_ALIASES = ["new-alias"] class TestModelServiceWithDefaultProjectIdHook: @@ -131,6 +132,103 @@ def test_upload_model(self, mock_client) -> None: ) mock_client.return_value.common_location_path.assert_called_once_with(TEST_PROJECT_ID, TEST_REGION) + @mock.patch(MODEL_SERVICE_STRING.format("ModelServiceHook.get_model_service_client")) + def test_list_model_versions(self, mock_client) -> None: + self.hook.list_model_versions( + project_id=TEST_PROJECT_ID, region=TEST_REGION, model_id=TEST_MODEL_NAME + ) + mock_client.assert_called_once_with(TEST_REGION) + mock_client.return_value.list_model_versions.assert_called_once_with( + request=dict( + name=mock_client.return_value.model_path.return_value, + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(MODEL_SERVICE_STRING.format("ModelServiceHook.get_model_service_client")) + def test_delete_model_version(self, mock_client) -> None: + self.hook.delete_model_version( + project_id=TEST_PROJECT_ID, region=TEST_REGION, model_id=TEST_MODEL_NAME + ) + mock_client.assert_called_once_with(TEST_REGION) + mock_client.return_value.delete_model_version.assert_called_once_with( + request=dict( + name=mock_client.return_value.model_path.return_value, + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(MODEL_SERVICE_STRING.format("ModelServiceHook.get_model_service_client")) + def test_get_model(self, mock_client) -> None: + self.hook.get_model(project_id=TEST_PROJECT_ID, region=TEST_REGION, model_id=TEST_MODEL_NAME) + mock_client.assert_called_once_with(TEST_REGION) + mock_client.return_value.get_model.assert_called_once_with( + request=dict( + name=mock_client.return_value.model_path.return_value, + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(MODEL_SERVICE_STRING.format("ModelServiceHook.get_model_service_client")) + def test_set_version_as_default(self, mock_client) -> None: + self.hook.set_version_as_default( + project_id=TEST_PROJECT_ID, region=TEST_REGION, model_id=TEST_MODEL_NAME + ) + mock_client.assert_called_once_with(TEST_REGION) + mock_client.return_value.merge_version_aliases.assert_called_once_with( + request=dict( + name=mock_client.return_value.model_path.return_value, + version_aliases=["default"], + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(MODEL_SERVICE_STRING.format("ModelServiceHook.get_model_service_client")) + def test_add_version_aliases(self, mock_client) -> None: + self.hook.add_version_aliases( + project_id=TEST_PROJECT_ID, + region=TEST_REGION, + model_id=TEST_MODEL_NAME, + version_aliases=TEST_VERSION_ALIASES, + ) + mock_client.assert_called_once_with(TEST_REGION) + mock_client.return_value.merge_version_aliases.assert_called_once_with( + request=dict( + name=mock_client.return_value.model_path.return_value, + version_aliases=TEST_VERSION_ALIASES, + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(MODEL_SERVICE_STRING.format("ModelServiceHook.get_model_service_client")) + def test_delete_version_aliases(self, mock_client) -> None: + self.hook.delete_version_aliases( + project_id=TEST_PROJECT_ID, + region=TEST_REGION, + model_id=TEST_MODEL_NAME, + version_aliases=TEST_VERSION_ALIASES, + ) + mock_client.assert_called_once_with(TEST_REGION) + mock_client.return_value.merge_version_aliases.assert_called_once_with( + request=dict( + name=mock_client.return_value.model_path.return_value, + version_aliases=["-" + alias for alias in TEST_VERSION_ALIASES], + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + class TestModelServiceWithoutDefaultProjectIdHook: def setup_method(self): @@ -219,3 +317,100 @@ def test_upload_model(self, mock_client) -> None: timeout=None, ) mock_client.return_value.common_location_path.assert_called_once_with(TEST_PROJECT_ID, TEST_REGION) + + @mock.patch(MODEL_SERVICE_STRING.format("ModelServiceHook.get_model_service_client")) + def test_list_model_versions(self, mock_client) -> None: + self.hook.list_model_versions( + project_id=TEST_PROJECT_ID, region=TEST_REGION, model_id=TEST_MODEL_NAME + ) + mock_client.assert_called_once_with(TEST_REGION) + mock_client.return_value.list_model_versions.assert_called_once_with( + request=dict( + name=mock_client.return_value.model_path.return_value, + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(MODEL_SERVICE_STRING.format("ModelServiceHook.get_model_service_client")) + def test_delete_model_version(self, mock_client) -> None: + self.hook.delete_model_version( + project_id=TEST_PROJECT_ID, region=TEST_REGION, model_id=TEST_MODEL_NAME + ) + mock_client.assert_called_once_with(TEST_REGION) + mock_client.return_value.delete_model_version.assert_called_once_with( + request=dict( + name=mock_client.return_value.model_path.return_value, + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(MODEL_SERVICE_STRING.format("ModelServiceHook.get_model_service_client")) + def test_get_model(self, mock_client) -> None: + self.hook.get_model(project_id=TEST_PROJECT_ID, region=TEST_REGION, model_id=TEST_MODEL_NAME) + mock_client.assert_called_once_with(TEST_REGION) + mock_client.return_value.get_model.assert_called_once_with( + request=dict( + name=mock_client.return_value.model_path.return_value, + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(MODEL_SERVICE_STRING.format("ModelServiceHook.get_model_service_client")) + def test_set_version_as_default(self, mock_client) -> None: + self.hook.set_version_as_default( + project_id=TEST_PROJECT_ID, region=TEST_REGION, model_id=TEST_MODEL_NAME + ) + mock_client.assert_called_once_with(TEST_REGION) + mock_client.return_value.merge_version_aliases.assert_called_once_with( + request=dict( + name=mock_client.return_value.model_path.return_value, + version_aliases=["default"], + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(MODEL_SERVICE_STRING.format("ModelServiceHook.get_model_service_client")) + def test_add_version_aliases(self, mock_client) -> None: + self.hook.add_version_aliases( + project_id=TEST_PROJECT_ID, + region=TEST_REGION, + model_id=TEST_MODEL_NAME, + version_aliases=TEST_VERSION_ALIASES, + ) + mock_client.assert_called_once_with(TEST_REGION) + mock_client.return_value.merge_version_aliases.assert_called_once_with( + request=dict( + name=mock_client.return_value.model_path.return_value, + version_aliases=TEST_VERSION_ALIASES, + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(MODEL_SERVICE_STRING.format("ModelServiceHook.get_model_service_client")) + def test_delete_version_aliases(self, mock_client) -> None: + self.hook.delete_version_aliases( + project_id=TEST_PROJECT_ID, + region=TEST_REGION, + model_id=TEST_MODEL_NAME, + version_aliases=TEST_VERSION_ALIASES, + ) + mock_client.assert_called_once_with(TEST_REGION) + mock_client.return_value.merge_version_aliases.assert_called_once_with( + request=dict( + name=mock_client.return_value.model_path.return_value, + version_aliases=["-" + alias for alias in TEST_VERSION_ALIASES], + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) diff --git a/tests/providers/google/cloud/operators/test_vertex_ai.py b/tests/providers/google/cloud/operators/test_vertex_ai.py index f2921942ab41f..91e1d4718a8a8 100644 --- a/tests/providers/google/cloud/operators/test_vertex_ai.py +++ b/tests/providers/google/cloud/operators/test_vertex_ai.py @@ -65,9 +65,15 @@ ListHyperparameterTuningJobOperator, ) from airflow.providers.google.cloud.operators.vertex_ai.model_service import ( + AddVersionAliasesOnModelOperator, DeleteModelOperator, + DeleteModelVersionOperator, + DeleteVersionAliasesOnModelOperator, ExportModelOperator, + GetModelOperator, ListModelsOperator, + ListModelVersionsOperator, + SetDefaultVersionOnModelOperator, UploadModelOperator, ) @@ -168,6 +174,7 @@ TEST_CREATE_REQUEST_TIMEOUT = 100.5 TEST_BATCH_SIZE = 4000 +TEST_VERSION_ALIASES = ["new-alias"] class TestVertexAICreateCustomContainerTrainingJobOperator: @@ -219,6 +226,7 @@ def test_execute(self, mock_hook): test_fraction_split=TEST_FRACTION_SPLIT, region=GCP_LOCATION, project_id=GCP_PROJECT, + parent_model=None, model_serving_container_predict_route=None, model_serving_container_health_route=None, model_serving_container_command=None, @@ -249,6 +257,9 @@ def test_execute(self, mock_hook): timestamp_split_column_name=None, tensorboard=None, sync=True, + is_default_version=None, + model_version_aliases=None, + model_version_description=None, ) @@ -303,6 +314,10 @@ def test_execute(self, mock_hook): test_fraction_split=TEST_FRACTION_SPLIT, region=GCP_LOCATION, project_id=GCP_PROJECT, + parent_model=None, + is_default_version=None, + model_version_aliases=None, + model_version_description=None, model_serving_container_predict_route=None, model_serving_container_health_route=None, model_serving_container_command=None, @@ -378,6 +393,7 @@ def test_execute(self, mock_hook): training_fraction_split=None, validation_fraction_split=None, test_fraction_split=None, + parent_model=None, region=GCP_LOCATION, project_id=GCP_PROJECT, model_serving_container_predict_route=None, @@ -410,6 +426,9 @@ def test_execute(self, mock_hook): timestamp_split_column_name=None, tensorboard=None, sync=True, + is_default_version=None, + model_version_aliases=None, + model_version_description=None, ) @@ -708,6 +727,7 @@ def test_execute(self, mock_hook, mock_dataset): forecast_horizon=TEST_TRAINING_FORECAST_HORIZON, data_granularity_unit=TEST_TRAINING_DATA_GRANULARITY_UNIT, data_granularity_count=TEST_TRAINING_DATA_GRANULARITY_COUNT, + parent_model=None, optimization_objective=None, column_specs=None, column_transformations=None, @@ -730,6 +750,9 @@ def test_execute(self, mock_hook, mock_dataset): model_display_name=None, model_labels=None, sync=True, + is_default_version=None, + model_version_aliases=None, + model_version_description=None, ) @@ -760,6 +783,7 @@ def test_execute(self, mock_hook, mock_dataset): display_name=DISPLAY_NAME, dataset=mock_dataset.return_value, prediction_type="classification", + parent_model=None, multi_label=False, model_type="CLOUD", base_model=None, @@ -777,6 +801,9 @@ def test_execute(self, mock_hook, mock_dataset): model_labels=None, disable_early_stopping=False, sync=True, + is_default_version=None, + model_version_aliases=None, + model_version_description=None, ) @@ -812,6 +839,7 @@ def test_execute(self, mock_hook, mock_dataset): region=GCP_LOCATION, display_name=DISPLAY_NAME, dataset=mock_dataset.return_value, + parent_model=None, target_column=None, optimization_prediction_type=None, optimization_objective=None, @@ -836,6 +864,9 @@ def test_execute(self, mock_hook, mock_dataset): export_evaluated_data_items_bigquery_destination_uri=None, export_evaluated_data_items_override_destination=False, sync=True, + is_default_version=None, + model_version_aliases=None, + model_version_description=None, ) @@ -865,6 +896,7 @@ def test_execute(self, mock_hook, mock_dataset): region=GCP_LOCATION, display_name=DISPLAY_NAME, dataset=mock_dataset.return_value, + parent_model=None, prediction_type=None, multi_label=False, sentiment_max=10, @@ -880,6 +912,9 @@ def test_execute(self, mock_hook, mock_dataset): model_display_name=None, model_labels=None, sync=True, + is_default_version=None, + model_version_aliases=None, + model_version_description=None, ) @@ -908,6 +943,7 @@ def test_execute(self, mock_hook, mock_dataset): region=GCP_LOCATION, display_name=DISPLAY_NAME, dataset=mock_dataset.return_value, + parent_model=None, prediction_type="classification", model_type="CLOUD", labels=None, @@ -920,6 +956,9 @@ def test_execute(self, mock_hook, mock_dataset): model_display_name=None, model_labels=None, sync=True, + is_default_version=None, + model_version_aliases=None, + model_version_description=None, ) @@ -1505,3 +1544,173 @@ def test_execute(self, mock_hook, to_dict_mock): timeout=TIMEOUT, metadata=METADATA, ) + + +class TestVertexAIGetModelsOperator: + @mock.patch(VERTEX_AI_PATH.format("model_service.Model.to_dict")) + @mock.patch(VERTEX_AI_PATH.format("model_service.ModelServiceHook")) + def test_execute(self, mock_hook, to_dict_mock): + + op = GetModelOperator( + task_id=TASK_ID, + model_id=TEST_MODEL_NAME, + gcp_conn_id=GCP_CONN_ID, + impersonation_chain=IMPERSONATION_CHAIN, + region=GCP_LOCATION, + project_id=GCP_PROJECT, + retry=RETRY, + timeout=TIMEOUT, + metadata=METADATA, + ) + op.execute(context={"ti": mock.MagicMock()}) + mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) + mock_hook.return_value.get_model.assert_called_once_with( + region=GCP_LOCATION, + project_id=GCP_PROJECT, + model_id=TEST_MODEL_NAME, + retry=RETRY, + timeout=TIMEOUT, + metadata=METADATA, + ) + + +class TestVertexAIListModelVersionsOperator: + @mock.patch(VERTEX_AI_PATH.format("model_service.Model.to_dict")) + @mock.patch(VERTEX_AI_PATH.format("model_service.ModelServiceHook")) + def test_execute(self, mock_hook, to_dict_mock): + + op = ListModelVersionsOperator( + task_id=TASK_ID, + model_id=TEST_MODEL_NAME, + gcp_conn_id=GCP_CONN_ID, + impersonation_chain=IMPERSONATION_CHAIN, + region=GCP_LOCATION, + project_id=GCP_PROJECT, + retry=RETRY, + timeout=TIMEOUT, + metadata=METADATA, + ) + op.execute(context={"ti": mock.MagicMock()}) + mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) + mock_hook.return_value.list_model_versions.assert_called_once_with( + region=GCP_LOCATION, + project_id=GCP_PROJECT, + model_id=TEST_MODEL_NAME, + retry=RETRY, + timeout=TIMEOUT, + metadata=METADATA, + ) + + +class TestVertexAISetDefaultVersionOnModelOperator: + @mock.patch(VERTEX_AI_PATH.format("model_service.Model.to_dict")) + @mock.patch(VERTEX_AI_PATH.format("model_service.ModelServiceHook")) + def test_execute(self, mock_hook, to_dict_mock): + + op = SetDefaultVersionOnModelOperator( + task_id=TASK_ID, + model_id=TEST_MODEL_NAME, + gcp_conn_id=GCP_CONN_ID, + impersonation_chain=IMPERSONATION_CHAIN, + region=GCP_LOCATION, + project_id=GCP_PROJECT, + retry=RETRY, + timeout=TIMEOUT, + metadata=METADATA, + ) + op.execute(context={"ti": mock.MagicMock()}) + mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) + mock_hook.return_value.set_version_as_default.assert_called_once_with( + region=GCP_LOCATION, + project_id=GCP_PROJECT, + model_id=TEST_MODEL_NAME, + retry=RETRY, + timeout=TIMEOUT, + metadata=METADATA, + ) + + +class TestVertexAIAddVersionAliasesOnModelOperator: + @mock.patch(VERTEX_AI_PATH.format("model_service.Model.to_dict")) + @mock.patch(VERTEX_AI_PATH.format("model_service.ModelServiceHook")) + def test_execute(self, mock_hook, to_dict_mock): + + op = AddVersionAliasesOnModelOperator( + task_id=TASK_ID, + model_id=TEST_MODEL_NAME, + version_aliases=TEST_VERSION_ALIASES, + gcp_conn_id=GCP_CONN_ID, + impersonation_chain=IMPERSONATION_CHAIN, + region=GCP_LOCATION, + project_id=GCP_PROJECT, + retry=RETRY, + timeout=TIMEOUT, + metadata=METADATA, + ) + op.execute(context={"ti": mock.MagicMock()}) + mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) + mock_hook.return_value.add_version_aliases.assert_called_once_with( + region=GCP_LOCATION, + project_id=GCP_PROJECT, + model_id=TEST_MODEL_NAME, + version_aliases=TEST_VERSION_ALIASES, + retry=RETRY, + timeout=TIMEOUT, + metadata=METADATA, + ) + + +class TestVertexAIDeleteVersionAliasesOnModelOperator: + @mock.patch(VERTEX_AI_PATH.format("model_service.Model.to_dict")) + @mock.patch(VERTEX_AI_PATH.format("model_service.ModelServiceHook")) + def test_execute(self, mock_hook, to_dict_mock): + op = DeleteVersionAliasesOnModelOperator( + task_id=TASK_ID, + gcp_conn_id=GCP_CONN_ID, + impersonation_chain=IMPERSONATION_CHAIN, + region=GCP_LOCATION, + project_id=GCP_PROJECT, + model_id=TEST_MODEL_ID, + version_aliases=TEST_VERSION_ALIASES, + retry=RETRY, + timeout=TIMEOUT, + metadata=METADATA, + ) + op.execute(context={"ti": mock.MagicMock()}) + mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) + mock_hook.return_value.delete_version_aliases.assert_called_once_with( + region=GCP_LOCATION, + project_id=GCP_PROJECT, + model_id=TEST_MODEL_ID, + version_aliases=TEST_VERSION_ALIASES, + retry=RETRY, + timeout=TIMEOUT, + metadata=METADATA, + ) + + +class TestVertexAIDeleteModelVersionOperator: + @mock.patch(VERTEX_AI_PATH.format("model_service.Model.to_dict")) + @mock.patch(VERTEX_AI_PATH.format("model_service.ModelServiceHook")) + def test_execute(self, mock_hook, to_dict_mock): + op = DeleteModelVersionOperator( + task_id=TASK_ID, + gcp_conn_id=GCP_CONN_ID, + impersonation_chain=IMPERSONATION_CHAIN, + region=GCP_LOCATION, + project_id=GCP_PROJECT, + model_id=TEST_MODEL_ID, + retry=RETRY, + timeout=TIMEOUT, + metadata=METADATA, + ) + op.execute(context={"ti": mock.MagicMock()}) + mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) + mock_hook.return_value.delete_model_version.assert_called_once_with( + region=GCP_LOCATION, + project_id=GCP_PROJECT, + model_id=TEST_MODEL_ID, + retry=RETRY, + timeout=TIMEOUT, + metadata=METADATA, + ) diff --git a/tests/system/providers/google/cloud/ml_engine/example_mlengine.py b/tests/system/providers/google/cloud/ml_engine/example_mlengine.py index 8fac67aab5403..4ae63696f302b 100644 --- a/tests/system/providers/google/cloud/ml_engine/example_mlengine.py +++ b/tests/system/providers/google/cloud/ml_engine/example_mlengine.py @@ -21,300 +21,263 @@ from __future__ import annotations import os -import pathlib from datetime import datetime -from math import ceil -from airflow.decorators import task -from airflow.models.dag import DAG +from google.cloud.aiplatform import schema +from google.protobuf.json_format import ParseDict +from google.protobuf.struct_pb2 import Value + +from airflow import models from airflow.operators.bash import BashOperator -from airflow.providers.google.cloud.operators.gcs import GCSCreateBucketOperator, GCSDeleteBucketOperator -from airflow.providers.google.cloud.operators.mlengine import ( - MLEngineCreateModelOperator, - MLEngineCreateVersionOperator, - MLEngineDeleteModelOperator, - MLEngineDeleteVersionOperator, - MLEngineGetModelOperator, - MLEngineListVersionsOperator, - MLEngineSetDefaultVersionOperator, - MLEngineStartBatchPredictionJobOperator, - MLEngineStartTrainingJobOperator, +from airflow.providers.google.cloud.operators.gcs import ( + GCSCreateBucketOperator, + GCSDeleteBucketOperator, + GCSSynchronizeBucketsOperator, +) +from airflow.providers.google.cloud.operators.vertex_ai.batch_prediction_job import ( + CreateBatchPredictionJobOperator, + DeleteBatchPredictionJobOperator, +) +from airflow.providers.google.cloud.operators.vertex_ai.custom_job import ( + CreateCustomPythonPackageTrainingJobOperator, +) +from airflow.providers.google.cloud.operators.vertex_ai.dataset import ( + CreateDatasetOperator, + DeleteDatasetOperator, +) +from airflow.providers.google.cloud.operators.vertex_ai.model_service import ( + DeleteModelOperator, + DeleteModelVersionOperator, + GetModelOperator, + ListModelVersionsOperator, + SetDefaultVersionOnModelOperator, ) -from airflow.providers.google.cloud.transfers.local_to_gcs import LocalFilesystemToGCSOperator -from airflow.providers.google.cloud.utils import mlengine_operator_utils from airflow.utils.trigger_rule import TriggerRule PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT", "default") -ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID") - +ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default") DAG_ID = "example_gcp_mlengine" -PREDICT_FILE_NAME = "predict.json" -MODEL_NAME = f"model_{DAG_ID}_{ENV_ID}".replace("_", "-") -BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}".replace("_", "-") -BUCKET_PATH = f"gs://{BUCKET_NAME}" -JOB_DIR = f"{BUCKET_PATH}/job-dir" -SAVED_MODEL_PATH = f"{JOB_DIR}/" -PREDICTION_INPUT = f"{BUCKET_PATH}/{PREDICT_FILE_NAME}" -PREDICTION_OUTPUT = f"{BUCKET_PATH}/prediction_output/" -TRAINER_URI = "gs://airflow-system-tests-resources/ml-engine/trainer-0.2.tar.gz" -TRAINER_PY_MODULE = "trainer.task" -SUMMARY_TMP = f"{BUCKET_PATH}/tmp/" -SUMMARY_STAGING = f"{BUCKET_PATH}/staging/" +REGION = "us-central1" -BASE_DIR = pathlib.Path(__file__).parent.resolve() -PATH_TO_PREDICT_FILE = BASE_DIR / PREDICT_FILE_NAME +PACKAGE_DISPLAY_NAME = f"package-{DAG_ID}-{ENV_ID}".replace("_", "-") +MODEL_DISPLAY_NAME = f"model-{DAG_ID}-{ENV_ID}".replace("_", "-") +JOB_DISPLAY_NAME = f"batch_job_{DAG_ID}_{ENV_ID}".replace("-", "_") +RESOURCE_DATA_BUCKET = "airflow-system-tests-resources" +CUSTOM_PYTHON_GCS_BUCKET_NAME = f"bucket_python_{DAG_ID}_{ENV_ID}".replace("_", "-") -def generate_model_predict_input_data() -> list[int]: - return [1, 4, 9, 16, 25, 36] +BQ_SOURCE = "bq://bigquery-public-data.ml_datasets.penguins" +TABULAR_DATASET = { + "display_name": f"tabular-dataset-{ENV_ID}", + "metadata_schema_uri": schema.dataset.metadata.tabular, + "metadata": ParseDict( + {"input_config": {"bigquery_source": {"uri": BQ_SOURCE}}}, + Value(), + ), +} +REPLICA_COUNT = 1 +MACHINE_TYPE = "n1-standard-4" +ACCELERATOR_TYPE = "ACCELERATOR_TYPE_UNSPECIFIED" +ACCELERATOR_COUNT = 0 +TRAINING_FRACTION_SPLIT = 0.7 +TEST_FRACTION_SPLIT = 0.15 +VALIDATION_FRACTION_SPLIT = 0.15 -with DAG( +PYTHON_PACKAGE_GCS_URI = f"gs://{CUSTOM_PYTHON_GCS_BUCKET_NAME}/vertex-ai/penguins_trainer_script-0.1.zip" +PYTHON_MODULE_NAME = "penguins_trainer_script.task" + +TRAIN_IMAGE = "us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-8:latest" +DEPLOY_IMAGE = "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:latest" + + +with models.DAG( dag_id=DAG_ID, schedule="@once", start_date=datetime(2021, 1, 1), catchup=False, tags=["example", "ml_engine"], - params={"model_name": MODEL_NAME}, ) as dag: create_bucket = GCSCreateBucketOperator( - task_id="create-bucket", - bucket_name=BUCKET_NAME, + task_id="create_bucket", + bucket_name=CUSTOM_PYTHON_GCS_BUCKET_NAME, + storage_class="REGIONAL", + location=REGION, ) - @task(task_id="write-predict-data-file") - def write_predict_file(path_to_file: str): - predict_data = generate_model_predict_input_data() - with open(path_to_file, "w") as file: - for predict_value in predict_data: - file.write(f'{{"input_layer": [{predict_value}]}}\n') - - write_data = write_predict_file(path_to_file=PATH_TO_PREDICT_FILE) - - upload_file = LocalFilesystemToGCSOperator( - task_id="upload-predict-file", - src=[PATH_TO_PREDICT_FILE], - dst=PREDICT_FILE_NAME, - bucket=BUCKET_NAME, + move_data_files = GCSSynchronizeBucketsOperator( + task_id="move_files_to_bucket", + source_bucket=RESOURCE_DATA_BUCKET, + source_object="vertex-ai/penguins-data", + destination_bucket=CUSTOM_PYTHON_GCS_BUCKET_NAME, + destination_object="vertex-ai", + recursive=True, ) - - # [START howto_operator_gcp_mlengine_training] - training = MLEngineStartTrainingJobOperator( - task_id="training", + create_tabular_dataset = CreateDatasetOperator( + task_id="tabular_dataset", + dataset=TABULAR_DATASET, + region=REGION, project_id=PROJECT_ID, - region="us-central1", - job_id="training-job-{{ ts_nodash }}-{{ params.model_name }}", - runtime_version="1.15", - python_version="3.7", - job_dir=JOB_DIR, - package_uris=[TRAINER_URI], - training_python_module=TRAINER_PY_MODULE, - training_args=[], - labels={"job_type": "training"}, ) - # [END howto_operator_gcp_mlengine_training] - - # [START howto_operator_gcp_mlengine_create_model] - create_model = MLEngineCreateModelOperator( - task_id="create-model", + tabular_dataset_id = create_tabular_dataset.output["dataset_id"] + + # [START howto_operator_create_custom_python_training_job_v1] + create_custom_python_package_training_job = CreateCustomPythonPackageTrainingJobOperator( + task_id="create_custom_python_package_training_job", + staging_bucket=f"gs://{CUSTOM_PYTHON_GCS_BUCKET_NAME}", + display_name=PACKAGE_DISPLAY_NAME, + python_package_gcs_uri=PYTHON_PACKAGE_GCS_URI, + python_module_name=PYTHON_MODULE_NAME, + container_uri=TRAIN_IMAGE, + model_serving_container_image_uri=DEPLOY_IMAGE, + bigquery_destination=f"bq://{PROJECT_ID}", + # run params + dataset_id=tabular_dataset_id, + model_display_name=MODEL_DISPLAY_NAME, + replica_count=REPLICA_COUNT, + machine_type=MACHINE_TYPE, + accelerator_type=ACCELERATOR_TYPE, + accelerator_count=ACCELERATOR_COUNT, + training_fraction_split=TRAINING_FRACTION_SPLIT, + validation_fraction_split=VALIDATION_FRACTION_SPLIT, + test_fraction_split=TEST_FRACTION_SPLIT, + region=REGION, project_id=PROJECT_ID, - model={ - "name": MODEL_NAME, - }, ) - # [END howto_operator_gcp_mlengine_create_model] + # [END howto_operator_create_custom_python_training_job_v1] + model_id_v1 = create_custom_python_package_training_job.output["model_id"] # [START howto_operator_gcp_mlengine_get_model] - get_model = MLEngineGetModelOperator( - task_id="get-model", - project_id=PROJECT_ID, - model_name=MODEL_NAME, + get_model = GetModelOperator( + task_id="get_model", region=REGION, project_id=PROJECT_ID, model_id=model_id_v1 ) # [END howto_operator_gcp_mlengine_get_model] # [START howto_operator_gcp_mlengine_print_model] get_model_result = BashOperator( bash_command=f"echo {get_model.output}", - task_id="get-model-result", + task_id="get_model_result", ) # [END howto_operator_gcp_mlengine_print_model] - # [START howto_operator_gcp_mlengine_create_version1] - create_version_v1 = MLEngineCreateVersionOperator( - task_id="create-version-v1", - project_id=PROJECT_ID, - model_name=MODEL_NAME, - version={ - "name": "v1", - "description": "First-version", - "deployment_uri": JOB_DIR, - "runtime_version": "1.15", - "machineType": "mls1-c1-m2", - "framework": "TENSORFLOW", - "pythonVersion": "3.7", - }, - ) - # [END howto_operator_gcp_mlengine_create_version1] - - # [START howto_operator_gcp_mlengine_create_version2] - create_version_v2 = MLEngineCreateVersionOperator( - task_id="create-version-v2", + # [START howto_operator_create_custom_python_training_job_v2] + create_custom_python_package_training_job_v2 = CreateCustomPythonPackageTrainingJobOperator( + task_id="create_custom_python_package_training_job_v2", + staging_bucket=f"gs://{CUSTOM_PYTHON_GCS_BUCKET_NAME}", + display_name=PACKAGE_DISPLAY_NAME, + python_package_gcs_uri=PYTHON_PACKAGE_GCS_URI, + python_module_name=PYTHON_MODULE_NAME, + container_uri=TRAIN_IMAGE, + model_serving_container_image_uri=DEPLOY_IMAGE, + bigquery_destination=f"bq://{PROJECT_ID}", + parent_model=model_id_v1, + # run params + dataset_id=tabular_dataset_id, + model_display_name=MODEL_DISPLAY_NAME, + replica_count=REPLICA_COUNT, + machine_type=MACHINE_TYPE, + accelerator_type=ACCELERATOR_TYPE, + accelerator_count=ACCELERATOR_COUNT, + training_fraction_split=TRAINING_FRACTION_SPLIT, + validation_fraction_split=VALIDATION_FRACTION_SPLIT, + test_fraction_split=TEST_FRACTION_SPLIT, + region=REGION, project_id=PROJECT_ID, - model_name=MODEL_NAME, - version={ - "name": "v2", - "description": "Second version", - "deployment_uri": JOB_DIR, - "runtime_version": "1.15", - "machineType": "mls1-c1-m2", - "framework": "TENSORFLOW", - "pythonVersion": "3.7", - }, ) - # [END howto_operator_gcp_mlengine_create_version2] + # [END howto_operator_create_custom_python_training_job_v2] + model_id_v2 = create_custom_python_package_training_job_v2.output["model_id"] # [START howto_operator_gcp_mlengine_default_version] - set_defaults_version = MLEngineSetDefaultVersionOperator( - task_id="set-default-version", + set_default_version = SetDefaultVersionOnModelOperator( + task_id="set_default_version", project_id=PROJECT_ID, - model_name=MODEL_NAME, - version_name="v2", + region=REGION, + model_id=model_id_v2, ) # [END howto_operator_gcp_mlengine_default_version] # [START howto_operator_gcp_mlengine_list_versions] - list_version = MLEngineListVersionsOperator( - task_id="list-version", - project_id=PROJECT_ID, - model_name=MODEL_NAME, + list_model_versions = ListModelVersionsOperator( + task_id="list_model_versions", region=REGION, project_id=PROJECT_ID, model_id=model_id_v2 ) # [END howto_operator_gcp_mlengine_list_versions] - # [START howto_operator_gcp_mlengine_print_versions] - list_version_result = BashOperator( - bash_command=f"echo {list_version.output}", - task_id="list-version-result", - ) - # [END howto_operator_gcp_mlengine_print_versions] - - # [START howto_operator_gcp_mlengine_get_prediction] - prediction = MLEngineStartBatchPredictionJobOperator( - task_id="prediction", + # [START howto_operator_start_batch_prediction] + create_batch_prediction_job = CreateBatchPredictionJobOperator( + task_id="create_batch_prediction_job", + job_display_name=JOB_DISPLAY_NAME, + model_name=model_id_v2, + predictions_format="bigquery", + bigquery_source=BQ_SOURCE, + bigquery_destination_prefix=f"bq://{PROJECT_ID}", + region=REGION, project_id=PROJECT_ID, - job_id="prediction-{{ ts_nodash }}-{{ params.model_name }}", - region="us-central1", - model_name=MODEL_NAME, - data_format="TEXT", - input_paths=[PREDICTION_INPUT], - output_path=PREDICTION_OUTPUT, - labels={"job_type": "prediction"}, + machine_type=MACHINE_TYPE, ) - # [END howto_operator_gcp_mlengine_get_prediction] + # [END howto_operator_start_batch_prediction] # [START howto_operator_gcp_mlengine_delete_version] - delete_version_v1 = MLEngineDeleteVersionOperator( - task_id="delete-version-v1", + delete_model_version_1 = DeleteModelVersionOperator( + task_id="delete_model_version_1", project_id=PROJECT_ID, - model_name=MODEL_NAME, - version_name="v1", + region=REGION, + model_id=model_id_v2, trigger_rule=TriggerRule.ALL_DONE, ) # [END howto_operator_gcp_mlengine_delete_version] - delete_version_v2 = MLEngineDeleteVersionOperator( - task_id="delete-version-v2", - project_id=PROJECT_ID, - model_name=MODEL_NAME, - version_name="v2", - trigger_rule=TriggerRule.ALL_DONE, - ) - # [START howto_operator_gcp_mlengine_delete_model] - delete_model = MLEngineDeleteModelOperator( - task_id="delete-model", + delete_model = DeleteModelOperator( + task_id="delete_model", project_id=PROJECT_ID, - model_name=MODEL_NAME, - delete_contents=True, + region=REGION, + model_id=model_id_v1, trigger_rule=TriggerRule.ALL_DONE, ) # [END howto_operator_gcp_mlengine_delete_model] - delete_bucket = GCSDeleteBucketOperator( - task_id="delete-bucket", - bucket_name=BUCKET_NAME, + delete_batch_prediction_job = DeleteBatchPredictionJobOperator( + task_id="delete_batch_prediction_job", + batch_prediction_job_id=create_batch_prediction_job.output["batch_prediction_job_id"], + region=REGION, + project_id=PROJECT_ID, trigger_rule=TriggerRule.ALL_DONE, ) - # [START howto_operator_gcp_mlengine_get_metric] - def get_metric_fn_and_keys(): - """ - Gets metric function and keys used to generate summary - """ - - def normalize_value(inst: dict): - val = int(inst["output_layer"][0]) - return tuple([val]) # returns a tuple. - - return normalize_value, ["val"] # key order must match. - - # [END howto_operator_gcp_mlengine_get_metric] - - # [START howto_operator_gcp_mlengine_validate_error] - def validate_err_and_count(summary: dict) -> dict: - """ - Validate summary result - """ - summary = summary.get("val", 0) - initial_values = generate_model_predict_input_data() - initial_summary = sum(initial_values) / len(initial_values) - - multiplier = ceil(summary / initial_summary) - if multiplier != 2: - raise ValueError(f"Multiplier is not equal 2; multiplier: {multiplier}") - return summary - - # [END howto_operator_gcp_mlengine_validate_error] - - # [START howto_operator_gcp_mlengine_evaluate] - evaluate_prediction, evaluate_summary, evaluate_validation = mlengine_operator_utils.create_evaluate_ops( - task_prefix="evaluate-ops", - data_format="TEXT", - input_paths=[PREDICTION_INPUT], - prediction_path=PREDICTION_OUTPUT, - metric_fn_and_keys=get_metric_fn_and_keys(), - validate_fn=validate_err_and_count, - batch_prediction_job_id="evaluate-ops-{{ ts_nodash }}-{{ params.model_name }}", + delete_tabular_dataset = DeleteDatasetOperator( + task_id="delete_tabular_dataset", + dataset_id=tabular_dataset_id, + region=REGION, project_id=PROJECT_ID, - region="us-central1", - dataflow_options={ - "project": PROJECT_ID, - "tempLocation": SUMMARY_TMP, - "stagingLocation": SUMMARY_STAGING, - }, - model_name=MODEL_NAME, - version_name="v1", - py_interpreter="python3", + trigger_rule=TriggerRule.ALL_DONE, ) - # [END howto_operator_gcp_mlengine_evaluate] - - # TEST SETUP - create_bucket >> write_data >> upload_file - upload_file >> [prediction, evaluate_prediction] - create_bucket >> training >> create_version_v1 - - # TEST BODY - create_model >> get_model >> [get_model_result, delete_model] - create_model >> create_version_v1 >> create_version_v2 >> set_defaults_version >> list_version - create_version_v1 >> prediction - create_version_v1 >> evaluate_prediction - create_version_v2 >> prediction - - list_version >> [list_version_result, delete_version_v1] - prediction >> delete_version_v1 + delete_bucket = GCSDeleteBucketOperator( + task_id="delete_bucket", + bucket_name=CUSTOM_PYTHON_GCS_BUCKET_NAME, + trigger_rule=TriggerRule.ALL_DONE, + ) - # TEST TEARDOWN - evaluate_validation >> delete_version_v1 >> delete_version_v2 >> delete_model >> delete_bucket + ( + # TEST SETUP + create_bucket + >> move_data_files + >> create_tabular_dataset + # TEST BODY + >> create_custom_python_package_training_job + >> create_custom_python_package_training_job_v2 + >> create_batch_prediction_job + >> get_model + >> get_model_result + >> list_model_versions + >> set_default_version + # TEST TEARDOWN + >> delete_model_version_1 + >> delete_model + >> delete_batch_prediction_job + >> delete_tabular_dataset + >> delete_bucket + ) from tests.system.utils.watcher import watcher diff --git a/tests/system/providers/google/cloud/ml_engine/example_mlengine_async.py b/tests/system/providers/google/cloud/ml_engine/example_mlengine_async.py deleted file mode 100644 index 301fb1ffa2b1a..0000000000000 --- a/tests/system/providers/google/cloud/ml_engine/example_mlengine_async.py +++ /dev/null @@ -1,329 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -""" -Example Airflow DAG for Google ML Engine service. -""" -from __future__ import annotations - -import os -import pathlib -from datetime import datetime -from math import ceil - -from airflow.decorators import task -from airflow.models.dag import DAG -from airflow.operators.bash import BashOperator -from airflow.providers.google.cloud.operators.gcs import GCSCreateBucketOperator, GCSDeleteBucketOperator -from airflow.providers.google.cloud.operators.mlengine import ( - MLEngineCreateModelOperator, - MLEngineCreateVersionOperator, - MLEngineDeleteModelOperator, - MLEngineDeleteVersionOperator, - MLEngineGetModelOperator, - MLEngineListVersionsOperator, - MLEngineSetDefaultVersionOperator, - MLEngineStartBatchPredictionJobOperator, - MLEngineStartTrainingJobOperator, -) -from airflow.providers.google.cloud.transfers.local_to_gcs import LocalFilesystemToGCSOperator -from airflow.providers.google.cloud.utils import mlengine_operator_utils -from airflow.utils.trigger_rule import TriggerRule - -PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT", "default") -ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID") - -DAG_ID = "async_example_gcp_mlengine" -PREDICT_FILE_NAME = "async_predict.json" -MODEL_NAME = f"model_{DAG_ID}_{ENV_ID}".replace("-", "_") -BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}".replace("_", "-") -BUCKET_PATH = f"gs://{BUCKET_NAME}" -JOB_DIR = f"{BUCKET_PATH}/job-dir" -SAVED_MODEL_PATH = f"{JOB_DIR}/" -PREDICTION_INPUT = f"{BUCKET_PATH}/{PREDICT_FILE_NAME}" -PREDICTION_OUTPUT = f"{BUCKET_PATH}/prediction_output/" -TRAINER_URI = "gs://airflow-system-tests-resources/ml-engine/async-trainer-0.2.tar.gz" -TRAINER_PY_MODULE = "trainer.task" -SUMMARY_TMP = f"{BUCKET_PATH}/tmp/" -SUMMARY_STAGING = f"{BUCKET_PATH}/staging/" - -BASE_DIR = pathlib.Path(__file__).parent.resolve() -PATH_TO_PREDICT_FILE = BASE_DIR / PREDICT_FILE_NAME - - -def generate_model_predict_input_data() -> list[int]: - return [1, 4, 9, 16, 25, 36] - - -with DAG( - dag_id=DAG_ID, - schedule="@once", - start_date=datetime(2021, 1, 1), - catchup=False, - tags=["example", "ml_engine", "deferrable"], - params={"model_name": MODEL_NAME}, -) as dag: - create_bucket = GCSCreateBucketOperator( - task_id="create-bucket", - bucket_name=BUCKET_NAME, - ) - - @task(task_id="write-predict-data-file") - def write_predict_file(path_to_file: str): - predict_data = generate_model_predict_input_data() - with open(path_to_file, "w") as file: - for predict_value in predict_data: - file.write(f'{{"input_layer": [{predict_value}]}}\n') - - write_data = write_predict_file(path_to_file=PATH_TO_PREDICT_FILE) - - upload_file = LocalFilesystemToGCSOperator( - task_id="upload-predict-file", - src=[PATH_TO_PREDICT_FILE], - dst=PREDICT_FILE_NAME, - bucket=BUCKET_NAME, - ) - - # [START howto_operator_gcp_mlengine_training_async] - training = MLEngineStartTrainingJobOperator( - task_id="training", - project_id=PROJECT_ID, - region="us-central1", - job_id="async_training-job-{{ ts_nodash }}-{{ params.model_name }}", - runtime_version="1.15", - python_version="3.7", - job_dir=JOB_DIR, - package_uris=[TRAINER_URI], - training_python_module=TRAINER_PY_MODULE, - training_args=[], - labels={"job_type": "training"}, - deferrable=True, - ) - # [END howto_operator_gcp_mlengine_training_async] - - # [START howto_operator_gcp_mlengine_create_model] - create_model = MLEngineCreateModelOperator( - task_id="create-model", - project_id=PROJECT_ID, - model={ - "name": MODEL_NAME, - }, - ) - # [END howto_operator_gcp_mlengine_create_model] - - # [START howto_operator_gcp_mlengine_get_model] - get_model = MLEngineGetModelOperator( - task_id="get-model", - project_id=PROJECT_ID, - model_name=MODEL_NAME, - ) - # [END howto_operator_gcp_mlengine_get_model] - - # [START howto_operator_gcp_mlengine_print_model] - get_model_result = BashOperator( - bash_command=f"echo {get_model.output}", - task_id="get-model-result", - ) - # [END howto_operator_gcp_mlengine_print_model] - - # [START howto_operator_gcp_mlengine_create_version1] - create_version_v1 = MLEngineCreateVersionOperator( - task_id="create-version-v1", - project_id=PROJECT_ID, - model_name=MODEL_NAME, - version={ - "name": "v1", - "description": "First-version", - "deployment_uri": JOB_DIR, - "runtime_version": "1.15", - "machineType": "mls1-c1-m2", - "framework": "TENSORFLOW", - "pythonVersion": "3.7", - }, - ) - # [END howto_operator_gcp_mlengine_create_version1] - - # [START howto_operator_gcp_mlengine_create_version2] - create_version_v2 = MLEngineCreateVersionOperator( - task_id="create-version-v2", - project_id=PROJECT_ID, - model_name=MODEL_NAME, - version={ - "name": "v2", - "description": "Second version", - "deployment_uri": JOB_DIR, - "runtime_version": "1.15", - "machineType": "mls1-c1-m2", - "framework": "TENSORFLOW", - "pythonVersion": "3.7", - }, - ) - # [END howto_operator_gcp_mlengine_create_version2] - - # [START howto_operator_gcp_mlengine_default_version] - set_defaults_version = MLEngineSetDefaultVersionOperator( - task_id="set-default-version", - project_id=PROJECT_ID, - model_name=MODEL_NAME, - version_name="v2", - ) - # [END howto_operator_gcp_mlengine_default_version] - - # [START howto_operator_gcp_mlengine_list_versions] - list_version = MLEngineListVersionsOperator( - task_id="list-version", - project_id=PROJECT_ID, - model_name=MODEL_NAME, - ) - # [END howto_operator_gcp_mlengine_list_versions] - - # [START howto_operator_gcp_mlengine_print_versions] - list_version_result = BashOperator( - bash_command=f"echo {list_version.output}", - task_id="list-version-result", - ) - # [END howto_operator_gcp_mlengine_print_versions] - - # [START howto_operator_gcp_mlengine_get_prediction] - prediction = MLEngineStartBatchPredictionJobOperator( - task_id="prediction", - project_id=PROJECT_ID, - job_id="async-prediction-{{ ts_nodash }}-{{ params.model_name }}", - region="us-central1", - model_name=MODEL_NAME, - data_format="TEXT", - input_paths=[PREDICTION_INPUT], - output_path=PREDICTION_OUTPUT, - labels={"job_type": "prediction"}, - ) - # [END howto_operator_gcp_mlengine_get_prediction] - - # [START howto_operator_gcp_mlengine_delete_version] - delete_version_v1 = MLEngineDeleteVersionOperator( - task_id="delete-version-v1", - project_id=PROJECT_ID, - model_name=MODEL_NAME, - version_name="v1", - trigger_rule=TriggerRule.ALL_DONE, - ) - # [END howto_operator_gcp_mlengine_delete_version] - - delete_version_v2 = MLEngineDeleteVersionOperator( - task_id="delete-version-v2", - project_id=PROJECT_ID, - model_name=MODEL_NAME, - version_name="v2", - trigger_rule=TriggerRule.ALL_DONE, - ) - - # [START howto_operator_gcp_mlengine_delete_model] - delete_model = MLEngineDeleteModelOperator( - task_id="delete-model", - project_id=PROJECT_ID, - model_name=MODEL_NAME, - delete_contents=True, - trigger_rule=TriggerRule.ALL_DONE, - ) - # [END howto_operator_gcp_mlengine_delete_model] - - delete_bucket = GCSDeleteBucketOperator( - task_id="delete-bucket", - bucket_name=BUCKET_NAME, - trigger_rule=TriggerRule.ALL_DONE, - ) - - # [START howto_operator_gcp_mlengine_get_metric] - def get_metric_fn_and_keys(): - """ - Gets metric function and keys used to generate summary - """ - - def normalize_value(inst: dict): - val = int(inst["output_layer"][0]) - return tuple([val]) # returns a tuple. - - return normalize_value, ["val"] # key order must match. - - # [END howto_operator_gcp_mlengine_get_metric] - - # [START howto_operator_gcp_mlengine_validate_error] - def validate_err_and_count(summary: dict) -> dict: - """ - Validate summary result - """ - summary = summary.get("val", 0) - initial_values = generate_model_predict_input_data() - initial_summary = sum(initial_values) / len(initial_values) - - multiplier = ceil(summary / initial_summary) - if multiplier != 2: - raise ValueError(f"Multiplier is not equal 2; multiplier: {multiplier}") - return summary - - # [END howto_operator_gcp_mlengine_validate_error] - - # [START howto_operator_gcp_mlengine_evaluate] - evaluate_prediction, evaluate_summary, evaluate_validation = mlengine_operator_utils.create_evaluate_ops( - task_prefix="evaluate-ops", - data_format="TEXT", - input_paths=[PREDICTION_INPUT], - prediction_path=PREDICTION_OUTPUT, - metric_fn_and_keys=get_metric_fn_and_keys(), - validate_fn=validate_err_and_count, - batch_prediction_job_id="async-evaluate-ops-{{ ts_nodash }}-{{ params.model_name }}", - project_id=PROJECT_ID, - region="us-central1", - dataflow_options={ - "project": PROJECT_ID, - "tempLocation": SUMMARY_TMP, - "stagingLocation": SUMMARY_STAGING, - }, - model_name=MODEL_NAME, - version_name="v1", - py_interpreter="python3", - ) - # [END howto_operator_gcp_mlengine_evaluate] - - # TEST SETUP - create_bucket >> write_data >> upload_file - upload_file >> [prediction, evaluate_prediction] - create_bucket >> training >> create_version_v1 - - # TEST BODY - create_model >> get_model >> [get_model_result, delete_model] - create_model >> create_version_v1 >> create_version_v2 >> set_defaults_version >> list_version - - create_version_v1 >> prediction - create_version_v1 >> evaluate_prediction - create_version_v2 >> prediction - - list_version >> [list_version_result, delete_version_v1] - prediction >> delete_version_v1 - - # TEST TEARDOWN - evaluate_validation >> delete_version_v1 >> delete_version_v2 >> delete_model >> delete_bucket - - from tests.system.utils.watcher import watcher - - # This test needs watcher in order to properly mark success/failure - # when "tearDown" task with trigger rule is part of the DAG - list(dag.tasks) >> watcher() - -from tests.system.utils import get_test_run # noqa: E402 - -# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) -test_run = get_test_run(dag) diff --git a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_auto_ml_video_training.py b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_auto_ml_video_training.py index c6810ed36ec3c..ec83aa9faef60 100644 --- a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_auto_ml_video_training.py +++ b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_auto_ml_video_training.py @@ -119,8 +119,23 @@ region=REGION, project_id=PROJECT_ID, ) + model_id_v1 = create_auto_ml_video_training_job.output["model_id"] # [END how_to_cloud_vertex_ai_create_auto_ml_video_training_job_operator] + # [START how_to_cloud_vertex_ai_create_auto_ml_video_training_job_v2_operator] + create_auto_ml_video_training_job_v2 = CreateAutoMLVideoTrainingJobOperator( + task_id="auto_ml_video_v2_task", + display_name=VIDEO_DISPLAY_NAME, + prediction_type="classification", + model_type="CLOUD", + dataset_id=video_dataset_id, + model_display_name=MODEL_DISPLAY_NAME, + parent_model=model_id_v1, + region=REGION, + project_id=PROJECT_ID, + ) + # [END how_to_cloud_vertex_ai_create_auto_ml_video_training_job_v2_operator] + delete_auto_ml_video_training_job = DeleteAutoMLTrainingJobOperator( task_id="delete_auto_ml_video_training_job", training_pipeline_id="{{ task_instance.xcom_pull(task_ids='auto_ml_video_task', " @@ -153,6 +168,7 @@ >> import_video_dataset # TEST BODY >> create_auto_ml_video_training_job + >> create_auto_ml_video_training_job_v2 # TEST TEARDOWN >> delete_auto_ml_video_training_job >> delete_video_dataset diff --git a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_job.py b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_job.py index 364cbad67334f..6fa52e873ab34 100644 --- a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_job.py +++ b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_job.py @@ -134,8 +134,29 @@ def TABULAR_DATASET(bucket_name): region=REGION, project_id=PROJECT_ID, ) + model_id_v1 = create_custom_training_job.output["model_id"] # [END how_to_cloud_vertex_ai_create_custom_training_job_operator] + # [START how_to_cloud_vertex_ai_create_custom_training_job_v2_operator] + create_custom_training_job_v2 = CreateCustomTrainingJobOperator( + task_id="custom_task_v2", + staging_bucket=f"gs://{CUSTOM_GCS_BUCKET_NAME}", + display_name=CUSTOM_DISPLAY_NAME, + script_path=LOCAL_TRAINING_SCRIPT_PATH, + container_uri=CONTAINER_URI, + requirements=["gcsfs==0.7.1"], + model_serving_container_image_uri=MODEL_SERVING_CONTAINER_URI, + parent_model=model_id_v1, + # run params + dataset_id=tabular_dataset_id, + replica_count=REPLICA_COUNT, + model_display_name=MODEL_DISPLAY_NAME, + sync=False, + region=REGION, + project_id=PROJECT_ID, + ) + # [END how_to_cloud_vertex_ai_create_custom_training_job_v2_operator] + # [START how_to_cloud_vertex_ai_delete_custom_training_job_operator] delete_custom_training_job = DeleteCustomTrainingJobOperator( task_id="delete_custom_training_job", @@ -168,6 +189,7 @@ def TABULAR_DATASET(bucket_name): >> create_tabular_dataset # TEST BODY >> create_custom_training_job + >> create_custom_training_job_v2 # TEST TEARDOWN >> delete_custom_training_job >> delete_tabular_dataset diff --git a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_model_service.py b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_model_service.py index fccc3a610fded..e52d42edaa93d 100644 --- a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_model_service.py +++ b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_model_service.py @@ -46,9 +46,15 @@ DeleteDatasetOperator, ) from airflow.providers.google.cloud.operators.vertex_ai.model_service import ( + AddVersionAliasesOnModelOperator, DeleteModelOperator, + DeleteModelVersionOperator, + DeleteVersionAliasesOnModelOperator, ExportModelOperator, + GetModelOperator, ListModelsOperator, + ListModelVersionsOperator, + SetDefaultVersionOnModelOperator, UploadModelOperator, ) from airflow.providers.google.cloud.transfers.gcs_to_local import GCSToLocalFilesystemOperator @@ -161,6 +167,57 @@ region=REGION, project_id=PROJECT_ID, ) + model_id_v1 = create_custom_training_job.output["model_id"] + + create_custom_training_job_v2 = CreateCustomTrainingJobOperator( + task_id="custom_task_v2", + staging_bucket=f"gs://{DATA_SAMPLE_GCS_BUCKET_NAME}", + display_name=TRAIN_DISPLAY_NAME, + script_path=LOCAL_TRAINING_SCRIPT_PATH, + container_uri=CONTAINER_URI, + requirements=["gcsfs==0.7.1"], + model_serving_container_image_uri=MODEL_SERVING_CONTAINER_URI, + parent_model=model_id_v1, + # run params + dataset_id=tabular_dataset_id, + replica_count=1, + model_display_name=MODEL_DISPLAY_NAME, + sync=False, + region=REGION, + project_id=PROJECT_ID, + ) + model_id_v2 = create_custom_training_job_v2.output["model_id"] + + # [START how_to_cloud_vertex_ai_get_model_operator] + get_model = GetModelOperator( + task_id="get_model", region=REGION, project_id=PROJECT_ID, model_id=model_id_v1 + ) + # [END how_to_cloud_vertex_ai_get_model_operator] + + # [START how_to_cloud_vertex_ai_list_model_versions_operator] + list_model_versions = ListModelVersionsOperator( + task_id="list_model_versions", region=REGION, project_id=PROJECT_ID, model_id=model_id_v1 + ) + # [END how_to_cloud_vertex_ai_list_model_versions_operator] + + # [START how_to_cloud_vertex_ai_set_version_as_default_operator] + set_default_version = SetDefaultVersionOnModelOperator( + task_id="set_default_version", + project_id=PROJECT_ID, + region=REGION, + model_id=model_id_v2, + ) + # [END how_to_cloud_vertex_ai_set_version_as_default_operator] + + # [START how_to_cloud_vertex_ai_add_version_aliases_operator] + add_version_alias = AddVersionAliasesOnModelOperator( + task_id="add_version_alias", + project_id=PROJECT_ID, + region=REGION, + version_aliases=["new-version", "beta"], + model_id=model_id_v2, + ) + # [END how_to_cloud_vertex_ai_add_version_aliases_operator] # [START how_to_cloud_vertex_ai_upload_model_operator] upload_model = UploadModelOperator( @@ -199,6 +256,26 @@ ) # [END how_to_cloud_vertex_ai_list_models_operator] + # [START how_to_cloud_vertex_ai_delete_version_aliases_operator] + delete_version_alias = DeleteVersionAliasesOnModelOperator( + task_id="delete_version_alias", + project_id=PROJECT_ID, + region=REGION, + version_aliases=["new-version"], + model_id=model_id_v2, + ) + # [END how_to_cloud_vertex_ai_delete_version_aliases_operator] + + # [START how_to_cloud_vertex_ai_delete_version_operator] + delete_model_version = DeleteModelVersionOperator( + task_id="delete_model_version", + project_id=PROJECT_ID, + region=REGION, + model_id=model_id_v1, + trigger_rule=TriggerRule.ALL_DONE, + ) + # [END how_to_cloud_vertex_ai_delete_version_operator] + delete_custom_training_job = DeleteCustomTrainingJobOperator( task_id="delete_custom_training_job", training_pipeline_id="{{ task_instance.xcom_pull(task_ids='custom_task', key='training_id') }}", @@ -229,12 +306,19 @@ >> download_training_script_file >> create_tabular_dataset >> create_custom_training_job + >> create_custom_training_job_v2 # TEST BODY + >> get_model + >> list_model_versions + >> set_default_version + >> add_version_alias >> upload_model >> export_model >> delete_model >> list_models # TEST TEARDOWN + >> delete_version_alias + >> delete_model_version >> delete_custom_training_job >> delete_tabular_dataset >> delete_bucket