scripts/train_grpo.py

"""GRPO training for OpsGuard via TRL's environment_factory integration.

Mirrors the official TRL openenv wordle/sudoku pattern:
  - one env class per generation
  - environment_factory=<class> on GRPOTrainer
  - reward_funcs reads env.reward set during the rollout

Our env exposes a single tool method `triage(...)` that takes the structured
fields of OpsguardAction. We deliberately use ONE method (not one-per-action)
because action_type is a tagged-union — splitting it would require 9 nearly
identical methods that all just dispatch to the same env.step().

This script supports two modes:
  1. Plain TRL + vLLM (matches the wordle.py reference) — most reliable.
  2. Unsloth FastLanguageModel preload — faster 4-bit/LoRA path. Note that
     Unsloth's GRPO support sometimes lags TRL's environment_factory API; if
     it fails, fall back to mode 1.

Run on HF Jobs:
    hf jobs uv run --flavor a100-large \\
        --with "trl>=0.18,unsloth,openenv-core,peft,bitsandbytes,vllm" \\
        --secrets HF_TOKEN \\
        scripts/train_grpo.py \\
        --model unsloth/Qwen2.5-7B-Instruct-bnb-4bit \\
        --env-url http://0.0.0.0:8001 \\
        --hub-repo me/opsguard-grpo \\
        --num-steps 200 --num-generations 4

Run locally (single GPU, vLLM colocate):
    # 1. start the env server
    uv run --with openenv-core uvicorn server.app:app --host 0.0.0.0 --port 8001
    # 2. run training
    python scripts/train_grpo.py --env-url http://0.0.0.0:8001
"""
from __future__ import annotations

import argparse
import inspect
import os
import sys
from pathlib import Path
from typing import Any

# Make project root importable when run as `python scripts/train_grpo.py`
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(_PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(_PROJECT_ROOT))

# These are CPU-safe imports (pure python).
from models import ActionType, OpsguardAction  # noqa: E402
try:
    from scripts.system_prompt import SYSTEM_PROMPT, format_observation  # noqa: E402
except ImportError:
    sys.path.insert(0, str(_PROJECT_ROOT / "scripts"))
    from system_prompt import SYSTEM_PROMPT, format_observation  # type: ignore  # noqa: E402


def _safe_kwargs(cls_or_callable, kwargs: dict) -> dict:
    """Drop kwargs not accepted by the target signature.

    TRL's GRPOConfig / GRPOTrainer signatures drift between releases; this is
    the same trick used in handoff §6 — instead of pinning a TRL version, we
    introspect at runtime and silently drop unknown kwargs (after warning).
    """
    try:
        sig = inspect.signature(cls_or_callable)
        params = sig.parameters
    except (TypeError, ValueError):
        return kwargs
    has_var_kw = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
    if has_var_kw:
        return kwargs
    valid = set(params.keys())
    out = {}
    dropped = []
    for k, v in kwargs.items():
        if k in valid:
            out[k] = v
        else:
            dropped.append(k)
    if dropped:
        print(f"  [WARN] dropping unsupported kwargs for {getattr(cls_or_callable, '__name__', cls_or_callable)}: {dropped}",
              flush=True)
    return out


