Define as_cuda_array

Provides a function to let us coerce our underlying
`__cuda_array_interface__` objects into something that behaves more like
an array. Prefers CuPy if possible, but will fallback to Numba if its
not available.

Author: jakirkham

distributed/comm/ucx.py

@@ -35,6 +35,7 @@
 ucp = None
 host_array = None
 device_array = None
+as_device_array = None
 
 
 def synchronize_stream(stream=0):
@@ -47,7 +48,7 @@ def synchronize_stream(stream=0):
 
 
 def init_once():
-    global ucp, host_array, device_array
+    global ucp, host_array, device_array, as_device_array
     if ucp is not None:
         return
 
@@ -100,6 +101,23 @@ def device_array(n):
                     "In order to send/recv CUDA arrays, Numba or RMM is required"
                 )
 
+    # Find the function, `as_device_array()`
+    try:
+        import cupy
+
+        as_device_array = lambda a: cupy.asarray(a)
+    except ImportError:
+        try:
+            import numba.cuda
+
+            as_device_array = lambda a: numba.cuda.as_cuda_array(a)
+        except ImportError:
+
+            def as_device_array(n):
+                raise RuntimeError(
+                    "In order to send/recv CUDA arrays, CuPy or Numba is required"
+                )
+
     pool_size_str = dask.config.get("rmm.pool-size")
     if pool_size_str is not None:
         pool_size = parse_bytes(pool_size_str)