langfun 0.1.2.dev202509120804__py3-none-any.whl → 0.1.2.dev202512150805__py3-none-any.whl

langfun/core/eval/v2/{runners.py → runners/base.py}

@@ -11,18 +11,18 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Evaluation experiment runners."""
+"""Base experiment runner."""
+
 import abc
-import collections
 import concurrent.futures
 import random
 import threading
-import time
 import traceback
-from typing import Any, Annotated, Callable, Iterator
+from typing import Any, Annotated, Callable, Iterator, Literal
 
 from langfun import core as lf
 from langfun.core.eval.v2 import checkpointing
+from langfun.core.eval.v2 import config_saver
 from langfun.core.eval.v2 import evaluation as evaluation_lib
 from langfun.core.eval.v2 import example as example_lib
 from langfun.core.eval.v2 import experiment as experiment_lib
@@ -38,35 +38,57 @@ Experiment = experiment_lib.Experiment
 Plugin = experiment_lib.Plugin
 
 
-_RUN_MANIFEST = 'run.json'
-
-
 class RunnerBase(Runner):
-  """A simple runner that runs evaluations and their examples sequentially."""
+  """Base class for runners with plugin support and IO pooling.
+
+  `RunnerBase` provides the basic runner functionalities such as plugin
+  integration for checkpointing, reporting and progress tracking.
+  It also manages a thread pool for background IO operations.
+  Subclasses should implement `_run` and `_evaluate_items` for different
+  execution strategies.
+  """
 
-  tqdm: Annotated[
-      bool,
+  progress_tracker: Annotated[
+      Literal['tqdm', 'html', 'auto', None],
       (
-          'If True, force using tqdm for progress update. Otherwise, determine '
-          'it automatically based on the running environment (console vs. '
-          'notebook)'
+          'If `tqdm`, force using tqdm for progress update. '
+          'If `html`, force using html for progress update. '
+          'If `auto`, determine it automatically based on the running '
+          'environment (console vs. notebook)'
+          'If `none`, disable progress update.'
       )
-  ] = False
+  ] = 'auto'
 
   plugins = [
       checkpointing.BulkCheckpointer(),
       reporting.HtmlReporter(),
+      config_saver.RunConfigSaver(),
   ]
 
+  max_background_threads: Annotated[
+      int,
+      'Max number of background threads for IO operations.'
+  ] = 128
+
   def _on_bound(self):
     super()._on_bound()
 
     # Install the tqdm plugin if needed.
-    with pg.notify_on_change(False):
-      self.plugins.append(progress_tracking.progress_tracker(self.tqdm))
+    if self.progress_tracker is not None:
+      with pg.notify_on_change(False):
+        self.plugins.append(
+            progress_tracking.progress_tracker(self.progress_tracker)
+        )
+
+    if self.max_background_threads > 0:
+      self._io_pool_lock = threading.Lock()
+      self._io_pool = concurrent.futures.ThreadPoolExecutor(
+          max_workers=self.max_background_threads
+      )
+    else:
+      self._io_pool_lock = None
+      self._io_pool = None
 
-    self._io_pool_lock = threading.Lock()
-    self._io_pool = concurrent.futures.ThreadPoolExecutor(max_workers=16)
     # TODO(daiyip): render background errors.
     self._background_last_error = None
 
@@ -78,9 +100,12 @@ class RunnerBase(Runner):
       except Exception as e:  # pylint: disable=broad-except
         self._background_last_error = e
 
-    with self._io_pool_lock:
-      if self._io_pool is not None:
-        self._io_pool.submit(_background_run, *args, **kwargs)
+    if self.max_background_threads > 0:
+      with self._io_pool_lock:
+        if self._io_pool is not None:
+          self._io_pool.submit(_background_run, *args, **kwargs)
+    else:
+      _background_run(*args, **kwargs)
 
   def _all_plugins(self, experiment: Experiment) -> Iterator[Plugin]:
     """Returns all plugins for the experiment."""
@@ -89,24 +114,8 @@ class RunnerBase(Runner):
     for plugin in experiment.plugins:
       yield plugin
 
-  #
-  # IO operations for saving running files.
-  #
-
-  def _save_run_manifest(self) -> None:
-    def _save():
-      pg.symbolic.deref(self.current_run.clone(), recursive=True).save(
-          self.current_run.output_path_for(
-              self.current_run.experiment, _RUN_MANIFEST
-          ),
-          hide_default_values=True
-      )
-    self.background_run(_save)
-
   def on_run_start(self) -> None:
     """Called when a runner is started."""
-    self._save_run_manifest()
-
     for plugin in self._all_plugins(self.current_run.experiment):
       plugin.on_run_start(self, self.current_run.experiment)
 
@@ -126,11 +135,11 @@ class RunnerBase(Runner):
     num_examples_to_evaluate = 0
     if experiment.is_leaf:
       assert isinstance(experiment, Evaluation)
-      num_examples_to_evaluate = (
-          len(self.current_run.example_ids)
-          if self.current_run.example_ids else experiment.num_examples
+      num_examples_to_evaluate = len(
+          self.current_run.examples_to_evaluate(experiment)
       )
       experiment.progress.start(total=num_examples_to_evaluate)
+      pg.io.mkdirs(self.current_run.output_dir(experiment))
     else:
       experiment.progress.start(total=len(experiment.leaf_nodes))
 
@@ -139,6 +148,7 @@ class RunnerBase(Runner):
       plugin.on_experiment_start(self, experiment)
 
     if experiment.is_leaf:
+      pg.io.mkdirs(self.current_run.output_dir(experiment))
       experiment.info(
           f'Starting evaluation {experiment.id!r} with '
           f'{num_examples_to_evaluate} examples to evaluate.'
@@ -180,10 +190,7 @@ class RunnerBase(Runner):
       self._log_experiment_completion(experiment)
 
   def _log_experiment_completion(self, experiment: Experiment):
-    example_ids = (
-        self.current_run.example_ids if self.current_run.example_ids else
-        list(range(1, experiment.num_examples + 1))
-    )
+    example_ids = sorted(self.current_run.examples_to_evaluate(experiment))
     num_from_checkpoint, num_processed = 0, 0
     for example_id in example_ids:
       status = experiment.state.get_status(example_id)
@@ -220,7 +227,7 @@ class RunnerBase(Runner):
       else:
         # A evaluation could be considered as done if it has processed all the
         # examples specified by `example_ids`.
-        assert progress.is_completed
+        assert progress.is_completed, progress
         parent_progress.increment_processed()
 
       if parent_progress.is_completed:
@@ -235,6 +242,8 @@ class RunnerBase(Runner):
       example: Example
   ) -> None:
     """Called when an evaluation example is started."""
+    assert isinstance(experiment, Evaluation), experiment
+    experiment.state.update(example, in_progress=True)
     for plugin in self._all_plugins(experiment):
       plugin.on_example_start(self, experiment, example)
     experiment.info(f'Starting to evaluate example {example.id}.')
@@ -245,6 +254,8 @@ class RunnerBase(Runner):
       example: Example
   ) -> None:
     """Called when an evaluation example is complete."""
+    assert isinstance(experiment, Evaluation), experiment
+    experiment.state.update(example, in_progress=False)
     if example.newly_processed:
       if example.error is None:
         experiment.progress.increment_processed()
@@ -256,7 +267,7 @@ class RunnerBase(Runner):
         experiment.progress.increment_failed()
         experiment.error(
             (
-                f'Failed to evaluate example {example.id} in'
+                f'Failed to evaluate example {example.id} in '
                 f'{example.elapse:.2f} seconds.'
             ),
             error=example.error
@@ -316,7 +327,7 @@ class RunnerBase(Runner):
         self._run(targets)
 
       self.on_run_complete()
-    except Exception as e:  # pylint: disable=broad-except
+    except BaseException as e:  # pylint: disable=broad-except
       self.on_run_abort(e)
       raise e
     finally:
@@ -324,9 +335,10 @@ class RunnerBase(Runner):
         self.background_run(cache.save)
 
       # Wait for the background tasks to finish.
-      with self._io_pool_lock:
-        self._io_pool, io_pool = None, self._io_pool
-      io_pool.shutdown(wait=True)
+      if self.max_background_threads > 0:
+        with self._io_pool_lock:
+          self._io_pool, io_pool = None, self._io_pool
+        io_pool.shutdown(wait=True)
 
   @abc.abstractmethod
   def _run(self, evaluations: list[Evaluation]) -> None:
@@ -335,6 +347,7 @@ class RunnerBase(Runner):
   def run_evaluation(self, evaluation: Evaluation) -> None:
     """Runs the evaluation."""
     try:
+      evaluation.setup()
       self.on_experiment_start(evaluation)
 
       per_evaluation_settings = {}
@@ -344,18 +357,14 @@ class RunnerBase(Runner):
         per_evaluation_settings['cache'] = cache
 
       with lf.use_settings(**per_evaluation_settings):
-        if self.current_run.example_ids is None:
-          items = (
-              Example(id=i + 1, input=ex) for i, ex in enumerate(
-                  evaluation.example_inputs)
-          )
-        else:
-          items = (
-              Example(
-                  id=example_id,
-                  input=evaluation.example_input_by_id(example_id)
-              ) for example_id in self.current_run.example_ids
-          )
+        items = (
+            Example(
+                id=example_id,
+                input=evaluation.example_input_by_id(example_id)
+            ) for example_id in sorted(
+                self.current_run.examples_to_evaluate(evaluation)
+            )
+        )
         if self.current_run.shuffle_inputs:
           items = list(items)
           random.shuffle(items)
@@ -367,6 +376,8 @@ class RunnerBase(Runner):
     except BaseException as e:  # pylint: disable=broad-except
       self.on_experiment_abort(evaluation, e)
       raise e
+    finally:
+      evaluation.teardown()
 
   @abc.abstractmethod
   def _evaluate_items(
@@ -394,121 +405,3 @@ class RunnerBase(Runner):
     return in_memory.InMemory(
         self.current_run.output_path_for(experiment, 'cache.json')
     )
-
-
-class SequentialRunner(RunnerBase):
-  """Sequential runner.
-
-  Sequential runner runs all evaluations and their examples in sequence,
-  as well as the background tasks, it allows the developer to catch all
-  exceptions thrown from the background tasks, making it easier to debug.
-  """
-
-  NAME = 'sequential'
-
-  def background_run(
-      self, func: Callable[..., Any], *args: Any, **kwargs: Any
-  ) -> None:
-    """Runs the function with the IO pool."""
-    func(*args, **kwargs)
-
-  def _run(self, evaluations: list[Evaluation]) -> None:
-    """Runs the experiment in sequence."""
-    for e in evaluations:
-      self.run_evaluation(e)
-
-  def _evaluate_items(
-      self, evaluation: Evaluation, items: Iterator[Example]
-  ) -> None:
-    """Runs the evaluation items in sequence."""
-    for item in items:
-      self.evaluate_item(evaluation, item)
-
-
-class DebugRunner(SequentialRunner):
-  """Debug runner."""
-
-  NAME = 'debug'
-
-  # Do not use the checkpointer for debug runner.
-  plugins = []
-
-  def _on_bound(self):
-    super()._on_bound()
-    if self.current_run.example_ids is None:
-      self.current_run.rebind(example_ids=[1], skip_notification=True)
-    self.current_run.rebind(raise_if_has_error=True, skip_notification=True)
-
-  def _save_run_manifest(self) -> None:
-    """Do nothing to avoid overriden existing runs."""
-
-
-class ParallelRunner(RunnerBase):
-  """Parallel runner."""
-
-  NAME = 'parallel'
-
-  timeout: Annotated[
-      int | None,
-      'Timeout for each evaluation example.'
-  ] = None
-
-  concurrent_startup_delay: Annotated[
-      tuple[int, int] | None,
-      (
-          'A range of seconds to delay the initial evaluation of each thread '
-          'in the thread pool, helping to prevent a burst in LLM QPS at '
-          'startup. If set to None, no delay will be applied.'
-      )
-  ] = None
-
-  def _run(self, evaluations: list[Evaluation]) -> None:
-    """Runs the evaluations in parallel."""
-    def _run_group(evaluation_group: list[Evaluation]):
-      for e in evaluation_group:
-        self.run_evaluation(e)
-
-    # Run evaluations in parallel groupped by resource key.
-    groups: dict[str, list[Evaluation]] = collections.defaultdict(list)
-    for e in evaluations:
-      resource_ids = e.resource_ids()
-      if not resource_ids:
-        group_id = e.id
-      else:
-        # TODO(daiyip): support group that requires multiple resources.
-        group_id = resource_ids.pop()
-      groups[group_id].append(e)
-
-    for _, _, _ in lf.concurrent_map(
-        _run_group,
-        groups.values(),
-        max_workers=max(64, len(groups)),
-        timeout=self.timeout,
-        silence_on_errors=None,
-    ):
-      pass
-
-  def _evaluate_items(
-      self, evaluation: Evaluation, items: Iterator[Example]
-  ) -> None:
-    """Override run items to run in parallel."""
-    if self.concurrent_startup_delay is not None:
-      thread_delayed = {}
-      def _evaluate_item(item: Example):
-        thread_id = threading.current_thread().ident
-        if thread_id not in thread_delayed:
-          thread_delayed[thread_id] = True
-          time.sleep(random.randint(*self.concurrent_startup_delay))
-        return self.evaluate_item(evaluation, item)
-    else:
-      def _evaluate_item(item: Example):
-        return self.evaluate_item(evaluation, item)
-
-    for _, _, _ in lf.concurrent_map(
-        _evaluate_item,
-        items,
-        max_workers=evaluation.max_workers,
-        timeout=self.timeout,
-        silence_on_errors=None,
-    ):
-      pass