langfun 0.1.2.dev202412190804__tar.gz → 0.1.2.dev202412210804__tar.gz

{langfun-0.1.2.dev202412190804 → langfun-0.1.2.dev202412210804}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.1.2.dev202412190804
+Version: 0.1.2.dev202412210804
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.1.2.dev202412190804 → langfun-0.1.2.dev202412210804}/langfun/core/eval/v2/checkpointing.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Checkpointing evaluation runs."""
 import threading
+import traceback
 
 import langfun.core as lf
 from langfun.core.eval.v2 import example as example_lib
@@ -27,6 +28,21 @@ Runner = experiment_lib.Runner
 class Checkpointer(experiment_lib.Plugin):
   """Base class for checkpointing evaluation examples."""
 
+  def on_experiment_start(self, experiment: Experiment):
+    if experiment.state.evaluated_examples:
+      experiment.info(
+          'Loaded %d examples from checkpoint files. Example IDs: %s' %
+          (
+              len(experiment.state.evaluated_examples),
+              list(sorted(experiment.state.evaluated_examples.keys()))
+          ),
+      )
+    else:
+      experiment.info(
+          'No previous evaluated examples are loaded. '
+          f'Experiment {experiment.id} starts from scratch.'
+      )
+
 
 class PerExampleCheckpointer(Checkpointer):
   """Checkpointer that saves each example to a separate file."""
@@ -50,6 +66,10 @@ class PerExampleCheckpointer(Checkpointer):
 
     # For refresh runs, we don't want to load the previous state.
     if not runner.current_run.refresh:
+      if runner.current_run.input_root != runner.current_run.output_root:
+        experiment.info(
+            f'Warm starting from directory: {runner.current_run.input_root}.'
+        )
       def _load_state(ckpt_file):
         experiment.load_state(ckpt_file)
 
@@ -68,10 +88,11 @@ class PerExampleCheckpointer(Checkpointer):
           _load_state, ckpt_files, max_workers=64,
       ):
         if error is not None:
-          pg.logging.warning(
-              'Failed to load checkpoint file %s: %s. Skipping the file.',
-              ckpt_file, error
+          experiment.warning(
+              f'Failed to load checkpoint file {ckpt_file}: {error}. '
+              'Skipping the file.'
           )
+    super().on_experiment_start(experiment)
 
   def on_example_complete(
       self,
@@ -80,7 +101,11 @@ class PerExampleCheckpointer(Checkpointer):
       example: Example,
   ) -> None:
     """Saves the example to the checkpoint file."""
-    if not example.has_error:
+    if example.has_error:
+      experiment.warning(
+          f'Example {example.id} has error. Skipping checkpointing.'
+      )
+    else:
       def save_state(example: Example):
         writer = SequenceWriter(
             runner.current_run.output_path_for(
@@ -91,8 +116,18 @@ class PerExampleCheckpointer(Checkpointer):
                 )
             )
         )
-        writer.add(example)
-        writer.close()
+        try:
+          writer.add(example)
+          writer.close()
+          experiment.info(
+              f'Example {example.id} is saved to {writer.path}.',
+          )
+        except BaseException as e:  # pylint: disable=broad-except
+          experiment.error(
+              f'Failed to save 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]:
@@ -150,6 +185,10 @@ class BulkCheckpointer(Checkpointer):
       return
     # For refresh runs, we don't want to load the previous state.
     if not runner.current_run.refresh:
+      if runner.current_run.input_root != runner.current_run.output_root:
+        experiment.info(
+            f'Warm starting from directory: {runner.current_run.input_root}.'
+        )
       experiment.load_state(
           runner.current_run.input_path_for(
               experiment, self.checkpoint_filename
@@ -164,6 +203,7 @@ class BulkCheckpointer(Checkpointer):
     with self._lock:
       if self._sequence_writer is not None:
         self._sequence_writer[experiment.id] = sequence_writer
+    super().on_experiment_start(experiment)
 
   def on_experiment_complete(
       self,
@@ -178,8 +218,12 @@ class BulkCheckpointer(Checkpointer):
       if self._sequence_writer is not None:
         # Make sure the writer is closed without delay so the file will be
         # available immediately.
-        self._sequence_writer[experiment.id].close()
-        del self._sequence_writer[experiment.id]
+        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 on_example_complete(
       self,
@@ -189,8 +233,22 @@ class BulkCheckpointer(Checkpointer):
   ) -> None:
     """Saves the example to the checkpoint file."""
     assert experiment.id in self._sequence_writer
-    if not example.has_error:
-      runner.background_run(self._sequence_writer[experiment.id].add, example)
+    if example.has_error:
+      experiment.warning(
+          f'Example {example.id} has error. Skipping checkpointing.'
+      )
+    else:
+      def _save_example(example: Example):
+        writer = self._sequence_writer[experiment.id]
+        try:
+          writer.add(example)
+        except BaseException as e:  # pylint: disable=broad-except
+          experiment.error(
+              f'Failed to save example {example.id} to {writer.path}. '
+              f'Error: {e}, Stacktrace: \n{traceback.format_exc()}.',
+          )
+          raise e
+      runner.background_run(_save_example, example)
 
 
 class SequenceWriter:
@@ -198,8 +256,13 @@ class SequenceWriter:
 
   def __init__(self, path: str):
     self._lock = threading.Lock()
+    self._path = path
     self._sequence_writer = pg.io.open_sequence(path, 'w')
 
+  @property
+  def path(self) -> str:
+    return self._path
+
   def add(self, example: Example):
     example_blob = pg.to_json_str(
         example,

{langfun-0.1.2.dev202412190804 → langfun-0.1.2.dev202412210804}/langfun/core/eval/v2/checkpointing_test.py

@@ -16,9 +16,9 @@ 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
-from langfun.core.eval.v2 import test_helper
 import pyglove as pg
 
 Example = example_lib.Example
@@ -56,7 +56,7 @@ class PerExampleCheckpointerTest(unittest.TestCase):
 
   def test_checkpointing(self):
     root_dir = os.path.join(tempfile.gettempdir(), 'per_example_checkpointer')
-    experiment = test_helper.test_experiment()
+    experiment = eval_test_helper.test_experiment()
     checkpoint_filename = 'checkpoint.jsonl'
     checkpointer = checkpointing.PerExampleCheckpointer(checkpoint_filename)
     run = experiment.run(
@@ -89,7 +89,7 @@ class BulkCheckpointerTest(unittest.TestCase):
 
   def test_checkpointing(self):
     root_dir = os.path.join(tempfile.gettempdir(), 'test_bulk_checkpointer')
-    experiment = test_helper.test_experiment()
+    experiment = eval_test_helper.test_experiment()
     checkpoint_filename = 'checkpoint.jsonl'
     checkpointer = checkpointing.BulkCheckpointer(checkpoint_filename)
     run = experiment.run(

{langfun-0.1.2.dev202412190804 → langfun-0.1.2.dev202412210804}/langfun/core/eval/v2/evaluation.py

@@ -285,36 +285,43 @@ class Evaluation(experiment_lib.Experiment):
   # Evaluation-level logging.
   #
 
-  def _log(self, level: lf.logging.LogLevel, message: str, **kwargs):
+  def _log(self, log_func, level: lf.logging.LogLevel, message: str, **kwargs):
+    # Write to external logging system.
+    log_message = f'{self.id}: {message}'
+    if kwargs:
+      log_message = f'{log_message} (metadata: {kwargs!r})'
+    log_func(log_message)
+
+    # Add to experiment log history.
+    log_entry = lf.logging.LogEntry(
+        level=level,
+        time=datetime.datetime.now(),
+        message=message,
+        metadata=kwargs,
+    )
     with self._log_lock:
-      self._log_entries.append(
-          lf.logging.LogEntry(
-              level=level,
-              time=datetime.datetime.now(),
-              message=message,
-              metadata=kwargs,
-          )
-      )
+      self._log_entries.append(log_entry)
 
   def debug(self, message: str, **kwargs):
     """Logs a debug message to the session."""
-    self._log('debug', message, **kwargs)
+    self._log(pg.logging.debug, 'debug', message, **kwargs)
 
   def info(self, message: str, **kwargs):
     """Logs an info message to the session."""
-    self._log('info', message, **kwargs)
+    self._log(pg.logging.info, 'info', message, **kwargs)
 
   def warning(self, message: str, **kwargs):
     """Logs a warning message to the session."""
-    self._log('warning', message, **kwargs)
+    self._log(pg.logging.warning, 'warning', message, **kwargs)
 
   def error(self, message: str, **kwargs):
     """Logs an error message to the session."""
-    self._log('error', message, **kwargs)
+    self._log(pg.logging.error, 'error', message, **kwargs)
 
   def fatal(self, message: str, **kwargs):
     """Logs a fatal message to the session."""
-    self._log('fatal', message, **kwargs)
+    # We use error level for fatal message, which does not trigger assertion.
+    self._log(pg.logging.error, 'fatal', message, **kwargs)
 
   #
   # HTML views.

{langfun-0.1.2.dev202412190804 → langfun-0.1.2.dev202412210804}/langfun/core/eval/v2/evaluation_test.py

@@ -15,12 +15,11 @@ import os
 import tempfile
 import unittest
 
+from langfun.core.eval.v2 import eval_test_helper
 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
 
-from langfun.core.eval.v2 import test_helper
-
 import pyglove as pg
 
 Example = example_lib.Example
@@ -32,17 +31,23 @@ Run = experiment_lib.Run
 class EvaluationTest(unittest.TestCase):
 
   def test_hyper_evaluation(self):
-    exp = test_helper.TestEvaluation(
-        lm=test_helper.TestLLM(offset=pg.oneof(range(3)))
+    exp = eval_test_helper.TestEvaluation(
+        lm=eval_test_helper.TestLLM(offset=pg.oneof(range(3)))
     )
     self.assertFalse(exp.is_leaf)
     self.assertTrue(
         pg.eq(
             exp.children,
             [
-                test_helper.TestEvaluation(lm=test_helper.TestLLM(offset=0)),
-                test_helper.TestEvaluation(lm=test_helper.TestLLM(offset=1)),
-                test_helper.TestEvaluation(lm=test_helper.TestLLM(offset=2)),
+                eval_test_helper.TestEvaluation(
+                    lm=eval_test_helper.TestLLM(offset=0)
+                ),
+                eval_test_helper.TestEvaluation(
+                    lm=eval_test_helper.TestLLM(offset=1)
+                ),
+                eval_test_helper.TestEvaluation(
+                    lm=eval_test_helper.TestLLM(offset=2)
+                ),
             ]
         )
     )
@@ -57,19 +62,21 @@ class EvaluationTest(unittest.TestCase):
     )
 
   def test_input(self):
-    exp = test_helper.TestEvaluation()
+    exp = eval_test_helper.TestEvaluation()
     self.assertEqual(exp.num_examples, 10)
-    exp = test_helper.TestEvaluation(inputs=test_helper.test_inputs(None))
+    exp = eval_test_helper.TestEvaluation(
+        inputs=eval_test_helper.test_inputs(None)
+    )
     self.assertEqual(exp.num_examples, 20)
     @pg.functor
     def my_inputs():
       yield pg.Dict(x=1, y=2)
       yield pg.Dict(x=3, y=4)
-    exp = test_helper.TestEvaluation(inputs=my_inputs())
+    exp = eval_test_helper.TestEvaluation(inputs=my_inputs())
     self.assertEqual(exp.num_examples, 2)
 
   def test_evaluate(self):
-    exp = test_helper.TestEvaluation()
+    exp = eval_test_helper.TestEvaluation()
     example = exp.evaluate(Example(id=3))
     self.assertIs(exp.state.get(3), example)
     self.assertTrue(example.newly_processed)
@@ -85,7 +92,7 @@ class EvaluationTest(unittest.TestCase):
     self.assertIsNotNone(example.start_time)
     self.assertIsNotNone(example.end_time)
 
-    exp = test_helper.TestEvaluation(lm=test_helper.TestLLM(offset=1))
+    exp = eval_test_helper.TestEvaluation(lm=eval_test_helper.TestLLM(offset=1))
     example = exp.evaluate(3)
     self.assertTrue(example.newly_processed)
     self.assertEqual(example.input, pg.Dict(x=2, y=4, groundtruth=6))
@@ -109,7 +116,7 @@ class EvaluationTest(unittest.TestCase):
     pg.io.mkdirs(eval_dir, exist_ok=True)
     state_file = os.path.join(eval_dir, 'state.jsonl')
     with pg.io.open_sequence(state_file, 'w') as f:
-      exp = test_helper.TestEvaluation()
+      exp = eval_test_helper.TestEvaluation()
       example = exp.evaluate(3)
       self.assertTrue(example.newly_processed)
       self.assertEqual(example.input, pg.Dict(x=2, y=4, groundtruth=6))
@@ -132,7 +139,7 @@ class EvaluationTest(unittest.TestCase):
     self.assertEqual(example.usage_summary.uncached.total.num_requests, 0)
 
   def test_html_view(self):
-    exp = test_helper.TestEvaluation()
+    exp = eval_test_helper.TestEvaluation()
     exp.debug('debug message')
     exp.info('info message')
     exp.warning('warning message', x=1)

{langfun-0.1.2.dev202412190804 → langfun-0.1.2.dev202412210804}/langfun/core/eval/v2/experiment.py

@@ -959,6 +959,14 @@ class Plugin(lf.Component):
   ) -> None:
     """Called when an experiment (both leaf and non-leaf) is complete."""
 
+  def on_experiment_abort(
+      self,
+      runner: Runner,
+      experiment: Experiment,
+      error: BaseException,
+  ) -> None:
+    """Called when an experiment (both leaf and non-leaf) is aborted."""
+
   def on_example_start(
       self,
       runner: Runner,

{langfun-0.1.2.dev202412190804 → langfun-0.1.2.dev202412210804}/langfun/core/eval/v2/progress_tracking_test.py

@@ -18,9 +18,9 @@ import tempfile
 import unittest
 
 from langfun.core import console as lf_console
+from langfun.core.eval.v2 import eval_test_helper
 from langfun.core.eval.v2 import progress_tracking    # pylint: disable=unused-import
 from langfun.core.eval.v2 import runners as runners_lib  # pylint: disable=unused-import
-from langfun.core.eval.v2 import test_helper
 import pyglove as pg
 
 
@@ -35,7 +35,7 @@ class HtmlProgressTrackerTest(unittest.TestCase):
         display=display
     )
     root_dir = os.path.join(tempfile.gettempdir(), 'test_html_progress_tracker')
-    experiment = test_helper.test_experiment()
+    experiment = eval_test_helper.test_experiment()
     _ = experiment.run(root_dir, 'new', plugins=[])
     self.assertIsInstance(result['view'], pg.Html)
     lf_console._notebook = None
@@ -45,7 +45,7 @@ class TqdmProgressTrackerTest(unittest.TestCase):
 
   def test_basic(self):
     root_dir = os.path.join(tempfile.gettempdir(), 'test_tqdm_progress_tracker')
-    experiment = test_helper.test_experiment()
+    experiment = eval_test_helper.test_experiment()
     string_io = io.StringIO()
     with contextlib.redirect_stderr(string_io):
       _ = experiment.run(root_dir, 'new', plugins=[])
@@ -55,7 +55,7 @@ class TqdmProgressTrackerTest(unittest.TestCase):
     root_dir = os.path.join(
         tempfile.gettempdir(), 'test_tqdm_progress_tracker_with_example_ids'
     )
-    experiment = test_helper.test_experiment()
+    experiment = eval_test_helper.test_experiment()
     string_io = io.StringIO()
     with contextlib.redirect_stderr(string_io):
       _ = experiment.run(root_dir, 'new', example_ids=[1], plugins=[])

{langfun-0.1.2.dev202412190804 → langfun-0.1.2.dev202412210804}/langfun/core/eval/v2/reporting.py

@@ -14,6 +14,7 @@
 """Reporting evaluation results."""
 
 import time
+import traceback
 from typing import Annotated
 
 from langfun.core.eval.v2 import example as example_lib
@@ -61,6 +62,14 @@ class HtmlReporter(experiment_lib.Plugin):
   ) -> None:
     self._maybe_update_summary(runner, force=True)
 
+  def on_run_abort(
+      self,
+      runner: Runner,
+      root: Experiment,
+      error: BaseException
+  ) -> None:
+    self._maybe_update_summary(runner, force=True)
+
   def on_experiment_start(
       self,
       runner: Runner,
@@ -75,6 +84,16 @@ class HtmlReporter(experiment_lib.Plugin):
     if experiment.is_leaf:
       self._maybe_update_experiment_html(runner, experiment, force=True)
 
+  def on_experiment_abort(
+      self,
+      runner: Runner,
+      experiment: Experiment,
+      error: BaseException
+  ) -> None:
+    del error
+    assert experiment.is_leaf
+    self._maybe_update_experiment_html(runner, experiment, force=True)
+
   def on_example_complete(
       self, runner: Runner, experiment: Experiment, example: Example
   ):
@@ -103,19 +122,26 @@ class HtmlReporter(experiment_lib.Plugin):
       self, runner: Runner, experiment: Experiment, force: bool = False
   ) -> None:
     def _save():
-      html = experiment.to_html(
-          collapse_level=None,
-          extra_flags=dict(
-              current_run=runner.current_run,
-              interactive=False,
-              card_view=False,
-          ),
-      )
-      html.save(
-          runner.current_run.output_path_for(
-              experiment, _EVALULATION_DETAIL_FILE
-          )
+      index_html_path = runner.current_run.output_path_for(
+          experiment, _EVALULATION_DETAIL_FILE
       )
+      try:
+        html = experiment.to_html(
+            collapse_level=None,
+            extra_flags=dict(
+                current_run=runner.current_run,
+                interactive=False,
+                card_view=False,
+            ),
+        )
+        html.save(index_html_path)
+      except BaseException as e:  # pylint: disable=broad-except
+        experiment.error(
+            f'Failed to save HTML {index_html_path!r}. '
+            f'Error: {e}, Stacktrace: \n{traceback.format_exc()}.',
+        )
+        raise e
+
     if force or (
         time.time() - self._last_experiment_report_time[experiment.id]
         > self.experiment_report_interval
@@ -128,17 +154,24 @@ class HtmlReporter(experiment_lib.Plugin):
   ) -> None:
     """Saves the example."""
     def _save():
-      html = example.to_html(
-          collapse_level=None,
-          enable_summary_tooltip=False,
-          extra_flags=dict(
-              # For properly rendering the next link.
-              num_examples=getattr(experiment, 'num_examples', None)
-          ),
-      )
-      html.save(
-          runner.current_run.output_path_for(
-              experiment, f'{example.id}.html'
-          )
-      )
+      try:
+        html = example.to_html(
+            collapse_level=None,
+            enable_summary_tooltip=False,
+            extra_flags=dict(
+                # For properly rendering the next link.
+                num_examples=getattr(experiment, 'num_examples', None)
+            ),
+        )
+        html.save(
+            runner.current_run.output_path_for(
+                experiment, f'{example.id}.html'
+            )
+        )
+      except BaseException as e:  # pylint: disable=broad-except
+        experiment.error(
+            f'Failed to save HTML {example.id}.html. '
+            f'Error: {e}, Stacktrace: \n{traceback.format_exc()}.',
+        )
+        raise e
     runner.background_run(_save)

{langfun-0.1.2.dev202412190804 → langfun-0.1.2.dev202412210804}/langfun/core/eval/v2/reporting_test.py

@@ -15,9 +15,9 @@ import os
 import tempfile
 import unittest
 
+from langfun.core.eval.v2 import eval_test_helper
 from langfun.core.eval.v2 import reporting
 from langfun.core.eval.v2 import runners as runners_lib  # pylint: disable=unused-import
-from langfun.core.eval.v2 import test_helper
 import pyglove as pg
 
 
@@ -25,7 +25,7 @@ class ReportingTest(unittest.TestCase):
 
   def test_reporting(self):
     root_dir = os.path.join(tempfile.gettempdir(), 'test_reporting')
-    experiment = test_helper.test_experiment()
+    experiment = eval_test_helper.test_experiment()
     reporter = reporting.HtmlReporter()
     run = experiment.run(root_dir, 'new', plugins=[reporter])
     pg.io.path_exists(run.output_path_for(experiment, 'summary.html'))