From 39d82b08b8553b57f88f0d2351699ac8d55699eb Mon Sep 17 00:00:00 2001 From: Lucas Tonon Date: Sat, 18 Jul 2026 23:16:52 -0300 Subject: [PATCH 1/2] feat(agent): agent-selectable GPU/CPU per execute-code call with project allowance MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - SandboxConfig gains allowed_gpus (canonical labels, validated) and max_timeout; shared resolver in services/compute_allowance.py used by both the handler and the prompt renderer so they never disagree. Effective allowance = {cpu} ∪ profile GPUs ∪ allowed_gpus; absent config keeps existing projects unchanged. - execute-code gains optional gpu/timeout args. Precedence: explicit gpu > heavy profile > default profile. Denied GPUs return a structured tool error naming the allowed set; timeouts clamp to the project cap. tool_start events now carry the gpu/timeout actually used. - _format_compute_env rewritten: renders the allowance with per-option guidance and ~$/hr from the active provider's rates, the timeout cap, and choice guidance — injected into the system prompt at session start for every execute-code agent. - Project Settings UI: allowed-GPU checkbox grid + max-timeout field; GPU_OPTIONS fixed to canonical labels (stale 'A100' broke billing — legacy labels are normalized leniently on read). - trainer.yaml / SKILL.md guidance updated to prefer explicit gpu=. Closes #146 Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_012GrGFBtdZv7qtgn8KQEQEB --- backend/agents/trainer.yaml | 10 +- backend/schemas.py | 38 +++- backend/services/agent/runner.py | 103 +++++++---- backend/services/compute_allowance.py | 92 ++++++++++ backend/skills/execute-code/SKILL.md | 18 +- backend/skills/execute-code/handler.py | 54 +++++- backend/skills/execute-code/schema.yaml | 20 +++ backend/skills/train-tabular/SKILL.md | 3 +- backend/tests/test_compute_allowance.py | 168 ++++++++++++++++++ backend/tests/test_execute_code_handler.py | 154 ++++++++++++++++ .../src/components/ProjectSettingsModal.tsx | 78 +++++++- frontend/src/lib/types.ts | 5 + 12 files changed, 689 insertions(+), 54 deletions(-) create mode 100644 backend/services/compute_allowance.py create mode 100644 backend/tests/test_compute_allowance.py create mode 100644 backend/tests/test_execute_code_handler.py diff --git a/backend/agents/trainer.yaml b/backend/agents/trainer.yaml index d8a1030..06fddf4 100644 --- a/backend/agents/trainer.yaml +++ b/backend/agents/trainer.yaml @@ -49,10 +49,12 @@ system: | If a file isn't where you expect, either scan `/data/sessions/{session_id}/` or use `inspect-agent-context` against the data_prep agent (see `list-session-agents`). - Project dataset files (reference only): `/data/projects/{project_id}/datasets/` - - ALWAYS set heavy=true on execute_code for training runs, hyperparameter - tuning, and large-scale transforms. This activates the project's training - sandbox profile (GPU + extended timeout). Use heavy=false only for quick - data inspection or plotting. + - For training runs, hyperparameter tuning, and large-scale transforms, + pick explicit compute: pass a `gpu=` value from the allowed list in your + "Compute environment" section, sized to the job (cheapest that fits), + plus a `timeout=` when the default is too short. When unsure, fall back + to heavy=true (the project's training profile: GPU + extended timeout). + Use heavy=false / gpu="cpu" for quick data inspection or plotting. ## Your workspace Your session workspace is `/data/sessions/{session_id}/`. You share this diff --git a/backend/schemas.py b/backend/schemas.py index 16f5f4b..ecba4d7 100644 --- a/backend/schemas.py +++ b/backend/schemas.py @@ -4,7 +4,7 @@ from typing import Literal, Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, field_validator # Generous-but-not-infinite caps to stop runaway inputs from swamping the # database or bloating an agent's context window. Calibrated so legitimate @@ -20,6 +20,21 @@ _GPU_MAX = 32 +# Canonical compute labels ("cpu" = no GPU). Must match the rate keys in +# services/sandbox.yml, the provider GPU mappings +# (services/compute/runpod_provider/gpu.py) and the execute-code / +# create-serving-app schema enums. +CANONICAL_GPUS: tuple[str, ...] = ( + "cpu", + "T4", + "L4", + "A10G", + "A100-40GB", + "A100-80GB", + "H100", +) + + class SandboxProfile(BaseModel): gpu: Optional[str] = Field(default=None, max_length=_GPU_MAX) timeout: Optional[int] = Field(default=None, ge=10, le=7200) @@ -28,6 +43,27 @@ class SandboxProfile(BaseModel): class SandboxConfig(BaseModel): default: Optional[SandboxProfile] = None training: Optional[SandboxProfile] = None + # GPUs the agent may explicitly request via execute-code's `gpu` arg. + # "cpu" and the profiles' GPUs are always implicitly allowed (the + # agent can already reach those via heavy=True/False) — this list + # widens the choice beyond the profiles. Absent/empty = profiles only. + allowed_gpus: Optional[list[str]] = Field(default=None, max_length=16) + # Hard cap on any agent-requested per-call timeout (seconds). Absent = + # capped at the largest owner-configured profile timeout. + max_timeout: Optional[int] = Field(default=None, ge=10, le=7200) + + @field_validator("allowed_gpus") + @classmethod + def _canonical_gpus(cls, v: Optional[list[str]]) -> Optional[list[str]]: + if v is None: + return v + v = list(dict.fromkeys(v)) # dedupe, preserve order + bad = [g for g in v if g not in CANONICAL_GPUS] + if bad: + raise ValueError( + f"Unknown GPU label(s) {bad}; allowed: {list(CANONICAL_GPUS)}" + ) + return v or None # [] normalizes to "not configured" class ExperimentCreate(BaseModel): diff --git a/backend/services/agent/runner.py b/backend/services/agent/runner.py index de39cd6..e174de2 100644 --- a/backend/services/agent/runner.py +++ b/backend/services/agent/runner.py @@ -244,54 +244,85 @@ async def _load_project_context(experiment_id: str) -> tuple[str, str, str, dict return project_id, project_name, files_listing, sandbox_config +# One-line hardware guidance per canonical label, shown next to each +# allowed option in the compute-environment prompt block. +_GPU_BLURBS: dict[str, str] = { + "cpu": "CPU only — EDA, plotting, sklearn/xgboost/lightgbm", + "T4": "16GB entry GPU — small fine-tunes, light inference", + "L4": "24GB — best price/perf for medium GPU work", + "A10G": "24GB mid-tier — solid training workhorse", + "A100-40GB": "40GB — large models, big batches", + "A100-80GB": "80GB — very large models / long contexts", + "H100": "80GB top-tier — only when speed or memory demands it", +} + + +def _gpu_hourly_usd(gpu: str) -> float | None: + """Approx $/hr for a canonical label on the active provider; None when + pricing is unavailable (the prompt then omits prices).""" + try: + from services.usage import _resolve_compute_rate + + rate = _resolve_compute_rate(settings.compute_provider, gpu) + return rate * 3600 if rate > 0 else None + except Exception: + return None + + def _format_compute_env(sandbox_config: dict) -> str: - """Render the project's per-profile sandbox config as a prompt block the - agent can read before deciding how to dimension execute-code calls. + """Render the agent's compute allowance as a prompt block: which + hardware it may request per execute-code call (`gpu` arg), the max + per-call timeout, and how the heavy/default profile fallback works. - Mirrors the runtime fallback in services/sandbox.py: - gpu = profile.get("gpu") or None → CPU only - timeout = profile.get("timeout") or settings.sandbox_timeout (default 600) + Uses the same resolver as the execute-code handler + (services/compute_allowance.py) so the prompt never advertises + hardware the handler would reject. """ + from services.compute_allowance import resolve_compute_allowance + + allowance = resolve_compute_allowance(sandbox_config) fallback_timeout = settings.sandbox_timeout - def _profile_line(label: str, profile: dict | None, default_to_used: int) -> str: - p = profile or {} - gpu = p.get("gpu") - timeout = p.get("timeout") or fallback_timeout - gpu_part = f"GPU={gpu}" if gpu else "CPU only (no GPU)" - timeout_part = f"timeout={timeout}s ({timeout // 60}m{timeout % 60:02d}s)" - return f" - **{label}**: {gpu_part}, {timeout_part}" + default_profile = sandbox_config.get("default") or {} + training_profile = sandbox_config.get("training") or {} + default_gpu = default_profile.get("gpu") or "cpu" + training_gpu = training_profile.get("gpu") or "cpu" + default_timeout = default_profile.get("timeout") or fallback_timeout + training_timeout = training_profile.get("timeout") or fallback_timeout - default_profile = sandbox_config.get("default") - training_profile = sandbox_config.get("training") + hw_lines = [] + for gpu in allowance.allowed_gpus: + blurb = _GPU_BLURBS.get(gpu, "") + price = _gpu_hourly_usd(gpu) + price_part = f" (~${price:.2f}/hr)" if price is not None else "" + hw_lines.append(f" - `{gpu}` — {blurb}{price_part}") lines = [ "## Compute environment for `execute-code`", "", - "Your sandbox is provisioned per call by Modal. Two profiles are", - "configured at the project level — pick the right one when you call", - "the skill:", + "Each call provisions a fresh sandbox. You choose the compute per", + "call with the optional `gpu` argument:", + "", + "**Allowed hardware** (values accepted for `gpu`):", + *hw_lines, "", - _profile_line( - "default profile (`heavy=False`, the default)", default_profile, 600 - ), - _profile_line("training profile (`heavy=True`)", training_profile, 1800), + f"**Timeout**: pass `timeout` (seconds, per call; max {allowance.max_timeout}s" + f" — higher values are clamped). Defaults: {default_timeout}s" + f" (default profile) / {training_timeout}s (`heavy=True`).", "", - "Dimension your code to fit:", - "- **Timeout is per call**, not per session. If a single fit / sweep", - " would exceed it, split the work across multiple calls (one fold,", - " one trial, one epoch chunk per call) and persist intermediate", - " state to the session workspace between calls.", - "- **No GPU configured for a profile** → don't import torch.cuda or", - " rely on `device='cuda'`. Stay on CPU-friendly libraries (xgboost,", - " lightgbm, sklearn) or use small models.", - "- **GPU configured** → free to use torch / GPU-accelerated paths.", - " Match batch size and model size to the GPU's memory class.", - "- Use `heavy=True` when calling `execute-code` for any work that", - " needs the training profile (long-running fit, hyperparameter sweep,", - " GPU-bound code). The default profile is for inspection / quick checks.", - "- The user can change these settings live in the Project Settings", - " modal — your next call will pick up the new values automatically.", + "**How to choose**:", + f"- Omit `gpu` → profile fallback: `heavy=False` = default profile" + f" ({default_gpu}), `heavy=True` = training profile ({training_gpu}).", + "- Prefer the cheapest hardware that fits. CPU for EDA / plots /", + " inspection; small GPUs for modest fine-tunes; big GPUs only when", + " memory or speed demands it. On CPU, don't rely on `device='cuda'`.", + "- **Timeout is per call**, not per session — split long fits across", + " calls (one fold / trial / epoch chunk each) and persist state to", + " the session workspace between calls.", + "- Requesting hardware outside the list returns an error naming the", + " allowed set.", + "- The user can change this allowance in Project Settings — your", + " next call picks up the new values automatically.", ] return "\n".join(lines) diff --git a/backend/services/compute_allowance.py b/backend/services/compute_allowance.py new file mode 100644 index 0000000..954e2ed --- /dev/null +++ b/backend/services/compute_allowance.py @@ -0,0 +1,92 @@ +"""Compute allowance — which GPUs/CPUs the agent may request, and how long +a single execution may run. + +One resolver, two consumers: the execute-code handler (enforcement) and +the system-prompt renderer in services/agent/runner.py (presentation). +Sharing it guarantees the prompt never advertises hardware the handler +would reject. + +Semantics: + allowed = {"cpu"} ∪ {profile GPUs} ∪ (sandbox_config.allowed_gpus or ∅) + Profile GPUs are implicitly allowed — the agent can already reach + them via heavy=True, so denying them on the explicit arg would be + security theater. "cpu" is always allowed (it's the floor). + max_timeout = sandbox_config.max_timeout + or max(profile timeouts, settings.sandbox_timeout) + so the agent can't self-grant a longer run than the owner ever + configured. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from config import settings +from schemas import CANONICAL_GPUS + +# Legacy / sloppy labels seen in old project configs and the pre-canonical +# frontend dropdown. Maps lowercase alias -> canonical label. +_LABEL_ALIASES: dict[str, str] = { + "a100": "A100-40GB", # the old frontend 'A100' option + "none": "cpu", +} + +_CANONICAL_BY_LOWER = {g.lower(): g for g in CANONICAL_GPUS} + + +def normalize_gpu(label: str | None) -> str | None: + """Case-insensitive map of a user/agent-supplied label to its + canonical form. Returns None for unknown labels.""" + if not label: + return None + key = str(label).strip().lower() + key = _LABEL_ALIASES.get(key, key).lower() + return _CANONICAL_BY_LOWER.get(key) + + +@dataclass(frozen=True) +class ComputeAllowance: + allowed_gpus: list[str] # canonical labels incl. "cpu", cost-ascending + max_timeout: int # seconds + explicit: bool # owner set allowed_gpus (vs profile-implied only) + + def permits(self, label: str) -> bool: + return label in self.allowed_gpus + + +def resolve_compute_allowance(sandbox_config: dict | None) -> ComputeAllowance: + cfg = sandbox_config or {} + profiles = [cfg.get("default") or {}, cfg.get("training") or {}] + + allowed: set[str] = {"cpu"} + for profile in profiles: + normalized = normalize_gpu(profile.get("gpu")) + if normalized: + allowed.add(normalized) + + configured = cfg.get("allowed_gpus") or None + if configured: + for label in configured: + normalized = normalize_gpu(label) + if normalized: + allowed.add(normalized) + + max_timeout = cfg.get("max_timeout") or max( + [p.get("timeout") or 0 for p in profiles] + [settings.sandbox_timeout] + ) + + order = {g: i for i, g in enumerate(CANONICAL_GPUS)} + return ComputeAllowance( + allowed_gpus=sorted(allowed, key=lambda g: order.get(g, 99)), + max_timeout=int(max_timeout), + explicit=bool(configured), + ) + + +def clamp_timeout(requested, allowance: ComputeAllowance) -> int | None: + """Clamp an agent-requested timeout into [10, allowance.max_timeout]. + Returns None for garbage input (caller falls back to the profile).""" + try: + return max(10, min(int(requested), allowance.max_timeout)) + except (TypeError, ValueError): + return None diff --git a/backend/skills/execute-code/SKILL.md b/backend/skills/execute-code/SKILL.md index 11d8af6..a30bdea 100644 --- a/backend/skills/execute-code/SKILL.md +++ b/backend/skills/execute-code/SKILL.md @@ -19,10 +19,20 @@ Use os.makedirs(path, exist_ok=True) before saving. Print all results to stdout. Each execution has a 10-minute timeout by default. -Set heavy=true for GPU-intensive workloads (model training, -hyperparameter tuning, large-scale data processing). This uses the -project's training sandbox profile which may have a GPU attached -and a longer timeout. +Choosing compute — three levers, in precedence order: + +- `gpu` (optional): explicitly pick hardware for this call from the + allowed list shown in your system prompt's "Compute environment" + section (e.g. `gpu="cpu"`, `gpu="L4"`). Overrides `heavy`. Values + outside the allowance return an error naming the allowed set. +- `heavy=true`: fall back to the project's training sandbox profile + (GPU + extended timeout) for GPU-intensive workloads (model training, + hyperparameter tuning, large-scale data processing). +- `timeout` (optional): per-call timeout in seconds, clamped to the + project's max. Use it instead of splitting work when a single fit + slightly exceeds the profile default. + +Prefer the cheapest hardware that fits the job. ## When to use Run any Python in an isolated Modal sandbox — EDA, modeling, validation. diff --git a/backend/skills/execute-code/handler.py b/backend/skills/execute-code/handler.py index ec21897..83a6b21 100644 --- a/backend/skills/execute-code/handler.py +++ b/backend/skills/execute-code/handler.py @@ -9,6 +9,11 @@ import modal.exception as modal_exc from services.compute.base import SandboxTimeoutError +from services.compute_allowance import ( + clamp_timeout, + normalize_gpu, + resolve_compute_allowance, +) from services.sandbox import run_code from services.skills.state import ( _known_files, @@ -74,8 +79,14 @@ def create_handler( async def handler(args: dict): code = args.get("code", "") if isinstance(args, dict) else str(args) heavy = args.get("heavy", False) if isinstance(args, dict) else False - - # Pick the right sandbox profile based on heavy flag + requested_gpu = args.get("gpu") if isinstance(args, dict) else None + requested_timeout = args.get("timeout") if isinstance(args, dict) else None + + # Compute selection precedence: explicit `gpu` arg > heavy profile + # > default profile. The explicit path is validated against the + # project's compute allowance (see services/compute_allowance.py — + # the same resolver renders the allowance into the system prompt). + allowance = resolve_compute_allowance(_sandbox_config) profile_key = "training" if heavy else "default" profile = _sandbox_config.get(profile_key) or {} gpu = profile.get("gpu") @@ -83,10 +94,43 @@ async def handler(args: dict): start = time.time() + if requested_gpu: + label = normalize_gpu(requested_gpu) + if label is None or not allowance.permits(label): + error_msg = ( + f"GPU '{requested_gpu}' is not available in this project. " + f"Allowed: {', '.join(allowance.allowed_gpus)}. " + f"Re-call execute-code with one of those values, or omit " + f"`gpu` to use the {profile_key} profile." + ) + await publish_fn( + session_id, + "tool_end", + {"tool": "execute_code", "output": error_msg, "duration": 0}, + role="tool", + ) + return { + "content": [{"type": "text", "text": error_msg}], + "is_error": True, + } + # "cpu" is an explicit request for no GPU; run_code(gpu=None) + # schedules on the CPU pool. + gpu = None if label == "cpu" else label + + if requested_timeout is not None: + clamped = clamp_timeout(requested_timeout, allowance) + if clamped is not None: + timeout = clamped + await publish_fn( session_id, "tool_start", - {"tool": "execute_code", "input": {"code": code[:500], "heavy": heavy}}, + { + "tool": "execute_code", + "input": {"code": code[:500], "heavy": heavy}, + "gpu": gpu or "cpu", + "timeout": timeout, + }, role="tool", ) @@ -130,7 +174,9 @@ async def handler(args: dict): f"The Python process was killed mid-execution and partial " f"output (if any) is lost. Options: (a) split the work into " f"smaller chunks, (b) reduce data size or iterations, " - f"(c) re-run with heavy=true for the GPU/training profile. " + f"(c) re-run with heavy=true for the GPU/training profile, " + f"(d) re-run with an explicit `timeout=` up to " + f"{allowance.max_timeout}s. " f"Underlying error: {e.__class__.__name__}" ) logger.warning("Sandbox timeout (session=%s): %s", session_id, e) diff --git a/backend/skills/execute-code/schema.yaml b/backend/skills/execute-code/schema.yaml index d1821a7..6284c3f 100644 --- a/backend/skills/execute-code/schema.yaml +++ b/backend/skills/execute-code/schema.yaml @@ -8,8 +8,28 @@ properties: description: 'Set to true for heavy ML workloads (training, tuning, large transforms). Uses the project''s training sandbox profile with GPU and extended timeout. Leave false for lightweight tasks (EDA, plotting, data inspection). + Ignored for GPU selection when `gpu` is given. ' default: false + gpu: + type: string + enum: [cpu, T4, L4, A10G, A100-40GB, A100-80GB, H100] + description: 'Optional. Explicitly choose the compute for THIS call, overriding + the profile picked by `heavy`. Only the values listed under "Compute + environment" in your system prompt are permitted for this project; anything + else returns an error naming the allowed set. Use "cpu" to force a CPU-only + sandbox (cheapest). Omit to fall back to the heavy/default profile. + + ' + timeout: + type: integer + minimum: 10 + maximum: 7200 + description: 'Optional. Per-call timeout in seconds. Values above the project''s + max (shown in your system prompt) are silently clamped down. Omit to use the + selected profile''s timeout. + + ' required: - code diff --git a/backend/skills/train-tabular/SKILL.md b/backend/skills/train-tabular/SKILL.md index bc8f0ab..751f36b 100644 --- a/backend/skills/train-tabular/SKILL.md +++ b/backend/skills/train-tabular/SKILL.md @@ -48,7 +48,8 @@ The report MUST include: - `scripts/lightgbm_baseline.py` — same, lightgbm flavor - `scripts/sweep_xgb.py` — Optuna sweep with a sensible default search space -Run them via execute_code (`heavy=true` is recommended for sweeps): +Run them via execute_code (`heavy=true` — or an explicit `gpu=` from your +allowed list — is recommended for sweeps): ```python import sys diff --git a/backend/tests/test_compute_allowance.py b/backend/tests/test_compute_allowance.py new file mode 100644 index 0000000..04040a4 --- /dev/null +++ b/backend/tests/test_compute_allowance.py @@ -0,0 +1,168 @@ +"""Tests for services/compute_allowance.py and the SandboxConfig allowance +schema — the shared resolver that both the execute-code handler and the +system-prompt renderer consume.""" + +import pytest +from pydantic import ValidationError + +from config import settings +from schemas import CANONICAL_GPUS, SandboxConfig +from services.compute_allowance import ( + clamp_timeout, + normalize_gpu, + resolve_compute_allowance, +) + + +class TestNormalizeGpu: + def test_canonical_passthrough(self): + for label in CANONICAL_GPUS: + assert normalize_gpu(label) == label + + def test_case_insensitive(self): + assert normalize_gpu("h100") == "H100" + assert normalize_gpu("CPU") == "cpu" + assert normalize_gpu("a100-80gb") == "A100-80GB" + + def test_legacy_aliases(self): + # The old frontend dropdown wrote 'A100' into project configs. + assert normalize_gpu("A100") == "A100-40GB" + assert normalize_gpu("none") == "cpu" + + def test_unknown_returns_none(self): + assert normalize_gpu("B9000") is None + assert normalize_gpu("") is None + assert normalize_gpu(None) is None + + +class TestResolveAllowance: + def test_absent_config_allows_cpu_only(self): + allowance = resolve_compute_allowance(None) + assert allowance.allowed_gpus == ["cpu"] + assert allowance.max_timeout == settings.sandbox_timeout + assert allowance.explicit is False + + def test_profile_gpus_implicitly_allowed(self): + allowance = resolve_compute_allowance( + {"default": {"gpu": None}, "training": {"gpu": "A10G"}} + ) + assert allowance.allowed_gpus == ["cpu", "A10G"] + assert allowance.explicit is False + + def test_explicit_list_unions_profiles_and_cpu(self): + allowance = resolve_compute_allowance( + { + "training": {"gpu": "A10G"}, + "allowed_gpus": ["T4", "H100"], + } + ) + assert allowance.allowed_gpus == ["cpu", "T4", "A10G", "H100"] + assert allowance.explicit is True + + def test_legacy_profile_label_normalized(self): + allowance = resolve_compute_allowance({"training": {"gpu": "A100"}}) + assert "A100-40GB" in allowance.allowed_gpus + + def test_result_sorted_cost_ascending(self): + allowance = resolve_compute_allowance( + {"allowed_gpus": ["H100", "cpu", "L4", "T4"]} + ) + assert allowance.allowed_gpus == ["cpu", "T4", "L4", "H100"] + + def test_max_timeout_defaults_to_largest_profile(self): + allowance = resolve_compute_allowance( + {"default": {"timeout": 300}, "training": {"timeout": 3600}} + ) + assert allowance.max_timeout == 3600 + + def test_max_timeout_never_below_settings_default(self): + allowance = resolve_compute_allowance({"default": {"timeout": 60}}) + assert allowance.max_timeout == settings.sandbox_timeout + + def test_explicit_max_timeout_wins(self): + allowance = resolve_compute_allowance( + {"training": {"timeout": 3600}, "max_timeout": 900} + ) + assert allowance.max_timeout == 900 + + +class TestClampTimeout: + def _allowance(self, max_timeout=1000): + return resolve_compute_allowance({"max_timeout": max_timeout}) + + def test_above_cap_clamped_down(self): + assert clamp_timeout(5000, self._allowance()) == 1000 + + def test_below_floor_raised_to_ten(self): + assert clamp_timeout(1, self._allowance()) == 10 + + def test_in_range_passthrough(self): + assert clamp_timeout(500, self._allowance()) == 500 + + def test_garbage_returns_none(self): + assert clamp_timeout("soon", self._allowance()) is None + assert clamp_timeout(None, self._allowance()) is None + + +class TestSandboxConfigSchema: + def test_unknown_label_rejected(self): + with pytest.raises(ValidationError, match="Unknown GPU label"): + SandboxConfig(allowed_gpus=["T4", "GTX-9090"]) + + def test_dedupe_preserves_order(self): + cfg = SandboxConfig(allowed_gpus=["T4", "H100", "T4"]) + assert cfg.allowed_gpus == ["T4", "H100"] + + def test_empty_list_normalizes_to_none(self): + assert SandboxConfig(allowed_gpus=[]).allowed_gpus is None + + def test_max_timeout_bounds(self): + with pytest.raises(ValidationError): + SandboxConfig(max_timeout=5) + with pytest.raises(ValidationError): + SandboxConfig(max_timeout=10_000) + assert SandboxConfig(max_timeout=600).max_timeout == 600 + + def test_old_shape_still_valid(self): + cfg = SandboxConfig( + default={"gpu": None, "timeout": 600}, + training={"gpu": "A10G", "timeout": 1800}, + ) + assert cfg.allowed_gpus is None + assert cfg.max_timeout is None + + +class TestPromptRendering: + def test_renders_allowed_labels_and_cap(self): + from services.agent.runner import _format_compute_env + + block = _format_compute_env( + { + "default": {"gpu": None, "timeout": 600}, + "training": {"gpu": "A10G", "timeout": 1800}, + "allowed_gpus": ["T4"], + "max_timeout": 3600, + } + ) + assert "`cpu`" in block + assert "`T4`" in block + assert "`A10G`" in block + assert "`H100`" not in block + assert "max 3600s" in block + assert "heavy=True" in block + + def test_prices_omitted_gracefully(self, monkeypatch): + import services.agent.runner as runner_mod + + monkeypatch.setattr( + runner_mod, + "_gpu_hourly_usd", + lambda gpu: (_ for _ in ()).throw(RuntimeError("no rates")), + ) + # A rate failure must not break prompt assembly. + with pytest.raises(RuntimeError): + runner_mod._gpu_hourly_usd("T4") + monkeypatch.setattr(runner_mod, "_gpu_hourly_usd", lambda gpu: None) + block = runner_mod._format_compute_env({}) + assert "$" not in block + assert "`cpu`" in block diff --git a/backend/tests/test_execute_code_handler.py b/backend/tests/test_execute_code_handler.py new file mode 100644 index 0000000..da728fb --- /dev/null +++ b/backend/tests/test_execute_code_handler.py @@ -0,0 +1,154 @@ +"""Tests for the execute-code handler's agent-selectable compute — the +`gpu`/`timeout` args, allowance enforcement, and profile precedence.""" + +import importlib.util +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest + +HANDLER_PATH = ( + Path(__file__).resolve().parents[1] / "skills" / "execute-code" / "handler.py" +) + + +def _load_handler_module(): + spec = importlib.util.spec_from_file_location("execute_code_handler", HANDLER_PATH) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +@pytest.fixture +def handler_mod(monkeypatch): + mod = _load_handler_module() + mod.run_code = AsyncMock( + return_value={"stdout": "ok", "stderr": "", "returncode": 0} + ) + mod.write_to_volume = AsyncMock() + mod.detect_new_files = AsyncMock() + return mod + + +def _make_handler(mod, sandbox_config=None): + publish = AsyncMock() + handler = mod.create_handler( + session_id="sess-1", + stage="train", + publish_fn=publish, + sandbox_config=sandbox_config, + ) + return handler, publish + + +CONFIG = { + "default": {"gpu": None, "timeout": 600}, + "training": {"gpu": "A10G", "timeout": 1800}, + "allowed_gpus": ["T4", "L4"], + "max_timeout": 3600, +} + + +class TestGpuSelection: + @pytest.mark.asyncio + async def test_allowed_gpu_passes_through(self, handler_mod): + handler, _ = _make_handler(handler_mod, CONFIG) + result = await handler({"code": "print(1)", "gpu": "T4"}) + assert not result.get("is_error") + assert handler_mod.run_code.await_args.kwargs["gpu"] == "T4" + + @pytest.mark.asyncio + async def test_cpu_maps_to_none(self, handler_mod): + handler, _ = _make_handler(handler_mod, CONFIG) + await handler({"code": "print(1)", "gpu": "cpu", "heavy": True}) + # Explicit cpu overrides the heavy profile's A10G. + assert handler_mod.run_code.await_args.kwargs["gpu"] is None + + @pytest.mark.asyncio + async def test_case_insensitive_label(self, handler_mod): + handler, _ = _make_handler(handler_mod, CONFIG) + await handler({"code": "print(1)", "gpu": "t4"}) + assert handler_mod.run_code.await_args.kwargs["gpu"] == "T4" + + @pytest.mark.asyncio + async def test_denied_gpu_returns_structured_error(self, handler_mod): + handler, publish = _make_handler(handler_mod, CONFIG) + result = await handler({"code": "print(1)", "gpu": "H100"}) + assert result["is_error"] is True + text = result["content"][0]["text"] + assert "H100" in text + assert "cpu, T4, L4, A10G" in text + handler_mod.run_code.assert_not_awaited() + + @pytest.mark.asyncio + async def test_profile_gpu_implicitly_allowed(self, handler_mod): + # A10G comes from the training profile, not allowed_gpus. + handler, _ = _make_handler(handler_mod, CONFIG) + result = await handler({"code": "print(1)", "gpu": "A10G"}) + assert not result.get("is_error") + assert handler_mod.run_code.await_args.kwargs["gpu"] == "A10G" + + +class TestPrecedence: + @pytest.mark.asyncio + async def test_explicit_gpu_beats_heavy_profile(self, handler_mod): + handler, _ = _make_handler(handler_mod, CONFIG) + await handler({"code": "x", "heavy": True, "gpu": "L4"}) + assert handler_mod.run_code.await_args.kwargs["gpu"] == "L4" + + @pytest.mark.asyncio + async def test_heavy_without_gpu_uses_training_profile(self, handler_mod): + handler, _ = _make_handler(handler_mod, CONFIG) + await handler({"code": "x", "heavy": True}) + kwargs = handler_mod.run_code.await_args.kwargs + assert kwargs["gpu"] == "A10G" + assert kwargs["timeout"] == 1800 + + @pytest.mark.asyncio + async def test_neither_uses_default_profile(self, handler_mod): + handler, _ = _make_handler(handler_mod, CONFIG) + await handler({"code": "x"}) + kwargs = handler_mod.run_code.await_args.kwargs + assert kwargs["gpu"] is None + assert kwargs["timeout"] == 600 + + +class TestTimeout: + @pytest.mark.asyncio + async def test_timeout_above_cap_clamped(self, handler_mod): + handler, _ = _make_handler(handler_mod, CONFIG) + await handler({"code": "x", "timeout": 99999}) + assert handler_mod.run_code.await_args.kwargs["timeout"] == 3600 + + @pytest.mark.asyncio + async def test_timeout_below_floor_raised(self, handler_mod): + handler, _ = _make_handler(handler_mod, CONFIG) + await handler({"code": "x", "timeout": 3}) + assert handler_mod.run_code.await_args.kwargs["timeout"] == 10 + + @pytest.mark.asyncio + async def test_garbage_timeout_falls_back_to_profile(self, handler_mod): + handler, _ = _make_handler(handler_mod, CONFIG) + await handler({"code": "x", "timeout": "soon"}) + assert handler_mod.run_code.await_args.kwargs["timeout"] == 600 + + +class TestEventPayload: + @pytest.mark.asyncio + async def test_tool_start_carries_gpu_and_timeout(self, handler_mod): + handler, publish = _make_handler(handler_mod, CONFIG) + await handler({"code": "x", "gpu": "T4", "timeout": 120}) + start_calls = [c for c in publish.await_args_list if c.args[1] == "tool_start"] + payload = start_calls[0].args[2] + assert payload["gpu"] == "T4" + assert payload["timeout"] == 120 + + +class TestNoConfig: + @pytest.mark.asyncio + async def test_absent_config_only_cpu_allowed(self, handler_mod): + handler, _ = _make_handler(handler_mod, None) + result = await handler({"code": "x", "gpu": "T4"}) + assert result["is_error"] is True + ok = await handler({"code": "x", "gpu": "cpu"}) + assert not ok.get("is_error") diff --git a/frontend/src/components/ProjectSettingsModal.tsx b/frontend/src/components/ProjectSettingsModal.tsx index c0087f0..9fc1a47 100644 --- a/frontend/src/components/ProjectSettingsModal.tsx +++ b/frontend/src/components/ProjectSettingsModal.tsx @@ -4,14 +4,23 @@ import { useEffect, useState } from 'react'; import { Settings, X } from 'lucide-react'; import type { SandboxConfig, SandboxProfile } from '@/lib/types'; +// Canonical labels — must match backend/schemas.py CANONICAL_GPUS and the +// (provider, gpu) rate keys in backend/services/sandbox.yml. The old +// 'A100' value silently broke billing (rates are keyed 'A100-40GB'). const GPU_OPTIONS = [ { value: '', label: 'None (CPU only)' }, { value: 'T4', label: 'T4 — 16 GB' }, { value: 'L4', label: 'L4 — 24 GB' }, { value: 'A10G', label: 'A10G — 24 GB' }, - { value: 'A100', label: 'A100 — 40 GB' }, + { value: 'A100-40GB', label: 'A100 — 40 GB' }, + { value: 'A100-80GB', label: 'A100 — 80 GB' }, + { value: 'H100', label: 'H100 — 80 GB' }, ]; +// Options for the agent allowance checkboxes ("cpu" is always allowed and +// shown as locked). +const ALLOWANCE_OPTIONS = GPU_OPTIONS.filter((o) => o.value !== ''); + interface Props { isOpen: boolean; projectName: string; @@ -85,6 +94,8 @@ export default function ProjectSettingsModal({ const [defaultTimeout, setDefaultTimeout] = useState(600); const [trainingGpu, setTrainingGpu] = useState(''); const [trainingTimeout, setTrainingTimeout] = useState(1800); + const [allowedGpus, setAllowedGpus] = useState([]); + const [maxTimeout, setMaxTimeout] = useState(''); useEffect(() => { if (isOpen) { @@ -94,9 +105,17 @@ export default function ProjectSettingsModal({ setDefaultTimeout(d?.timeout ?? 600); setTrainingGpu(t?.gpu || ''); setTrainingTimeout(t?.timeout ?? 1800); + setAllowedGpus(sandboxConfig.allowed_gpus ?? []); + setMaxTimeout(sandboxConfig.max_timeout ?? ''); } }, [isOpen, sandboxConfig]); + const toggleAllowedGpu = (value: string) => { + setAllowedGpus((prev) => + prev.includes(value) ? prev.filter((g) => g !== value) : [...prev, value] + ); + }; + useEffect(() => { if (!isOpen) return; const onKey = (e: KeyboardEvent) => { @@ -122,6 +141,8 @@ export default function ProjectSettingsModal({ onSave({ default: buildProfile(defaultGpu, defaultTimeout), training: buildProfile(trainingGpu, trainingTimeout), + allowed_gpus: allowedGpus.length ? allowedGpus : null, + max_timeout: maxTimeout === '' ? null : maxTimeout, }); onClose(); }; @@ -155,7 +176,7 @@ export default function ProjectSettingsModal({ {/* Body */}

