langfun 0.1.2.dev202509120804__py3-none-any.whl → 0.1.2.dev202512150805__py3-none-any.whl

langfun/core/eval/v2/evaluation.py

@@ -32,17 +32,63 @@ import pyglove as pg
 
 
 class Evaluation(experiment_lib.Experiment):
-  """Evaluation.
-
-  An evaluation can be a leaf node or a container of other evaluations,
-  depending on whether the current evaluation object is configured with
-  any `pg.oneof`.
-
-  For example, `MyEval(lm=pg.oneof([lf.llms.Gpt4(), lf.llms.Gemini1_5Pro()]))`
-  is a container of two sub-experiments, one for each LLM. In such case, the
-  evaluation object with `pg.oneof` is called a hyper evaluation, which
-  represents a search space of evaluations, and each sub-evaluation is called
-  a leaf evaluation, which will perform the actual evaluation.
+  """Base class for Langfun evaluations.
+
+  `lf.eval.Evaluation` is the base class for defining evaluation tasks in
+  Langfun. Users typically subclass it to implement custom evaluation logic by
+  overriding `inputs` and `process` methods.
+
+  An `Evaluation` object encapsulates:
+
+  *   **`inputs`**: A callable that returns an iterable of input examples to be
+      processed. This is usually provided by implementing an `inputs(self)`
+      method in the subclass, which yields input items for evaluation one by
+      one.
+  *   **`process(self, example)`**: An abstract method that processes one
+      example and returns the output, or a tuple of (output, metadata).
+      The output will be used for computing metrics.
+  *   **`metrics`**: A list of metrics (e.g., `lf.metrics.Accuracy`) to compute
+      based on the outputs from `process`. Some metrics may require users to
+      implement a `ground_truth(self, example)` method in the subclass to
+      compute metrics against ground truth.
+  *   **Hyperparameters**: Any other attributes of the class serve as
+      hyperparameters for the evaluation (e.g., the language model to use).
+
+  **Running Evaluations:**
+
+  Evaluations are executed via `lf.eval.Suite` or by calling the `.run()`
+  method on an `Evaluation` instance, which returns a `Run` object
+  containing the evaluation run information and results. If an evaluation
+  contains sweeable parameters (using `pg.oneof`), `.run()` will expand it
+  into multiple evaluation sub-tasks -- one for each combination of
+  hyperparameters -- all managed within the same `Run`.
+
+  **Example:**
+
+  ```python
+  import langfun as lf
+  import pyglove as pg
+
+  class MyEval(lf.eval.Evaluation):
+    lm: lf.LanguageModel
+    prompt: str = '1 + 1 = '
+
+    def inputs(self):
+      yield 2
+
+    def process(self, example: lf.eval.Example):
+      return int(lf.query(self.prompt, lm=self.lm))
+
+    def ground_truth(self, example: lf.eval.Example) -> int:
+      return example.input
+
+  # Run evaluation using two different LMs
+  evaluation = MyEval(
+      lm=pg.oneof([lf.llms.Gpt4(), lf.llms.Gemini()]),
+      metrics=[lf.metrics.Accuracy()]
+  )
+  run_info = evaluation.run()
+  ```
   """
 
   inputs: Annotated[
@@ -68,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).
   #
@@ -126,6 +179,20 @@ class Evaluation(experiment_lib.Experiment):
   # Evaluation logics.
   #
 
+  def setup(self) -> None:
+    """Sets up resources required by the evaluation.
+
+    Subclasses should always call the super().setup() method to ensure the
+    proper initialization of the evaluation.
+    """
+
+  def teardown(self) -> None:
+    """Tears down resources used by the evaluation.
+
+    Subclasses should always call the super().teardown() method to ensure the
+    proper cleanup of the evaluation.
+    """
+
   @abc.abstractmethod
   def process(
       self,
@@ -137,7 +204,7 @@ class Evaluation(experiment_lib.Experiment):
 
     Args:
       example: An example object to process. `example.input` is an object
-        returned from `Evaluable.inputs`.
+        yielded from `inputs()` method.
 
     Returns:
       A processed output. Or a tuple of (output, metadata).
@@ -150,6 +217,7 @@ class Evaluation(experiment_lib.Experiment):
       example: example_lib.Example | int,
       raise_if_has_error: bool = False,
       reevaluate_upon_previous_errors: bool = True,
+      force_recompute_metrics: bool = False
   ) -> example_lib.Example:
     """Evaluates a single example input.
 
@@ -158,6 +226,8 @@ class Evaluation(experiment_lib.Experiment):
       raise_if_has_error: Whether to raise an error if the example has error.
       reevaluate_upon_previous_errors: Whether to reevaluate the example if
         the previous checkpointed run has error.
