agent.py

"""
Experiment agent for selecting benchmark configurations.
 
The agent proposes the next configuration to evaluate, tracks completed
trials, and selects the best configuration under a given constraint.
 
It supports two proposal strategies:
- LLM-based selection using the Hugging Face OpenAI-compatible router when
  HF_TOKEN is available.
- A deterministic heuristic fallback that orders candidates by workload
  and compile mode when the LLM is unavailable or fails.
 
The fallback ensures the experiment can run end-to-end without an LLM.
"""
from __future__ import annotations

import itertools
import json
import os
from typing import List, Optional

import httpx
from dotenv import load_dotenv
from openai import AsyncOpenAI

from config import BenchmarkResult, CompileMode, DType, ExperimentConfig

load_dotenv()


def _build_candidate_grid() -> List[ExperimentConfig]:
    batch_sizes = (1, 8, 32)
    seq_lens = (128, 512)
    dtypes = (DType.FP16, DType.BF16)
    compile_modes = (CompileMode.EAGER, CompileMode.DEFAULT, CompileMode.MAX_AUTOTUNE)
    return [
        ExperimentConfig(batch_size=bs, seq_len=sl, dtype=dt, compile_mode=cm)
        for bs, sl, dt, cm in itertools.product(batch_sizes, seq_lens, dtypes, compile_modes)
    ]


CANDIDATE_GRID = _build_candidate_grid()


class ExperimentAgent:
    def __init__(
        self,
        constraint: str,
        trial_budget: int = 8,
        use_llm: bool = True,
        model_name: Optional[str] = None,
    ):
        """Initialize the experiment agent.
 
        Args:
            constraint: Constraint used when evaluating experiment results.
            trial_budget: Configured budget for the experiment. Not enforced
                internally by this class (propose_next will keep proposing
                candidates until the grid is exhausted); the caller is
                responsible for stopping after `trial_budget` trials, as
                main.py does.
            use_llm: Whether to use the LLM proposal strategy when available.
            model_name: Optional model identifier for the LLM proposal
                strategy.
        """
        self.constraint = constraint
        self.trial_budget = trial_budget
        self.results: List[BenchmarkResult] = []
        self._tried_keys: set[str] = set()

        self.use_llm = use_llm and bool(os.getenv("HF_TOKEN"))
        self._client: Optional[AsyncOpenAI] = None
        self._model_name = model_name or os.getenv(
            "AGENT_MODEL", "meta-llama/Llama-3.1-8B-Instruct"
        )
        if self.use_llm:
            http_client = httpx.AsyncClient(timeout=httpx.Timeout(20.0, connect=5.0))
            self._client = AsyncOpenAI(
                http_client=http_client,
                base_url=os.getenv("AGENT_BASE_URL", "https://router.huggingface.co/v1"),
                api_key=os.getenv("HF_TOKEN"),
            )

    async def propose_next(self) -> Optional[ExperimentConfig]:
        """Propose the next untried experiment configuration.
 
        Uses the LLM strategy when enabled and available. If the LLM is
        unavailable or the proposal fails, falls back to the deterministic
        heuristic strategy.
 
        Returns:
            The next configuration to evaluate, or None if all candidates
            have already been tried.
        """
        remaining = [c for c in CANDIDATE_GRID if c.key() not in self._tried_keys]
        if not remaining:
            return None

        if self.use_llm:
            try:
                return await self._propose_with_llm(remaining)
            except Exception as e:  
                print(f"[agent] LLM proposal failed ({e}); falling back to heuristic search.")

        return self._propose_heuristic(remaining)

    def _propose_heuristic(self, remaining: List[ExperimentConfig]) -> ExperimentConfig:
        remaining = sorted(
            remaining,
            key=lambda c: (c.batch_size, c.seq_len, c.compile_mode != CompileMode.EAGER),
        )
        return remaining[0]

    async def _propose_with_llm(self, remaining: List[ExperimentConfig]) -> ExperimentConfig:
        """Select the next configuration using the LLM proposal policy.
 
        Args:
            remaining: Untried experiment configurations.
 
        Returns:
            The configuration selected by the LLM.
 
        Raises:
            Exception: If the LLM request or response parsing fails. Callers
                (propose_next) catch this broadly and fall back to the
                heuristic policy rather than crash the run.
        """
        assert self._client is not None
        history = [r.model_dump(mode="json") for r in self.results]
        candidates = [c.model_dump(mode="json") for c in remaining[:12]]

        system = (
            "You are a GPU inference optimization agent tuning a PyTorch model's "
            "runtime configuration. Your objective/constraint is:\n"
            f"{self.constraint}\n\n"
            "You are given the results of trials so far and a list of untried "
            "candidate configs. Choose exactly ONE candidate most likely to be "
            "informative for the objective (e.g. fill a gap in what's been tried, "
            "or push toward the constraint). "
            "Respond with ONLY a JSON object with keys: batch_size, seq_len, dtype, "
            "compile_mode, use_triton_softmax. No prose, no markdown fences."
        )
        user = json.dumps({"history": history, "candidates": candidates})

        resp = await self._client.chat.completions.create(
            model=self._model_name,
            messages=[
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
            temperature=0.2,
        )
        content = (resp.choices[0].message.content or "").strip()
        content = content.strip("`")
        if content.lower().startswith("json"):
            content = content[4:].strip()
        picked = json.loads(content)
        config = ExperimentConfig(**picked)

        # Guard against a hallucinated / duplicate pick: the caller's contract
        # (see propose_next / main.py's trial loop) is that every proposal is
        # untried. Without this check a bad LLM response could silently repeat
        # a config or invent one outside the candidate grid.
        remaining_keys = {c.key() for c in remaining}
        if config.key() not in remaining_keys:
            raise ValueError(
                f"LLM picked a config that is not an untried candidate: {config.key()}"
            )
        return config

    def record(self, result: BenchmarkResult) -> None:
        self.results.append(result)
        self._tried_keys.add(result.config.key())

    def select_best(self) -> Optional[BenchmarkResult]:
        """Pick the best trial under the stated constraint.

        This MVP objective is: lowest latency among trials that ran
        successfully and passed the correctness check. Swap this out (or make
        it constraint-aware, e.g. filter by peak_memory_mb / throughput
        thresholds parsed from `self.constraint`) as the search space grows.
        """
        valid = [r for r in self.results if r.ok and r.correctness_passed is not False]
        if not valid:
            return None
        return min(valid, key=lambda r: r.mean_latency_ms)