fix logging too much

Author: Dylan Huang

eval_protocol/pytest/evaluation_test.py

@@ -374,7 +374,6 @@ def _log_eval_error(
                         # has to be done in the pytest main process since it's
                         # used to determine whether this eval has stopped
                         row.pid = os.getpid()
-                        active_logger.log(row)
 
                     # Prepare rollout processor config once; we will generate fresh outputs per run
                     config = RolloutProcessorConfig(
@@ -401,6 +400,10 @@ def _log_eval_error(
                         for row in fresh_dataset:
                             row.rollout_id = generate_id()
 
+                        # log the fresh_dataset
+                        for row in fresh_dataset:
+                            active_logger.log(row)
+
                         processed_dataset = execute_function(rollout_processor, rows=fresh_dataset, config=config)
 
                         if mode == "pointwise":

tests/pytest/test_pytest_ids.py

@@ -1,33 +1,67 @@
-import eval_protocol.pytest.evaluation_test as evaluation_test_module
+from typing import List
+
+import eval_protocol.dataset_logger as dataset_logger
+from eval_protocol.dataset_logger.dataset_logger import DatasetLogger
 from eval_protocol.models import EvaluationRow
 from eval_protocol.pytest.default_no_op_rollout_process import default_no_op_rollout_processor
-from eval_protocol.pytest.evaluation_test import evaluation_test as evaluation_decorator
 from tests.pytest.test_markdown_highlighting import markdown_dataset_to_evaluation_row
 
 
-class InMemoryLogger:
+class InMemoryLogger(DatasetLogger):
     def __init__(self):
-        self._rows = []
+        self._rows: dict[str, EvaluationRow] = {}
 
-    def log(self, row):
-        self._rows.append(row)
+    def log(self, row: EvaluationRow):
+        print(row.run_id, row.rollout_id)
+        self._rows[row.rollout_id] = row
 
     def read(self):
-        return list(self._rows)
+        return list(self._rows.values())
+
+
+def test_evaluation_test_decorator(monkeypatch):
+    from eval_protocol.pytest.evaluation_test import evaluation_test
+
+    logger = InMemoryLogger()
+
+    @evaluation_test(
+        input_dataset=[
+            "tests/pytest/data/markdown_dataset.jsonl",
+        ],
+        model=["dummy/local-model"],
+        dataset_adapter=markdown_dataset_to_evaluation_row,
+        rollout_processor=default_no_op_rollout_processor,
+        mode="pointwise",
+        combine_datasets=False,
+        num_runs=2,
+        logger=logger,
+    )
+    def eval_fn(row: EvaluationRow) -> EvaluationRow:
+        return row
+
+    dataset_paths = [
+        "tests/pytest/data/markdown_dataset.jsonl",
+    ]
+
+    # Manually invoke all parameter combinations within a single test
+    for ds_path in dataset_paths:
+        eval_fn(model="dummy/local-model", dataset_path=[ds_path])
+
+    # Assertions on IDs generated by the decorator logic
+    assert len(logger.read()) == 38
 
 
 def test_evaluation_test_decorator_ids_single(monkeypatch):
-    # Use an in-memory logger to avoid sqlite side effects
     in_memory_logger = InMemoryLogger()
-    monkeypatch.setattr(evaluation_test_module, "default_logger", in_memory_logger, raising=False)
-
     unique_run_ids = set()
     unique_cohort_ids = set()
     unique_rollout_ids = set()
     unique_invocation_ids = set()
     unique_row_ids = set()
 
-    @evaluation_decorator(
+    from eval_protocol.pytest.evaluation_test import evaluation_test
+
+    @evaluation_test(
         input_dataset=[
             "tests/pytest/data/markdown_dataset.jsonl",
             "tests/pytest/data/markdown_dataset.jsonl",
@@ -39,6 +73,7 @@ def test_evaluation_test_decorator_ids_single(monkeypatch):
         mode="pointwise",
         combine_datasets=False,
         num_runs=5,
+        logger=InMemoryLogger(),
     )
     def eval_fn(row: EvaluationRow) -> EvaluationRow:
         unique_run_ids.add(row.run_id)