langfun 0.0.2.dev20240429__py3-none-any.whl → 0.1.2.dev202501150804__py3-none-any.whl

langfun/core/eval/v2/checkpointing.py

@@ -0,0 +1,380 @@
+# 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.
+"""Checkpointing evaluation runs."""
+import abc
+import re
+import threading
+import traceback
+from typing import Annotated
+
+import langfun.core as lf
+from langfun.core.eval.v2 import example as example_lib
+from langfun.core.eval.v2 import experiment as experiment_lib
+import pyglove as pg
+
+Example = example_lib.Example
+Experiment = experiment_lib.Experiment
+Runner = experiment_lib.Runner
+
+
+class Checkpointer(experiment_lib.Plugin):
+  """Base class for checkpointing evaluation examples."""
+
+  checkpoint_filename: Annotated[
+      str,
+      'Checkpoint file pattern.'
+  ] = 'checkpoint.bagz'
+
+  def on_experiment_start(
+      self,
+      runner: Runner,
+      experiment: Experiment
+  ) -> None:
+    if not experiment.is_leaf:
+      return
+
+    current_run = runner.current_run
+    if current_run.reprocess is not True:  # pylint: disable=g-bool-id-comparison
+      if current_run.input_root != current_run.output_root:
+        experiment.info(
+            f'Warm starting from directory: {current_run.input_root}.'
+        )
+      self._load_experiment(runner, experiment)
+
+    example_ids_to_evaluate = current_run.examples_to_evaluate(experiment)
+    if experiment.state.evaluated_examples:
+      loaded_example_ids = list(
+          sorted(experiment.state.evaluated_examples.keys())
+      )
+      example_ids_to_evaluate -= set(loaded_example_ids)
+      example_ids_to_evaluate = list(sorted(example_ids_to_evaluate))
+      experiment.info(
+          f'{len(experiment.state.evaluated_examples)} examples '
+          'loaded from checkpoint files. Their outputs will be used '
+          f'for recomputing metrics. Example IDs: {loaded_example_ids}.'
+      )
+      experiment.info(
+          f'{len(example_ids_to_evaluate)} examples will be processed from '
+          f'scratch. Example IDs: {example_ids_to_evaluate}.'
+      )
+    else:
+      experiment.info(
+          'No examples are loaded from checkpoint files. '
+          f'{len(example_ids_to_evaluate)} examples will be processed from '
+          f'scratch. Example IDs: {example_ids_to_evaluate}.'
+      )
+
+  def on_example_complete(
+      self,
+      runner: Runner,
+      experiment: Experiment,
+      example: Example,
+  ) -> None:
+    """Saves the example to the checkpoint file."""
+    if example.has_error:
+      experiment.warning(
+          f'Example {example.id} has error. Skipping checkpointing.'
+      )
+    elif example.newly_processed:
+      self._save_example(runner, experiment, example)
+
+  def _load_experiment(
+      self,
+      runner: Runner,
+      experiment: Experiment,
+  ) -> None:
+    """Creates the checkpoint file."""
+    ckpt_files = self._list_checkpoint_filenames(runner, experiment)
+    experiment.info(f'Found {len(ckpt_files)} checkpoint files to load.')
+
+    # Load the checkpoint files in parallel.
+    current_run = runner.current_run
+    examples_to_load = current_run.examples_to_load(experiment)
+    examples_to_load_metadata = current_run.examples_to_load_metadata(
+        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),
+              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
+          error = e
+        finally:
+          with context['counter_lock']:
+            context['counter'] += 1
+
+          progress_str = f'{context["counter"]}/{len(ckpt_files)}'
+          if error is None:
+            experiment.info(
+                f'Checkpoint file {ckpt_file!r} loaded in {t.elapse:.2f} '
+                f'seconds. ({progress_str})'
+            )
+          else:
+            experiment.warning(
+                f'Failed to load checkpoint file {ckpt_file!r}: {error}. '
+                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}.'
+          )
+
+    _ = list(
+        lf.concurrent_map(
+            _load_state, ckpt_files, max_workers=16, silence_on_errors=None
+        )
+    )
+
+  @abc.abstractmethod
+  def _list_checkpoint_filenames(
+      self, runner: Runner, experiment: Experiment
+  ) -> list[str]:
+    """Lists the checkpoint filenames to restore."""
+
+  @abc.abstractmethod
+  def _save_example(
+      self,
+      runner: Runner,
+      experiment: Experiment,
+      example: Example,
+  ) -> None:
+    """Saves an evaluated example."""
+
+
+class PerExampleCheckpointer(Checkpointer):
+  """Checkpointer that saves each example to a separate file."""
+
+  def _on_bound(self):
+    super()._on_bound()
+    prefix, ext = self._file_prefix_and_ext(self.checkpoint_filename)
+    self._checkpoint_file_prefix = prefix
+    self._checkpoint_file_ext = ext
+
+  def _list_checkpoint_filenames(
+      self, runner: Runner, experiment: Experiment
+  ) -> list[str]:
+    experiment_dir = runner.current_run.input_dir(experiment)
+    filenames = []
+    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('.', '\\.')
+      )
+      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
+
+  def _save_example(
+      self,
+      runner: Runner,
+      experiment: Experiment,
+      example: Example,
+  ) -> None:
+    """Saves the example to the checkpoint file."""
+    def save_state(example: Example):
+      writer = SequenceWriter(
+          runner.current_run.output_path_for(
+              experiment,
+              (
+                  f'{self._checkpoint_file_prefix}_{example.id}'
+                  f'{self._checkpoint_file_ext}'
+              )
+          )
+      )
+      try:
+        writer.add(example)
+        writer.close()
+        experiment.info(
+            f'Example {example.id} checkpointed to {writer.path}.',
+        )
+      except BaseException as e:  # pylint: disable=broad-except
+        experiment.error(
+            f'Failed to checkpoint example {example.id} to {writer.path}. '
+            f'Error: {e}, Stacktrace: \n{traceback.format_exc()}.',
+        )
+        raise e
+    runner.background_run(save_state, example)
+
+  def _file_prefix_and_ext(self, filename: str) -> tuple[str, str]:
+    ext_index = filename.rfind('.')
+    if ext_index == -1:
+      return filename, ''
+    else:
+      return filename[:ext_index], filename[ext_index:]
+
+
+class BulkCheckpointer(Checkpointer):
+  """Checkpointer that saves all examples to a single file."""
+
+  def _on_bound(self):
+    super()._on_bound()
+    self._lock = threading.Lock()
+    self._sequence_writer = None
+
+  def on_run_start(
+      self,
+      runner: Runner,
+      root: Experiment,
+  ) -> None:
+    self._sequence_writer = {}
+
+  def on_run_abort(
+      self,
+      runner: Runner,
+      root: Experiment,
+      error: BaseException
+  ) -> None:
+    with self._lock:
+      if self._sequence_writer is not None:
+        for writer in self._sequence_writer.values():
+          writer.close()
+        self._sequence_writer.clear()
+
+  def on_run_complete(
+      self,
+      runner: Runner,
+      root: Experiment,
+  ) -> None:
+    with self._lock:
+      assert self._sequence_writer is not None and not self._sequence_writer
+
+  def on_experiment_start(
+      self,
+      runner: Runner,
+      experiment: Experiment,
+  ) -> None:
+    super().on_experiment_start(runner, experiment)
+
+    # Prepare the sequence writer for the experiment.
+    if experiment.is_leaf:
+      sequence_writer = SequenceWriter(
+          runner.current_run.output_path_for(
+              experiment, self.checkpoint_filename
+          )
+      )
+      with self._lock:
+        if self._sequence_writer is not None:
+          self._sequence_writer[experiment.id] = sequence_writer
+
+  def _list_checkpoint_filenames(
+      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]
+    return []
+
+  def on_experiment_complete(
+      self,
+      runner: Runner,
+      experiment: Experiment,
+  ) -> None:
+    """Closes the checkpoint file."""
+    if not experiment.is_leaf:
+      return
+    assert experiment.id in self._sequence_writer
+    with self._lock:
+      if self._sequence_writer is not None:
+        # Make sure the writer is closed without delay so the file will be
+        # available immediately.
+        writer = self._sequence_writer.pop(experiment.id)
+        writer.close()
+        experiment.info(
+            f'{len(experiment.state.evaluated_examples)} examples are '
+            f'checkpointed to {writer.path}.'
+        )
+
+  def _save_example(
+      self,
+      runner: Runner,
+      experiment: Experiment,
+      example: Example,
+  ) -> None:
+    """Saves the example to the checkpoint file."""
+    assert experiment.id in self._sequence_writer
+    def _save_example(example: Example):
+      writer = self._sequence_writer[experiment.id]
+      try:
+        writer.add(example)
+        experiment.info(
+            f'Example {example.id} checkpointed to {writer.path}.',
+        )
+      except BaseException as e:  # pylint: disable=broad-except
+        experiment.error(
+            f'Failed to checkpoint example {example.id} to {writer.path}. '
+            f'Error: {e}, Stacktrace: \n{traceback.format_exc()}.',
+        )
+        raise e
+    runner.background_run(_save_example, example)
+
+
+class SequenceWriter:
+  """Thread safe sequence writer."""
+
+  def __init__(self, path: str):
+    self._lock = threading.Lock()
+    self._path = path
+    self._sequence_writer = pg.io.open_sequence(path, 'a')
+
+  @property
+  def path(self) -> str:
+    return self._path
+
+  def add(self, example: Example):
+    example_blob = pg.to_json_str(
+        example,
+        hide_default_values=True,
+        save_ref_value=True,
+        exclude_input=True
+    )
+    with self._lock:
+      if self._sequence_writer is None:
+        return
+      self._sequence_writer.add(example_blob)
+
+  def close(self):
+    # Make sure there is no write in progress.
+    with self._lock:
+      if self._sequence_writer is None:
+        return
+      self._sequence_writer.close()
+      self._sequence_writer = None
+
+  def __del__(self):
+    self.close()

