From 5929cfdf3ebe0e21ca0e53ef5494dfa9c29fd229 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ole=20Sch=C3=BCtt?= Date: Fri, 2 Jul 2021 13:25:13 +0200 Subject: [PATCH] regtesting: Set HIP_VISIBLE_DEVICES --- tools/regtesting/do_regtest.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tools/regtesting/do_regtest.py b/tools/regtesting/do_regtest.py index 116440f4b7..16bb635d7f 100755 --- a/tools/regtesting/do_regtest.py +++ b/tools/regtesting/do_regtest.py @@ -198,8 +198,11 @@ class Config: self.debug = args.debug self.max_errors = args.maxerrors self.restrictdirs = args.restrictdir if args.restrictdir else [".*"] - cmd = "nvidia-smi --query-gpu=gpu_name --format=csv,noheader | wc -l" - self.num_gpus = int(subprocess.run(cmd, shell=True, capture_output=True).stdout) + nv_cmd = "nvidia-smi --query-gpu=gpu_name --format=csv,noheader | wc -l" + nv_gpus = int(subprocess.run(nv_cmd, shell=True, capture_output=True).stdout) + amd_cmd = "rocm-smi --showid --csv | grep card | wc -l" + amd_gpus = int(subprocess.run(amd_cmd, shell=True, capture_output=True).stdout) + self.num_gpus = nv_gpus + amd_gpus self.next_gpu = 0 # Used to assign devices round robin to processes. def launch_exe( @@ -207,10 +210,12 @@ class Config: ) -> Coroutine[Any, Any, Process]: env = os.environ.copy() if self.num_gpus > self.mpiranks: - env["CUDA_VISIBLE_DEVICES"] = "" + visible_gpu_devices = [] for _ in range(self.mpiranks): # Utilize all available GPU devices. self.next_gpu = (self.next_gpu + 1) % self.num_gpus - env["CUDA_VISIBLE_DEVICES"] += f"{self.next_gpu}," + visible_gpu_devices.append(f"{self.next_gpu}") + env["CUDA_VISIBLE_DEVICES"] = ",".join(visible_gpu_devices) + env["HIP_VISIBLE_DEVICES"] = ",".join(visible_gpu_devices) env["OMP_NUM_THREADS"] = str(self.ompthreads) exe = self.cp2k_root / "exe" / self.arch / f"{exe_stem}.{self.version}" cmd = ["mpiexec", f"-np={self.mpiranks}", exe] if self.use_mpi else [exe]