Fix logic when defining enviroment variable

mlcommons · Nov 28, 2023 · 7a5fea8 · 7a5fea8
1 parent 81aebee
commit 7a5fea8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/mlcube/mlcube/parser.py b/mlcube/mlcube/parser.py
@@ -128,9 +128,6 @@ def parse_extra_arg(
             if parsed_args.get("gpus", None):
                 if platform == "docker":
                     runner_run_args["--gpus"] = parsed_args["gpus"]
-                    os.environ["CUDA_VISIBLE_DEVICES"] = parsed_args[
-                        "gpus"
-                    ]
                 else:
                     runner_run_args["--nv"] = ""
                     os.environ["SINGULARITYENV_CUDA_VISIBLE_DEVICES"] = parsed_args[

diff --git a/runners/mlcube_docker/mlcube_docker/docker_run.py b/runners/mlcube_docker/mlcube_docker/docker_run.py
@@ -266,6 +266,9 @@ def run(self) -> None:
         if extra_args:
             run_args += " " + extra_args
 
+        cuda_visible_devices = self.mlcube.runner["--gpus"] if "--gpus" in self.mlcube.runner else num_gpus
+        run_args += f" --env CUDA_VISIBLE_DEVICES={cuda_visible_devices}"
+
         if "entrypoint" in self.mlcube.tasks[self.task]:
             logger.info(
                 "Using custom task entrypoint: task=%s, entrypoint='%s'",