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

langfun/core/eval/v2/checkpointing.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 """Checkpointing evaluation runs."""
 import abc
+import datetime
+import os
 import re
 import threading
 import traceback
@@ -29,12 +31,32 @@ Runner = experiment_lib.Runner
 
 
 class Checkpointer(experiment_lib.Plugin):
-  """Base class for checkpointing evaluation examples."""
+  """Base class for checkpointing evaluation examples.
+
+  `Checkpointer` is a plugin that saves the state of processed examples
+  incrementally during an experiment run, allowing the experiment to be resumed
+  later. When an experiment starts, the checkpointer loads any previously saved
+  examples from an earlier run (or a warm-start run) into `experiment.state`,
+  so the runner can skip processing them again.
+  Subclasses should implement `_list_checkpoint_files` to identify
+  checkpoint files to load, and `_save_example` to save a newly processed
+  example.
+  """
 
   checkpoint_filename: Annotated[
       str,
       'Checkpoint file pattern.'
-  ] = 'checkpoint.bagz'
+  ] = 'checkpoint.jsonl'
+
+  enable_inprogress_file: Annotated[
+      bool,
+      'If True, write file "<example_id>.inprogress" when example gets started.'
+  ] = True
+
+  max_ckpt_loading_threads: Annotated[
+      int,
+      'Max number of workers for loading checkpoint files at startup.'
+  ] = 128
 
   def on_experiment_start(
       self,
@@ -75,6 +97,24 @@ class Checkpointer(experiment_lib.Plugin):
           f'scratch. Example IDs: {example_ids_to_evaluate}.'
       )
 
+  def on_example_start(
+      self,
+      runner: Runner,
+      experiment: Experiment,
+      example: Example,
+  ) -> None:
+    """Saves the example to the checkpoint file."""
+    if self.enable_inprogress_file:
+      def _save_inprogress_file(example: Example):
+        inprogress_file = runner.current_run.output_path_for(
+            experiment, f'{example.id}.inprogress'
+        )
+        pg.io.writefile(
+            inprogress_file,
+            datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+        )
+      runner.background_run(_save_inprogress_file, example)
+
   def on_example_complete(
       self,
       runner: Runner,
@@ -91,7 +131,7 @@ class Checkpointer(experiment_lib.Plugin):
       experiment: Experiment,
   ) -> None:
     """Creates the checkpoint file."""
-    ckpt_files = self._list_checkpoint_filenames(runner, experiment)
+    ckpt_files = self._list_checkpoint_files(runner, experiment)
     experiment.info(f'Found {len(ckpt_files)} checkpoint files to load.')
 
     # Load the checkpoint files in parallel.
@@ -101,18 +141,18 @@ class Checkpointer(experiment_lib.Plugin):
         experiment
     )
     context = dict(counter=0, counter_lock=threading.Lock())
-    copy_ckpt = current_run.input_root != current_run.output_root
 
     def _load_state(ckpt_file):
       error = None
       with pg.timeit() as t:
         try:
-          experiment.load_state(
-              current_run.input_path_for(experiment, ckpt_file),
+          loaded_examples = experiment.load_state(
+              ckpt_file,
               filter=lambda x: x.id in examples_to_load,
               load_example_metadata=lambda x: x.id in examples_to_load_metadata,
           )
         except BaseException as e:  # pylint: disable=broad-except
+          loaded_examples = []
           error = e
         finally:
           with context['counter_lock']:
@@ -130,34 +170,33 @@ class Checkpointer(experiment_lib.Plugin):
                 f'Skipping the file. ({progress_str})'
             )
 
