benchmax 0.1.2.dev14__tar.gz → 0.1.2.dev16__tar.gz

{benchmax-0.1.2.dev14 → benchmax-0.1.2.dev16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: benchmax
-Version: 0.1.2.dev14
+Version: 0.1.2.dev16
 Summary: Framework-Agnostic RL Environments for LLM Fine-Tuning
 Author: cgft.io
 Classifier: Programming Language :: Python :: 3
@@ -12,6 +12,7 @@ Requires-Dist: aiohttp>=3.13.1
 Requires-Dist: asyncio>=4.0.0
 Requires-Dist: cloudpickle>=3.0.0
 Requires-Dist: datasets>=4.0.0
+Requires-Dist: expt-logger>=0.1.0.dev20
 Provides-Extra: mcp
 Requires-Dist: fastmcp~=2.12.0; extra == "mcp"
 Requires-Dist: pyjwt>=2.10.1; extra == "mcp"

{benchmax-0.1.2.dev14 → benchmax-0.1.2.dev16}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "benchmax"
-version = "0.1.2.dev14"
+version = "0.1.2.dev16"
 description = "Framework-Agnostic RL Environments for LLM Fine-Tuning"
 readme = "README.md"
 authors = [{ name = "cgft.io" }]
@@ -10,6 +10,7 @@ dependencies = [
     "asyncio>=4.0.0",
     "cloudpickle>=3.0.0",
     "datasets>=4.0.0",
+    "expt-logger>=0.1.0.dev20",
 ]
 classifiers = [
     "Programming Language :: Python :: 3",

{benchmax-0.1.2.dev14 → benchmax-0.1.2.dev16}/src/benchmax/envs/base_env.py

@@ -1,8 +1,9 @@
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Dict, List, Any, Optional, Tuple
 from pathlib import Path
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 
-from benchmax.envs.types import ToolDefinition, StandardizedExample
+from benchmax.envs.tracking import TrackingConfig, log_env, with_tracking
+from benchmax.envs.types import StandardizedExample, ToolDefinition
 from benchmax.prompts.tools import render_tools_prompt
 
 if TYPE_CHECKING:
@@ -13,6 +14,38 @@ class BaseEnv(ABC):
     """Base benchmax environment for tool execution and reward computation"""
 
     system_prompt: str = ""
+    _tracking_config: TrackingConfig | None = None
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+
+        compute_reward = cls.__dict__.get("compute_reward")
+        if compute_reward is None:
+            return
+        if getattr(compute_reward, "__benchmax_tracking_wrapped__", False):
+            return
+
+        wrapped = with_tracking(lambda self, *a, **kw: self.get_tracking_config())(
+            compute_reward
+        )
+        setattr(wrapped, "__benchmax_tracking_wrapped__", True)
+        setattr(cls, "compute_reward", wrapped)
+
+    def __init__(
+        self,
+        experiment_id: Optional[str] = None,
+        api_key: Optional[str] = None,
+        **kwargs,
+    ):
+        self._tracking_config = TrackingConfig(
+            experiment_id=experiment_id, api_key=api_key
+        )
+
+    def get_tracking_config(self) -> TrackingConfig | None:
+        return self._tracking_config
+
+    def log_env(self, rollout_id: str, message: str) -> None:
+        log_env(rollout_id, message)
 
     # Override this method if your example does not match the default structure
     @classmethod

{benchmax-0.1.2.dev14 → benchmax-0.1.2.dev16}/src/benchmax/envs/excel/workdir/reward_fn.py

@@ -34,17 +34,23 @@ def spreadsheet_comparison_reward(
 
     output_path = workspace / output_filename
     ground_truth_path = workspace / ground_truth_filename
+    rollout_id = kwargs.get("rollout_id", "unknown_rollout")
+    log_env(
+        rollout_id, f"excel_reward:compare_files={ground_truth_filename}:{output_filename}:{answer_position}"
+    )
 
     # Return 1.0 score if the output completely matches the ground truth
     try:
         match, _ = compare_excel_cells(
             str(ground_truth_path), str(output_path), answer_position
         )
+        log_env(rollout_id, f"excel_reward:spreadsheet_match={float(match)}")
         return 1.0 if match else 0.0
     except Exception as e:
         print(
             f"Error comparing spreadsheets {ground_truth_path} and {output_path}: {e}"
         )
+        log_env(rollout_id, f"excel_reward:error={str(e)}")
         return 0.0
 
 

{benchmax-0.1.2.dev14 → benchmax-0.1.2.dev16}/src/benchmax/envs/mcp/parallel_mcp_env.py

@@ -19,6 +19,7 @@ except ModuleNotFoundError as e:
     ) from e
 
 from benchmax.envs.base_env import BaseEnv
+from benchmax.envs.tracking import to_tracking_payload
 from benchmax.envs.types import ToolDefinition
 from .server_pool import ServerPool
 from .provisioners.base_provisioner import BaseProvisioner
@@ -96,7 +97,7 @@ class ParallelMcpEnv(BaseEnv):
             provision_at_init: Whether to launch a server at the point of initialization
             **kwargs: Additional keyword arguments (currently unused).
         """
-        super().__init__()
+        super().__init__(**kwargs)
 
         self._workdir_path = Path(workdir_path).absolute()
         self._provisioner = provisioner
@@ -373,6 +374,8 @@ class ParallelMcpEnv(BaseEnv):
         payload = {
             "completion": completion or "",
             "ground_truth": ground_truth or "",
+            **to_tracking_payload(self.get_tracking_config()),
+            "rollout_id": rollout_id,
             **kwargs,
         }
 

{benchmax-0.1.2.dev14 → benchmax-0.1.2.dev16}/src/benchmax/envs/mcp/provisioners/utils.py

@@ -15,7 +15,8 @@ def setup_sync_dir(workdir_path: Path) -> Path:
 
     This creates a temp directory and copies:
     1. proxy_server.py from the mcp/ directory
-    2. All contents of the provided workdir_path
+    2. env tracking helper for reward logging
+    3. All contents of the provided workdir_path
 
     Args:
         workdir_path: Path to workdir containing mcp_config.yaml, setup.sh, etc.
@@ -38,6 +39,14 @@ def setup_sync_dir(workdir_path: Path) -> Path:
             )
         shutil.copy(src_server_path, sync_dir / "proxy_server.py")
 
+        # Copy shared env tracking helper for reward_fn logging.
+        src_tracking_path = Path(__file__).parents[2] / "tracking.py"
+        if not src_tracking_path.exists():
+            raise FileNotFoundError(
+                f"Expected tracking helper at {src_tracking_path}, but not found."
+            )
+        shutil.copy(src_tracking_path, sync_dir / "env_tracking.py")
+
         # Validate workdir exists and is a directory
         if not workdir_path.exists():
             raise FileNotFoundError(
@@ -91,7 +100,7 @@ def get_setup_command() -> str:
 # Install uv
 curl -LsSf https://astral.sh/uv/install.sh | sh
 UV_VENV_CLEAR=1 uv venv ~/venv && source ~/venv/bin/activate
-uv pip install fastmcp~=2.12.0 pyyaml psutil
+uv pip install fastmcp~=2.12.0 pyyaml psutil expt-logger
 bash setup.sh
 """
 

{benchmax-0.1.2.dev14 → benchmax-0.1.2.dev16}/src/benchmax/envs/mcp/proxy_server.py

@@ -24,12 +24,30 @@ from starlette.requests import Request
 from starlette.responses import PlainTextResponse, FileResponse, JSONResponse, Response
 from starlette.datastructures import UploadFile
 
-from reward_fn import reward_functions  # type: ignore
+try:
+    from benchmax.envs.tracking import log_env, pop_tracking_config, tracking_context
+except Exception:
+    # In provisioned MCP servers, this helper is copied as env_tracking.py.
+    from env_tracking import log_env, pop_tracking_config, tracking_context  # type: ignore
+
+from reward_fn import reward_functions as imported_reward_functions  # type: ignore
 
 RewardFunction = Callable[..., Union[float, Awaitable[float]]]
 DEFAULT_API_SECRET = "dev_default_api_secret_please_change_me_32chars!"
 
 
+def _with_log_env(func: RewardFunction) -> RewardFunction:
+    """Decorator that binds the shared log_env callable into reward_fn globals."""
+    func.__globals__["log_env"] = log_env
+    return func
+
+
+reward_functions: Dict[str, RewardFunction] = {
+    name: _with_log_env(func)
+    for name, func in (imported_reward_functions or {}).items()
+}
+
+
 # ---------------- Utility Functions ---------------- #
 def setup_workspace(base_dir: Path) -> Path:
     """Create a unique workspace directory."""
@@ -329,14 +347,17 @@ class ProxyServer:
                 status_code=400,
             )
 
+        payload_kwargs: Dict[str, Any] = {
+            k: v for k, v in data.items() if k not in ("completion", "ground_truth")
+        }
+        tracking_config = pop_tracking_config(payload_kwargs)
+
         kwargs: Dict[str, Any] = {
             "completion": completion,
             "ground_truth": ground_truth,
             "workspace": self.workspace,
             "mcp_client": self.client,
-            **{
-                k: v for k, v in data.items() if k not in ("completion", "ground_truth")
-            },
+            **payload_kwargs,
         }
 
         async def _call_reward(name: str, func: RewardFunction) -> Tuple[str, float]:
@@ -357,8 +378,9 @@ class ProxyServer:
         rf: Dict[str, RewardFunction] = reward_functions or {}
 
         try:
-            tasks = [_call_reward(name, func) for name, func in rf.items()]
-            results_list: List[Tuple[str, float]] = await asyncio.gather(*tasks)
+            with tracking_context(tracking_config):
+                tasks = [_call_reward(name, func) for name, func in rf.items()]
+                results_list: List[Tuple[str, float]] = await asyncio.gather(*tasks)
             results: Dict[str, float] = dict(results_list)
             return JSONResponse(results)
         except Exception as e:

benchmax-0.1.2.dev16/src/benchmax/envs/tracking.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import inspect
+import logging
+import os
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass
+from functools import wraps
+from typing import Any, Callable, Dict, Iterator, Optional
+
+LOGGER = logging.getLogger(__name__)
+
+TRACKING_EXPERIMENT_ID_KEY = "__benchmax_expt_logger_experiment_id"
+TRACKING_API_KEY_KEY = "__benchmax_expt_logger_api_key"
+
+_ACTIVE_TRACKER: ContextVar[Any | None] = ContextVar(
+    "benchmax_active_expt_logger_tracker", default=None
+)
+_TRACKER_CACHE: Dict[tuple[Optional[str], Optional[str]], Any | None] = {}
+
+
+@dataclass(frozen=True)
+class TrackingConfig:
+    experiment_id: Optional[str] = None
+    api_key: Optional[str] = None
+
+    def resolved_experiment_id(self) -> Optional[str]:
+        return self.experiment_id or os.getenv("EXPT_LOGGER_EXPERIMENT_ID")
+
+    def is_enabled(self) -> bool:
+        return bool(self.resolved_experiment_id())
+
+
+def _build_tracker(config: TrackingConfig) -> Any | None:
+    if not config.is_enabled():
+        return None
+
+    try:
+        import expt_logger
+    except Exception as e:
+        LOGGER.debug("expt_logger import failed; env tracking disabled: %s", e)
+        return None
+
+    try:
+        run = expt_logger.init(
+            experiment_id=config.resolved_experiment_id(),
+            api_key=config.api_key,
+        )
+    except Exception as e:
+        LOGGER.debug("expt_logger init failed; env tracking disabled: %s", e)
+        return None
+
+    if hasattr(expt_logger, "log_environment"):
+        return expt_logger
+    if hasattr(run, "log_environment"):
+        return run
+
+    LOGGER.debug("expt_logger has no log_environment; env tracking disabled")
+    return None
+
+
+def get_tracker(config: TrackingConfig | None) -> Any | None:
+    if config is None:
+        return None
+
+    key = (config.resolved_experiment_id(), config.api_key)
+    if key not in _TRACKER_CACHE:
+        _TRACKER_CACHE[key] = _build_tracker(config)
+    return _TRACKER_CACHE[key]
+
+
+@contextmanager
+def tracking_context(config: TrackingConfig | None) -> Iterator[None]:
+    token = _ACTIVE_TRACKER.set(get_tracker(config))
+    try:
+        yield
+    finally:
+        _ACTIVE_TRACKER.reset(token)
+
+
+def log_env(rollout_id: str, message: str) -> None:
+    tracker = _ACTIVE_TRACKER.get()
+    if tracker is None:
+        return
+
+    try:
+        tracker.log_environment(rollout_id, str(message))
+    except Exception as e:
+        LOGGER.debug("log_environment failed: %s", e)
+
+
+def with_tracking(
+    config_resolver: Callable[..., TrackingConfig | None],
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """Wrap a function so calls run with an active env tracking context."""
+
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        if inspect.iscoroutinefunction(func):
+
+            @wraps(func)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                with tracking_context(config_resolver(*args, **kwargs)):
+                    return await func(*args, **kwargs)
+
+            return async_wrapper
+
+        @wraps(func)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+            with tracking_context(config_resolver(*args, **kwargs)):
+                return func(*args, **kwargs)
+
+        return sync_wrapper
+
+    return decorator
+
+
+def to_tracking_payload(config: TrackingConfig | None) -> Dict[str, str]:
+    if config is None:
+        return {}
+
+    payload: Dict[str, str] = {}
+    resolved_experiment_id = config.resolved_experiment_id()
+    if resolved_experiment_id:
+        payload[TRACKING_EXPERIMENT_ID_KEY] = resolved_experiment_id
+    if config.api_key:
+        payload[TRACKING_API_KEY_KEY] = config.api_key
+    return payload
+
+
+def pop_tracking_config(payload: Dict[str, Any]) -> TrackingConfig:
+    experiment_id = payload.pop(TRACKING_EXPERIMENT_ID_KEY, None)
+    api_key = payload.pop(TRACKING_API_KEY_KEY, None)
+    return TrackingConfig(experiment_id=experiment_id, api_key=api_key)

{benchmax-0.1.2.dev14 → benchmax-0.1.2.dev16}/src/benchmax/envs/wikipedia/wiki_env.py

@@ -4,6 +4,7 @@ import re
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
 from benchmax.envs.base_env import BaseEnv
+from benchmax.envs.tracking import log_env
 from benchmax.envs.types import ToolDefinition, StandardizedExample
 from benchmax.envs.wikipedia.utils import APIKeyRotator, clean_html, safe_request
 
@@ -13,7 +14,7 @@ Write your complete answer on the final line only as a concise entity, within th
 """
 
 
-def text_match_reward_function(completion: str, ground_truth: str, **kwargs) -> float:
+def text_match_reward_function(completion: str, ground_truth: str, rollout_id: str, **kwargs) -> float:
     """
     Score 1.0 if ground truth appears in <answer> tags, else 0.0.
 
@@ -31,10 +32,13 @@ def text_match_reward_function(completion: str, ground_truth: str, **kwargs) ->
         r"<answer>(.*?)</answer>", completion, flags=re.IGNORECASE | re.DOTALL
     )
     if not m:
+        log_env(rollout_id, "wikipedia_reward:no_answer_tag")
         return 0.0
 
     answer_text = unescape(m.group(1)).strip().lower()
-    return float(ground_truth.lower() == answer_text)
+    score = float(ground_truth.lower() == answer_text)
+    log_env(rollout_id, f"wikipedia_reward:text_match={score}")
+    return score
 
 
 def _make_wikipedia_tools(key_rotator: APIKeyRotator):
@@ -264,5 +268,5 @@ class WikipediaEnv(BaseEnv):
     ) -> Dict[str, float]:
         """Compute rewards using the text match reward function."""
         return {
-            "text_match": text_match_reward_function(completion, ground_truth, **kwargs)
+            "text_match": text_match_reward_function(completion, ground_truth, rollout_id, **kwargs)
         }

