def test_multiple_listeners():
lf1 = ucp.start_listener(hello, is_coroutine=True)
lf2 = ucp.start_listener(hello, is_coroutine=True)
assert lf1.port > 0
assert lf2.port > 0
assert lf1.port != lf2.port
lf1.coroutine.close()
lf2.coroutine.close()
ucp.stop_listener(lf1)
ucp.stop_listener(lf2)
def start(self):
async def serve_forever(client_ep, listener_instance):
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,
)
self.listener_instance = listener_instance
if self.comm_handler:
await self.comm_handler(ucx)
ucp.init()
self.ucp_server = ucp.start_listener(serve_forever,
listener_port=self._input_port,
is_coroutine=True)
try:
loop = asyncio.get_running_loop()
except (RuntimeError, AttributeError):
loop = asyncio.get_event_loop()
t = loop.create_task(self.ucp_server.coroutine)
self._task = t
async def _func_ucp_create_listener(sessionId, r):
"""
Creates a UCP listener for incoming endpoint connections.
This function runs in a loop asynchronously in the background
on the worker
:param sessionId: uuid Unique id for current instance
:param r: float a random number to stop the function from being cached
"""
if "ucp_listener" in worker_state(sessionId):
print("Listener already started for sessionId=" +
str(sessionId))
else:
ucp.init()
listener = ucp.start_listener(_connection_func, 0,
is_coroutine=True)
worker_state(sessionId)["ucp_listener"] = listener
while not listener.done():
await listener.coroutine
await asyncio.sleep(1)
del worker_state(sessionId)["ucp_listener"]
del listener
ucp.fin()
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 test_zero_port():
ucp.init()
listener = ucp.start_listener(talk_to_client,
listener_port=0,
is_coroutine=True)
assert 0 < listener.port < 2**16
ip = ucp.get_address()
await asyncio.gather(listener.coroutine,
talk_to_server(ip.encode(), listener.port))
ucp.fin()
async def main(args=None):
ucp.init()
args = parse_args(args)
host = args.host.encode()
message = args.message
port = 13337
print("1. Calling connect")
client = connect(host, port, message=message, type_=args.type)
print("2. Calling start_server")
server = ucp.start_listener(serve_wrapper(message, args.type),
port,
is_coroutine=True)
await asyncio.gather(server.coroutine, client)
print("-" * 80)
port = 13338
print("1. Calling connect")
client = connect(host, port, message, type_=args.type)
print("2. Calling start_server")
server = ucp.start_listener(serve_wrapper(message, args.type),
port,
is_coroutine=True)
await asyncio.gather(server.coroutine, client)
def serve(port, n_bytes, n_iter, recv, np, verbose, increment):
arr = np.zeros(n_bytes, dtype='u1')
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)
ep.close()
ucp.stop_listener(lf)
lf = ucp.start_listener(inc, port, is_coroutine=True)
return lf.coroutine
parser.add_argument("-n",
"--n-bytes",
help="Size of the messages (in bytes)",
type=int,
default=1024)
args = parser.parse_args()
# initiate ucp
init_str = ""
server = False
if args.server is None:
server = True
else:
server = False
init_str = args.server
ucp.init()
loop = asyncio.get_event_loop()
# coro points to either client or server-side coroutine
if server:
listener = ucp.start_listener(talk_to_client, is_coroutine=True)
coro = listener.coroutine
else:
coro = talk_to_server(init_str.encode(), int(args.port))
loop.run_until_complete(coro)
loop.close()
ucp.fin()
init_str = ""
server = False
cb_not_done = True
if args.server is None:
server = True
else:
server = False
init_str = args.server
ucp.init()
if not args.use_asyncio:
if server:
if args.intra_node:
ucp.set_cuda_dev(1)
ucp.start_listener(talk_to_client, is_coroutine=False)
while cb_not_done:
ucp.progress()
else:
talk_to_server(init_str.encode(), int(args.port))
else:
loop = asyncio.get_event_loop()
if server:
if args.intra_node:
ucp.set_cuda_dev(1)
listener = ucp.start_listener(talk_to_client_async, is_coroutine=True)
coro = listener.coroutine
else:
coro = talk_to_server_async(init_str.encode(), int(args.port))
loop.run_until_complete(coro)
action="store_true")
parser.add_argument('-b',
'--blind_recv',
help='Use blind recv. Default = False',
action="store_true")
args = parser.parse_args()
## initiate ucp
init_str = ""
server = False
if args.server is None:
server = True
else:
server = False
init_str = args.server
ucp.init()
loop = asyncio.get_event_loop()
# coro points to either client or server-side coroutine
listener = ucp.start_listener(talk_to_client,
listener_port=int(args.my_port),
is_coroutine=True)
coro_server = listener.coroutine
time.sleep(5)
coro_client = talk_to_server(init_str.encode(), int(args.port))
loop.run_until_complete(asyncio.gather(coro_server, coro_client))
loop.close()
ucp.fin()
async def main():
ucp.init()
a = ucp.start_listener(handler, is_coroutine=True)
b = spin()
await asyncio.gather(a.coroutine, b)
print("Got on", value.decode())
await ep.send_obj(value * 10)
ucp.destroy_ep(ep)
ucp.stop_listener(lf)
print('stopped serving', value.decode())
return serve
ucp.init()
loop = asyncio.get_event_loop()
listener1 = ucp.start_listener(make_server(b'a'),
listener_port=13337,
is_coroutine=True)
coro1 = listener1.coroutine
listener2 = ucp.start_listener(make_server(b'b'),
listener_port=13338,
is_coroutine=True)
coro2 = listener2.coroutine
async def main():
task1 = asyncio.create_task(coro1)
task2 = asyncio.create_task(coro2)
await task1
await task2