langfun 0.1.2.dev202412180804__py3-none-any.whl → 0.1.2.dev202412230804__py3-none-any.whl

langfun/core/eval/v2/reporting.py

@@ -13,11 +13,14 @@
 # limitations under the License.
 """Reporting evaluation results."""
 
+import threading
 import time
+import traceback
 from typing import Annotated
 
 from langfun.core.eval.v2 import example as example_lib
 from langfun.core.eval.v2 import experiment as experiment_lib
+import pyglove as pg
 
 Runner = experiment_lib.Runner
 Experiment = experiment_lib.Experiment
@@ -39,12 +42,15 @@ class HtmlReporter(experiment_lib.Plugin):
   experiment_report_interval: Annotated[
       int,
       'The interval of writing report for inidividual experiments in seconds.'
-  ] = 60
+  ] = 120
 
   def _on_bound(self):
     super()._on_bound()
     self._last_summary_time = 0
     self._last_experiment_report_time = {}
+    self._update_thread = None
+    self._stop_update = False
+    self._stop_update_experiment_ids = set()
 
   def on_run_start(
       self,
@@ -53,14 +59,41 @@ class HtmlReporter(experiment_lib.Plugin):
   ) -> None:
     self._maybe_update_summary(runner)
     self._last_experiment_report_time = {leaf.id: 0 for leaf in root.leaf_nodes}
+    self._stop_update = False
+    self._stop_update_experiment_ids = set()
+    self._update_thread = threading.Thread(
+        target=self._update_thread_func, args=(runner,)
+    )
+    self._update_thread.start()
 
   def on_run_complete(
       self,
       runner: Runner,
       root: Experiment
   ) -> None:
+    self._stop_update = True
     self._maybe_update_summary(runner, force=True)
 
+  def on_run_abort(
+      self,
+      runner: Runner,
+      root: Experiment,
+      error: BaseException
+  ) -> None:
+    self._stop_update = True
+    self._maybe_update_summary(runner, force=True)
+
+  def _update_thread_func(self, runner: Runner):
+    while not self._stop_update:
+      self._maybe_update_summary(runner, background=False)
+      for leaf in runner.current_run.experiment.leaf_nodes:
+        if leaf.id in self._stop_update_experiment_ids:
+          continue
+        self._maybe_update_experiment_html(runner, leaf, background=False)
+        if leaf.progress.is_stopped:
+          self._stop_update_experiment_ids.add(leaf.id)
+      time.sleep(5)
+
   def on_experiment_start(
       self,
       runner: Runner,
@@ -75,6 +108,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
   ):
@@ -82,7 +125,11 @@ class HtmlReporter(experiment_lib.Plugin):
     self._maybe_update_experiment_html(runner, experiment)
     self._maybe_update_summary(runner)
 
-  def _maybe_update_summary(self, runner: Runner, force: bool = False) -> None:
+  def _maybe_update_summary(
+      self,
+      runner: Runner,
+      background: bool = True,
+      force: bool = False) -> None:
     """Maybe update the summary of current run."""
     run = runner.current_run
     def _summary():
@@ -96,31 +143,52 @@ class HtmlReporter(experiment_lib.Plugin):
       )
 
     if force or (time.time() - self._last_summary_time > self.summary_interval):
-      runner.background_run(_summary)
+      if background:
+        runner.background_run(_summary)
+      else:
+        _summary()
       self._last_summary_time = time.time()
 
   def _maybe_update_experiment_html(
-      self, runner: Runner, experiment: Experiment, force: bool = False
+      self,
+      runner: Runner,
+      experiment: Experiment,
+      force: bool = False,
+      background: bool = True,
   ) -> None:
     def _save():
