Support MIG parsing during CUDA context creation in UCX initialization (#6720)

Author: pentschev

codecov.yml

@@ -1,6 +1,11 @@
 codecov:
   require_ci_to_pass: yes
 
+ignore:
+  # Files that exercise GPU-only functionality or are only tested by gpuCI
+  # but don't interact with codecov are ignored.
+  - "distributed/comm/ucx.py"
+
 coverage:
   precision: 2
   round: down

distributed/comm/tests/test_ucx.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import asyncio
+import os
+from unittest.mock import patch
 
 import pytest
 
@@ -15,7 +17,12 @@
 from distributed.comm.core import CommClosedError
 from distributed.comm.registry import backends, get_backend
 from distributed.deploy.local import LocalCluster
-from distributed.diagnostics.nvml import has_cuda_context
+from distributed.diagnostics.nvml import (
+    device_get_count,
+    get_device_index_and_uuid,
+    get_device_mig_mode,
+    has_cuda_context,
+)
 from distributed.protocol import to_serialize
 from distributed.utils_test import gen_test, inc
 
@@ -320,20 +327,40 @@ async def test_simple(
             assert await client.submit(lambda x: x + 1, 10) == 11
 
 
+@pytest.mark.xfail(reason="If running on Docker, requires --pid=host")
 @gen_test()
 async def test_cuda_context(
     ucx_loop,
 ):
-    with dask.config.set({"distributed.comm.ucx.create-cuda-context": True}):
-        async with LocalCluster(
-            protocol="ucx", n_workers=1, asynchronous=True
-        ) as cluster:
-            async with Client(cluster, asynchronous=True) as client:
-                assert cluster.scheduler_address.startswith("ucx://")
-                assert has_cuda_context() == 0
-                worker_cuda_context = await client.run(has_cuda_context)
-                assert len(worker_cuda_context) == 1
-                assert list(worker_cuda_context.values())[0] == 0
+    try:
+        device_info = get_device_index_and_uuid(
+            next(
+                filter(
+                    lambda i: get_device_mig_mode(i)[0] == 0, range(device_get_count())
+                )
+            )
+        )
+    except StopIteration:
+        pytest.skip("No CUDA device in non-MIG mode available")
+
+    with patch.dict(
+        os.environ, {"CUDA_VISIBLE_DEVICES": device_info.uuid.decode("utf-8")}
+    ):
+        with dask.config.set({"distributed.comm.ucx.create-cuda-context": True}):
+            async with LocalCluster(
+                protocol="ucx", n_workers=1, asynchronous=True
+            ) as cluster:
+                async with Client(cluster, asynchronous=True) as client:
+                    assert cluster.scheduler_address.startswith("ucx://")
+                    ctx = has_cuda_context()
+                    assert ctx.has_context and ctx.device_info == device_info
+                    worker_cuda_context = await client.run(has_cuda_context)
+                    assert len(worker_cuda_context) == 1
+                    worker_cuda_context = list(worker_cuda_context.values())
+                    assert (
+                        worker_cuda_context[0].has_context
+                        and worker_cuda_context[0].device_info == device_info
+                    )
 
 
 @gen_test()

distributed/comm/ucx.py

@@ -28,7 +28,11 @@
     host_array,
     to_frames,
 )
-from distributed.diagnostics.nvml import has_cuda_context
+from distributed.diagnostics.nvml import (
+    CudaDeviceInfo,
+    get_device_index_and_uuid,
+    has_cuda_context,
+)
 from distributed.utils import ensure_ip, get_ip, get_ipv6, log_errors, nbytes
 
 logger = logging.getLogger(__name__)
@@ -57,17 +61,27 @@
 )
 
 
-def _warn_existing_cuda_context(ctx, pid):
+def _get_device_and_uuid_str(device_info: CudaDeviceInfo) -> str:
+    return f"{device_info.device_index} ({str(device_info.uuid)})"
+
+
+def _warn_existing_cuda_context(device_info: CudaDeviceInfo, pid: int) -> None:
+    device_uuid_str = _get_device_and_uuid_str(device_info)
     logger.warning(
-        f"A CUDA context for device {ctx} already exists on process ID {pid}. {_warning_suffix}"
+        f"A CUDA context for device {device_uuid_str} already exists "
+        f"on process ID {pid}. {_warning_suffix}"
     )
 
 
-def _warn_cuda_context_wrong_device(ctx_expected, ctx_actual, pid):
+def _warn_cuda_context_wrong_device(
+    device_info_expected: CudaDeviceInfo, device_info_actual: CudaDeviceInfo, pid: int
+) -> None:
+    expected_device_uuid_str = _get_device_and_uuid_str(device_info_expected)
+    actual_device_uuid_str = _get_device_and_uuid_str(device_info_actual)
     logger.warning(
         f"Worker with process ID {pid} should have a CUDA context assigned to device "
-        f"{ctx_expected}, but instead the CUDA context is on device {ctx_actual}. "
-        f"{_warning_suffix}"
+        f"{expected_device_uuid_str}, but instead the CUDA context is on device "
+        f"{actual_device_uuid_str}. {_warning_suffix}"
     )
 
 
