# Copyright 2021 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. """A utility to submit a Vertex Training T5X job.""" import gcsfs import fsspec import os from absl import flags from absl import app from absl import logging from datetime import datetime from importlib import import_module from google.cloud import aiplatform as vertex_ai from google.cloud.aiplatform import CustomJob from typing import List from typing import Dict from typing import Any from typing import Union from typing import Optional MACHINE_TYPE = 'cloud-tpu' flags.DEFINE_string('project_id', None, 'GCP Project') flags.DEFINE_string('region', None, 'Vertex Pipelines region') flags.DEFINE_string('staging_bucket', None, 'Staging bucket') flags.DEFINE_string('training_sa', None, 'Training SA') flags.DEFINE_string('image_uri', None, 'Training image') flags.DEFINE_string('job_name_prefix', 't5x_job', 'Job name prefix') flags.DEFINE_list('gin_files', None, 'Gin configuration files') flags.DEFINE_list('gin_overwrites', None, 'Gin overwrites') flags.DEFINE_list('gin_search_paths', None, 'Gin search paths') flags.DEFINE_string('accelerator_type', 'TPU_V2', 'Accelerator type') flags.DEFINE_integer('accelerator_count', 8, 'Number of cores') flags.DEFINE_string('run_mode', 'train', 'Run mode') flags.DEFINE_string('tfds_data_dir', None, 'TFDS data dir') flags.DEFINE_bool('sync', True, 'Execute synchronously') flags.mark_flag_as_required('project_id') flags.mark_flag_as_required('region') flags.mark_flag_as_required('staging_bucket') flags.mark_flag_as_required('gin_files') flags.mark_flag_as_required('image_uri') FLAGS = flags.FLAGS def _create_t5x_custom_job( display_name: str, machine_type: str, accelerator_type: str, accelerator_count: int, image_uri: str, run_mode: str, gin_files: List[str], model_dir: str, gin_search_paths: Optional[List[str]] = None, tfds_data_dir: Optional[str] = None, replica_count: int = 1, gin_overwrites: Optional[List[str]] = None, base_output_dir: Optional[str] = None, ) -> CustomJob: """Creates a Vertex AI custom T5X training job. It copies the configuration files (.gin) to GCS, creates a worker_pool_spec and returns an aiplatform.CustomJob. Args: display_name (str): Required. User defined display name for the Vertex AI custom T5X job. machine_type (str): Required. The type of machine for running the custom training job on dedicated resources. For TPUs, use `cloud-tpu`. accelerator_type (str): Required. The type of accelerator(s) that may be attached to the machine as per `accelerator_count`. Only used if `machine_type` is set. Options: `TPU_V2` or `TPU_V3`. accelerator_count (int): Required. The number of accelerators to attach to the `machine_type`. Only used if `machine_type` is set. For TPUs, this is the number of cores to be provisioned. Example: 8, 128, 512, etc. image_uri (str): Required. Full image path to be used as the execution environment of the custom T5X training job. Example: 'gcr.io/{PROJECT_ID}/{IMAGE_NAME}' run_mode (str): Required. The mode to run T5X under. Options: `train`, `eval`, `infer`. gin_files (List[str]): Required. Full path to gin configuration file on local filesystem. Multiple paths may be passed and will be imported in the given order, with later configurations overriding earlier ones. gin_search_paths (List[str]): List of gin config path prefixes to be prepended to gin suffixes in gin includes and gin_files model_dir (str): Required. Path on Google Cloud Storage to store all the artifacts generated by the custom T5X training job. The path must be in this format: `gs://{bucket name}/{your folder}/...`. Example: gs://my_bucket/experiments/model1/ tfds_data_dir (Optional[str] = None): Optional. If set, this directory will be used to store datasets prepared by TensorFlow Datasets that are not available in the public TFDS GCS bucket. Note that this flag overrides the `tfds_data_dir` attribute of all Task`s. This path must be a valid GCS path. Example: gs://my_bucket/datasets/my_dataset/ replica_count (int = 1): Optional. The number of worker replicas. If replica count = 1 then one chief replica will be provisioned. For TPUs this must be set to 1. gin_overwrites (Optional[List[str]] = None): Optional. List of arguments to overwrite gin configurations. Argument must be enclosed in parentheses. Example: --gin.TRAIN_PATH=\"gs://my_bucket/folder\" base_output_dir (Optional[str] = None): Returns: (aiplatform.CustomJob): Return an instance of a Vertex AI training CustomJob. """ local_fs = fsspec.filesystem('file') gcs_fs = gcsfs.GCSFileSystem() # Check if gin files exists if not gin_files or not all([local_fs.isfile(f) for f in gin_files]): raise FileNotFoundError( 'Provide a list of valid gin files.' ) # Try to copy files to GCS bucket try: gcs_gin_files = [] for gin_file in gin_files: gcs_path = os.path.join(model_dir, gin_file.split(sep='/')[-1]) gcs_fs.put(gin_file, gcs_path) gcs_gin_files.append(gcs_path.replace('gs://', '/gcs/')) except: raise RuntimeError('Could not copy gin files to GCS.') container_spec = {"image_uri": image_uri} # Temporary mitigation to address issues with t5x/main.py # and inference on tfrecord files if run_mode == 'infer': args = [ f'--gin.MODEL_DIR="{model_dir}"', f'--tfds_data_dir={tfds_data_dir}', ] container_spec['command'] = ["python", "./t5x/t5x/infer.py"] else: args = [ f'--run_mode={run_mode}', f'--gin.MODEL_DIR="{model_dir}"', f'--tfds_data_dir={tfds_data_dir}', ] if gin_search_paths: args.append(f'--gin_search_paths={",".join(gin_search_paths)}') args += [f'--gin_file={gcs_path}' for gcs_path in gcs_gin_files] if gin_overwrites: args += [f'--gin.{overwrite}' for overwrite in gin_overwrites] container_spec["args"] = args worker_pool_specs = [ { "machine_spec": { "machine_type": machine_type, "accelerator_type": accelerator_type, "accelerator_count": accelerator_count, }, "replica_count": replica_count, "container_spec": container_spec, } ] job = vertex_ai.CustomJob( display_name=display_name, worker_pool_specs=worker_pool_specs, base_output_dir=base_output_dir ) return job def _main(argv): vertex_ai.init( project=FLAGS.project_id, location=FLAGS.region, staging_bucket=FLAGS.staging_bucket, ) job_name = f'{FLAGS.job_name_prefix}_{datetime.now().strftime("%Y%m%d%H%M%S")}' job_dir = f'{FLAGS.staging_bucket}/t5x_jobs/{job_name}' job = _create_t5x_custom_job( display_name=job_name, machine_type=MACHINE_TYPE, accelerator_type=FLAGS.accelerator_type, accelerator_count=FLAGS.accelerator_count, image_uri=FLAGS.image_uri, run_mode=FLAGS.run_mode, gin_files=FLAGS.gin_files, gin_overwrites=FLAGS.gin_overwrites, gin_search_paths=FLAGS.gin_search_paths, tfds_data_dir=FLAGS.tfds_data_dir, model_dir=job_dir, ) logging.info(f'Starting job: {job_name}') job.run(sync=FLAGS.sync) if __name__ == "__main__": app.run(_main)