-      html = experiment.to_html(
-          collapse_level=None,
-          extra_flags=dict(
-              current_run=runner.current_run,
-              interactive=False,
-              card_view=False,
-          ),
+      index_html_path = runner.current_run.output_path_for(
+          experiment, _EVALULATION_DETAIL_FILE
       )
-      html.save(
-          runner.current_run.output_path_for(
-              experiment, _EVALULATION_DETAIL_FILE
+      try:
+        with pg.timeit() as t:
+          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)
+          experiment.info(
+              f'Generated HTML {index_html_path!r} in {t.elapse:.2f} seconds.',
+          )
+      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
     ):
-      runner.background_run(_save)
+      if background:
+        runner.background_run(_save)
+      else:
+        _save()
       self._last_experiment_report_time[experiment.id] = time.time()
 
   def _save_example_html(
@@ -128,17 +196,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/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'))

langfun/core/eval/v2/runners.py

@@ -18,6 +18,7 @@ import concurrent.futures
 import random
 import threading
 import time
+import traceback
 from typing import Any, Annotated, Callable, Iterator
 
 from langfun import core as lf
@@ -64,6 +65,7 @@ class RunnerBase(Runner):
     with pg.notify_on_change(False):
       self.plugins.append(progress_tracking.progress_tracker(self.tqdm))
 
+    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
@@ -75,7 +77,10 @@ class RunnerBase(Runner):
         func(*args, **kwargs)
       except Exception as e:  # pylint: disable=broad-except
         self._background_last_error = e
-    self._io_pool.submit(_background_run, *args, **kwargs)
+
+    with self._io_pool_lock:
+      if self._io_pool is not None:
+        self._io_pool.submit(_background_run, *args, **kwargs)
 
   def _all_plugins(self, experiment: Experiment) -> Iterator[Plugin]:
     """Returns all plugins for the experiment."""
@@ -120,9 +125,14 @@ class RunnerBase(Runner):
     # Start the progress of the evaluation.
     if experiment.is_leaf:
       assert isinstance(experiment, Evaluation)
-      experiment.progress.start(
-          total=(len(self.current_run.example_ids)
-                 if self.current_run.example_ids else experiment.num_examples)
+      num_examples_to_evaluate = (
+          len(self.current_run.example_ids)
+          if self.current_run.example_ids else experiment.num_examples
+      )
+      experiment.progress.start(total=num_examples_to_evaluate)
+      experiment.info(
+          'Starting evaluation %s with %d examples to evaluate.'
+          % (experiment.id, num_examples_to_evaluate)
       )
     else:
       experiment.progress.start(total=len(experiment.leaf_nodes))
@@ -144,8 +154,7 @@ class RunnerBase(Runner):
 
     # Only leaf evaluations will trigger the complete notification of the
     # ancestors.
-    if experiment.is_leaf:
-      self._update_ancestor_progresses(experiment)
+    self._update_ancestor_progresses(experiment)
 
   def on_experiment_complete(self, experiment: Experiment) -> None:
     """Called when an evaluation is complete."""
@@ -160,6 +169,35 @@ class RunnerBase(Runner):
     # ancestors.
     if experiment.is_leaf:
       self._update_ancestor_progresses(experiment)
+      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))
+    )
+    num_from_checkpoint, num_processed = 0, 0
+    for example_id in example_ids:
+      example = experiment.state.get(example_id)
+      if example.newly_processed:
+        num_processed += 1
+      else:
+        num_from_checkpoint += 1
+    experiment.info(
+        f'{experiment.id} completed with {num_from_checkpoint + num_processed} '
+        f'examples evaluated ({num_from_checkpoint} from checkpoint, '
+        f'{num_processed} newly processed).'
+    )
+
+  def on_experiment_abort(
+      self, experiment: Experiment, error: BaseException) -> None:
+    """Called when an evaluation is complete."""
+    assert experiment.is_leaf
+    experiment.fatal(f'{error}\n\n{traceback.format_exc()}')
+
+    # Notify the plugins of the experiment abort.
+    for plugin in self._all_plugins(experiment):
+      plugin.on_experiment_abort(self, experiment, error)
 
   def _update_ancestor_progresses(self, experiment: Experiment):
     """Updates the progresses of the parent nodes of the experiment."""