# ---------------------------------------------------------------------------
# OpsGuard tool-env wrapper (one per generation, per TRL doc)
# ---------------------------------------------------------------------------
def make_opsguard_env_class(env_url: str, max_steps: int):
    """Factory that builds a class, capturing env_url/max_steps in closure.

    TRL invokes `environment_factory()` per generation; the factory must be
    callable with no args. Capturing config in the enclosing scope (the
    pattern recommended by the TRL OpenEnv docs) is how we inject env_url.
    """
    # Deferred import — only available once openenv-opsguard is installed.
    from client import OpsguardEnv  # type: ignore

    class OpsguardToolEnv:
        """One instance per parallel generation. Wraps the HTTP env client."""

        def __init__(self):
            # Lazy: create the websocket client per episode (in reset), to
            # avoid stale connections when TRL re-uses the instance.
            self._client: OpsguardEnv | None = None
            self.reward = 0.0
            self.cumulative_reward = 0.0
            self.done = False
            self._steps = 0
            self._initial_obs: str | None = None

        def _ensure_client(self):
            if self._client is None:
                self._client = OpsguardEnv(base_url=env_url)
            return self._client

        def _close_client(self):
            if self._client is not None:
                try:
                    close = getattr(self._client, "close", None)
                    if callable(close):
                        close()
                except Exception:
                    pass
                self._client = None

        def reset(self, **kwargs) -> str | None:
            """Called by TRL at episode start. Returns initial observation string."""
            self._close_client()
            self.reward = 0.0
            self.cumulative_reward = 0.0
            self.done = False
            self._steps = 0
            client = self._ensure_client()

            scenario_id = kwargs.get("scenario_id")
            seed = kwargs.get("seed")
            try:
                if scenario_id is not None or seed is not None:
                    result = client.reset(scenario_id=scenario_id, seed=seed)
                else:
                    result = client.reset()
            except TypeError:
                # Older client signatures don't accept kwargs
                result = client.reset()

            obs = result.observation if hasattr(result, "observation") else result
            self._initial_obs = format_observation(obs)
            return self._initial_obs

        def triage(
            self,
            action_type: str,
            target_issue_id: int = 0,
            label: str = "",
            duplicate_of_id: int = 0,
            assignee_login: str = "",
            comment_body: str = "",
            query: str = "",
            reasoning: str = "",
        ) -> str:
            """
            Take ONE triage action against the current issue in the queue and
            advance to the next.

            Args:
                action_type: One of 'label', 'close_spam', 'request_info',
                    'link_duplicate', 'assign', 'comment', 'merge_pr',
                    'query_history', 'wait'.
                target_issue_id: The issue_id of the current issue (0 to use current).
                label: Label name (only for 'label' action).
                duplicate_of_id: Existing issue id (only for 'link_duplicate').
                assignee_login: Maintainer login (only for 'assign').
                comment_body: Comment text (only for 'comment'/'request_info').
                query: Free-text search query (only for 'query_history').
                reasoning: Short rationale shown to the grader.

            Returns:
                The next observation rendered as JSON, or 'Episode done.' on terminal.
            """
            if self.done:
                raise ValueError("Episode is over. Stop calling tools.")

            try:
                at = ActionType(action_type.lower().strip())
            except (ValueError, AttributeError):
                at = ActionType.WAIT
                reasoning = f"unknown action_type={action_type!r}; defaulting to wait"

            action = OpsguardAction(
                action_type=at,
                target_issue_id=target_issue_id or None,
                label=label or None,
                duplicate_of_id=duplicate_of_id or None,
                assignee_login=assignee_login or None,
                comment_body=comment_body or None,
                query=query or None,
                reasoning=reasoning or None,
            )

            client = self._ensure_client()
            result = client.step(action)

            obs = result.observation if hasattr(result, "observation") else result
            step_reward = float(result.reward or 0.0) if hasattr(result, "reward") else 0.0
            self.reward = step_reward  # last-step reward (for shaping experiments)
            self.cumulative_reward += step_reward
            self._steps += 1

            self.done = bool(getattr(result, "done", False)) or self._steps >= max_steps
            if self.done:
                # Final state: include cumulative reward in the message, no further tool call
                return (f"Episode done. cumulative_reward={self.cumulative_reward:.3f} "
                        f"steps={self._steps}.")
            return format_observation(obs)

    return OpsguardToolEnv


# ---------------------------------------------------------------------------
# Reward function
# ---------------------------------------------------------------------------
def make_reward_func(reduce: str = "cumulative"):
    """Return a TRL-compatible reward function.

    `reduce`:
      - 'cumulative' : sum of all step rewards in the episode
      - 'last'       : reward of the final step only
      - 'mean'       : cumulative / steps
    """
    def reward_func(environments, **kwargs):
        out = []
        for env in environments:
            cum = float(getattr(env, "cumulative_reward", 0.0))
            steps = max(1, int(getattr(env, "_steps", 1)))
            last = float(getattr(env, "reward", 0.0))
            if reduce == "last":
                out.append(last)
            elif reduce == "mean":
                out.append(cum / steps)
            else:
                out.append(cum)
        return out
    return reward_func


