# -*- coding: utf-8 -*-

# Copyright 2023 Google LLC

#

# Licensed 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.

#

import datetime

import time

from typing import Dict, List, Literal, Optional, Sequence, Tuple, Union

from google.protobuf import json_format

import abc

from google.auth import credentials as auth_credentials

from google.cloud.aiplatform import base

from google.cloud.aiplatform.constants import base as constants

from google.cloud.aiplatform import datasets

from google.cloud.aiplatform import explain

from google.cloud.aiplatform import hyperparameter_tuning as hpt

from google.cloud.aiplatform import initializer

from google.cloud.aiplatform import models

from google.cloud.aiplatform import jobs

from google.cloud.aiplatform import schema

from google.cloud.aiplatform import utils

from google.cloud.aiplatform.utils import console_utils

from google.cloud.aiplatform.compat.types import env_var as gca_env_var

from google.cloud.aiplatform.compat.types import io as gca_io

from google.cloud.aiplatform.compat.types import model as gca_model

from google.cloud.aiplatform.compat.types import (

pipeline_state as gca_pipeline_state,

)

from google.cloud.aiplatform.compat.types import (

training_pipeline as gca_training_pipeline,

study as gca_study_compat,

custom_job as gca_custom_job_compat,

service_networking as gca_service_networking,

)

from google.cloud.aiplatform.utils import _timestamped_gcs_dir

from google.cloud.aiplatform.utils import source_utils

from google.cloud.aiplatform.utils import worker_spec_utils

from google.cloud.aiplatform.utils import column_transformations_utils

from google.cloud.aiplatform.utils import _explanation_utils

from google.cloud.aiplatform.v1.schema.trainingjob import (

definition_v1 as training_job_inputs,

)

from google.rpc import code_pb2

from google.rpc import status_pb2

import proto

_LOGGER = base.Logger(__name__)

_PIPELINE_COMPLETE_STATES = set(

[

gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED,

gca_pipeline_state.PipelineState.PIPELINE_STATE_FAILED,

gca_pipeline_state.PipelineState.PIPELINE_STATE_CANCELLED,

gca_pipeline_state.PipelineState.PIPELINE_STATE_PAUSED,

]

)

# _block_until_complete wait times

_JOB_WAIT_TIME = 5 # start at five seconds

_LOG_WAIT_TIME = 5

_MAX_WAIT_TIME = 60 * 5 # 5 minute wait

_WAIT_TIME_MULTIPLIER = 2 # scale wait by 2 every iteration

class _TrainingJob(base.VertexAiStatefulResource):

client_class = utils.PipelineClientWithOverride

_resource_noun = "trainingPipelines"

_getter_method = "get_training_pipeline"

_list_method = "list_training_pipelines"

_delete_method = "delete_training_pipeline"

_parse_resource_name_method = "parse_training_pipeline_path"

_format_resource_name_method = "training_pipeline_path"

# Required by the done() method

_valid_done_states = _PIPELINE_COMPLETE_STATES

def __init__(

self,

display_name: Optional[str] = None,

project: Optional[str] = None,

location: Optional[str] = None,

credentials: Optional[auth_credentials.Credentials] = None,

labels: Optional[Dict[str, str]] = None,

training_encryption_spec_key_name: Optional[str] = None,

model_encryption_spec_key_name: Optional[str] = None,

):

"""Constructs a Training Job.

Args:

display_name (str):

Optional. The user-defined name of the training pipeline. The

name must contain 128 or fewer UTF-8 characters.

project (str):

Optional. The name of a Google Cloud project from which to

retrieve the model. This overrides the project that was set by

`aiplatform.init`.

location (str):

Optional. The Google Cloud region where this from where the

model is retrieved. This region overrides the region that was

set by `aiplatform.init`.

credentials (auth_credentials.Credentials):

Optional. The credentials that are used to retrieve the model.

These credentials override the credentials set by

`aiplatform.init`.

labels (Dict[str, str]):

Optional. Labels with user-defined metadata to organize your

training pipeline. The maximum length of a key and value is 64

unicode characters. Labels and keys can contain only lowercase

letters, numeric characters, underscores, and dashes.

International characters are allowed. No more than 64 user

labels can be associated with one Tensorboard (system labels are

excluded). For more information and examples of using labels,

see [Using labels to organize Google Cloud Platform

resources](https://goo.gl/xmQnxf). System reserved label keys

are prefixed with `aiplatform.googleapis.com/` and are

immutable.

training_encryption_spec_key_name (Optional[str]):

Optional. The Cloud KMS resource identifier of the customer

managed encryption key that's used to protect the training

pipeline. The format of the key is

`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 `training_encryption_spec_key_name` is set, this training

pipeline is secured by this key. The model trained by this

training pipeline is also secured by this key if

`model_to_upload` isn't set.

This `training_encryption_spec_key_name` overrides the

`encryption_spec_key_name` set by `aiplatform.init`.

model_encryption_spec_key_name (Optional[str]):

Optional. The Cloud KMS resource identifier of the customer

managed encryption key that's used to protect the model. The

format of the key is

`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 `model_encryption_spec_key_name` is set, the traind model is

secured by this key.

This `model_encryption_spec_key_name` overrides the

`encryption_spec_key_name` set by `aiplatform.init`.

"""