langfun/core/eval/v2/checkpointing_test.py

@@ -0,0 +1,228 @@
+# 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.
+import os
+import tempfile
+import unittest
+
+from langfun.core.eval.v2 import checkpointing
+from langfun.core.eval.v2 import eval_test_helper
+from langfun.core.eval.v2 import example as example_lib
+from langfun.core.eval.v2 import runners as runners_lib  # pylint: disable=unused-import
+import pyglove as pg
+
+Example = example_lib.Example
+
+
+class SequenceWriterTest(unittest.TestCase):
+
+  def test_basic(self):
+    file = os.path.join(tempfile.gettempdir(), 'test.jsonl')
+    writer = checkpointing.SequenceWriter(file)
+    example = Example(id=1, input=pg.Dict(x=1), output=2)
+    writer.add(example)
+    del writer
+    self.assertTrue(pg.io.path_exists(file))
+
+  def test_error_handling(self):
+    file = os.path.join(tempfile.gettempdir(), 'test_error_handling.jsonl')
+    writer = checkpointing.SequenceWriter(file)
+    writer.add(Example(id=1, input=pg.Dict(x=1), output=2))
+
+    def f():
+      raise ValueError('Intentional error')
+
+    try:
+      writer.add(f())
+    except ValueError:
+      del writer
+
+    self.assertTrue(pg.io.path_exists(file))
+    with pg.io.open_sequence(file, 'r') as f:
+      self.assertEqual(len(list(iter(f))), 1)
+
+
+class CheckpointerTest(unittest.TestCase):
+
+  def assert_found_in_log(self, experiment, message):
+    found_error_log = False
+    for log_entry in experiment._log_entries:
+      if log_entry.message.startswith(message):
+        found_error_log = True
+        break
+    self.assertTrue(found_error_log)
+
+
+class PerExampleCheckpointerTest(CheckpointerTest):
+
+  def test_checkpointing(self):
+    root_dir = os.path.join(tempfile.gettempdir(), 'per_example_checkpointer')
+    experiment = eval_test_helper.test_experiment()
+    checkpoint_filename = 'checkpoint.jsonl'
+    checkpointer = checkpointing.PerExampleCheckpointer(checkpoint_filename)
+    run = experiment.run(
+        root_dir, 'new', runner='sequential', plugins=[checkpointer]
+    )
+    num_processed = {}
+    for leaf in experiment.leaf_nodes:
+      for i in range(leaf.num_examples):
+        example = leaf.state.get(i + 1)
+        ckpt = run.output_path_for(leaf, f'checkpoint_{example.id}.jsonl')
+        if example.has_error:
+          self.assertFalse(pg.io.path_exists(ckpt))
+        else:
+          self.assertTrue(pg.io.path_exists(ckpt))
+          with pg.io.open_sequence(ckpt) as f:
+            self.assertEqual(len(list(iter(f))), 1)
+      if leaf.id not in num_processed:
+        self.assertEqual(leaf.progress.num_skipped, 0)
+        num_processed[leaf.id] = leaf.progress.num_processed
+
+    # Run again, should skip existing.
+    _ = experiment.run(
+        root_dir, 'latest', runner='sequential', plugins=[checkpointer]
+    )
+    for leaf in experiment.leaf_nodes:
+      self.assertEqual(leaf.progress.num_skipped, num_processed[leaf.id])
+
+    # Test warm start without reprocess.
+    root_dir = os.path.join(tempfile.gettempdir(), 'per_example_checkpointer2')
+    experiment = eval_test_helper.test_experiment()
+    _ = experiment.run(
+        root_dir, 'new', runner='sequential', plugins=[checkpointer],
+        warm_start_from=run.output_root
+    )
+    for leaf in experiment.leaf_nodes:
+      self.assertEqual(leaf.progress.num_skipped, num_processed[leaf.id])
+
+    # Test warm start with reprocess.
+    root_dir = os.path.join(tempfile.gettempdir(), 'per_example_checkpointer3')
+    experiment = eval_test_helper.test_experiment()
+    _ = experiment.run(
+        root_dir, 'new', runner='sequential', plugins=[checkpointer],
+        warm_start_from=run.output_root,
+        reprocess=True
+    )
+    for leaf in experiment.leaf_nodes:
+      self.assertEqual(leaf.progress.num_skipped, 0)
+
+    root_dir = os.path.join(tempfile.gettempdir(), 'per_example_checkpointer4')
+    experiment = eval_test_helper.test_experiment()
+    _ = experiment.run(
+        root_dir, 'new', runner='sequential', plugins=[checkpointer],
+        warm_start_from=run.output_root,
+        reprocess=[1, 2, 3]
+    )
+    for leaf in experiment.leaf_nodes:
+      self.assertEqual(leaf.progress.num_skipped, num_processed[leaf.id] - 3)
+
+  def test_loading_corrupted_checkpoint(self):
+    root_dir = os.path.join(
+        tempfile.gettempdir(),
+        'per_example_checkpointer_with_corrupted_checkpoint'
+    )
+    experiment = eval_test_helper.TestEvaluation()
+    checkpoint_filename = 'checkpoint.jsonl'
+    checkpointer = checkpointing.PerExampleCheckpointer(checkpoint_filename)
+    run = experiment.run(
+        root_dir, 'new', runner='sequential', plugins=[checkpointer]
+    )
+    num_processed = {}
+    for i in range(experiment.num_examples):
+      example = experiment.state.get(i + 1)
+      ckpt = run.output_path_for(experiment, f'checkpoint_{example.id}.jsonl')
+      if not example.has_error:
+        self.assertTrue(pg.io.path_exists(ckpt))
+        with pg.io.open_sequence(ckpt) as f:
+          self.assertEqual(len(list(iter(f))), 1)
+
+        # Simulate corrupting the first checkpoint.
+        if i == 0:
+          pg.io.writefile(ckpt, 'bad file')
+        num_processed[example.id] = i + 1
+
+    root_dir = os.path.join(
+        tempfile.gettempdir(),
+        'per_example_checkpointer_with_corrupted_checkpoint_warm_start'
+    )
+    experiment = eval_test_helper.TestEvaluation()
+    _ = experiment.run(
+        root_dir, 'new', runner='sequential', plugins=[checkpointer],
+        warm_start_from=run.output_root,
+    )
+    for leaf in experiment.leaf_nodes:
+      self.assertEqual(leaf.progress.num_skipped, len(num_processed) - 1)
+    self.assert_found_in_log(experiment, 'Failed to load checkpoint')
+
+  def test_checkpointing_error(self):
+    root_dir = os.path.join(
+        tempfile.gettempdir(),
+        'per_example_checkpointer_with_checkpointing_error'
+    )
+    experiment = (eval_test_helper
+                  .test_experiment_with_example_checkpointing_error())
+    checkpointer = checkpointing.PerExampleCheckpointer('checkpoint.jsonl')
+    _ = experiment.run(
+        root_dir, 'new', runner='parallel', plugins=[checkpointer]
+    )
+    self.assert_found_in_log(experiment, 'Failed to checkpoint')
+
+
+class BulkCheckpointerTest(CheckpointerTest):
+
+  def test_checkpointing(self):
+    root_dir = os.path.join(tempfile.gettempdir(), 'test_bulk_checkpointer')
+    experiment = eval_test_helper.test_experiment()
+    checkpoint_filename = 'checkpoint.jsonl'
+    checkpointer = checkpointing.BulkCheckpointer(checkpoint_filename)
+    run = experiment.run(
+        root_dir, 'new', runner='sequential', plugins=[checkpointer]
+    )
+    self.assertEqual(len(checkpointer._sequence_writer), 0)
+    num_processed = {}
+    for leaf in experiment.leaf_nodes:
+      ckpt = run.output_path_for(leaf, checkpoint_filename)
+      self.assertTrue(pg.io.path_exists(ckpt))
+      with pg.io.open_sequence(ckpt) as f:
+        self.assertEqual(
+            len(list(iter(f))),
+            leaf.progress.num_completed - leaf.progress.num_failed
+        )
+      if leaf.id not in num_processed:
+        self.assertEqual(leaf.progress.num_skipped, 0)
+        num_processed[leaf.id] = leaf.progress.num_processed
+
+    # Run again, should skip existing.
+    _ = experiment.run(
+        root_dir, 'latest', runner='sequential', plugins=[checkpointer]
+    )
+    self.assertEqual(len(checkpointer._sequence_writer), 0)
+    for leaf in experiment.leaf_nodes:
+      self.assertEqual(leaf.progress.num_skipped, num_processed[leaf.id])
+
+  def test_checkpointing_error(self):
+    root_dir = os.path.join(
+        tempfile.gettempdir(),
+        'bulk_checkpointer_with_checkpointing_error'
+    )
+    experiment = (eval_test_helper
+                  .test_experiment_with_example_checkpointing_error())
+    checkpointer = checkpointing.BulkCheckpointer('checkpoint.jsonl')
+    _ = experiment.run(
+        root_dir, 'new', runner='parallel', plugins=[checkpointer]
+    )
+    self.assert_found_in_log(experiment, 'Failed to checkpoint')
+
+
+if __name__ == '__main__':
+  unittest.main()