Join the SIG TFX-Addons community and help make TFX even better!

tfx.extensions.google_cloud_ai_platform.training_clients.UCAIPJobClient

Class for interacting with uCAIP CustomJob.

Inherits From: AbstractJobClient

Methods

create_client

View source

Creates the Gapic job client.

Can also be used for recreating the job client (e.g. in the case of communication failure).

Multiple job requests can be done in parallel if needed, by creating an instance of the class for each job. Note that one class instance should only be used for one job, as each instance stores variables (e.g. job_id) specific to each job.

create_training_args

View source

Get training args for runner._launch_aip_training.

The training args contain the inputs/outputs/exec_properties to the tfx.scripts.run_executor module.

Args
input_dict Passthrough input dict for tfx.components.Trainer.executor.
output_dict Passthrough input dict for tfx.components.Trainer.executor.
exec_properties Passthrough input dict for tfx.components.Trainer.executor.
executor_class_path class path for TFX core default trainer.
training_inputs Spec for CustomJob for AI Platform (Unified) custom training job. See https://cloud.google.com/ai-platform-unified/docs/reference/rest/v1/CustomJobSpec for the detailed schema.
job_id Display name for AI Platform (Unified) custom training job. If not supplied, system-determined unique ID is given. Refer to https://cloud.google.com/ai-platform-unified/docs/reference/rest/v1/projects.locations.customJobs

Returns
A dict containing the training arguments

get_job

View source

Gets the long-running job.

get_job_name

View source

Gets the job name.

get_job_state

View source

Gets the state of the long-running job.

Args
response The response from get_job

Returns
The job state.

launch_job

View source

Launches a long-running job.

Args
job_id The display name of the AI Platform (Unified) custom training job.
project The GCP project under which the training job will be executed.
training_input Spec for CustomJob for AI Platform (Unified) custom training job. See https://cloud.google.com/ai-platform-unified/docs/reference/rest/v1/CustomJobSpec for the detailed schema.
job_labels The dict of labels that will be attached to this job.

JOB_STATES_COMPLETED (<JobState.JOB_STATE_SUCCEEDED: 4>, <JobState.JOB_STATE_FAILED: 5>, <JobState.JOB_STATE_CANCELLED: 7>)
JOB_STATES_FAILED (<JobState.JOB_STATE_FAILED: 5>, <JobState.JOB_STATE_CANCELLED: 7>)