feat: progress bar for custom dask.distributed setups (#32)

* add progress bar for xrdcp notebook
* add progress bar for main notebook

Author: alexander-held

atlas/analysis.ipynb

@@ -267,12 +267,11 @@
     "\n",
     "if USE_CUSTOM_PROCESSING:\n",
     "    # configure here whether to preload branches\n",
-    "    columns_to_preload = json.load(pathlib.Path(\"columns_to_preload.json\").open())[\"JET_JER_Effective\"]\n",
-    "    columns_to_preload = []\n",
+    "    columns_to_preload = json.load(pathlib.Path(\"columns_to_preload.json\").open())[\"JET_JER_Effective\"]  # or []\n",
     "\n",
     "    with performance_report(filename=\"process_custom.html\"):\n",
     "        out, report = utils.custom_process(preprocess_output, processor_class=Analysis, schema=run.schema, client=client, preload=columns_to_preload)\n",
-    "    print(f\"preloaded columns: {len(columns_to_preload)}, {columns_to_preload} {\"etc.\" if len(columns_to_preload) > 4 else \"\"}\")\n",
+    "    print(f\"preloaded columns: {len(columns_to_preload)}, {columns_to_preload[:4]} {\"etc.\" if len(columns_to_preload) > 4 else \"\"}\")\n",
     "    print(f\"preloaded but unused columns: {len([c for c in columns_to_preload if c not in report[\"columns\"]])}\")\n",
     "    print(f\"used but not preloaded columns: {len([c for c in report[\"columns\"] if c not in columns_to_preload])}\")\n",
     "\n",

atlas/ntuple_production/distributed_xrdcp.ipynb

@@ -18,6 +18,7 @@
     "import numpy as np\n",
     "import matplotlib.dates as mdates\n",
     "import matplotlib.pyplot as plt\n",
+    "import tqdm.notebook\n",
     "\n",
     "client = Client(\"tls://localhost:8786\")"
    ]
@@ -65,7 +66,13 @@
     "\n",
     "t0 = time.time()\n",
     "tasks = [dask.delayed(run_xrdcp)(fname, size) for fname, size in zip(all_files, all_sizes_GB)]\n",
-    "res = dask.compute(*tasks)\n",
+    "futures = client.compute(tasks)\n",
+    "\n",
+    "with tqdm.notebook.tqdm(total=len(futures)) as pbar:\n",
+    "  for future in dask.distributed.as_completed(futures):\n",
+    "    pbar.update(1)\n",
+    "\n",
+    "res = [f.result() for f in futures]\n",
     "t1 = time.time()"
    ]
   },

atlas/utils.py

@@ -19,6 +19,7 @@
 import matplotlib.dates as mdates
 import matplotlib.pyplot as plt
 import numpy as np
+import tqdm.notebook
 import uproot
 
 
@@ -311,13 +312,16 @@ def extract_metadata(fname_and_treename: str, custom_func) -> dict:
                 meta.update({"custom_meta": custom_func(f)})
         return {fname: meta}
 
-    preprocess_input = dask.bag.from_sequence(files_to_preprocess, partition_size=1)
     print(f"pre-processing {len(files_to_preprocess)} file(s)")
-    futures = preprocess_input.map(functools.partial(extract_metadata, custom_func=custom_func))
-    result = client.compute(futures).result()
+    tasks = client.map(functools.partial(extract_metadata, custom_func=custom_func), files_to_preprocess)
+    futures = client.compute(tasks)
+
+    with tqdm.notebook.tqdm(total=len(futures)) as pbar:
+      for _ in dask.distributed.as_completed(futures):
+        pbar.update(1)
 
     # turn into dict for easier use
-    result_dict = {k: v for res in result for k, v in res.items()}
+    result_dict = {k: v for res in [f.result() for f in futures] for k, v in res.items()}
 
     # join back together per-file information with fileset-level information and turn into WorkItem list for coffea
     workitems = []
@@ -410,5 +414,13 @@ def sum_output(a, b):
         )
 
     workitems_bag = dask.bag.from_sequence(workitems, partition_size=1)
-    futures = workitems_bag.map(run_analysis).fold(sum_output)
-    return client.compute(futures).result()
+    tasks = workitems_bag.map(run_analysis).to_delayed()
+    futures = client.compute(tasks)
+    workitems_bag = dask.bag.from_delayed(futures)
+    res = client.compute(workitems_bag.fold(sum_output))
+
+    with tqdm.notebook.tqdm(total=len(futures)) as pbar:
+        for _ in dask.distributed.as_completed(futures):
+            pbar.update(1)
+
+    return res.result()