if not display_name:

display_name = self.__class__._generate_display_name()

utils.validate_display_name(display_name)

if labels:

utils.validate_labels(labels)

super().__init__(project=project, location=location, credentials=credentials)

self._display_name = display_name

self._labels = labels

self._training_encryption_spec = initializer.global_config.get_encryption_spec(

encryption_spec_key_name=training_encryption_spec_key_name

)

self._model_encryption_spec = initializer.global_config.get_encryption_spec(

encryption_spec_key_name=model_encryption_spec_key_name

)

self._gca_resource = None

@property

@classmethod

@abc.abstractmethod

def _supported_training_schemas(cls) -> Tuple[str]:

"""The list of supported schemas for this training job."""

pass

@property

def start_time(self) -> Optional[datetime.datetime]:

"""Optional. The time when the training job first entered the

`PIPELINE_STATE_RUNNING` state."""

self._sync_gca_resource()

return getattr(self._gca_resource, "start_time")

@property

def end_time(self) -> Optional[datetime.datetime]:

"""Optional. The time when the training job entered the

`PIPELINE_STATE_SUCCEEDED`, `PIPELINE_STATE_FAILED`, or

`PIPELINE_STATE_CANCELLED` state."""

self._sync_gca_resource()

return getattr(self._gca_resource, "end_time")

@property

def error(self) -> Optional[status_pb2.Status]:

"""Optional. Detailed error information for this training job resource.

Error information is created only when the state of the training job is

`PIPELINE_STATE_FAILED` or `PIPELINE_STATE_CANCELLED`."""

self._sync_gca_resource()

return getattr(self._gca_resource, "error")

@classmethod

def get(

cls,

resource_name: str,

project: Optional[str] = None,

location: Optional[str] = None,

credentials: Optional[auth_credentials.Credentials] = None,

) -> "_TrainingJob":

"""Gets a training job using the `resource_name` that's passed in.

Args:

resource_name (str):

Required. A fully-qualified resource name or ID.

project (str):

Optional. The name of the Google Cloud project to retrieve the

training job from. This overrides the project that was set by

`aiplatform.init`.

location (str):

Optional. The Google Cloud region from where the training job is

retrieved. This region overrides the region that was set by

`aiplatform.init`.

credentials (auth_credentials.Credentials):

Optional. The credentials that are used to upload this model.

These credentials override the credentials set by

`aiplatform.init`.

Raises:

ValueError: A `ValueError` is raised if the task definition of the

retrieved training job doesn't match the custom training task

definition.

Returns:

A Vertex AI training job.

"""

# Create job with dummy parameters

# These parameters won't be used as user can not run the job again.

# If they try, an exception will be raised.

self = cls._empty_constructor(

project=project,

location=location,

credentials=credentials,

resource_name=resource_name,

)

self._gca_resource = self._get_gca_resource(resource_name=resource_name)

if (

self._gca_resource.training_task_definition

not in cls._supported_training_schemas

):

raise ValueError(

f"The task definition of the retrieved training job "

f"is {self._gca_resource.training_task_definition}, "

f"which is not compatible with {cls.__name__}."

)

return self

@classmethod

def _get_and_return_subclass(

cls,

resource_name: str,

project: Optional[str] = None,

location: Optional[str] = None,

credentials: Optional[auth_credentials.Credentials] = None,

) -> "_TrainingJob":

