def get_distributed_data(input_ddf):
ddf = input_ddf
comms = Comms.get_comms()
data = DistributedDataHandler.create(data=ddf)
if data.worker_info is None and comms is not None:
data.calculate_worker_and_rank_info(comms)
return data
def _replicate_adjlist(self):
client = mg_utils.get_client()
comms = Comms.get_comms()
# FIXME: There might be a better way to control it
if client is None:
return
weights = None
offsets_futures = replication.replicate_cudf_series(
self.adjlist.offsets,
client=client,
comms=comms)
indices_futures = replication.replicate_cudf_series(
self.adjlist.indices,
client=client,
comms=comms)
if self.adjlist.weights is not None:
weights = replication.replicate_cudf_series(self.adjlist.weights)
else:
weights = {worker: None for worker in offsets_futures}
merged_futures = {worker: [offsets_futures[worker],
indices_futures[worker], weights[worker]]
for worker in offsets_futures}
self.batch_adjlists = merged_futures
def _replicate_edgelist(self):
client = mg_utils.get_client()
comms = Comms.get_comms()
# FIXME: There might be a better way to control it
if client is None:
return
work_futures = replication.replicate_cudf_dataframe(
self.edgelist.edgelist_df, client=client, comms=comms)
self.batch_edgelists = work_futures
def get_local_data(input_graph, by, load_balance=True):
input_graph.compute_renumber_edge_list(transposed=(by == 'dst'))
_ddf = input_graph.edgelist.edgelist_df
ddf = _ddf.sort_values(by=by, ignore_index=True)
if load_balance:
ddf = load_balance_func(ddf, by=by)
comms = Comms.get_comms()
data = DistributedDataHandler.create(data=ddf)
data.calculate_local_data(comms, by)
return data
def enable_batch(self):
client = mg_utils.get_client()
comms = Comms.get_comms()
if client is None or comms is None:
msg = "MG Batch needs a Dask Client and the " \
"Communicator needs to be initialized."
raise Exception(msg)
self.batch_enabled = True
if self.edgelist is not None:
if self.batch_edgelists is None:
self._replicate_edgelist()
if self.adjlist is not None:
if self.batch_adjlists is None:
self._replicate_adjlist()
if self.transposedadjlist is not None:
if self.batch_transposed_adjlists is None:
self._replicate_transposed_adjlist()
