Update to support CUDA context warnings for all protocols (#1548)

Update to support CUDA context warnings for all protocols, even after removal of `distributed.comm.ucx`. Depends on #1546. Partially addresses #1517.

Authors:
  - Peter Andreas Entschev (https://github.com/pentschev)

Approvers:
  - Tom Augspurger (https://github.com/TomAugspurger)

URL: #1548

Author: pentschev

dask_cuda/initialize.py

@@ -8,13 +8,52 @@
 import cuda.core.experimental
 
 import dask
-from distributed.diagnostics.nvml import get_device_index_and_uuid, has_cuda_context
+from distributed.diagnostics.nvml import (
+    CudaDeviceInfo,
+    get_device_index_and_uuid,
+    has_cuda_context,
+)
 
 from .utils import get_ucx_config
 
 logger = logging.getLogger(__name__)
 
 
+pre_existing_cuda_context = None
+cuda_context_created = None
+
+
+_warning_suffix = (
+    "This is often the result of a CUDA-enabled library calling a CUDA runtime "
+    "function before Dask-CUDA can spawn worker processes. Please make sure any such "
+    "function calls don't happen at import time or in the global scope of a program."
+)
+
+
+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 {device_uuid_str} already exists "
+        f"on process ID {pid}. {_warning_suffix}"
+    )
+
+
+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"{expected_device_uuid_str}, but instead the CUDA context is on device "
+        f"{actual_device_uuid_str}. {_warning_suffix}"
+    )
+
+
 def _create_cuda_context_handler():
     if int(os.environ.get("DASK_CUDA_TEST_SINGLE_GPU", "0")) != 0:
         try:
@@ -25,98 +64,76 @@ def _create_cuda_context_handler():
         cuda.core.experimental.Device().set_current()
 
 
-def _warn_generic():
-    try:
-        # TODO: update when UCX-Py is removed, see
-        # https://github.com/rapidsai/dask-cuda/issues/1517
-        import distributed.comm.ucx
+def _create_cuda_context_and_warn():
+    """Create CUDA context and warn depending on certain conditions.
 
-        # Added here to ensure the parent `LocalCUDACluster` process creates the CUDA
-        # context directly from the UCX module, thus avoiding a similar warning there.
-        cuda_visible_device = get_device_index_and_uuid(
-            os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")[0]
-        )
-        ctx = has_cuda_context()
-        if (
-            ctx.has_context
-            and not distributed.comm.ucx.cuda_context_created.has_context
-        ):
-            distributed.comm.ucx._warn_existing_cuda_context(ctx, os.getpid())
-
-        _create_cuda_context_handler()
-
-        if (
-            distributed.comm.ucx.cuda_context_created is False
-            or distributed.comm.ucx.cuda_context_created.has_context
-        ):
-            ctx = has_cuda_context()
-            if ctx.has_context and ctx.device_info != cuda_visible_device:
-                distributed.comm.ucx._warn_cuda_context_wrong_device(
-                    cuda_visible_device, ctx.device_info, os.getpid()
-                )
+    Warns if a pre-existing CUDA context already existed or if the resulting CUDA
+    context was created in the wrong device.
 
-    except Exception:
-        logger.error("Unable to start CUDA Context", exc_info=True)
+    This function is almost an identical duplicate from
+    `distributed_ucxx.ucxx.init_once`, the duplication is necessary because Dask-CUDA
+    needs to support `protocol="tcp"` as well, even when distributed-ucxx is not
+    installed, but this here runs _after_ comms have started, which is fine for TCP
+    because the time when CUDA context is created is not important. The code needs to
+    live also in distributed-ucxx because there the time when a CUDA context is created
+    matters, and it needs to happen _before_ UCX is initialized, but comms in
+    Distributed is initialized before preload, and thus only after this function
+    executes.
 
+    Raises
+    ------
+    Exception
+        If anything wrong happened during context initialization.
 
-def _initialize_ucx():
-    try:
-        import distributed.comm.ucx
+    Returns
+    -------
+    None
+    """
+    global pre_existing_cuda_context, cuda_context_created
 
-        distributed.comm.ucx.init_once()
-    except ModuleNotFoundError:
-        # UCX initialization has to be delegated to Distributed, it will take care
-        # of setting correct environment variables and importing `ucp` after that.
-        # Therefore if ``import ucp`` fails we can just continue here.
-        pass
+    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.has_context:
+        _warn_existing_cuda_context(pre_existing_cuda_context.device_info, os.getpid())
+
+    _create_cuda_context_handler()
+
+    cuda_context_created = has_cuda_context()
+    if (
+        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.device_info, os.getpid()
+        )
 
 
-def _initialize_ucxx():
+def _create_cuda_context():
     try:
         # Added here to ensure the parent `LocalCUDACluster` process creates the CUDA
         # context directly from the UCX module, thus avoiding a similar warning there.
         import distributed_ucxx.ucxx
+    except ImportError:
+        pass
+    else:
+        if distributed_ucxx.ucxx.ucxx is not None:
+            # UCXX has already initialized (and warned if necessary)
+            return
 
-        distributed_ucxx.ucxx.init_once()
-
-        cuda_visible_device = get_device_index_and_uuid(
-            os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")[0]
-        )
-        ctx = has_cuda_context()
-        if (
-            ctx.has_context
-            and not distributed_ucxx.ucxx.cuda_context_created.has_context
-        ):
-            distributed_ucxx.ucxx._warn_existing_cuda_context(ctx, os.getpid())
-
-        _create_cuda_context_handler()
-
-        if not distributed_ucxx.ucxx.cuda_context_created.has_context:
-            ctx = has_cuda_context()
-            if ctx.has_context and ctx.device_info != cuda_visible_device:
-                distributed_ucxx.ucxx._warn_cuda_context_wrong_device(
-                    cuda_visible_device, ctx.device_info, os.getpid()
-                )
-
+    try:
+        _create_cuda_context_and_warn()
     except Exception:
         logger.error("Unable to start CUDA Context", exc_info=True)
 
 
-def _create_cuda_context(protocol="ucx"):
-    if protocol in ("ucx", "ucxx"):
-        _initialize_ucxx()
-    else:
-        # Not a UCX protocol, just raise CUDA context warnings if needed.
-        _warn_generic()
-
-
 def initialize(
     create_cuda_context=True,
     enable_tcp_over_ucx=None,
     enable_infiniband=None,
     enable_nvlink=None,
     enable_rdmacm=None,
-    protocol="ucx",
 ):
     """Create CUDA context and initialize UCXX configuration.
 
@@ -167,12 +184,11 @@ def initialize(
         enable_infiniband=enable_infiniband,
         enable_nvlink=enable_nvlink,
         enable_rdmacm=enable_rdmacm,
-        protocol=protocol,
     )
     dask.config.set({"distributed-ucxx": ucx_config})
 
     if create_cuda_context:
-        _create_cuda_context(protocol=protocol)
+        _create_cuda_context()
 
 
 @click.command()
@@ -185,6 +201,5 @@ def dask_setup(
     worker,
     create_cuda_context,
 ):
-    protocol = worker._protocol.split("://")[0]
     if create_cuda_context:
-        _create_cuda_context(protocol=protocol)
+        _create_cuda_context()

dask_cuda/tests/test_dask_setup.py

@@ -17,12 +17,10 @@
 from dask_cuda.utils import wait_workers
 
 
-@pytest.mark.parametrize("protocol", ["tcp", "ucx", "ucxx"])
-def test_dask_setup_function_with_mock_worker(protocol):
+def test_dask_setup_function_with_mock_worker():
     """Test the dask_setup function directly with mock worker."""
     # Create a mock worker object
     mock_worker = Mock()
-    mock_worker._protocol = protocol
 
     with patch("dask_cuda.initialize._create_cuda_context") as mock_create_context:
         # Test with create_cuda_context=True
@@ -33,7 +31,7 @@ def test_dask_setup_function_with_mock_worker(protocol):
             create_cuda_context=True,
         )
 
-        mock_create_context.assert_called_once_with(protocol=protocol)
+        mock_create_context.assert_called_once_with()
 
         mock_create_context.reset_mock()
 

dask_cuda/tests/test_dgx.py

@@ -152,7 +152,6 @@ def _test_ucx_infiniband_nvlink(
             cm_tls_priority = ["tcp"]
 
     initialize(
-        protocol="ucx",
         enable_tcp_over_ucx=enable_tcp_over_ucx,
         enable_infiniband=enable_infiniband,
         enable_nvlink=enable_nvlink,