Fix parser logic

mlcommons · Dec 19, 2023 · c79b896 · c79b896
1 parent b6bc2df
commit c79b896
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/mlcube/mlcube/parser.py b/mlcube/mlcube/parser.py
@@ -126,13 +126,17 @@ def parse_extra_arg(
                 key = "--security-opt" if platform == "docker" else "--security"
                 runner_run_args[key] = parsed_args["security"]
             if parsed_args.get("gpus", None):
+                cuda_visible_devices = parsed_args["gpus"]
+                if "device" in cuda_visible_devices:
+                    cuda_visible_devices = cuda_visible_devices.replace("device=", "")
+                elif str(cuda_visible_devices).isnumeric():
+                    cuda_visible_devices = str(list(range(int(cuda_visible_devices))))
+                    cuda_visible_devices = cuda_visible_devices.replace(" ", "")[1:-1]
                 if platform == "docker":
-                    runner_run_args["--gpus"] = parsed_args["gpus"]
+                    runner_run_args["--gpus"] = cuda_visible_devices
                 else:
                     runner_run_args["--nv"] = ""
-                    os.environ["SINGULARITYENV_CUDA_VISIBLE_DEVICES"] = parsed_args[
-                        "gpus"
-                    ]
+                    os.environ["SINGULARITYENV_CUDA_VISIBLE_DEVICES"] = cuda_visible_devices
             if parsed_args.get("memory", None):
                 key = "--memory" if platform == "docker" else "--vm-ram"
                 runner_run_args[key] = parsed_args["memory"]

diff --git a/runners/mlcube_docker/mlcube_docker/docker_run.py b/runners/mlcube_docker/mlcube_docker/docker_run.py
@@ -268,11 +268,8 @@ def run(self) -> None:
 
         valid_gpu_flag = "--gpus" in self.mlcube.runner and self.mlcube.runner["--gpus"] is not None
 
-
         if valid_gpu_flag:
             cuda_visible_devices = self.mlcube.runner["--gpus"]
-            if "device" in cuda_visible_devices:
-                cuda_visible_devices = cuda_visible_devices.replace("device=", "")
         else:
             cuda_visible_devices = num_gpus