-        if not copy_ckpt:
-          return
-
-        # Copy the checkpoint records to the output directory.
-        try:
-          with pg.io.open_sequence(
-              current_run.output_path_for(experiment, ckpt_file), 'w'
-          ) as o, pg.io.open_sequence(
-              current_run.input_path_for(experiment, ckpt_file), 'r'
-          ) as i:
-            for x in i:
-              o.add(x)
-        except BaseException as e:  # pylint: disable=broad-except
-          experiment.warning(
-              f'Failed to copy checkpoint {ckpt_file!r}: {e}.'
-          )
+        output_ckpt_file = current_run.output_path_for(
+            experiment, os.path.basename(ckpt_file)
+        )
+        if ckpt_file != output_ckpt_file and any(
+            e for e in loaded_examples if not e.has_error
+        ):
+          # Write the error-free warm-start examples to the output checkpoint
+          # file.
+          with SequenceWriter(output_ckpt_file) as writer:
+            for example in loaded_examples:
+              if not example.has_error:
+                writer.add(example)
 
     _ = list(
         lf.concurrent_map(
-            _load_state, ckpt_files, max_workers=16, silence_on_errors=None
+            _load_state,
+            ckpt_files,
+            max_workers=self.max_ckpt_loading_threads,
+            silence_on_errors=None
         )
     )
 
   @abc.abstractmethod
