Eval runner to wait for background IO before finish.

Before this CL, `lf.eval.v2.Experiment.run` can finish before background IO completes, which could lead to saved file corruption due to premature process exit. PiperOrigin-RevId: 702756055
google · Dec 4, 2024 · fca081b · fca081b
1 parent 1220ac6
commit fca081b
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/langfun/core/eval/v2/runners.py b/langfun/core/eval/v2/runners.py
@@ -261,6 +261,9 @@ def run(self) -> None:
       if cache is not None:
         self.background_run(cache.save)
 
+      # Wait for the background tasks to finish.
+      self._io_pool.shutdown(wait=True)
+
   @abc.abstractmethod
   def _run(self, evaluations: list[Evaluation]) -> None:
     """Runs multiple evaluations."""