async def test_two_servers_many_clients():
lf1 = ucp.create_listener(server_node)
lf2 = ucp.create_listener(server_node)
clients = []
for _ in range(100):
clients.append(client_node(lf1.port))
clients.append(client_node(lf2.port))
await asyncio.gather(*clients, loop=asyncio.get_event_loop())
async def test_multiple_nodes():
lf1 = ucp.create_listener(server_node)
lf2 = ucp.create_listener(server_node)
assert lf1.port != lf2.port
nodes = []
for _ in range(10):
nodes.append(client_node(lf1.port))
nodes.append(client_node(lf2.port))
await asyncio.gather(*nodes, loop=asyncio.get_event_loop())
async def test_message_probe(transfer_api):
msg = bytearray(b"0" * 10)
async def server_node(ep):
# Wait for remote endpoint to close before probing the endpoint for
# in-transit message and receiving it.
while not ep.closed():
await asyncio.sleep(0) # Yield task
if transfer_api == "am":
assert ep._ep.am_probe() is True
received = await ep.am_recv()
else:
assert ep._ctx.worker.tag_probe(ep._tags["msg_recv"]) is True
received = bytearray(10)
await ep.recv(received)
assert received == msg
async def client_node(port):
ep = await ucp.create_endpoint(
ucp.get_address(),
port,
)
if transfer_api == "am":
await ep.am_send(msg)
else:
await ep.send(msg)
listener = ucp.create_listener(server_node, )
await client_node(listener.port)
async def test_close_callback(server_close_callback):
endpoint_error_handling = ucp.get_ucx_version() >= (1, 10, 0)
closed = [False]
def _close_callback():
closed[0] = True
async def server_node(ep):
if server_close_callback is True:
ep.set_close_callback(_close_callback)
msg = bytearray(10)
await ep.recv(msg)
if server_close_callback is False:
await ep.close()
async def client_node(port):
ep = await ucp.create_endpoint(
ucp.get_address(), port, endpoint_error_handling=endpoint_error_handling
)
if server_close_callback is False:
ep.set_close_callback(_close_callback)
await ep.send(bytearray(b"0" * 10))
if server_close_callback is True:
await ep.close()
listener = ucp.create_listener(
server_node, endpoint_error_handling=endpoint_error_handling
)
await client_node(listener.port)
assert closed[0] is True
async def run():
async def server_handler(ep):
times = []
msg_recv_list = []
if not args.reuse_alloc:
for _ in range(args.n_iter):
msg_recv_list.append(np.zeros(args.n_bytes, dtype="u1"))
else:
t = np.zeros(args.n_bytes, dtype="u1")
for _ in range(args.n_iter):
msg_recv_list.append(t)
assert msg_recv_list[0].nbytes == args.n_bytes
for i in range(args.n_iter):
await ep.recv(msg_recv_list[i], args.n_bytes)
await ep.send(msg_recv_list[i], args.n_bytes)
await ep.close()
lf.close()
lf = ucp.create_listener(server_handler)
queue.put(lf.port)
while not lf.closed():
await asyncio.sleep(0.5)
async def f(listener_port):
# coroutine shows up when the client asks
# to connect
set_rmm()
async def write(ep):
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"))
while True:
for i in range(ITERATIONS):
print("ITER: ", i)
# Send meta data
await send(ep, frames)
frames, msg = await recv(ep)
print("CONFIRM RECEIPT")
await ep.close()
break
# lf.close()
del msg
del frames
lf = ucp.create_listener(write, port=listener_port)
try:
while not lf.closed():
await asyncio.sleep(0.1)
except ucp.UCXCloseError:
pass
async def test_close_callback(server_close_callback):
closed = [False]
def _close_callback():
closed[0] = True
async def server_node(ep):
if server_close_callback is True:
ep.set_close_callback(_close_callback)
if server_close_callback is False:
await ep.close()
async def client_node(port):
ep = await ucp.create_endpoint(
ucp.get_address(),
port,
)
if server_close_callback is False:
ep.set_close_callback(_close_callback)
if server_close_callback is True:
await ep.close()
listener = ucp.create_listener(
server_node,
)
await client_node(listener.port)
assert closed[0] is True
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"))
# Send meta data
try:
await send(ep, frames)
except Exception:
# Avoids process hanging on "Endpoint timeout"
pass
print("Shutting Down Server...")
await ep.close()
lf.close()
lf = ucp.create_listener(
write, port=listener_port, endpoint_error_handling=endpoint_error_handling
)
try:
while not lf.closed():
await asyncio.sleep(0.1)
except ucp.UCXCloseError:
pass
async def test_send_recv_am(size, blocking_progress_mode, recv_wait, data):
rndv_thresh = 8192
ucp.init(
options={"RNDV_THRESH": str(rndv_thresh)},
blocking_progress_mode=blocking_progress_mode,
)
ucp.register_am_allocator(data["allocator"], data["memory_type"])
msg = data["generator"](size)
recv = []
listener = ucp.create_listener(simple_server(size, recv))
num_clients = 1
clients = [
await ucp.create_endpoint(ucp.get_address(), listener.port)
for i in range(num_clients)
]
for c in clients:
if recv_wait:
# By sleeping here we ensure that the listener's
# ep.am_recv call will have to wait, rather than return
# immediately as receive data is already available.
await asyncio.sleep(1)
await c.am_send(msg)
for c in clients:
await c.close()
listener.close()
if data["memory_type"] == "cuda" and msg.nbytes < rndv_thresh:
# Eager messages are always received on the host, if no host
# allocator is registered UCX-Py defaults to `bytearray`.
assert recv[0] == bytearray(msg.get())
else:
data["validator"](recv[0], msg)
async def run():
async def server_handler(ep):
if not args.enable_am:
msg_recv_list = []
if not args.reuse_alloc:
for _ in range(args.n_iter):
msg_recv_list.append(xp.zeros(args.n_bytes,
dtype="u1"))
else:
t = xp.zeros(args.n_bytes, dtype="u1")
for _ in range(args.n_iter):
msg_recv_list.append(t)
assert msg_recv_list[0].nbytes == args.n_bytes
for i in range(args.n_iter):
if args.enable_am is True:
recv = await ep.am_recv()
await ep.am_send(recv)
else:
await ep.recv(msg_recv_list[i])
await ep.send(msg_recv_list[i])
await ep.close()
lf.close()
lf = ucp.create_listener(server_handler, port=args.port)
queue.put(lf.port)
while not lf.closed():
await asyncio.sleep(0.5)
async def test_tag_match():
msg1 = bytes("msg1", "utf-8")
msg2 = bytes("msg2", "utf-8")
async def server_node(ep):
f1 = ep.send(msg1, tag="msg1")
await asyncio.sleep(1) # Let msg1 finish
f2 = ep.send(msg2, tag="msg2")
await asyncio.gather(f1, f2)
lf = ucp.create_listener(server_node)
ep = await ucp.create_endpoint(ucp.get_address(), lf.port)
m1, m2 = (bytearray(len(msg1)), bytearray(len(msg2)))
# May be dropped in favor of `asyncio.create_task` only
# once Python 3.6 is dropped.
if hasattr(asyncio, "create_future"):
f2 = asyncio.create_task(ep.recv(m2, tag="msg2"))
else:
f2 = asyncio.ensure_future(ep.recv(m2, tag="msg2"))
# At this point f2 shouldn't be able to finish because its
# tag "msg2" doesn't match the servers send tag "msg1"
done, pending = await asyncio.wait({f2}, timeout=0.01)
assert f2 in pending
# "msg1" should be ready
await ep.recv(m1, tag="msg1")
assert m1 == msg1
await f2
assert m2 == msg2
async def test_close_after_n_recv(message_type):
"""The Endpoint.close_after_n_recv()"""
async def server_node(ep):
for _ in range(10):
await _shutdown_send(ep, message_type)
async def client_node(port):
ep = await ucp.create_endpoint(
ucp.get_address(),
port,
)
ep.close_after_n_recv(10)
for _ in range(10):
await _shutdown_recv(ep, message_type)
assert ep.closed()
ep = await ucp.create_endpoint(
ucp.get_address(),
port,
)
for _ in range(5):
await _shutdown_recv(ep, message_type)
ep.close_after_n_recv(5)
for _ in range(5):
await _shutdown_recv(ep, message_type)
assert ep.closed()
ep = await ucp.create_endpoint(
ucp.get_address(),
port,
)
for _ in range(5):
await _shutdown_recv(ep, message_type)
ep.close_after_n_recv(10, count_from_ep_creation=True)
for _ in range(5):
await _shutdown_recv(ep, message_type)
assert ep.closed()
ep = await ucp.create_endpoint(
ucp.get_address(),
port,
)
for _ in range(10):
await _shutdown_recv(ep, message_type)
with pytest.raises(
ucp.exceptions.UCXError,
match="`n` cannot be less than current recv_count",
):
ep.close_after_n_recv(5, count_from_ep_creation=True)
ep.close_after_n_recv(1)
with pytest.raises(
ucp.exceptions.UCXError,
match="close_after_n_recv has already been set to",
):
ep.close_after_n_recv(1)
listener = ucp.create_listener(server_node, )
await client_node(listener.port)
def test_mismatch(server_guarantee_msg_order):
# We use an exception handle to catch errors raised by the server
def handle_exception(loop, context):
msg = str(context.get("exception", context["message"]))
loop.test_failed = msg.find(loop.error_msg_expected) == -1
loop = asyncio.get_event_loop()
loop.set_exception_handler(handle_exception)
loop.test_failed = False
loop.error_msg_expected = "Both peers must set guarantee_msg_order identically"
with pytest.raises(ValueError, match=loop.error_msg_expected):
lt = ucp.create_listener(
lambda x: x, guarantee_msg_order=server_guarantee_msg_order)
loop.run_until_complete(
ucp.create_endpoint(
ucp.get_address(),
lt.port,
guarantee_msg_order=(not server_guarantee_msg_order),
))
loop.run_until_complete(
asyncio.sleep(0.1)) # Give the server time to finish
assert not loop.test_failed, "expected error message not raised by the server"
async def test_get_ucp_worker():
worker = ucp.get_ucp_worker()
assert isinstance(worker, int)
def server(ep):
assert ep.get_ucp_worker() == worker
lt = ucp.create_listener(server)
ep = await ucp.create_endpoint(ucp.get_address(), lt.port)
assert ep.get_ucp_worker() == worker
async def test_mismatch():
def server(ep):
pass
lt = ucp.create_listener(server, guarantee_msg_order=True)
with pytest.raises(
ValueError,
match="Both peers must set guarantee_msg_order identically"):
await ucp.create_endpoint(ucp.get_address(),
lt.port,
guarantee_msg_order=False)
lt = ucp.create_listener(server, guarantee_msg_order=False)
with pytest.raises(
ValueError,
match="Both peers must set guarantee_msg_order identically"):
await ucp.create_endpoint(ucp.get_address(),
lt.port,
guarantee_msg_order=True)
async def test_get_endpoint():
def server(ep):
ucp_ep = ep.get_ucp_endpoint()
assert isinstance(ucp_ep, int)
assert ucp_ep > 0
lt = ucp.create_listener(server)
ep = await ucp.create_endpoint(ucp.get_address(), lt.port)
ucp_ep = ep.get_ucp_endpoint()
assert isinstance(ucp_ep, int)
assert ucp_ep > 0
async def test_reset():
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
del ep
reset()
async def test_listener_del():
"""The client delete the listener"""
async def server_node(ep):
await ep.send(np.arange(100, dtype=np.int64))
await ep.send(np.arange(100, dtype=np.int64))
listener = ucp.create_listener(server_node)
ep = await ucp.create_endpoint(ucp.get_address(), listener.port)
msg = np.empty(100, dtype=np.int64)
await ep.recv(msg)
assert listener.closed() is False
del listener
await ep.recv(msg)
async def test_send_recv_bytes(size, blocking_progress_mode):
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
async def start(self):
async def serve_forever(client_ep):
ucx = UCX(
client_ep,
local_addr=self.address,
peer_addr=self.address,
deserialize=self.deserialize,
)
if self.comm_handler:
await self.comm_handler(ucx)
init_once()
self.ucp_server = ucp.create_listener(serve_forever, port=self._input_port)
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
async def run():
async def server_node(ep):
try:
await ep.send(np.arange(100, dtype=np.int64))
# Waiting for signal to close the endpoint
await mp_queue_get_nowait(server_queue)
await ep.close()
finally:
listener.close()
listener = ucp.create_listener(server_node)
client_queue.put(listener.port)
while not listener.closed():
await asyncio.sleep(0.1)
async def test_send_recv_error(blocking_progress_mode):
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)
async def test_send_recv_numpy(size, dtype, blocking_progress_mode):
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_close_after_n_recv():
"""The Endpoint.close_after_n_recv()"""
async def server_node(ep):
for _ in range(10):
await ep.send(np.arange(10))
async def client_node(port):
ep = await ucp.create_endpoint(ucp.get_address(), port)
ep.close_after_n_recv(10)
for _ in range(10):
msg = np.empty(10)
await ep.recv(msg)
assert ep.closed()
ep = await ucp.create_endpoint(ucp.get_address(), port)
for _ in range(5):
msg = np.empty(10)
await ep.recv(msg)
ep.close_after_n_recv(5)
for _ in range(5):
msg = np.empty(10)
await ep.recv(msg)
assert ep.closed()
ep = await ucp.create_endpoint(ucp.get_address(), port)
for _ in range(5):
msg = np.empty(10)
await ep.recv(msg)
ep.close_after_n_recv(10, count_from_ep_creation=True)
for _ in range(5):
msg = np.empty(10)
await ep.recv(msg)
assert ep.closed()
ep = await ucp.create_endpoint(ucp.get_address(), port)
for _ in range(10):
msg = np.empty(10)
await ep.recv(msg)
with pytest.raises(ucp.exceptions.UCXError,
match="`n` cannot be less than current recv_count"):
ep.close_after_n_recv(5, count_from_ep_creation=True)
ep.close_after_n_recv(1)
with pytest.raises(ucp.exceptions.UCXError,
match="close_after_n_recv has already been set to"):
ep.close_after_n_recv(1)
listener = ucp.create_listener(server_node)
await client_node(listener.port)
def start(self):
async def serve_forever(client_ep):
ucx = UCX(
client_ep,
local_addr=self.address,
peer_addr=self.
address, # TODO: https://github.com/Akshay-Venkatesh/ucx-py/issues/111
deserialize=self.deserialize,
)
if self.comm_handler:
await self.comm_handler(ucx)
self.ucp_server = ucp.create_listener(serve_forever,
port=self._input_port)
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_send_recv_obj(blocking_progress_mode):
ucp.init(blocking_progress_mode=blocking_progress_mode)
async def echo_obj_server(ep):
obj = await ep.recv_obj()
await ep.send_obj(obj)
listener = ucp.create_listener(echo_obj_server)
client = await ucp.create_endpoint(ucp.get_address(), listener.port)
msg = bytearray(b"hello")
await client.send_obj(msg)
got = await client.recv_obj()
assert msg == got
async def f(listener_port, frames):
# coroutine shows up when the client asks
# to connect
if enable_rmm:
set_rmm()
# Use a global so the `write` callback function can read frames
global _frames
global _connected
global _disconnected
global _lock
_connected = 0
_disconnected = 0
_lock = threading.Lock()
_frames = frames
async def write(ep):
global _connected
global _disconnected
_lock.acquire()
_connected += 1
_lock.release()
for i in range(TRANSFER_ITERATIONS):
print("ITER: ", i)
# Send meta data
await send(ep, _frames)
print("CONFIRM RECEIPT")
await ep.close()
_lock.acquire()
_disconnected += 1
_lock.release()
# break
lf = ucp.create_listener(write, port=listener_port)
proc_conn.send("initialized")
proc_conn.close()
try:
while _disconnected < num_workers * EP_ITERATIONS:
await asyncio.sleep(0.1)
print("Closing listener")
lf.close()
except ucp.UCXCloseError:
pass
async def test_send_recv_numba(size, dtype, blocking_progress_mode):
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))