vertex_nas_cli.py

# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
r"""Vertex NAS CLI.

To build containers, run the following:

python3 vertex_nas_cli.py build \
--project_id=${PROJECT_ID} \
--trainer_docker_id=${TRAINER_DOCKER_ID} \
--latency_calculator_docker_id=${LATENCY_CALCULATOR_DOCKER_ID} \
--trainer_docker_file=tf_vision/nas_multi_trial.Dockerfile \
--latency_calculator_docker_file=tf_vision/latency_computation_using_saved_model.Dockerfile

To run search locally with prebuilt trainer / search space, run the following:
(the flags to the container are passed in after `--`):

python3 vertex_nas_cli.py search_in_local \
--project_id=${PROJECT_ID} \
--trainer_docker_id=${TRAINER_DOCKER_ID} \
--prebuilt_search_space=<> \
--search_docker_flags \
flag_1_to_trainer_docker=<>
flag_2_to_trainer_docker=<>


Example 1: Run the mnist search codelab:
https://github.com/google/vertex-ai-nas/blob/main/third_party/tutorial/vertex_tutorial3.md


Example 2: Run with NAS prebuilt search-space under constraint:

# Build the training container:
python3 vertex_nas_cli.py build --project_id=${PROJECT_ID} \
--trainer_docker_id=${TRAINER_DOCKER_ID} \
--trainer_docker_file=tf_vision/nas_multi_trial.Dockerfile \
--latency_calculator_docker_id=${LATENCY_CALCULATOR_DOCKER_ID}

# Run in cloud (with use_prebuilt_trainer=True, the pass-in flags will be
validated):
python3 vertex_nas_cli.py search --project_id=${PROJECT_ID} \
--trainer_docker_id=${TRAINER_DOCKER_ID} \
--latency_calculator_docker_id=${LATENCY_CALCULATOR_DOCKER_ID} \
--use_prebuilt_trainer=True \
--prebuilt_search_space=nasfpn \
--root_output_dir=${GCS_ROOT_DIR} \
--nas_target_reward_metric=AP \
--job_id=cli_test \
--search_docker_flags \
params_override=configs/nasfpn_search.yaml \
training_data_path=${TRAINING_DATA_PATH} \
validation_data_path=${VALIDATION_DATA_PATH} \
target_device_latency_ms=85 \
target_memory_mb=5000 \
model=retinanet


Example 3: Run with NAS prebuilt search-space under constraint and with stage-2
training as well:

# Build the training container:
python3 vertex_nas_cli.py build --project_id=${PROJECT_ID} \
--trainer_docker_id=${TRAINER_DOCKER_ID} \
--trainer_docker_file=tf_vision/nas_multi_trial.Dockerfile \
--latency_calculator_docker_id=${LATENCY_CALCULATOR_DOCKER_ID}

# Run in cloud (with use_prebuilt_trainer=True, the pass-in flags will be
validated):
python3 vertex_nas_cli.py search --project_id=${PROJECT_ID} \
--trainer_docker_id=${TRAINER_DOCKER_ID} \
--latency_calculator_docker_id=${LATENCY_CALCULATOR_DOCKER_ID} \
--use_prebuilt_trainer=True \
--prebuilt_search_space=nasfpn \
--root_output_dir=${GCS_ROOT_DIR} \
--nas_target_reward_metric=AP \
--job_id=cli_test \
--search_docker_flags \
params_override=configs/nasfpn_search.yaml \
training_data_path=${TRAINING_DATA_PATH} \
validation_data_path=${VALIDATION_DATA_PATH} \
target_device_latency_ms=85 \
target_memory_mb=5000 \
model=retinanet \
--train_docker_flags \
params_override=configs/nasfpn_search_train.yaml \
training_data_path=${TRAINING_DATA_PATH} \
validation_data_path=${VALIDATION_DATA_PATH} \
model=retinanet \
--train_max_parallel_trial=2 \
--train_frequency=2


Example 4: List trial details from a search job:

python3 vertex_nas_cli.py list_trials \
--project_id=${PROJECT_ID} \
--job_id=<numeric job id generated by Vertex AI> \
--trials_output_file=${OUTPUT_FILE_PATH}
--max_trials=5
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import copy
import datetime
import importlib
import json
import logging
import os
import pprint
import sys
from typing import Any, List, Mapping

import vertex_client_utils as client_utils
import vertex_nas_cli_parser as nas_cli_parser

from gcs_utils import gcs_path_utils

from proxy_task import proxy_task_controller_utils_constants as controller_constants
from proxy_task import proxy_task_model_selection_lib_constants as model_selection_constants
from proxy_task import proxy_task_variance_measurement_lib_constants as variance_measurement_constants

import pyglove as pg

import yaml

# The directory for model outputs to be mounted in the docker container.
_LOCAL_RUN_OUTPUT_DIR = "/tmp/nas_job_output"
# The directory for model inputs to be mounted in the docker container.
_LOCAL_RUN_DATA_DIR = "/test_data"

_PREBUILT_SEARCH_SPACES_MAP = {
    "nasfpn": "search_spaces.nasfpn_search_space",
    "spinenet": "search_spaces.spinenet_search_space",
    "spinenet_v2": "search_spaces.spinenet_v2_search_space",
    "spinenet_mbconv": "search_spaces.spinenet_mbconv_search_space",
    "mnasnet": "search_spaces.mnasnet_search_space",
    "efficientnet_v2": "search_spaces.efficientnet_v2_search_space",
    "pointpillars": "search_spaces.pointpillars_search_space",
    "randaugment_detection": "search_spaces.randaugment_detection_search_space",
    "randaugment_segmentation": (
        "search_spaces.randaugment_segmentation_search_space"
    ),
    "autoaugment_detection": "search_spaces.autoaugment_detection_search_space",
    "autoaugment_segmentation": (
        "search_spaces.autoaugment_segmentation_search_space"
    ),
    "augment_3d_basic": "search_spaces.augment_3d_basic_search_space",
    "spinenet_scaling": "search_spaces.spinenet_scaling_search_space",
}

_PREBUILT_TRAINER_FLAGS = frozenset(
    {"training_data_path", "validation_data_path", "model", "search_space"})

# This is the default directory where `application_default_credentials.json` is
# saved by `gcloud auth application-default login` command.
_DEFAULT_LOCAL_CREDENTIAL_DIR = "~/.config/gcloud"

# Default value for jobSpec.scheduling.timeout. After the timeout, the requested
# job gets cancelled. Note that this is not the timeout for trials.
_JOB_SCHEDULING_TIMEOUT = "1209600s"  # 14 days

# Default NAS artifact-repository name.
_NAS_ARTIFACT_REGISTRY_REPOSITORY = "nas"


def get_docker_uri(region, project_id, docker_id):
  """Returns the artifact registry docker uri."""
  return "{}-docker.pkg.dev/{}/{}/{}:latest".format(
      region, project_id, _NAS_ARTIFACT_REGISTRY_REPOSITORY, docker_id)


def build_containers(args):
  """Builds containers."""
  if args.trainer_docker_id:
    trainer_docker_uri = get_docker_uri(
        region=args.region,
        project_id=args.project_id,
        docker_id=args.trainer_docker_id,
    )
    logging.info(
        "Starting building trainer docker %s with %s.",
        trainer_docker_uri,
        args.trainer_docker_file,
    )
    build_and_push_docker(
        trainer_docker_uri, args.trainer_docker_file, args.use_cache
    )

  if args.latency_calculator_docker_id:
    latency_calculator_docker_uri = get_docker_uri(
        region=args.region,
        project_id=args.project_id,
        docker_id=args.latency_calculator_docker_id)
    logging.info("Starting building latency calculator docker %s with %s.",
                 latency_calculator_docker_uri,
                 args.latency_calculator_docker_file)
    build_and_push_docker(latency_calculator_docker_uri,
                          args.latency_calculator_docker_file,
                          args.use_cache)

  if args.proxy_task_model_selection_docker_id:
    proxy_task_model_selection_docker_uri = get_docker_uri(
        region=args.region,
        project_id=args.project_id,
        docker_id=args.proxy_task_model_selection_docker_id)
    logging.info(
        "Starting building proxy-task model selection docker %s with %s.",
        proxy_task_model_selection_docker_uri,
        args.proxy_task_model_selection_docker_file)
    build_and_push_docker(proxy_task_model_selection_docker_uri,
                          args.proxy_task_model_selection_docker_file,
                          args.use_cache)
  if args.proxy_task_search_controller_docker_id:
    proxy_task_search_controller_docker_uri = get_docker_uri(
        region=args.region,
        project_id=args.project_id,
        docker_id=args.proxy_task_search_controller_docker_id)
    logging.info(
        "Starting building proxy-task search controller docker %s with %s.",
        proxy_task_search_controller_docker_uri,
        args.proxy_task_search_controller_docker_file)
    build_and_push_docker(proxy_task_search_controller_docker_uri,
                          args.proxy_task_search_controller_docker_file,
                          args.use_cache)
  if args.proxy_task_variance_measurement_docker_id:
    proxy_task_variance_measurement_docker_uri = get_docker_uri(
        region=args.region,
        project_id=args.project_id,
        docker_id=args.proxy_task_variance_measurement_docker_id)
    logging.info(
        "Starting building proxy-task variance measurement docker %s with %s.",
        proxy_task_variance_measurement_docker_uri,
        args.proxy_task_variance_measurement_docker_file)
    build_and_push_docker(proxy_task_variance_measurement_docker_uri,
                          args.proxy_task_variance_measurement_docker_file,
                          args.use_cache)


def get_search_space(args):
  """Returns the search space definition from the args."""

  search_space_module = None
  if args.prebuilt_search_space and args.search_space_module:
    raise ValueError(
        "Only one of prebuilt_search_space search_space_module can be set.")

  if args.prebuilt_search_space:
    search_space_module = _PREBUILT_SEARCH_SPACES_MAP.get(
        args.prebuilt_search_space, args.prebuilt_search_space)
    if not search_space_module:
      raise ValueError("search_space_module does not exist.")
  elif args.search_space_module:
    search_space_module = args.search_space_module
  else:
    raise ValueError(
        "Either prebuilt_search_space or search_space_module should be set.")

  logging.info("Using search_space_module: %s", search_space_module)
  search_space_file, search_space_mthod_name = search_space_module.rsplit(
      ".", 1)
  module = importlib.import_module(search_space_file)
  search_space_mthod = getattr(module, search_space_mthod_name)
  return search_space_mthod()


def add_machine_configurations(worker_pool_spec, accelerator_type, num_gpus,
                               master_machine_type):
  """Sets machine configurations for cloud training jobs."""

  # More checks for A100 GPUs.
  # https://cloud.google.com/vertex-ai/docs/training/configure-compute#specifying_gpus
  if (accelerator_type == "NVIDIA_TESLA_A100" or
      master_machine_type.startswith("a2")):
    middle = "highgpu" if num_gpus != 16 else "megagpu"
    expected_machine_type = "a2-{}-{}g".format(middle, num_gpus)
    if master_machine_type != expected_machine_type:
      raise ValueError(
          "master_machine_type, accelerator_type and num_gpus should be set "
          "consistently for A100 GPUs. Expect master_machine_type to be "
          "{} but is set as {}.".format(expected_machine_type,
                                        master_machine_type))

  if accelerator_type.lower() == "tpu":
    worker_pool_spec["machineSpec"] = {
        "machineType": master_machine_type,
        "acceleratorType": "TPU_V2",
        "acceleratorCount": 8
    }
  elif accelerator_type.lower().startswith("nvidia"):
    worker_pool_spec["machineSpec"] = {
        "machineType": master_machine_type,
        "acceleratorType": accelerator_type,
        "acceleratorCount": num_gpus
    }
  elif not accelerator_type:
    worker_pool_spec["machineSpec"] = {
        "machineType": master_machine_type,
        "acceleratorType": "ACCELERATOR_TYPE_UNSPECIFIED"
    }

  worker_pool_spec.update({
      "replicaCount": 1,
      "diskSpec": {
          "bootDiskType": "pd-ssd",
          "bootDiskSizeGb": 100
      }
  })


def add_prebuilt_trainer_args(args, args_map, is_train_args_map):
  """Adds pre-built trainer args and returns args_map."""
  for k, v in args_map.items():
    if not k or not v:
      raise ValueError(
          "Container-flag key/value must be non-empty: {}={}".format(k, v))

  args_map["search_space"] = args.prebuilt_search_space

  if args.prebuilt_search_space in ("randaugment_detection",
                                    "randaugment_segmentation",
                                    "augment_3d_basic"):
    args_map["multiple_eval_during_search"] = True

  if args.command == "search" and args.accelerator_type == "TPU":
    # Respects `use_tpu` field when running in cloud (search or train mode)
    args_map["use_tpu"] = True
  else:
    args_map["use_tpu"] = False

  if "params_override" in args_map:
    if not args_map["params_override"].endswith(".yaml"):
      raise ValueError("params_override needs to be a yaml file.")
    with open(args_map["params_override"]) as f:
      params_override = yaml.load(f, Loader=yaml.FullLoader)
      args_map["params_override"] = json.dumps(
          params_override, indent=2, sort_keys=True)
      logging.info("Using params_override: %s\n",
                   pprint.PrettyPrinter().pformat(params_override))
  if is_train_args_map:
    args_map["job_mode"] = "train"

  # Ensures all required prebuilt trainer flags are set.
  if not _PREBUILT_TRAINER_FLAGS.issubset(args_map.keys()):
    raise ValueError("Prebuilt trainer needs to set flags: {}".format(
        _PREBUILT_TRAINER_FLAGS - args_map.keys()))

  return args_map


def create_container_args_map(container_flags, args, is_train_args_map):
  """Creates a dict for the flags to be passed to the container."""
  args_map = client_utils.extract_container_flags(container_flags)
  if args.use_prebuilt_trainer:
    return add_prebuilt_trainer_args(
        args=args, args_map=args_map, is_train_args_map=is_train_args_map)
  else:
    return args_map


def sample_nas_trial_param_str(search_space):
  """Returns JSON string value for nas trial param."""
  algorithm = pg.geno.Random()
  algorithm.setup(search_space)
  sample_model = algorithm.propose()
  sample_model.use_spec(search_space)
  return json.dumps(sample_model.parameters(use_literal_values=True))


def run_latency_calculator_local(latency_args, docker_uri, local_output_dir):
  """Runs latency calculator container locally."""

  if local_output_dir:
    local_job_dir_cmd = [
        "-v", "{}:{}".format(local_output_dir, _LOCAL_RUN_OUTPUT_DIR)
    ]
  else:
    local_job_dir_cmd = []

  # Local credential path to mount to the container.
  credential_path = os.path.expanduser(_DEFAULT_LOCAL_CREDENTIAL_DIR)
  cmd = [
      "docker",
      "run",
      "--ipc",
      "host",
      "-v",
      "{}:/root/.config/gcloud".format(credential_path),
  ] + local_job_dir_cmd + ["-t", docker_uri] + latency_args

  client_utils.run_command_with_stdout(
      cmd, error_message="Fail to run latency calculator locally.")


def run_container_local(flag_map, docker_uri, args):
  """Runs container locally."""
  training_args = client_utils.convert_flag_map_to_list(flag_map)

  # Mount volume for local data directory.
  mount_dir_cmds = []
  if args.local_data_dir:
    mount_dir_cmds.extend(
        ["-v", "{}:{}".format(args.local_data_dir, _LOCAL_RUN_DATA_DIR)])

  # Mount volume for local output directory.
  if args.local_output_dir:
    mount_dir_cmds.extend(
        ["-v", "{}:{}".format(args.local_output_dir, _LOCAL_RUN_OUTPUT_DIR)])

  # Mount volume for application-default credential directory.
  credential_path = os.path.expanduser(_DEFAULT_LOCAL_CREDENTIAL_DIR)
  mount_dir_cmds.extend(
      ["-v", "{}:/root/.config/gcloud".format(credential_path)])

  cmd = (["docker", "run", "--ipc", "host"] + mount_dir_cmds +
         ["-t", docker_uri] + training_args)
  client_utils.run_command_with_stdout(
      cmd, error_message="Failed to run docker locally"
  )


def run_binary_local(local_binary, flag_map):
  """Runs python script locally."""
  training_args = client_utils.convert_flag_map_to_list(flag_map)
  cmd = ["python3", local_binary] + training_args
  logging.info("Run local binary with command:\n%s", " ".join(cmd))
  client_utils.run_command_with_stdout(
      cmd, error_message="Failed to run binary locally"
  )


def get_nas_target_reward_metric(args, search_args_map):
  """Returns the NAS target reward metric."""
  if args.use_prebuilt_trainer:
    if search_args_map["model"] == "segmentation":
      nas_target_reward_metric = "miou"
    elif search_args_map["model"] == "classification":
      nas_target_reward_metric = "top_1_accuracy"
    else:
      nas_target_reward_metric = "AP"
  else:
    nas_target_reward_metric = args.nas_target_reward_metric
  if not nas_target_reward_metric:
    raise ValueError("nas_target_reward_metric must be set.")
  return nas_target_reward_metric


def get_multi_trial_algorithm(search_args_map):
  """Returns the multi-trial-algorithm."""
  if search_args_map.get("search_space") in [
      "randaugment_detection", "randaugment_segmentation", "augment_3d_basic",
      "spinenet_scaling"
  ]:
    return "GRID_SEARCH"
  else:
    return "REINFORCEMENT_LEARNING"


def create_mirror_machines(
    worker_pool_spec,
    total_num_mirror_machines):
  """Returns worker-pool-spec list for mirrored multi-machine distributed training.
  """
  if total_num_mirror_machines < 1:
    raise ValueError("total_num_mirror_machines can not be less than 1.")
  if total_num_mirror_machines == 1:
    # No distributed training.
    return [worker_pool_spec]

  # For mirrored distributed training, worker-pool-0 by default
  # has 1 machine always. The extra machines (apart from the one above)
  # are then added to worker-pool-1 as "extra" replicas:
  # https://cloud.google.com/vertex-ai/docs/training/distributed-training#configure_a_distributed_training_job
  replica_worker_pool_spec = copy.deepcopy(worker_pool_spec)
  replica_worker_pool_spec["replicaCount"] = total_num_mirror_machines - 1
  return [worker_pool_spec, replica_worker_pool_spec]


def get_max_trial_counts(args):
  """Returns max-trial-count, max-parallel-trial-count, and max-failed-trial-count.
  """

  if args.command == "select_proxy_task_models":
    max_nas_trial = model_selection_constants.START_NUM_MODELS
    max_parallel_nas_trial = args.max_parallel_nas_trial
    max_failed_nas_trial = model_selection_constants.MAX_ALLOWED_FAILURES
    return max_nas_trial, max_parallel_nas_trial, max_failed_nas_trial

  if args.command == "search_proxy_task":
    max_nas_trial = 0
    max_parallel_nas_trial = 0
    max_failed_nas_trial = 0
    return max_nas_trial, max_parallel_nas_trial, max_failed_nas_trial

  if args.command == "measure_proxy_task_variance":
    max_nas_trial = variance_measurement_constants.NUM_TRIALS_FOR_MEASUREMENT
    max_parallel_nas_trial = variance_measurement_constants.NUM_TRIALS_FOR_MEASUREMENT
    max_failed_nas_trial = variance_measurement_constants.NUM_TRIALS_FOR_MEASUREMENT
    return max_nas_trial, max_parallel_nas_trial, max_failed_nas_trial

  max_nas_trial = args.max_nas_trial
  max_parallel_nas_trial = args.max_parallel_nas_trial
  max_failed_nas_trial = args.max_failed_nas_trial
  return max_nas_trial, max_parallel_nas_trial, max_failed_nas_trial


def create_search_and_train_job_spec(args, search_args_map, search_space_spec,
                                     train_args_map):
  """Constructs NAS-job for NAS search or train job."""
  # TODO: update document link.
  # See documentation:
  # https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#traininginput

  expected_commands = ["search"]

  expected_commands += [
      "select_proxy_task_models", "search_proxy_task",
      "measure_proxy_task_variance"
  ]

  if args.command not in expected_commands:
    raise ValueError(
        "Unexpected command for _create_search_and_train_job_spec: {}".format(
            args.command))
  if not search_space_spec:
    raise ValueError("Need a search_space_spec to start the search job.")

  # Get job directories.
  job_name = args.job_name
  dir_name = job_name + datetime.datetime.now().strftime("_%Y%m%d_%H%M%S")
  search_job_dir = os.path.join(args.root_output_dir, dir_name, "nas", "search")
  train_job_dir = os.path.join(args.root_output_dir, dir_name, "nas", "train")

  # By default, we use the `latest` built/pushed docker image uri.
  image_uri = get_docker_uri(
      region=args.region,
      project_id=args.project_id,
      docker_id=args.trainer_docker_id)

  # Create worker-pool-spec for search-job.
  search_job_args = client_utils.convert_flag_map_to_list(search_args_map)
  search_worker_pool_spec = {
      "containerSpec": {
          "imageUri": image_uri,
          "args": search_job_args
      }
  }
  add_machine_configurations(
      search_worker_pool_spec,
      args.accelerator_type,
      num_gpus=args.num_gpus,
      master_machine_type=args.master_machine_type)

  # Create worker-pool-spec for train-job.
  if train_args_map is not None:
    train_job_args = client_utils.convert_flag_map_to_list(train_args_map)
    train_worker_pool_spec = {
        "containerSpec": {
            "imageUri": image_uri,
            "args": train_job_args
        }
    }
    add_machine_configurations(
        train_worker_pool_spec,
        args.train_accelerator_type,
        num_gpus=args.train_num_gpus,
        master_machine_type=args.train_master_machine_type)

  max_nas_trial, max_parallel_nas_trial, max_failed_nas_trial = get_max_trial_counts(
      args)
  # Fill the nas-job-spec.
  nas_job_spec = {
      "searchSpaceSpec": search_space_spec,
      "multiTrialAlgorithmSpec": {
          "multiTrialAlgorithm": get_multi_trial_algorithm(search_args_map),
          "metric": {
              "metricId":
                  get_nas_target_reward_metric(
                      args=args, search_args_map=search_args_map),
              "goal":
                  "MAXIMIZE",
          },
          # Search job spec.
          "searchTrialSpec": {
              "searchTrialJobSpec": {
                  "workerPoolSpecs":
                      create_mirror_machines(
                          worker_pool_spec=search_worker_pool_spec,
                          total_num_mirror_machines=args.num_mirror_machines),
                  "baseOutputDirectory": {
                      "outputUriPrefix": search_job_dir
                  },
                  "scheduling": {
                      "timeout": _JOB_SCHEDULING_TIMEOUT
                  },
              },
              "maxTrialCount": max_nas_trial,
              "maxParallelTrialCount": max_parallel_nas_trial,
              "maxFailedTrialCount": max_failed_nas_trial,
          },
      }
  }

  if train_args_map is not None:
    # Finetuning job spec.
    nas_job_spec["multiTrialAlgorithmSpec"]["trainTrialSpec"] = {
        "trainTrialJobSpec": {
            "workerPoolSpecs":
                create_mirror_machines(
                    worker_pool_spec=train_worker_pool_spec,
                    total_num_mirror_machines=args.train_num_mirror_machines),
            "baseOutputDirectory": {
                "outputUriPrefix": train_job_dir
            },
            "scheduling": {
                "timeout": _JOB_SCHEDULING_TIMEOUT
            }
        },
        "maxParallelTrialCount": args.train_max_parallel_trial,
        "frequency": args.train_frequency,
    }

  nas_job = {
      "displayName": job_name,
      "labels": {
          "search_space": search_args_map.get("search_space", ""),
      },
      "nasJobSpec": nas_job_spec
  }
  return nas_job


def create_search_and_train_job_spec_from_previous_job(args):
  """Constructs NAS-job for NAS search or train job form a previous search job.
  """
  if args.command != "search_resume":
    raise ValueError(
        "Unexpected command for "
        "create_search_and_train_job_spec_from_previous_job: {}".format(
            args.command))

  if not args.previous_nas_job_id:
    raise ValueError("Need a 'previous_nas_job_id' to resume search")

  # Get previous job spec.
  previous_nas_job = client_utils.get_job(
      vertex_ai_endpoint=client_utils.get_service_endpoint(
          service_endpoint=args.service_endpoint, region=args.region),
      project_id=args.project_id,
      location=args.region,
      job_id=args.previous_nas_job_id)

  # Create new job-spec from the previous job-spec.
  # Will only add a new GCS output directory, job-name,
  # total-trials, parallel trials, and max-failed trials.
  # The rest of the job settings will be reused.
  job_name = args.job_name
  dir_name = job_name + datetime.datetime.now().strftime("_%Y%m%d_%H%M%S")
  search_job_dir = os.path.join(args.root_output_dir, dir_name, "nas", "search")
  train_job_dir = os.path.join(args.root_output_dir, dir_name, "nas", "train")
  nas_job = {
      "displayName": job_name,
      "labels": previous_nas_job["labels"],
      "nasJobSpec": previous_nas_job["nasJobSpec"]
  }
  nas_job["nasJobSpec"]["searchSpaceSpec"] = None
  nas_job["nasJobSpec"]["resume_nas_job_id"] = args.previous_nas_job_id
  nas_job["nasJobSpec"]["multiTrialAlgorithmSpec"]["multiTrialAlgorithm"] = None
  nas_job["nasJobSpec"]["multiTrialAlgorithmSpec"]["metric"] = None

  search_trial_spec = nas_job["nasJobSpec"]["multiTrialAlgorithmSpec"][
      "searchTrialSpec"]
  search_trial_spec["searchTrialJobSpec"]["baseOutputDirectory"][
      "outputUriPrefix"] = search_job_dir
  search_trial_spec["maxTrialCount"] = args.max_nas_trial
  search_trial_spec["maxParallelTrialCount"] = args.max_parallel_nas_trial
  search_trial_spec["maxFailedTrialCount"] = args.max_failed_nas_trial
  if "trainTrialSpec" in nas_job["nasJobSpec"]["multiTrialAlgorithmSpec"]:
    nas_job["nasJobSpec"]["multiTrialAlgorithmSpec"]["trainTrialSpec"][
        "trainTrialJobSpec"]["baseOutputDirectory"][
            "outputUriPrefix"] = train_job_dir
  return nas_job


def create_train_only_job_spec(args, search_space_spec, train_args_map):
  """Constructs NAS-job for train-only job."""

  if args.command != "train":
    raise ValueError(
        "Unexpected command for _create_train_only_job_spec: {}".format(
            args.command))
  if not search_space_spec:
    raise ValueError("Need a search_space_spec to start the job.")

  # Add more flags (search-job-dir and trials to retrain) to the
  # training-docker (train_args_map) for the stage2 train only run.
  prev_search_job = client_utils.get_job(
      vertex_ai_endpoint=client_utils.get_service_endpoint(
          service_endpoint=args.service_endpoint,
          region=args.search_job_region),
      project_id=args.project_id,
      location=args.search_job_region,
      job_id=args.search_job_id)
  prev_search_job_dir = client_utils.get_job_dir_for_nas_job(prev_search_job)
  train_args_map["retrain_search_job_dir"] = prev_search_job_dir
  train_args_map["retrain_search_job_trials"] = args.train_nas_trial_numbers
  retrain_nas_trial_count = len(args.train_nas_trial_numbers.split(","))

  # Get job directory.
  job_name = "search_{}_{}".format(args.search_job_id, args.train_job_suffix)
  dir_name = job_name + datetime.datetime.now().strftime("_%Y%m%d_%H%M%S")
  train_job_dir = os.path.join(args.root_output_dir, dir_name, "nas", "train")

  # By default, we use the `latest` built/pushed docker image uri.
  image_uri = get_docker_uri(
      region=args.region,
      project_id=args.project_id,
      docker_id=args.trainer_docker_id)

  # Create worker-pool-spec for train-job.
  train_job_args = client_utils.convert_flag_map_to_list(train_args_map)
  train_worker_pool_spec = {
      "containerSpec": {
          "imageUri": image_uri,
          "args": train_job_args
      }
  }
  add_machine_configurations(
      train_worker_pool_spec,
      args.train_accelerator_type,
      num_gpus=args.train_num_gpus,
      master_machine_type=args.train_master_machine_type)

  # Fill the nas-job-spec.
  # NOTE: Although this is a train-only job, we will still set it up as a
  # a dummy search job to launch multiple stage2 trials.
  nas_job_spec = {
      "searchSpaceSpec": search_space_spec,
      "multiTrialAlgorithmSpec": {
          "multiTrialAlgorithm": get_multi_trial_algorithm(train_args_map),
          "metric": {
              "metricId":
                  get_nas_target_reward_metric(
                      args=args, search_args_map=train_args_map),
              "goal":
                  "MAXIMIZE",
          },
          # Search job spec.
          "searchTrialSpec": {
              "searchTrialJobSpec": {
                  "workerPoolSpecs":
                      create_mirror_machines(
                          worker_pool_spec=train_worker_pool_spec,
                          total_num_mirror_machines=args
                          .train_num_mirror_machines),
                  "baseOutputDirectory": {
                      "outputUriPrefix": train_job_dir
                  },
                  "scheduling": {
                      "timeout": _JOB_SCHEDULING_TIMEOUT
                  },
              },
              "maxTrialCount": retrain_nas_trial_count,
              "maxParallelTrialCount": retrain_nas_trial_count,
              "maxFailedTrialCount": retrain_nas_trial_count,
          },
      }
  }

  nas_job = {
      "displayName": job_name,
      "labels": {
          "search_space": train_args_map.get("search_space", ""),
      },
      "nasJobSpec": nas_job_spec
  }
  return nas_job


def list_trials(args):
  """Lists trials from NAS search job."""
  trials = client_utils.list_trials(
      vertex_ai_endpoint=client_utils.get_service_endpoint(
          service_endpoint=args.service_endpoint, region=args.region),
      project_id=args.project_id,
      location=args.region,
      job_id=args.job_id,
      max_trials=args.max_trials)
  if not trials:
    logging.info("Unable to find trials.")
  else:
    if args.trials_output_file:
      with open(args.trials_output_file, "w") as f:
        json.dump(trials, f, indent=4)
        logging.info("%d trials were written to: %s", len(trials),
                     args.trials_output_file)
    else:
      logging.info(json.dumps(trials, indent=4))


def get_latency_container_args(args, prebuilt_trainer_model, nas_job_name):
  """Gets the args for latency calculator in a dict."""

  # The users can pass parameters through the `latency_docker_flags`.
  flags_map = client_utils.extract_container_flags(args.latency_docker_flags)

  if args.use_prebuilt_latency_calculator:
    # Set flags used only by pre-built latency calculator.
    flags_map["image_width"] = args.prebuilt_latency_image_width
    flags_map["image_height"] = args.prebuilt_latency_image_height
    flags_map["input_node"] = args.prebuilt_latency_input_node
    flags_map["num_repetitions_for_latency_computation"] = (
        args.prebuilt_num_repetitions_for_latency_computation)

    if prebuilt_trainer_model == "retinanet":
      flags_map["output_nodes"] = ("DetectionBoxes:0,DetectionClasses:0,"
                                   "DetectionScores:0")
    elif prebuilt_trainer_model == "classification":
      flags_map["output_nodes"] = "ClassId:0,Probabilities:0"
    else:
      flags_map["output_nodes"] = args.prebuilt_latency_output_nodes

    if args.target_device_type == "CPU":
      flags_map["device_type"] = "CPU"
      flags_map["use_tensorrt_conversion_on_gpu"] = False
    else:
      flags_map["device_type"] = "GPU"
      flags_map["use_tensorrt_conversion_on_gpu"] = (
          args.prebuilt_latency_tensorrt_conversion_on_gpu)

  flags_map.update({
      "service_endpoint":
          client_utils.get_service_endpoint(
              service_endpoint=args.service_endpoint, region=args.region),
      "project_id":
          args.project_id,
      "nas_job_id":
          nas_job_name,
  })

  return flags_map


def get_local_latency_container_args(args, prebuilt_trainer_model,
                                     nas_job_name):
  """Gets the args for local latency calculator in a dict."""
  flags_map = get_latency_container_args(args, prebuilt_trainer_model,
                                         nas_job_name)
  # Add local latency calculator specific flags.
  flags_map["latency_worker_id"] = args.latency_worker_id
  flags_map["num_latency_workers"] = args.num_latency_workers
  return flags_map


def create_latency_calculation_job_spec(args, prebuilt_trainer_model,
                                        nas_job_name):
  """Constructs CustomJob for latency calculator container."""
  # See documentation:
  # https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#traininginput

  if not args.latency_calculator_docker_id:
    raise ValueError("latency_calculator_docker_id flag should be set.")
  latency_calculation_docker = get_docker_uri(
      region=args.region,
      project_id=args.project_id,
      docker_id=args.latency_calculator_docker_id)
  logging.info("latency_calculation_docker is %s", latency_calculation_docker)

  num_cloud_latency_workers = args.num_cloud_latency_workers

  # Tell each latency worker its ID and the total number of workers.
  worker_pool_specs = []
  for latency_worker_id in range(num_cloud_latency_workers):
    flags_map = get_latency_container_args(args, prebuilt_trainer_model,
                                           nas_job_name)
    flags_map["latency_worker_id"] = latency_worker_id
    flags_map["num_latency_workers"] = num_cloud_latency_workers
    latency_args = client_utils.convert_flag_map_to_list(flags_map)

    worker_pool_spec = {
        "machineSpec": {
            "machineType": "n1-highmem-8"
        },
        "containerSpec": {
            "imageUri": latency_calculation_docker,
            "args": latency_args,
        },
        "replicaCount": 1
    }

    if args.target_device_type != "CPU":
      worker_pool_spec["machineSpec"].update({
          "acceleratorType": args.target_device_type,
          "acceleratorCount": 1
      })

    worker_pool_specs.append(worker_pool_spec)

  custom_job = {
      "displayName":
          client_utils.get_latency_calculator_display_name(nas_job_name),
      "labels": {
          "nas_job_type": "latency_calculation",
      },
      "jobSpec": {
          "workerPoolSpecs": worker_pool_specs,
          "scheduling": {
              "timeout": _JOB_SCHEDULING_TIMEOUT
          },
      }
  }
  return custom_job


def create_latency_calculation_job_spec_from_previous_job(args, nas_job_name):
  """Constructs CustomJob for latency calculator container using a previous job.
  """
  if not args.previous_latency_job_id:
    raise ValueError("previous_latency_job_id flag should be set.")

  # Get previous latency job spec.
  previous_custom_job = client_utils.get_job(
      vertex_ai_endpoint=client_utils.get_service_endpoint(
          service_endpoint=args.service_endpoint, region=args.region),
      project_id=args.project_id,
      location=args.region,
      job_id=args.previous_latency_job_id,
      job_type="custom",
      check_job_failed=False)

  # Create new job spec.
  custom_job_spec = {
      "labels": previous_custom_job["labels"],
      "jobSpec": previous_custom_job["jobSpec"],
  }
  client_utils.reset_nas_job_name_in_latency_job_spec(
      latency_job_spec=custom_job_spec, nas_job_name=nas_job_name)

  return custom_job_spec


def build_and_push_docker(docker_uri, docker_file="Dockerfile", use_cache=True):
  """Builds and pushes docker to cloud gcr."""

  cache_option = [] if use_cache else ["--no-cache"]
  client_utils.run_command_with_stdout(
      ["docker", "build"]
      + cache_option
      + ["-f", docker_file, "-t", docker_uri, "."],
      error_message="Fail to build docker image",
  )

  client_utils.run_command_with_stdout(
      ["docker", "push", docker_uri], error_message="Fail to push docker image")

  print("Successfully built/pushed container: {}".format(docker_uri))


def create_proxy_task_model_selection_job_spec(args, search_job_spec,
                                               latency_calculator_job_spec):
  """Constructs job spec for proxy-task model-selection."""

  # Set proxy task model selection docker path.
  # This docker will run on cloud CPU.
  if not args.proxy_task_model_selection_docker_id:
    raise ValueError("proxy_task_model_selection_docker_id flag should be set.")
  proxy_task_model_selection_docker = get_docker_uri(
      region=args.region,
      project_id=args.project_id,
      docker_id=args.proxy_task_model_selection_docker_id)
  logging.info("proxy_task_model_selection_docker is %s",
               proxy_task_model_selection_docker)

  model_selection_job_name = "Model_Selection_" + args.job_name
  # Create a gcs working directory for the model-selection job.
  if args.previous_model_selection_dir:
    model_selection_dir = args.previous_model_selection_dir
  else:
    model_selection_dir_name = model_selection_job_name + datetime.datetime.now(
    ).strftime("_%Y%m%d_%H%M%S")
    model_selection_dir = os.path.join(args.root_output_dir,
                                       model_selection_dir_name)

  # Save search-job-spec to model-selection-dir on GCS.
  client_utils.save_search_job_spec_file(
      project_id=args.project_id,
      search_job_spec=search_job_spec,
      controller_dir=model_selection_dir)
  # Save latency-calculator-job-spec to model-selection-dir on GCS.
  if latency_calculator_job_spec:
    client_utils.save_latency_calculator_job_spec_file(
        project_id=args.project_id,
        latency_calculator_job_spec=latency_calculator_job_spec,
        controller_dir=model_selection_dir)

  # Set docker arguments.
  flags_map = {}
  flags_map["model_selection_dir"] = model_selection_dir
  flags_map["region"] = args.region
  flags_map["project_id"] = args.project_id
  flags_map["service_endpoint"] = client_utils.get_service_endpoint(
      service_endpoint=args.service_endpoint, region=args.region)
  flags_map["accuracy_metric_id"] = args.accuracy_metric_id
  flags_map["latency_metric_id"] = args.latency_metric_id
  docker_args = client_utils.convert_flag_map_to_list(flags_map)
  # Set up machine for docker.
  worker_pool_spec = {
      "machineSpec": {
          "machineType": "n1-highmem-8"
      },
      "containerSpec": {
          "imageUri": proxy_task_model_selection_docker,
          "args": docker_args,
      },
      "replicaCount": 1
  }

  # Create cloud job-spec.
  cloud_job_spec = {
      "displayName": model_selection_job_name,
      "labels": {
          "nas_job_type": "proxy_task_model_selection",
      },
      "jobSpec": {
          "workerPoolSpecs": worker_pool_spec,
          "scheduling": {
              "timeout": _JOB_SCHEDULING_TIMEOUT
          },
          "serviceAccount": args.service_account,
      }
  }
  return cloud_job_spec, model_selection_dir


def create_proxy_task_search_controller_job_spec(args, search_job_spec,
                                                 latency_calculator_job_spec):
  """Constructs job spec for proxy-task search-controller."""

  # Set proxy-task search controller docker path.
  # This docker will run on cloud CPU.
  if not args.proxy_task_search_controller_docker_id:
    raise ValueError(
        "proxy_task_search_controller_docker_id flag should be set.")
  proxy_task_search_controller_docker = get_docker_uri(
      region=args.region,
      project_id=args.project_id,
      docker_id=args.proxy_task_search_controller_docker_id)
  logging.info("proxy_task_search_controller_docker is %s",
               proxy_task_search_controller_docker)

  search_controller_job_name = "Search_controller_" + args.job_name
  # Create a gcs working directory for the job.
  if args.previous_proxy_task_search_dir:
    search_controller_dir = args.previous_proxy_task_search_dir
  else:
    search_controller_dir_name = search_controller_job_name + datetime.datetime.now(
    ).strftime("_%Y%m%d_%H%M%S")
    search_controller_dir = os.path.join(args.root_output_dir,
                                         search_controller_dir_name)
  # Save search-job-spec to GCS.
  client_utils.save_search_job_spec_file(
      project_id=args.project_id,
      search_job_spec=search_job_spec,
      controller_dir=search_controller_dir)
  # Save latency-calculator-job-spec to GCS.
  if latency_calculator_job_spec:
    client_utils.save_latency_calculator_job_spec_file(
        project_id=args.project_id,
        latency_calculator_job_spec=latency_calculator_job_spec,
        controller_dir=search_controller_dir)

  # Set docker arguments.
  flags_map = {}
  flags_map[
      "proxy_task_model_selection_job_id"] = args.proxy_task_model_selection_job_id
  flags_map[
      "proxy_task_model_selection_job_region"] = args.proxy_task_model_selection_job_region
  flags_map[
      "proxy_task_model_selection_service_endpoint"] = client_utils.get_service_endpoint(
          service_endpoint=args.service_endpoint,
          region=args.proxy_task_model_selection_job_region)
  flags_map[
      "proxy_task_config_generator_module"] = args.proxy_task_config_generator_module
  flags_map["search_controller_dir"] = search_controller_dir
  flags_map["region"] = args.region
  flags_map["project_id"] = args.project_id
  flags_map["service_endpoint"] = client_utils.get_service_endpoint(
      service_endpoint=args.service_endpoint, region=args.region)
  if args.desired_accuracy_correlation:
    flags_map[
        "desired_accuracy_correlation"] = args.desired_accuracy_correlation
  if args.desired_accuracy:
    flags_map["desired_accuracy"] = args.desired_accuracy
  if args.training_time_hrs_limit:
    flags_map["training_time_hrs_limit"] = args.training_time_hrs_limit
  if args.desired_latency_correlation:
    flags_map["desired_latency_correlation"] = args.desired_latency_correlation
  flags_map[
      "early_stop_proxy_task_if_not_best"] = args.early_stop_proxy_task_if_not_best
  docker_args = client_utils.convert_flag_map_to_list(flags_map)
  # Set up machine for docker.
  worker_pool_spec = {
      "machineSpec": {
          "machineType": "n1-highmem-8"
      },
      "containerSpec": {
          "imageUri": proxy_task_search_controller_docker,
          "args": docker_args,
      },
      "replicaCount": 1
  }

  # Create cloud job-spec.
  cloud_job_spec = {
      "displayName": search_controller_job_name,
      "labels": {
          "nas_job_type": "proxy_task_search_controller",
      },
      "jobSpec": {
          "workerPoolSpecs": worker_pool_spec,
          "scheduling": {
              "timeout": _JOB_SCHEDULING_TIMEOUT
          },
          "serviceAccount": args.service_account,
      }
  }
  return cloud_job_spec, search_controller_dir


def create_proxy_task_variance_measurement_job_spec(args, search_job_spec):
  """Constructs job spec for proxy-task variance measurement."""

  # Create a gcs working directory for the variance-measurement job.
  measurement_job_name = "Variance_Measurement_" + args.job_name
  measurement_job_dir_name = measurement_job_name + datetime.datetime.now(
  ).strftime("_%Y%m%d_%H%M%S")
  measurement_job_dir = os.path.join(args.root_output_dir,
                                     measurement_job_dir_name)
  # Save search-job-spec to variance-measurement-dir on GCS.
  client_utils.save_search_job_spec_file(
      project_id=args.project_id,
      search_job_spec=search_job_spec,
      controller_dir=measurement_job_dir)
  logging.info("Saved search_job_spec to %s", measurement_job_dir)

  # Set proxy task variance measurement docker path.
  if not args.proxy_task_variance_measurement_docker_id:
    raise ValueError(
        "proxy_task_variance_measurement_docker_id flag should be set.")
  variance_docker = get_docker_uri(
      region=args.region,
      project_id=args.project_id,
      docker_id=args.proxy_task_variance_measurement_docker_id)
  logging.info("proxy_task_variance_measurement_docker is %s", variance_docker)

  # Set docker arguments.
  flags_map = {}
  flags_map["variance_measurement_dir"] = measurement_job_dir
  flags_map["region"] = args.region
  flags_map["project_id"] = args.project_id
  flags_map["service_endpoint"] = client_utils.get_service_endpoint(
      service_endpoint=args.service_endpoint, region=args.region)
  docker_args = client_utils.convert_flag_map_to_list(flags_map)

  # Set up machine for docker.
  # This docker will run on cloud CPU.
  worker_pool_spec = {
      "machineSpec": {
          "machineType": "n1-highmem-8"
      },
      "containerSpec": {
          "imageUri": variance_docker,
          "args": docker_args,
      },
      "replicaCount": 1
  }

  # Create cloud job-spec.
  cloud_job_spec = {
      "displayName": measurement_job_name,
      "labels": {
          "nas_job_type": "proxy_task_variance_measurement",
      },
      "jobSpec": {
          "workerPoolSpecs": worker_pool_spec,
          "scheduling": {"timeout": _JOB_SCHEDULING_TIMEOUT},
          "serviceAccount": args.service_account,
      },
  }
  return cloud_job_spec, measurement_job_dir


def save_controller_job_name_to_gcs(
    project_id, controller_dir, controller_job_name
):
  """Saves controller job name to controller job directory."""
  # Save controller job name to controller-dir too.
  controller_job_name_filepath = os.path.join(
      controller_dir,
      controller_constants.CONTROLLER_JOB_NAME_FILE,
  )
  client_utils.write_json_object_to_gcs(
      project_id=project_id,
      json_object={
          controller_constants.CONTROLLER_JOB_NAME_KEY: controller_job_name
      },
      gcs_filename=controller_job_name_filepath,
  )


def main(argv):
  del argv

  logging.basicConfig(stream=sys.stdout, level=logging.INFO)
  parser = nas_cli_parser.create_nas_cli_parser()
  args = parser.parse_args()

  commands_requiring_search_space_spec = ["search", "search_in_local", "train"]

  commands_requiring_search_space_spec += [
      "search_proxy_task", "select_proxy_task_models",
      "measure_proxy_task_variance"
  ]

  if args.command in commands_requiring_search_space_spec:
    search_space = pg.search_space(get_search_space(args))
    search_space_spec = search_space.to_json_str(hide_default_values=True)
  else:
    search_space = None
    search_space_spec = None

  if args.command == "build":
    build_containers(args)
    return

  elif args.command == "search":
    search_args_map = create_container_args_map(
        args.search_docker_flags, args, is_train_args_map=False
    )
    # NOTE: An empty list ([]) for train_docker_flags is a valid value.
    if args.train_docker_flags is not None:
      train_args_map = create_container_args_map(
          args.train_docker_flags, args, is_train_args_map=True
      )
    else:
      train_args_map = None
    search_and_train_job_spec = create_search_and_train_job_spec(
        args=args,
        search_args_map=search_args_map,
        search_space_spec=search_space_spec,
        train_args_map=train_args_map,
    )
    nas_job_resource_name, _ = client_utils.launch_cloud_nas_job(
        service_endpoint=client_utils.get_service_endpoint(
            service_endpoint=args.service_endpoint, region=args.region
        ),
        project_id=args.project_id,
        location=args.region,
        job_spec=search_and_train_job_spec,
        job_type="NAS Search",
    )
    if args.latency_calculator_docker_id:
      latency_calculator_job_spec = create_latency_calculation_job_spec(
          args=args,
          prebuilt_trainer_model=search_args_map.get("model", ""),
          nas_job_name=nas_job_resource_name,
      )
      client_utils.launch_cloud_custom_job(
          service_endpoint=client_utils.get_service_endpoint(
              service_endpoint=args.service_endpoint, region=args.region
          ),
          project_id=args.project_id,
          location=args.region,
          job_spec=latency_calculator_job_spec,
          job_type="Latency Calculator",
      )

  elif args.command == "search_resume":
    search_and_train_job_spec = (
        create_search_and_train_job_spec_from_previous_job(args=args)
    )
    nas_job_resource_name, _ = client_utils.launch_cloud_nas_job(
        service_endpoint=client_utils.get_service_endpoint(
            service_endpoint=args.service_endpoint, region=args.region
        ),
        project_id=args.project_id,
        location=args.region,
        job_spec=search_and_train_job_spec,
        job_type="NAS Search",
    )
    if args.previous_latency_job_id:
      latency_calculator_job_spec = (
          create_latency_calculation_job_spec_from_previous_job(
              args, nas_job_name=nas_job_resource_name
          )
      )
      client_utils.launch_cloud_custom_job(
          service_endpoint=client_utils.get_service_endpoint(
              service_endpoint=args.service_endpoint, region=args.region
          ),
          project_id=args.project_id,
          location=args.region,
          job_spec=latency_calculator_job_spec,
          job_type="Latency Calculator",
      )

  elif args.command == "train":
    train_args_map = create_container_args_map(
        args.train_docker_flags, args, is_train_args_map=True
    )
    train_job_spec = create_train_only_job_spec(
        args=args,
        search_space_spec=search_space_spec,
        train_args_map=train_args_map,
    )
    nas_job_resource_name, _ = client_utils.launch_cloud_nas_job(
        service_endpoint=client_utils.get_service_endpoint(
            service_endpoint=args.service_endpoint, region=args.region
        ),
        project_id=args.project_id,
        location=args.region,
        job_spec=train_job_spec,
        job_type="NAS Train",
    )
    if args.latency_calculator_docker_id:
      latency_calculator_job_spec = create_latency_calculation_job_spec(
          args=args,
          prebuilt_trainer_model=train_args_map.get("model", ""),
          nas_job_name=nas_job_resource_name,
      )
      client_utils.launch_cloud_custom_job(
          service_endpoint=client_utils.get_service_endpoint(
              service_endpoint=args.service_endpoint, region=args.region
          ),
          project_id=args.project_id,
          location=args.region,
          job_spec=latency_calculator_job_spec,
          job_type="Latency Calculator",
      )

  elif args.command == "select_proxy_task_models":
    # Get search_job_spec and latency_job_spec for the customer job.
    # The job will be launched and managed by the model-selection code.
    search_args_map = create_container_args_map(
        args.search_docker_flags, args, is_train_args_map=False
    )
    search_job_spec = create_search_and_train_job_spec(
        args=args,
        search_args_map=search_args_map,
        search_space_spec=search_space_spec,
        train_args_map=None,
    )
    if args.latency_calculator_docker_id:
      # NOTE: Since the search job will be launched by the
      # model-selection code later, the `nas_job_name`
      # is passed an empty string for now. The model-selection code will
      # later manage this.
      latency_calculator_job_spec = create_latency_calculation_job_spec(
          args=args,
          prebuilt_trainer_model=search_args_map.get("model", ""),
          nas_job_name="",
      )
    else:
      latency_calculator_job_spec = None

    job_spec, controller_dir = create_proxy_task_model_selection_job_spec(
        args=args,
        search_job_spec=search_job_spec,
        latency_calculator_job_spec=latency_calculator_job_spec,
    )
    controller_job, _ = client_utils.launch_cloud_custom_job(
        service_endpoint=client_utils.get_service_endpoint(
            service_endpoint=args.service_endpoint, region=args.region
        ),
        project_id=args.project_id,
        location=args.region,
        job_spec=job_spec,
        job_type="Proxy Task Model Selection",
    )
    logging.info("Model-selection job is %s", controller_job)
    save_controller_job_name_to_gcs(
        project_id=args.project_id,
        controller_dir=controller_dir,
        controller_job_name=controller_job,
    )

  elif args.command == "measure_proxy_task_variance":
    search_args_map = create_container_args_map(
        args.search_docker_flags, args, is_train_args_map=False
    )
    search_job_spec = create_search_and_train_job_spec(
        args=args,
        search_args_map=search_args_map,
        search_space_spec=search_space_spec,
        train_args_map=None,
    )
    job_spec, controller_dir = create_proxy_task_variance_measurement_job_spec(
        args=args, search_job_spec=search_job_spec
    )
    controller_job, _ = client_utils.launch_cloud_custom_job(
        service_endpoint=client_utils.get_service_endpoint(
            service_endpoint=args.service_endpoint, region=args.region
        ),
        project_id=args.project_id,
        location=args.region,
        job_spec=job_spec,
        job_type="Proxy Task Variance Measurement",
    )
    save_controller_job_name_to_gcs(
        project_id=args.project_id,
        controller_dir=controller_dir,
        controller_job_name=controller_job,
    )

  elif args.command == "search_proxy_task":
    # Get search_job_spec and latency_job_spec for a proxy-task template job.
    # The actual proxy-task jobs will be launched and managed by
    # the proxy-task search controller code.
    # For now use search_args_map for the past full-training job.
    # This will get overwritten later before the actual proxy-tasks
    # will be launched.
    search_args_map = create_container_args_map(
        args.search_docker_flags, args, is_train_args_map=False)
    search_job_spec = create_search_and_train_job_spec(
        args=args,
        search_args_map=search_args_map,
        search_space_spec=search_space_spec,
        train_args_map=None)
    if args.latency_calculator_docker_id:
      # NOTE: Since the search job will be launched by the
      # proxy-task search controller code later, the `nas_job_name`
      # is passed an empty string for now. The proxy-task search controller
      # code will later manage this.
      latency_calculator_job_spec = create_latency_calculation_job_spec(
          args=args,
          prebuilt_trainer_model=search_args_map.get("model", ""),
          nas_job_name="")
    else:
      latency_calculator_job_spec = None

    job_spec, controller_dir = create_proxy_task_search_controller_job_spec(
        args=args,
        search_job_spec=search_job_spec,
        latency_calculator_job_spec=latency_calculator_job_spec)
    controller_job, _ = client_utils.launch_cloud_custom_job(
        service_endpoint=client_utils.get_service_endpoint(
            service_endpoint=args.service_endpoint, region=args.region),
        project_id=args.project_id,
        location=args.region,
        job_spec=job_spec,
        job_type="Proxy Task Search Controller")
    logging.info("Search-controller job is %s", controller_job)
    save_controller_job_name_to_gcs(
        project_id=args.project_id,
        controller_dir=controller_dir,
        controller_job_name=controller_job,
    )

  elif args.command == "search_in_local":
    local_args = (
        args.local_binary_flags
        if args.run_local_binary else args.search_docker_flags)
    search_args_map = create_container_args_map(
        local_args, args, is_train_args_map=False)
    search_args_map["nas_params_str"] = sample_nas_trial_param_str(search_space)

    if args.run_local_binary:
      search_args_map["job-dir"] = args.local_output_dir
      run_binary_local(args.local_binary, search_args_map)
    else:
      search_args_map["job-dir"] = _LOCAL_RUN_OUTPUT_DIR
      image_uri = get_docker_uri(
          region=args.region,
          project_id=args.project_id,
          docker_id=args.trainer_docker_id)
      run_container_local(search_args_map, image_uri, args)
    return

  elif args.command == "list_trials":
    list_trials(args)

  elif args.command == "run_latency_calculator_local":
    job_resource_name = ""
    if args.search_job_id:
      job_resource_name = "projects/{}/locations/{}/nasJobs/{}".format(
          args.project_id, args.region, args.search_job_id)

    if args.controller_job_id:
      job_resource_name = "projects/{}/locations/{}/customJobs/{}".format(
          args.project_id, args.region, args.controller_job_id)

    if not job_resource_name:
      raise ValueError("You need to provide parent job id.")
    flags_map = get_local_latency_container_args(args,
                                                 args.prebuilt_trainer_model,
                                                 job_resource_name)
    job_flags = client_utils.convert_flag_map_to_list(flags_map)
    latency_calculation_docker = get_docker_uri(
        region=args.region,
        project_id=args.project_id,
        docker_id=args.latency_calculator_docker_id)
    logging.info("latency_calculation_docker is %s", latency_calculation_docker)
    run_latency_calculator_local(job_flags, latency_calculation_docker,
                                 args.local_output_dir)

  else:
    raise ValueError("Unexpected command: {}".format(args.command))


if __name__ == "__main__":
  main("")