diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 7da3ba3..025552b 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -54,6 +54,7 @@ def __init__( output_directory=None, perturbed_system=None, xml_filenames=None, + gpu_devices=None, ): """ Constructor. @@ -89,6 +90,11 @@ def __init__( xml_filenames: list of str A list of file paths for the OpenMM XML output, one per replica. If None, XML files are not written. + + gpu_devices: list + The physical devices backing each OpenMM device index, i.e. the + entries of CUDA_VISIBLE_DEVICES. Used to query the memory of the + right device. If None, the OpenMM index is used directly. """ # Warn if the number of replicas is not a multiple of the number of GPUs. @@ -101,6 +107,7 @@ def __init__( # Initialise attributes. self._lambdas = lambdas self._rest2_scale_factors = rest2_scale_factors + self._gpu_devices = gpu_devices self._states = _np.array(range(len(lambdas))) self._time = None self._openmm_states = [None] * len(lambdas) @@ -250,7 +257,9 @@ def _create_dynamics( # Record baseline memory before the first replica on this device. if device not in device_mem: - used_before, _, total_mem = self._check_device_memory(device) + used_before, _, total_mem = self._check_device_memory( + self._physical_device(device) + ) device_mem[device] = { "before": used_before, "total": total_mem, @@ -335,7 +344,9 @@ def _create_dynamics( # Estimate memory after the first or second replica. if info["count"] == 1: - used_mem, _, _ = self._check_device_memory(device) + used_mem, _, _ = self._check_device_memory( + self._physical_device(device) + ) info["after_first"] = used_mem if num_contexts == 1: @@ -346,7 +357,9 @@ def _create_dynamics( est_total = None elif info["count"] == 2: - used_mem, _, _ = self._check_device_memory(device) + used_mem, _, _ = self._check_device_memory( + self._physical_device(device) + ) # The first replica includes one-time context overhead. # The marginal cost of subsequent replicas is the difference # between the second and first. @@ -595,24 +608,65 @@ def get_swaps(self): """ return self._num_swaps + def _physical_device(self, device): + """ + Map an OpenMM device index to the physical device backing it. + + Parameters + ---------- + + device: int + The OpenMM device index, which is relative to the visible set. + + Returns + ------- + + int, str + The physical device, i.e. the matching entry from + CUDA_VISIBLE_DEVICES. Falls back to the OpenMM index if the visible + set is unknown, which is correct whenever it starts at zero and is + contiguous. + """ + gpu_devices = getattr(self, "_gpu_devices", None) + + if gpu_devices is None or device >= len(gpu_devices): + return device + + return gpu_devices[device] + @staticmethod - def _check_device_memory(device_index=0): + def _check_device_memory(device=0): """ Check the memory usage of the specified GPU device. Parameters ---------- - index: int - The index of the GPU device. + device: int, str + The device to query. This is the physical device, i.e. an entry + from CUDA_VISIBLE_DEVICES (or the equivalent for other platforms), + not the index used by OpenMM. OpenMM numbers devices relative to + the visible set, whereas pynvml and pyopencl enumerate all devices + on the machine, so the two only agree when the visible set starts + at zero and is contiguous. CUDA_VISIBLE_DEVICES entries may be + either an index or a UUID. """ + device = str(device).strip() + + # A UUID cannot be used to index into the OpenCL device list. + is_uuid = device.startswith("GPU-") or device.startswith("MIG-") + device_index = None if is_uuid else int(device) + # Try to use pyopencl to detect the GPU vendor. vendor = None ocl_device = None try: import pyopencl as cl + if device_index is None: + raise ValueError("Cannot index OpenCL devices by UUID") + platforms = cl.get_platforms() all_devices = [] for platform in platforms: @@ -642,15 +696,20 @@ def _check_device_memory(device_index=0): import pynvml pynvml.nvmlInit() - handle = pynvml.nvmlDeviceGetHandleByIndex(device_index) + if is_uuid: + handle = pynvml.nvmlDeviceGetHandleByUUID(device.encode()) + else: + handle = pynvml.nvmlDeviceGetHandleByIndex(device_index) memory = pynvml.nvmlDeviceGetMemoryInfo(handle) pynvml.nvmlShutdown() return (memory.used, memory.free, memory.total) except Exception as e: if vendor is None: - msg = f"Could not get GPU memory info for device {device_index} via OpenCL or pynvml: {e}" + msg = f"Could not get GPU memory info for device {device} via OpenCL or pynvml: {e}" else: - msg = f"Could not get NVIDIA GPU memory info for device {device_index}: {e}" + msg = ( + f"Could not get NVIDIA GPU memory info for device {device}: {e}" + ) _logger.error(msg) raise RuntimeError(msg) from e @@ -668,9 +727,7 @@ def _check_device_memory(device_index=0): used = total - free return (used, free, total) except Exception as e: - msg = ( - f"Could not get AMD GPU memory info for device {device_index}: {e}" - ) + msg = f"Could not get AMD GPU memory info for device {device}: {e}" _logger.error(msg) raise RuntimeError(msg) from e @@ -740,6 +797,10 @@ def __init__(self, system, config): else: self._num_gpus = min(self._config.max_gpus, len(gpu_devices)) + # The physical devices backing each OpenMM device index. OpenMM numbers + # devices relative to the visible set, so index i is gpu_devices[i]. + self._gpu_devices = list(gpu_devices)[: self._num_gpus] + # Store the name of the dynamics cache pickle file. self._repex_state = self._config.output_directory / "repex_state.pkl" @@ -813,6 +874,7 @@ def __init__(self, system, config): perturbed_system=self._perturbed_system, output_directory=self._config.output_directory, xml_filenames=xml_filenames, + gpu_devices=self._gpu_devices, ) else: @@ -883,6 +945,10 @@ def __init__(self, system, config): if not isinstance(self._system, list): self._system.set_time(time) + # The physical device list is not pickled, since the run may be + # restarted against a different set of GPUs. + self._dynamics_cache._gpu_devices = self._gpu_devices + # Create the dynamics objects. self._dynamics_cache._create_dynamics( self._system, diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 9336053..e7b970d 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -148,3 +148,90 @@ def test_rest2_selection(ethane_methanol, rest2_selection, is_valid): else: with pytest.raises(ValueError): runner = RunnerBase(ethane_methanol, Config(**config)) + + +@pytest.mark.parametrize( + "gpu_devices, expected", + [ + # Visible set starts at zero and is contiguous, so the OpenMM index and + # the physical device agree. + (["0", "1", "2"], ["0", "1", "2"]), + # Offset visible set: OpenMM index 0 is physical device 1. + (["1", "2"], ["1", "2"]), + # A single device that is not device zero. + (["3"], ["3"]), + # CUDA_VISIBLE_DEVICES may hold UUIDs rather than indices. + (["GPU-abc123", "GPU-def456"], ["GPU-abc123", "GPU-def456"]), + # Unknown visible set falls back to the OpenMM index. + (None, [0, 1]), + ], +) +def test_physical_device_mapping(gpu_devices, expected): + """ + Validate that an OpenMM device index is mapped to the physical device + backing it. + + OpenMM numbers devices relative to the visible set, whereas pynvml and + pyopencl enumerate every device on the machine. Querying the memory of a + device by its OpenMM index therefore reports the wrong GPU whenever the + visible set does not start at zero. + """ + from somd2.runner._repex import DynamicsCache + + cache = object.__new__(DynamicsCache) + cache._gpu_devices = gpu_devices + + assert [cache._physical_device(i) for i in range(len(expected))] == expected + + +@pytest.mark.parametrize( + "device, key, value", + [ + ("2", "index", 2), + (2, "index", 2), + ("GPU-abc123", "uuid", b"GPU-abc123"), + ], +) +def test_check_device_memory_queries_requested_device(monkeypatch, device, key, value): + """ + Validate that the memory query is made against the device it was asked + for, by index or by UUID. + """ + import sys + import types + + from somd2.runner._repex import DynamicsCache + + pynvml = pytest.importorskip("pynvml") + + # Force the OpenCL branch to fail so that the pynvml path is always taken, + # regardless of what the machine running the tests has installed. + broken = types.SimpleNamespace() + + def get_platforms(): + raise RuntimeError("no OpenCL") + + broken.get_platforms = get_platforms + monkeypatch.setitem(sys.modules, "pyopencl", broken) + + requested = {} + + class Memory: + used, free, total = 1, 2, 3 + + def by_index(index): + requested["index"] = index + return "handle" + + def by_uuid(uuid): + requested["uuid"] = uuid + return "handle" + + monkeypatch.setattr(pynvml, "nvmlInit", lambda: None) + monkeypatch.setattr(pynvml, "nvmlShutdown", lambda: None) + monkeypatch.setattr(pynvml, "nvmlDeviceGetHandleByIndex", by_index) + monkeypatch.setattr(pynvml, "nvmlDeviceGetHandleByUUID", by_uuid) + monkeypatch.setattr(pynvml, "nvmlDeviceGetMemoryInfo", lambda handle: Memory) + + assert DynamicsCache._check_device_memory(device) == (1, 2, 3) + assert requested == {key: value}