"""Retrieves a training job subclass for the `resource_name` that's

passed in without knowing the `training_task_definition`.

Example usage:

```

aiplatform.training_jobs._TrainingJob._get_and_return_subclass(

'projects/.../locations/.../trainingPipelines/12345'

)

# Returns: <google.cloud.aiplatform.training_jobs.AutoMLImageTrainingJob>

```

Args:

resource_name (str):

Required. A fully-qualified resource name or ID.

project (str):

Optional. The name of the Google Cloud project to retrieve the

training job from. This overrides the project that was set by

`aiplatform.init`.

location (str):

Optional. The Google Cloud region from where the training job is

retrieved. This region overrides the region that was set by

`aiplatform.init`.

credentials (auth_credentials.Credentials):

Optional. The credentials that are used to upload this model.

These credentials override the credentials set by

`aiplatform.init`.

Returns:

A Vertex AI training job.

"""

# Retrieve training pipeline resource before class construction

client = cls._instantiate_client(location=location, credentials=credentials)

gca_training_pipeline = getattr(client, cls._getter_method)(name=resource_name)

schema_uri = gca_training_pipeline.training_task_definition

# Collect all AutoML training job classes and CustomTrainingJob

class_list = [

c for c in cls.__subclasses__() if c.__name__.startswith("AutoML")

] + [CustomTrainingJob]

# Identify correct training job subclass, construct and return object

for c in class_list:

if schema_uri in c._supported_training_schemas:

return c._empty_constructor(

project=project,

location=location,

credentials=credentials,

resource_name=resource_name,

)

@property

@abc.abstractmethod

def _model_upload_fail_string(self) -> str:

"""Helper property for model upload failure."""

pass

@abc.abstractmethod

def run(self) -> Optional[models.Model]:

"""Runs the training job.

Should call _run_job internally

"""

pass

@staticmethod

def _create_input_data_config(

dataset: Optional[datasets._Dataset] = None,

annotation_schema_uri: Optional[str] = None,

training_fraction_split: Optional[float] = None,

validation_fraction_split: Optional[float] = None,

test_fraction_split: Optional[float] = None,

training_filter_split: Optional[str] = None,

validation_filter_split: Optional[str] = None,

test_filter_split: Optional[str] = None,

predefined_split_column_name: Optional[str] = None,

timestamp_split_column_name: Optional[str] = None,

gcs_destination_uri_prefix: Optional[str] = None,

bigquery_destination: Optional[str] = None,

) -> Optional[gca_training_pipeline.InputDataConfig]:

