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()
def ucp_init(): try: set_env() ucp.init() yield finally: ucp.fin()
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 echo_pair(cuda_info=None): ucp.init() loop = asyncio.get_event_loop() listener = ucp.start_listener(ucp.make_server(cuda_info), is_coroutine=True) t = loop.create_task(listener.coroutine) address = ucp.get_address() client = await ucp.get_endpoint(address.encode(), listener.port) try: yield listener, client finally: ucp.destroy_ep(client) await t ucp.fin()
async def main(args=None): args = parse_args(args) if args.object_type == 'numpy': import numpy as xp else: import cupy as xp if args.server: if args.object_type == 'cupy': xp.cuda.runtime.setDevice(0) print(xp.cuda.runtime.getDevice()) await connect(args.server, args.port, args.n_bytes, args.n_iter, args.recv, xp, args.verbose, args.inc) else: if args.object_type == 'cupy': xp.cuda.runtime.setDevice(1) print(xp.cuda.runtime.getDevice()) await serve(args.port, args.n_bytes, args.n_iter, args.recv, xp, args.verbose, args.inc) ucp.fin()
def test_multiple_init_fin(): ucp.init() ucp.fin() ucp.init() ucp.fin() ucp.init() ucp.fin()
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()
async def test_timeout(): ucp.init() ip = ucp.get_address() await asyncio.gather(talk_to_server(ip.encode(), 9999, timeout=0.1)) ucp.fin()