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