def main(
address,
enable_nvlink,
enable_infiniband,
):
enable_rdmacm = False
ucx_net_devices = None
if enable_infiniband:
# enable_rdmacm = True # RDMACM not working right now
ucx_net_devices = "mlx5_0:1"
# set up environment
initialize(
enable_tcp_over_ucx=True,
enable_nvlink=enable_nvlink,
enable_infiniband=enable_infiniband,
enable_rdmacm=enable_rdmacm,
net_devices=ucx_net_devices,
)
# initialize client
client = Client(address)
# user code here
rs = da.random.RandomState(RandomState=cupy.random.RandomState)
x = rs.random((10000, 10000), chunks=1000)
x.sum().compute()
# shutdown cluster
client.shutdown()
class DaskHandler(IProcessingHandler):
"""This class wraps all Dask related functions."""
def __init__(self, number_of_workers, class_cb: Callable, brain_class, worker_log_level=logging.WARNING):
super().__init__(number_of_workers)
self._client: Optional[Client] = None
self._cluster: Optional[LocalCluster] = None
self.class_cb = class_cb
self.brain_class = brain_class
self.worker_log_level = worker_log_level
def init_framework(self):
if self._client:
raise RuntimeError("Dask client already initialized.")
# threads_per_worker must be one, because atari-env is not thread-safe.
# And because lower the thread-count from the default, we must increase the number of workers
self._cluster = LocalCluster(processes=True, asynchronous=False, threads_per_worker=1,
silence_logs=self.worker_log_level,
n_workers=self.number_of_workers,
memory_pause_fraction=False,
lifetime='1 hour', lifetime_stagger='5 minutes', lifetime_restart=True,
interface="lo")
self._client = Client(self._cluster)
self._client.register_worker_plugin(_CreatorPlugin(self.class_cb, self.brain_class), name="creator-plugin")
logging.info("Dask dashboard available at port: " + str(self._client.scheduler_info()["services"]["dashboard"]))
def map(self, func, *iterable):
if not self._client:
raise RuntimeError("Dask client not initialized. Call \"init_framework\" before calling \"map\"")
return self._client.gather(self._client.map(func, *iterable))
def cleanup_framework(self):
self._client.shutdown()
def main():
"""Evaluation loop finding the best trackmate parameters in three steps.
1. Find set of thresholds based on median-normalized quantiles
2. Process all files distributed outputting the F1 score across thresholds
3. Meaning across images and determination of best parameters
"""
args = _parse_args()
basedir = args.basedir
threshold = float(args.threshold)
# Start dask client with default: localhost:8787
client = Client()
print("Dask client started.")
files = glob.glob(os.path.join(basedir, "test_predictions", "*.csv"))
def f(files):
"""Processes all files by calling delayed functions for each step."""
results = []
for file in files:
fname, pred = load_prediction(file, threshold)
true = load_true(basedir, fname)
df = process(fname, pred, true)
result = save(basedir, fname, df)
results.append(result)
return results
result = dask.compute(f(files))[0] # Returns tuple
print("Files processed.")
client.shutdown()
def run_simulations_dask(daylist, posxs, moduleWiths, kwargs):
# Create client
scheduler_file = '/scratch/sayala/dask_testing/scheduler.json'
client = Client(scheduler_file=scheduler_file)
# Iterate over inputs
futures = []
# Add Iterations HERE
for daydate in daylist:
for posx in posxs:
for moduleWith in moduleWiths:
futures.append(
client.submit(simulate_single,
daydate=daydate,
posx=posx,
moduleWith=moduleWith,
**kwargs))
# Get results for all simulations
res = client.gather(futures)
# Close all dask workers and scheduler
try:
client.shutdown()
except:
pass
# Close client
client.close()
res = 'FINISHED!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'
return res
def run_simulations_dask(xgaps, numpanelss, sensorsxs, kwargs):
# Create client
scheduler_file = '/scratch/sayala/dask_testing/scheduler.json'
client = Client(scheduler_file=scheduler_file)
# Iterate over inputs
futures = []
for nn in range(0, len(numpanelss)):
numpanels = numpanelss[nn]
for xx in range(0, len(xgaps)):
xgap = xgaps[xx]
for ii in sensorsxs:
futures.append(
client.submit(simulate_single,
xgap=xgap,
numpanels=numpanels,
sensorx=ii,
**kwargs))
# Get results for all simulations
res = client.gather(futures)
# Close all dask workers and scheduler
try:
client.shutdown()
except:
pass
# Close client
client.close()
res = 'FINISHED!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'
return res
def run_simulations_dask(clearance_heights, xgaps, Ds, tilts, kwargs):
# Create client
scheduler_file = '/scratch/sayala/dask_testing/scheduler.json'
client = Client(scheduler_file=scheduler_file)
# Iterate over inputs
futures = []
for ch in range (0, len(clearance_heights)):
clearance_height = clearance_heights[ch]
for xx in range (0, len(xgaps)):
xgap = xgaps[xx]
for tt in range (0, len(tilts)):
tilt = tilts[tt]
for dd in range (0, len(Ds)):
D = Ds[dd]
futures.append(client.submit(simulate_single, clearance_height=clearance_height,
xgap=xgap, tilt=tilt, D=D, **kwargs))
# Get results for all simulations
res = client.gather(futures)
# Close all dask workers and scheduler
try:
client.shutdown()
except:
pass
# Close client
client.close()
res = 'FINISHED!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'
return res
def run_simulations_dask(tilts, kwargs):
# Create client
scheduler_file = '/scratch/sayala/dask_testing/scheduler.json'
client = Client(scheduler_file=scheduler_file)
# Iterate over inputs
futures = []
# Add Iterations HERE
for tilt in tilts:
futures.append(client.submit(simulate_single, tilt=tilt, **kwargs))
# Get results for all simulations
res = client.gather(futures)
# Close all dask workers and scheduler
try:
client.shutdown()
except:
pass
# Close client
client.close()
res = 'FINISHED!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'
return res
def DASK_batch_mult(matrix_input, vector_input, workers, batch_size,
input_size, output_channels):
client = Client(n_workers=workers)
results = []
batch_no = matrix_input.shape[0] // batch_size
for i in range(batch_no):
batch = client.scatter(matrix_input[i * batch_size:i * batch_size +
batch_size])
results.append(
client.submit(convolution_mean, batch, vector_input, batch_size,
vector_input.shape[0]))
wait(results)
data = client.gather(results)
out_tensor = np.empty(
(batch_size * batch_no, output_channels, input_size, input_size))
for i in range(batch_no):
out_tensor[i * batch_size:i * batch_size +
batch_size] = data[i].reshape(batch_size, output_channels,
input_size, input_size)
client.shutdown()
return out_tensor
class GenerateHindcastEnsembleDaskDistributed(GenerateHindcastEnsembleDask):
def setup(self, *args, **kwargs):
"""The same tests but on spatially chunked data with dask.distributed.Client."""
_skip_slow()
requires_dask()
super().setup(**kwargs)
self.client = Client()
def cleanup(self):
self.client.shutdown()
def _stop_dask_scheduler(args):
print("Start stop scheduler.")
client = Client("{}:{}".format(args.master, args.port))
try:
client.shutdown()
client.close()
del client
# print("Scheduler stopped.")
time.sleep(8)
except Exception:
print("Failed to stop scheduler, please stop it manually.")
def _stop_dask_scheduler(args):
global _port
_print("Start stop scheduler.")
client = Client(f"{args.master}:{_port}")
try:
client.shutdown()
client.close()
del client
time.sleep(8)
except Exception:
_print("Failed to stop scheduler, please stop it manually.")
def managed_client(devices, device_limit, protocol):
client = Client(
LocalCUDACluster(
protocol=protocol,
n_workers=len(devices.split(",")),
enable_nvlink=(protocol == "ucx"),
device_memory_limit=device_limit,
))
try:
yield client
finally:
client.shutdown()
class ComputeDaskDistributed(ComputeDask):
def setup(self, *args, **kwargs):
"""Benchmark time and peak memory of `compute_perfect_model` and
`bootstrap_perfect_model`. This executes the same tests as `Compute` but
on chunked data with dask.distributed.Client."""
requires_dask()
# magic taken from
# https://github.com/pydata/xarray/blob/stable/asv_bench/benchmarks/rolling.py
super().setup(**kwargs)
self.client = Client()
def cleanup(self):
self.client.shutdown()
def main():
"""Evaluation loop finding the best trackmate parameters in three steps.
1. Find set of thresholds based on median-normalized quantiles
2. Process all files distributed outputting the F1 score across thresholds
3. Meaning across images and determination of best parameters
"""
args = _parse_args()
basedir = args.basedir
# Start dask client with default: localhost:8787
client = Client()
print("Dask client started.")
# Calculate thresholds
files = glob.glob(os.path.join(basedir, "train_predictions", "*.csv"))
thresholds = get_thresholds(basedir, files)
print("Thresholds calculated / loaded.")
# Process files
def f(files):
"""Processes all files by calling delayed functions for each step."""
results = []
for file in files:
fname, detector, radius, pred_all = load_prediction(file)
true = load_true(basedir, fname)
df = process(fname, detector, radius, pred_all, true, thresholds)
result = save(basedir, fname, df)
results.append(result)
return results
# Compute files in batch
length = len(files)
batch_size = 100
for idx in range(0, length, batch_size):
dask.compute(f(files[idx:min(idx + batch_size, length)]))
print("Files processed.")
# Compute maximal scores
df_max = get_max_scores(basedir)
with open(os.path.join(basedir, "info.txt"), "a") as f:
f.write(f"Optimal parameters found are:")
f.write(f'... F1 score: {df_max["f1_score"][0]}')
f.write(f'... Detector: {df_max["detector"][0]}')
f.write(f'... Radius: {df_max["radius"][0]}')
f.write(f'... Threshold: {df_max["threshold"][0]}')
client.shutdown()
class DaskCluster:
def __init__(self):
self.client = None
def start_local_cluster(self):
cluster = LocalCluster(
n_workers=dask_options_dict['n_workers'],
threads_per_worker=dask_options_dict['threads_per_worker'],
memory_limit=dask_options_dict['memory_limit']
) # threads_per_worker=1 needed if using numba :(
self.client = Client(cluster)
def stop_local_cluster(self):
self.client.shutdown()
self.client = None
def code_range(path_dict):
"""
path_dict = RANGE_CATEGORIES[5]
"""
# Create a numeric dictionary for these keys
key_dict = {key: i + 1 for i, key in enumerate(path_dict.keys())}
number_dict = {k: i for i, k in key_dict.items()}
vals = key_dict.values()
combos = [[c for c in combinations(vals, i + 1)] for i in range(len(vals))]
combos = [c for sc in combos for c in sc]
combo_keys = {}
for combo in combos:
key = "-".join([number_dict[c] for c in combo])
value = seq_int(combo)
combo_keys[key] = value
# Assign each raster a unique value
arrays = []
for key, path in path_dict.items():
value = key_dict[key]
full_path = DP.join(path)
array = xr.open_rasterio(full_path, chunks=CHUNKS)[0].data
array[da.isnan(array)] = 0
array[array > 0] = value
arrays.append(array)
# Stack everything together - we might have to save this a temporary file
stack = da.stack(arrays, axis=0)
stack = stack.rechunk((stack.shape[0], 5000, 5000))
stack = stack[:, 4000:10000, 4000:10000]
# Try to map the function to each point
client = Client()
codes = da.apply_along_axis(seq_int, 0, stack, dtype="uint8")
future = client.compute(codes)
result = future.result()
client.shutdown()
client.close()
# Save to temp and delete
template = rasterio.open(full_path)
temp_path = DP2.join("test.tif")
with rasterio.Env():
profile = template.profile
profile.update(dtype=rasterio.uint8, count=1, compress='lzw')
with rasterio.open(temp_path, 'w', **profile) as dst:
dst.write(result)
def test_do_not_close_external_client(self):
tmp_dir = tempfile.mkdtemp()
single_worker_mock = unittest.mock.Mock()
client = Client()
parallel_runner = DaskParallelRunner(
single_worker=single_worker_mock, dask_client=client, n_workers=1, output_directory=tmp_dir
) # noqa F841
del parallel_runner
self.assertFalse(os.path.exists(os.path.join(tmp_dir, '.dask_scheduler_file')))
self.assertEqual(client.status, 'running')
parallel_runner = DaskParallelRunner(
single_worker=single_worker_mock, dask_client=client, n_workers=1, output_directory=tmp_dir
) # noqa F841
del parallel_runner
self.assertEqual(client.status, 'running')
client.shutdown()
def main():
object = ps.preprocess()
X_train, X_test, y_train, y_test = object.cleaning()
param_grid = {
'objective': ['binary:logistic'],
'nround': [1000],
'max_depth': [8]
}
estimator = xgb.XGBRegressor()
grid_search = GridSearchCV(estimator, param_grid, verbose=2, cv=2, n_jobs=-1)
client = Client(processes=False)
start_time = time.time()
with joblib.parallel_backend("dask"):
grid_search.fit(X_train, y_train)
end_time = time.time()
grid_search.predict(X_test)
print ("time difference in GridSearchCV first XGBRegressor is %d seconds" % end_time)
client.shutdown()
def cache_sofk(fout, fjson, Sk, nx, verbose=True, para=True):
if para:
import dask
from dask.distributed import Client, progress
client = Client(processes=False)
Sk = dask.delayed(Sk)
import pandas as pd
mdf = pd.read_json(fjson)
box = mdf.iloc[0]['box']
kvecs = legal_kvecs(nx, box)
nk = len(kvecs)
if verbose and (not para):
from progressbar import ProgressBar, Bar, ETA
widgets = [Bar('>'), ETA()]
bar = ProgressBar(widgets=widgets, maxval=len(mdf))
bar.start()
skl = []
sk2l = []
isk = 0
for label, row in mdf.iterrows():
com = np.array(row['positions'])
sk = Sk(kvecs, com)
skl.append(sk)
sk2l.append(sk**2)
isk += 1
if verbose and (not para):
bar.update(isk)
skm = np.mean(skl, axis=0)
ske = (np.mean(sk2l, axis=0)-skm**2)**0.5/len(skl)**0.5
if para:
skm, ske = dask.persist(skm, ske)
if verbose:
progress(skm, ske)
skm, ske = dask.compute(skm, ske)
client.shutdown()
# spherical average
uk, uskm, uske = shavg(kvecs, skm, ske)
np.savetxt(fout, np.array([uk, uskm, uske]).T)
def start_dask(num_workers, msg, logger):
"""Context manager used for starting/shutting down dask
Args:
num_workers (`int`): Number of dask workers
msg (`str`): Message for timer
logger: The logger being used
Yields:
client: Dask client
"""
# Update dask
with open("dask-config.yaml") as f:
config = yaml.load(f, Loader=SafeLoader)
dask.config.update(dask.config.config, config)
cluster_type = next(iter(dask.config.config['jobqueue']))
set_local_directory(cluster_type)
if cluster_type == 'local':
from dask.distributed import LocalCluster
cluster = LocalCluster(n_workers=num_workers, threads_per_worker=1)
else:
if cluster_type == 'lsf':
from dask_jobqueue import LSFCluster
cluster = LSFCluster()
elif cluster_type == 'slurm':
from dask_jobqueue import SLURMCluster
cluster = SLURMCluster()
elif cluster_type == 'sge':
from dask_jobqueue import SGECluster
cluster = SGECluster()
cluster.scale(num_workers)
try:
with io_util.Timing_Messager(f"Starting dask cluster for {msg}",
logger):
client = Client(cluster)
io_util.print_with_datetime(
f"Check {client.cluster.dashboard_link} for {msg} status.", logger)
yield client
finally:
client.shutdown()
client.close()
def connection():
"""
Creates a mock Dask cluster. Returns the corresponding connection object.
Note that the name of this fixture will be used to inject the result into an
equally named variable that should be used as input argument to all tests
that need it. Example:
```
def test_my_feature(connection):
df = RDataFrame(10, daskclient=connection)
```
"""
connection = Client(
LocalCluster(n_workers=2,
threads_per_worker=1,
processes=True,
memory_limit="2GiB"))
yield connection
connection.shutdown()
connection.close()
def main():
object = ps.preprocess()
X_train, X_test, y_train, y_test = object.cleaning()
params = {'objective': 'binary:logistic',
'max_depth': 8, 'eta': 0.01, 'subsample': 0.5,
'min_child_weight': 1}
print ("Start training dxgb")
cluster = LocalCluster(n_workers=8, threads_per_worker=1)
client = Client(cluster)
start_time = time.time()
bst = dxgb.train(client, params, X_train, y_train)
end_time = time.time()
#time difference in dXGB is 1588108665
print ("time difference in dXGB is %d seonds" % end_time)
predictions = dxgb.predict(client, bst, X_test)
#Accuracy = 0.6968888393419537
print ("Accuracy score is : ")
print(roc_auc_score(y_test.compute(),
predictions.compute()))
client.shutdown()
def dask_client():
dask_scheduler_file = os.environ.get("SCHEDULER_FILE")
cluster = None
client = None
tempdir_object = None
if dask_scheduler_file:
# Env var UCX_MAX_RNDV_RAILS=1 must be set too.
initialize(
enable_tcp_over_ucx=True,
enable_nvlink=True,
enable_infiniband=True,
enable_rdmacm=True,
# net_devices="mlx5_0:1",
)
client = Client(scheduler_file=dask_scheduler_file)
print("\ndask_client fixture: client created using "
f"{dask_scheduler_file}")
else:
# The tempdir created by tempdir_object should be cleaned up once
# tempdir_object goes out-of-scope and is deleted.
tempdir_object = tempfile.TemporaryDirectory()
cluster = LocalCUDACluster(local_directory=tempdir_object.name)
client = Client(cluster)
client.wait_for_workers(len(get_visible_devices()))
print("\ndask_client fixture: client created using LocalCUDACluster")
Comms.initialize(p2p=True)
yield client
Comms.destroy()
# Shut down the connected scheduler and workers
# therefore we will no longer rely on killing the dask cluster ID
# for MNMG runs
client.shutdown()
if cluster:
cluster.close()
print("\ndask_client fixture: client.close() called")
def run_simulations_dask(arraysimulations, kwargs):
# Create client
scheduler_file = '/scratch/sayala/dask_testing/scheduler.json'
client = Client(scheduler_file=scheduler_file)
# Iterate over inputs
futures = []
for ii in range(0, len(arraysimulations)):
idx = arraysimulations.iloc[ii].idx
wavelength = arraysimulations.iloc[ii].wavelength
test_folder = test_folder_fmt.format(f'{idx:04}', f'{wavelength:04}')
if not os.path.exists(test_folder):
futures.append(
client.submit(simulate_single,
idx=idx,
wavelength=wavelength,
**kwargs))
else:
print("\n\nAlready simulated ***********************\n\n", idx,
wavelength)
# Get results for all simulations
res = client.gather(futures)
# Close all dask workers and scheduler
try:
client.shutdown()
except:
pass
# Close client
client.close()
res = 'FINISHED!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'
return res
def managed_client(dask_workdir, devices, device_limit, protocol):
if protocol == "tcp":
client = Client(
LocalCUDACluster(
protocol=protocol,
n_workers=len(devices.split(",")),
CUDA_VISIBLE_DEVICES=devices,
device_memory_limit=device_limit,
local_directory=dask_workdir,
))
else:
client = Client(
LocalCUDACluster(
protocol=protocol,
n_workers=len(devices.split(",")),
CUDA_VISIBLE_DEVICES=devices,
enable_nvlink=True,
device_memory_limit=device_limit,
local_directory=dask_workdir,
))
try:
yield client
finally:
client.shutdown()
def connection():
"""
Creates a mock Dask cluster as well as a mock Spark cluster. Returns a tuple
with both connection objects.
Note that the name of this fixture will be used to inject the result into an
equally named variable that should be used as input argument to all tests
that need it. Example:
```
def test_my_feature(connection):
# The tests in this folder need both dask and spark connections
daskclient, sparkcontext = connection
df1 = RDataFrame(10, daskclient=daskclient)
df2 = RDataFrame(10, sparkcontext=sparkcontext)
```
"""
daskconn = Client(
LocalCluster(n_workers=2,
threads_per_worker=1,
processes=True,
memory_limit="2GiB"))
conf = {
"spark.master": "local[2]",
"spark.driver.memory": "4g",
"spark.app.name": "roottest-distrdf-common"
}
sparkconf = pyspark.SparkConf().setAll(conf.items())
sparkconn = pyspark.SparkContext(conf=sparkconf)
yield daskconn, sparkconn
daskconn.shutdown()
daskconn.close()
sparkconn.stop()
print("Merging outputs")
hists = {ds[0]: [] for ds in datasets}
numevents = {ds[0]: 0 for ds in datasets}
for r, _args in zip(ret, arglist):
rh = r["hists"]
ds = _args[2]
hists[ds] += [Results(r["hists"])]
numevents[ds] += r["numevents"]
timing = sum([r["timing"] for r in ret], Results({}))
timing["cuda"] = use_cuda
timing["njec"] = args.njec
timing["nthreads"] = args.nthreads
timing["walltime"] = walltime_t1 - walltime_t0
for k, v in hists.items():
hists[k] = sum(hists[k], Results({}))
print("Writing output pkl")
with open(args.out, "wb") as fi:
pickle.dump({
"hists": hists,
"numevents": numevents,
"timing": timing
}, fi)
client.shutdown()
sum_numev = sum(numevents.values())
print("Processed {0} events in {1:.1f} seconds, {2:.2E} Hz".format(
sum_numev, timing["walltime"], sum_numev / timing["walltime"]))
def main(args):
cluster_options = get_cluster_options(args)
Cluster = cluster_options["class"]
cluster_args = cluster_options["args"]
cluster_kwargs = cluster_options["kwargs"]
scheduler_addr = cluster_options["scheduler_addr"]
if args.sched_addr:
client = Client(args.sched_addr)
else:
filterwarnings("ignore",
message=".*NVLink.*rmm_pool_size.*",
category=UserWarning)
cluster = Cluster(*cluster_args, **cluster_kwargs)
if args.multi_node:
import time
# Allow some time for workers to start and connect to scheduler
# TODO: make this a command-line argument?
time.sleep(15)
client = Client(scheduler_addr if args.multi_node else cluster)
if args.type == "gpu":
client.run(
setup_memory_pool,
pool_size=args.rmm_pool_size,
disable_pool=args.disable_rmm_pool,
log_directory=args.rmm_log_directory,
)
# Create an RMM pool on the scheduler due to occasional deserialization
# of CUDA objects. May cause issues with InfiniBand otherwise.
client.run_on_scheduler(
setup_memory_pool,
pool_size=1e9,
disable_pool=args.disable_rmm_pool,
log_directory=args.rmm_log_directory,
)
scheduler_workers = client.run_on_scheduler(get_scheduler_workers)
n_workers = len(scheduler_workers)
client.wait_for_workers(n_workers)
if args.all_to_all:
all_to_all(client)
took_list = []
for _ in range(args.runs - 1):
took_list.append(run(client, args, n_workers, write_profile=None))
took_list.append(
run(client, args, n_workers,
write_profile=args.profile)) # Only profiling the last run
# Collect, aggregate, and print peer-to-peer bandwidths
incoming_logs = client.run(
lambda dask_worker: dask_worker.incoming_transfer_log)
bandwidths = defaultdict(list)
total_nbytes = defaultdict(list)
for k, L in incoming_logs.items():
for d in L:
if d["total"] >= args.ignore_size:
bandwidths[k, d["who"]].append(d["bandwidth"])
total_nbytes[k, d["who"]].append(d["total"])
bandwidths = {(scheduler_workers[w1].name, scheduler_workers[w2].name): [
"%s/s" % format_bytes(x)
for x in numpy.quantile(v, [0.25, 0.50, 0.75])
]
for (w1, w2), v in bandwidths.items()}
total_nbytes = {(
scheduler_workers[w1].name,
scheduler_workers[w2].name,
): format_bytes(sum(nb))
for (w1, w2), nb in total_nbytes.items()}
t_runs = numpy.empty(len(took_list))
if args.markdown:
print("```")
print("Shuffle benchmark")
print("-------------------------------")
print(f"backend | {args.backend}")
print(f"partition-size | {format_bytes(args.partition_size)}")
print(f"in-parts | {args.in_parts}")
print(f"protocol | {args.protocol}")
print(f"device(s) | {args.devs}")
if args.device_memory_limit:
print(f"memory-limit | {format_bytes(args.device_memory_limit)}")
print(f"rmm-pool | {(not args.disable_rmm_pool)}")
if args.protocol == "ucx":
print(f"tcp | {args.enable_tcp_over_ucx}")
print(f"ib | {args.enable_infiniband}")
print(f"nvlink | {args.enable_nvlink}")
print(f"data-processed | {format_bytes(took_list[0][0])}")
print("===============================")
print("Wall-clock | Throughput")
print("-------------------------------")
for idx, (data_processed, took) in enumerate(took_list):
throughput = int(data_processed / took)
m = format_time(took)
m += " " * (15 - len(m))
print(f"{m}| {format_bytes(throughput)}/s")
t_runs[idx] = float(format_bytes(throughput).split(" ")[0])
print("===============================")
if args.markdown:
print("\n```")
if args.plot is not None:
plot_benchmark(t_runs, args.plot, historical=True)
if args.backend == "dask":
if args.markdown:
print(
"<details>\n<summary>Worker-Worker Transfer Rates</summary>\n\n```"
)
print("(w1,w2) | 25% 50% 75% (total nbytes)")
print("-------------------------------")
for (d1, d2), bw in sorted(bandwidths.items()):
fmt = ("(%s,%s) | %s %s %s (%s)" if args.multi_node or
args.sched_addr else "(%02d,%02d) | %s %s %s (%s)")
print(fmt % (d1, d2, bw[0], bw[1], bw[2], total_nbytes[(d1, d2)]))
if args.markdown:
print("```\n</details>\n")
if args.benchmark_json:
bandwidths_json = {
"bandwidth_({d1},{d2})_{i}" if args.multi_node or args.sched_addr
else "(%02d,%02d)_%s" % (d1, d2, i): parse_bytes(v.rstrip("/s"))
for (d1, d2), bw in sorted(bandwidths.items()) for i, v in zip(
["25%", "50%", "75%", "total_nbytes"],
[bw[0], bw[1], bw[2], total_nbytes[(d1, d2)]],
)
}
with open(args.benchmark_json, "a") as fp:
for data_processed, took in took_list:
fp.write(
dumps(
dict(
{
"backend": args.backend,
"partition_size": args.partition_size,
"in_parts": args.in_parts,
"protocol": args.protocol,
"devs": args.devs,
"device_memory_limit":
args.device_memory_limit,
"rmm_pool": not args.disable_rmm_pool,
"tcp": args.enable_tcp_over_ucx,
"ib": args.enable_infiniband,
"nvlink": args.enable_nvlink,
"data_processed": data_processed,
"wall_clock": took,
"throughput": data_processed / took,
},
**bandwidths_json,
)) + "\n")
if args.multi_node:
client.shutdown()
client.close()
c = None
if len(sys.argv) > 1 and int(sys.argv[1]) > 1:
print(f"Starting dask Client with {sys.argv[1]} processors")
c = Client(threads_per_worker=1, n_workers=int(sys.argv[1]))
print("client started...")
# load in gadget, set a selection
ds = yt.load_sample("snapshot_033")
ptf = {'PartType0': ['Mass']}
sp = ds.sphere(ds.domain_center, (2, 'code_length'))
select_time = time.time()
glob_stats, chunk_stats = ga.selector_stats(ds, ptf, sp.selector)
# glob_stats, chunk_stats = ga.selector_stats(ds,ptf) # for no selector
select_time = time.time() - select_time
print("\nlocal stats for each chunk:")
for ch in chunk_stats:
print(ch)
print("\nglobal, cross-chunk stats")
print(glob_stats)
print(
f"\nCompute time (neglecting Client spinup and yt initialization): {select_time}s"
)
if c is not None:
print("\nshutting down dask client")
c.shutdown()
def test_prefect_executors(train_data, grid_search, parallel_columns):
try:
from prefect.engine.executors import DaskExecutor
from prefect.engine.executors import LocalDaskExecutor
from prefect.engine.executors import LocalExecutor
from dask.distributed import Client
except:
print("`prefect` not installed, skipping the test...")
pass
else:
client = Client()
executors = {
"dask_already_running":
DaskExecutor(address=client.scheduler.address),
"local": LocalExecutor(),
"local_dask": LocalDaskExecutor(),
"dask_create_on_call": DaskExecutor(
), # this spins up LocalDaskExecutor, but just to check the interface
}
for executor_name, executor in executors.items():
flow, state = run_model_selection(
df=train_data,
grid_search=grid_search,
target_col_name="Quantity",
frequency="D",
partition_columns=["Product"],
parallel_over_columns=parallel_columns,
include_rules=None,
exclude_rules=None,
country_code_column="Holidays_code",
output_path="",
persist_cv_results=False,
persist_cv_data=False,
persist_model_reprs=False,
persist_best_model=False,
persist_partition=False,
persist_model_selector_results=False,
visualize_success=False,
executor=executor,
)
assert state.is_successful()
results = select_model_general(
df=train_data,
grid_search=grid_search,
target_col_name="Quantity",
frequency="D",
partition_columns=["Product"],
parallel_over_columns=parallel_columns,
executor=executor,
include_rules=None,
exclude_rules=None,
country_code_column="Holidays_code",
output_path="",
persist_cv_results=False,
persist_cv_data=False,
persist_model_reprs=False,
persist_best_model=False,
persist_partition=False,
persist_model_selector_results=False,
)
assert len(results) == len(
train_data[parallel_columns + ["Product"]].drop_duplicates())
assert isinstance(results[0], ModelSelectorResult)
if executor_name == "dask_already_running":
client.shutdown()
if client.status != "closed":
client.shutdown()