@@ -262,7 +300,9 @@ class RunnerBase(Runner):
         self.background_run(cache.save)
 
       # Wait for the background tasks to finish.
-      self._io_pool.shutdown(wait=True)
+      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:
@@ -270,31 +310,36 @@ class RunnerBase(Runner):
 
   def run_evaluation(self, evaluation: Evaluation) -> None:
     """Runs the evaluation."""
-    self.on_experiment_start(evaluation)
-
-    per_evaluation_settings = {}
-    cache = None
-    if self.current_run.use_cache == 'per_dataset':
-      cache = self._load_or_create_cache(evaluation)
-      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
-        )
-      self._evaluate_items(evaluation, items)
-
-    if cache:
-      self.background_run(cache.save)
-    self.on_experiment_complete(evaluation)
+    try:
+      self.on_experiment_start(evaluation)
+
+      per_evaluation_settings = {}
+      cache = None
+      if self.current_run.use_cache == 'per_dataset':
+        cache = self._load_or_create_cache(evaluation)
+        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
+          )
+        self._evaluate_items(evaluation, items)
+
+      if cache:
+        self.background_run(cache.save)
+      self.on_experiment_complete(evaluation)
+    except BaseException as e:  # pylint: disable=broad-except
+      self.on_experiment_abort(evaluation, e)
+      raise e
 
   @abc.abstractmethod
   def _evaluate_items(
@@ -410,9 +455,7 @@ class ParallelRunner(RunnerBase):
         groups.values(),
         max_workers=max(64, len(groups)),
         timeout=self.timeout,
-        silence_on_errors=(
-            None if self.current_run.raise_if_has_error else BaseException
-        )
+        silence_on_errors=None,
     ):
       pass
 
@@ -437,8 +480,6 @@ class ParallelRunner(RunnerBase):
         items,
         max_workers=evaluation.max_workers,
         timeout=self.timeout,
-        silence_on_errors=(
-            None if self.current_run.raise_if_has_error else BaseException
-        )
+        silence_on_errors=None,
     ):
       pass

langfun/core/eval/v2/runners_test.py

@@ -18,10 +18,11 @@ import time
 from typing import Any
 import unittest
 
+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 experiment as experiment_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
 
 
@@ -101,7 +102,7 @@ class RunnerTest(unittest.TestCase):
 
   def test_basic(self):
     plugin = TestPlugin()
-    exp = test_helper.test_experiment()
+    exp = eval_test_helper.test_experiment()
     root_dir = os.path.join(tempfile.gettempdir(), 'test_sequential_runner')
     run = exp.run(root_dir, runner='sequential', plugins=[plugin])
 
@@ -143,7 +144,7 @@ class RunnerTest(unittest.TestCase):
 
   def test_raise_if_has_error(self):
     root_dir = os.path.join(tempfile.gettempdir(), 'test_raise_if_has_error')
-    exp = test_helper.TestEvaluation()
+    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
@@ -154,7 +155,7 @@ class RunnerTest(unittest.TestCase):
 
   def test_example_ids(self):
     root_dir = os.path.join(tempfile.gettempdir(), 'test_example_ids')
-    exp = test_helper.test_experiment()
+    exp = eval_test_helper.test_experiment()
     plugin = TestPlugin()
     _ = exp.run(
         root_dir, runner='sequential', plugins=[plugin], example_ids=[5, 7, 9]
@@ -164,7 +165,7 @@ class RunnerTest(unittest.TestCase):
 
   def test_filter(self):
     plugin = TestPlugin()
-    exp = test_helper.test_experiment()
+    exp = eval_test_helper.test_experiment()
     root_dir = os.path.join(tempfile.gettempdir(), 'test_filter')
 
     _ = exp.run(
@@ -193,7 +194,7 @@ class RunnerTest(unittest.TestCase):
           ) for i in range(num_examples)
       ]
 