@@ -116,22 +130,24 @@ def init_once():
                 "CUDA support with UCX requires Numba for context management"
             )
 
-        cuda_visible_device = int(
+        cuda_visible_device = get_device_index_and_uuid(
             os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")[0]
         )
         pre_existing_cuda_context = has_cuda_context()
-        if pre_existing_cuda_context is not False:
-            _warn_existing_cuda_context(pre_existing_cuda_context, os.getpid())
+        if pre_existing_cuda_context.has_context:
+            _warn_existing_cuda_context(
+                pre_existing_cuda_context.device_info, os.getpid()
+            )
 
         numba.cuda.current_context()
 
         cuda_context_created = has_cuda_context()
         if (
-            cuda_context_created is not False
-            and cuda_context_created != cuda_visible_device
+            cuda_context_created.has_context
+            and cuda_context_created.device_info.uuid != cuda_visible_device.uuid
         ):
             _warn_cuda_context_wrong_device(
-                cuda_visible_device, cuda_context_created, os.getpid()
+                cuda_visible_device, cuda_context_created.device_info, os.getpid()
             )
 
     import ucp as _ucp

distributed/diagnostics/nvml.py

@@ -3,6 +3,7 @@
 import os
 from enum import IntEnum, auto
 from platform import uname
+from typing import NamedTuple
 
 from packaging.version import parse as parse_version
 
@@ -29,6 +30,17 @@ class NVMLState(IntEnum):
     """PyNVML and NVML available, but on WSL and the driver version is insufficient"""
 
 
+class CudaDeviceInfo(NamedTuple):
+    uuid: bytes | None = None
+    device_index: int | None = None
+    mig_index: int | None = None
+
+
+class CudaContext(NamedTuple):
+    has_context: bool
+    device_info: CudaDeviceInfo | None = None
+
+
 # Initialisation must occur per-process, so an initialised state is a
 # (state, pid) pair
 NVML_STATE = (
@@ -147,27 +159,138 @@ def _pynvml_handles():
         return pynvml.nvmlDeviceGetHandleByIndex(gpu_idx)
 
 
+def _running_process_matches(handle):
+    """Check whether the current process is same as that of handle
+
+    Parameters
+    ----------
+    handle : pyvnml.nvml.LP_struct_c_nvmlDevice_t
+        NVML handle to CUDA device
+
+    Returns
+    -------
+    out : bool
+        Whether the device handle has a CUDA context on the running process.
+    """
+    init_once()
+    if hasattr(pynvml, "nvmlDeviceGetComputeRunningProcesses_v2"):
+        running_processes = pynvml.nvmlDeviceGetComputeRunningProcesses_v2(handle)
+    else:
+        running_processes = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
+    return any(os.getpid() == proc.pid for proc in running_processes)
+
+
 def has_cuda_context():
     """Check whether the current process already has a CUDA context created.
 
     Returns
     -------
-    ``False`` if current process has no CUDA context created, otherwise returns the
-    index of the device for which there's a CUDA context.
+    out : CudaContext
+        Object containing information as to whether the current process has a CUDA
+        context created, and in the positive case containing also information about
+        the device the context belongs to.
     """
     init_once()
-    if not is_initialized():
-        return False
-    for index in range(device_get_count()):
-        handle = pynvml.nvmlDeviceGetHandleByIndex(index)
-        if hasattr(pynvml, "nvmlDeviceGetComputeRunningProcesses_v2"):
-            running_processes = pynvml.nvmlDeviceGetComputeRunningProcesses_v2(handle)
-        else:
-            running_processes = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
-        for proc in running_processes:
-            if os.getpid() == proc.pid:
-                return index
-    return False
+    if is_initialized():
+        for index in range(device_get_count()):
+            handle = pynvml.nvmlDeviceGetHandleByIndex(index)
+            try:
+                mig_current_mode, mig_pending_mode = pynvml.nvmlDeviceGetMigMode(handle)
+            except pynvml.NVMLError_NotSupported:
+                mig_current_mode = pynvml.NVML_DEVICE_MIG_DISABLE
+            if mig_current_mode == pynvml.NVML_DEVICE_MIG_ENABLE:
+                for mig_index in range(pynvml.nvmlDeviceGetMaxMigDeviceCount(handle)):
+                    try:
+                        mig_handle = pynvml.nvmlDeviceGetMigDeviceHandleByIndex(
+                            handle, mig_index
+                        )
+                    except pynvml.NVMLError_NotFound:
+                        # No MIG device with that index
+                        continue
+                    if _running_process_matches(mig_handle):
+                        uuid = pynvml.nvmlDeviceGetUUID(mig_handle)
+                        return CudaContext(
+                            has_context=True,
+                            device_info=CudaDeviceInfo(
+                                uuid=uuid, device_index=index, mig_index=mig_index
+                            ),
+                        )
+            else:
+                if _running_process_matches(handle):
+                    uuid = pynvml.nvmlDeviceGetUUID(handle)
+                    return CudaContext(
+                        has_context=True,
+                        device_info=CudaDeviceInfo(uuid=uuid, device_index=index),
+                    )
+    return CudaContext(has_context=False)
+
+
+def get_device_index_and_uuid(device):
+    """Get both device index and UUID from device index or UUID
+
+    Parameters
+    ----------
+    device : int, bytes or str
+        An ``int`` with the index of a GPU, or ``bytes`` or ``str`` with the UUID
+        of a CUDA (either GPU or MIG) device.
+
+    Returns
+    -------
+    out : CudaDeviceInfo
+        Object containing information about the device.
+
+    Examples
+    --------
+    >>> get_device_index_and_uuid(0)  # doctest: +SKIP
+    {'device-index': 0, 'uuid': b'GPU-e1006a74-5836-264f-5c26-53d19d212dfe'}
+
+    >>> get_device_index_and_uuid('GPU-e1006a74-5836-264f-5c26-53d19d212dfe')  # doctest: +SKIP
+    {'device-index': 0, 'uuid': b'GPU-e1006a74-5836-264f-5c26-53d19d212dfe'}
+
+    >>> get_device_index_and_uuid('MIG-7feb6df5-eccf-5faa-ab00-9a441867e237')  # doctest: +SKIP
+    {'device-index': 0, 'uuid': b'MIG-7feb6df5-eccf-5faa-ab00-9a441867e237'}
+    """
+    init_once()
+    try:
+        device_index = int(device)
+        device_handle = pynvml.nvmlDeviceGetHandleByIndex(device_index)
+        uuid = pynvml.nvmlDeviceGetUUID(device_handle)
+    except ValueError:
+        uuid = device if isinstance(device, bytes) else bytes(device, "utf-8")
+
+        # Validate UUID, get index and UUID as seen with `nvidia-smi -L`
+        uuid_handle = pynvml.nvmlDeviceGetHandleByUUID(uuid)
+        device_index = pynvml.nvmlDeviceGetIndex(uuid_handle)
+        uuid = pynvml.nvmlDeviceGetUUID(uuid_handle)
+
+    return CudaDeviceInfo(uuid=uuid, device_index=device_index)
+
+
+def get_device_mig_mode(device):
+    """Get MIG mode for a device index or UUID
+
+    Parameters
+    ----------
+    device: int, bytes or str
+        An ``int`` with the index of a GPU, or ``bytes`` or ``str`` with the UUID
+        of a CUDA (either GPU or MIG) device.
+
+    Returns
+    -------
+    out : list
+        A ``list`` with two integers ``[current_mode, pending_mode]``.
+    """
+    init_once()
+    try:
+        device_index = int(device)
+        handle = pynvml.nvmlDeviceGetHandleByIndex(device_index)
+    except ValueError:
+        uuid = device if isinstance(device, bytes) else bytes(device, "utf-8")
+        handle = pynvml.nvmlDeviceGetHandleByUUID(uuid)
+    try:
+        return pynvml.nvmlDeviceGetMigMode(handle)
+    except pynvml.NVMLError_NotSupported:
+        return [0, 0]
 
 
 def _get_utilization(h):

distributed/diagnostics/tests/test_nvml.py

@@ -56,19 +56,25 @@ def test_wsl_monitoring_enabled():
 
 def run_has_cuda_context(queue):
     try:
-        assert not nvml.has_cuda_context()
+        assert not nvml.has_cuda_context().has_context
 
         import numba.cuda
 
         numba.cuda.current_context()
-        assert nvml.has_cuda_context() == 0
+        ctx = nvml.has_cuda_context()
+        assert (
+            ctx.has_context
+            and ctx.device_info.device_index == 0
+            and isinstance(ctx.device_info.uuid, bytes)
+        )
 
         queue.put(None)
 
     except Exception as e:
         queue.put(e)
 
 
+@pytest.mark.xfail(reason="If running on Docker, requires --pid=host")
 def test_has_cuda_context():
     if nvml.device_get_count() < 1:
         pytest.skip("No GPUs available")

distributed/utils_test.py

@@ -2130,6 +2130,23 @@ def ucx_loop():
     ucp.reset()
     loop.close()
 
+    # Reset also Distributed's UCX initialization, i.e., revert the effects of
+    # `distributed.comm.ucx.init_once()`.
+    import distributed.comm.ucx
+
+    distributed.comm.ucx.ucp = None
+    # If the test created a context, clean it up.
+    # TODO: should we check if there's already a context _before_ the test runs?
+    # I think that would be useful.
+    from distributed.diagnostics.nvml import has_cuda_context
+
+    ctx = has_cuda_context()
+    if ctx.has_context:
+        import numba.cuda
+
+        ctx = numba.cuda.current_context()
+        ctx.device.reset()
+
 
 def wait_for_log_line(
     match: bytes, stream: IO[bytes] | None, max_lines: int | None = 10