async def test_send_recv_addr(blocking_progress_mode):
ucp.init(blocking_progress_mode=blocking_progress_mode)
msg = ucp.get_worker_address()
msg_size = np.array([len(bytes(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(msg_size[0])
await client.recv(resp)
assert resp == bytes(msg)
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")
async def run():
# Read local worker address
address = ucp.get_worker_address()
# Receive worker address from server via multiprocessing.Queue, create
# endpoint to server
remote_address = queue.get()
ep = await ucp.create_endpoint_from_worker_address(remote_address)
# Send local address to server on tag 0
await ep.send(np.array(address.length, np.int64),
tag=0,
force_tag=True)
await ep.send(address, tag=0, force_tag=True)
# Receive message from server
recv_msg = np.empty(10, dtype=np.int64)
await ep.recv(recv_msg, tag=1, force_tag=True)
np.testing.assert_array_equal(recv_msg, np.arange(10, dtype=np.int64))
async def run():
# Send worker address to client process via multiprocessing.Queue
address = ucp.get_worker_address()
queue.put(address)
# Receive address size
address_size = np.empty(1, dtype=np.int64)
await ucp.recv(address_size, tag=0)
# Receive address buffer on tag 0 and create UCXAddress from it
remote_address = bytearray(address_size[0])
await ucp.recv(remote_address, tag=0)
remote_address = ucp.get_ucx_address_from_buffer(remote_address)
# Create endpoint to remote worker using the received address
ep = await ucp.create_endpoint_from_worker_address(remote_address)
# Send data to client's endpoint
send_msg = np.arange(10, dtype=np.int64)
await ep.send(send_msg, tag=1, force_tag=True)
async def run():
async def _handle_client(packed_remote_address):
# Unpack the fixed-size address+tag buffer
unpacked = _unpack_address_and_tag(packed_remote_address)
remote_address = ucp.get_ucx_address_from_buffer(
unpacked["address"])
# Create endpoint to remote worker using the received address
ep = await ucp.create_endpoint_from_worker_address(remote_address)
# Send data to client's endpoint
send_msg = np.arange(10, dtype=np.int64)
await ep.send(send_msg, tag=unpacked["send_tag"], force_tag=True)
# Receive data from client's endpoint
recv_msg = np.empty(20, dtype=np.int64)
await ep.recv(recv_msg, tag=unpacked["recv_tag"], force_tag=True)
np.testing.assert_array_equal(recv_msg,
np.arange(20, dtype=np.int64))
# Send worker address to client processes via multiprocessing.Queue,
# one entry for each client.
address = ucp.get_worker_address()
for i in range(num_nodes):
queue.put(address)
address_info = _get_address_info()
server_tasks = []
for i in range(num_nodes):
# Receive fixed-size address+tag buffer on tag 0
packed_remote_address = bytearray(address_info["frame_size"])
await ucp.recv(packed_remote_address, tag=0)
# Create an async task for client
server_tasks.append(_handle_client(packed_remote_address))
# Await handling each client request
await asyncio.gather(*server_tasks)
async def run():
# Read local worker address
address = ucp.get_worker_address()
recv_tag = ucp.utils.hash64bits(os.urandom(16))
send_tag = ucp.utils.hash64bits(os.urandom(16))
packed_address = _pack_address_and_tag(address, recv_tag, send_tag)
# Receive worker address from server via multiprocessing.Queue, create
# endpoint to server
remote_address = queue.get()
ep = await ucp.create_endpoint_from_worker_address(remote_address)
# Send local address to server on tag 0
await ep.send(packed_address, tag=0, force_tag=True)
# Receive message from server
recv_msg = np.empty(10, dtype=np.int64)
await ep.recv(recv_msg, tag=recv_tag, force_tag=True)
np.testing.assert_array_equal(recv_msg, np.arange(10, dtype=np.int64))
# Send message to server
send_msg = np.arange(20, dtype=np.int64)
await ep.send(send_msg, tag=send_tag, force_tag=True)
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()
