def stop(self):
# What all should this do?
if self._task:
self._task.cancel()
if self._ep:
ucp.destroy_ep(self._ep)
async def talk_to_client_async(ep, listener):
print("{}\t{}\t{}\t{}".format("Size (bytes)", "Latency (us)", "BW (GB/s)",
"Issue (us)", "Progress (us)"))
for i in range(msg_log):
msg_len = 2**i
send_obj = b'0' * msg_len
for j in range(warmup_iters):
await ep.send_obj(send_obj)
await recv(ep, msg_len)
start = time.time()
for j in range(max_iters):
await ep.send_obj(send_obj)
await recv(ep, msg_len)
end = time.time()
lat = end - start
print("{}\t\t{:.2f}\t\t{:.2f}".format(msg_len,
get_avg_us(lat, max_iters),
((msg_len /
(lat / 2)) / 1000000)))
ucp.destroy_ep(ep)
ucp.stop_listener(listener)
async def talk_to_server_async(ip, port):
print("{}\t{}\t{}\t{}".format("Size (bytes)", "Latency (us)", "BW (GB/s)",
"Issue (us)", "Progress (us)"))
ep = ucp.get_endpoint(ip, port)
for i in range(msg_log):
msg_len = 2**i
send_obj = b'0' * msg_len
for j in range(warmup_iters):
await recv(ep, msg_len)
await ep.send_obj(send_obj)
start = time.time()
for j in range(max_iters):
await recv(ep, msg_len)
await ep.send_obj(send_obj)
end = time.time()
lat = end - start
print("{}\t\t{:.2f}\t\t{:.2f}".format(msg_len,
get_avg_us(lat, max_iters),
((msg_len /
(lat / 2)) / 1000000)))
print("past iters")
ucp.destroy_ep(ep)
print("past ep destroy")
def talk_to_client(client_ep):
global args
global cb_not_done
msg_log = max_msg_log
iters = max_iters
comm_ep = client_ep
send_buffer_region = ucp.buffer_region()
if args.mem_type == 'cuda':
send_buffer_region.alloc_cuda(1 << msg_log)
else:
send_buffer_region.alloc_host(1 << msg_log)
print("{}\t\t{}".format("Size (bytes)", "Uni-Bandwidth (GB/s)"))
for i in range(msg_log):
msg_len = 2**i
warmup_iters = int((0.1 * iters))
for j in range(warmup_iters):
pending_list = []
for k in range(window_size):
send_msg = ucp.ucp_msg(send_buffer_region)
send_ft = comm_ep.send(send_msg, msg_len)
pending_list.append(send_ft)
while len(pending_list) > 0:
for ft in pending_list:
if ft.done() == True:
pending_list.remove(ft)
start = time.time()
for j in range(iters):
pending_list = []
for k in range(window_size):
send_msg = ucp.ucp_msg(send_buffer_region)
send_ft = comm_ep.send(send_msg, msg_len)
pending_list.append(send_ft)
while len(pending_list) > 0:
for ft in pending_list:
if ft.done() == True:
pending_list.remove(ft)
end = time.time()
lat = end - start
#lat = ((lat/2) / iters)* 1000000
bw = (iters * window_size * msg_len) / lat
bw = bw / 1e9 #GB/s
print("{}\t\t{}".format(msg_len, bw))
if args.mem_type == 'cuda':
send_buffer_region.free_cuda()
else:
send_buffer_region.free_host()
ucp.destroy_ep(client_ep)
cb_not_done = False
ucp.stop_server()
async def talk_to_server(ip, port):
global args
msg_log = max_msg_log
iters = max_iters
server_ep = ucp.get_endpoint(ip, port)
comm_ep = server_ep
send_buffer_region = ucp.buffer_region()
recv_buffer_region = ucp.buffer_region()
if args.mem_type == 'cuda':
send_buffer_region.alloc_cuda(1 << msg_log)
recv_buffer_region.alloc_cuda(1 << msg_log)
else:
send_buffer_region.alloc_host(1 << msg_log)
recv_buffer_region.alloc_host(1 << msg_log)
for i in range(msg_log):
msg_len = 2**i
warmup_iters = int((0.1 * iters))
for j in range(warmup_iters):
pending_list = []
for k in range(window_size):
send_msg = ucp.ucp_msg(send_buffer_region)
recv_msg = ucp.ucp_msg(recv_buffer_region)
send_ft = comm_ep.send(send_msg, msg_len)
recv_ft = comm_ep.recv(recv_msg, msg_len)
pending_list.append(send_ft)
pending_list.append(recv_ft)
await asyncio.wait(pending_list)
start = time.time()
for j in range(iters):
pending_list = []
for k in range(window_size):
send_msg = ucp.ucp_msg(send_buffer_region)
recv_msg = ucp.ucp_msg(recv_buffer_region)
send_ft = comm_ep.send(send_msg, msg_len)
recv_ft = comm_ep.recv(recv_msg, msg_len)
pending_list.append(send_ft)
pending_list.append(recv_ft)
await asyncio.wait(pending_list)
end = time.time()
lat = end - start
lat = ((lat / 2) / iters) * 1000000
if args.mem_type == 'cuda':
send_buffer_region.free_cuda()
recv_buffer_region.free_cuda()
else:
send_buffer_region.free_host()
recv_buffer_region.free_host()
ucp.destroy_ep(server_ep)
async def talk_to_server(ip, port):
global args
global max_msg_log
start_string = "in talk_to_server using " + args.object_type
if args.blind_recv:
start_string += " + blind recv"
print(start_string)
ep1 = ucp.get_endpoint(ip, port)
ep2 = ucp.get_endpoint(ip, port)
recv_msg = None
if not args.blind_recv:
recv_string1 = "hello from ucx client @" + socket.gethostname()
if args.validate:
recv_string1 = 'c' * (2 ** max_msg_log)
recv_string2 = "hello from ucx client @" + socket.gethostname()
if args.validate:
recv_string2 = 'c' * (2 ** max_msg_log)
recv_msg1 = get_msg(recv_string1, args.object_type)
recv_msg2 = get_msg(recv_string2, args.object_type)
recv_req1 = await ep1.recv_obj(recv_msg1, sys.getsizeof(recv_msg1))
recv_req2 = await ep2.recv_obj(recv_msg2, sys.getsizeof(recv_msg2))
else:
recv_req1 = await ep1.recv_future()
recv_req2 = await ep2.recv_future()
recv_msg1 = ucp.get_obj_from_msg(recv_req1)
recv_msg2 = ucp.get_obj_from_msg(recv_req2)
print("about to send")
send_string1 = "hello from ucx client ep2 @" + socket.gethostname()
if args.validate:
send_string = 'd' * (2 ** max_msg_log)
send_string2 = "hello from ucx client ep1 @" + socket.gethostname()
if args.validate:
send_string = 'd' * (2 ** max_msg_log)
send_msg1 = get_msg(send_string1, args.object_type)
send_msg2 = get_msg(send_string2, args.object_type)
send_req1 = await ep1.send_obj(send_msg1, sys.getsizeof(send_msg1))
send_req2 = await ep2.send_obj(send_msg2, sys.getsizeof(send_msg2))
if not args.validate:
print_msg("client sent: ", send_msg1, args.object_type)
print_msg("client sent: ", send_msg2, args.object_type)
print_msg("client received: ", recv_msg1, args.object_type)
print_msg("client received: ", recv_msg2, args.object_type)
else:
assert(recv_msg1 == get_msg('a' * (2 ** max_msg_log), args.object_type))
assert(recv_msg2 == get_msg('a' * (2 ** max_msg_log), args.object_type))
ucp.destroy_ep(ep1)
ucp.destroy_ep(ep2)
print('talk_to_server done')
async def serve(ep, lf):
print('serving', ep)
await ep.recv_future()
print("Got on", value.decode())
await ep.send_obj(value * 10)
ucp.destroy_ep(ep)
ucp.stop_listener(lf)
print('stopped serving', value.decode())
def talk_to_server(ip, port):
global args
msg_log = max_msg_log
iters = max_iters
server_ep = ucp.get_endpoint(ip, port)
comm_ep = server_ep
send_buffer_region = ucp.buffer_region()
recv_buffer_region = ucp.buffer_region()
if args.mem_type == 'cuda':
send_buffer_region.alloc_cuda(1 << msg_log)
recv_buffer_region.alloc_cuda(1 << msg_log)
else:
send_buffer_region.alloc_host(1 << msg_log)
recv_buffer_region.alloc_host(1 << msg_log)
for i in range(msg_log):
msg_len = 2 ** i
warmup_iters = int((0.1 * iters))
for j in range(warmup_iters):
send_msg = ucp.ucp_msg(send_buffer_region)
recv_msg = ucp.ucp_msg(recv_buffer_region)
recv_req = comm_ep.recv(recv_msg, msg_len)
recv_req.result()
send_req = comm_ep.send(send_msg, msg_len)
send_req.result()
send_msg = []
recv_msg = []
for j in range(iters):
send_msg.append(ucp.ucp_msg(send_buffer_region))
recv_msg.append(ucp.ucp_msg(recv_buffer_region))
start = time.time()
for j in range(iters):
recv_req = comm_ep.recv(recv_msg[j], msg_len)
recv_req.result()
send_req = comm_ep.send(send_msg[j], msg_len)
send_req.result()
end = time.time()
lat = end - start
lat = ((lat/2) / iters)* 1000000
if args.mem_type == 'cuda':
send_buffer_region.free_cuda()
recv_buffer_region.free_cuda()
else:
send_buffer_region.free_host()
recv_buffer_region.free_host()
ucp.destroy_ep(server_ep)
def _func_destroy_all(sessionId, comms_p2p):
worker_state(sessionId)["nccl"].destroy()
del worker_state(sessionId)["nccl"]
if comms_p2p:
for ep in worker_state(sessionId)["ucp_eps"]:
if ep is not None:
ucp.destroy_ep(ep)
del ep
del worker_state(sessionId)["ucp_eps"]
del worker_state(sessionId)["handle"]
async def echo_pair(cuda_info=None):
loop = asyncio.get_event_loop()
listener = ucp.start_listener(ucp.make_server(cuda_info),
is_coroutine=True)
t = loop.create_task(listener.coroutine)
client = ucp.get_endpoint(address.encode(), listener.port)
try:
yield listener, client
finally:
t.cancel()
ucp.destroy_ep(client)
async def talk_to_server(ip, port):
global max_msg_log
global args
start_string = "in talk_to_server"
if args.blind_recv:
start_string += " + blind recv"
if args.check_data:
start_string += " + data validity check"
print(start_string)
msg_log = max_msg_log
ep = ucp.get_endpoint(ip, port)
send_buffer_region = ucp.BufferRegion()
send_buffer_region.alloc_host(1 << msg_log)
send_msg = ucp.Message(send_buffer_region)
recv_msg = None
recv_buffer_region = None
recv_req = None
if not args.blind_recv:
recv_buffer_region = ucp.BufferRegion()
recv_buffer_region.alloc_host(1 << msg_log)
recv_msg = ucp.Message(recv_buffer_region)
if args.check_data:
send_msg.set_mem(1, 1 << msg_log)
if not args.blind_recv:
recv_msg.set_mem(1, 1 << msg_log)
if not args.blind_recv:
recv_req = await ep.recv(recv_msg, 1 << msg_log)
else:
recv_req = await ep.recv_future()
send_req = await ep.send(send_msg, 1 << msg_log)
if args.check_data:
errs = 0
errs = recv_req.check_mem(0, 1 << msg_log)
print("num errs: " + str(errs))
send_buffer_region.free_host()
if not args.blind_recv:
recv_buffer_region.free_host()
ucp.destroy_ep(ep)
print("done with talk_to_server")
async def talk_to_client(ep, listener):
global max_msg_log
global args
start_string = "in talk_to_client"
if args.blind_recv:
start_string += " + blind recv"
if args.check_data:
start_string += " + data validity check"
print(start_string)
msg_log = max_msg_log
send_buffer_region = ucp.buffer_region()
send_buffer_region.alloc_host(1 << msg_log)
send_msg = ucp.ucp_msg(send_buffer_region)
recv_msg = None
recv_buffer_region = None
recv_req = None
if not args.blind_recv:
recv_buffer_region = ucp.buffer_region()
recv_buffer_region.alloc_host(1 << msg_log)
recv_msg = ucp.ucp_msg(recv_buffer_region)
if args.check_data:
send_msg.set_mem(0, 1 << msg_log)
if not args.blind_recv:
recv_msg.set_mem(0, 1 << msg_log)
send_req = await ep.send(send_msg, 1 << msg_log)
if not args.blind_recv:
recv_req = await ep.recv(recv_msg, 1 << msg_log)
else:
recv_req = await ep.recv_future()
if args.check_data:
errs = 0
errs = recv_req.check_mem(1, 1 << msg_log)
print("num errs: " + str(errs))
send_buffer_region.free_host()
if not args.blind_recv:
recv_buffer_region.free_host()
ucp.destroy_ep(ep)
print("done with talk_to_client")
ucp.stop_listener(listener)
async def talk_to_client(client_ep):
global args
msg_log = max_msg_log
iters = max_iters
comm_ep = client_ep
send_buffer_region = ucp.buffer_region()
if args.mem_type == 'cuda':
send_buffer_region.alloc_cuda(1 << msg_log)
else:
send_buffer_region.alloc_host(1 << msg_log)
print("{}\t\t{}".format("Size (bytes)", "Bi-Bandwidth (GB/s)"))
for i in range(msg_log):
msg_len = 2**i
warmup_iters = int((0.1 * iters))
for j in range(warmup_iters):
pending_list = []
for k in range(window_size):
send_msg = ucp.ucp_msg(send_buffer_region)
send_ft = comm_ep.send(send_msg, msg_len)
pending_list.append(send_ft)
await asyncio.wait(pending_list)
start = time.time()
for j in range(iters):
pending_list = []
for k in range(window_size):
send_msg = ucp.ucp_msg(send_buffer_region)
send_ft = comm_ep.send(send_msg, msg_len)
pending_list.append(send_ft)
await asyncio.wait(pending_list)
end = time.time()
lat = end - start
bw = (iters * window_size * msg_len) / lat
bw = bw / 1e9 #GB/s
print("{}\t\t{}".format(msg_len, bw))
if args.mem_type == 'cuda':
send_buffer_region.free_cuda()
else:
send_buffer_region.free_host()
ucp.destroy_ep(client_ep)
ucp.stop_server()
async def talk_to_server(ip, port):
global args
global max_msg_log
msg_log = max_msg_log
start_string = "in talk_to_server"
if args.blind_recv:
start_string += " + blind recv"
if args.use_fast:
start_string += " + using fast ops"
print(start_string)
ep = ucp.get_endpoint(ip, port)
print("got endpoint")
send_buffer_region = ucp.buffer_region()
send_buffer_region.alloc_cuda(1 << msg_log)
recv_msg = None
recv_buffer_region = None
recv_req = None
if not args.blind_recv:
recv_buffer_region = ucp.buffer_region()
recv_buffer_region.alloc_cuda(1 << msg_log)
recv_msg = ucp.ucp_msg(recv_buffer_region)
send_msg = ucp.ucp_msg(send_buffer_region)
if not args.blind_recv:
if args.use_fast:
recv_req = await ep.recv_fast(recv_msg, 1 << msg_log)
else:
recv_req = await ep.recv(recv_msg, 1 << msg_log)
else:
recv_req = await ep.recv_future()
if args.use_fast:
send_req = await ep.send_fast(send_msg, 1 << msg_log)
else:
send_req = await ep.send(send_msg, 1 << msg_log)
send_buffer_region.free_cuda()
if not args.blind_recv:
recv_buffer_region.free_cuda()
ucp.destroy_ep(ep)
print("passed talk_to_server")
async def serve(ep, lf):
print("5. Starting serve")
if message:
# msg = await ep.recv_future()
size = len(client_msg)
msg = await ep.recv_obj(size)
msg = ucp.get_obj_from_msg(msg)
print("6. Server got message", bytes(msg).decode())
# response = "Got: {}".format(server_msg.decode()).encode()
await ep.send_obj(box(server_msg), name='serve-send')
print('7. Stopping server')
ucp.destroy_ep(ep)
ucp.stop_listener(lf)
def talk_to_server(ip, port):
global args
msg_log = max_msg_log
iters = max_iters
server_ep = ucp.get_endpoint(ip, port)
comm_ep = server_ep
recv_buffer_region = ucp.buffer_region()
if args.mem_type == 'cuda':
recv_buffer_region.alloc_cuda(1 << msg_log)
else:
recv_buffer_region.alloc_host(1 << msg_log)
for i in range(msg_log):
msg_len = 2**i
warmup_iters = int((0.1 * iters))
for j in range(warmup_iters):
pending_list = []
for k in range(window_size):
recv_msg = ucp.ucp_msg(recv_buffer_region)
recv_ft = comm_ep.recv(recv_msg, msg_len)
pending_list.append(recv_ft)
while len(pending_list) > 0:
for ft in pending_list:
if ft.done() == True:
pending_list.remove(ft)
for j in range(iters):
pending_list = []
for k in range(window_size):
recv_msg = ucp.ucp_msg(recv_buffer_region)
recv_ft = comm_ep.recv(recv_msg, msg_len)
pending_list.append(recv_ft)
while len(pending_list) > 0:
for ft in pending_list:
if ft.done() == True:
pending_list.remove(ft)
if args.mem_type == 'cuda':
recv_buffer_region.free_cuda()
else:
recv_buffer_region.free_host()
ucp.destroy_ep(server_ep)
async def talk_to_server(ip, port):
# recv, send
ep = await ucp.get_endpoint(ip, port, timeout=10)
if not args.blind_recv:
recv_req = await ep.recv_obj(args.n_bytes)
else:
recv_req = await ep.recv_future()
br = recv_req.get_buffer_region()
print("about to reply")
await ep.send_obj(br)
ucp.destroy_ep(ep)
print("talk_to_server done")
async def talk_to_client(client_ep):
global args
msg_log = max_msg_log
iters = max_iters
send_buffer_region = ucp.buffer_region()
recv_buffer_region = ucp.buffer_region()
if args.mem_type == 'cuda':
send_buffer_region.alloc_cuda(1 << msg_log)
recv_buffer_region.alloc_cuda(1 << msg_log)
else:
send_buffer_region.alloc_host(1 << msg_log)
recv_buffer_region.alloc_host(1 << msg_log)
send_msg = ucp.ucp_msg(send_buffer_region)
recv_msg = ucp.ucp_msg(recv_buffer_region)
print("{}\t\t{}".format("Size (bytes)", "Latency (us)"))
for i in range(msg_log):
msg_len = 2**i
warmup_iters = int((0.1 * iters))
for j in range(warmup_iters):
send_req = await client_ep.send_fast(send_msg, msg_len)
recv_req = await client_ep.recv_fast(recv_msg, msg_len)
start = time.time()
for j in range(iters):
send_req = await client_ep.send_fast(send_msg, msg_len)
recv_req = await client_ep.recv_fast(recv_msg, msg_len)
end = time.time()
lat = end - start
lat = ((lat / 2) / iters) * 1000000
print("{}\t\t{}".format(msg_len, lat))
if args.mem_type == 'cuda':
send_buffer_region.free_cuda()
recv_buffer_region.free_cuda()
else:
send_buffer_region.free_host()
recv_buffer_region.free_host()
ucp.destroy_ep(client_ep)
ucp.stop_server()
async def talk_to_client(ep):
global args
global max_msg_log
start_string = "in talk_to_client"
if args.blind_recv:
start_string += " + blind recv"
if args.use_fast:
start_string += " + using fast ops"
print(start_string)
msg_log = max_msg_log
send_buffer_region = ucp.BufferRegion()
send_buffer_region.alloc_cuda(1 << msg_log)
send_msg = ucp.Message(send_buffer_region)
recv_msg = None
recv_buffer_region = None
recv_req = None
if not args.blind_recv:
recv_buffer_region = ucp.BufferRegion()
recv_buffer_region.alloc_cuda(1 << msg_log)
recv_msg = ucp.Message(recv_buffer_region)
if args.use_fast:
send_req = await ep.send_fast(send_msg, 1 << msg_log)
else:
send_req = await ep.send(send_msg, 1 << msg_log)
if not args.blind_recv:
if args.use_fast:
recv_req = await ep.recv_fast(recv_msg, 1 << msg_log)
else:
recv_req = await ep.recv(recv_msg, 1 << msg_log)
else:
recv_req = await ep.recv_future()
send_buffer_region.free_cuda()
if not args.blind_recv:
recv_buffer_region.free_cuda()
ucp.destroy_ep(ep)
print("passed talk_to_client")
ucp.stop_listener()
async def talk_to_client(ep, listener):
global args
global max_msg_log
global count
start_string = "in talk_to_client using " + args.object_type
if args.blind_recv:
start_string += " + blind recv"
print(start_string)
print("about to send")
send_string = "hello from ucx server @" + socket.gethostname()
if args.validate:
send_string = 'a' * (2 ** max_msg_log)
send_msg = get_msg(send_string, args.object_type)
send_req = await ep.send_obj(send_msg, sys.getsizeof(send_msg))
recv_msg = None
print("about to recv")
if not args.blind_recv:
recv_string = "hello from ucx server @" + socket.gethostname()
if args.validate:
recv_string = 'b' * (2 ** max_msg_log)
recv_msg = get_msg(recv_string, args.object_type)
recv_req = await ep.recv_obj(recv_msg, sys.getsizeof(recv_msg))
else:
recv_req = await ep.recv_future()
recv_msg = ucp.get_obj_from_msg(recv_req)
if not args.validate:
print_msg("server sent: ", send_msg, args.object_type)
print_msg("server received: ", recv_msg, args.object_type)
else:
assert(recv_msg == get_msg('d' * (2 ** max_msg_log), args.object_type))
ucp.destroy_ep(ep)
print('talk_to_client done')
count += 1
if 2 == count:
ucp.stop_listener(listener)
print('past attempt to stop listener')
async def talk_to_server(ip, port):
global args
msg_log = max_msg_log
iters = max_iters
server_ep = ucp.get_endpoint(ip, port)
send_buffer_region = ucp.buffer_region()
if args.mem_type == 'cuda':
send_buffer_region.alloc_cuda(1 << msg_log)
else:
send_buffer_region.alloc_host(1 << msg_log)
send_msg = ucp.ucp_msg(send_buffer_region)
for i in range(msg_log):
msg_len = 2**i
warmup_iters = int((0.1 * iters))
for j in range(warmup_iters):
recv_req = server_ep.recv_future()
await recv_req
send_req = await server_ep.send(send_msg, msg_len)
start = time.time()
for j in range(iters):
recv_req = server_ep.recv_future()
await recv_req
send_req = await server_ep.send(send_msg, msg_len)
end = time.time()
lat = end - start
lat = ((lat / 2) / iters) * 1000000
if args.mem_type == 'cuda':
send_buffer_region.free_cuda()
else:
send_buffer_region.free_host()
ucp.destroy_ep(server_ep)
async def connect(host, port=13337, message=True, type_='bytes'):
if type_ == 'memoryview':
box = memoryview
else:
box = bytes
print("3. Starting connect")
ep = ucp.get_endpoint(host, port)
if message:
print("4. Client send")
msg = box(client_msg)
await ep.send_obj(msg, name='connect-send')
# resp = await ep.recv_future()
size = len(client_msg)
resp = await ep.recv_obj(size)
r_msg = ucp.get_obj_from_msg(resp)
print("8. Client got message: {}".format(bytes(r_msg).decode()))
print("9. Stopping client")
ucp.destroy_ep(ep)
async def connect(host, port, n_bytes, n_iter, recv, np, verbose,
increment):
ep = await ucp.get_endpoint(host.encode(), port, timeout=10)
arr = np.zeros(n_bytes, dtype='u1')
start = clock()
for i in range(n_iter):
await ep.send_obj(arr)
if recv == 'recv_into':
await ep.recv_into(arr, arr.nbytes)
else:
# This is failing right now
msg = await ep.recv_obj(arr.nbytes, cuda=np.__name__ == 'cupy')
arr = np.asarray(msg.get_obj())
stop = clock()
expected = np.ones(n_bytes, dtype='u1')
# 0 or n_iter
expected *= (int(increment) * n_iter)
np.testing.assert_array_equal(arr, expected)
took = stop - start
# 2 for round-trip, n_iter for number of trips.
print("Roundtrip benchmark")
print("-------------------")
print(f"n_iter | {n_iter}")
print(f"n_bytes | {format_bytes(n_bytes)}")
print(f"recv | {recv}")
print(f"object | {np.__name__}")
print(f"inc | {increment}")
print("\n===================")
print(format_bytes(2 * n_iter * arr.nbytes / took), '/ s')
print("===================")
await ep.recv_future()
await ep.send_obj(np.ones(1))
ucp.destroy_ep(ep)
async def inc(ep, lf):
nonlocal arr
times = []
tstart = clock()
cuda = np.__name__ == 'cupy'
for i in range(n_iter):
t0 = clock()
if recv == 'recv_into':
await ep.recv_into(arr, n_bytes)
t1 = t2 = clock()
else:
obj = await ep.recv_obj(n_bytes, cuda=cuda)
t1 = clock()
arr = np.asarray(obj.get_obj())
t2 = clock()
if increment:
arr += 1
await ep.send_obj(arr)
t3 = clock()
times.append(
(t1 - t0, t2 - t1, t3 - t2, t3 - tstart)
)
tstart = t3
if verbose:
import pandas as pd
df = pd.DataFrame(times,
columns=[recv, 'asarray', 'send', 'total'])
print('\n')
print(df)
await ep.send_obj(np.ones(1))
await ep.recv_future()
await asyncio.sleep(2)
ucp.destroy_ep(ep)
ucp.stop_listener(lf)
async def talk_to_client(ep, listener):
print("about to send")
base = b"0" * args.n_bytes
send_msg = get_msg(base, args.object_type)
await ep.send_obj(send_msg)
print("about to recv")
if not args.blind_recv:
recv_req = await ep.recv_obj(args.n_bytes)
recv_msg = get_msg(recv_req.get_obj(), args.object_type)
else:
recv_req = await ep.recv_future()
recv_msg = ucp.get_obj_from_msg(recv_req)
if not args.validate:
print("server sent: ", reprlib.repr(send_msg), type(send_msg))
print("server recv: ", reprlib.repr(recv_msg), type(recv_msg))
else:
check(send_msg, recv_msg, args.object_type)
ucp.destroy_ep(ep)
print("talk_to_client done")
ucp.stop_listener(listener)
def abort(self):
if self._ep:
ucp.destroy_ep(self._ep)
logger.debug("Destroyed UCX endpoint")
self._ep = None
def talk_to_client(client_ep):
global args
global cb_not_done
msg_log = max_msg_log
iters = max_iters
comm_ep = client_ep
send_buffer_region = ucp.buffer_region()
recv_buffer_region = ucp.buffer_region()
if args.mem_type == 'cuda':
send_buffer_region.alloc_cuda(1 << msg_log)
recv_buffer_region.alloc_cuda(1 << msg_log)
else:
send_buffer_region.alloc_host(1 << msg_log)
recv_buffer_region.alloc_host(1 << msg_log)
print("{}\t\t{}\t\t{}\t\t{}".format("Size (bytes)", "Latency (us)",
"Issue (us)", "Progress (us)"))
for i in range(msg_log):
msg_len = 2 ** i
warmup_iters = int((0.1 * iters))
for j in range(warmup_iters):
send_msg = ucp.ucp_msg(send_buffer_region)
recv_msg = ucp.ucp_msg(recv_buffer_region)
send_req = comm_ep.send(send_msg, msg_len)
recv_req = comm_ep.recv(recv_msg, msg_len)
send_req.result()
recv_req.result()
send_msg = []
recv_msg = []
for j in range(iters):
send_msg.append(ucp.ucp_msg(send_buffer_region))
recv_msg.append(ucp.ucp_msg(recv_buffer_region))
start = time.time()
issue_lat = 0
progress_lat = 0
for j in range(iters):
tmp_start = time.time()
send_req = comm_ep.send(send_msg[j], msg_len)
tmp_end = time.time()
issue_lat += (tmp_end - tmp_start)
tmp_start = time.time()
send_req.result()
tmp_end = time.time()
progress_lat += (tmp_end - tmp_start)
tmp_start = time.time()
recv_req = comm_ep.recv(recv_msg[j], msg_len)
tmp_end = time.time()
issue_lat += (tmp_end - tmp_start)
tmp_start = time.time()
recv_req.result()
tmp_end = time.time()
progress_lat += (tmp_end - tmp_start)
end = time.time()
lat = end - start
lat = ((lat/2) / iters)* 1000000
issue_lat = ((issue_lat/2) / iters)* 1000000
progress_lat = ((progress_lat/2) / iters)* 1000000
print("{}\t\t{}\t\t{}\t\t{}".format(msg_len, lat, issue_lat,
progress_lat))
if args.mem_type == 'cuda':
send_buffer_region.free_cuda()
recv_buffer_region.free_cuda()
else:
send_buffer_region.free_host()
recv_buffer_region.free_host()
ucp.destroy_ep(client_ep)
cb_not_done = False
async def talk_to_server(ip, port):
client = await ucp.get_endpoint(ip, port)
await client.send_obj(bytes(b"42"))
ucp.destroy_ep(client)
async def talk_to_client(ep, listener):
recv_req = await ep.recv_future()
recv_msg = ucp.get_obj_from_msg(recv_req)
ucp.destroy_ep(ep)
ucp.stop_listener(listener)
async def talk_to_server(ip, port):
ep = ucp.get_endpoint(ip, port)
await ep.send_obj(bytes(b"42"))
ucp.destroy_ep(ep)
