Prevent CUDA context errors when testing on single-GPU (#737)

Add a `DASK_CUDA_TEST_SINGLE_GPU` environment variable that allows
informing a single-GPU system is used for testing (such as gpuCI). This
then prevents throwing errors when attempting to mock create CUDA
context on devices that are specified via `CUDA_VISIBLE_DEVICES` but are
unavailable in the system.

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

Approvers:
  - Mads R. B. Kristensen (https://github.com/madsbk)
  - AJ Schmidt (https://github.com/ajschmidt8)

URL: #737

ci/gpu/build.sh

@@ -106,7 +106,7 @@ else
     gpuci_logger "Python pytest for dask-cuda"
     cd "$WORKSPACE"
     ls dask_cuda/tests/
-    UCXPY_IFNAME=eth0 UCX_WARN_UNUSED_ENV_VARS=n UCX_MEMTYPE_CACHE=n pytest -vs -Werror::DeprecationWarning -Werror::FutureWarning --cache-clear --basetemp="$WORKSPACE/dask-cuda-tmp" --junitxml="$WORKSPACE/junit-dask-cuda.xml" --cov-config=.coveragerc --cov=dask_cuda --cov-report=xml:"$WORKSPACE/dask-cuda-coverage.xml" --cov-report term dask_cuda/tests/
+    DASK_CUDA_TEST_SINGLE_GPU=1 UCXPY_IFNAME=eth0 UCX_WARN_UNUSED_ENV_VARS=n UCX_MEMTYPE_CACHE=n pytest -vs -Werror::DeprecationWarning -Werror::FutureWarning --cache-clear --basetemp="$WORKSPACE/dask-cuda-tmp" --junitxml="$WORKSPACE/junit-dask-cuda.xml" --cov-config=.coveragerc --cov=dask_cuda --cov-report=xml:"$WORKSPACE/dask-cuda-coverage.xml" --cov-report term dask_cuda/tests/
 
     logger "Run local benchmark..."
     python dask_cuda/benchmarks/local_cudf_shuffle.py --partition-size="1 KiB" -d 0  --runs 1 --backend dask

dask_cuda/initialize.py

@@ -14,6 +14,16 @@
 logger = logging.getLogger(__name__)
 
 
+def _create_cuda_context_handler():
+    if int(os.environ.get("DASK_CUDA_TEST_SINGLE_GPU", "0")) != 0:
+        try:
+            numba.cuda.current_context()
+        except numba.cuda.cudadrv.error.CudaSupportError:
+            pass
+    else:
+        numba.cuda.current_context()
+
+
 def _create_cuda_context():
     try:
         # Added here to ensure the parent `LocalCUDACluster` process creates the CUDA
@@ -39,7 +49,7 @@ def _create_cuda_context():
                 "import time or in the global scope of a program."
             )
 
-        numba.cuda.current_context()
+        _create_cuda_context_handler()
 
         if distributed.comm.ucx.cuda_context_created is False:
             ctx = has_cuda_context()