diff --git a/docs/v6/guides/training-agents.mdx b/docs/v6/guides/training-agents.mdx index 11f1abf13..187188444 100644 --- a/docs/v6/guides/training-agents.mdx +++ b/docs/v6/guides/training-agents.mdx @@ -94,6 +94,11 @@ One job spans the session; each step appends a batch and trains on it: - **Open the job** with `group=8` - 8 rollouts per task, so the rewards are comparable (next). - **Roll out** the batch, the same eval as the [previous guide](/v6/guides/running-an-eval). The [runtime](/v6/reference/runtime) sets where it runs; swap `LocalRuntime` for `HUDRuntime()` unchanged. + Prefer **`HUDRuntime`** (or `LocalRuntime`) for TrainingClient loops: the agent stays local so + `create_agent(...)` gateway clients, `return_token_ids`, system prompts, and Tinker/Qwen + `completion_kwargs` all apply in-process. **`HostedRuntime`** also works with `create_agent` + (the gateway client is rebuilt remotely and those kwargs still travel), but a true BYOK / + custom `model_client` cannot be serialized — use `HUDRuntime` for that. - **Nudge** with `trainer.step` - the one line that learns. It scores each rollout against its group, shifts the weights, then **promotes** them so the gateway serves the new ones at once. diff --git a/docs/v6/reference/agents.mdx b/docs/v6/reference/agents.mdx index 1f32b2961..301194882 100644 --- a/docs/v6/reference/agents.mdx +++ b/docs/v6/reference/agents.mdx @@ -35,7 +35,14 @@ The HUD gateway is an OpenAI-compatible endpoint (`settings.hud_gateway_url`, de `https://inference.beta.hud.ai`) that fronts every provider behind your single `HUD_API_KEY`, so you switch between Claude, GPT, and Gemini by name alone, with unified tracing. `create_agent` accepts any id the gateway knows (`claude-...`, `gpt-...`, `gemini-...`); extra kwargs pass through to the agent's -config. +config — including `system_prompt` and `completion_kwargs` (for example +`extra_body={"return_token_ids": True}` for training). + +For local loops (`HUDRuntime` / `LocalRuntime`) the agent keeps a gateway `model_client`. For +[`HostedRuntime`](/v6/reference/runtime#hostedruntime) that client is stripped and rebuilt on the +platform from the model name, so the same `create_agent(...)` call works for hosted evals and +training. A true BYOK / custom `model_client` cannot travel with hosted submission — keep the agent +loop local with `HUDRuntime` instead. Built-in agents are **catalog-driven**: each run they read the environment's manifest, open the capabilities they support, build the matching provider tools, and loop against `run.prompt_messages`. diff --git a/docs/v6/reference/runtime.mdx b/docs/v6/reference/runtime.mdx index 102760d83..74495a4b8 100644 --- a/docs/v6/reference/runtime.mdx +++ b/docs/v6/reference/runtime.mdx @@ -173,6 +173,14 @@ Where `HUDRuntime` runs the agent loop locally against a tunneled env, `HostedRu runs the agent right next to it. This process only submits the rollout and polls its trace to completion. It requires a gateway agent that can serialize its identity (Claude/OpenAI/Gemini). +Agents from [`create_agent`](/v6/reference/agents#create-agent) are supported: the local gateway +`model_client` is stripped and rebuilt on the platform, while `system_prompt`, +`completion_kwargs` (including `return_token_ids` and Tinker/Qwen request kwargs), and other +config travel intact. A true BYOK / custom `model_client` cannot be serialized — use +`HUDRuntime` (or `LocalRuntime`) so the agent loop stays on your machine. For +[`TrainingClient`](/v6/reference/training#trainingclient) workflows, prefer `HUDRuntime` unless +you specifically need co-located hosted execution. + ### `Runtime` ```python diff --git a/docs/v6/reference/training.mdx b/docs/v6/reference/training.mdx index d9a11877f..3536d87a8 100644 --- a/docs/v6/reference/training.mdx +++ b/docs/v6/reference/training.mdx @@ -38,6 +38,18 @@ Gradients accumulate across `forward_backward` / `forward_backward_custom` / `ba | `TrainingClient` | Drives managed training for one model slug. | | `CheckpointResponse` | One node in the model's checkpoint tree. | +## Runtime placement + +Training reuses the same [runtime](/v6/reference/runtime) placement as eval. Which one you pick +matters for how the agent is constructed: + +| Placement | Agent loop | Use for training when | +| --- | --- | --- | +| `LocalRuntime` / `HUDRuntime` | Local (this process) | Default. `create_agent` gateway clients, `system_prompt`, and `completion_kwargs` (`return_token_ids`, Tinker/Qwen kwargs) apply in-process. Required for a true BYOK / custom `model_client`. | +| `HostedRuntime` | Remote (serialized agent) | Co-located env+agent on a leased box. `create_agent(...)` is supported: the gateway client is rebuilt remotely and config kwargs still travel. Custom / BYOK clients fail early — switch to `HUDRuntime`. | + +Cookbooks that load a platform taskset (`Taskset.from_api`) use `HUDRuntime()` for this reason. + ## TrainingClient A `TrainingClient` drives managed training for one model: it accumulates gradients from rewarded diff --git a/hud/agents/__init__.py b/hud/agents/__init__.py index b816634ab..ae582bc8f 100644 --- a/hud/agents/__init__.py +++ b/hud/agents/__init__.py @@ -30,7 +30,16 @@ def create_agent(model: str, **kwargs: Any) -> GatewayAgent: """Create an agent routed through the HUD gateway. - For direct API access with provider API keys, instantiate the agent classes directly. + Attaches a gateway ``model_client`` for local loops (``HUDRuntime`` / + ``LocalRuntime``). For ``HostedRuntime``, that client is stripped by + :meth:`~hud.agents.tool_agent.ToolAgent.hosted_spec` and rebuilt on the + platform from the model name — ``system_prompt``, ``completion_kwargs`` + (including ``return_token_ids`` / Tinker kwargs), and other config still + travel. True BYOK / custom clients are not serializable; use + ``HUDRuntime`` for those TrainingClient workflows. + + For direct API access with provider API keys, instantiate the agent classes + directly. """ requested_model = model model = normalize_gateway_model_id(model) diff --git a/hud/agents/tool_agent.py b/hud/agents/tool_agent.py index fb538248d..b0105299a 100644 --- a/hud/agents/tool_agent.py +++ b/hud/agents/tool_agent.py @@ -90,12 +90,17 @@ def __init_subclass__(cls, **kwargs: Any) -> None: def hosted_spec(self) -> dict[str, Any]: """HUD-hosted execution runs the agent remotely, so it is reconstructed there from this identity (type, model, step budget, system - prompt) with the model resolved through the HUD gateway. + prompt, provider kwargs) with the model resolved through the HUD gateway. """ - if self.config.model_client is not None: + from hud.utils.gateway import is_gateway_client + + if not is_gateway_client(self.config.model_client): raise ValueError( "hosted execution cannot serialize a custom model_client; " - "set the model by name and let the hosted runner build the gateway client" + "use create_agent(model, ...) so the hosted runner rebuilds the " + "gateway client, or run the agent loop locally with HUDRuntime() " + "/ LocalRuntime (recommended for TrainingClient workflows that " + "attach a BYOK client)" ) agent_type = AgentType.of(self) if agent_type is None: diff --git a/hud/eval/tests/test_hosted.py b/hud/eval/tests/test_hosted.py index 374aa5481..e18ac556a 100644 --- a/hud/eval/tests/test_hosted.py +++ b/hud/eval/tests/test_hosted.py @@ -92,13 +92,83 @@ def test_hosted_spec_serializes_full_config() -> None: assert "hosted_tools" not in config +def test_hosted_spec_strips_create_agent_gateway_client(monkeypatch: pytest.MonkeyPatch) -> None: + """create_agent attaches a gateway client for local loops; HostedRuntime must + still serialize system_prompt + completion_kwargs (HUD-2100).""" + from hud.agents import create_agent + from hud.utils.gateway import GatewayModelInfo, GatewayProviderInfo, mark_gateway_client + + class _GatewayStub: + """Attribute-capable stand-in for a provider SDK client.""" + + model = GatewayModelInfo( + id="arith-rl", + model_name="arith-rl", + sdk_agent_type="openai_compatible", + provider=GatewayProviderInfo(name="openai"), + ) + monkeypatch.setattr("hud.agents.list_gateway_models", lambda: [model]) + monkeypatch.setattr( + "hud.agents.build_gateway_client", + lambda _provider: mark_gateway_client(_GatewayStub()), + ) + + agent = create_agent( + "arith-rl", + system_prompt="/no_think", + completion_kwargs={ + "extra_body": { + "return_token_ids": True, + "chat_template_kwargs": {"enable_thinking": False}, + } + }, + ) + assert agent.config.model_client is not None + + spec = agent.hosted_spec() + config = spec["config"] + assert spec["type"] == "openai_compatible" + assert config["model"] == "arith-rl" + assert config["system_prompt"] == "/no_think" + assert config["completion_kwargs"]["extra_body"]["return_token_ids"] is True + assert config["completion_kwargs"]["extra_body"]["chat_template_kwargs"] == { + "enable_thinking": False + } + assert "model_client" not in config + + def test_hosted_spec_rejects_custom_model_client() -> None: agent = _agent() agent.config = OpenAIChatConfig(model="m", model_client=object()) - with pytest.raises(ValueError, match="model_client"): + with pytest.raises(ValueError, match="custom model_client"): + agent.hosted_spec() + with pytest.raises(ValueError, match="HUDRuntime"): agent.hosted_spec() +def test_mark_gateway_client_round_trip() -> None: + from hud.utils.gateway import is_gateway_client, mark_gateway_client + + class _Stub: + pass + + tagged = mark_gateway_client(_Stub()) + assert is_gateway_client(tagged) + assert not is_gateway_client(_Stub()) + assert is_gateway_client(None) + + +def test_mark_gateway_client_works_on_async_openai() -> None: + """Provider SDK clients must accept the attribute mark (no id-registry fallback).""" + from openai import AsyncOpenAI + + from hud.utils.gateway import is_gateway_client, mark_gateway_client + + client = AsyncOpenAI(api_key="test", base_url="http://localhost") + assert is_gateway_client(mark_gateway_client(client)) + assert not is_gateway_client(AsyncOpenAI(api_key="test", base_url="http://localhost")) + + @pytest.mark.asyncio async def test_run_rejects_non_gateway_agent() -> None: """An agent that can't serialize its identity yields a failed Run, not a crash.""" diff --git a/hud/utils/gateway.py b/hud/utils/gateway.py index 4fe5212bd..310f90df9 100644 --- a/hud/utils/gateway.py +++ b/hud/utils/gateway.py @@ -8,7 +8,7 @@ from __future__ import annotations from functools import lru_cache -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, TypeVar from openai import AsyncOpenAI from pydantic import BaseModel, Field @@ -25,6 +25,8 @@ GatewayClient: TypeAlias = AsyncAnthropic | AsyncAnthropicBedrock | GenaiClient | AsyncOpenAI +_T = TypeVar("_T") + class GatewayProviderInfo(BaseModel): name: str | None = None @@ -88,22 +90,52 @@ def build_gateway_client(provider: str) -> GatewayClient: if provider == "anthropic": from anthropic import AsyncAnthropic - return AsyncAnthropic(api_key=settings.api_key, base_url=settings.hud_gateway_url) + client: GatewayClient = AsyncAnthropic( + api_key=settings.api_key, base_url=settings.hud_gateway_url + ) + return mark_gateway_client(client) if provider == "gemini": from google import genai from google.genai.types import HttpOptions - return genai.Client( + client = genai.Client( api_key=settings.api_key, http_options=HttpOptions( api_version="v1beta", base_url=settings.hud_gateway_url, ), ) + return mark_gateway_client(client) # OpenAI-compatible (openai, azure, together, groq, fireworks, etc.) - return AsyncOpenAI(api_key=settings.api_key, base_url=settings.hud_gateway_url) + client = AsyncOpenAI(api_key=settings.api_key, base_url=settings.hud_gateway_url) + return mark_gateway_client(client) + + +_HUD_GATEWAY_CLIENT_ATTR = "_hud_gateway_client" + + +def mark_gateway_client(client: _T) -> _T: + """Tag a client built for the HUD gateway so hosted serialization can drop it. + + The mark lives on the instance (``setattr``). Clients that reject attributes + are not supported — real provider SDK clients accept them. + """ + setattr(client, _HUD_GATEWAY_CLIENT_ATTR, True) + return client + + +def is_gateway_client(client: object | None) -> bool: + """True when ``client`` is missing or was built by :func:`build_gateway_client`. + + Hosted reconstruction rebuilds a gateway client from the model name, so a + tagged gateway client is safe to strip from ``hosted_spec``. A true BYOK / + custom client is not. + """ + if client is None: + return True + return bool(getattr(client, _HUD_GATEWAY_CLIENT_ATTR, False)) @lru_cache(maxsize=1)