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

@@ -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()
 

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)