From c1243e3b1e00744a46c09bda0cf8110e91ab5d1e Mon Sep 17 00:00:00 2001 From: Barbara Bertozzi Date: Fri, 22 Aug 2025 18:18:44 +0200 Subject: [PATCH] feat: improve cluster shutdown and cleanup logic --- scripts/sp2xr_pipeline.py | 606 ++++++++++++++++++++------------------ src/sp2xr/helpers.py | 2 +- 2 files changed, 313 insertions(+), 295 deletions(-) diff --git a/scripts/sp2xr_pipeline.py b/scripts/sp2xr_pipeline.py index 5dfd7d6..bbbb37a 100644 --- a/scripts/sp2xr_pipeline.py +++ b/scripts/sp2xr_pipeline.py @@ -1,6 +1,8 @@ from __future__ import annotations import yaml +import signal +import sys import time import dask.dataframe as dd import pandas as pd @@ -39,171 +41,139 @@ def main(): args = parse_args() run_config = load_and_resolve_config(args) - client = initialize_cluster(run_config) + client, cluster = initialize_cluster(run_config) - # -1. chunking - pbp_times = extract_partitioned_datetimes(run_config["input_pbp"]) - hk_times = extract_partitioned_datetimes(run_config["input_hk"]) - global_start = min(min(pbp_times), min(hk_times)) - global_end = max(max(pbp_times), max(hk_times)) - chunk_freq = run_config["chunking"]["freq"] # e.g. "6h", "3d" - time_chunks = get_time_chunks_from_range(global_start, global_end, chunk_freq) - - # 0. calibration stage -------------------------------------------- - instr_config = yaml.safe_load(open(run_config["instr_cfg"])) - - # 1. Bins - inc_mass_bin_lims = np.logspace( - np.log10(run_config["histo"]["inc"]["min_mass"]), - np.log10(run_config["histo"]["inc"]["max_mass"]), - run_config["histo"]["inc"]["n_bins"], - ) - inc_mass_bin_ctrs = bin_lims_to_ctrs(inc_mass_bin_lims) - - scatt_bin_lims = np.logspace( - np.log10(run_config["histo"]["scatt"]["min_D"]), - np.log10(run_config["histo"]["scatt"]["max_D"]), - run_config["histo"]["scatt"]["n_bins"], - ) - scatt_bin_ctrs = bin_lims_to_ctrs(scatt_bin_lims) - - timelag_bins_lims = np.linspace( - run_config["histo"]["timelag"]["min"], - run_config["histo"]["timelag"]["max"], - run_config["histo"]["timelag"]["n_bins"], - ) - timelag_bin_ctrs = bin_lims_to_ctrs(timelag_bins_lims) - - for chunk_start, chunk_end in time_chunks: - print(f"Processing: {chunk_start} to {chunk_end}") - - pbp_filters = [ - ("date", ">=", chunk_start.date().strftime("%Y-%m-%d")), - ("date", "<", chunk_end.date().strftime("%Y-%m-%d")), - ] - if "hour" in run_config["chunking"]["freq"]: # optionally filter by hour - pbp_filters.append(("hour", ">=", chunk_start.hour)) - pbp_filters.append(("hour", "<", chunk_end.hour)) - - # 2. HK processing -------------------------------------- + def handle_sigterm(signum, frame): + print( + f"\nSIGTERM received (signal {signum}), shutting down Dask...", flush=True + ) try: - ddf_hk = dd.read_parquet( - run_config["input_hk"], - engine="pyarrow", - filters=pbp_filters, - calculate_divisions=True, - ) - except (FileNotFoundError, OSError): - print(" → no HK files for this chunk; skipping.") - continue + client.close() + cluster.close() + except Exception as e: + print(f"Error during cleanup: {e}", flush=True) + sys.exit(0) - if ddf_hk.npartitions == 0 or partition_rowcount(ddf_hk) == 0: - print(" → HK frame is empty; skipping.") - continue - ddf_hk = ddf_hk.map_partitions(lambda pdf: pdf.sort_index()) - if not ddf_hk.known_divisions: - ddf_hk = ( - ddf_hk.reset_index().set_index( # 'calculated_time' becomes a column + signal.signal(signal.SIGTERM, handle_sigterm) + try: + # -1. chunking + pbp_times = extract_partitioned_datetimes(run_config["input_pbp"]) + hk_times = extract_partitioned_datetimes(run_config["input_hk"]) + global_start = min(min(pbp_times), min(hk_times)) + global_end = max(max(pbp_times), max(hk_times)) + chunk_freq = run_config["chunking"]["freq"] # e.g. "6h", "3d" + time_chunks = get_time_chunks_from_range(global_start, global_end, chunk_freq) + + # 0. calibration stage -------------------------------------------- + instr_config = yaml.safe_load(open(run_config["instr_cfg"])) + + # 1. Bins + inc_mass_bin_lims = np.logspace( + np.log10(run_config["histo"]["inc"]["min_mass"]), + np.log10(run_config["histo"]["inc"]["max_mass"]), + run_config["histo"]["inc"]["n_bins"], + ) + inc_mass_bin_ctrs = bin_lims_to_ctrs(inc_mass_bin_lims) + + scatt_bin_lims = np.logspace( + np.log10(run_config["histo"]["scatt"]["min_D"]), + np.log10(run_config["histo"]["scatt"]["max_D"]), + run_config["histo"]["scatt"]["n_bins"], + ) + scatt_bin_ctrs = bin_lims_to_ctrs(scatt_bin_lims) + + timelag_bins_lims = np.linspace( + run_config["histo"]["timelag"]["min"], + run_config["histo"]["timelag"]["max"], + run_config["histo"]["timelag"]["n_bins"], + ) + timelag_bin_ctrs = bin_lims_to_ctrs(timelag_bins_lims) + + for chunk_start, chunk_end in time_chunks: + print(f"Processing: {chunk_start} to {chunk_end}") + + pbp_filters = [ + ("date", ">=", chunk_start.date().strftime("%Y-%m-%d")), + ("date", "<", chunk_end.date().strftime("%Y-%m-%d")), + ] + if "hour" in run_config["chunking"]["freq"]: # optionally filter by hour + pbp_filters.append(("hour", ">=", chunk_start.hour)) + pbp_filters.append(("hour", "<", chunk_end.hour)) + + # 2. HK processing -------------------------------------- + try: + ddf_hk = dd.read_parquet( + run_config["input_hk"], + engine="pyarrow", + filters=pbp_filters, + calculate_divisions=True, + ) + except (FileNotFoundError, OSError): + print(" → no HK files for this chunk; skipping.") + continue + + if ddf_hk.npartitions == 0 or partition_rowcount(ddf_hk) == 0: + print(" → HK frame is empty; skipping.") + continue + ddf_hk = ddf_hk.map_partitions(lambda pdf: pdf.sort_index()) + if not ddf_hk.known_divisions: + ddf_hk = ddf_hk.reset_index().set_index( # 'calculated_time' becomes a column "calculated_time", sorted=False, shuffle="tasks" ) # Dask now infers divisions + ddf_hk = ddf_hk.repartition(freq="1h") + meta = pd.DataFrame( + { + "Sample Flow Controller Read (sccm)": pd.Series(dtype="float64"), + "Sample Flow Controller Read (vccm)": pd.Series(dtype="float64"), + "date": pd.Series(dtype="datetime64[ns]"), + "hour": pd.Series(dtype="int64"), + }, + index=pd.DatetimeIndex([]), ) - ddf_hk = ddf_hk.repartition(freq="1h") - meta = pd.DataFrame( - { - "Sample Flow Controller Read (sccm)": pd.Series(dtype="float64"), - "Sample Flow Controller Read (vccm)": pd.Series(dtype="float64"), - "date": pd.Series(dtype="datetime64[ns]"), - "hour": pd.Series(dtype="int64"), - }, - index=pd.DatetimeIndex([]), - ) - ddf_hk_dt = ddf_hk.map_partitions( - resample_hk_partition, dt=f"{run_config['dt']}s", meta=meta - ) - - flow_dt = ddf_hk_dt["Sample Flow Controller Read (vccm)"].compute() - - # 3. PBP processing -------------------------------------- - try: - ddf_raw = dd.read_parquet( - run_config["input_pbp"], - engine="pyarrow", - filters=pbp_filters, - calculate_divisions=True, + ddf_hk_dt = ddf_hk.map_partitions( + resample_hk_partition, dt=f"{run_config['dt']}s", meta=meta ) - except (FileNotFoundError, OSError): - print(" → no PbP files for this chunk; skipping.") - continue - if ddf_raw.npartitions == 0 or partition_rowcount(ddf_raw) == 0: - print(" → PbP frame is empty; skipping.") - continue + flow_dt = ddf_hk_dt["Sample Flow Controller Read (vccm)"].compute() - ddf_raw = ddf_raw.map_partitions(lambda pdf: pdf.sort_index()) - if not ddf_raw.known_divisions: - ddf_raw = ( - ddf_raw.reset_index().set_index( # 'calculated_time' becomes a column + # 3. PBP processing -------------------------------------- + try: + ddf_raw = dd.read_parquet( + run_config["input_pbp"], + engine="pyarrow", + filters=pbp_filters, + calculate_divisions=True, + ) + except (FileNotFoundError, OSError): + print(" → no PbP files for this chunk; skipping.") + continue + + if ddf_raw.npartitions == 0 or partition_rowcount(ddf_raw) == 0: + print(" → PbP frame is empty; skipping.") + continue + + ddf_raw = ddf_raw.map_partitions(lambda pdf: pdf.sort_index()) + if not ddf_raw.known_divisions: + ddf_raw = ddf_raw.reset_index().set_index( # 'calculated_time' becomes a column "calculated_time", sorted=False, shuffle="tasks" ) # Dask now infers divisions + ddf_raw = ddf_raw.repartition(freq="1h") + + ddf_cal = calibrate_single_particle(ddf_raw, instr_config, run_config) + + ddf_pbp_with_flow = join_pbp_with_flow(ddf_cal, flow_dt, run_config) + + delete_partition_if_exists( + output_path=f"{run_config['output']}/pbp_calibrated", + partition_values={ + "date": chunk_start.strftime("%Y-%m-%d 00:00:00"), + "hour": chunk_start.hour, + }, ) - ddf_raw = ddf_raw.repartition(freq="1h") + ddf_pbp_with_flow = enforce_schema(ddf_pbp_with_flow) - ddf_cal = calibrate_single_particle(ddf_raw, instr_config, run_config) - - ddf_pbp_with_flow = join_pbp_with_flow(ddf_cal, flow_dt, run_config) - - delete_partition_if_exists( - output_path=f"{run_config['output']}/pbp_calibrated", - partition_values={ - "date": chunk_start.strftime("%Y-%m-%d 00:00:00"), - "hour": chunk_start.hour, - }, - ) - ddf_pbp_with_flow = enforce_schema(ddf_pbp_with_flow) - - ddf_pbp_with_flow.to_parquet( - path=f"{run_config['output']}/pbp_calibrated", - partition_on=["date", "hour"], - engine="pyarrow", - write_index=True, - write_metadata_file=True, - append=True, - schema="infer", - ) - - # 4. Aggregate PBP --------------------------------------------- - ddf_pbp_dt = ddf_cal.map_partitions( - build_dt_summary, - dt_s=run_config["dt"], - meta=build_dt_summary(ddf_cal._meta), - ) - - ddf_pbp_hk_dt = aggregate_dt(ddf_pbp_dt, ddf_hk_dt, run_config) - - # 4. (optional) dt bulk conc -------------------------- - if run_config["do_conc"]: - meta_conc = add_concentrations(ddf_pbp_hk_dt._meta, dt=run_config["dt"]) - meta_conc = meta_conc.astype( - {c: CANONICAL_DTYPES.get(c, DEFAULT_FLOAT) for c in meta_conc.columns}, - copy=False, - ).convert_dtypes(dtype_backend="pyarrow") - - ddf_conc = ddf_pbp_hk_dt.map_partitions( - add_concentrations, dt=run_config["dt"], meta=meta_conc - ).map_partitions(cast_and_arrow, meta=meta_conc) - - idx_target = "datetime64[ns]" - ddf_conc = ddf_conc.map_partitions( - lambda pdf: pdf.set_index(pdf.index.astype(idx_target, copy=False)), - meta=ddf_conc._meta, - ) - - # 2) cast partition columns *before* Dask strips them off - ddf_conc["date"] = dd.to_datetime(ddf_conc["date"]).astype("datetime64[ns]") - ddf_conc["hour"] = ddf_conc["hour"].astype("int64") - - ddf_conc.to_parquet( - f"{run_config['output']}/conc_{run_config['dt']}s", + ddf_pbp_with_flow.to_parquet( + path=f"{run_config['output']}/pbp_calibrated", partition_on=["date", "hour"], engine="pyarrow", write_index=True, @@ -212,172 +182,220 @@ def main(): schema="infer", ) - # 5. (optional) dt histograms -------------------------- + # 4. Aggregate PBP --------------------------------------------- + ddf_pbp_dt = ddf_cal.map_partitions( + build_dt_summary, + dt_s=run_config["dt"], + meta=build_dt_summary(ddf_cal._meta), + ) - if run_config["do_BC_hist"]: - print("Computing BC distributions...") - # --- Mass histogram - BC_hist_configs = [ - {"flag_col": None, "flag_value": None}, - {"flag_col": "cnts_thin", "flag_value": 1}, - {"flag_col": "cnts_thin_noScatt", "flag_value": 1}, - {"flag_col": "cnts_thick", "flag_value": 1}, - {"flag_col": "cnts_thick_sat", "flag_value": 1}, - {"flag_col": "cnts_thin_sat", "flag_value": 1}, - {"flag_col": "cnts_ntl_sat", "flag_value": 1}, - {"flag_col": "cnts_ntl", "flag_value": 1}, - { - "flag_col": "cnts_extreme_positive_timelag", - "flag_value": 1, - }, - { - "flag_col": "cnts_thin_low_inc_scatt_ratio", - "flag_value": 1, - }, - {"flag_col": "cnts_thin_total", "flag_value": 1}, - {"flag_col": "cnts_thick_total", "flag_value": 1}, - {"flag_col": "cnts_unclassified", "flag_value": 1}, - ] + ddf_pbp_hk_dt = aggregate_dt(ddf_pbp_dt, ddf_hk_dt, run_config) - results = [] + # 4. (optional) dt bulk conc -------------------------- + if run_config["do_conc"]: + meta_conc = add_concentrations(ddf_pbp_hk_dt._meta, dt=run_config["dt"]) + meta_conc = meta_conc.astype( + { + c: CANONICAL_DTYPES.get(c, DEFAULT_FLOAT) + for c in meta_conc.columns + }, + copy=False, + ).convert_dtypes(dtype_backend="pyarrow") - for cfg_hist in BC_hist_configs[:2]: - meta_hist = ( - make_hist_meta( - bin_ctrs=inc_mass_bin_ctrs, - kind="mass", + ddf_conc = ddf_pbp_hk_dt.map_partitions( + add_concentrations, dt=run_config["dt"], meta=meta_conc + ).map_partitions(cast_and_arrow, meta=meta_conc) + + idx_target = "datetime64[ns]" + ddf_conc = ddf_conc.map_partitions( + lambda pdf: pdf.set_index(pdf.index.astype(idx_target, copy=False)), + meta=ddf_conc._meta, + ) + + # 2) cast partition columns *before* Dask strips them off + ddf_conc["date"] = dd.to_datetime(ddf_conc["date"]).astype( + "datetime64[ns]" + ) + ddf_conc["hour"] = ddf_conc["hour"].astype("int64") + + ddf_conc.to_parquet( + f"{run_config['output']}/conc_{run_config['dt']}s", + partition_on=["date", "hour"], + engine="pyarrow", + write_index=True, + write_metadata_file=True, + append=True, + schema="infer", + ) + + # 5. (optional) dt histograms -------------------------- + + if run_config["do_BC_hist"]: + print("Computing BC distributions...") + # --- Mass histogram + BC_hist_configs = [ + {"flag_col": None, "flag_value": None}, + {"flag_col": "cnts_thin", "flag_value": 1}, + {"flag_col": "cnts_thin_noScatt", "flag_value": 1}, + {"flag_col": "cnts_thick", "flag_value": 1}, + {"flag_col": "cnts_thick_sat", "flag_value": 1}, + {"flag_col": "cnts_thin_sat", "flag_value": 1}, + {"flag_col": "cnts_ntl_sat", "flag_value": 1}, + {"flag_col": "cnts_ntl", "flag_value": 1}, + { + "flag_col": "cnts_extreme_positive_timelag", + "flag_value": 1, + }, + { + "flag_col": "cnts_thin_low_inc_scatt_ratio", + "flag_value": 1, + }, + {"flag_col": "cnts_thin_total", "flag_value": 1}, + {"flag_col": "cnts_thick_total", "flag_value": 1}, + {"flag_col": "cnts_unclassified", "flag_value": 1}, + ] + + results = [] + + for cfg_hist in BC_hist_configs[:2]: + meta_hist = ( + make_hist_meta( + bin_ctrs=inc_mass_bin_ctrs, + kind="mass", + flag_col=cfg_hist["flag_col"], + rho_eff=run_config["rho_eff"], + BC_type=run_config["BC_type"], + ) + .astype(DEFAULT_FLOAT, copy=False) + .convert_dtypes(dtype_backend="pyarrow") + ) + ddf_out = ddf_pbp_with_flow.map_partitions( + process_hist_and_dist_partition, + col="BC mass within range", flag_col=cfg_hist["flag_col"], + flag_value=cfg_hist["flag_value"], + bin_lims=inc_mass_bin_lims, + bin_ctrs=inc_mass_bin_ctrs, + dt=run_config["dt"], + calculate_conc=True, + flow=None, rho_eff=run_config["rho_eff"], BC_type=run_config["BC_type"], + t=1, + meta=meta_hist, + ).map_partitions(cast_and_arrow, meta=meta_hist) + results.append(ddf_out) + + # --- Scattering histogram + if run_config["do_scatt_hist"]: + print("Computing scattering distribution...") + meta_hist = ( + make_hist_meta( + bin_ctrs=scatt_bin_ctrs, + kind="scatt", + flag_col=None, + rho_eff=None, + BC_type=None, ) .astype(DEFAULT_FLOAT, copy=False) .convert_dtypes(dtype_backend="pyarrow") ) - ddf_out = ddf_pbp_with_flow.map_partitions( + ddf_scatt = ddf_pbp_with_flow.map_partitions( process_hist_and_dist_partition, - col="BC mass within range", - flag_col=cfg_hist["flag_col"], - flag_value=cfg_hist["flag_value"], - bin_lims=inc_mass_bin_lims, - bin_ctrs=inc_mass_bin_ctrs, + col="Opt diam scatt only", + flag_col=None, + flag_value=None, + bin_lims=scatt_bin_lims, + bin_ctrs=scatt_bin_ctrs, dt=run_config["dt"], calculate_conc=True, flow=None, - rho_eff=run_config["rho_eff"], - BC_type=run_config["BC_type"], + rho_eff=None, + BC_type=None, t=1, meta=meta_hist, ).map_partitions(cast_and_arrow, meta=meta_hist) - results.append(ddf_out) + results.append(ddf_scatt) - # --- Scattering histogram - if run_config["do_scatt_hist"]: - print("Computing scattering distribution...") - meta_hist = ( - make_hist_meta( - bin_ctrs=scatt_bin_ctrs, - kind="scatt", - flag_col=None, - rho_eff=None, - BC_type=None, + # --- Timelag histogram + if run_config["do_timelag_hist"]: + print("Computing time delay distribution...") + mass_bins = ( + ddf_pbp_with_flow[["BC mass bin"]] + .compute() + .astype("Int64") + .drop_duplicates() + .dropna() ) - .astype(DEFAULT_FLOAT, copy=False) - .convert_dtypes(dtype_backend="pyarrow") - ) - ddf_scatt = ddf_pbp_with_flow.map_partitions( - process_hist_and_dist_partition, - col="Opt diam scatt only", - flag_col=None, - flag_value=None, - bin_lims=scatt_bin_lims, - bin_ctrs=scatt_bin_ctrs, - dt=run_config["dt"], - calculate_conc=True, - flow=None, - rho_eff=None, - BC_type=None, - t=1, - meta=meta_hist, - ).map_partitions(cast_and_arrow, meta=meta_hist) - results.append(ddf_scatt) - # --- Timelag histogram - if run_config["do_timelag_hist"]: - print("Computing time delay distribution...") - mass_bins = ( - ddf_pbp_with_flow[["BC mass bin"]] - .compute() - .astype("Int64") - .drop_duplicates() - .dropna() + for idx, mass_bin in enumerate(mass_bins[:1]): + ddf_bin = ddf_pbp_with_flow[ + ddf_pbp_with_flow["BC mass bin"] == mass_bin + ] + + name_prefix = f"dNdlogDmev_{inc_mass_bin_ctrs[idx]:.2f}_timelag" + + meta_hist = make_hist_meta( + bin_ctrs=timelag_bin_ctrs, + kind="timelag", + flag_col="cnts_particles_for_tl_dist", + name_prefix=name_prefix, + rho_eff=None, + BC_type=None, + ) + + tl_ddf = ddf_bin.map_partitions( + process_hist_and_dist_partition, + col="time_lag", + flag_col="cnts_particles_for_tl_dist", + flag_value=1, + bin_lims=timelag_bins_lims, + bin_ctrs=timelag_bin_ctrs, + dt=run_config["dt"], + calculate_conc=True, + flow=None, + rho_eff=None, + BC_type=None, + t=1, + name_prefix=name_prefix, + meta=meta_hist, + ) + + results.append(tl_ddf) + + # --- Merge all hists + merged_ddf = dd.concat(results, axis=1, interleave_partitions=True) + + idx_target = "datetime64[ns]" + merged_ddf = merged_ddf.map_partitions( + lambda pdf: pdf.set_index(pdf.index.astype(idx_target, copy=False)), + meta=merged_ddf._meta, ) - for idx, mass_bin in enumerate(mass_bins[:1]): - ddf_bin = ddf_pbp_with_flow[ - ddf_pbp_with_flow["BC mass bin"] == mass_bin - ] + index_as_dt = dd.to_datetime(merged_ddf.index.to_series()) + merged_ddf["date"] = index_as_dt.map_partitions( + lambda s: s.dt.normalize(), meta=("date", "datetime64[ns]") + ) - name_prefix = f"dNdlogDmev_{inc_mass_bin_ctrs[idx]:.2f}_timelag" + # --- Save hists to parquet - meta_hist = make_hist_meta( - bin_ctrs=timelag_bin_ctrs, - kind="timelag", - flag_col="cnts_particles_for_tl_dist", - name_prefix=name_prefix, - rho_eff=None, - BC_type=None, - ) - - tl_ddf = ddf_bin.map_partitions( - process_hist_and_dist_partition, - col="time_lag", - flag_col="cnts_particles_for_tl_dist", - flag_value=1, - bin_lims=timelag_bins_lims, - bin_ctrs=timelag_bin_ctrs, - dt=run_config["dt"], - calculate_conc=True, - flow=None, - rho_eff=None, - BC_type=None, - t=1, - name_prefix=name_prefix, - meta=meta_hist, - ) - - results.append(tl_ddf) - - # --- Merge all hists - merged_ddf = dd.concat(results, axis=1, interleave_partitions=True) - - idx_target = "datetime64[ns]" - merged_ddf = merged_ddf.map_partitions( - lambda pdf: pdf.set_index(pdf.index.astype(idx_target, copy=False)), - meta=merged_ddf._meta, - ) - - index_as_dt = dd.to_datetime(merged_ddf.index.to_series()) - merged_ddf["date"] = index_as_dt.map_partitions( - lambda s: s.dt.normalize(), meta=("date", "datetime64[ns]") - ) - - # --- Save hists to parquet - - delete_partition_if_exists( - output_path=f"{run_config['output']}/hists_{run_config['dt']}s", - partition_values={ - "date": chunk_start.strftime("%Y-%m-%d"), - "hour": chunk_start.hour, - }, - ) - merged_ddf.to_parquet( - f"{run_config['output']}/hists_{run_config['dt']}s", - partition_on=["date"], - append=True, - schema="infer", - ) - - client.close() + delete_partition_if_exists( + output_path=f"{run_config['output']}/hists_{run_config['dt']}s", + partition_values={ + "date": chunk_start.strftime("%Y-%m-%d"), + "hour": chunk_start.hour, + }, + ) + merged_ddf.to_parquet( + f"{run_config['output']}/hists_{run_config['dt']}s", + partition_on=["date"], + append=True, + schema="infer", + ) + finally: + print("Final cleanup...", flush=True) + client.close() + cluster.close() if __name__ == "__main__": diff --git a/src/sp2xr/helpers.py b/src/sp2xr/helpers.py index 6ee379f..c5379ee 100644 --- a/src/sp2xr/helpers.py +++ b/src/sp2xr/helpers.py @@ -114,7 +114,7 @@ def make_slurm_cluster(config): cluster.scale(1) client.wait_for_workers(1, timeout=600) print(f"Dask SLURM dashboard: {client.dashboard_link}") - return client + return client, cluster def make_local_cluster(config):