Fix issue #182

CHANGELOG.md

@@ -12,6 +12,7 @@ Changelog
 * Allow `oversubscription_factor` to change on restart [#177](https://github.com/OpenBioSim/somd2/pull/177).
 * Restrict energy component decomposition to force groups that are used for integration [#180](https://github.com/OpenBioSim/somd2/pull/180).
 * Parallelise replica mixing [#181](https://github.com/OpenBioSim/somd2/pull/181).
+* Fixed the replica exchange GPU memory check querying the wrong device when `CUDA_VISIBLE_DEVICES` does not start at zero, since OpenMM numbers devices relative to the visible set whereas `pynvml` enumerates all of them [#182](https://github.com/OpenBioSim/somd2/issues/182).
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------

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.
@@ -619,24 +632,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:
@@ -666,15 +720,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
 
@@ -692,9 +751,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
 
@@ -764,6 +821,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]
+
         # Auto-generate a Boresch restraint for ABFE runs with no user-supplied
         # restraint. This must happen before the dynamics cache is built below,
         # since the per-replica OpenMM contexts it creates are fixed at
@@ -851,6 +912,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:
@@ -921,6 +983,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,

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}