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()