langfun 0.1.2.dev202511200805__tar.gz → 0.1.2.dev202512150805__tar.gz

{langfun-0.1.2.dev202511200805 → langfun-0.1.2.dev202512150805}/PKG-INFO

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

{langfun-0.1.2.dev202511200805 → langfun-0.1.2.dev202512150805}/langfun/core/agentic/action.py

@@ -1907,7 +1907,6 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
 
   def _on_bound(self):
     super()._on_bound()
-    self._event_handler = None
     self._tls = threading.local()
     self._current_action = self.root
     self._current_execution = self.root.execution

{langfun-0.1.2.dev202511200805 → langfun-0.1.2.dev202512150805}/langfun/core/data/conversion/gemini.py

@@ -138,6 +138,8 @@ class GeminiMessageConverter(lf.MessageConverter):
                 self._safe_read(data, 'mimeType')
             ).from_uri(self._safe_read(data, 'fileUri'))
         )
+      elif 'functionCall' in part or 'functionResponse' in part:
+        pass
       else:
         raise ValueError(f'Unsupported content part: {part!r}.')
     message = message_cls.from_chunks(chunks)

{langfun-0.1.2.dev202511200805 → langfun-0.1.2.dev202512150805}/langfun/core/data/conversion/gemini_test.py

@@ -242,6 +242,42 @@ class GeminiConversionTest(unittest.TestCase):
         'https://www.youtube.com/watch?v=abcd'
     )
 
+  def test_from_value_with_function_call(self):
+    message = lf.Message.from_value(
+        {
+            'role': 'model',
+            'parts': [
+                {'text': 'Let me search for that.'},
+                {
+                    'functionCall': {
+                        'name': 'search',
+                        'args': {'query': 'test'},
+                    }
+                },
+            ],
+        },
+        format='gemini',
+    )
+    self.assertEqual(message.text, 'Let me search for that.')
+
+  def test_from_value_with_function_response(self):
+    message = lf.Message.from_value(
+        {
+            'role': 'user',
+            'parts': [
+                {
+                    'functionResponse': {
+                        'name': 'search',
+                        'response': {'results': ['a', 'b']},
+                    }
+                },
+                {'text': 'Here are the results.'},
+            ],
+        },
+        format='gemini',
+    )
+    self.assertEqual(message.text, 'Here are the results.')
+
 
 if __name__ == '__main__':
   unittest.main()

{langfun-0.1.2.dev202511200805 → langfun-0.1.2.dev202512150805}/langfun/core/eval/v2/__init__.py

@@ -35,11 +35,13 @@ from langfun.core.eval.v2.experiment import Runner
 from langfun.core.eval.v2 import runners
 
 # Plugins
+from langfun.core.eval.v2.config_saver import RunConfigSaver
 from langfun.core.eval.v2.checkpointing import BulkCheckpointer
 from langfun.core.eval.v2.checkpointing import PerExampleCheckpointer
 from langfun.core.eval.v2.reporting import HtmlReporter
 from langfun.core.eval.v2.reporting import ExampleHtmlGenerator
 
+# Google-internal imports.
 
 # pylint: enable=g-bad-import-order
 # pylint: enable=g-importing-member

{langfun-0.1.2.dev202511200805 → langfun-0.1.2.dev202512150805}/langfun/core/eval/v2/checkpointing.py

@@ -14,6 +14,7 @@
 """Checkpointing evaluation runs."""
 import abc
 import datetime
+import os
 import re
 import threading
 import traceback
