def mac_gpu_check(job_mode: conf.JobMode, command: str) -> None: """If the command depends on 'docker run' and is running on a Mac, fail fast.""" if conf.gpu(job_mode) and command in ("shell", "notebook", "run"): u.err("\n'caliban {}' doesn't support GPU usage on Macs! Please pass ". format(command) + "--nogpu to use this command.\n\n") u.err( "(GPU mode is fine for 'caliban cloud' from a Mac; just nothing that runs " "locally.)\n\n") sys.exit(1)
def tf_base_image(job_mode: c.JobMode, tensorflow_version: str) -> str: """Returns the base image to use, depending on whether or not we're using a GPU. This is JUST for building our base images for Blueshift; not for actually using in a job. List of available tags: https://hub.docker.com/r/tensorflow/tensorflow/tags """ if tensorflow_version not in TF_VERSIONS: raise Exception("""{} is not a valid tensorflow version. Try one of: {}""".format(tensorflow_version, TF_VERSIONS)) gpu = "-gpu" if c.gpu(job_mode) else "" return "tensorflow/tensorflow:{}{}-py3".format(tensorflow_version, gpu)
def _run_cmd(job_mode: c.JobMode, run_args: Optional[List[str]] = None) -> List[str]: """Returns the sequence of commands for the subprocess run functions required to execute `docker run`. in CPU or GPU mode, depending on the value of job_mode. Keyword args: - run_args: list of args to pass to docker run. """ if run_args is None: run_args = [] runtime = ["--runtime", "nvidia"] if c.gpu(job_mode) else [] return ["docker", "run"] + runtime + ["--ipc", "host"] + run_args
def base_extras(job_mode: c.JobMode, path: str, extras: Optional[List[str]]) -> Optional[List[str]]: """Returns None if the supplied path doesn't exist (it's assumed it points to a setup.py file). If the path DOES exist, generates a list of extras to install. gpu or cpu are always added to the beginning of the list, depending on the mode. """ ret = None if os.path.exists(path): base = extras or [] extra = 'gpu' if c.gpu(job_mode) else 'cpu' ret = base if extra in base else [extra] + base return ret
def test_gpu(): assert c.gpu(c.JobMode.GPU) assert not c.gpu(c.JobMode.CPU) assert not c.gpu("face")
def base_image_suffix(job_mode: c.JobMode) -> str: return "gpu" if c.gpu(job_mode) else "cpu"
def base_image_suffix(job_mode: c.JobMode) -> str: return DEFAULT_GPU_TAG if c.gpu(job_mode) else DEFAULT_CPU_TAG