Move host_array to comm.utils

Author: jakirkham

distributed/comm/asyncio_tcp.py

@@ -17,7 +17,7 @@
 
 import dask
 
-from ..utils import ensure_ip, get_ip, get_ipv6
+from ..utils import ensure_ip, get_ip, get_ipv6, host_array
 from .addressing import parse_host_port, unparse_host_port
 from .core import Comm, CommClosedError, Connector, Listener
 from .registry import Backend
@@ -29,15 +29,6 @@
 _COMM_CLOSED = object()
 
 
-# Find the function, `host_array()`, to use when allocating new host arrays
-try:
-    import numpy
-
-    host_array = lambda n: memoryview(numpy.empty((n,), dtype="u1"))  # type: ignore
-except ImportError:
-    host_array = lambda n: memoryview(bytearray(n))  # type: ignore
-
-
 def coalesce_buffers(
     buffers: list[bytes],
     target_buffer_size: int = 64 * 1024,

distributed/comm/tcp.py

@@ -32,7 +32,13 @@
 from .addressing import parse_host_port, unparse_host_port
 from .core import Comm, CommClosedError, Connector, FatalCommClosedError, Listener
 from .registry import Backend
-from .utils import ensure_concrete_host, from_frames, get_tcp_server_address, to_frames
+from .utils import (
+    ensure_concrete_host,
+    from_frames,
+    get_tcp_server_address,
+    host_array,
+    to_frames,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -41,15 +47,6 @@
 MAX_BUFFER_SIZE = MEMORY_LIMIT / 2
 
 
-# Find the function, `host_array()`, to use when allocating new host arrays
-try:
-    import numpy
-
-    host_array = lambda n: memoryview(numpy.empty((n,), dtype="u1"))  # type: ignore
-except ImportError:
-    host_array = lambda n: memoryview(bytearray(n))  # type: ignore
-
-
 def set_tcp_timeout(comm):
     """
     Set kernel-level TCP timeout on the stream.

distributed/comm/ucx.py

@@ -17,7 +17,7 @@
 from dask.utils import parse_bytes
 
 from ..diagnostics.nvml import has_cuda_context
-from ..utils import ensure_ip, get_ip, get_ipv6, log_errors, nbytes
+from ..utils import ensure_ip, get_ip, get_ipv6, host_array, log_errors, nbytes
 from .addressing import parse_host_port, unparse_host_port
 from .core import Comm, CommClosedError, Connector, Listener
 from .registry import Backend, backends
@@ -41,7 +41,6 @@
     ucx_create_endpoint = None  # type: ignore
     ucx_create_listener = None  # type: ignore
 
-host_array = None
 device_array = None
 pre_existing_cuda_context = False
 cuda_context_created = False
@@ -57,7 +56,7 @@ def synchronize_stream(stream=0):
 
 
 def init_once():
-    global ucp, host_array, device_array
+    global ucp, device_array
     global ucx_create_endpoint, ucx_create_listener
     global pre_existing_cuda_context, cuda_context_created
 
@@ -115,14 +114,6 @@ def init_once():
 
     ucp.init(options=ucx_config, env_takes_precedence=True)
 
-    # Find the function, `host_array()`, to use when allocating new host arrays
-    try:
-        import numpy
-
-        host_array = lambda n: numpy.empty((n,), dtype="u1")
-    except ImportError:
-        host_array = lambda n: bytearray(n)
-
     # Find the function, `cuda_array()`, to use when allocating new CUDA arrays
     try:
         import rmm

distributed/comm/utils.py

@@ -18,6 +18,15 @@
     OFFLOAD_THRESHOLD = parse_bytes(OFFLOAD_THRESHOLD)
 
 
+# Find the function, `host_array()`, to use when allocating new host arrays
+try:
+    import numpy
+
+    host_array = lambda n: memoryview(numpy.empty((n,), dtype="u1"))  # type: ignore
+except ImportError:
+    host_array = lambda n: memoryview(bytearray(n))  # type: ignore
+
+
 async def to_frames(
     msg,
     allow_offload=True,