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
