configure monitor via CLI (chz)

Author: thesofakillers

project/paperbench/paperbench/monitor/create_monitor.py

@@ -1,10 +1,13 @@
+from __future__ import annotations
+
 import re
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import Any
 
 import blobfile as bf
 import structlog.stdlib
+from pydantic import BaseModel
 from unidecode import unidecode
 
 from paperbench.paper_registry import Paper
@@ -51,6 +54,14 @@ def to_dict(self) -> dict[str, Any]:
 class Monitor(ABC):
     """Base class for monitoring agent behavior through logs."""
 
+    class Config(BaseModel, ABC):
+        """Serializable configuration for a :class:`Monitor`."""
+
+        @abstractmethod
+        def build(self, paper: Paper) -> Monitor:
+            """Instantiate the monitor for the provided paper."""
+            ...
+
     def __init__(
         self,
         paper: Paper,
@@ -74,12 +85,18 @@ def __init__(
     @abstractmethod
     def check_log(self, log_file: str) -> MonitorResult:
         """Check a log file for violations of monitoring rules."""
-        raise NotImplementedError()
+        ...
 
 
 class BasicMonitor(Monitor):
     """Simple implementation that checks for occurrences of blacklisted terms with git clone, curl, or wget commands in agent logs."""
 
+    class Config(Monitor.Config):
+        """Configuration for :class:`BasicMonitor`."""
+
+        def build(self, paper: Paper) -> BasicMonitor:
+            return BasicMonitor(paper=paper)
+
     def _normalize_url(self, url: str) -> str:
         """Normalize URL by removing protocol, parameters, and anchors."""
         # Remove protocol (http:// or https://)

project/paperbench/paperbench/monitor/monitor.py

@@ -18,6 +18,7 @@
 from nanoeval.solvers.computer_tasks.steps import FinalResult
 from nanoeval.solvers.computer_tasks.task import ComputerTask
 from paperbench.metrics import compute_agg_stats, per_paper_results
+from paperbench.monitor.monitor import BasicMonitor, Monitor
 from paperbench.nano.structs import (
     JudgeConfig,
     PaperBenchGrade,
@@ -48,6 +49,7 @@
 class PaperBench(PythonCodingEval):
     reproduction: ReproductionConfig = chz.field(default_factory=ReproductionConfig)
     judge: JudgeConfig = chz.field(default_factory=JudgeConfig)
+    monitor_config: Monitor.Config = chz.field(default_factory=BasicMonitor.Config)
 
     # task args
     paper_split: Literal["debug", "dev", "human", "testing", "all"] = chz.field(
@@ -153,6 +155,7 @@ async def get_instances(self) -> list[PBTask]:
                         target_duration_hr=self.target_duration_hr,
                         judge=self.judge,
                         reproduction=self.reproduction,
+                        monitor_config=self.monitor_config,
                         save_cluster_output_to_host=self.save_cluster_output_to_host,
                         network_mode=NetworkMode.UNPROXIED
                         if self.allow_internet

project/paperbench/paperbench/nano/eval.py

@@ -38,8 +38,7 @@
     WORKSPACE_BASE,
 )
 from paperbench.grade import JudgeOutput, grade_submission
-from paperbench.monitor.create_monitor import create_monitor
-from paperbench.monitor.monitor import MonitorResult
+from paperbench.monitor.monitor import Monitor, MonitorResult
 from paperbench.nano.structs import (
     JudgeConfig,
     PaperBenchGrade,
@@ -68,6 +67,7 @@ class PBTask(ComputerTask):
     target_duration_hr: int | None
     reproduction: ReproductionConfig
     judge: JudgeConfig
+    monitor_config: Monitor.Config
     skipped_rollout: bool = False  # whether rollouts were skipped (e.g. if we're resuming)
 
     save_cluster_output_to_host: bool
@@ -334,10 +334,7 @@ def _should_monitor(self, log_file_path: str) -> bool:
             return False
 
     def _run_monitor(self, log_file_path: str) -> MonitorResult:
-        """
-        Runs the monitor on an given log file
-        TODO: make this configurable through chz in `PaperBenchEval`
-        """
+        """Run the configured monitor on the given log file."""
         ctx_logger = logger.bind(
             run_group_id=self.run_group_id, run_id=self.run_id, runs_dir=self.runs_dir
         )
@@ -346,11 +343,7 @@ def _run_monitor(self, log_file_path: str) -> MonitorResult:
             f"Running monitor on {self.run_id} agent.log", destinations=["run"], _print=True
         )
         paper = paper_registry.get_paper(self.paper_id)
-        monitor = create_monitor(
-            monitor_type="basic",
-            paper=paper,
-            monitor_kwargs={},
-        )
+        monitor = self.monitor_config.build(paper=paper)
         monitor_result = monitor.check_log(log_file_path)
         return monitor_result
 

project/paperbench/paperbench/nano/task.py

@@ -1,4 +1,5 @@
-import argparse
+from __future__ import annotations
+
 import asyncio
 import datetime
 import json
@@ -11,20 +12,46 @@
 import structlog.stdlib
 from tqdm.asyncio import tqdm_asyncio
 
-from paperbench.monitor.create_monitor import create_monitor
+import chz
+from paperbench.monitor.monitor import BasicMonitor, Monitor
 from paperbench.paper_registry import paper_registry
 
 logger = structlog.stdlib.get_logger(component=__name__)
 
 
+def _describe_monitor_config(monitor_config: Monitor.Config) -> str:
+    return f"{monitor_config.__class__.__module__}.{monitor_config.__class__.__qualname__}"
+
+
+@chz.chz
+class MonitorCLIArgs:
+    """Monitor agent logs for violations."""
+
+    logs_dir: Path = chz.field(
+        doc="Directory containing multiple run groups.",
+    )
+    run_groups: list[str] = chz.field(
+        default_factory=list,
+        doc="List of run group IDs to monitor.",
+    )
+    monitor_config: Monitor.Config = chz.field(
+        default_factory=BasicMonitor.Config,
+        doc="Specify the monitor to use (default: BasicMonitor).",
+    )
+    out_dir: Path | None = chz.field(
+        default=None,
+        doc="Directory to save the monitor results JSON file (default: current directory).",
+    )
+
+
 def get_paper_id_from_run_id(run_id: str) -> str:
     """Extract paper ID from run ID (e.g. 'rice_508398cb-0825-4bf0-b647-a9200ac03d21' -> 'rice')"""
     return run_id.split("_")[0]
 
 
 async def monitor_single_log(
     run_dir: Path,
-    monitor_type: str,
+    monitor_config: Monitor.Config,
 ) -> dict[str, Any] | None:
     """
     Monitor a single run's log with the specified monitor.
@@ -67,22 +94,24 @@ async def monitor_single_log(
         logger.warning(f"Log file not found at {log_file}")
         return None
 
-    logger.info(f"Running monitor on agent.log from {run_id}")
+    monitor_config_payload = monitor_config.model_dump(mode="json")
+    logger.info(
+        f"Running monitor on agent.log from {run_id}",
+        monitor=_describe_monitor_config(monitor_config),
+        monitor_config_json=json.dumps(monitor_config_payload, indent=2),
+    )
 
     # Create monitor
     paper = paper_registry.get_paper(paper_id)
-    monitor = create_monitor(
-        monitor_type=monitor_type,
-        paper=paper,
-        monitor_kwargs={},
-    )
+    monitor = monitor_config.build(paper=paper)
 
     # Run monitor on the log file
     result = await asyncio.to_thread(monitor.check_log, log_file.as_posix())
 
     return {
         "run_group_id": run_dir.parent.name,
-        "monitor_type": monitor_type,
+        "monitor_type": _describe_monitor_config(monitor_config),
+        "monitor_config": monitor_config_payload,
         "paper_id": paper_id,
         "log_file": str(log_file),
         "run_id": run_id,
@@ -103,7 +132,7 @@ async def monitor_single_log(
 
 async def monitor_run_group(
     group_dir: Path,
-    monitor_type: str,
+    monitor_config: Monitor.Config,
 ) -> list[dict[str, Any] | None]:
     """Monitor all runs in a run group directory."""
     run_group_id = group_dir.name
@@ -115,7 +144,7 @@ async def monitor_run_group(
     tasks = [
         monitor_single_log(
             run_dir=run_dir,
-            monitor_type=monitor_type,
+            monitor_config=monitor_config,
         )
         for run_dir in run_dirs
     ]
@@ -126,7 +155,7 @@ async def monitor_run_group(
 
 async def monitor_multiple_run_groups(
     logs_dir: Path,
-    monitor_type: str,
+    monitor_config: Monitor.Config,
     run_groups: list[str] | None = None,
 ) -> dict[str, Any] | None:
     """Run monitor on multiple run groups that are in a directory of run groups."""
@@ -154,7 +183,7 @@ async def monitor_multiple_run_groups(
     tasks = [
         monitor_run_group(
             group_dir=logs_dir / run_group_id,
-            monitor_type=monitor_type,
+            monitor_config=monitor_config,
         )
         for run_group_id in run_groups
     ]
@@ -169,10 +198,13 @@ async def monitor_multiple_run_groups(
     flagged_results = [result for result in all_results if len(result["results"]["violations"]) > 0]
     other_results = [result for result in all_results if len(result["results"]["violations"]) == 0]
 
+    monitor_config_payload = monitor_config.model_dump(mode="json")
+
     # Create final output with results and summary
     return {
         "timestamp": datetime.datetime.now().isoformat(),
-        "monitor_type": monitor_type,
+        "monitor_type": _describe_monitor_config(monitor_config),
+        "monitor_config": monitor_config_payload,
         "logs_dir": str(logs_dir.absolute()),
         "run_groups": run_groups,
         "total_runs": len(all_results),
@@ -184,21 +216,23 @@ async def monitor_multiple_run_groups(
 
 
 async def main(
-    monitor_type: str,
     logs_dir: Path,
+    monitor_config: Monitor.Config,
     run_groups: list[str] | None = None,
     out_dir: Path | None = None,
 ) -> None:
     """
     Main function to run the monitor on a directory of logs.
     """
 
+    monitor_config = monitor_config.model_copy()
+
     if out_dir:
         out_dir.mkdir(parents=True, exist_ok=True)
 
     results = await monitor_multiple_run_groups(
         logs_dir=logs_dir,
-        monitor_type=monitor_type,
+        monitor_config=monitor_config,
         run_groups=run_groups,
     )
 
@@ -211,41 +245,14 @@ async def main(
         logger.info(f"All monitor results written to {output_file}")
 
 
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Monitor agent logs for violations.")
-    parser.add_argument(
-        "--logs-dir",
-        type=Path,
-        help="Directory containing multiple run groups.",
-        required=True,
-    )
-    parser.add_argument(
-        "--run-groups",
-        nargs="+",
-        help="List of run group IDs to monitor.",
-        required=False,
-    )
-    parser.add_argument(
-        "-m",
-        "--monitor",
-        choices=["basic"],
-        default="basic",
-        help="Specify the monitor to use (default: basic).",
-    )
-    parser.add_argument(
-        "--out-dir",
-        type=Path,
-        help="Directory to save the monitor results JSON file (default: current directory).",
-        required=False,
+async def _run_from_cli(args: MonitorCLIArgs) -> None:
+    await main(
+        logs_dir=args.logs_dir,
+        monitor_config=args.monitor_config,
+        run_groups=args.run_groups or None,
+        out_dir=args.out_dir,
     )
 
-    args = parser.parse_args()
 
-    asyncio.run(
-        main(
-            monitor_type=args.monitor,
-            logs_dir=args.logs_dir,
-            run_groups=args.run_groups,
-            out_dir=args.out_dir,
-        )
-    )
+if __name__ == "__main__":
+    asyncio.run(chz.nested_entrypoint(_run_from_cli))