[Tracing the single evaluation name] (#67)

Author: ruiqi-zhong

tinker_cookbook/rl/train.py

@@ -262,6 +262,19 @@ class Config:
     num_groups_to_log: int = 4  # Number of groups to log per iteration (0 = disable logging)
 
 
+@scope
+async def run_single_evaluation(evaluator, cfg, i_batch, sampling_client):
+    ev_name = _get_evaluator_name(evaluator)
+    with _get_logtree_scope(
+        log_path=cfg.log_path,
+        num_groups_to_log=cfg.num_groups_to_log,
+        f_name=f"eval_{ev_name}_iteration_{i_batch:06d}",
+        scope_name=f"Running evaluation {ev_name} {i_batch}",
+    ):
+        eval_metrics = await evaluator(sampling_client)
+        return {f"test/{k}": v for k, v in eval_metrics.items()}
+
+
 @scope
 async def run_evaluations_parallel(
     evaluators: list[SamplingClientEvaluator],
@@ -271,17 +284,6 @@ async def run_evaluations_parallel(
 ) -> dict[str, Any]:
     """Run all evaluators in parallel and return aggregated metrics."""
 
-    async def run_single_evaluation(evaluator, cfg, i_batch, sampling_client):
-        ev_name = _get_evaluator_name(evaluator)
-        with _get_logtree_scope(
-            log_path=cfg.log_path,
-            num_groups_to_log=cfg.num_groups_to_log,
-            f_name=f"eval_{ev_name}_iteration_{i_batch:06d}",
-            scope_name=f"Running evaluation {ev_name} {i_batch}",
-        ):
-            eval_metrics = await evaluator(sampling_client)
-            return {f"test/{k}": v for k, v in eval_metrics.items()}
-
     # Create tasks for all evaluators with names for better traceability
     tasks = []
     for i, evaluator in enumerate(evaluators):