- Modal Sandbox + Compute Sandbox

+
+ + {/* Agent GPU allowance */} +
+
+

Agent GPU allowance

+ + hardware the agent may request per call + +
+
+ + {ALLOWANCE_OPTIONS.map((opt) => ( + + ))} +
+
+ + + setMaxTimeout(e.target.value === '' ? '' : Number(e.target.value)) + } + className="w-full px-2.5 py-1.5 rounded-lg bg-white/[0.04] border border-white/[0.08] text-xs text-white focus:outline-none focus:border-blue-500/50 transition-colors" + /> +
+
+

- Agents automatically select the right profile. The training profile is used when{' '} - heavy=true is set on code execution. + Agents pick the cheapest allowed hardware per call via the{' '} + gpu argument; the training profile + remains the heavy=true fallback. + Profile GPUs are always implicitly allowed.

diff --git a/frontend/src/lib/types.ts b/frontend/src/lib/types.ts index a1cd5ac..60977a2 100644 --- a/frontend/src/lib/types.ts +++ b/frontend/src/lib/types.ts @@ -6,6 +6,11 @@ export interface SandboxProfile { export interface SandboxConfig { default?: SandboxProfile | null; training?: SandboxProfile | null; + // GPUs the agent may explicitly request per execute-code call. + // "cpu" and the profiles' GPUs are always implicitly allowed. + allowed_gpus?: string[] | null; + // Hard cap (seconds) on agent-requested per-call timeouts. + max_timeout?: number | null; } export interface Project { From 7175a5680c16dcc7c1e553af105d6e67e35f7d1d Mon Sep 17 00:00:00 2001 From: Lucas Tonon Date: Wed, 29 Jul 2026 16:57:42 -0300 Subject: [PATCH 2/2] =?UTF-8?q?fix(agent):=20gpu-choice=20review=20follow-?= =?UTF-8?q?ups=20=E2=80=94=20tool=20event=20order,=20owner=20timeout=20cap?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - GPU denial now emits tool_start before tool_end (every other handler path does; the UI's active-tool state machine depends on the pair) - an owner-set max_timeout now caps profile-fallback timeouts too, not just the explicit timeout= arg (heavy=True without timeout= used to bypass it) + regression tests - _gpu_hourly_usd logs when pricing resolution fails instead of silently dropping prices from the prompt --- backend/services/agent/runner.py | 6 +++- backend/skills/execute-code/handler.py | 21 +++++++++++++ backend/tests/test_execute_code_handler.py | 36 ++++++++++++++++++++++ 3 files changed, 62 insertions(+), 1 deletion(-) diff --git a/backend/services/agent/runner.py b/backend/services/agent/runner.py index e174de2..d7796a4 100644 --- a/backend/services/agent/runner.py +++ b/backend/services/agent/runner.py @@ -261,11 +261,15 @@ def _gpu_hourly_usd(gpu: str) -> float | None: """Approx $/hr for a canonical label on the active provider; None when pricing is unavailable (the prompt then omits prices).""" try: + # Private import on purpose: sandbox.yml rate resolution has no + # public API yet. A signature/name change there degrades to + # price-less prompts — logged below so it isn't invisible. from services.usage import _resolve_compute_rate rate = _resolve_compute_rate(settings.compute_provider, gpu) return rate * 3600 if rate > 0 else None - except Exception: + except Exception as e: + logger.debug("GPU pricing unavailable for %s: %s", gpu, e) return None diff --git a/backend/skills/execute-code/handler.py b/backend/skills/execute-code/handler.py index 83a6b21..47a22b8 100644 --- a/backend/skills/execute-code/handler.py +++ b/backend/skills/execute-code/handler.py @@ -8,6 +8,7 @@ import modal.exception as modal_exc +from config import settings from services.compute.base import SandboxTimeoutError from services.compute_allowance import ( clamp_timeout, @@ -103,6 +104,20 @@ async def handler(args: dict): f"Re-call execute-code with one of those values, or omit " f"`gpu` to use the {profile_key} profile." ) + # tool_start before tool_end — every other path in this + # handler emits the pair in order, and the UI's active-tool + # state machine depends on it. + await publish_fn( + session_id, + "tool_start", + { + "tool": "execute_code", + "input": {"code": code[:500], "heavy": heavy}, + "gpu": requested_gpu, + "timeout": timeout, + }, + role="tool", + ) await publish_fn( session_id, "tool_end", @@ -122,6 +137,12 @@ async def handler(args: dict): if clamped is not None: timeout = clamped + # An owner-set max_timeout caps EVERY execution — including the + # profile fallback (heavy=True without an explicit `timeout=`) and + # the sandbox default. When max_timeout is profile-derived this is + # a no-op (it is the max of the profile timeouts and the default). + timeout = min(timeout or settings.sandbox_timeout, allowance.max_timeout) + await publish_fn( session_id, "tool_start", diff --git a/backend/tests/test_execute_code_handler.py b/backend/tests/test_execute_code_handler.py index da728fb..3eb07e3 100644 --- a/backend/tests/test_execute_code_handler.py +++ b/backend/tests/test_execute_code_handler.py @@ -143,6 +143,42 @@ async def test_tool_start_carries_gpu_and_timeout(self, handler_mod): assert payload["gpu"] == "T4" assert payload["timeout"] == 120 + @pytest.mark.asyncio + async def test_denied_gpu_emits_tool_start_before_tool_end(self, handler_mod): + """The UI's active-tool state machine needs the start/end pair in + order — a denial must not emit an orphaned tool_end.""" + handler, publish = _make_handler(handler_mod, CONFIG) + await handler({"code": "x", "gpu": "H100"}) + events = [ + c.args[1] + for c in publish.await_args_list + if c.args[1] in ("tool_start", "tool_end") + ] + assert events == ["tool_start", "tool_end"] + + +class TestOwnerMaxTimeout: + """An owner-set max_timeout caps EVERY execution, including the + profile fallback (heavy=True without an explicit `timeout=`).""" + + OWNER_CAPPED = { + "default": {"gpu": None, "timeout": 600}, + "training": {"gpu": "A10G", "timeout": 1800}, + "max_timeout": 100, + } + + @pytest.mark.asyncio + async def test_profile_timeout_capped_by_owner_max(self, handler_mod): + handler, _ = _make_handler(handler_mod, self.OWNER_CAPPED) + await handler({"code": "x", "heavy": True}) + assert handler_mod.run_code.await_args.kwargs["timeout"] == 100 + + @pytest.mark.asyncio + async def test_default_fallback_capped_by_owner_max(self, handler_mod): + handler, _ = _make_handler(handler_mod, {"max_timeout": 100}) + await handler({"code": "x"}) + assert handler_mod.run_code.await_args.kwargs["timeout"] == 100 + class TestNoConfig: @pytest.mark.asyncio