"""Mloq docker command implementation."""
from pathlib import Path
from typing import Optional
import click
from omegaconf import DictConfig, MISSING, OmegaConf
from mloq.command import Command
from mloq.commands.requirements import ( # REQUIREMENT_CHOICES,
pytorch_req,
RequirementsCMD,
tensorflow_req,
)
from mloq.config.param_patch import param
from mloq.files import ASSETS_PATH, file
DOCKER_ASSETS_PATH = ASSETS_PATH / "docker"
dockerfile = file("Dockerfile", DOCKER_ASSETS_PATH, description="Docker container for the project")
makefile_docker = file(
"Makefile.docker",
DOCKER_ASSETS_PATH,
description="Makefile for the Docker container setup",
is_static=True,
)
DOCKER_FILES = [dockerfile, makefile_docker]
[docs]class DockerCMD(Command):
"""Implement the functionality of the docker Command."""
cmd_name = "docker"
files = tuple(DOCKER_FILES)
disable = param.Boolean(default=None, doc="Disable docker command?")
cuda = param.Boolean(None, doc="Install CUDA?")
cuda_image_type = param.String(MISSING, doc="Type of cuda docker container")
cuda_version = param.String("11.2", doc="CUDA version installed in the container")
ubuntu_version = param.String("20.04", doc="Ubuntu version of the base image")
project_name = param.String("${globals.project_name}", doc="Select project name")
docker_org = param.String("${globals.owner}", doc="Name of your Docker organization")
python_version = param.String("3.8", doc="Python version installed in the container")
base_image = param.String(MISSING, doc="Base Docker image used to build the container")
test = param.Boolean(True, doc="Install requirements-test.txt?")
lint = param.Boolean(True, doc="Install requirements-lint.txt?")
jupyter = param.Boolean(True, doc="Install a jupyter notebook server?")
jupyter_password = param.String(
"${docker.project_name}",
doc="password for the Jupyter notebook server",
)
requirements = param.List(default=["none"], doc="Project requirements")
extra = param.String("", doc="Extra code to add to Dockerfile")
makefile = param.Boolean(True, doc="Add docker commands to makefile")
# requirements = param.ListSelector(
# default="none", doc="Project requirements", objects=REQUIREMENT_CHOICES,
# )
[docs] @staticmethod
def require_cuda_from_requirements(project_config: Optional[DictConfig] = None) -> bool:
"""Return True if any of the project dependencies require CUDA."""
project_config = {} if project_config is None else project_config
if "requirements" not in project_config:
return False
options = project_config.get("requirements", [])
if RequirementsCMD.requirements_is_empty(options):
return False
elif isinstance(options, str):
options = [options]
tf_alias = RequirementsCMD.REQUIREMENTS_ALIASES[tensorflow_req]
torch_alias = RequirementsCMD.REQUIREMENTS_ALIASES[pytorch_req]
for option in options:
if option in tf_alias or option in torch_alias:
return True
return False
[docs] def requires_cuda(self) -> bool:
"""Return True if the Docker container requires CUDA."""
if not OmegaConf.is_missing(self.config, "cuda") and self.cuda is not None:
return self.config.cuda
try:
_ = self.config.requirements
except Exception:
self.config.requirements = []
return self.require_cuda_from_requirements(self.config)
[docs] def get_base_image(self):
"""Return the name of the base image for the project Docker container."""
# Check value is not missing and it is not None due to a bad interpolation resolution
if (
not OmegaConf.is_missing(self.config, "base_image")
and self.config.base_image is not None
):
return self.config.base_image
elif self.config.cuda:
cuda_image = (
f"nvidia/cuda:{self.config.cuda_version}-{self.config.cuda_image_type}"
f"-ubuntu{self.config.ubuntu_version}"
)
return cuda_image
return f"ubuntu:{self.config.ubuntu_version}"
[docs] def parse_config(self) -> DictConfig:
"""Update the configuration dictionary from the data entered by the user."""
super(DockerCMD, self).parse_config()
if self.cuda is None:
self.cuda = self.requires_cuda()
self.base_image = self.get_base_image()
return super(DockerCMD, self).parse_config()
[docs] def interactive_config(self) -> DictConfig:
"""Generate the configuration of the project interactively."""
click.echo("Provide the values to generate the project Docker container.")
return self.parse_config()
[docs] def record_files(self) -> None:
"""Register the files that will be generated by mloq."""
self.record.register_file(file=dockerfile, path=Path())
self.record.register_file(file=makefile_docker, path=Path())