+      force_recompute_metrics: If True, force recompute the metrics even if
+        metric metadata is already present from previous checkpoint.
 
     Returns:
       The evaluated example with the output and metric metadata populated.
@@ -206,6 +276,7 @@ class Evaluation(experiment_lib.Experiment):
         # Use the output and metadata obtained from the previous processing.
         example.output = checkpointed.output
         example.metadata = checkpointed.metadata
+        example.metric_metadata = checkpointed.metric_metadata
         example.error = checkpointed.error
         example.newly_processed = False
         example.execution_status = checkpointed.execution_status
@@ -225,8 +296,16 @@ class Evaluation(experiment_lib.Experiment):
         self.info(f'Starting metric computation for example {example.id}.')
         metric_metadata = {}
         for metric in self.metrics:
-          metric_metadata.update(metric.audit(example))
-        example.metric_metadata = metric_metadata
+          metric_metadata[metric.name] = metric.update(
+              example, force_recompute=force_recompute_metrics
+          )
+
+        if example.metric_metadata is None:
+          example.metric_metadata = metric_metadata
+        else:
+          # Accumulate the metric metadata as there might be existing metadata
+          # from previous metric computation runs.
+          example.metric_metadata.update(metric_metadata)
         self.info(f'Completed metric computation for example {example.id}.')
 
     # For previously processed examples, we keep the execution status for the
@@ -287,7 +366,7 @@ class Evaluation(experiment_lib.Experiment):
       A unique string representing the resource required.
     """
     return {
-        v.resource_id for _, v in self.sym_init_args.items()
+        v.resource_id for _, v in self.sym_init_args.sym_items()
         if isinstance(v, lf.LanguageModel)
     }
 
@@ -307,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,
@@ -318,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."""
@@ -760,7 +840,7 @@ class Evaluation(experiment_lib.Experiment):
 
 
 class EvaluationState:
-  """Evaluation state."""
+  """In-memory state of an evaluation."""
 
   class ExampleStatus(pg.Object):
     """Example state."""
@@ -808,8 +888,9 @@ class EvaluationState:
       load_example_metadata: bool | Callable[
           [example_lib.Example], bool] = True,
       filter: Callable[[example_lib.Example], bool] | None = None,  # pylint: disable=redefined-builtin
-  ) -> None:
+  ) -> list[example_lib.Example]:
     """Loads the state from the example sequence file."""
