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

langfun/core/eval/v2/runners/parallel.py

@@ -0,0 +1,243 @@
+# Copyright 2025 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""Parallel runner."""
+
+import collections
+import math
+import random
+import threading
+import time
+
+from typing import Annotated, Iterator
+import langfun.core as lf
+from langfun.core.eval.v2 import checkpointing
+from langfun.core.eval.v2 import experiment as experiment_lib
+from langfun.core.eval.v2.runners import base
+from langfun.core.eval.v2.runners import ckpt_monitor
+import pyglove as pg
+
+
+class ParallelRunner(base.RunnerBase):
+  """A runner that executes evaluations and examples in parallel.
+
+  The parallel runner groups evaluations by their required resources
+  (e.g., specific LLMs) and runs evaluations that do not share resources in
+  parallel. Within each evaluation, examples are also processed in parallel
+  using threads, up to `Evaluation.max_workers`.
+  """
+
+  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[base.Evaluation]) -> None:
+    """Runs the evaluations in parallel."""
+    def _run_group(evaluation_group: list[base.Evaluation]):
+      for e in evaluation_group:
+        self.run_evaluation(e)
+
+    # Run evaluations in parallel groupped by resource key.
+    groups: dict[str, list[base.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: base.Evaluation, items: Iterator[base.Example]
+  ) -> None:
+    """Override run items to run in parallel."""
+    if self.concurrent_startup_delay is not None:
+      thread_delayed = {}
+      def _evaluate_item(item: base.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: base.Example):
+        return self.evaluate_item(evaluation, item)
+
+    for _, _, _ in lf.concurrent_map(
+        _evaluate_item,
+        items,
+        max_workers=self._max_workers(evaluation),
+        timeout=self.timeout,
+        silence_on_errors=None,
+    ):
+      pass
+
+  def _max_workers(self, evaluation: base.Evaluation) -> int | None:
+    return evaluation.max_workers
+
+
+class _SingleSliceRunner(ParallelRunner):
+  """A single slice runner."""
+
+  NAME = '__single_slice_runner__'
+
+  # Do not track progress in single slice runner.
+  progress_tracker = None
+
+  num_slices: Annotated[
+      int,
+      'The number of slices to run the evaluations in.'
+  ] = 1
+
+  def _max_workers(self, evaluation: base.Evaluation) -> int | None:
+    max_workers = super()._max_workers(evaluation)
+    if max_workers is None:
+      return None
+    return max(1, math.ceil(max_workers / self.num_slices))
+
+
+class MultiSliceParallelRunner(experiment_lib.Runner):
+  """A sliced parallel runner.
+
+  An evaluation is split into `num_slices` slices. Each MultiSliceParallelRunner
+  instance is responsible for evaluating a single slice. The instance with
+  `slice_id` 0 will also aggregate checkpoints from all slices.
+
+  Sliced parallel runner allows running multiple instances across different
+  machines/hosts parallelly. This can be utilize for scaling the evaluation
+  jobs to run on multiple machines, or running evaluations in a fault-tolerant
+  way by splitting each evaluation into multiple slices.
+  """
+
+  NAME = 'sliced-parallel'
+
+  slice_id: Annotated[
+      int,
+      (
+          'The slice ID of the runner. If 0, it will also run as the '
+          'aggregator for collecting results from other slices. '
+      )
+  ] = 0
+
+  num_slices: Annotated[
+      int,
+      'The number of slices to run the evaluations in parallel.'
+  ] = 1
+
+  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
+
+  ckpt_format: Annotated[
+      str,
+      'The file extension of the checkpoint files.'
+  ] = 'bagz'
+
+  max_aggregation_threads: Annotated[
+      int,
+      'The maximum number of threads to aggregate checkpoints.'
+  ] = 32
+
+  def _on_bound(self):
+    super()._on_bound()
+    if self.current_run.use_cache != 'no':
+      raise ValueError(
+          'Cache is not supported in MultiProcessParallelRunner. '
+          f'Encountered: {self.current_run.use_cache}'
+      )
+    monitor_plugins = []
+    worker_plugins = [
+        checkpointing.PerExampleCheckpointer(
+            checkpoint_filename=f'checkpoint.{self.ckpt_format}'
+        ),
+    ]
+    for plugin in self.plugins:
+      if isinstance(plugin, checkpointing.Checkpointer):
+        pg.logging.warning(
+            'Built-in checkpointing is enabled on MultiProcessParallelRunner. '
+            f'Ignoring checkpointer: {plugin!r}.'
+        )
+      elif plugin.is_per_example():
+        worker_plugins.append(pg.Ref(plugin))
+      else:
+        monitor_plugins.append(pg.Ref(plugin))
+
+    if self.slice_id == 0:
+      self._ckpt_monitor = ckpt_monitor.CheckpointMonitor(
+          pg.Ref(self.current_run),
+          plugins=monitor_plugins,
+          monitor_inprogress_files=True,
+          checkpoint_pattern=f'checkpoint_*.{self.ckpt_format}',
+          max_aggregation_threads=self.max_aggregation_threads,
+      )
+    else:
+      self._ckpt_monitor = None
+
+    self._slice_runner = _SingleSliceRunner(
+        current_run=self.current_run.clone(
+            override=dict(
+                # Clone the experiment to avoid updating the original one.
+                experiment=self.current_run.experiment.clone(),
+                example_ids=self._examples_to_evaluate,
+            )
+        ),
+        plugins=worker_plugins,
+        timeout=self.timeout,
+        concurrent_startup_delay=self.concurrent_startup_delay,
+    )
+
+  def _examples_to_evaluate(
+      self,
+      experiment: experiment_lib.Experiment
+  ) -> list[int]:
+    all_ids = self.current_run.examples_to_evaluate(experiment)
+    return [x for x in all_ids if x % self.num_slices == self.slice_id]
+
+  def run(self) -> None:
+    if self._ckpt_monitor is not None:
+      self._ckpt_monitor.start()
+    self._slice_runner.run()
+    if self._ckpt_monitor is not None:
+      self._ckpt_monitor.join()

langfun/core/eval/v2/runners/parallel_test.py

@@ -0,0 +1,182 @@
+# Copyright 2024 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""Tests for parallel runner."""
+import os
+import tempfile
+import threading
+from typing import Any
+import unittest
+
+from langfun.core.eval.v2 import checkpointing
+from langfun.core.eval.v2 import eval_test_helper
+from langfun.core.eval.v2 import reporting
+from langfun.core.eval.v2.runners import parallel  # pylint: disable=unused-import
+
+import pyglove as pg
+
+
+class ParallelRunnerTest(unittest.TestCase):
+
+  def assert_same_list(self, actual: list[Any], expected: list[Any]):
+    self.assertEqual(len(actual), len(expected))
+    for i, (x, y) in enumerate(zip(actual, expected)):
+      if x is not y:
+        print(i, pg.diff(x, y))
+      self.assertIs(x, y)
+
+  def test_parallel_runner(self):
+    plugin = eval_test_helper.TestPlugin()
+    exp = eval_test_helper.test_experiment()
+    root_dir = os.path.join(tempfile.mkdtemp(), 'test_parallel_runner')
+    _ = exp.run(root_dir, runner='parallel', plugins=[plugin])
+
+    self.assertIsNotNone(plugin.start_time)
+    self.assertIsNotNone(plugin.complete_time)
+    self.assertGreater(plugin.complete_time, plugin.start_time)
+
+    self.assertEqual(
+        len(plugin.started_experiments), len(exp.nodes)
+    )
+    self.assertEqual(
+        len(plugin.completed_experiments), len(exp.nodes)
+    )
+    self.assertEqual(
+        len(plugin.started_example_ids), 6 * 10
+    )
+    self.assertEqual(
+        len(plugin.completed_example_ids), 6 * 10
+    )
+    self.assert_same_list(plugin.skipped_experiments, [])
+
+    for node in exp.nodes:
+      self.assertTrue(node.progress.is_started)
+      self.assertTrue(node.progress.is_completed)
+      if node.is_leaf:
+        self.assertEqual(node.progress.num_skipped, 0)
+        self.assertEqual(node.progress.num_completed, 10)
+        self.assertEqual(node.progress.num_failed, 1)
+      else:
+        self.assertEqual(node.progress.num_skipped, 0)
+        self.assertEqual(node.progress.num_failed, 0)
+        self.assertEqual(node.progress.num_processed, node.progress.num_total)
+
+  def test_raise_if_has_error(self):
+    root_dir = os.path.join(tempfile.mkdtemp(), 'test_raise_if_has_error')
+    exp = eval_test_helper.TestEvaluation()
+    with self.assertRaisesRegex(ValueError, 'x should not be 5'):
+      exp.run(root_dir, runner='parallel', plugins=[], raise_if_has_error=True)
+
+  def test_concurrent_startup_delay(self):
+    plugin = eval_test_helper.TestPlugin()
+    exp = eval_test_helper.test_experiment()
+    root_dir = os.path.join(
+        tempfile.mkdtemp(), 'test_concurrent_startup_delay'
+    )
+    _ = exp.run(
+        root_dir,
+        runner='parallel',
+        plugins=[plugin],
+        concurrent_startup_delay=(0, 5),
+    )
+
+
+class MultiProcessParallelRunnerTest(unittest.TestCase):
+
+  def assert_same_list(self, actual: list[Any], expected: list[Any]):
+    self.assertEqual(len(actual), len(expected))
+    for i, (x, y) in enumerate(zip(actual, expected)):
+      if x is not y:
+        print(i, pg.diff(x, y))
+      self.assertIs(x, y)
+
+  def test_basic(self):
+    plugin = eval_test_helper.TestPlugin()
+    exp = eval_test_helper.test_experiment()
+    root_dir = os.path.join(tempfile.mkdtemp(), 'test_parallel_runner')
+    runs = [None, None]
+
+    def run_slice(slice_id: int):
+      runs[slice_id] = exp.run(
+          root_dir,
+          runner='sliced-parallel',
+          plugins=[
+              pg.Ref(plugin),
+              # Ignored as PerExampleCheckpointer is used.
+              checkpointing.BulkCheckpointer(),
+              reporting.ExampleHtmlGenerator(),
+          ],
+          use_cache='no',
+          slice_id=slice_id,
+          num_slices=2,
+          ckpt_format='jsonl',
+      )
+
+    # We simulate two slices running in parallel.
+    threads = [threading.Thread(target=run_slice, args=(i,)) for i in range(2)]
+    for t in threads:
+      t.start()
+    for t in threads:
+      t.join()
+
+    self.assertIsNotNone(plugin.start_time)
+    self.assertIsNotNone(plugin.complete_time)
+    self.assertGreater(plugin.complete_time, plugin.start_time)
+
+    self.assertEqual(
+        len(plugin.started_experiments), len(exp.nodes)
+    )
+    self.assertEqual(
+        len(plugin.completed_experiments), len(exp.nodes)
+    )
+    self.assertEqual(
+        len(plugin.started_example_ids), 6 * 10
+    )
+    self.assertEqual(
+        len(plugin.completed_example_ids), 6 * 10
+    )
+    self.assert_same_list(plugin.skipped_experiments, [])
+
+    for node in exp.nodes:
+      self.assertTrue(node.progress.is_started)
+      self.assertTrue(node.progress.is_completed)
+      if node.is_leaf:
+        self.assertEqual(node.progress.num_skipped, 0)
+        self.assertEqual(node.progress.num_completed, 10)
+        self.assertEqual(node.progress.num_failed, 1)
+        for example_id in runs[0].examples_to_evaluate(node):
+          self.assertTrue(
+              pg.io.path_exists(
+                  runs[0].output_path_for(node, f'{example_id}.html')
+              )
+          )
+      else:
+        self.assertEqual(node.progress.num_skipped, 0)
+        self.assertEqual(node.progress.num_failed, 0)
+        self.assertEqual(node.progress.num_processed, node.progress.num_total)
+
+  def test_parallel_mp_runner_does_not_support_cache(self):
+    exp = eval_test_helper.test_experiment()
+    root_dir = os.path.join(tempfile.mkdtemp(), 'test_parallel_mp_runner_cache')
+    with self.assertRaisesRegex(ValueError, 'Cache is not supported'):
+      exp.run(
+          root_dir,
+          runner='sliced-parallel',
+          use_cache='global',
+          slice_id=0,
+          num_slices=1,
+      )
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/eval/v2/runners/sequential.py

@@ -0,0 +1,47 @@
+# Copyright 2024 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""Sequential runner."""
+
+from typing import Any, Callable, Iterator
+from langfun.core.eval.v2.runners import base
+
+
+class SequentialRunner(base.RunnerBase):
+  """A runner that executes evaluations and examples sequentially.
+
+  The sequential runner executes all evaluations and their examples in the
+  calling thread. Background tasks are also run sequentially, which makes it
+  easier to debug as exceptions from background tasks will be raised
+  immediately.
+  """
+
+  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[base.Evaluation]) -> None:
+    """Runs the experiment in sequence."""
+    for e in evaluations:
+      self.run_evaluation(e)
+
+  def _evaluate_items(
+      self, evaluation: base.Evaluation, items: Iterator[base.Example]
+  ) -> None:
+    """Runs the evaluation items in sequence."""
+    for item in items:
+      self.evaluate_item(evaluation, item)