"""Constructs an input data config to pass to the training pipeline.

Args:

dataset (datasets._Dataset):

The dataset within the same Project from which data will be used to train the Model. The

Dataset must use schema compatible with Model being trained,

and what is compatible should be described in the used

TrainingPipeline's [training_task_definition]

[google.cloud.aiplatform.v1beta1.TrainingPipeline.training_task_definition].

For tabular Datasets, all their data is exported to

training, to pick and choose from.

annotation_schema_uri (str):

Google Cloud Storage URI points to a YAML file describing

annotation schema. The schema is defined as an OpenAPI 3.0.2

[Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schema-object) The schema files

that can be used here are found in

gs://google-cloud-aiplatform/schema/dataset/annotation/,

note that the chosen schema must be consistent with

``metadata``

of the Dataset specified by

``dataset_id``.

Only Annotations that both match this schema and belong to

DataItems not ignored by the split method are used in

respectively training, validation or test role, depending on

the role of the DataItem they are on.

When used in conjunction with

``annotations_filter``,

the Annotations used for training are filtered by both

``annotations_filter``

and

``annotation_schema_uri``.

training_fraction_split (float):

Optional. The fraction of the input data that is to be used to train

the Model. This is ignored if Dataset is not provided.

validation_fraction_split (float):

Optional. The fraction of the input data that is to be used to validate

the Model. This is ignored if Dataset is not provided.

test_fraction_split (float):

Optional. The fraction of the input data that is to be used to evaluate

the Model. This is ignored if Dataset is not provided.

training_filter_split (str):

Optional. A filter on DataItems of the Dataset. DataItems that match

this filter are used to train the Model. A filter with same syntax

as the one used in DatasetService.ListDataItems may be used. If a

single DataItem is matched by more than one of the FilterSplit filters,

then it is assigned to the first set that applies to it in the training,

validation, test order. This is ignored if Dataset is not provided.

validation_filter_split (str):

Optional. A filter on DataItems of the Dataset. DataItems that match

this filter are used to validate the Model. A filter with same syntax

as the one used in DatasetService.ListDataItems may be used. If a

single DataItem is matched by more than one of the FilterSplit filters,

then it is assigned to the first set that applies to it in the training,

validation, test order. This is ignored if Dataset is not provided.

test_filter_split (str):

Optional. A filter on DataItems of the Dataset. DataItems that match

this filter are used to test the Model. A filter with same syntax

as the one used in DatasetService.ListDataItems may be used. If a

single DataItem is matched by more than one of the FilterSplit filters,

then it is assigned to the first set that applies to it in the training,

validation, test order. This is ignored if Dataset is not provided.

predefined_split_column_name (str):

Optional. The key is a name of one of the Dataset's data

columns. The value of the key (either the label's value or

value in the column) must be one of {``training``,

``validation``, ``test``}, and it defines to which set the

given piece of data is assigned. If for a piece of data the

key is not present or has an invalid value, that piece is

ignored by the pipeline.

Supported only for tabular and time series Datasets.

timestamp_split_column_name (str):

Optional. The key is a name of one of the Dataset's data

columns. The value of the key values of the key (the values in

the column) must be in RFC 3339 `date-time` format, where

`time-offset` = `"Z"` (e.g. 1985-04-12T23:20:50.52Z). If for a

piece of data the key is not present or has an invalid value,

that piece is ignored by the pipeline.

Supported only for tabular and time series Datasets.

This parameter must be used with training_fraction_split,

validation_fraction_split, and test_fraction_split.

gcs_destination_uri_prefix (str):

Optional. The Google Cloud Storage location.

The Vertex AI environment variables representing Google

Cloud Storage data URIs will always be represented in the

Google Cloud Storage wildcard format to support sharded

data.

- AIP_DATA_FORMAT = "jsonl".

- AIP_TRAINING_DATA_URI = "gcs_destination/training-*"

- AIP_VALIDATION_DATA_URI = "gcs_destination/validation-*"

- AIP_TEST_DATA_URI = "gcs_destination/test-*".

bigquery_destination (str):

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"

Raises:

ValueError: When more than 1 type of split configuration is passed or when

the split configuration passed is incompatible with the dataset schema.

"""

input_data_config = None

if dataset:

# Initialize all possible splits

filter_split = None

predefined_split = None

timestamp_split = None

fraction_split = None

# Create filter split

if any(

[

training_filter_split is not None,

validation_filter_split is not None,

test_filter_split is not None,

]

):

if all(

[

training_filter_split is not None,

validation_filter_split is not None,

test_filter_split is not None,

]

):

filter_split = gca_training_pipeline.FilterSplit(

training_filter=training_filter_split,

validation_filter=validation_filter_split,

test_filter=test_filter_split,

)

else:

raise ValueError(

"All filter splits must be passed together or not at all"

)

# Create predefined split

if predefined_split_column_name:

predefined_split = gca_training_pipeline.PredefinedSplit(

key=predefined_split_column_name

)

# Create timestamp split or fraction split

if timestamp_split_column_name:

timestamp_split = gca_training_pipeline.TimestampSplit(

training_fraction=training_fraction_split,

validation_fraction=validation_fraction_split,

test_fraction=test_fraction_split,

key=timestamp_split_column_name,

)

elif any(

[

training_fraction_split is not None,

validation_fraction_split is not None,

test_fraction_split is not None,

]

):

fraction_split = gca_training_pipeline.FractionSplit(

training_fraction=training_fraction_split,

validation_fraction=validation_fraction_split,

test_fraction=test_fraction_split,

)

splits = [

split

for split in [

filter_split,

predefined_split,

timestamp_split_column_name,

fraction_split,

]

if split is not None

]

# Fallback to fraction split if nothing else is specified

if len(splits) == 0:

_LOGGER.info(

"No dataset split provided. The service will use a default split."

)

elif len(splits) > 1:

raise ValueError(

"You've specified too many split types. You can specify"

" only one of the following:"

" 1. `training_filter_split`, `validation_filter_split`,"

" `test_filter_split`"

" 2. `predefined_split_column_name`"

" 3. `timestamp_split_column_name`,"

" `training_fraction_split`, `validation_fraction_split`,"

" `test_fraction_split`"

" 4.`training_fraction_split`,"

" `validation_fraction_split`, `test_fraction_split`"

)

