diff --git a/runners/mlcube_docker/mlcube_docker/docker_run.py b/runners/mlcube_docker/mlcube_docker/docker_run.py index 0f64d4e..1e2cf9f 100644 --- a/runners/mlcube_docker/mlcube_docker/docker_run.py +++ b/runners/mlcube_docker/mlcube_docker/docker_run.py @@ -1,23 +1,29 @@ +import json +import logging import os import shlex -import logging import typing as t from pathlib import Path + from omegaconf import DictConfig, OmegaConf -from mlcube.shell import Shell -from mlcube.runner import Runner, RunnerConfig -from mlcube.errors import IllegalParameterValueError, ExecutionError, ConfigurationError +from mlcube.errors import ( + ConfigurationError, + ExecutionError, + IllegalParameterValueError, + MLCubeError, +) +from mlcube.runner import Runner, RunnerConfig +from mlcube.shell import Shell +from mlcube.validate import Validate __all__ = ["Config", "DockerRun"] -from mlcube.validate import Validate - logger = logging.getLogger(__name__) class Config(RunnerConfig): - """ Helper class to manage `docker` environment configuration.""" + """Helper class to manage `docker` environment configuration.""" class BuildStrategy(object): """MLCube docker runner configuration strategy. @@ -96,7 +102,7 @@ def merge(mlcube: DictConfig) -> None: @staticmethod def validate(mlcube: DictConfig) -> None: - """ Initialize configuration from user config + """Initialize configuration from user config Args: mlcube: MLCube `container` configuration, possible merged with user local configuration. Return: @@ -120,7 +126,7 @@ def validate(mlcube: DictConfig) -> None: class DockerRun(Runner): - """ Docker runner. """ + """Docker runner.""" CONFIG = Config @@ -192,7 +198,7 @@ def configure(self) -> None: ) def run(self) -> None: - """ Run a cube. """ + """Run a cube.""" docker: t.Text = self.mlcube.runner.docker image: t.Text = self.mlcube.runner.image @@ -228,7 +234,7 @@ def run(self) -> None: ) if mounts_opts: for host_path, mount_type in mounts_opts.items(): - mounts[host_path] += f':{mount_type}' + mounts[host_path] += f":{mount_type}" except ConfigurationError as err: raise ExecutionError.mlcube_run_error( self.__class__.__name__, @@ -240,16 +246,20 @@ def run(self) -> None: volumes = Shell.to_cli_args(mounts, sep=":", parent_arg="--volume") env_args = self.mlcube.runner.env_args - num_gpus: int = self.mlcube.get("platform", {}).get( - "accelerator_count", None - ) or 0 + num_gpus: int = ( + self.mlcube.get("platform", {}).get("accelerator_count", None) or 0 + ) - run_args: t.Text = self.mlcube.runner.cpu_args if num_gpus == 0 else self.mlcube.runner.gpu_args + run_args: t.Text = ( + self.mlcube.runner.cpu_args + if num_gpus == 0 + else self.mlcube.runner.gpu_args + ) extra_args_list = [ f"{key}={value}" for key, value in self.mlcube.runner.items() - if key.startswith('--') and value is not None + if key.startswith("--") and value is not None ] extra_args = " ".join(extra_args_list) if extra_args: @@ -339,3 +349,31 @@ def run(self) -> None: f"volumes={volumes}, image={image}, task_args={task_args}).", **err.context, ) + + def inspect(self, force: bool = False) -> t.Dict: + docker: str = self.mlcube.runner.docker + image: str = self.mlcube.runner.image + if not Shell.docker_image_exists(docker, image): + if not force: + raise MLCubeError( + "MLCube does not exist. Either configure MLCube (e.g., `mlcube configure ...`) or set `force` " + "argument to True (e.g., `mlcube inspect --force ...`)." + ) + self.configure() + + docker_inspect_cmd = [docker, "inspect", image] + exit_code, output = Shell.run_and_capture_output(docker_inspect_cmd) + if exit_code != 0: + raise MLCubeError(output) + + image_info: t.List[t.Dict] = json.loads(output) + if ( + not isinstance(image_info, list) + or len(image_info) != 1 + or not isinstance(image_info[0], dict) + ): + raise MLCubeError( + f"Unexpected output from `{' '.join(docker_inspect_cmd)}`. Expected a list of dicts of length 1." + ) + + return {"hash": image_info[0].get("Id", None)} diff --git a/runners/mlcube_docker/mlcube_docker/tests/test_docker_runner.py b/runners/mlcube_docker/mlcube_docker/tests/test_docker_runner.py index 32aa0b4..20389be 100644 --- a/runners/mlcube_docker/mlcube_docker/tests/test_docker_runner.py +++ b/runners/mlcube_docker/mlcube_docker/tests/test_docker_runner.py @@ -1,15 +1,15 @@ +import typing as t import unittest from unittest import TestCase -from unittest.mock import (mock_open, patch) +from unittest.mock import mock_open, patch + +from mlcube_docker.docker_run import Config, DockerRun +from omegaconf import DictConfig, OmegaConf from mlcube.config import MLCubeConfig from mlcube.shell import Shell -from mlcube_docker.docker_run import (Config, DockerRun) - -from omegaconf import DictConfig, OmegaConf - -_HAVE_DOCKER: bool = Shell.run(['docker', '--version'], on_error='ignore') == 0 +_HAVE_DOCKER: bool = Shell.run(["docker", "--version"], on_error="ignore") == 0 _MLCUBE_DEFAULT_ENTRY_POINT = """ docker: @@ -29,6 +29,10 @@ class TestDockerRunner(TestCase): + def _check_inspect_output(self, info: t.Dict) -> None: + self.assertIsInstance(info, dict) + self.assertIn("hash", info) + self.assertTrue(info["hash"].startswith("sha256:")) @staticmethod def noop(*args, **kwargs) -> None: @@ -54,36 +58,45 @@ def tearDown(self) -> None: def test_mlcube_default_entrypoints(self): with patch("io.open", mock_open(read_data=_MLCUBE_DEFAULT_ENTRY_POINT)): mlcube: DictConfig = MLCubeConfig.create_mlcube_config( - "/some/path/to/mlcube.yaml", runner_config=Config.DEFAULT, runner_cls=DockerRun + "/some/path/to/mlcube.yaml", + runner_config=Config.DEFAULT, + runner_cls=DockerRun, ) - self.assertEqual(mlcube.runner.image, 'ubuntu:18.04') + self.assertEqual(mlcube.runner.image, "ubuntu:18.04") self.assertDictEqual( OmegaConf.to_container(mlcube.tasks), { - 'ls': {'parameters': {'inputs': {}, 'outputs': {}}}, - 'pwd': {'parameters': {'inputs': {}, 'outputs': {}}} - } + "ls": {"parameters": {"inputs": {}, "outputs": {}}}, + "pwd": {"parameters": {"inputs": {}, "outputs": {}}}, + }, ) DockerRun(mlcube, task=None).configure() - DockerRun(mlcube, task='ls').run() - DockerRun(mlcube, task='pwd').run() + self._check_inspect_output(DockerRun(mlcube, task=None).inspect()) + DockerRun(mlcube, task="ls").run() + DockerRun(mlcube, task="pwd").run() @unittest.skipUnless(_HAVE_DOCKER, reason="No docker available.") def test_mlcube_custom_entrypoints(self): with patch("io.open", mock_open(read_data=_MLCUBE_CUSTOM_ENTRY_POINTS)): mlcube: DictConfig = MLCubeConfig.create_mlcube_config( - "/some/path/to/mlcube.yaml", runner_config=Config.DEFAULT, runner_cls=DockerRun + "/some/path/to/mlcube.yaml", + runner_config=Config.DEFAULT, + runner_cls=DockerRun, ) - self.assertEqual(mlcube.runner.image, 'ubuntu:18.04') + self.assertEqual(mlcube.runner.image, "ubuntu:18.04") self.assertDictEqual( OmegaConf.to_container(mlcube.tasks), { - 'ls': {'parameters': {'inputs': {}, 'outputs': {}}}, - 'free': {'entrypoint': '/usr/bin/free', 'parameters': {'inputs': {}, 'outputs': {}}} - } + "ls": {"parameters": {"inputs": {}, "outputs": {}}}, + "free": { + "entrypoint": "/usr/bin/free", + "parameters": {"inputs": {}, "outputs": {}}, + }, + }, ) DockerRun(mlcube, task=None).configure() - DockerRun(mlcube, task='ls').run() - DockerRun(mlcube, task='free').run() + self._check_inspect_output(DockerRun(mlcube, task=None).inspect()) + DockerRun(mlcube, task="ls").run() + DockerRun(mlcube, task="free").run()