+    examples = []
     for example in example_lib.Example.iter_ckpts(
         state_file,
         example_input_by_id=example_input_by_id,
@@ -819,6 +900,8 @@ class EvaluationState:
         continue
       example.newly_processed = False
       self._ckpt_examples[example.id] = example
+      examples.append(example)
+    return examples
 
   @property
   def evaluation_status(self) -> dict[int, ExampleStatus]:

langfun/core/eval/v2/evaluation_test.py

@@ -88,7 +88,7 @@ class EvaluationTest(unittest.TestCase):
     self.assertEqual(example.output, 6)
     self.assertIsNone(example.error)
     self.assertEqual(example.metadata, {})
-    self.assertEqual(example.metric_metadata, dict(match=True))
+    self.assertEqual(example.metric_metadata, dict(match=dict(is_correct=True)))
     self.assertIsNotNone(example.usage_summary)
     self.assertGreater(example.usage_summary.total.total_tokens, 0)
     self.assertEqual(example.usage_summary.total.num_requests, 1)
@@ -103,7 +103,10 @@ class EvaluationTest(unittest.TestCase):
     self.assertEqual(example.output, 7)
     self.assertIsNone(example.error)
     self.assertEqual(example.metadata, {})
-    self.assertEqual(example.metric_metadata, dict(mismatch=True))
+    self.assertEqual(
+        example.metric_metadata,
+        dict(match=dict(is_correct=False))
+    )
 
     with self.assertRaisesRegex(ValueError, 'x should not be 5'):
       _ = exp.evaluate(6, raise_if_has_error=True)
@@ -113,7 +116,10 @@ class EvaluationTest(unittest.TestCase):
     self.assertEqual(pg.MISSING_VALUE, example.output)
     self.assertEqual(example.error.tag, 'ValueError')
     self.assertEqual(example.metadata, {})
-    self.assertEqual(example.metric_metadata, dict(error='ValueError'))
+    self.assertEqual(
+        example.metric_metadata,
+        dict(match=dict(error='ValueError'))
+    )
 
   def test_evaluate_withstate(self):
     eval_dir = os.path.join(tempfile.mkdtemp(), 'test_eval')

langfun/core/eval/v2/example.py

@@ -22,19 +22,30 @@ import pyglove as pg
 
 @dataclasses.dataclass
 class Example(pg.JSONConvertible, pg.views.HtmlTreeView.Extension):
-  """An item for the evaluation.
+  """An example for evaluation.
+
+  An evaluation example contains the input and output of an evaluation task,
+  as well as metadata about the evaluation process, such as execution time,
+  LLM usage, and metric results.
 
   Attributes:
-    id: The 1-based ID of the item in the evaluation set.
-    input: An element returned from the `Evaluable.inputs` functor.
-    output: The output of the `process` method. If `pg.MISSING_VALUE`, it has
-      not been processed yet.
-    metadata: The metadata of the item produced by the `process` method.
-    metric_metadata: The dictionary returned from `Metric.audit`.
-    start_time: The start time of the evaluation item.
-    end_time: The end time of the evaluation item.
-    usage_summary: The summary of LLM usages of the evaluation item.
-    execution_status: The timeit status of the evaluation item.
+    id: The 1-based ID of the example in the evaluation set.
+    input: An element returned from the `Evaluable.inputs` functor, which serves
+      as the input for `lf.Evaluable.process`.
+    output: The output of `lf.Evaluable.process` method. If `pg.MISSING_VALUE`,
+      it indicates the example has not been processed yet.
+    error: The error raised from `lf.Evaluable.process`. If None, it
+      indicates the process was successful.
+    metadata: The metadata of the example produced by `lf.Evaluable.process`.
+    metric_metadata: The dictionary returned from `Metric.audit`, which contains
+      metadata about metric computation for this example.
+    newly_processed: Whether this example is processed in the current run. If
+      False, it indicates the example was loaded from a checkpoint from previous
+      runs.
+    start_time: The start time of processing this example.
+    end_time: The end time of processing this example.
+    usage_summary: The summary of LLM usages for processing this example.
+    execution_status: The timeit status of processing this example.
   """
   id: int
   input: Any = pg.MISSING_VALUE
@@ -49,14 +60,6 @@ class Example(pg.JSONConvertible, pg.views.HtmlTreeView.Extension):
   usage_summary: lf.UsageSummary | None = None
   execution_status: dict[str, pg.utils.TimeIt.Status] | None = None
 
-  def __post_init__(self):
-    if self.execution_status is not None:
-      for status in self.execution_status.values():
-        if status.has_error:
-          assert isinstance(status.error, pg.ErrorInfo)
-          self.error = status.error
-          break
-
   @property
   def is_processed(self) -> bool:
     """Returns whether the item has been processed."""
@@ -152,6 +155,8 @@ class Example(pg.JSONConvertible, pg.views.HtmlTreeView.Extension):
       ckpt_file: str | list[str],
       example_input_by_id: Callable[[int], Any] | None = None,
       load_example_metadata: bool = True,
+      convert_unknown: bool = True,
+      **kwargs
   ) -> Iterator['Example']:
     """Iterates Examples from the checkpoint files."""
     ckpt_files = [ckpt_file] if isinstance(ckpt_file, str) else ckpt_file
@@ -161,7 +166,9 @@ class Example(pg.JSONConvertible, pg.views.HtmlTreeView.Extension):
           example = pg.from_json_str(
               record,
               example_input_by_id=example_input_by_id,
-              load_example_metadata=load_example_metadata
+              load_example_metadata=load_example_metadata,
+              convert_unknown=convert_unknown,
+              **kwargs
           )
           assert isinstance(example, cls), example
           yield example
@@ -182,15 +189,23 @@ class Example(pg.JSONConvertible, pg.views.HtmlTreeView.Extension):
     extra_flags = extra_flags or {}
     num_examples = extra_flags.get('num_examples', None)
 
-    def _metric_metadata_badge(key, value):
-      if isinstance(value, bool) and bool:
-        text = key
-      else:
-        text = f'{key}:{value}'
-      return pg.views.html.controls.Badge(
-          text,
-          css_classes=[pg.utils.camel_to_snake(key, '-')],
-      )
+    def _metric_label_group(metric_metadata: dict[str, Any] | None):
+      """Renders a label group for metric metadata."""
+      badges = []
+      if metric_metadata:
+        for metric_name, metadata in metric_metadata.items():
+          assert isinstance(metadata, dict), (metric_name, metadata)
+          for k, v in metadata.items():
+            css_class = k
+            if isinstance(v, bool):
+              css_class += '_true' if v else '_false'
+            badge = pg.views.html.controls.Badge(
+                f'{k}:{v}',
+                tooltip=f'{metric_name}: {k}',
+                css_classes=[css_class],
+            )
+            badges.append(badge)
+      return pg.views.html.controls.LabelGroup(badges)
 
     def _render_header():
       return pg.Html.element(
@@ -229,12 +244,7 @@ class Example(pg.JSONConvertible, pg.views.HtmlTreeView.Extension):
                           extra_flags=dict(as_badge=True)
                       ) if self.usage_summary is not None else None,
                       # Metric metadata.
-                      pg.views.html.controls.LabelGroup(
-                          [   # pylint: disable=g-long-ternary
-                              _metric_metadata_badge(k, v)
-                              for k, v in self.metric_metadata.items()
-                          ] if self.metric_metadata else []
-                      ),
+                      _metric_label_group(self.metric_metadata)
                   ],
                   css_classes=['example-container'],
               )
@@ -305,18 +315,18 @@ class Example(pg.JSONConvertible, pg.views.HtmlTreeView.Extension):
           color: black;
         }
         /* Badge styles. */
-        .eval-example .badge.match {
+        .eval-example .badge.is_correct_true {
           color: green;
           background-color: #dcefbe;
         }
+        .eval-example .badge.is_correct_false {
+          color: orange;
+          background-color: #ffefc4;
+        }
         .eval-example .badge.error {
           color: red;
           background-color: #fdcccc;
         }
-        .eval-example .badge.mismatch {
-          color: orange;
-          background-color: #ffefc4;
-        }
         .eval-example .badge.score {
           color: blue;
           background-color: #c4dced;

langfun/core/eval/v2/example_test.py

@@ -32,9 +32,9 @@ class ExampleTest(unittest.TestCase):
             name='evaluation', elapse=1.0, error=error
         )
     })
-    self.assertEqual(ex.error, error)
+    self.assertIsNone(ex.error)
     self.assertFalse(ex.is_processed)
-    self.assertTrue(ex.has_error)
+    self.assertFalse(ex.has_error)
     self.assertEqual(ex.elapse, 1.0)
 
     ex = Example(id=2, output=1)
@@ -94,15 +94,23 @@ class ExampleTest(unittest.TestCase):
     pg.JSONConvertible._TYPE_REGISTRY._type_to_cls_map.pop(
         inputs[0].b.__type_name__
     )
-    v = pg.from_json_str(json_str, auto_dict=True, load_example_metadata=True)
-    v.output.pop('type_name')
-    v.metadata.b.pop('type_name')
+    v = pg.from_json_str(
+        json_str,
+        convert_unknown=True,
+        load_example_metadata=True
+    )
     self.assertEqual(
         v,
         Example(
             id=1,
-            output=pg.Dict(x=1),
-            metadata=dict(b=pg.Dict(x=1, y=2)),
+            output=pg.symbolic.UnknownTypedObject(
+                inputs[0].a.__type_name__, x=1
+            ),
+            metadata=dict(
+                b=pg.symbolic.UnknownTypedObject(
+                    inputs[0].b.__type_name__, x=1, y=2
+                )
+            ),
         )
     )
     # Serialize with input.
@@ -116,7 +124,7 @@ class ExampleTest(unittest.TestCase):
         input=pg.Dict(a=1, b=2),
         output=3,
         metadata=dict(sum=3),
-        metric_metadata=dict(match=True),
+        metric_metadata=dict(match=dict(match=True)),
     )
     self.assertNotIn(
         'next',

langfun/core/eval/v2/experiment.py

@@ -139,10 +139,10 @@ class Experiment(lf.Component, pg.views.HtmlTreeView.Extension):
 
   # Checkpointing
 
-  Experiments support checkpointing, which is enabled by default. It allows 
+  Experiments support checkpointing, which is enabled by default. It allows
   users to resume their experiments from a saved state. When an experiment runs,
-  it creates a new directory for that run and saves the current state to a
-  checkpoint file. If the experiment is interrupted or fails, users can resume
+  it creates a new directory for that run and saves its progress to checkpoint
+  files. If the experiment is interrupted or fails, users can resume
   it by specifying the 'id' or 'warm_start_from' argument (shown above) to
   seamlessly continue from previously saved state without starting over.
 
@@ -169,7 +169,7 @@ class Experiment(lf.Component, pg.views.HtmlTreeView.Extension):
 
   # Experiment Plugins
 
-  Experiment can be extended by plugins. Plugins can listen to the events of
+  Experiments can be extended by plugins. Plugins can listen to the events of
   experiment execution and produce additional outputs. For example, a plugin
   can be added to an experiment to generate additional metrics or to save
   additional data to a database. More details will be added in the future.
@@ -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:
@@ -657,7 +657,30 @@ class Experiment(lf.Component, pg.views.HtmlTreeView.Extension):
 
 @pg.use_init_args(['children'])
 class Suite(Experiment):
-  """A suite of evaluations."""
+  """A suite of evaluations.
+
+  `lf.eval.Suite` groups multiple `lf.eval.Evaluation` or other `Suite`
+  objects into a single experiment, allowing them to be run, managed, and
+  reported together.
+
+  **Example:**
+
+  ```python
+  import langfun as lf
+
+  suite = lf.eval.Suite([
+      MyEval(lm=lf.llms.Gpt4()),
+      MyEval(lm=lf.llms.Gemini()),
+      lf.eval.Suite([
+          AnotherEval(lm=lf.llms.Gpt4()),
+          AnotherEval(lm=lf.llms.Gemini())
+      ])
+  ])
+
+  # Run all evaluations in the suite
+  run_info = suite.run('/path/to/my/suite_run')
+  ```
+  """
 
   children: Annotated[
       list[Experiment], 'A list of child experiments.'
@@ -668,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."""
@@ -791,7 +820,14 @@ class RunId(pg.Object):
 
 
 class Run(pg.Object, pg.views.html.HtmlTreeView.Extension):
-  """A run of an experiment."""
+  """Represents a single run of an experiment.
+
+  A `Run` object holds all the configurations for executing an experiment,
+  such as the experiment definition, input/output directories, and flags
+  controlling the execution behavior (e.g., error handling, checkpointing).
+  It also provides utility methods for accessing run-specific paths and
+  filtering examples for evaluation.
+  """
 
   root_dir: Annotated[
       str,
@@ -818,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
 
@@ -937,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."""
@@ -971,7 +1010,13 @@ class Run(pg.Object, pg.views.html.HtmlTreeView.Extension):
 
 
 class Runner(pg.Object):
-  """Interface for experiment runner."""
+  """Interface for experiment runner.
+
+  A runner is responsible for executing the evaluations within an experiment
+  based on the configuration specified in a `Run` object. Different runners
+  can implement different execution strategies, such as sequential or parallel
+  processing of examples and evaluations.
+  """
 
   # Class-level variable for registering the runner.
   NAME = None
@@ -1010,7 +1055,37 @@ class Runner(pg.Object):
 
 
 class Plugin(lf.Component):
-  """Base class for experiment plugins."""
+  """Base class for experiment plugins.
+
+  Plugins provide a mechanism to extend the behavior of an experiment run
+  by hooking into various events during the lifecycle of experiment and
+  example execution, such as `on_run_start`, `on_experiment_complete`,
+  `on_example_start`, etc. They can be used for custom logging, monitoring,
+  or result processing.
+  """
+
+  @classmethod
+  def is_per_example(cls) -> bool:
+    """Returns whether the plugin is per example only.
+
+    Per-example plugins can be installed on individual workers when examples
+    are evaluated by multiple processes in parallel.
+    """
+
+    def same_code(method1, method2):
+      return method1.__code__ == method2.__code__
+    return all(
+        same_code(method1, method2)
+        for method1, method2 in [
+            (Plugin.on_run_start, cls.on_run_start),
+            (Plugin.on_run_complete, cls.on_run_complete),
+            (Plugin.on_run_abort, cls.on_run_abort),
+            (Plugin.on_experiment_start, cls.on_experiment_start),
+            (Plugin.on_experiment_skipped, cls.on_experiment_skipped),
+            (Plugin.on_experiment_complete, cls.on_experiment_complete),
+            (Plugin.on_experiment_abort, cls.on_experiment_abort),
+        ]
+    )
 
   def on_run_start(
       self,

langfun/core/eval/v2/experiment_test.py

@@ -433,5 +433,24 @@ class RunnerTest(unittest.TestCase):
           pass
 
 
+class PluginTest(unittest.TestCase):
+
+  def test_per_example_only(self):
+
+    class PerExamplePlugin(experiment_lib.Plugin):
+
+      def on_example_complete(self, runner, experiment, example):
+        print('on_example_complete')
+
+    self.assertTrue(PerExamplePlugin.is_per_example())
+
+    class NonPerExamplePlugin(experiment_lib.Plugin):
+
+      def on_experiment_complete(self, runner, experiment):
+        print('on_example_complete')
+
+    self.assertFalse(NonPerExamplePlugin.is_per_example())
+
+
 if __name__ == '__main__':
   unittest.main()