-    exp = test_helper.TestEvaluation(
+    exp = eval_test_helper.TestEvaluation(
         inputs=test_inputs(num_examples=pg.oneof([2, 4]))
     )
     # Global cache.
@@ -234,7 +235,7 @@ class ParallelRunnerTest(RunnerTest):
 
   def test_parallel_runner(self):
     plugin = TestPlugin()
-    exp = test_helper.test_experiment()
+    exp = eval_test_helper.test_experiment()
     root_dir = os.path.join(tempfile.gettempdir(), 'test_parallel_runner')
     run = exp.run(root_dir, runner='parallel', plugins=[plugin])
 
@@ -274,7 +275,7 @@ class ParallelRunnerTest(RunnerTest):
 
   def test_concurrent_startup_delay(self):
     plugin = TestPlugin()
-    exp = test_helper.test_experiment()
+    exp = eval_test_helper.test_experiment()
     root_dir = os.path.join(
         tempfile.gettempdir(), 'test_concurrent_startup_delay'
     )
@@ -290,7 +291,7 @@ class DebugRunnerTest(RunnerTest):
 
   def test_debug_runner(self):
     plugin = TestPlugin()
-    exp = test_helper.test_experiment()
+    exp = eval_test_helper.test_experiment()
     root_dir = os.path.join(tempfile.gettempdir(), 'test_debug_runner')
     run = exp.run(root_dir, runner='debug', plugins=[plugin])
 

langfun/core/logging.py

@@ -54,6 +54,25 @@ class LogEntry(pg.Object, pg.views.HtmlTreeView.Extension):
   def should_output(self, min_log_level: LogLevel) -> bool:
     return _LOG_LEVELS.index(self.level) >= _LOG_LEVELS.index(min_log_level)
 
+  def format(self,
+             compact: bool = False,
+             verbose: bool = True,
+             root_indent: int = 0,
+             *,
+             text_format: bool = True,
+             **kwargs):
+    if text_format:
+      s = f"""{self.time.strftime('%H:%M:%S')} {self.level.upper()} - {self.message}"""
+      if self.metadata:
+        s += f' (metadata: {self.metadata!r})'
+      return s
+    return super().format(
+        compact=compact,
+        verbose=verbose,
+        root_indent=root_indent,
+        **kwargs
+    )
+
   def _html_tree_view_summary(
       self,
       view: pg.views.HtmlTreeView,

langfun/core/logging_test.py

@@ -61,6 +61,25 @@ class LoggingTest(unittest.TestCase):
       print(actual)
     self.assertEqual(actual, expected)
 
+  def test_format(self):
+    time = datetime.datetime(2024, 10, 10, 12, 30, 45)
+    self.assertEqual(
+        str(
+            logging.LogEntry(
+                level='info', message='hello\nworld',
+                time=time, metadata=dict(x=1),
+            )
+        ),
+        '12:30:45 INFO - hello\nworld (metadata: {x=1})',
+    )
+    self.assertIn(
+        'LogEntry(',
+        logging.LogEntry(
+            level='info', message='hello\nworld',
+            time=time, metadata=dict(x=1),
+        ).format(text_format=False),
+    )
+
   def test_html(self):
     time = datetime.datetime(2024, 10, 10, 12, 30, 45)
     self.assert_html_content(

langfun/core/structured/parsing.py

@@ -270,24 +270,31 @@ def call(
   if schema in (str, None):
     return lm_output if returns_message else lm_output.text
 
+  def _chain_nl_output_message(parsing_message: lf.Message):
+    """Chain the source of the parsed output to the LM output."""
+    parsing_message.root.source = lm_output
+    parsing_message.tag('parsing-lm-output')
+    parsing_message.lm_input.tag('parsing-lm-input')
+
   # Call `parsing_lm` for structured parsing.
-  parsing_message = querying.query(
-      lm_output.text,
-      schema,
-      examples=parsing_examples,
-      lm=parsing_lm or lm,
-      include_context=parsing_include_context,
-      cache_seed=cache_seed,
-      autofix=autofix,
-      autofix_lm=autofix_lm or lm,
-      protocol=protocol,
-      returns_message=True,
-      **kwargs,
-  )
-  # Chain the source of the parsed output to the LM output.
-  parsing_message.root.source = lm_output
-  parsing_message.tag('parsing-lm-output')
-  parsing_message.lm_input.tag('parsing-lm-input')
+  try:
+    parsing_message = querying.query(
+        lm_output.text,
+        schema,
+        examples=parsing_examples,
+        lm=parsing_lm or lm,
+        include_context=parsing_include_context,
+        cache_seed=cache_seed,
+        autofix=autofix,
+        autofix_lm=autofix_lm or lm,
+        protocol=protocol,
+        returns_message=True,
+        **kwargs,
+    )
+    _chain_nl_output_message(parsing_message)
+  except mapping.MappingError as e:
+    _chain_nl_output_message(e.lm_response)
+    raise e
   return parsing_message if returns_message else parsing_message.result
 
 

langfun/core/structured/parsing_test.py

@@ -686,6 +686,31 @@ class CallTest(unittest.TestCase):
         ],
         returns_message=True,
     )
+    self.assertIn('parsing-lm-output', output.tags)
+    self.assertIn('parsing-lm-input', output.source.tags)
+    self.assertEqual(output.root.text, 'Compute 1 + 2')
+
+  def test_call_with_parsing_message_chaining_on_parsing_error(self):
+    try:
+      output = parsing.call(
+          'Compute 1 + 2',
+          int,
+          lm=fake.StaticSequence(['three']),
+          parsing_lm=fake.StaticSequence(['abc']),
+          parsing_examples=[
+              mapping.MappingExample(
+                  context='Multiple four and five',
+                  input='twenty',
+                  schema=int,
+                  output=20,
+              )
+          ],
+          returns_message=True,
+      )
+    except mapping.MappingError as e:
+      output = e.lm_response
+    self.assertIn('parsing-lm-output', output.tags)
+    self.assertIn('parsing-lm-input', output.source.tags)
     self.assertEqual(output.root.text, 'Compute 1 + 2')
 
   def test_call_with_autofix(self):

langfun/core/structured/querying.py

@@ -583,7 +583,16 @@ class QueryInvocation(pg.Object, pg.views.HtmlTreeView.Extension):
 
   @functools.cached_property
   def output(self) -> Any:
-    return query_output(self.lm_response, self.schema)
+    """The output of `lf.query`. If it failed, returns the `MappingError`."""
+    try:
+      return query_output(self.lm_response, self.schema)
+    except mapping.MappingError as e:
+      return e
+
+  @property
+  def has_error(self) -> bool:
+    """Returns True if the query failed to generate a valid output."""
+    return isinstance(self.output, BaseException)
 
   @property
   def elapse(self) -> float:

langfun/core/structured/querying_test.py

@@ -1051,6 +1051,16 @@ class QueryStructureJsonTest(unittest.TestCase):
 
 class QueryInvocationTest(unittest.TestCase):
 
+  def test_basics(self):
+    lm = fake.StaticSequence([
+        'Activity(description="hi"',
+    ])
+    with querying.track_queries() as queries:
+      querying.query('foo', Activity, default=None, lm=lm)
+
+    self.assertTrue(queries[0].has_error)
+    self.assertIsInstance(queries[0].output, mapping.MappingError)
+
   def test_to_html(self):
     lm = fake.StaticSequence([
         'Activity(description="hi")',

{langfun-0.1.2.dev202412180804.dist-info → langfun-0.1.2.dev202412230804.dist-info}/METADATA

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