def test_pickle(n_streams):
a = Handle(n_streams=n_streams)
ap = pickle.dumps(a)
b = pickle.loads(ap)
assert isinstance(b, Handle)
assert b.getNumInternalStreams() == n_streams
def _func_build_handle(sessionId, streams_per_handle):
"""
Builds a cumlHandle on the current worker given the initialized comms
:param nccl_comm: ncclComm_t Initialized NCCL comm
:param nWorkers: int number of workers in cluster
:param workerId: int Rank of current worker
:return:
"""
handle = Handle(streams_per_handle)
session_state = worker_state(sessionId)
workerId = session_state["wid"]
nWorkers = session_state["nworkers"]
nccl_comm = session_state["nccl"]
inject_comms_on_handle_coll_only(handle, nccl_comm, nWorkers, workerId)
session_state["handle"] = handle
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