def _create_job_spec_dict( experiment: Experiment, job_mode: c.JobMode, image_id: str, index: int, caliban_config: Dict[str, Any], run_args: Optional[List[str]] = None, ) -> Dict[str, Any]: '''creates a job spec dictionary for a local job''' # Without the unbuffered environment variable, stderr and stdout won't be # emitted in the proper order from inside the container. terminal_cmds = ["-e", "PYTHONUNBUFFERED=1"] + window_size_env_cmds() base_cmd = _run_cmd(job_mode, run_args) + terminal_cmds + [image_id] launcher_args = um.mlflow_args( caliban_config=caliban_config, experiment_name=experiment.xgroup.name, index=index, tags={ um.GPU_ENABLED_TAG: str(job_mode == c.JobMode.GPU).lower(), um.TPU_ENABLED_TAG: 'false', um.DOCKER_IMAGE_TAG: image_id, um.PLATFORM_TAG: Platform.LOCAL.value, }, ) cmd_args = ce.experiment_to_args(experiment.kwargs, experiment.args) # cmd args *must* be last in order for the launcher to pass them through command = base_cmd + launcher_args + cmd_args return {'command': command, 'container': image_id}
def _print_flags(expanded: List[ce.Experiment]) -> None: """Print the flags associated with each experiment in the supplied expansion list. """ for m in expanded: flags = ce.experiment_to_args(m) print(' '.join(flags))
def log_job_spec_instance(job_spec: JobSpec, i: int) -> JobSpec: """Prints logging as a side effect for the supplied sequence of job specs generated from an experiment definition; returns the input job spec. """ args = ce.experiment_to_args(job_spec.experiment.kwargs, job_spec.experiment.args) logging.info("") logging.info("Job {} - Experiment args: {}".format(i, t.yellow(str(args)))) return job_spec
def local_callback(idx: int, job: Job) -> None: """Provides logging feedback for jobs run locally. If the return code is 0, logs success; else, logs the failure as an error and logs the script args that provided the failure. """ if job.status == JobStatus.SUCCEEDED: logging.info(t.green(f'Job {idx} succeeded!')) else: logging.error( t.red( f'Job {idx} failed with return code {job.details["ret_code"]}.' )) args = ce.experiment_to_args(job.spec.experiment.kwargs, job.spec.experiment.args) logging.error(t.red(f'Failing args for job {idx}: {args}'))
def _create_job_spec_dict( experiment: Experiment, job_mode: c.JobMode, image_id: str, run_args: Optional[List[str]] = None, ) -> Dict[str, Any]: '''creates a job spec dictionary for a local job''' # Without the unbuffered environment variable, stderr and stdout won't be # emitted in the proper order from inside the container. terminal_cmds = ["-e", "PYTHONUNBUFFERED=1"] + window_size_env_cmds() base_cmd = _run_cmd(job_mode, run_args) + terminal_cmds + [image_id] command = base_cmd + ce.experiment_to_args(experiment.kwargs, experiment.args) return {'command': command, 'container': image_id}
def _job_specs( job_name: str, training_input: Dict[str, Any], labels: Dict[str, str], experiments: Iterable[ht.Experiment], caliban_config: Optional[Dict[str, Any]] = None, ) -> Iterable[ht.JobSpec]: """Returns a generator that yields a JobSpec instance for every possible combination of parameters in the supplied experiment config. All other arguments parametrize every JobSpec that's generated; labels, arguments and job id change for each JobSpec. This is lower-level than build_job_specs below. """ caliban_config = caliban_config or {} for idx, m in enumerate(experiments, 1): launcher_args = um.mlflow_args( caliban_config=caliban_config, experiment_name=m.xgroup.name, index=idx, tags={ um.PLATFORM_TAG: ht.Platform.CAIP.value, **labels, }, ) cmd_args = ce.experiment_to_args(m.kwargs, m.args) # cmd args *must* be last in order for the launcher to pass them through args = launcher_args + cmd_args yield _job_spec(job_name=job_name, idx=idx, training_input={ **training_input, "args": args }, labels=labels, experiment=m)
def _job_specs( job_name: str, training_input: Dict[str, Any], labels: Dict[str, str], experiments: Iterable[ht.Experiment], ) -> Iterable[ht.JobSpec]: """Returns a generator that yields a JobSpec instance for every possible combination of parameters in the supplied experiment config. All other arguments parametrize every JobSpec that's generated; labels, arguments and job id change for each JobSpec. This is lower-level than build_job_specs below. """ for idx, m in enumerate(experiments, 1): args = ce.experiment_to_args(m.kwargs, m.args) yield _job_spec(job_name=job_name, idx=idx, training_input={ **training_input, "args": args }, labels=labels, experiment=m)
def create_simple_job_spec( self, experiment: Experiment, name: str, image: str, min_cpu: int, min_mem: int, command: Optional[List[str]] = None, env: Dict[str, str] = {}, accelerator: Optional[Accelerator] = None, accelerator_count: int = 1, namespace: str = k.DEFAULT_NAMESPACE, machine_type: Optional[MachineType] = None, preemptible: bool = True, preemptible_tpu: bool = True, tpu_driver: str = k.DEFAULT_TPU_DRIVER) -> Optional[JobSpec]: """creates a simple kubernetes job (1 container, 1 pod) JobSpec for this cluster Args: name: job name image: container image url (gcr.io/...) min_cpu: minimum cpu needed, in milli-cpu min_mem: minimum memory needed, in MB command: command to execute, None = container entrypoint args: args to pass to command env: environment vars for container accelerator: accelerator type, None=cpu only accelerator_count: accelerator count namespace: kubernetes namespace machine_type: machine type, None=default for mode (cpu/gpu) preemptible: use preemptible instance preemptible_tpu: use preemptible tpus tpu_driver: tpu driver to use Returns: JobSpec on success, None otherwise """ args = ce.experiment_to_args(experiment.kwargs, experiment.args) # ------------------------------------------------------------------------ # container # tpu/gpu resources container_resources = V1ResourceRequirements( requests=Cluster.container_requests(min_cpu, min_mem), limits=Cluster.container_limits( accelerator, accelerator_count, preemptible_tpu, ), ) container_env = [V1EnvVar(name=k, value=v) for k, v in env.items()] # this is a simple 1-container, 1-pod job, so we just name the # container the same thing (minus the generated suffix) as the job itself container = V1Container( name=name, image=image, command=command, args=args, resources=container_resources, env=container_env, image_pull_policy='Always', ) # ------------------------------------------------------------------------ # template # todo: should we support anything other than a 'never' restart policy? # see this for discussion # https://kubernetes.io/docs/concepts/workloads/controllers/jobs-run-to-completion/#pod-backoff-failure-policy tolerations = Cluster.tolerations(preemptible=preemptible) # backoff count plus 'OnFailure' may be correct here template_spec = V1PodSpec( restart_policy='Never', containers=[container], tolerations=tolerations, node_selector=Cluster.node_selector( preemptible=preemptible, machine_type=machine_type, accelerator=accelerator, ), host_ipc=True, ) template = V1PodTemplateSpec( metadata=Cluster.template_metadata( accelerator=accelerator, tpu_driver=tpu_driver, ), spec=template_spec, ) # ------------------------------------------------------------------------ # job job_spec = V1JobSpec(template=template, backoff_limit=4) return JobSpec.get_or_create( experiment=experiment, spec=ApiClient().sanitize_for_serialization(job_spec), platform=Platform.GKE, )