# create GCS destination

gcs_destination = None

if gcs_destination_uri_prefix:

gcs_destination = gca_io.GcsDestination(

output_uri_prefix=gcs_destination_uri_prefix

)

# TODO(b/177416223) validate managed BQ dataset is passed in

bigquery_destination_proto = None

if bigquery_destination:

bigquery_destination_proto = gca_io.BigQueryDestination(

output_uri=bigquery_destination

)

# create input data config

input_data_config = gca_training_pipeline.InputDataConfig(

fraction_split=fraction_split,

filter_split=filter_split,

predefined_split=predefined_split,

timestamp_split=timestamp_split,

dataset_id=dataset.name,

annotation_schema_uri=annotation_schema_uri,

gcs_destination=gcs_destination,

bigquery_destination=bigquery_destination_proto,

)

return input_data_config

def _run_job(

self,

training_task_definition: str,

training_task_inputs: Union[dict, proto.Message],

dataset: Optional[datasets._Dataset],

training_fraction_split: Optional[float] = None,

validation_fraction_split: Optional[float] = None,

test_fraction_split: Optional[float] = None,

training_filter_split: Optional[str] = None,

validation_filter_split: Optional[str] = None,

test_filter_split: Optional[str] = None,

predefined_split_column_name: Optional[str] = None,

timestamp_split_column_name: Optional[str] = None,

annotation_schema_uri: Optional[str] = None,

model: Optional[gca_model.Model] = None,

model_id: Optional[str] = None,

parent_model: Optional[str] = None,

is_default_version: Optional[bool] = True,

model_version_aliases: Optional[Sequence[str]] = None,

model_version_description: Optional[str] = None,

gcs_destination_uri_prefix: Optional[str] = None,

bigquery_destination: Optional[str] = None,

create_request_timeout: Optional[float] = None,

block: Optional[bool] = True,

) -> Optional[models.Model]:

