async def test_get_ucp_worker():
worker = ucp.get_ucp_worker()
assert isinstance(worker, int)
def server(ep):
assert ep.get_ucp_worker() == worker
lt = ucp.create_listener(server)
ep = await ucp.create_endpoint(ucp.get_address(), lt.port)
assert ep.get_ucp_worker() == worker
def _func_build_handle_p2p(sessionId, streams_per_handle, verbose):
"""
Builds a cumlHandle on the current worker given the initialized comms
:param nccl_comm: ncclComm_t Initialized NCCL comm
:param eps: size_t initialized endpoints
:param nWorkers: int number of workers in cluster
:param workerId: int Rank of current worker
:return:
"""
ucp_worker = ucp.get_ucp_worker()
session_state = worker_state(sessionId)
handle = Handle(streams_per_handle)
nccl_comm = session_state["nccl"]
eps = session_state["ucp_eps"]
nWorkers = session_state["nworkers"]
workerId = session_state["wid"]
inject_comms_on_handle(handle, nccl_comm, ucp_worker, eps, nWorkers,
workerId, verbose)
worker_state(sessionId)["handle"] = handle
def get_worker(self):
return ucp.get_ucp_worker()
