From 5fb056c72f74614c1a08dcb80c6a37560c0c3a1d Mon Sep 17 00:00:00 2001 From: hasan7n <78664424+hasan7n@users.noreply.github.com> Date: Fri, 8 Dec 2023 18:19:31 +0100 Subject: [PATCH] Use current user ID when running MLCubes (#512) * use current user ID when running docker container * fix conflicts in unittests --- cli/medperf/entities/cube.py | 38 ++++++++++++++++++++----- cli/medperf/tests/entities/test_cube.py | 35 +++++++++++++++++++++++ 2 files changed, 66 insertions(+), 7 deletions(-) diff --git a/cli/medperf/entities/cube.py b/cli/medperf/entities/cube.py index 2e33f6406..caa04763e 100644 --- a/cli/medperf/entities/cube.py +++ b/cli/medperf/entities/cube.py @@ -301,6 +301,14 @@ def run( for k, v in kwargs.items(): cmd_arg = f'{k}="{v}"' cmd = " ".join([cmd, cmd_arg]) + + cpu_args = self.get_config("docker.cpu_args") or "" + gpu_args = self.get_config("docker.gpu_args") or "" + cpu_args = " ".join([cpu_args, "-u $(id -u):$(id -g)"]).strip() + gpu_args = " ".join([gpu_args, "-u $(id -u):$(id -g)"]).strip() + cmd += f' -Pdocker.cpu_args="{cpu_args}"' + cmd += f' -Pdocker.gpu_args="{gpu_args}"' + logging.info(f"Running MLCube command: {cmd}") proc = pexpect.spawn(cmd, timeout=timeout) proc_out = combine_proc_sp_text(proc) @@ -328,14 +336,10 @@ def get_default_output(self, task: str, out_key: str, param_key: str = None) -> str: the path as specified in the mlcube.yaml file for the desired output for the desired task. Defaults to None if out_key not found """ - with open(self.cube_path, "r") as f: - cube = yaml.safe_load(f) + out_path = self.get_config(f"tasks.{task}.parameters.outputs.{out_key}") + if out_path is None: + return - out_params = cube["tasks"][task]["parameters"]["outputs"] - if out_key not in out_params: - return None - - out_path = cube["tasks"][task]["parameters"]["outputs"][out_key] if isinstance(out_path, dict): # output is specified as a dict with type and default values out_path = out_path["default"] @@ -350,6 +354,26 @@ def get_default_output(self, task: str, out_key: str, param_key: str = None) -> return out_path + def get_config(self, identifier): + """ + Returns the output parameter specified in the mlcube.yaml file + + Args: + identifier (str): `.` separated keys to traverse the mlcube dict + Returns: + str: the parameter value, None if not found + """ + with open(self.cube_path, "r") as f: + cube = yaml.safe_load(f) + + keys = identifier.split(".") + for key in keys: + if key not in cube: + return + cube = cube[key] + + return cube + def todict(self) -> Dict: return self.extended_dict() diff --git a/cli/medperf/tests/entities/test_cube.py b/cli/medperf/tests/entities/test_cube.py index eb0196a44..a2a147ef8 100644 --- a/cli/medperf/tests/entities/test_cube.py +++ b/cli/medperf/tests/entities/test_cube.py @@ -160,9 +160,12 @@ def test_cube_runs_command(self, mocker, timeout, setup, task): spy = mocker.patch( PATCH_CUBE.format("pexpect.spawn"), side_effect=mpexpect.spawn ) + mocker.patch(PATCH_CUBE.format("Cube.get_config"), side_effect=["", ""]) expected_cmd = ( f"mlcube run --mlcube={self.manifest_path} --task={task} " + f"--platform={self.platform} --network=none --mount=ro" + + ' -Pdocker.cpu_args="-u $(id -u):$(id -g)"' + + ' -Pdocker.gpu_args="-u $(id -u):$(id -g)"' ) # Act @@ -176,9 +179,15 @@ def test_cube_runs_command_with_rw_access(self, mocker, setup, task): # Arrange mpexpect = MockPexpect(0, "expected_hash") spy = mocker.patch("pexpect.spawn", side_effect=mpexpect.spawn) + mocker.patch( + PATCH_CUBE.format("Cube.get_config"), + side_effect=["", ""], + ) expected_cmd = ( f"mlcube run --mlcube={self.manifest_path} --task={task} " + f"--platform={self.platform} --network=none" + + ' -Pdocker.cpu_args="-u $(id -u):$(id -g)"' + + ' -Pdocker.gpu_args="-u $(id -u):$(id -g)"' ) # Act @@ -192,9 +201,12 @@ def test_cube_runs_command_with_extra_args(self, mocker, setup, task): # Arrange mpexpect = MockPexpect(0, "expected_hash") spy = mocker.patch("pexpect.spawn", side_effect=mpexpect.spawn) + mocker.patch(PATCH_CUBE.format("Cube.get_config"), side_effect=["", ""]) expected_cmd = ( f"mlcube run --mlcube={self.manifest_path} --task={task} " + f'--platform={self.platform} --network=none --mount=ro test="test"' + + ' -Pdocker.cpu_args="-u $(id -u):$(id -g)"' + + ' -Pdocker.gpu_args="-u $(id -u):$(id -g)"' ) # Act @@ -204,10 +216,33 @@ def test_cube_runs_command_with_extra_args(self, mocker, setup, task): # Assert spy.assert_any_call(expected_cmd, timeout=None) + def test_cube_runs_command_and_preserves_runtime_args(self, mocker, setup, task): + # Arrange + mpexpect = MockPexpect(0, "expected_hash") + spy = mocker.patch("pexpect.spawn", side_effect=mpexpect.spawn) + mocker.patch( + PATCH_CUBE.format("Cube.get_config"), + side_effect=["cpuarg cpuval", "gpuarg gpuval"], + ) + expected_cmd = ( + f"mlcube run --mlcube={self.manifest_path} --task={task} " + + f"--platform={self.platform} --network=none --mount=ro" + + ' -Pdocker.cpu_args="cpuarg cpuval -u $(id -u):$(id -g)"' + + ' -Pdocker.gpu_args="gpuarg gpuval -u $(id -u):$(id -g)"' + ) + + # Act + cube = Cube.get(self.id) + cube.run(task) + + # Assert + spy.assert_any_call(expected_cmd, timeout=None) + def test_run_stops_execution_if_child_fails(self, mocker, setup, task): # Arrange mpexpect = MockPexpect(1, "expected_hash") mocker.patch("pexpect.spawn", side_effect=mpexpect.spawn) + mocker.patch(PATCH_CUBE.format("Cube.get_config"), side_effect=["", ""]) # Act & Assert cube = Cube.get(self.id)