"""Runs the training job.

Args:

training_task_definition (str):

Required. A Google Cloud Storage path to the

YAML file that defines the training task which

is responsible for producing the model artifact,

and may also include additional auxiliary work.

The definition files that can be used here are

found in gs://google-cloud-

aiplatform/schema/trainingjob/definition/. 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.

training_task_inputs (Union[dict, proto.Message]):

Required. The training task's input that corresponds to the training_task_definition parameter.

dataset (datasets._Dataset):

The dataset within the same Project from which data will be used to train the Model. The

Dataset must use schema compatible with Model being trained,

and what is compatible should be described in the used

TrainingPipeline's [training_task_definition]

[google.cloud.aiplatform.v1beta1.TrainingPipeline.training_task_definition].

For tabular Datasets, all their data is exported to

training, to pick and choose from.

annotation_schema_uri (str):

Google Cloud Storage URI points to a YAML file describing

annotation schema. The schema is defined as an OpenAPI 3.0.2

[Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schema-object) The schema files

that can be used here are found in

gs://google-cloud-aiplatform/schema/dataset/annotation/,

note that the chosen schema must be consistent with

``metadata``

of the Dataset specified by

``dataset_id``.

Only Annotations that both match this schema and belong to

DataItems not ignored by the split method are used in

respectively training, validation or test role, depending on

the role of the DataItem they are on.

When used in conjunction with

``annotations_filter``,

the Annotations used for training are filtered by both

``annotations_filter``

and

``annotation_schema_uri``.

training_fraction_split (float):

Optional. The fraction of the input data that is to be used to train

the Model. This is ignored if Dataset is not provided.

validation_fraction_split (float):

Optional. The fraction of the input data that is to be used to validate

the Model. This is ignored if Dataset is not provided.

test_fraction_split (float):

Optional. The fraction of the input data that is to be used to evaluate

the Model. This is ignored if Dataset is not provided.

training_filter_split (str):

Optional. A filter on DataItems of the Dataset. DataItems that match

this filter are used to train the Model. A filter with same syntax

as the one used in DatasetService.ListDataItems may be used. If a

single DataItem is matched by more than one of the FilterSplit filters,

then it is assigned to the first set that applies to it in the training,

validation, test order. This is ignored if Dataset is not provided.

validation_filter_split (str):

Optional. A filter on DataItems of the Dataset. DataItems that match

this filter are used to validate the Model. A filter with same syntax

as the one used in DatasetService.ListDataItems may be used. If a

single DataItem is matched by more than one of the FilterSplit filters,

then it is assigned to the first set that applies to it in the training,

validation, test order. This is ignored if Dataset is not provided.

test_filter_split (str):

Optional. A filter on DataItems of the Dataset. DataItems that match

this filter are used to test the Model. A filter with same syntax

as the one used in DatasetService.ListDataItems may be used. If a

single DataItem is matched by more than one of the FilterSplit filters,

then it is assigned to the first set that applies to it in the training,

validation, test order. This is ignored if Dataset is not provided.

predefined_split_column_name (str):

Optional. The key is a name of one of the Dataset's data

columns. The value of the key (either the label's value or

value in the column) must be one of {``training``,

``validation``, ``test``}, and it defines to which set the

given piece of data is assigned. If for a piece of data the

key is not present or has an invalid value, that piece is

ignored by the pipeline.

Supported only for tabular and time series Datasets.

timestamp_split_column_name (str):

Optional. The key is a name of one of the Dataset's data

columns. The value of the key values of the key (the values in

the column) must be in RFC 3339 `date-time` format, where

`time-offset` = `"Z"` (e.g. 1985-04-12T23:20:50.52Z). If for a

piece of data the key is not present or has an invalid value,

that piece is ignored by the pipeline.

Supported only for tabular and time series Datasets.

This parameter must be used with training_fraction_split,

validation_fraction_split, and test_fraction_split.

model (~.model.Model):

Optional. Describes the Model that may be uploaded (via

[ModelService.UploadMode][]) by this TrainingPipeline. The

TrainingPipeline's

``training_task_definition``

should make clear whether this Model description should be

populated, and if there are any special requirements

regarding how it should be filled. If nothing is mentioned

in the

``training_task_definition``,

then it should be assumed that this field should not be

filled and the training task either uploads the Model

without a need of this information, or that training task

does not support uploading a Model as part of the pipeline.

When the Pipeline's state becomes

``PIPELINE_STATE_SUCCEEDED`` and the trained Model had been

uploaded into Vertex AI, then the model_to_upload's

resource ``name``

is populated. The Model is always uploaded into the Project

and Location in which this pipeline is.

model_id (str):

Optional. The ID to use for the Model produced by this job,

which will become the final component of the model resource name.

This value may be up to 63 characters, and valid characters

are `[a-z0-9_-]`. The first character cannot be a number or hyphen.

parent_model (str):

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.

is_default_version (bool):

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.

model_version_aliases (Sequence[str]):

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]

model_version_description (str):

Optional. The description of the model version being uploaded by this job.

gcs_destination_uri_prefix (str):

Optional. The Google Cloud Storage location.

The Vertex AI environment variables representing Google

Cloud Storage data URIs will always be represented in the

Google Cloud Storage wildcard format to support sharded

data.

- AIP_DATA_FORMAT = "jsonl".

- AIP_TRAINING_DATA_URI = "gcs_destination/training-*"

- AIP_VALIDATION_DATA_URI = "gcs_destination/validation-*"

- AIP_TEST_DATA_URI = "gcs_destination/test-*".

bigquery_destination (str):

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"

create_request_timeout (float):

Optional. The timeout for the create request in seconds.

block (bool):

Optional. If True, block until complete.

"""

input_data_config = self._create_input_data_config(

dataset=dataset,

annotation_schema_uri=annotation_schema_uri,

training_fraction_split=training_fraction_split,

validation_fraction_split=validation_fraction_split,

test_fraction_split=test_fraction_split,

training_filter_split=training_filter_split,

validation_filter_split=validation_filter_split,

test_filter_split=test_filter_split,

predefined_split_column_name=predefined_split_column_name,

timestamp_split_column_name=timestamp_split_column_name,

gcs_destination_uri_prefix=gcs_destination_uri_prefix,

bigquery_destination=bigquery_destination,

)

parent_model = models.ModelRegistry._get_true_version_parent(

parent_model=parent_model,

project=self.project,

location=self.location,

)

if model:

model.version_aliases = models.ModelRegistry._get_true_alias_list(

version_aliases=model_version_aliases,

is_default_version=is_default_version,

)