-  def _list_checkpoint_filenames(
+  def _list_checkpoint_files(
       self, runner: Runner, experiment: Experiment
   ) -> list[str]:
-    """Lists the checkpoint filenames to restore."""
+    """Lists the checkpoint file paths to restore."""
 
   @abc.abstractmethod
   def _save_example(
@@ -170,7 +209,12 @@ class Checkpointer(experiment_lib.Plugin):
 
 
 class PerExampleCheckpointer(Checkpointer):
-  """Checkpointer that saves each example to a separate file."""
+  """Checkpointer that saves each example to a separate file.
+
+  This checkpointer saves each processed example to its own checkpoint file,
+  named using the pattern `<checkpoint_filename_prefix>_<example_id>.<ext>`.
+  For example, `checkpoint_1.bagz`, `checkpoint_2.bagz`, etc.
+  """
 
   def _on_bound(self):
     super()._on_bound()
@@ -178,22 +222,41 @@ class PerExampleCheckpointer(Checkpointer):
     self._checkpoint_file_prefix = prefix
     self._checkpoint_file_ext = ext
 
-  def _list_checkpoint_filenames(
+  def _list_checkpoint_files(
       self, runner: Runner, experiment: Experiment
   ) -> list[str]:
-    experiment_dir = runner.current_run.input_dir(experiment)
-    filenames = []
+
+    def _list_checkpoints_from(ckpt_dir: str, examples_to_load: set[int]):
+      ckpt_files = []
+      if pg.io.path_exists(ckpt_dir):
+        regex = re.compile(
+            f'{self._checkpoint_file_prefix}_(\\d+){self._checkpoint_file_ext}'
+            .replace('.', '\\.')
+        )
+        for filename in pg.io.listdir(ckpt_dir):
+          match = regex.match(filename)
+          if match and int(match.group(1)) in examples_to_load:
+            examples_to_load.remove(int(match.group(1)))
+            ckpt_files.append(os.path.join(ckpt_dir, filename))
+      return ckpt_files
+
     examples_to_load = runner.current_run.examples_to_load(experiment)
-    if pg.io.path_exists(experiment_dir):
-      regex = re.compile(
-          f'{self._checkpoint_file_prefix}_(\\d+){self._checkpoint_file_ext}'
-          .replace('.', '\\.')
+
+    # Take output directory as the first priority to checkpoints processed in
+    # this run.
+    ckpt_files = _list_checkpoints_from(
+        runner.current_run.output_dir(experiment), examples_to_load
+    )
+    # If the input and output directories are different, also load from the
+    # input directory.
+    if (examples_to_load
+        and runner.current_run.input_root != runner.current_run.output_root):
+      ckpt_files.extend(
+          _list_checkpoints_from(
+              runner.current_run.input_dir(experiment), examples_to_load
+          )
       )
-      for filename in pg.io.listdir(experiment_dir):
-        match = regex.match(filename)
-        if match and int(match.group(1)) in examples_to_load:
-          filenames.append(filename)
-    return filenames
+    return ckpt_files
 
   def _save_example(
       self,
@@ -235,7 +298,13 @@ class PerExampleCheckpointer(Checkpointer):
 
 
 class BulkCheckpointer(Checkpointer):
-  """Checkpointer that saves all examples to a single file."""
+  """Checkpointer that saves all examples of an evaluation to a single file.
+
+  This checkpointer appends newly processed examples of an evaluation to a
+  single sequence file (e.g., `checkpoint.bagz`). This is often more efficient
+  than `PerExampleCheckpointer` when dealing with a large number of examples
+  or when file system overhead is a concern.
+  """
 
   def _on_bound(self):
     super()._on_bound()
@@ -287,13 +356,24 @@ class BulkCheckpointer(Checkpointer):
         if self._sequence_writer is not None:
           self._sequence_writer[experiment.id] = sequence_writer
 
-  def _list_checkpoint_filenames(
+  def _list_checkpoint_files(
       self, runner: Runner, experiment: Experiment
   ) -> list[str]:
-    if pg.io.path_exists(
-        runner.current_run.input_path_for(experiment, self.checkpoint_filename)
-    ):
-      return [self.checkpoint_filename]
+    # Always honor the output directory if it's present, as it contains both
+    # the warm-started examples and newly processed examples.
+    output_ckpt_file = runner.current_run.output_path_for(
+        experiment, self.checkpoint_filename
+    )
+    if pg.io.path_exists(output_ckpt_file):
+      return [output_ckpt_file]
+
+    if runner.current_run.input_root != runner.current_run.output_root:
+      input_ckpt_file = runner.current_run.input_path_for(
+          experiment, self.checkpoint_filename
+      )
+      if pg.io.path_exists(input_ckpt_file):
+        return [input_ckpt_file]
+    print('CCC', experiment.hash, [])
     return []
 
   def on_experiment_complete(
@@ -341,12 +421,26 @@ class BulkCheckpointer(Checkpointer):
 
 
 class SequenceWriter:
-  """Thread safe sequence writer."""
+  """A thread-safe writer for sequence files (e.g., Bagz) with atomic write.
+
+  `SequenceWriter` wraps a `pg.io.SequenceWriter` to provide thread-safe
+  `add` and `close` operations, ensuring that examples can be written
+  concurrently from multiple threads without corrupting the sequence file.
+  It writes to a temporary file and renames it to target path on `close` to
+  achieve atomic write. If the target path exists, new examples are appended
+  to existing content.
+  """
 
   def __init__(self, path: str):
     self._lock = threading.Lock()
     self._path = path
-    self._sequence_writer = pg.io.open_sequence(path, 'a')
+    basename = os.path.basename(path)
+    self._tmp_path = os.path.join(
+        os.path.dirname(path), f'tmp.{basename}'
+    )
+    if pg.io.path_exists(self._path):
+      pg.io.copy(self._path, self._tmp_path)
+    self._sequence_writer = pg.io.open_sequence(self._tmp_path, 'a')
 
   @property
   def path(self) -> str:
@@ -371,6 +465,14 @@ class SequenceWriter:
         return
       self._sequence_writer.close()
       self._sequence_writer = None
+      pg.io.rename(self._tmp_path, self._path)
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, *args, **kwargs):
+    del args, kwargs
+    self.close()
 
   def __del__(self):
     self.close()

langfun/core/eval/v2/checkpointing_test.py

@@ -65,7 +65,7 @@ class ExampleCollector(experiment_lib.Plugin):
     return self._examples
 
   def on_example_complete(
-      self, runner: runners_lib.Runner,
+      self, runner: experiment_lib.Runner,
       experiment: experiment_lib.Experiment,
       example: example_lib.Example,
   ):
@@ -90,7 +90,10 @@ class PerExampleCheckpointerTest(CheckpointerTest):
     root_dir = os.path.join(tempfile.mkdtemp(), 'per_example_checkpointer')
     experiment = eval_test_helper.test_experiment()
     checkpoint_filename = 'checkpoint.jsonl'
-    checkpointer = checkpointing.PerExampleCheckpointer(checkpoint_filename)
+    checkpointer = checkpointing.PerExampleCheckpointer(
+        checkpoint_filename,
+        enable_inprogress_file=True
+    )
     collector = ExampleCollector()
     run = experiment.run(
         root_dir, 'new', runner='sequential', plugins=[checkpointer, collector]
@@ -102,6 +105,10 @@ class PerExampleCheckpointerTest(CheckpointerTest):
         example = collector.examples[i + 1]
         ckpt = run.output_path_for(leaf, f'checkpoint_{example.id}.jsonl')
         self.assertTrue(pg.io.path_exists(ckpt))
+        inprogress_file = run.output_path_for(
+            leaf, f'{example.id}.inprogress'
+        )
+        self.assertTrue(pg.io.path_exists(inprogress_file))
         with pg.io.open_sequence(ckpt) as f:
           examples_from_ckpt = list(iter(f))
           # `eval_test_helper.test_experiment` has two TestEvaluation with

langfun/core/eval/v2/config_saver.py

@@ -0,0 +1,37 @@
+# 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.
+"""Config saver plugins."""
+
+import os
+from langfun.core.eval.v2 import experiment as experiment_lib
+
+
+class RunConfigSaver(experiment_lib.Plugin):
+  """Saves the current run."""
+
+  def on_run_start(
+      self,
+      runner: experiment_lib.Runner,
+      root: experiment_lib.Experiment
+  ) -> None:
+    del root  # Unused.
+    self._save_run_config(runner)
+
+  def _save_run_config(self, runner: experiment_lib.Runner) -> None:
+    def _save():
+      runner.current_run.save(
+          os.path.join(runner.current_run.output_root, 'run.json'),
+          hide_default_values=True,
+      )
+    runner.background_run(_save)

langfun/core/eval/v2/config_saver_test.py

@@ -0,0 +1,36 @@
+# 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.
+"""Config saver test."""
+
+import os
+import tempfile
+import unittest
+from langfun.core.eval.v2 import config_saver
+from langfun.core.eval.v2 import eval_test_helper
+from langfun.core.eval.v2.runners import parallel  # pylint: disable=unused-import
+
+
+class RunConfigSaverTest(unittest.TestCase):
+
+  def test_save_run_config(self):
+    root_dir = os.path.join(tempfile.mkdtemp(), 'test_run_config_saver')
+    experiment = eval_test_helper.test_evaluation()
+    run = experiment.run(
+        root_dir, 'new', plugins=[config_saver.RunConfigSaver()]
+    )
+    self.assertTrue(os.path.exists(os.path.join(run.output_root, 'run.json')))
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/eval/v2/eval_test_helper.py

@@ -13,6 +13,9 @@
 # limitations under the License.
 """Helper classes and functions for evaluation tests."""
 
+import threading
+import time
+
 from langfun.core import language_model
 from langfun.core import llms
 from langfun.core import message as message_lib
@@ -47,6 +50,8 @@ class TestLLM(llms.Fake):
 
   offset: int = 0
 
+  __test__ = False
+
   def _response_from(self, prompt: message_lib.Message) -> message_lib.Message:
     return message_lib.AIMessage(
         str(prompt.metadata.x + prompt.metadata.y + self.offset)
@@ -63,6 +68,8 @@ class TestEvaluation(Evaluation):
   metrics = [metrics_lib.Match()]
   lm: language_model.LanguageModel = TestLLM()
 
+  __test__ = False
+
   def process(self, example):
     v = example.input
     if v.x == 5:
@@ -75,7 +82,7 @@ class TestEvaluation(Evaluation):
 
 class BadJsonConvertible(pg.Object):
 
-  def to_json(self, *args, **kwargs):
+  def sym_jsonify(self, *args, **kwargs):
     raise ValueError('Cannot convert to JSON.')
 
 
@@ -84,6 +91,8 @@ class TestEvaluationWithExampleCheckpointingError(TestEvaluation):
   inputs = test_inputs()
   metrics = [metrics_lib.Match()]
 
+  __test__ = False
+
   def process(self, example):
     return 1, dict(
         x=BadJsonConvertible()
@@ -101,6 +110,8 @@ class TestEvaluationWithExampleHtmlGenerationError(Evaluation):
   inputs = test_inputs()
   metrics = [metrics_lib.Match()]
 
+  __test__ = False
+
   def process(self, example):
     return 1, dict(
         x=BadHtmlConvertible()
@@ -110,15 +121,22 @@ class TestEvaluationWithExampleHtmlGenerationError(Evaluation):
 class TestEvaluationWithIndexHtmlGenerationError(TestEvaluation):
   """Test evaluation class with bad index HTML generation."""
 
+  __test__ = False
+
   def _html_tree_view(self, *args, **kwargs):
     raise ValueError('Cannot render HTML.')
 
 
+def test_evaluation(offset: int | pg.hyper.OneOf = 0):
+  """Returns a test evaluation."""
+  return TestEvaluation(lm=TestLLM(offset=offset))
+
+
 def test_experiment():
   """Returns a test experiment."""
   return Suite([
-      TestEvaluation(lm=TestLLM(offset=0)),
-      TestEvaluation(lm=TestLLM(offset=pg.oneof(range(5)))),
+      test_evaluation(),
+      test_evaluation(pg.oneof(range(5))),
   ])
 
 
@@ -135,3 +153,86 @@ def test_experiment_with_example_html_generation_error():
 def test_experiment_with_index_html_generation_error():
   """Returns a test experiment with bad index HTML."""
   return TestEvaluationWithIndexHtmlGenerationError()
+
+
+class TestPlugin(experiment_lib.Plugin):
+  """Plugin for testing."""
+
+  started_experiments: list[experiment_lib.Experiment] = []
+  completed_experiments: list[experiment_lib.Experiment] = []
+  skipped_experiments: list[experiment_lib.Experiment] = []
+  started_example_ids: list[int] = []
+  completed_example_ids: list[int] = []
+  start_time: float | None = None
+  complete_time: float | None = None
+
+  __test__ = False
+
+  def _on_bound(self):
+    super()._on_bound()
+    self._lock = threading.Lock()
+
+  def on_run_start(
+      self,
+      runner: experiment_lib.Runner,
+      root: experiment_lib.Experiment
+  ) -> None:
+    del root
+    with pg.notify_on_change(False), pg.allow_writable_accessors(True):
+      self.start_time = time.time()
+
+  def on_run_complete(
+      self,
+      runner: experiment_lib.Runner,
+      root: experiment_lib.Experiment
+  ) -> None:
+    del root
+    with pg.notify_on_change(False), pg.allow_writable_accessors(True):
+      self.complete_time = time.time()
+
+  def on_experiment_start(
+      self,
+      runner: experiment_lib.Runner,
+      experiment: experiment_lib.Experiment
+  ) -> None:
+    del runner
+    with pg.notify_on_change(False), self._lock:
+      self.started_experiments.append(pg.Ref(experiment))
+
+  def on_experiment_skipped(
+      self,
+      runner: experiment_lib.Runner,
+      experiment: experiment_lib.Experiment
+  ) -> None:
+    del runner
+    with pg.notify_on_change(False), self._lock:
+      self.skipped_experiments.append(pg.Ref(experiment))
+
+  def on_experiment_complete(
+      self,
+      runner: experiment_lib.Runner,
+      experiment: experiment_lib.Experiment
+  ) -> None:
+    del runner
+    with pg.notify_on_change(False), self._lock:
+      self.completed_experiments.append(pg.Ref(experiment))
+
+  def on_example_start(
+      self,
+      runner: experiment_lib.Runner,
+      experiment: experiment_lib.Experiment,
+      example: Example
+  ) -> None:
+    del runner, experiment
+    with pg.notify_on_change(False), self._lock:
+      self.started_example_ids.append(example.id)
+
+  def on_example_complete(
+      self,
+      runner: experiment_lib.Runner,
+      experiment: experiment_lib.Experiment,
+      example: Example
+  ) -> None:
+    del runner, experiment
+    with pg.notify_on_change(False), self._lock:
+      self.completed_example_ids.append(example.id)