langfun 0.1.2.dev202412180804__tar.gz → 0.1.2.dev202412270804__tar.gz

{langfun-0.1.2.dev202412180804 → langfun-0.1.2.dev202412270804}/PKG-INFO

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

langfun-0.1.2.dev202412270804/langfun/core/eval/v2/checkpointing.py

@@ -0,0 +1,346 @@
+# 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 threading
+import traceback
+
+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."""
+
+  def on_experiment_start(
+      self,
+      runner: Runner,
+      experiment: Experiment
+  ) -> None:
+    if not experiment.is_leaf:
+      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}.'
+        )
+      self._load_experiment(runner, experiment)
+
+    if experiment.state.evaluated_examples:
+      loaded_example_ids = list(
+          sorted(experiment.state.evaluated_examples.keys())
+      )
+      example_ids_to_evaluate = (
+          set(runner.current_run.example_ids) if runner.current_run.example_ids
+          else set(range(1, experiment.num_examples + 1))
+      )
+      example_ids_to_evaluate -= set(loaded_example_ids)
+
+      experiment.info(
+          f'{len(experiment.state.evaluated_examples)} examples have been '
+          '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: {list(sorted(example_ids_to_evaluate))}'
+      )
+    else:
+      experiment.info(
+          'No examples are loaded from checkpoint files. '
+          f'Experiment {experiment.id} starts from scratch.'
+      )
+
+  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.'
+      )
+    else:
+      self._save_example(runner, experiment, example)
+
+  @abc.abstractmethod
+  def _load_experiment(self, runner: Runner, experiment: Experiment) -> None:
+    """Loads the experiment state from checkpoint files."""
+
+  @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."""
+
+  checkpoint_filename: str = 'checkpoint.bagz'
+
+  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 _load_experiment(
+      self,
+      runner: Runner,
+      experiment: Experiment,
+  ) -> None:
+    """Creates the checkpoint file."""
+    experiment_dir = runner.current_run.input_dir(experiment)
+    if pg.io.path_exists(experiment_dir):
+      ckpt_files = [
+          runner.current_run.input_path_for(experiment, filename)
+          for filename in pg.io.listdir(experiment_dir)
+          if filename.startswith(self._checkpoint_file_prefix)
+          and filename.endswith(self._checkpoint_file_ext)
+      ]
+    else:
+      ckpt_files = []
+
+    experiment.info(f'Found {len(ckpt_files)} checkpoint files to load.')
+
+    # Load the checkpoint files in parallel.
+    context = dict(counter=0, counter_lock=threading.Lock())
+    def _load_state(ckpt_file):
+      error = None
+      with pg.timeit() as t:
+        try:
+          experiment.load_state(ckpt_file)
+        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'Loaded checkpoint file {ckpt_file} in {t.elapse:.2f} '
+                f'seconds. ({progress_str})'
+            )
+          else:
+            experiment.warning(
+                f'Failed to load checkpoint file {ckpt_file}: {error}. '
+                f'Skipping the file. ({progress_str})'
+            )
+
+    _ = list(
+        lf.concurrent_map(
+            _load_state, ckpt_files, max_workers=16, silence_on_errors=None
+        )
+    )
+
+  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} 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]:
+    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."""
+
+  checkpoint_filename: str = 'checkpoint.bagz'
+
+  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 _load_experiment(
+      self,
+      runner: Runner,
+      experiment: Experiment,
+  ) -> None:
+    """Creates the checkpoint file."""
+    experiment.load_state(
+        runner.current_run.input_path_for(
+            experiment, self.checkpoint_filename
+        ),
+        raise_if_not_exist=False
+    )
+
+  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} added 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_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, 'w')
+
+  @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-0.1.2.dev202412180804 → langfun-0.1.2.dev202412270804}/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
@@ -55,8 +55,9 @@ class SequenceWriterTest(unittest.TestCase):
 class PerExampleCheckpointerTest(unittest.TestCase):
 
   def test_checkpointing(self):
+    pg.defaults.loggers.use_stdout()
     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 +90,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.dev202412180804 → langfun-0.1.2.dev202412270804}/langfun/core/eval/v2/evaluation.py

@@ -14,7 +14,9 @@
 """Base class for Langfun evaluation tasks."""
 
 import abc
+import datetime
 import functools
+import threading
 import time
 
 from typing import Annotated, Any, Callable, Iterable
@@ -63,6 +65,8 @@ class Evaluation(experiment_lib.Experiment):
     self.__dict__.pop('is_leaf', None)
     self.__dict__.pop('children', None)
     super()._on_bound()
+    self._log_entries = []
+    self._log_lock = threading.Lock()
 
   #
   # Handling evaluation hierarchy (materialized vs. hyper evaluations).
@@ -277,6 +281,48 @@ class Evaluation(experiment_lib.Experiment):
       for metric in self.metrics:
         metric.reset()
 
+  #
+  # Evaluation-level logging.
+  #
+
+  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(log_entry)
+
+  def debug(self, message: str, **kwargs):
+    """Logs a debug message to the session."""
+    self._log(pg.logging.debug, 'debug', message, **kwargs)
+
+  def info(self, message: str, **kwargs):
+    """Logs an info message to the session."""
+    self._log(pg.logging.info, 'info', message, **kwargs)
+
+  def warning(self, message: str, **kwargs):
+    """Logs a warning message to the session."""
+    self._log(pg.logging.warning, 'warning', message, **kwargs)
+
+  def error(self, message: str, **kwargs):
+    """Logs an error message to the session."""
+    self._log(pg.logging.error, 'error', message, **kwargs)
+
+  def fatal(self, message: str, **kwargs):
+    """Logs a fatal message to the session."""
+    # We use error level for fatal message, which does not trigger assertion.
+    self._log(pg.logging.error, 'fatal', message, **kwargs)
+
   #
   # HTML views.
   #
@@ -465,6 +511,25 @@ class Evaluation(experiment_lib.Experiment):
           )
       )
 
+    def _logs_tab() -> pg.views.html.controls.Tab:
+      """Renders a tab for the logs of the evaluation."""
+      with self._log_lock:
+        log_history = '\n'.join(str(l) for l in self._log_entries)
+      return pg.views.html.controls.Tab(
+          label='Logs',
+          content=pg.Html.element(
+              'div',
+              [
+                  pg.Html.element(
+                      'textarea',
+                      [pg.Html.escape(log_history)],
+                      readonly=True,
+                      css_classes=['logs-textarea'],
+                  )
+              ]
+          )
+      )
+
     def _main_tabs() -> pg.Html:
       return pg.Html.element(
           'div',
@@ -474,6 +539,8 @@ class Evaluation(experiment_lib.Experiment):
                       _definition_tab(),
                   ] + [
                       _metric_tab(m) for m in self.metrics
+                  ] + [
+                      _logs_tab()
                   ],
                   selected=1,
               )
@@ -593,6 +660,14 @@ class Evaluation(experiment_lib.Experiment):
           width:100%;
           height:100%;
         }
+        .logs-textarea {
+          width: 100%;
+          height: 500px;
+          padding: 5px;
+          border: 1px solid #DDD;
+          background-color: #EEE;
+          resize: vertical;
+        }
         """
     ]
 
