langfun 0.1.2.dev202412200804__py3-none-any.whl → 0.1.2.dev202412220804__py3-none-any.whl

langfun/core/eval/v2/checkpointing.py

@@ -66,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)
 
@@ -85,8 +89,8 @@ class PerExampleCheckpointer(Checkpointer):
       ):
         if error is not None:
           experiment.warning(
-              'Failed to load checkpoint file %s: %s. Skipping the file.',
-              ckpt_file, error
+              f'Failed to load checkpoint file {ckpt_file}: {error}. '
+              'Skipping the file.'
           )
     super().on_experiment_start(experiment)
 
@@ -181,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

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/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/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/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

@@ -65,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
@@ -76,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."""
@@ -296,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:

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-0.1.2.dev202412200804.dist-info → langfun-0.1.2.dev202412220804.dist-info}/METADATA

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

{langfun-0.1.2.dev202412200804.dist-info → langfun-0.1.2.dev202412220804.dist-info}/RECORD

@@ -58,10 +58,11 @@ langfun/core/eval/patching_test.py,sha256=8kCd54Egjju22FMgtJuxEsrXkW8ifs-UUBHtrC
 langfun/core/eval/scoring.py,sha256=B69IsIxiPs1xZcOBFIhZF70YmDue2Siik-CPL2bh33s,6254
 langfun/core/eval/scoring_test.py,sha256=O8olHbrUEg60gMxwOkWzKBJZpZoUlmVnBANX5Se2SXM,4546
 langfun/core/eval/v2/__init__.py,sha256=qoa6zKdFXOFyCX6vay6OdgPf1eUhYGoHYAxe35qECGk,1628
-langfun/core/eval/v2/checkpointing.py,sha256=ZYKHN91pALZjnoMdikc-JdJ-HiBRIqitOUWbFablFI8,8367
-langfun/core/eval/v2/checkpointing_test.py,sha256=dAERKQTW_PM1B0oUauB0YVQkMEI-cgJq0q-wAVlGYpU,4383
+langfun/core/eval/v2/checkpointing.py,sha256=zr2hxOjm6Hdq71sYTsbQtL_CwQOWr-Ir9T5TPUnhqMI,8741
+langfun/core/eval/v2/checkpointing_test.py,sha256=Imy96lwDkvmtj-1YFpP2DZukMOoYqpPov2J_MsQKxxI,4398
+langfun/core/eval/v2/eval_test_helper.py,sha256=pDpZTBnWRR5xjJv3Uy3NWEzArqlL8FTMOgeR4C53F5M,2348
 langfun/core/eval/v2/evaluation.py,sha256=NFBGAWw2BtW7H0zcoZhfWtz59Psra84eshJm73uAFwg,21807
-langfun/core/eval/v2/evaluation_test.py,sha256=ld8oBOjsfN-LNLL2eViSTu17wAq90GcsfURXX6oVlFo,6014
+langfun/core/eval/v2/evaluation_test.py,sha256=GmV1TiqX1V15st2qpcGWooM5hudomQVjW5kajovGDvE,6231
 langfun/core/eval/v2/example.py,sha256=fURrvdNmMsVMqoEErcsmLmC6Xq3ny16dYsnLH8HVlcY,9626
 langfun/core/eval/v2/example_test.py,sha256=WcJmU7IQQXvjFia63mokySC4CqxzVL9Wso1sC5F0YK8,3032
 langfun/core/eval/v2/experiment.py,sha256=xfk4aNZ3dH46y0lWSS_fC7JpfJCG77Z5qsakV4gHcOs,29762
@@ -73,12 +74,11 @@ langfun/core/eval/v2/metrics_test.py,sha256=p4FzLJsE8XAzAQuyP9hfEf9YeKWZ__PO_ue8
 langfun/core/eval/v2/progress.py,sha256=azZgssQgNdv3IgjKEaQBuGI5ucFDNbdi02P4z_nQ8GE,10292
 langfun/core/eval/v2/progress_test.py,sha256=YU7VHzmy5knPZwj9vpBN3rQQH2tukj9eKHkuBCI62h8,2540
 langfun/core/eval/v2/progress_tracking.py,sha256=l9fEkz4oP5McpZzf72Ua7PYm3lAWtRru7gRWNf8H0ms,6083
-langfun/core/eval/v2/progress_tracking_test.py,sha256=iO-DslCJWncU7-27XaMKxDeKrsGbwdk_tKfoRk3KboE,2271
+langfun/core/eval/v2/progress_tracking_test.py,sha256=fouMVJkFJqHjbhQJngGLGCmA9x3n0dU4USI2dY163mg,2291
 langfun/core/eval/v2/reporting.py,sha256=vsh45GLVnA7GMU-8cvNYOt4Nb7mEwvcguhO-BSXSzTE,5358
-langfun/core/eval/v2/reporting_test.py,sha256=JxffbUPWInUyLjo-AQVFrllga884Mdfm05R86FtxSss,1482
-langfun/core/eval/v2/runners.py,sha256=iTARDpPHPJKZL5Hu3k-O2LhK1ICOo1Ywbm5e7O2rNBA,15680
-langfun/core/eval/v2/runners_test.py,sha256=UeiUNygux_U6iGVG18rhp68ZE4hoWeoT6XsXvSjxNQg,11620
-langfun/core/eval/v2/test_helper.py,sha256=pDpZTBnWRR5xjJv3Uy3NWEzArqlL8FTMOgeR4C53F5M,2348
+langfun/core/eval/v2/reporting_test.py,sha256=4nobW6pcaatiZh8u4xciexciaiZNDlDoJci157Wp_RI,1492
+langfun/core/eval/v2/runners.py,sha256=t6_yHAJ4HWufK4wvh_OntKcok2KquA5ARIHIk1vvEwc,15870
+langfun/core/eval/v2/runners_test.py,sha256=A37fKK2MvAVTiShsg_laluJzJ9AuAQn52k7HPbfD0Ks,11666
 langfun/core/llms/__init__.py,sha256=lWXKjGHv66ShG7AE_Bc4QM7SDTxJdfoQMn3PF0lr0sU,6461
 langfun/core/llms/anthropic.py,sha256=afKZmdiLcosS_UEBlB8WKyf1K-zeXgwtPAx6ofg2Gww,13989
 langfun/core/llms/anthropic_test.py,sha256=-2U4kc_pgBM7wqxu8RuxzyHPGww1EAWqKUvN4PW8Btw,8058
@@ -148,8 +148,8 @@ langfun/core/templates/demonstration.py,sha256=vCrgYubdZM5Umqcgp8NUVGXgr4P_c-fik
 langfun/core/templates/demonstration_test.py,sha256=SafcDQ0WgI7pw05EmPI2S4v1t3ABKzup8jReCljHeK4,2162
 langfun/core/templates/selfplay.py,sha256=yhgrJbiYwq47TgzThmHrDQTF4nDrTI09CWGhuQPNv-s,2273
 langfun/core/templates/selfplay_test.py,sha256=Ot__1P1M8oJfoTp-M9-PQ6HUXqZKyMwvZ5f7yQ3yfyM,2326
-langfun-0.1.2.dev202412200804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202412200804.dist-info/METADATA,sha256=7IgyNjTrvkrYSyuwvSrBzrY8tzacbgL1qLaoonRRgxc,8281
-langfun-0.1.2.dev202412200804.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-langfun-0.1.2.dev202412200804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202412200804.dist-info/RECORD,,
+langfun-0.1.2.dev202412220804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202412220804.dist-info/METADATA,sha256=SKSVb1h469QIDbiOKMxz5HwTjTIOWoomol55YNkkFqU,8281
+langfun-0.1.2.dev202412220804.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+langfun-0.1.2.dev202412220804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202412220804.dist-info/RECORD,,