feat(agent): agent-selectable GPU/CPU per execute-code call with project allowance

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

backend/schemas.py

@@ -21,6 +21,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)
@@ -29,6 +44,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"
 
 
 # Model families a user can restrict the trainer to. Mirrors the `framework`

backend/services/agent/runner.py

@@ -357,54 +357,89 @@ def _format_training_constraints(training_config: dict) -> str:
     return "\n".join(lines)
 
 
+# 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:
+        # 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 as e:
+        logger.debug("GPU pricing unavailable for %s: %s", gpu, e)
+        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)
 

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

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.