show progress bar with client

ikrommyd · Oct 3, 2024 · 6d8af63 · 6d8af63
1 parent 661fb86
commit 6d8af63
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git 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 = client.compute(to_compute)
+            progress(futures)
+            out = client.gather(futures)
     else:
         with ProgressBar():
             (out,) = dask.compute(to_compute, scheduler=scheduler)

diff --git 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,
     )