From 3937b0a77ba5b98d3a6eeed90d52466f3416f216 Mon Sep 17 00:00:00 2001 From: Iason Krommydas Date: Sun, 6 Oct 2024 10:11:21 +0200 Subject: [PATCH] show progress bar with distributed client, avoid client.compute --- scripts/run_analysis.py | 6 ++++-- tests/test_cli.py | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/scripts/run_analysis.py b/scripts/run_analysis.py index aa252751..5a6970c0 100644 --- a/scripts/run_analysis.py +++ b/scripts/run_analysis.py @@ -9,7 +9,7 @@ import dask from dask.diagnostics import ProgressBar -from dask.distributed import Client, LocalCluster, performance_report +from dask.distributed import Client, LocalCluster, performance_report, progress from egamma_tnp.config import binning_manager from egamma_tnp.utils import runner_utils @@ -198,7 +198,9 @@ def main(): if client: with performance_report(filename="/tmp/dask-report.html"): logger.info("The performance report will be saved in /tmp/dask-report.html") - (out,) = dask.compute(to_compute, scheduler="distributed") + (futures,) = dask.persist(to_compute) + progress(futures) + (out,) = dask.compute(futures) else: with ProgressBar(): (out,) = dask.compute(to_compute, scheduler=scheduler) diff --git a/tests/test_cli.py b/tests/test_cli.py index 2d6c6d4a..bc7849a1 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -29,7 +29,7 @@ def assert_arrays_equal(a1, a2): def test_cli(): subprocess.run( - "run_analysis --config tests/example_runner.json --settings tests/example_settings.json --fileset tests/example_fileset.json --binning tests/example_binning.json --output tests/output --executor threads", + "run_analysis --config tests/example_runner.json --settings tests/example_settings.json --fileset tests/example_fileset.json --binning tests/example_binning.json --output tests/output --executor distributed --cores 1 --dashboard_address 8786", shell=True, check=True, )