def event_loop(scope="session"):
loop = asyncio.new_event_loop()
loop.set_exception_handler(handle_exception)
ucp.reset()
yield loop
ucp.reset()
loop.run_until_complete(asyncio.sleep(0))
def test_get_ucx_version():
ucp.reset()
version = ucp.get_ucx_version()
assert isinstance(version, tuple)
assert len(version) == 3
# Check UCX isn't initialized
assert ucp.core._ctx is None
def test_init_options():
ucp.reset()
options = {"SEG_SIZE": "3M"}
# environment specification should be ignored
ucp.init(options)
config = ucp.get_config()
assert config["SEG_SIZE"] == options["SEG_SIZE"]
def test_init_options():
ucp.reset()
os.environ["UCX_SEG_SIZE"] = "2M" # Should be ignored
options = {"SEG_SIZE": "3M"}
ucp.init(options)
config = ucp.get_config()
assert config["SEG_SIZE"] == options["SEG_SIZE"]
def test_init_options_and_env():
ucp.reset()
os.environ["UCX_SEG_SIZE"] = "4M"
options = {"SEG_SIZE": "3M"} # Should be ignored
ucp.init(options, env_takes_precedence=True)
config = ucp.get_config()
assert config["SEG_SIZE"] == options["SEG_SIZE"]
def test_init_options_and_env():
ucp.reset()
options = {"SEG_SIZE": "3M"} # Should be ignored
ucp.init(options, env_takes_precedence=True)
config = ucp.get_config()
assert config["SEG_SIZE"] == os.environ["UCX_SEG_SIZE"]
# Provided options dict was not modified.
assert options == {"SEG_SIZE": "3M"}
def test_get_config():
with patch.dict(os.environ):
# Unset to test default value
if os.environ.get("UCX_TLS") is not None:
del os.environ["UCX_TLS"]
ucp.reset()
config = ucp.get_config()
assert isinstance(config, dict)
assert config["TLS"] == "all"
def test_check_transport(transports):
transports_list = transports.split(",")
inactive_transports = list(set(["posix", "tcp"]) - set(transports_list))
ucp.reset()
options = {"TLS": transports, "NET_DEVICES": "all"}
ucp.init(options)
active_transports = ucp.get_active_transports()
for t in transports_list:
assert any([at.startswith(t) for at in active_transports])
for it in inactive_transports:
assert any([not at.startswith(it) for at in active_transports])
async def test_ep_still_in_scope_error():
reset = ResetAfterN(2)
def server(ep):
ep.abort()
reset()
lt = ucp.create_listener(server)
ep = await ucp.create_endpoint(ucp.get_address(), lt.port)
del lt
with pytest.raises(ucp.exceptions.UCXError, match="_ucp_endpoint"):
ucp.reset()
ep.abort()
ucp.reset()
async def test_lt_still_in_scope_error():
reset = ResetAfterN(2)
def server(ep):
ep.abort()
reset()
lt = ucp.create_listener(server)
ep = await ucp.create_endpoint(ucp.get_address(), lt.port)
del ep
with pytest.raises(ucp.exceptions.UCXError, match="ucp._libs.core._Listener"):
ucp.reset()
lt.close()
ucp.reset()
async def test_send_recv_bytes(size, blocking_progress_mode):
asyncio.get_event_loop().set_exception_handler(handle_exception)
ucp.reset()
ucp.init(blocking_progress_mode=blocking_progress_mode)
msg = bytearray(b"m" * size)
msg_size = np.array([len(msg)], dtype=np.uint64)
listener = ucp.create_listener(make_echo_server(lambda n: bytearray(n)))
client = await ucp.create_endpoint(ucp.get_address(), listener.port)
await client.send(msg_size)
await client.send(msg)
resp = bytearray(size)
await client.recv(resp)
assert resp == msg
def initialize(
create_cuda_context=True,
enable_tcp_over_ucx=False,
enable_infiniband=False,
enable_nvlink=False,
net_devices="",
):
if create_cuda_context:
try:
numba.cuda.current_context()
except Exception:
logger.error("Unable to start CUDA Context", exc_info=True)
if enable_tcp_over_ucx or enable_infiniband or enable_nvlink:
try:
import ucp
except ImportError:
logger.error(
"UCX protocol requested but ucp module is not available",
exc_info=True)
else:
options = {}
if enable_tcp_over_ucx or enable_infiniband or enable_nvlink:
tls = "tcp,sockcm,cuda_copy"
tls_priority = "sockcm"
if enable_infiniband:
tls = "rc," + tls
if enable_nvlink:
tls = tls + ",cuda_ipc"
options = {"TLS": tls, "SOCKADDR_TLS_PRIORITY": tls_priority}
if net_devices is not None and net_devices != "":
options["NET_DEVICES"] = net_devices
ucp.reset()
ucp.init(options=options)
ucx_env = {}
for k, v in ucp.get_config().items():
# Skip values that aren't actual environment variables (i.e., not strings)
if isinstance(v, str):
ucx_env["UCX_" + k] = v
# Set also UCX environment variables: required by Dask client. It may be best ti
# to have the client asking the scheduler for the proper variables.
os.environ.update(ucx_env)
async def test_send_recv_numpy(size, dtype, blocking_progress_mode):
asyncio.get_event_loop().set_exception_handler(handle_exception)
ucp.reset()
ucp.init(blocking_progress_mode=blocking_progress_mode)
msg = np.arange(size, dtype=dtype)
msg_size = np.array([msg.nbytes], dtype=np.uint64)
listener = ucp.create_listener(
make_echo_server(lambda n: np.empty(n, dtype=np.uint8)))
client = await ucp.create_endpoint(ucp.get_address(), listener.port)
await client.send(msg_size)
await client.send(msg)
resp = np.empty_like(msg)
await client.recv(resp)
np.testing.assert_array_equal(resp, msg)
async def test_send_recv_numba(size, dtype, blocking_progress_mode):
asyncio.get_event_loop().set_exception_handler(handle_exception)
ucp.reset()
ucp.init(blocking_progress_mode=blocking_progress_mode)
cuda = pytest.importorskip("numba.cuda")
ary = np.arange(size, dtype=dtype)
msg = cuda.to_device(ary)
msg_size = np.array([msg.nbytes], dtype=np.uint64)
listener = ucp.create_listener(
make_echo_server(lambda n: cuda.device_array((n, ), dtype=np.uint8)))
client = await ucp.create_endpoint(ucp.get_address(), listener.port)
await client.send(msg_size)
await client.send(msg)
resp = cuda.device_array_like(msg)
await client.recv(resp)
np.testing.assert_array_equal(np.array(resp), np.array(msg))
async def test_ep_still_in_scope_error():
reset = ResetAfterN(2)
def server(ep):
ep.abort()
reset()
lt = ucp.create_listener(server)
ep = await ucp.create_endpoint(ucp.get_address(), lt.port)
del lt
with pytest.raises(
ucp.exceptions.UCXError,
match="Trying to reset UCX but not all Endpoints and/or Listeners are closed()",
):
ucp.reset()
ep.abort()
ucp.reset()
async def run():
address = bytearray(ucp.get_worker_address())
if error_type == "unreachable":
# Shutdown worker, then send its address to client process via
# multiprocessing.Queue
ucp.reset()
q1.put(address)
else:
# Send worker address to client process via # multiprocessing.Queue,
# wait for client to connect, then shutdown worker.
q1.put(address)
ep_ready = q2.get()
assert ep_ready == "ready"
ucp.reset()
q1.put("disconnected")
def test_logging():
"""
Test default logging configuration.
"""
import logging
root = logging.getLogger("ucx")
# ucp.init will only print INFO LINES
with captured_logger(root, level=logging.INFO) as foreign_log:
ucp.reset()
options = {"SEG_SIZE": "3M"}
ucp.init(options)
assert len(foreign_log.getvalue()) > 0
with captured_logger(root, level=logging.ERROR) as foreign_log:
ucp.reset()
options = {"SEG_SIZE": "3M"}
ucp.init(options)
assert len(foreign_log.getvalue()) == 0
async def test_send_recv_error(blocking_progress_mode):
asyncio.get_event_loop().set_exception_handler(handle_exception)
ucp.reset()
ucp.init(blocking_progress_mode=blocking_progress_mode)
async def say_hey_server(ep):
await ep.send(bytearray(b"Hey"))
listener = ucp.create_listener(say_hey_server)
client = await ucp.create_endpoint(ucp.get_address(), listener.port)
msg = bytearray(100)
with pytest.raises(
ucp.exceptions.UCXError,
match=r"length mismatch: 3 \(got\) != 100 \(expected\)"):
await client.recv(msg)
await client.close()
listener.close()
del client
assert listener.closed() is True
del listener
def test_worker_address(blocking_progress_mode):
ucp.init(blocking_progress_mode=blocking_progress_mode)
addr = ucp.get_worker_address()
assert addr is not None
ucp.reset()
def test_get_config():
ucp.reset()
config = ucp.get_config()
assert isinstance(config, dict)
assert config["MEMTYPE_CACHE"] == "n"
def __call__(self):
self.count += 1
if self.count == self.n:
ucp.reset()
def test_init_unknown_option():
ucp.reset()
options = {"UNKNOWN_OPTION": "3M"}
with pytest.raises(ucp.exceptions.UCXConfigError):
ucp.init(options)
def test_init_invalid_option():
ucp.reset()
options = {"SEG_SIZE": "invalid-size"}
with pytest.raises(ucp.exceptions.UCXConfigError):
ucp.init(options)
async def test_send_recv_cudf(event_loop, g):
# requires numba=0.45 (.nbytes)
# or fix nbytes in distributed
cudf = pytest.importorskip("cudf")
class UCX:
def __init__(self, ep):
self.ep = ep
loop = asyncio.get_event_loop()
self.queue = asyncio.Queue(loop=loop)
async def write(self, cdf):
header, _frames = cdf.serialize()
frames = [pickle.dumps(header)] + _frames
# Send meta data
await self.ep.send(np.array([len(frames)], dtype=np.uint64))
await self.ep.send(
np.array(
[hasattr(f, "__cuda_array_interface__") for f in frames],
dtype=np.bool,
)
)
await self.ep.send(np.array([nbytes(f) for f in frames], dtype=np.uint64))
# Send frames
# breakpoint()
for frame in frames:
if nbytes(frame) > 0:
await self.ep.send(frame)
else:
await self.ep.send(np.empty(0, dtype=np.uint8))
async def read(self):
try:
# Recv meta data
nframes = np.empty(1, dtype=np.uint64)
await self.ep.recv(nframes)
is_cudas = np.empty(nframes[0], dtype=np.bool)
await self.ep.recv(is_cudas)
sizes = np.empty(nframes[0], dtype=np.uint64)
await self.ep.recv(sizes)
except (ucp.exceptions.UCXCanceled, ucp.exceptions.UCXCloseError) as e:
msg = "SOMETHING TERRIBLE HAS HAPPENED IN THE TEST"
raise e(msg)
else:
# Recv frames
frames = []
for is_cuda, size in zip(is_cudas.tolist(), sizes.tolist()):
if size > 0:
if is_cuda:
frame = cuda.device_array((size,), dtype=np.uint8)
else:
frame = np.empty(size, dtype=np.uint8)
await self.ep.recv(frame)
frames.append(frame)
else:
if is_cuda:
frames.append(cuda.device_array((0,), dtype=np.uint8))
else:
frames.append(b"")
return frames
class UCXListener:
def __init__(self):
self.comm = None
def start(self):
async def serve_forever(ep):
ucx = UCX(ep)
self.comm = ucx
self.ucp_server = ucp.create_listener(serve_forever, port=14339)
uu = UCXListener()
uu.start()
uu.address = ucp.get_address()
uu.client = await ucp.create_endpoint(uu.address, uu.ucp_server.port)
ucx = UCX(uu.client)
await asyncio.sleep(0.2)
msg = g(cudf)
frames, _ = await asyncio.gather(uu.comm.read(), ucx.write(msg))
ucx_header = pickle.loads(frames[0])
cudf_buffer = frames[1:]
typ = type(msg)
res = typ.deserialize(ucx_header, cudf_buffer)
from dask.dataframe.utils import assert_eq
assert_eq(res, msg)
await uu.comm.ep.close()
await uu.client.close()
assert uu.client.closed()
assert uu.comm.ep.closed()
del uu.ucp_server
ucp.reset()
def reset():
ucp.reset()
yield
ucp.reset()
def test_set_env():
ucp.reset()
os.environ["UCX_SEG_SIZE"] = "2M"
config = ucp.get_config()
assert config["SEG_SIZE"] == os.environ["UCX_SEG_SIZE"]
def server(env, port, func):
# create listener receiver
# write cudf object
# confirm message is sent correctly
ucp.reset()
os.environ.update(env)
ucp.init()
async def f(listener_port):
# coroutine shows up when the client asks
# to connect
async def write(ep):
import cupy
cupy.cuda.set_allocator(None)
print("CREATING CUDA OBJECT IN SERVER...")
cuda_obj_generator = cloudpickle.loads(func)
cuda_obj = cuda_obj_generator()
msg = {"data": to_serialize(cuda_obj)}
frames = await to_frames(msg,
serializers=("cuda", "dask", "pickle"))
for i in range(ITERATIONS):
# Send meta data
await ep.send(np.array([len(frames)], dtype=np.uint64))
await ep.send(
np.array(
[
hasattr(f, "__cuda_array_interface__")
for f in frames
],
dtype=np.bool,
))
await ep.send(
np.array([nbytes(f) for f in frames], dtype=np.uint64))
# Send frames
for frame in frames:
if nbytes(frame) > 0:
await ep.send(frame)
print("CONFIRM RECEIPT")
close_msg = b"shutdown listener"
msg_size = np.empty(1, dtype=np.uint64)
await ep.recv(msg_size)
msg = np.empty(msg_size[0], dtype=np.uint8)
await ep.recv(msg)
recv_msg = msg.tobytes()
assert recv_msg == close_msg
print("Shutting Down Server...")
await ep.close()
lf.close()
lf = ucp.create_listener(write, port=listener_port)
try:
while not lf.closed():
await asyncio.sleep(0.1)
except ucp.UCXCloseError:
pass
loop = asyncio.get_event_loop()
loop.run_until_complete(f(port))
async def test_flush(blocking_progress_mode):
ucp.init(blocking_progress_mode=blocking_progress_mode)
await ucp.flush()
ucp.reset()
def test_fence(blocking_progress_mode):
ucp.init(blocking_progress_mode=blocking_progress_mode)
# this should always succeed
ucp.fence()
ucp.reset()
def client(env, port, func):
# wait for server to come up
# receive cudf object
# deserialize
# assert deserialized msg is cdf
# send receipt
ucp.reset()
os.environ.update(env)
ucp.init()
# must create context before importing
# cudf/cupy/etc
before_rx, before_tx = total_nvlink_transfer()
async def read():
await asyncio.sleep(1)
ep = await get_ep("client", port)
msg = None
import cupy
cupy.cuda.set_allocator(None)
for i in range(ITERATIONS):
# storing cu objects in msg
# we delete to minimize GPU memory usage
# del msg
try:
# Recv meta data
nframes = np.empty(1, dtype=np.uint64)
await ep.recv(nframes)
is_cudas = np.empty(nframes[0], dtype=np.bool)
await ep.recv(is_cudas)
sizes = np.empty(nframes[0], dtype=np.uint64)
await ep.recv(sizes)
except (ucp.exceptions.UCXCanceled,
ucp.exceptions.UCXCloseError) as e:
msg = "SOMETHING TERRIBLE HAS HAPPENED IN THE TEST"
raise e(msg)
else:
# Recv frames
frames = []
for is_cuda, size in zip(is_cudas.tolist(), sizes.tolist()):
if size > 0:
if is_cuda:
frame = cuda_array(size)
else:
frame = np.empty(size, dtype=np.uint8)
await ep.recv(frame)
frames.append(frame)
else:
if is_cuda:
frames.append(cuda_array(size))
else:
frames.append(b"")
msg = await from_frames(frames)
close_msg = b"shutdown listener"
close_msg_size = np.array([len(close_msg)], dtype=np.uint64)
await ep.send(close_msg_size)
await ep.send(close_msg)
print("Shutting Down Client...")
return msg["data"]
rx_cuda_obj = asyncio.get_event_loop().run_until_complete(read())
rx_cuda_obj + rx_cuda_obj
num_bytes = nbytes(rx_cuda_obj)
print(f"TOTAL DATA RECEIVED: {num_bytes}")
# nvlink only measures in KBs
if num_bytes > 90000:
rx, tx = total_nvlink_transfer()
msg = f"RX BEFORE SEND: {before_rx} -- RX AFTER SEND: {rx} \
-- TOTAL DATA: {num_bytes}"
print(msg)
assert rx > before_rx
cuda_obj_generator = cloudpickle.loads(func)
pure_cuda_obj = cuda_obj_generator()
if isinstance(rx_cuda_obj, cupy.ndarray):
cupy.testing.assert_allclose(rx_cuda_obj, pure_cuda_obj)
else:
cudf.tests.utils.assert_eq(rx_cuda_obj, pure_cuda_obj)
