diff --git a/openmmtools/multistate/multistatesampler.py b/openmmtools/multistate/multistatesampler.py index 31d49ffa..5f6594cc 100644 --- a/openmmtools/multistate/multistatesampler.py +++ b/openmmtools/multistate/multistatesampler.py @@ -35,6 +35,7 @@ import inspect import logging import datetime +import subprocess import numpy as np @@ -1773,11 +1774,19 @@ def _update_timing(self, iteration_time, partial_total_time, run_initial_iterati @staticmethod def _display_cuda_devices(): """Query system nvidia-smi to get available GPUs indices and names in debug log.""" - # Read nvidia-smi query, should return empty strip if no GPU is found. - cuda_query_output = os.popen("nvidia-smi --query-gpu=index,gpu_name --format=csv,noheader").read().strip() - # Split by line jump and comma - cuda_devices_list = [entry.split(',') for entry in cuda_query_output.split('\n')] - logger.debug(f"CUDA devices available: {*cuda_devices_list,}") + + cuda_query_output = subprocess.run("nvidia-smi --query-gpu=gpu_uuid,gpu_name,compute_mode --format=csv", shell=True, capture_output=True, text=True) + # Check if command worked + if cuda_query_output.returncode == 0: + # Split by line jump and comma + cuda_devices_list = [entry for entry in cuda_query_output.stdout.splitlines()] + logger.debug(f"CUDA devices available: {*cuda_devices_list,}") + # We only support "Default" and not "Exclusive_Process" for the compute mode + if "Default" not in cuda_query_output.stdout: + logger.warning(f"GPU in 'Exclusive_Process' mode (or Prohibited), one context is allowed per device. This may prevent some openmmtools features from working. GPU must be in 'Default' compute mode") + # Handel the case where the command had some error + else: + logger.debug(f"nvidia-smi command failed: {cuda_query_output.stderr}, this is expected if there is no GPU available") def _flatten_moves_iterator(self): """Recursively flatten MCMC moves. Handles the cases where each move can be a set of moves, for example with