def _build_comms(index_handler, query_handler, streams_per_handle):
# Communicator clique needs to include the union of workers hosting
# query and index partitions
workers = set(index_handler.workers)
workers.update(query_handler.workers)
comms = CommsContext(comms_p2p=True,
streams_per_handle=streams_per_handle)
comms.init(workers=workers)
return comms
def test_comms_init_p2p_no_ucx():
cluster = LocalCUDACluster()
client = Client(cluster) # noqa
cb = CommsContext(comms_p2p=True)
cb.init()
assert cb.nccl_initialized is True
assert cb.ucx_initialized is False
def test_default_comms(client):
try:
cb = CommsContext(comms_p2p=True, client=client)
cb.init()
comms = default_comms()
assert (cb.sessionId == comms.sessionId)
finally:
comms.destroy()
def _build_comms(index_futures, query_futures, streams_per_handle,
verbose):
# Communicator clique needs to include the union of workers hosting
# query and index partitions
workers = set(map(lambda x: x[0], index_futures))
workers.update(list(map(lambda x: x[0], query_futures)))
comms = CommsContext(comms_p2p=True,
streams_per_handle=streams_per_handle,
verbose=verbose)
comms.init(workers=workers)
return comms
def test_default_comms():
cluster = LocalCUDACluster(threads_per_worker=1)
client = Client(cluster)
cb = CommsContext(comms_p2p=True, client=client)
cb.init()
comms = default_comms()
assert(cb.sessionId == comms.sessionId)
comms.destroy()
client.close()
cluster.close()
def _fit_with_colocality(self, X, y):
input_futures = self.client.sync(extract_colocated_ddf_partitions,
X, y, self.client)
workers = list(input_futures.keys())
comms = CommsContext(comms_p2p=False)
comms.init(workers=workers)
worker_info = comms.worker_info(comms.worker_addresses)
n_cols = X.shape[1]
n_rows = 0
self.rnks = dict()
partsToSizes = dict()
key = uuid1()
for w, futures in input_futures.items():
self.rnks[w] = worker_info[w]["r"]
parts = [(self.client.submit(
Ridge._func_get_size_cl,
future,
workers=[w],
key="%s-%s" % (key, idx)).result())
for idx, future in enumerate(futures)]
partsToSizes[worker_info[w]["r"]] = parts
for p in parts:
n_rows = n_rows + p
key = uuid1()
self.linear_models = [(w, self.client.submit(
Ridge._func_create_model,
comms.sessionId,
**self.kwargs,
workers=[w],
key="%s-%s" % (key, idx)))
for idx, w in enumerate(workers)]
key = uuid1()
linear_fit = dict([(worker_info[wf[0]]["r"], self.client.submit(
Ridge._func_fit_colocated,
wf[1],
input_futures[wf[0]],
n_rows, n_cols,
partsToSizes,
worker_info[wf[0]]["r"],
key="%s-%s" % (key, idx),
workers=[wf[0]]))
for idx, wf in enumerate(self.linear_models)])
wait(list(linear_fit.values()))
raise_exception_from_futures(list(linear_fit.values()))
comms.destroy()
self.local_model = self.linear_models[0][1].result()
self.coef_ = self.local_model.coef_
self.intercept_ = self.local_model.intercept_
def test_recv_any_rank(n_trials, ucx_cluster):
client = Client(ucx_cluster)
try:
cb = CommsContext(comms_p2p=True)
cb.init()
dfs = [
client.submit(func_test_recv_any_rank,
cb.sessionId,
n_trials,
random.random(),
workers=[w]) for w in cb.worker_addresses
]
wait(dfs)
result = [x.result() for x in dfs]
assert result
finally:
cb.destroy()
client.close()
def fit(self, X):
"""
Fits a distributed KMeans model
:param X: dask_cudf.Dataframe to fit
:return: This KMeans instance
"""
gpu_futures = self.client.sync(extract_ddf_partitions, X)
workers = list(map(lambda x: x[0], gpu_futures.items()))
comms = CommsContext(comms_p2p=False)
comms.init(workers=workers)
key = uuid1()
kmeans_fit = [
self.client.submit(KMeans.func_fit,
comms.sessionId,
wf[1],
**self.kwargs,
workers=[wf[0]],
key="%s-%s" % (key, idx))
for idx, wf in enumerate(gpu_futures.items())
]
wait(kmeans_fit)
self.raise_exception_from_futures(kmeans_fit)
comms.destroy()
self.local_model = kmeans_fit[0].result()
self.cluster_centers_ = self.local_model.cluster_centers_
return self
def fit(self, X):
"""
Fit a multi-node multi-GPU KMeans model
Parameters
----------
X : Dask cuDF DataFrame or CuPy backed Dask Array
Training data to cluster.
"""
data = DistributedDataHandler.create(X, client=self.client)
self.datatype = data.datatype
comms = CommsContext(comms_p2p=False)
comms.init(workers=data.workers)
kmeans_fit = [
self.client.submit(KMeans._func_fit,
comms.sessionId,
wf[1],
self.datatype,
**self.kwargs,
workers=[wf[0]],
pure=False)
for idx, wf in enumerate(data.worker_to_parts.items())
]
wait_and_raise_from_futures(kmeans_fit)
comms.destroy()
self._set_internal_model(kmeans_fit[0])
return self
def test_allreduce(cluster):
client = Client(cluster)
try:
cb = CommsContext()
cb.init()
start = time.time()
dfs = [
client.submit(func_test_allreduce,
cb.sessionId,
random.random(),
workers=[w]) for wid, w in
zip(range(len(cb.worker_addresses)), cb.worker_addresses)
]
wait(dfs)
print("Time: " + str(time.time() - start))
print(str(list(map(lambda x: x.result(), dfs))))
assert all(list(map(lambda x: x.result(), dfs)))
finally:
cb.destroy()
client.close()
def test_send_recv(n_trials):
cluster = LocalCUDACluster(threads_per_worker=1)
client = Client(cluster)
cb = CommsContext(comms_p2p=True)
cb.init()
cb = default_comms()
start = time.time()
dfs = [client.submit(func_test_send_recv,
cb.sessionId,
n_trials,
random.random(),
workers=[w])
for wid, w in zip(range(len(cb.worker_addresses)),
cb.worker_addresses)]
wait(dfs)
print("Time: " + str(time.time() - start))
result = list(map(lambda x: x.result(), dfs))
print(str(result))
assert(result)
cb.destroy()
client.close()
cluster.close()
def test_recv_any_rank(n_trials, ucx_cluster):
client = Client(ucx_cluster)
try:
cb = CommsContext(comms_p2p=True)
cb.init()
dfs = [
client.submit(func_test_recv_any_rank,
cb.sessionId,
n_trials,
random.random(),
workers=[w]) for wid, w in
zip(range(len(cb.worker_addresses)), cb.worker_addresses)
]
wait(dfs)
result = list(map(lambda x: x.result(), dfs))
assert (result)
finally:
cb.destroy()
client.close()
def test_comms_init_no_p2p():
cluster = LocalCUDACluster(threads_per_worker=1)
client = Client(cluster) # noqa
cb = CommsContext(comms_p2p=False)
cb.init()
assert cb.nccl_initialized is True
assert cb.ucx_initialized is False
cb.destroy()
def test_comms_init_no_p2p(client):
try:
cb = CommsContext(comms_p2p=False)
cb.init()
assert cb.nccl_initialized is True
assert cb.ucx_initialized is False
finally:
cb.destroy()
def test_comms_init_p2p_no_ucx(cluster):
client = Client(cluster)
try:
cb = CommsContext(comms_p2p=True)
cb.init()
assert cb.nccl_initialized is True
assert cb.ucx_initialized is False
finally:
cb.destroy()
client.close()
def test_allreduce(client):
try:
cb = CommsContext()
cb.init()
dfs = [
client.submit(func_test_allreduce,
cb.sessionId,
random.random(),
workers=[w]) for w in cb.worker_addresses
]
wait(dfs, timeout=5)
assert all([x.result() for x in dfs])
finally:
cb.destroy()
def _fit(self, model_func, data, **kwargs):
n_cols = data[0].shape[1]
data = DistributedDataHandler.create(data=data, client=self.client)
self.datatype = data.datatype
comms = CommsContext(comms_p2p=False, verbose=self.verbose)
comms.init(workers=data.workers)
data.calculate_parts_to_sizes(comms)
self.ranks = data.ranks
lin_models = dict([
(data.worker_info[wf[0]]["rank"],
self.client.submit(model_func,
comms.sessionId,
self.datatype,
**self.kwargs,
pure=False,
workers=[wf[0]]))
for idx, wf in enumerate(data.worker_to_parts.items())
])
lin_fit = dict([
(wf[0],
self.client.submit(
_func_fit,
lin_models[data.worker_info[wf[0]]["rank"]],
wf[1],
data.total_rows,
n_cols,
data.parts_to_sizes[data.worker_info[wf[0]]["rank"]],
data.worker_info[wf[0]]["rank"],
pure=False,
workers=[wf[0]]))
for idx, wf in enumerate(data.worker_to_parts.items())
])
wait(list(lin_fit.values()))
raise_exception_from_futures(list(lin_fit.values()))
comms.destroy()
return lin_models
def fit(self, X):
"""
Fit a multi-node multi-GPU KMeans model
Parameters
----------
X : dask_cudf.Dataframe
Returns
-------
self: KMeans model
"""
gpu_futures = self.client.sync(extract_ddf_partitions, X)
worker_to_parts = workers_to_parts(gpu_futures)
workers = list(map(lambda x: x[0], worker_to_parts.items()))
comms = CommsContext(comms_p2p=False)
comms.init(workers=workers)
key = uuid1()
kmeans_fit = [
self.client.submit(KMeans._func_fit,
comms.sessionId,
wf[1],
**self.kwargs,
workers=[wf[0]],
key="%s-%s" % (key, idx))
for idx, wf in enumerate(worker_to_parts.items())
]
wait(kmeans_fit)
self.raise_exception_from_futures(kmeans_fit)
comms.destroy()
self.local_model = kmeans_fit[0].result()
self.cluster_centers_ = self.local_model.cluster_centers_
return self
def test_allreduce(cluster):
client = Client(cluster)
try:
cb = CommsContext()
cb.init()
dfs = [
client.submit(func_test_allreduce,
cb.sessionId,
random.random(),
workers=[w]) for w in cb.worker_addresses
]
wait(dfs)
assert all([x.result() for x in dfs])
finally:
cb.destroy()
client.close()
def test_send_recv(n_trials, ucx_client):
try:
cb = CommsContext(comms_p2p=True)
cb.init()
dfs = [
ucx_client.submit(func_test_send_recv,
cb.sessionId,
n_trials,
random.random(),
workers=[w]) for w in cb.worker_addresses
]
wait(dfs, timeout=5)
assert (list(map(lambda x: x.result(), dfs)))
finally:
cb.destroy()
def fit(self, X):
"""
Fits a distributed KMeans model
:param X: dask_cudf.Dataframe to fit
:return: This KMeans instance
"""
gpu_futures = self.client.sync(extract_ddf_partitions, X)
workers = list(map(lambda x: x[0], gpu_futures.items()))
comms = CommsContext(comms_p2p=False)
comms.init(workers=workers)
kmeans_fit = [
self.client.submit(KMeans.func_fit,
comms.sessionId,
self.n_clusters,
self.max_iter,
self.tol,
self.verbose,
self.random_state,
self.precompute_distances,
self.init,
self.n_init,
self.algorithm,
f,
random.random(),
workers=[w]) for w, f in gpu_futures.items()
]
wait(kmeans_fit)
comms.destroy()
self.local_model = kmeans_fit[0].result()
self.cluster_centers_ = self.local_model.cluster_centers_
return self
def test_send_recv(n_trials, ucx_cluster):
client = Client(ucx_cluster)
try:
cb = CommsContext(comms_p2p=True, verbose=True)
cb.init()
dfs = [
client.submit(func_test_send_recv,
cb.sessionId,
n_trials,
random.random(),
workers=[w]) for w in cb.worker_addresses
]
wait(dfs)
assert (list(map(lambda x: x.result(), dfs)))
finally:
cb.destroy()
client.close()
def fit(self, X, _transform=False):
"""
Fit the model with X.
Parameters
----------
X : dask cuDF input
"""
gpu_futures = self.client.sync(extract_ddf_partitions, X)
self.rnks = dict()
rnk_counter = 0
worker_to_parts = OrderedDict()
for w, p in gpu_futures:
if w not in worker_to_parts:
worker_to_parts[w] = []
if w not in self.rnks.keys():
self.rnks[w] = rnk_counter
rnk_counter = rnk_counter + 1
worker_to_parts[w].append(p)
workers = list(map(lambda x: x[0], gpu_futures))
comms = CommsContext(comms_p2p=False)
comms.init(workers=workers)
worker_info = comms.worker_info(comms.worker_addresses)
key = uuid1()
partsToRanks = [(worker_info[wf[0]]["r"],
self.client.submit(TruncatedSVD._func_get_size,
wf[1],
workers=[wf[0]],
key="%s-%s" % (key, idx)).result())
for idx, wf in enumerate(gpu_futures)]
N = X.shape[1]
M = reduce(lambda a, b: a + b, map(lambda x: x[1], partsToRanks))
key = uuid1()
self.tsvd_models = [
(wf[0],
self.client.submit(TruncatedSVD._func_create_model,
comms.sessionId,
wf[1],
**self.kwargs,
workers=[wf[0]],
key="%s-%s" % (key, idx)))
for idx, wf in enumerate(worker_to_parts.items())
]
key = uuid1()
tsvd_fit = dict([(worker_info[wf[0]]["r"],
self.client.submit(TruncatedSVD._func_fit,
wf[1],
M,
N,
partsToRanks,
worker_info[wf[0]]["r"],
_transform,
key="%s-%s" % (key, idx),
workers=[wf[0]]))
for idx, wf in enumerate(self.tsvd_models)])
wait(list(tsvd_fit.values()))
raise_exception_from_futures(list(tsvd_fit.values()))
comms.destroy()
self.local_model = self.client.submit(TruncatedSVD._func_get_first,
self.tsvd_models[0][1]).result()
self.components_ = self.local_model.components_
self.explained_variance_ = self.local_model.explained_variance_
self.explained_variance_ratio_ = \
self.local_model.explained_variance_ratio_
self.singular_values_ = self.local_model.singular_values_
out_futures = []
if _transform:
completed_part_map = {}
for rank, size in partsToRanks:
if rank not in completed_part_map:
completed_part_map[rank] = 0
f = tsvd_fit[rank]
out_futures.append(
self.client.submit(TruncatedSVD._func_get_idx, f,
completed_part_map[rank]))
completed_part_map[rank] += 1
return to_dask_cudf(out_futures)
def _fit(self, X, _transform=False):
"""
Fit the model with X.
Parameters
----------
X : dask cuDF input
"""
n_cols = X.shape[1]
data = DistributedDataHandler.create(data=X, client=self.client)
self.datatype = data.datatype
if "svd_solver" in self.kwargs \
and self.kwargs["svd_solver"] == "tsqr":
comms = CommsContext(comms_p2p=True)
else:
comms = CommsContext(comms_p2p=False)
comms.init(workers=data.workers)
data.calculate_parts_to_sizes(comms)
worker_info = comms.worker_info(comms.worker_addresses)
parts_to_sizes, _ = parts_to_ranks(self.client, worker_info,
data.gpu_futures)
total_rows = data.total_rows
models = dict([(data.worker_info[wf[0]]["rank"],
self.client.submit(self._create_model,
comms.sessionId,
self._model_func,
self.datatype,
**self.kwargs,
pure=False,
workers=[wf[0]]))
for idx, wf in enumerate(data.worker_to_parts.items())])
pca_fit = dict([
(wf[0],
self.client.submit(DecompositionSyncFitMixin._func_fit,
models[data.worker_info[wf[0]]["rank"]],
wf[1],
total_rows,
n_cols,
parts_to_sizes,
data.worker_info[wf[0]]["rank"],
_transform,
pure=False,
workers=[wf[0]]))
for idx, wf in enumerate(data.worker_to_parts.items())
])
wait(list(pca_fit.values()))
raise_exception_from_futures(list(pca_fit.values()))
comms.destroy()
self._set_internal_model(list(models.values())[0])
if _transform:
out_futures = flatten_grouped_results(self.client,
data.gpu_futures, pca_fit)
return to_output(out_futures, self.datatype)
return self
def _fit(self, X, _transform=False):
"""
Fit the model with X.
Parameters
----------
X : dask cuDF input
"""
n_cols = X.shape[1]
data = DistributedDataHandler.create(data=X, client=self.client)
self.datatype = data.datatype
comms = CommsContext(comms_p2p=False)
comms.init(workers=data.workers)
data.calculate_parts_to_sizes(comms)
total_rows = data.total_rows
models = dict([(data.worker_info[wf[0]]["rank"],
self.client.submit(self._create_model,
comms.sessionId,
self._model_func,
self.datatype,
**self.kwargs,
pure=False,
workers=[wf[0]]))
for idx, wf in enumerate(data.worker_to_parts.items())])
pca_fit = dict([
(wf[0],
self.client.submit(
DecompositionSyncFitMixin._func_fit,
models[data.worker_info[wf[0]]["rank"]],
wf[1],
total_rows,
n_cols,
data.parts_to_sizes[data.worker_info[wf[0]]["rank"]],
data.worker_info[wf[0]]["rank"],
_transform,
pure=False,
workers=[wf[0]]))
for idx, wf in enumerate(data.worker_to_parts.items())
])
wait(list(pca_fit.values()))
raise_exception_from_futures(list(pca_fit.values()))
comms.destroy()
self.local_model = list(models.values())[0].result()
self.components_ = self.local_model.components_
self.explained_variance_ = self.local_model.explained_variance_
self.explained_variance_ratio_ = \
self.local_model.explained_variance_ratio_
self.singular_values_ = self.local_model.singular_values_
if _transform:
out_futures = flatten_grouped_results(self.client,
data.gpu_futures, pca_fit)
return to_output(out_futures, self.datatype)
return self
return self
def _fit(self, X, y):
X_futures = self.client.sync(extract_ddf_partitions, X)
y_futures = self.client.sync(extract_ddf_partitions, y)
X_partition_workers = [w for w, xc in X_futures]
y_partition_workers = [w for w, xc in y_futures]
if set(X_partition_workers) != set(y_partition_workers):
raise ValueError("""
X and y are not partitioned on the same workers expected \n_cols
Linear Regression""")
self.rnks = dict()
rnk_counter = 0
worker_to_parts = OrderedDict()
for w, p in X_futures:
if w not in worker_to_parts:
worker_to_parts[w] = []
if w not in self.rnks.keys():
self.rnks[w] = rnk_counter
rnk_counter = rnk_counter + 1
worker_to_parts[w].append(p)
worker_to_parts_y = OrderedDict()
for w, p in y_futures:
if w not in worker_to_parts_y:
worker_to_parts_y[w] = []
worker_to_parts_y[w].append(p)
workers = list(map(lambda x: x[0], X_futures))
comms = CommsContext(comms_p2p=False)
comms.init(workers=workers)
worker_info = comms.worker_info(comms.worker_addresses)
key = uuid1()
partsToSizes = [(worker_info[wf[0]]["r"], self.client.submit(
Ridge._func_get_size,
wf[1],
workers=[wf[0]],
key="%s-%s" % (key, idx)).result())
for idx, wf in enumerate(X_futures)]
n_cols = X.shape[1]
n_rows = reduce(lambda a, b: a+b, map(lambda x: x[1], partsToSizes))
key = uuid1()
self.linear_models = [(wf[0], self.client.submit(
Ridge._func_create_model,
comms.sessionId,
**self.kwargs,
workers=[wf[0]],
key="%s-%s" % (key, idx)))
for idx, wf in enumerate(worker_to_parts.items())]
key = uuid1()
linear_fit = dict([(worker_info[wf[0]]["r"], self.client.submit(
Ridge._func_fit,
wf[1],
worker_to_parts[wf[0]],
worker_to_parts_y[wf[0]],
n_rows, n_cols,
partsToSizes,
worker_info[wf[0]]["r"],
key="%s-%s" % (key, idx),
workers=[wf[0]]))
for idx, wf in enumerate(self.linear_models)])
wait(list(linear_fit.values()))
raise_exception_from_futures(list(linear_fit.values()))
comms.destroy()
self.local_model = self.linear_models[0][1].result()
self.coef_ = self.local_model.coef_
self.intercept_ = self.local_model.intercept_