Merge pull request #107 from ikrommyd/progress-with-client

feat: show progress bar with distributed client without using client.compute
ikrommyd · Oct 7, 2024 · 6512391 · 6512391
2 parents b152500 + 3937b0a
commit 6512391
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,) = 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
@@ -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,
  )