@@ -615,6 +690,11 @@ class EvaluationState:
         assert isinstance(example, example_lib.Example), example
         self._evaluated_examples[example.id] = example
 
+  @property
+  def evaluated_examples(self) -> dict[int, example_lib.Example]:
+    """Returns the examples in the state."""
+    return self._evaluated_examples
+
   def get(self, example_id: int) -> example_lib.Example | None:
     """Returns the example with the given ID."""
     return self._evaluated_examples.get(example_id)
@@ -622,9 +702,3 @@ class EvaluationState:
   def update(self, example: example_lib.Example) -> None:
     """Updates the state with the given example."""
     self._evaluated_examples[example.id] = example
-
-  @property
-  def evaluated_examples(self) -> dict[int, example_lib.Example]:
-    """Returns the examples in the state."""
-    return self._evaluated_examples
-

{langfun-0.1.2.dev202412180804 → langfun-0.1.2.dev202412270804}/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,13 @@ 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)
+    exp.error('error message', x=1)
+    exp.fatal('fatal message')
+
     self.assertIn(
         exp.id,
         exp.to_html(extra_flags=dict(card_view=True, current_run=None)).content

{langfun-0.1.2.dev202412180804 → langfun-0.1.2.dev202412270804}/langfun/core/eval/v2/experiment.py

@@ -81,7 +81,7 @@ class Experiment(lf.Component, pg.views.HtmlTreeView.Extension):
   directory (using the ID 'latest'). Users can specify 'new' to start a fresh
   run or provide a specific run ID (typically in the format %Y%m%d_%<number>).
   Additionally, when initiating a new run, users may specify a `warm_start_from`
-  ID to restore the experiment’s state from a previous run.
+  directory to restore the experiment’s state from a previous run.
 
   Examples:
 
@@ -97,9 +97,9 @@ class Experiment(lf.Component, pg.views.HtmlTreeView.Extension):
     # Start a new, clean run.
     experiment.run(root_dir, 'new')
 
-    # Start a new run with a warm start from the previous run located in
-    # 'run_20241031_1' of the root directory.
-    experiment.run(root_dir, 'new', warm_start_from='20241031_1')
+    # Start a new run with a warm start from the another run located at
+    # '/path/to/another/run' (e.g. /my_expreriment/run_20241031_1).
+    experiment.run(root_dir, 'new', warm_start_from='/path/to/another/run')
 
     # Resume run '20241031_1', re-running failed examples and recomputing
     # metrics as needed.
@@ -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.dev202412180804 → langfun-0.1.2.dev202412270804}/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=[])