langfun/core/eval/v2/runners/sequential_test.py

@@ -0,0 +1,169 @@
+# Copyright 2024 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""Tests for sequential runner."""
+import os
+import tempfile
+from typing import Any
+import unittest
+
+from langfun.core.eval.v2 import eval_test_helper
+from langfun.core.eval.v2.runners import sequential  # pylint: disable=unused-import
+
+import pyglove as pg
+
+
+class SequentialRunnerTest(unittest.TestCase):
+
+  def assert_same_list(self, actual: list[Any], expected: list[Any]):
+    self.assertEqual(len(actual), len(expected))
+    for i, (x, y) in enumerate(zip(actual, expected)):
+      if x is not y:
+        print(i, pg.diff(x, y))
+      self.assertIs(x, y)
+
+  def test_basic(self):
+    plugin = eval_test_helper.TestPlugin()
+    exp = eval_test_helper.test_experiment()
+    root_dir = os.path.join(tempfile.mkdtemp(), 'test_sequential_runner')
+    _ = exp.run(root_dir, runner='sequential', plugins=[plugin])
+
+    self.assertIsNotNone(plugin.start_time)
+    self.assertIsNotNone(plugin.complete_time)
+    self.assertGreater(plugin.complete_time, plugin.start_time)
+
+    self.assert_same_list(
+        plugin.started_experiments,
+        exp.nonleaf_nodes + exp.leaf_nodes
+    )
+    self.assert_same_list(
+        plugin.completed_experiments,
+        exp.leaf_nodes + list(reversed(exp.nonleaf_nodes))
+    )
+    self.assert_same_list(
+        plugin.started_example_ids, list(range(1, 11)) * 6
+    )
+    self.assert_same_list(
+        plugin.completed_example_ids, list(range(1, 11)) * 6
+    )
+    self.assert_same_list(plugin.skipped_experiments, [])
+
+    for node in exp.nodes:
+      self.assertTrue(node.progress.is_started)
+      self.assertTrue(node.progress.is_completed)
+      if node.is_leaf:
+        self.assertEqual(node.progress.num_skipped, 0)
+        self.assertEqual(node.progress.num_completed, 10)
+        self.assertEqual(node.progress.num_failed, 1)
+      else:
+        self.assertEqual(node.progress.num_skipped, 0)
+        self.assertEqual(node.progress.num_failed, 0)
+        self.assertEqual(node.progress.num_processed, node.progress.num_total)
+
+  def test_raise_if_has_error(self):
+    root_dir = os.path.join(tempfile.mkdtemp(), 'test_raise_if_has_error')
+    exp = eval_test_helper.TestEvaluation()
+    with self.assertRaisesRegex(ValueError, 'x should not be 5'):
+      exp.run(
+          root_dir, runner='sequential', plugins=[], raise_if_has_error=True
+      )
+
+  def test_example_ids(self):
+    root_dir = os.path.join(tempfile.mkdtemp(), 'test_example_ids')
+    exp = eval_test_helper.test_experiment()
+    plugin = eval_test_helper.TestPlugin()
+    _ = exp.run(
+        root_dir, runner='sequential', plugins=[plugin], example_ids=[5, 7, 9]
+    )
+    self.assertEqual(plugin.started_example_ids, [5, 7, 9] * 6)
+    self.assertEqual(plugin.completed_example_ids, [5, 7, 9] * 6)
+
+  def test_shuffle_inputs(self):
+    root_dir = os.path.join(tempfile.mkdtemp(), 'test_shuffle_inputs')
+    exp = eval_test_helper.test_experiment()
+    plugin = eval_test_helper.TestPlugin()
+    run = exp.run(
+        root_dir, runner='sequential', plugins=[plugin], shuffle_inputs=True
+    )
+    self.assertTrue(run.shuffle_inputs)
+
+  def test_filter(self):
+    plugin = eval_test_helper.TestPlugin()
+    exp = eval_test_helper.test_experiment()
+    root_dir = os.path.join(tempfile.mkdtemp(), 'test_filter')
+
+    _ = exp.run(
+        root_dir, runner='sequential', plugins=[plugin],
+        filter=lambda e: e.lm.offset != 0
+    )
+    self.assert_same_list(
+        plugin.started_experiments,
+        exp.nonleaf_nodes + exp.leaf_nodes[2:]
+    )
+    self.assert_same_list(
+        plugin.skipped_experiments, exp.leaf_nodes[:2]
+    )
+    self.assert_same_list(
+        plugin.completed_experiments,
+        exp.leaf_nodes[2:] + [exp.children[1], exp]
+    )
+
+  def test_use_cache(self):
+    @pg.functor()
+    def test_inputs(num_examples: int = 10):
+      return [
+          pg.Dict(
+              x=i // 2, y=(i // 2) ** 2,
+              groundtruth=(i // 2 + (i // 2) ** 2)
+          ) for i in range(num_examples)
+      ]
+
+    exp = eval_test_helper.TestEvaluation(
+        inputs=test_inputs(num_examples=pg.oneof([2, 4]))
+    )
+    # Global cache.
+    root_dir = os.path.join(tempfile.mkdtemp(), 'global_cache')
+    run = exp.run(
+        root_dir, 'new', runner='sequential', use_cache='global', plugins=[]
+    )
+    self.assertTrue(pg.io.path_exists(run.output_path_for(exp, 'cache.json')))
+    self.assertEqual(exp.usage_summary.cached.total.num_requests, 4)
+    self.assertEqual(exp.usage_summary.uncached.total.num_requests, 2)
+
+    # Per-dataset cache.
+    root_dir = os.path.join(tempfile.mkdtemp(), 'per_dataset')
+    run = exp.run(
+        root_dir, 'new', runner='sequential',
+        use_cache='per_dataset', plugins=[]
+    )
+    for leaf in exp.leaf_nodes:
+      self.assertTrue(
+          pg.io.path_exists(run.output_path_for(leaf, 'cache.json'))
+      )
+    self.assertEqual(exp.usage_summary.cached.total.num_requests, 3)
+    self.assertEqual(exp.usage_summary.uncached.total.num_requests, 3)
+
+    # No cache.
+    root_dir = os.path.join(tempfile.mkdtemp(), 'no')
+    run = exp.run(root_dir, runner='sequential', use_cache='no', plugins=[])
+    self.assertFalse(pg.io.path_exists(run.output_path_for(exp, 'cache.json')))
+    for leaf in exp.leaf_nodes:
+      self.assertFalse(
+          pg.io.path_exists(run.output_path_for(leaf, 'cache.json'))
+      )
+    self.assertEqual(exp.usage_summary.cached.total.num_requests, 0)
+    self.assertEqual(exp.usage_summary.uncached.total.num_requests, 6)
+
+
+if __name__ == '__main__':
+  unittest.main()