benchmax 0.2.2.dev0__tar.gz → 0.2.3__tar.gz

{benchmax-0.2.2.dev0 → benchmax-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: benchmax
-Version: 0.2.2.dev0
+Version: 0.2.3
 Summary: Platform-independent runtime for grouped LLM environments
 Author: benchmax Authors
 Classifier: Operating System :: OS Independent
@@ -98,6 +98,8 @@ environment class + constructor arguments
 
 benchmax creates the portable artifact so the same environment can be loaded outside the author's checkout. `castform` handles uploading the bundle, validating it remotely, and using it for training.
 
+Remote runtimes install `pip_dependencies` without enabling prereleases globally. If a dependency graph needs a prerelease, list that package explicitly even when it would normally be transitive. For example, use `pip_dependencies=["parent-package==1.0.0", "transitive-package==2.0.0rc1"]`.
+
 ## further reading
 
 - [base environment guide](src/benchmax/envs/base/README.md)

{benchmax-0.2.2.dev0 → benchmax-0.2.3}/README.md

@@ -83,6 +83,8 @@ environment class + constructor arguments
 
 benchmax creates the portable artifact so the same environment can be loaded outside the author's checkout. `castform` handles uploading the bundle, validating it remotely, and using it for training.
 
+Remote runtimes install `pip_dependencies` without enabling prereleases globally. If a dependency graph needs a prerelease, list that package explicitly even when it would normally be transitive. For example, use `pip_dependencies=["parent-package==1.0.0", "transitive-package==2.0.0rc1"]`.
+
 ## further reading
 
 - [base environment guide](src/benchmax/envs/base/README.md)

{benchmax-0.2.2.dev0 → benchmax-0.2.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "benchmax"
-version = "0.2.2.dev0"
+version = "0.2.3"
 description = "Platform-independent runtime for grouped LLM environments"
 readme = "README.md"
 authors = [{ name = "benchmax Authors" }]

{benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/auth.py

@@ -16,6 +16,7 @@ from dataclasses import dataclass, field
 from typing import Protocol, runtime_checkable
 
 import httpx
+from openai import AsyncOpenAI, OpenAI
 
 __all__ = [
     "InjectedAuth",
@@ -24,6 +25,8 @@ __all__ = [
     "RequestModelAuth",
     "StaticBearerAuth",
     "bind_model_auth",
+    "create_async_openai_client",
+    "create_openai_client",
 ]
 
 
@@ -99,6 +102,52 @@ class RequestModelAuth(httpx.Auth):
         yield request
 
 
+def create_openai_client(
+    *,
+    model: str,
+    base_url: str,
+    auth: ModelAuth,
+    request_id: str,
+    max_retries: int = 2,
+) -> OpenAI:
+    """Create a synchronous OpenAI-compatible client with explicit auth."""
+
+    context = ModelRequestContext(
+        base_url=base_url,
+        model=model,
+        rollout_id=request_id,
+    )
+    return OpenAI(
+        base_url=base_url,
+        api_key="benchmax-explicit-auth",
+        http_client=httpx.Client(auth=RequestModelAuth(auth, context)),
+        max_retries=max_retries,
+    )
+
+
+def create_async_openai_client(
+    *,
+    model: str,
+    base_url: str,
+    auth: ModelAuth,
+    request_id: str,
+    max_retries: int = 2,
+) -> AsyncOpenAI:
+    """Create an asynchronous OpenAI-compatible client with explicit auth."""
+
+    context = ModelRequestContext(
+        base_url=base_url,
+        model=model,
+        rollout_id=request_id,
+    )
+    return AsyncOpenAI(
+        base_url=base_url,
+        api_key="benchmax-explicit-auth",
+        http_client=httpx.AsyncClient(auth=RequestModelAuth(auth, context)),
+        max_retries=max_retries,
+    )
+
+
 def _resolve_headers_sync(
     auth: ModelAuth,
     context: ModelRequestContext,

{benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/bundle.py

@@ -203,6 +203,9 @@ def dump_bundle(
         env_class: A concrete Environment implementation.
         constructor_args: kwargs for ``env_class(**...)`` on load.
         pip_dependencies: Recorded in metadata. NOT installed by this call.
+            List every required prerelease explicitly, including a prerelease
+            that would otherwise only be a transitive dependency. Remote
+            runtimes do not enable prereleases globally.
         local_modules: Additional modules to pickle by value. Use this for
             local source outside the environment's own Python project. Such
             source otherwise fails loudly unless its installed distribution is

{benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/__init__.py

@@ -19,6 +19,7 @@ from benchmax.envs.shared_types import (
     RolloutFailure,
     RolloutOutcome,
     RolloutRequest,
+    ValidationDiagnostic,
 )
 
 __all__ = [
@@ -39,6 +40,7 @@ __all__ = [
     "RolloutFailure",
     "RolloutOutcome",
     "RolloutRequest",
+    "ValidationDiagnostic",
     "Tool",
     "StaticBearerAuth",
     "bind_model_auth",

{benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/environment.py

@@ -15,6 +15,7 @@ from benchmax.envs.shared_types import (
     RolloutFailure,
     RolloutOutcome,
     RolloutRequest,
+    ValidationDiagnostic,
 )
 
 __all__ = ["Environment"]
@@ -42,6 +43,11 @@ class Environment[Payload, Attempt: RolloutAttempt](ABC):
 
         return False
 
+    def validation_diagnostics(self) -> Sequence[ValidationDiagnostic]:
+        """Return static warnings and errors for environment configuration."""
+
+        return ()
+
     @abstractmethod
     async def create_dataset(
         self,

{benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/harbor/env.py

@@ -26,6 +26,7 @@ from benchmax.envs.shared_types import (
     RolloutAttempt,
     RolloutOutcome,
     RolloutRequest,
+    ValidationDiagnostic,
 )
 
 if TYPE_CHECKING:
@@ -59,6 +60,72 @@ _TERMINATION_REASON_BY_EXCEPTION = {
 _HARNESS_REPORTED_TERMINATION_REASONS = frozenset(
     {"context_exceeded", "output_exceeded", "max_turns_exceeded", "tool_budget_exceeded"}
 )
+_TRAINER_OWNED_MODEL_FIELDS = frozenset(
+    {
+        "temperature",
+        "top_p",
+        "top_k",
+        "presence_penalty",
+        "frequency_penalty",
+        "seed",
+        "stop",
+    }
+)
+_OUTPUT_CAP_FIELDS = frozenset({"max_tokens", "max_completion_tokens"})
+_UNSUPPORTED_MODEL_FIELDS = frozenset(
+    {
+        "best_of",
+        "do_sample",
+        "function_call",
+        "functions",
+        "grammar",
+        "max_new_tokens",
+        "min_p",
+        "min_tokens",
+        "num_beams",
+        "prediction",
+        "reasoning_effort",
+        "repetition_penalty",
+        "typical_p",
+        "verbosity",
+        "web_search_options",
+    }
+)
+
+
+def _walk_model_controls(
+    value: Mapping[str, object],
+    prefix: str = "agent.kwargs",
+) -> Sequence[tuple[str, str, object]]:
+    controls: list[tuple[str, str, object]] = []
+    for field, child in value.items():
+        path = f"{prefix}.{field}"
+        controls.append((path, field, child))
+        if isinstance(child, Mapping):
+            controls.extend(_walk_model_controls(child, path))
+    return controls
+
+
+def _unsupported_model_control(field: str, value: object) -> bool:
+    if field == "return_routed_experts":
+        return True
+    if field in _UNSUPPORTED_MODEL_FIELDS:
+        return value is not None
+    if field == "n":
+        return value not in (None, 1)
+    if field == "tool_choice":
+        return value not in (None, "auto")
+    if field == "logprobs":
+        return value not in (None, False)
+    if field == "top_logprobs":
+        return value is not None
+    if field == "parallel_tool_calls":
+        return value not in (None, True)
+    if field == "response_format":
+        return value not in (None, {}, {"type": "text"})
+    if field == "logit_bias":
+        return value not in (None, {})
+    return False
 
 
 class HarborEnv(Environment["TaskConfig", RolloutAttempt]):
@@ -74,6 +141,48 @@ class HarborEnv(Environment["TaskConfig", RolloutAttempt]):
 
         return self._requires_public_model_endpoint
 
+    def validation_diagnostics(self) -> Sequence[ValidationDiagnostic]:
+        """Flag Harbor harness controls that conflict with tracked training."""
+
+        agent = self._trial.agent
+        config = agent.config if isinstance(agent, BundledHarborAgent) else agent
+        kwargs = getattr(config, "kwargs", None)
+        if not isinstance(kwargs, Mapping):
+            return ()
+        diagnostics: list[ValidationDiagnostic] = []
+        for path, field, value in _walk_model_controls(kwargs):
+            if field in _OUTPUT_CAP_FIELDS:
+                diagnostics.append(
+                    ValidationDiagnostic(
+                        severity="warning",
+                        code="harness_output_cap",
+                        location=path,
+                        message=(
+                            f"{path} is a harness-requested output cap; Castform may clamp it "
+                            "to the remaining trainer context budget"
+                        ),
+                    )
+                )
+            elif field in _TRAINER_OWNED_MODEL_FIELDS:
+                diagnostics.append(
+                    ValidationDiagnostic(
+                        severity="error",
+                        code="trainer_owned_model_control",
+                        location=path,
+                        message=(f"{path} is trainer-owned and cannot be set by a Harbor harness"),
+                    )
+                )
+            elif _unsupported_model_control(field, value):
+                diagnostics.append(
+                    ValidationDiagnostic(
+                        severity="error",
+                        code="unsupported_model_control",
+                        location=path,
+                        message=f"{path}={value!r} is unsupported by tracked training sessions",
+                    )
+                )
+        return tuple(diagnostics)
+
     def __init__(
         self,
         *,

{benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/shared_types.py

@@ -13,12 +13,23 @@ __all__ = [
     "RolloutFailure",
     "RolloutOutcome",
     "RolloutRequest",
+    "ValidationDiagnostic",
 ]
 
 DatasetSplit = Literal["train", "eval"]
 type RewardMap = Mapping[str, float]
 
 
+@dataclass(frozen=True, slots=True)
+class ValidationDiagnostic:
+    """A configuration issue discoverable before executing a rollout."""
+
+    severity: Literal["warning", "error"]
+    code: str
+    message: str
+    location: str | None = None
+
+
 class RolloutFailure(RuntimeError):  # noqa: N818 — public exported name
     """Operational rollout failure that should become a terminal outcome.
 

benchmax-0.2.3/src/benchmax/rag/__init__.py

@@ -0,0 +1,13 @@
+"""Runtime building blocks for retrieval-augmented Benchmax environments."""
+
+from benchmax.rag.embed import DEFAULT_EMBED_MODEL, OpenAIEmbedder
+from benchmax.rag.env import RagEnv
+from benchmax.rag.search import SearchClient, SearchResult
+
+__all__ = [
+    "DEFAULT_EMBED_MODEL",
+    "OpenAIEmbedder",
+    "SearchClient",
+    "RagEnv",
+    "SearchResult",
+]

benchmax-0.2.3/src/benchmax/rag/embed.py

@@ -0,0 +1,79 @@
+"""OpenAI-compatible query embedding for RAG search clients.
+
+The rollout-facing provider search clients (Turbopuffer / Pinecone / Chroma)
+accept an optional async ``embed_fn``. Wiring one in makes vector / hybrid
+retrieval work regardless of how the user's index was built — turbopuffer
+vector/hybrid (no server-side embed), a pinecone index NOT on the hosted model,
+or a non-cloud chroma collection.
+
+``qa-gen`` does NOT need this (it reads chunks directly); it's only for retrieval.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from benchmax.auth import ModelAuth, create_async_openai_client
+
+# The served, non-hidden embeddings model in the llm-proxy catalog.
+DEFAULT_EMBED_MODEL = "text-embedding-3-large"
+
+
+@dataclass(frozen=True, slots=True)
+class OpenAIEmbedder:
+    """Pickle-safe embedding callable with explicit call-time model auth.
+
+    ``auth`` follows the same contract as :class:`benchmax.rewards.Judge`.
+    Managed environments use ``InjectedAuth("embedding")``; customer-owned
+    endpoints use ``StaticBearerAuth``. No platform or SDK environment
+    credential is inferred.
+    """
+
+    model: str
+    base_url: str
+    auth: ModelAuth
+    timeout: float | None = 60.0
+    max_retries: int = 2
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.auth, ModelAuth):
+            raise TypeError("embedding auth must implement ModelAuth")
+        if not isinstance(self.model, str) or not self.model.strip():
+            raise ValueError("embedding model must be non-empty")
+        if not isinstance(self.base_url, str) or not self.base_url.strip():
+            raise ValueError("embedding base_url must be non-empty")
+        if self.timeout is not None and (
+            isinstance(self.timeout, bool)
+            or not isinstance(self.timeout, (int, float))
+            or self.timeout <= 0
+        ):
+            raise ValueError("embedding timeout must be positive or None")
+        if (
+            isinstance(self.max_retries, bool)
+            or not isinstance(self.max_retries, int)
+            or self.max_retries < 0
+        ):
+            raise ValueError("embedding max_retries must be non-negative")
+
+    async def __call__(self, texts: list[str]) -> list[list[float]]:
+        if not isinstance(texts, list) or any(not isinstance(text, str) for text in texts):
+            raise TypeError("embedding input must be a list of strings")
+        if not texts:
+            return []
+
+        client = create_async_openai_client(
+            model=self.model,
+            base_url=self.base_url,
+            auth=self.auth,
+            request_id="rag-embedding",
+            max_retries=self.max_retries,
+        )
+        try:
+            response = await client.embeddings.create(
+                model=self.model,
+                input=texts,
+                timeout=self.timeout,
+            )
+            return [item.embedding for item in response.data]
+        finally:
+            await client.close()