[logtree for RLTestEvaluator] (#54)

Author: ruiqi-zhong

tinker_cookbook/rl/metric_util.py

@@ -10,6 +10,7 @@
 from tinker_cookbook.rl.rollouts import do_group_rollout
 from tinker_cookbook.rl.types import EnvGroupBuilder, RLDataset, TrajectoryGroup
 from tinker_cookbook.utils.misc_utils import all_same, dict_mean
+from tinker_cookbook.utils import logtree
 
 
 def _compute_by_group_metrics(trajectory_groups_P: List[TrajectoryGroup], good_thresh: float = 0.5):
@@ -102,15 +103,29 @@ def dataset_to_env_group_builders(dataset: RLDataset) -> list[EnvGroupBuilder]:
 
 
 class RLTestSetEvaluator(SamplingClientEvaluator):
-    def __init__(self, dataset: RLDataset, max_tokens: int, name: str | None = None):
+    def __init__(
+        self,
+        dataset: RLDataset,
+        max_tokens: int,
+        name: str | None = None,
+        num_groups_to_log: int = 4,
+        log_path: str | None = None,
+    ):
         self.env_group_builders_P = dataset_to_env_group_builders(dataset)
         self.max_tokens = max_tokens
         self.name = name
+        self.num_groups_to_log = num_groups_to_log
 
     async def __call__(self, sampling_client: tinker.SamplingClient) -> dict[str, float]:
         policy = TinkerTokenCompleter(sampling_client, max_tokens=self.max_tokens)
+
+        async def run_group_rollout(builder, i):
+            enable_logging = i < self.num_groups_to_log
+            with logtree.optional_enable_logging(enable=enable_logging):
+                return await do_group_rollout(builder, policy)
+
         trajectory_groups_P = await asyncio.gather(
-            *[do_group_rollout(builder, policy) for builder in self.env_group_builders_P]
+            *[run_group_rollout(builder, i) for i, builder in enumerate(self.env_group_builders_P)]
         )
         taglist_P = [builder.logging_tags() for builder in self.env_group_builders_P]
         metrics = compute_trajectory_metrics(trajectory_groups_P, taglist_P)

tinker_cookbook/rl/train.py

@@ -7,8 +7,7 @@
 import logging
 import os
 import time
-from contextlib import nullcontext
-from typing import Any, Callable, List, Literal, Sequence
+from typing import Any, Callable, List, Literal, Sequence, Iterator
 
 import chz
 import numpy as np
@@ -41,10 +40,24 @@
 from tinker_cookbook.utils import logtree, ml_log
 from tinker_cookbook.utils.misc_utils import safezip, split_list, timed
 from tinker_cookbook.utils.trace import scope, trace_init, get_scope_context
+from contextlib import contextmanager
+
 
 logger = logging.getLogger(__name__)
 
 
+@contextmanager
+def get_logtree_scope(
+    log_path: str | None, num_groups_to_log: int, f_name: str, scope_name: str
+) -> Iterator[None]:
+    if log_path is not None and num_groups_to_log > 0:
+        logtree_path = os.path.join(log_path, f"{f_name}.html")
+        with logtree.init_trace(scope_name, path=logtree_path):
+            yield
+    else:
+        yield
+
+
 @scope
 def _select_representative_inds(scores: list[float], num_inds: int) -> list[int]:
     assert num_inds <= len(scores)
@@ -272,19 +285,25 @@ async def do_sync_training_with_stream_minibatch(
         if (cfg.eval_every > 0 and i_batch % cfg.eval_every == 0) or i_batch == end_batch - 1:
             with timed("run_evals", metrics):
                 for evaluator in evaluators:
-                    eval_metrics = await evaluator(sampling_client)
-                    metrics.update({f"test/{k}": v for k, v in eval_metrics.items()})
+                    ev_name = (
+                        evaluator.name
+                        if isinstance(evaluator, RLTestSetEvaluator) and evaluator.name is not None
+                        else ""
+                    )
+                    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)
+                        metrics.update({f"test/{k}": v for k, v in eval_metrics.items()})
 
-        # Initialize logtree trace for this iteration if logging is enabled
-        logtree_path = (
-            os.path.join(cfg.log_path, f"iteration_{i_batch:06d}.html")
-            if cfg.num_groups_to_log > 0
-            else None
-        )
-        with (
-            logtree.init_trace(f"RL Iteration {i_batch}", path=logtree_path)
-            if logtree_path
-            else logtree.scope_disable()
+        with get_logtree_scope(
+            cfg.log_path,
+            cfg.num_groups_to_log,
+            f"train_iteration_{i_batch:06d}",
+            f"RL Iteration {i_batch}",
         ):
             # Samplers will produce trajectory groups asynchronously,
             # and the trainer will consume them as soon as they are ready
@@ -598,7 +617,8 @@ async def do_group_rollout_and_filter_constant_reward(
     enable_logging: bool = True,
 ) -> TrajectoryGroup | None:
     policy = TinkerTokenCompleter(sampling_client, max_tokens=max_tokens)
-    with nullcontext() if enable_logging else logtree.scope_disable():
+
+    with logtree.optional_enable_logging(enable_logging):
         trajectory_group = await do_group_rollout(env_group_builder, policy)
 
     # Remove if all trajectories have the same reward
@@ -900,23 +920,29 @@ async def do_sync_training(
         if cfg.eval_every > 0 and i_batch % cfg.eval_every == 0:
             with timed("run_evals", metrics):
                 for evaluator in evaluators:
-                    eval_metrics = await evaluator(sampling_client)
-                    metrics.update({f"test/{k}": v for k, v in eval_metrics.items()})
+                    ev_name = (
+                        evaluator.name
+                        if isinstance(evaluator, RLTestSetEvaluator) and evaluator.name is not None
+                        else ""
+                    )
+                    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)
+                        metrics.update({f"test/{k}": v for k, v in eval_metrics.items()})
 
         # Get batch and sample trajectories
         env_group_builders_P = dataset.get_batch(i_batch)
 
         # Initialize logtree trace for this iteration if logging is enabled
-        logtree_path = (
-            os.path.join(cfg.log_path, f"iteration_{i_batch:06d}.html")
-            if cfg.num_groups_to_log > 0
-            else None
-        )
-        with (
-            logtree.init_trace(f"RL Iteration {i_batch}", path=logtree_path)
-            if logtree_path
-            else logtree.scope_disable(),
-            timed("sample", metrics),
+        with get_logtree_scope(
+            log_path=cfg.log_path,
+            num_groups_to_log=cfg.num_groups_to_log,
+            f_name=f"train_iteration_{i_batch:06d}",
+            scope_name=f"RL Iteration {i_batch}",
         ):
             trajectory_groups_P = await asyncio.gather(
                 *[

tinker_cookbook/tests/test_logtree.py

@@ -175,26 +175,19 @@ def simple_function():
     def custom_title_function():
         logtree.log_text("Inside custom title function")
 
-    @logtree.scope_header_decorator(lambda x: f"Processing {x}")
-    def dynamic_title_function(x):
-        logtree.log_text(f"Value: {x}")
-
     with tempfile.TemporaryDirectory() as tmpdir:
         output_path = Path(tmpdir) / "decorator.html"
 
         with logtree.init_trace("Decorator Test", path=output_path):
             simple_function()
             custom_title_function()
-            dynamic_title_function(42)
 
         content = output_path.read_text()
 
         assert "simple_function" in content
         assert "Inside simple function" in content
         assert "Custom Title" in content
         assert "Inside custom title function" in content
-        assert "Processing 42" in content
-        assert "Value: 42" in content
 
 
 async def async_test_scope_header_decorator():

tinker_cookbook/utils/logtree.py

@@ -435,9 +435,6 @@ def init_trace(
         _current_trace.reset(tok_t)
 
 
-# Public API: Structure
-
-
 @contextmanager
 def scope_header(title: str, **attrs: Any) -> Iterator[None]:
     """
@@ -480,51 +477,37 @@ def scope_header(title: str, **attrs: Any) -> Iterator[None]:
 F = TypeVar("F", bound=Callable[..., Any])
 
 
-# Overloads to support both bare decorator and parameterized usage
-# More specific overloads must come first
+# Overloads the parameterized usage
 @overload
-def scope_header_decorator(title: str, **attrs: Any) -> Callable[[F], F]: ...  # String title
-
-
-@overload
-def scope_header_decorator(
-    title: Callable[..., str], **attrs: Any
-) -> Callable[[F], F]: ...  # Lambda title
-
-
-@overload
-def scope_header_decorator(title: None = None, **attrs: Any) -> Callable[[F], F]: ...  # No args
+def scope_header_decorator(title: str) -> Callable[[F], F]: ...  # String title
 
 
+# Overloads the bare usage
 @overload
 def scope_header_decorator(title: F) -> F: ...  # Bare: @scope_header_decorator
 
 
 def scope_header_decorator(
-    title: str | Callable[..., str] | F | None = None, **attrs: Any
+    title: str | F,
 ) -> F | Callable[[F], F]:
     """
     Decorator to wrap function in a scope_header.
 
     Args:
-        title: String, callable returning string, or None (use function name)
-        **attrs: HTML attributes
+        title: String or function returning string
 
     Examples:
         @logtree.scope_header_decorator
         async def process_batch():
             ...
 
-        @logtree.scope_header_decorator("Processing")
-        def work():
-            ...
-
-        @logtree.scope_header_decorator(lambda self, x: f"Item {x}")
-        def handle_item(self, x):
+        @logtree.scope_header_decorator("Handling item")
+        def handle_item():
             ...
     """
+    title_str = title if isinstance(title, str) else title.__name__
 
-    def _wrap(fn: F, title_fn: Callable[..., str]) -> F:
+    def _wrap(fn: F) -> F:
         if inspect.iscoroutinefunction(fn):
 
             @functools.wraps(fn)
@@ -533,7 +516,7 @@ async def aw(*args: Any, **kwargs: Any) -> Any:
                 if not _is_logging_enabled():
                     return await fn(*args, **kwargs)
 
-                with scope_header(title_fn(*args, **kwargs), **attrs):
+                with scope_header(title_str):
                     return await fn(*args, **kwargs)
 
             return aw  # type: ignore
@@ -545,38 +528,16 @@ def w(*args: Any, **kwargs: Any) -> Any:
                 if not _is_logging_enabled():
                     return fn(*args, **kwargs)
 
-                with scope_header(title_fn(*args, **kwargs), **attrs):
+                with scope_header(title_str):
                     return fn(*args, **kwargs)
 
             return w  # type: ignore
 
-    # Check if this is a bare decorator (no arguments)
-    # When used as @scope_header_decorator, title will be the decorated function.
-    # When used as @scope_header_decorator("string") or @scope_header_decorator(lambda ...),
-    # title will be a string or lambda, and this returns a decorator function.
-    # We distinguish by checking if it's a function but NOT a lambda.
-    if inspect.isfunction(title) and title.__name__ != "<lambda>" and not attrs:
-        # Bare decoration: @scope_header_decorator
+    if isinstance(title, str):
+        return _wrap
+    else:
         fn = title
-        return _wrap(fn, lambda *_args, **_kwargs: fn.__name__)  # type: ignore[arg-type]
-
-    # Parameterized decorator
-    def deco(fn: F) -> F:
-        if title is None:
-
-            def title_fn(*_args: Any, **_kwargs: Any) -> str:
-                return fn.__name__
-
-        elif callable(title):
-            title_fn = title
-        else:
-
-            def title_fn(*_args: Any, **_kwargs: Any) -> str:
-                return str(title)
-
-        return _wrap(fn, title_fn)
-
-    return deco
+        return _wrap(fn)
 
 
 @contextmanager
@@ -618,6 +579,16 @@ def scope_disable() -> Iterator[None]:
         _logging_disabled.reset(token)
 
 
+@contextmanager
+def optional_enable_logging(enable: bool) -> Iterator[None]:
+    """Context manager to optionally enable logging."""
+    if enable:
+        yield
+    else:
+        with scope_disable():
+            yield
+
+
 @contextmanager
 def scope_details(summary: str) -> Iterator[None]:
     """