"""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= 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()