@@ -37,7 +38,7 @@ class Checkpointer(experiment_lib.Plugin):
   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_filenames` to identify
+  Subclasses should implement `_list_checkpoint_files` to identify
   checkpoint files to load, and `_save_example` to save a newly processed
   example.
   """
@@ -130,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.
@@ -140,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']:
@@ -169,22 +170,18 @@ 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(
@@ -196,10 +193,10 @@ class Checkpointer(experiment_lib.Plugin):
     )
 
   @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(
@@ -225,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,
@@ -340,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(
@@ -394,17 +421,26 @@ class BulkCheckpointer(Checkpointer):
 
 
 class SequenceWriter:
-  """A thread-safe writer for sequence files (e.g., Bagz).
+  """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:
@@ -429,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-0.1.2.dev202512150805/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-0.1.2.dev202512150805/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-0.1.2.dev202511200805 → langfun-0.1.2.dev202512150805}/langfun/core/eval/v2/evaluation.py

@@ -114,6 +114,13 @@ class Evaluation(experiment_lib.Experiment):
     self._log_entries = []
     self._log_lock = threading.Lock()
 
+  def _identity(self) -> str:
+    """Returns the definition of the evaluation."""
+    return self.format(
+        compact=True, hide_default_values=True, use_inferred=True,
+        exclude_keys=('plugins', 'progress', 'usage_summary')
+    )
+
   #
   # Handling evaluation hierarchy (materialized vs. hyper evaluations).
   #
@@ -379,10 +386,10 @@ class Evaluation(experiment_lib.Experiment):
       load_example_metadata: bool = True,
       filter: Callable[[example_lib.Example], bool] | None = None,  # pylint: disable=redefined-builtin
       raise_if_not_exist: bool = False
-  ) -> None:
+  ) -> list[example_lib.Example]:
     """Loads saved state from a sequence IO file."""
     if pg.io.path_exists(state_file):
-      self._state.load(
+      return self._state.load(
           state_file,
           example_input_by_id=self.example_input_by_id,
           load_example_metadata=load_example_metadata,
@@ -390,6 +397,7 @@ class Evaluation(experiment_lib.Experiment):
       )
     elif raise_if_not_exist:
       raise ValueError(f'State file {state_file} does not exist.')
+    return []
 
   def _reset(self) -> None:
     """Resets the state of the evaluation."""

{langfun-0.1.2.dev202511200805 → langfun-0.1.2.dev202512150805}/langfun/core/eval/v2/experiment.py

@@ -268,11 +268,11 @@ class Experiment(lf.Component, pg.views.HtmlTreeView.Extension):
   @functools.cached_property
   def hash(self) -> str:
     """A 8-byte MD5 hash computed from experiment identity."""
-    identity = self.format(
-        compact=True, hide_default_values=True, use_inferred=True,
-        exclude_keys=('plugins', 'progress', 'usage_summary')
-    )
-    return hashlib.md5(identity.encode()).hexdigest()[:8]
+    return hashlib.md5(self._identity().encode()).hexdigest()[:8]
+
+  @abc.abstractmethod
+  def _identity(self) -> str:
+    """Returns the identity of the experiment."""
 
   @classmethod
   def link(cls, path: str) -> str:
@@ -691,6 +691,12 @@ class Suite(Experiment):
     """Returns whether the task is a leaf."""
     return False
 
+  def _identity(self) -> str:
+    """Returns the definition of the evaluation."""
+    return '[' + ', '.join(
+        [child._identity() for child in self.children]  # pylint: disable=protected-access
+    ) + ']'
+
 
 class RunId(pg.Object):
   """Structured repreesentation a experiment run ID."""
@@ -848,10 +854,10 @@ class Run(pg.Object, pg.views.html.HtmlTreeView.Extension):
   ] = None
 
   example_ids: Annotated[
-      list[int] | None,
+      list[int] | Callable[[Experiment], list[int]] | None,
       (
-          'The example IDs to run. If None, it will run all examples. '
-          'Though '
+          'The example IDs to run. Or a callable for determining the examples '
+          'to run based on the experiment. If None, it will run all examples. '
       )
   ] = None
 
@@ -967,10 +973,13 @@ class Run(pg.Object, pg.views.html.HtmlTreeView.Extension):
     """Returns the example IDs to evaluate."""
     if not experiment.is_leaf:
       return set()
-    return set(
-        self.example_ids if self.example_ids else
-        range(1, experiment.num_examples + 1)
-    )
+    if self.example_ids is None:
+      return set(range(1, experiment.num_examples + 1))
+    elif isinstance(self.example_ids, Callable):
+      return set(self.example_ids(experiment))
+    else:
+      assert isinstance(self.example_ids, list), self.example_ids
+      return set(self.example_ids)
 
   def examples_to_reprocess(self, experiment: Experiment) -> set[int]:
     """Returns the example IDs to reprocess per request."""

{langfun-0.1.2.dev202511200805 → langfun-0.1.2.dev202512150805}/langfun/core/eval/v2/progress.py

@@ -92,6 +92,7 @@ class Progress(pg.Object, pg.views.HtmlTreeView.Extension):
         stop_time=None,
         execution_summary=pg.object_utils.TimeIt.StatusSummary(),
     )
+    self._progress_bar = None
 
   @property
   def num_completed(self) -> int:

{langfun-0.1.2.dev202511200805 → langfun-0.1.2.dev202512150805}/langfun/core/eval/v2/progress_tracking.py

@@ -97,8 +97,7 @@ class _TqdmProgressTracker(experiment_lib.Plugin):
     self._leaf_progresses = {
         leaf.id: lf.concurrent.ProgressBar.install(
             label=f'[#{i + 1} - {leaf.id}]',
-            total=(len(runner.current_run.example_ids)
-                   if runner.current_run.example_ids else leaf.num_examples),
+            total=len(runner.current_run.examples_to_evaluate(leaf)),
             color='cyan',
             status=None
         )

{langfun-0.1.2.dev202511200805 → langfun-0.1.2.dev202512150805}/langfun/core/eval/v2/reporting.py

@@ -86,10 +86,8 @@ class ExampleHtmlGenerator(experiment_lib.Plugin):
         return
 
       try:
-        with pg.timeit() as t, pg.io.open(src_file, 'r') as src:
-          content = src.read()
-          with pg.io.open(dest_file, 'w') as dest:
-            dest.write(content)
+        with pg.timeit() as t:
+          pg.io.copy(src_file, dest_file)
         experiment.info(
             f'\'{example.id}.html\' copied in {t.elapse:.2f} seconds.'
         )

{langfun-0.1.2.dev202511200805 → langfun-0.1.2.dev202512150805}/langfun/core/eval/v2/runners/base.py

@@ -22,6 +22,7 @@ from typing import Any, Annotated, Callable, Iterator, Literal
 
 from langfun import core as lf
 from langfun.core.eval.v2 import checkpointing
+from langfun.core.eval.v2 import config_saver
 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
@@ -37,9 +38,6 @@ Experiment = experiment_lib.Experiment
 Plugin = experiment_lib.Plugin
 
 
-_RUN_MANIFEST = 'run.json'
-
-
 class RunnerBase(Runner):
   """Base class for runners with plugin support and IO pooling.
 
@@ -64,6 +62,7 @@ class RunnerBase(Runner):
   plugins = [
       checkpointing.BulkCheckpointer(),
       reporting.HtmlReporter(),
+      config_saver.RunConfigSaver(),
   ]
 
   max_background_threads: Annotated[
@@ -115,24 +114,8 @@ class RunnerBase(Runner):
     for plugin in experiment.plugins:
       yield plugin
 
-  #
-  # IO operations for saving running files.
-  #
-
-  def _save_run_manifest(self) -> None:
-    def _save():
-      pg.symbolic.deref(self.current_run.clone(), recursive=True).save(
-          self.current_run.output_path_for(
-              self.current_run.experiment, _RUN_MANIFEST
-          ),
-          hide_default_values=True
-      )
-    self.background_run(_save)
-
   def on_run_start(self) -> None:
     """Called when a runner is started."""
-    self._save_run_manifest()
-
     for plugin in self._all_plugins(self.current_run.experiment):
       plugin.on_run_start(self, self.current_run.experiment)
 
@@ -152,11 +135,11 @@ class RunnerBase(Runner):
     num_examples_to_evaluate = 0
     if experiment.is_leaf:
       assert isinstance(experiment, Evaluation)
-      num_examples_to_evaluate = (
-          len(self.current_run.example_ids)
-          if self.current_run.example_ids else experiment.num_examples
+      num_examples_to_evaluate = len(
+          self.current_run.examples_to_evaluate(experiment)
       )
       experiment.progress.start(total=num_examples_to_evaluate)
+      pg.io.mkdirs(self.current_run.output_dir(experiment))
     else:
       experiment.progress.start(total=len(experiment.leaf_nodes))
 
@@ -207,10 +190,7 @@ class RunnerBase(Runner):
       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))
-    )
+    example_ids = sorted(self.current_run.examples_to_evaluate(experiment))
     num_from_checkpoint, num_processed = 0, 0
     for example_id in example_ids:
       status = experiment.state.get_status(example_id)
@@ -377,18 +357,14 @@ class RunnerBase(Runner):
         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
-          )
+        items = (
+            Example(
+                id=example_id,
+                input=evaluation.example_input_by_id(example_id)
+            ) for example_id in sorted(
+                self.current_run.examples_to_evaluate(evaluation)
+            )
+        )
         if self.current_run.shuffle_inputs:
           items = list(items)
           random.shuffle(items)

{langfun-0.1.2.dev202511200805 → langfun-0.1.2.dev202512150805}/langfun/core/eval/v2/runners/beam.py

@@ -83,6 +83,9 @@ if beam is not None:
       self._runner = pg.from_json_str(self._runner_str)
       assert isinstance(self._runner, LeafNodeRunner)
       self._runner.setup()
+      self._input_dir = self._runner.current_run.input_dir(
+          self._runner.current_run.experiment
+      )
       self._output_dir = self._runner.current_run.output_dir(
           self._runner.current_run.experiment
       )
@@ -106,6 +109,16 @@ if beam is not None:
       )
       if pg.io.path_exists(ckpt_file):
         yield ckpt_file
+        return
+
+      if self._input_dir != self._output_dir:
+        warmup_ckpt_file = os.path.join(
+            self._input_dir, f'checkpoint_{example_id}.{self._ckpt_format}'
+        )
+        if pg.io.path_exists(warmup_ckpt_file):
+          pg.io.copy(warmup_ckpt_file, ckpt_file)
+          yield ckpt_file
+          return
 
       # Write the in-progress file to indicate that the example is being
       # processed.