# ---------------------------------------------------------------------------
# Dataset
# ---------------------------------------------------------------------------
def build_dataset(n_samples: int, scenarios: list[str]) -> Any:
    """One row per episode. The user message is just the system prompt body —
    the *actual* observation comes from env.reset() being returned as the
    first turn (TRL appends it as a tool/observation message)."""
    from datasets import Dataset  # type: ignore

    user_kickoff = (
        "You are about to start a triage session. Call the `triage` tool with "
        "exactly one action per turn, based on the observation that the "
        "environment provides after each call. Stop when you see 'Episode done.'"
    )

    rows = []
    for i in range(n_samples):
        scen = scenarios[i % len(scenarios)]
        rows.append({
            "prompt": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_kickoff},
            ],
            "scenario_id": scen,
            "seed": i,
        })
    return Dataset.from_list(rows)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", type=str, default="unsloth/Qwen2.5-7B-Instruct-bnb-4bit")
    ap.add_argument("--env-url", type=str, default="http://0.0.0.0:8001",
                    help="OpsGuard env server base URL")
    ap.add_argument("--hub-repo", type=str, default=None,
                    help="HF repo id (e.g. user/opsguard-grpo) to push final LoRA")
    ap.add_argument("--num-steps", type=int, default=200,
                    help="GRPO training steps (max_steps)")
    ap.add_argument("--num-generations", type=int, default=4,
                    help="GRPO group size — completions per prompt")
    ap.add_argument("--max-steps-per-episode", type=int, default=30,
                    help="Cap on env steps per episode (matches E0/E1 budgets)")
    ap.add_argument("--max-completion-length", type=int, default=4096)
    ap.add_argument("--per-device-batch-size", type=int, default=1)
    ap.add_argument("--grad-accum", type=int, default=8)
    ap.add_argument("--lr", type=float, default=5e-6)
    ap.add_argument("--scenarios", nargs="+",
                    default=["E0_quiet_day", "E1_release_week"])
    ap.add_argument("--n-prompts", type=int, default=1024,
                    help="Size of the synthetic prompt dataset (one row = one episode)")
    ap.add_argument("--output-dir", type=str,
                    default=str(_PROJECT_ROOT / "checkpoints" / "opsguard-grpo"))
    ap.add_argument("--sft-adapter", type=str, default=None,
                    help="Path to SFT-warmstart LoRA to load before GRPO")
    ap.add_argument("--use-vllm", action="store_true", default=True)
    ap.add_argument("--no-vllm", dest="use_vllm", action="store_false")
    ap.add_argument("--vllm-mode", type=str, default="colocate",
                    choices=["colocate", "server"])
    ap.add_argument("--vllm-server-url", type=str, default=None)
    ap.add_argument("--reward-reduce", type=str, default="cumulative",
                    choices=["cumulative", "mean", "last"])
    ap.add_argument("--use-unsloth", action="store_true",
                    help="Preload model via Unsloth (4-bit + LoRA)")
    args = ap.parse_args()

    print(f"=== OpsGuard GRPO training ===", flush=True)
    print(f"  model:         {args.model}", flush=True)
    print(f"  env_url:       {args.env_url}", flush=True)
    print(f"  num_steps:     {args.num_steps}", flush=True)
    print(f"  num_gens:      {args.num_generations}", flush=True)
    print(f"  max_steps_ep:  {args.max_steps_per_episode}", flush=True)
    print(f"  output_dir:    {args.output_dir}", flush=True)
    print(f"  sft_adapter:   {args.sft_adapter}", flush=True)
    print(f"  vllm:          {args.use_vllm} ({args.vllm_mode})", flush=True)

    # Heavy imports here so --help is fast.
    from trl import GRPOConfig, GRPOTrainer  # type: ignore

    OpsguardToolEnv = make_opsguard_env_class(
        env_url=args.env_url,
        max_steps=args.max_steps_per_episode,
    )

    dataset = build_dataset(args.n_prompts, args.scenarios)

    # Build the model: either let TRL load by name, or pre-load via Unsloth.
    model_or_name: Any = args.model
    tokenizer = None
    if args.use_unsloth:
        print("preloading via Unsloth...", flush=True)
        from unsloth import FastLanguageModel  # type: ignore
        model, tokenizer = FastLanguageModel.from_pretrained(
            model_name=args.model,
            max_seq_length=args.max_completion_length + 2048,
            load_in_4bit=True,
            dtype=None,
        )
        model = FastLanguageModel.get_peft_model(
            model,
            r=32,
            lora_alpha=64,
            lora_dropout=0.0,
            bias="none",
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                            "gate_proj", "up_proj", "down_proj"],
            use_gradient_checkpointing="unsloth",
        )
        if args.sft_adapter:
            print(f"loading SFT adapter from {args.sft_adapter}...", flush=True)
            try:
                model.load_adapter(args.sft_adapter, adapter_name="default")
            except Exception as e:
                print(f"  [WARN] could not load adapter: {e}", flush=True)
        model_or_name = model

    # GRPOConfig: build kwargs then prune unsupported ones for safety.
    grpo_kwargs = dict(
        output_dir=args.output_dir,
        per_device_train_batch_size=args.per_device_batch_size,
        gradient_accumulation_steps=args.grad_accum,
        learning_rate=args.lr,
        max_steps=args.num_steps,
        num_generations=args.num_generations,
        max_completion_length=args.max_completion_length,
        logging_steps=1,
        save_steps=max(20, args.num_steps // 5),
        save_total_limit=2,
        bf16=True,
        report_to="none",
        log_completions=True,
        push_to_hub=bool(args.hub_repo),
        hub_model_id=args.hub_repo,
        chat_template_kwargs={"enable_thinking": False},
        use_vllm=args.use_vllm,
        vllm_mode=args.vllm_mode,
        vllm_server_url=args.vllm_server_url,
    )
    grpo_config = GRPOConfig(**_safe_kwargs(GRPOConfig, grpo_kwargs))

    reward_func = make_reward_func(args.reward_reduce)

    trainer_kwargs = dict(
        model=model_or_name,
        args=grpo_config,
        train_dataset=dataset,
        reward_funcs=reward_func,
        environment_factory=OpsguardToolEnv,
    )
    if tokenizer is not None:
        trainer_kwargs["processing_class"] = tokenizer
    trainer = GRPOTrainer(**_safe_kwargs(GRPOTrainer, trainer_kwargs))

    print("starting trainer.train()...", flush=True)
    trainer.train()

    print(f"saving final adapter to {args.output_dir}...", flush=True)
    try:
        trainer.save_model(args.output_dir)
    except Exception as e:
        print(f"  [WARN] save_model failed: {e}", flush=True)

    if args.hub_repo:
        try:
            trainer.push_to_hub()
            print(f"pushed to https://huggingface.co/{args.hub_repo}", flush=True)
        except Exception as e:
            print(f"  [WARN] push_to_hub failed: {e}", flush=True)

    print("DONE.", flush=True)


if __name__ == "__main__":
    main()