{benchmax-0.1.2.dev14 → benchmax-0.1.2.dev16}/src/benchmax.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: benchmax
-Version: 0.1.2.dev14
+Version: 0.1.2.dev16
 Summary: Framework-Agnostic RL Environments for LLM Fine-Tuning
 Author: cgft.io
 Classifier: Programming Language :: Python :: 3
@@ -12,6 +12,7 @@ Requires-Dist: aiohttp>=3.13.1
 Requires-Dist: asyncio>=4.0.0
 Requires-Dist: cloudpickle>=3.0.0
 Requires-Dist: datasets>=4.0.0
+Requires-Dist: expt-logger>=0.1.0.dev20
 Provides-Extra: mcp
 Requires-Dist: fastmcp~=2.12.0; extra == "mcp"
 Requires-Dist: pyjwt>=2.10.1; extra == "mcp"

{benchmax-0.1.2.dev14 → benchmax-0.1.2.dev16}/src/benchmax.egg-info/SOURCES.txt

@@ -18,6 +18,7 @@ src/benchmax/bundle/payload.py
 src/benchmax/bundle/validator.py
 src/benchmax/envs/__init__.py
 src/benchmax/envs/base_env.py
+src/benchmax/envs/tracking.py
 src/benchmax/envs/types.py
 src/benchmax/envs/crm/crm_env.py
 src/benchmax/envs/crm/workdir/reward_fn.py

{benchmax-0.1.2.dev14 → benchmax-0.1.2.dev16}/src/benchmax.egg-info/requires.txt

@@ -2,6 +2,7 @@ aiohttp>=3.13.1
 asyncio>=4.0.0
 cloudpickle>=3.0.0
 datasets>=4.0.0
+expt-logger>=0.1.0.dev20
 
 [mcp]
 fastmcp~=2.12.0