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

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

@@ -61,6 +61,10 @@ class CustomTrainingJobBaseOperator(GoogleCloudBaseOperator):
         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 @@ class CustomTrainingJobBaseOperator(GoogleCloudBaseOperator):
         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 <https://tinyurl.com/y538mdwt#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 <https://tinyurl.com/y538mdwt#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 <https://tinyurl.com/y538mdwt#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 <https://tinyurl.com/y538mdwt#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 <https://tinyurl.com/y538mdwt#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 <https://tinyurl.com/y538mdwt#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 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
             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 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
             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 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
         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 <https://tinyurl.com/y538mdwt#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 <https://tinyurl.com/y538mdwt#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 <https://tinyurl.com/y538mdwt#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 <https://tinyurl.com/y538mdwt#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 <https://tinyurl.com/y538mdwt#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 <https://tinyurl.com/y538mdwt#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. <base_output_dir>/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_<dataset-id>_<annotation-type>_<timestamp-of-training-call>``
-            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_<dataset-id>_<annotation-type>_<timestamp-of-training-call>``
+        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 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator
             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 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator
             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 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator
         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 <https://tinyurl.com/y538mdwt#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 <https://tinyurl.com/y538mdwt#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 <https://tinyurl.com/y538mdwt#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 <https://tinyurl.com/y538mdwt#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 <https://tinyurl.com/y538mdwt#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 <https://tinyurl.com/y538mdwt#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 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
             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 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
             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 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
         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 @@ class DeleteCustomTrainingJobOperator(GoogleCloudBaseOperator):
 
 
 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.