model.version_description = model_version_description

# create training pipeline

training_pipeline = gca_training_pipeline.TrainingPipeline(

display_name=self._display_name,

training_task_definition=training_task_definition,

training_task_inputs=training_task_inputs,

model_to_upload=model,

model_id=model_id,

parent_model=parent_model,

input_data_config=input_data_config,

labels=self._labels,

encryption_spec=self._training_encryption_spec,

)

training_pipeline = self.api_client.create_training_pipeline(

parent=initializer.global_config.common_location_path(

self.project, self.location

),

training_pipeline=training_pipeline,

timeout=create_request_timeout,

)

self._gca_resource = training_pipeline

_LOGGER.info("View Training:\n%s" % self._dashboard_uri())

model = self._get_model(block=block)

if model is None:

_LOGGER.warning(

"Training did not produce a Managed Model returning None. "

+ self._model_upload_fail_string

)

return model

def _is_waiting_to_run(self) -> bool:

"""Returns `true` if the training job is pending upstream tasks,

otherwise it returns `false`."""

self._raise_future_exception()

if self._latest_future:

_LOGGER.info(

"The training job is waiting for upstream SDK tasks to complete before"

" launching."

)

return True

return False

@property

def state(self) -> Optional[gca_pipeline_state.PipelineState]:

"""Current training state."""

if self._assert_has_run():

return

self._sync_gca_resource()

return self._gca_resource.state

def get_model(self, sync=True) -> models.Model:

"""Returns the Vertex AI model produced by this training job.

Args:

sync (bool):

If set to `true`, this method runs synchronously. If `false`, this

method runs asynchronously.

Returns:

The Vertex AI model that was produced by this training job.

Raises:

RuntimeError: A runtime error is raised if the training job failed

or if a model wasn't produced by the training job.

"""

self._assert_has_run()

if not self._gca_resource.model_to_upload:

raise RuntimeError(self._model_upload_fail_string)

return self._force_get_model(sync=sync)

@base.optional_sync()

def _force_get_model(self, sync: bool = True) -> models.Model:

"""Returns the Vertex AI model produced by this training job.

Args:

sync (bool):

If set to `true`, this method runs synchronously. If `false`, this

method runs asynchronously.

Returns:

The Vertex AI model that was produced by this training job.

Raises:

RuntimeError: If training failed or if a model was not produced by this training.

"""

model = self._get_model()

if model is None:

raise RuntimeError(self._model_upload_fail_string)

return model

def _get_model(self, block: bool = True) -> Optional[models.Model]:

"""Helper method to get and instantiate the Model to Upload.

Returns:

model: Vertex AI Model if training succeeded and produced a Vertex AI

Model. None otherwise.

Raises:

RuntimeError: If Training failed.

"""

if block:

self._block_until_complete()

if self.has_failed:

raise RuntimeError(

f"Training Pipeline {self.resource_name} failed. No model available."

)

if not self._gca_resource.model_to_upload:

return None

if self._gca_resource.model_to_upload.name:

return models.Model(

model_name=self._gca_resource.model_to_upload.name,

version=self._gca_resource.model_to_upload.version_id,

)

def _wait_callback(self):

"""Callback performs custom logging during _block_until_complete. Override in subclass."""

pass

def _block_until_complete(self):

"""Helper method to block and check on job until complete."""

log_wait = _LOG_WAIT_TIME

previous_time = time.time()

while self.state not in _PIPELINE_COMPLETE_STATES:

current_time = time.time()

if current_time - previous_time >= log_wait:

_LOGGER.info(

"%s %s current state:\n%s"

% (

self.__class__.__name__,

self._gca_resource.name,

self._gca_resource.state,

)

)

log_wait = min(log_wait * _WAIT_TIME_MULTIPLIER, _MAX_WAIT_TIME)

previous_time = current_time

self._wait_callback()

time.sleep(_JOB_WAIT_TIME)

self._raise_failure()

_LOGGER.log_action_completed_against_resource("run", "completed", self)

if self._gca_resource.model_to_upload and not self.has_failed:

_LOGGER.info(

"Model available at %s" % self._gca_resource.model_to_upload.name

)

def _raise_failure(self):

"""Helper method to raise failure if TrainingPipeline fails.

Raises:

RuntimeError: If training failed.

"""