From 7a5fea8d5a9c23bc97c8eeb08b992b66f9530289 Mon Sep 17 00:00:00 2001 From: David Jurado Date: Tue, 28 Nov 2023 16:03:12 -0500 Subject: [PATCH] Fix logic when defining enviroment variable --- mlcube/mlcube/parser.py | 3 --- runners/mlcube_docker/mlcube_docker/docker_run.py | 3 +++ 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/mlcube/mlcube/parser.py b/mlcube/mlcube/parser.py index 0ef794d..17ee426 100644 --- 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 index a8f918f..d2c5870 100644 --- 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'",