Repository: sbi-dev/sbi

Issue:
fix progress bar and simulation batch size in simulate_in_batches
Two problems related to `simulate_in_batches`: 

1) The progress bar fills up by sending jobs to the workers, i.e., before these are actually done. E.g., when running 100 parameters with batch size 10 and 10 workers, the progress bar is done instantly, but simulations may continue for minutes or hours. We should find a way to fill up the pbar only when the jobs are done. 

2) The simulation batch size defaults to 1. This can make things very slow when one uses a simulator that cannot simulate multiple parameters at once (vectorised) and one use multiple workers. In this scenario each worker deals with a single simulation -- creating a lot of overhead. It would be much better (for me it gave a 10x speed up) to write a simple "batch"-simulator with a for loop and set the simulation batch size to 10 or 100 so that each worker call will run at least 10 or 100 simulations before going back. Actually, it could make sense to just set 
```
sim_batch_size = int(num_simulations / num_workers)
``` 
in order to minimise overhead. 
The only downside would be that the pbar becomes less informative. So one could think about reducing the batch size a bit to obtain more "steps" in the pbar. 

Any thoughts? 


Discussion:

Resolving patch:
diff --git a/sbi/simulators/simutils.py b/sbi/simulators/simutils.py
--- a/sbi/simulators/simutils.py
+++ b/sbi/simulators/simutils.py
@@ -8,6 +8,8 @@
 from torch import Tensor
 from tqdm.auto import tqdm
 from joblib import Parallel, delayed
+import contextlib
+import joblib
 
 
 def simulate_in_batches(
@@ -46,19 +48,19 @@ def simulate_in_batches(
 
         if num_workers > 1:
             # Parallelize the sequence of batches across workers.
-            # TODO: This usage of tqdm tracks the dispatching of jobs instead of the
-            # moment when they are done, resulting in waiting time at 100% in case the
-            # last jobs takes long. A potential solution can be found here: https://
-            # stackoverflow.com/a/61689175
-            simulation_outputs = Parallel(n_jobs=num_workers)(
-                delayed(simulator)(batch)
-                for batch in tqdm(
+            # We use the solution proposed here: https://stackoverflow.com/a/61689175
+            # to update the pbar only after the workers finished a task.
+            with tqdm_joblib(
+                tqdm(
                     batches,
                     disable=not show_progress_bars,
                     desc=f"Running {num_sims} simulations in {len(batches)} batches.",
                     total=len(batches),
                 )
-            )
+            ) as progress_bar:
+                simulation_outputs = Parallel(n_jobs=num_workers)(
+                    delayed(simulator)(batch) for batch in batches
+                )
         else:
             pbar = tqdm(
                 total=num_sims,
@@ -77,3 +79,28 @@ def simulate_in_batches(
         x = simulator(theta)
 
     return x
+
+
+@contextlib.contextmanager
+def tqdm_joblib(tqdm_object):
+    """Context manager to patch joblib to report into tqdm progress bar given as
+    argument
+
+    This wrapped context manager obtains the number of finished tasks from the tqdm
+    print function and uses it to update the pbar, as suggested in
+    https://stackoverflow.com/a/61689175. See #419, #421
+    """
+
+    def tqdm_print_progress(self):
+        if self.n_completed_tasks > tqdm_object.n:
+            n_completed = self.n_completed_tasks - tqdm_object.n
+            tqdm_object.update(n=n_completed)
+
+    original_print_progress = joblib.parallel.Parallel.print_progress
+    joblib.parallel.Parallel.print_progress = tqdm_print_progress
+
+    try:
+        yield tqdm_object
+    finally:
+        joblib.parallel.Parallel.print_progress = original_print_progress
+        tqdm_object.close()
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -24,7 +24,7 @@
 REQUIRES_PYTHON = ">=3.6.0"
 
 REQUIRED = [
-    "joblib",
+    "joblib>=1.0.0",
     "matplotlib",
     "numpy",
     "pillow",


Test patch: