Spaces:
Sleeping
Sleeping
fix: resolve merge conflict, restore HF frontmatter + full README
Browse files- .gitignore +1 -0
- BlogPost.md +68 -0
- client.py +3 -4
- models.py +1 -1
- server/pm_ops_environment.py +34 -17
- training/Competent Pink Train (1).ipynb +0 -0
- training/__init__.py +0 -0
- training/dataset.py +110 -0
- training/pm_ops_trainer.py +187 -0
- training/prompts.py +94 -0
- training/rewards.py +100 -0
- training/rollout.py +430 -0
- training/smoke_test.py +267 -0
- training/train.ipynb +669 -0
- training/train_v2.ipynb +715 -0
- training/train_v3.ipynb +814 -0
- training/train_v4.ipynb +847 -0
- training/triage_dataset.jsonl +150 -0
.gitignore
CHANGED
|
@@ -12,3 +12,4 @@ pdf_pages/
|
|
| 12 |
dist/
|
| 13 |
build/
|
| 14 |
.pytest_cache/
|
|
|
|
|
|
| 12 |
dist/
|
| 13 |
build/
|
| 14 |
.pytest_cache/
|
| 15 |
+
.vscode/
|
BlogPost.md
ADDED
|
@@ -0,0 +1,68 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
## PM-Ops: The Product Manager Benchmark
|
| 2 |
+
|
| 3 |
+
### The Problem
|
| 4 |
+
Most industries use frontier models with near-perfect benchmark scores. However, even with these high capabilities, when dropped into a real company's workflow, these models often fail basic junior-level tasks.
|
| 5 |
+
|
| 6 |
+
**Example:** When a database alert fires at 2:00 AM, a model trained on the internet knows that DB failures are serious and will likely create an internal incident ticket immediately. But what if the company in question outsources its DB infrastructure? In that case, the protocol is an email to the vendor, not an internal ticket. The model, relying on general knowledge, gets the organizational context completely wrong.
|
| 7 |
+
|
| 8 |
+
### Why System Prompts Don't Fix This
|
| 9 |
+
|
| 10 |
+
1. **Attention Decay:** LLMs are probabilistic systems. The next token is shaped by billions of parameters trained on averaged human text. If we fill the context window with organizational conventions, the model doesn't "memorize" them; rather, it attends to them with diminishing weight as the conversation grows.
|
| 11 |
+
2. **Dynamic Organizations:** A system prompt is a static snapshot. Channels get renamed, and teams get reorganized. A static prompt cannot follow these shifts, leading the model to post to dead channels or outdated aliases.
|
| 12 |
+
|
| 13 |
+
---
|
| 14 |
+
|
| 15 |
+
### The Environment
|
| 16 |
+
PM-Ops simulates a full software organization as three interconnected applications. The agent operates as a product manager navigating all three.
|
| 17 |
+
|
| 18 |
+
* **Ticketing App (Jira-like):** Tracks tickets, projects, teams, users, labels, and priorities. The agent interacts via a simulated API. The "ground truth" of the organizational state is stored in a database; rather than letting the LLM judge its own performance, we verify its actions against this database.
|
| 19 |
+
* **Codebase App (GitHub-like):** Contains simulated repositories, commit history, pull requests, file authors, and changed files. The goal here is detective work: when a bug report arrives, the agent must trace commit history to identify changes and the responsible authors.
|
| 20 |
+
* **Chat App (Slack-like):** Features channels, threaded conversations, direct messages, and user profiles. The agent can read history, search, and post. Channels are episode-specific; for example, a critical alert channel might be `#oncall-payments` in Episode 1, but be replaced by `#urgent-billing` in Episode 2.
|
| 21 |
+
|
| 22 |
+
**Key Tool:** `meta.read_runbook`. Before taking any action, the agent can call this tool to return the organization's current conventions as a structured document.
|
| 23 |
+
|
| 24 |
+
---
|
| 25 |
+
|
| 26 |
+
### Scenarios and Difficulty Tiers
|
| 27 |
+
LLM-generated scenarios often produce plausible-sounding outputs that fail formal verification. To solve this, we use an **org-generator** that is deterministic, parameterized, and seedable.
|
| 28 |
+
|
| 29 |
+
| Tier | Complexity | Constraints |
|
| 30 |
+
| :--- | :--- | :--- |
|
| 31 |
+
| **Easy** | 2 services, 3 labels, 2 priorities. | No ambiguity. |
|
| 32 |
+
| **Medium** | 4 labels, 4 priorities, 3 services. | Partial runbook; agent must infer info across systems. |
|
| 33 |
+
| **Hard** | ≥5 services, multi-owner configs. | High noise; outdated docs; requires resisting "panic" in general chat. |
|
| 34 |
+
|
| 35 |
+
### Task Types
|
| 36 |
+
1. **Triage:** A bug report arrives. The agent must create a ticket with the correct label, priority, and taxonomy, then assign it to the right team and notify the correct channel.
|
| 37 |
+
2. **Incident Routing:** A production incident occurs. The agent must identify affected services via the codebase and team map to route the escalation correctly and quickly.
|
| 38 |
+
3. **Release Notes:** The agent compiles and posts release notes for resolved tickets using the organization's specific format.
|
| 39 |
+
4. **Dependency Update:** A library change affects multiple services. The agent must identify all affected teams, notify them through their respective on-call channels, and create tickets for each—a multi-target coordination problem.
|
| 40 |
+
|
| 41 |
+
---
|
| 42 |
+
|
| 43 |
+
### Rewards and Penalties
|
| 44 |
+
Rewards are **delayed**; the score is revealed only when `meta.finish` is called. This prevents step-by-step optimization and forces the agent to commit to a coherent plan.
|
| 45 |
+
|
| 46 |
+
* **Rewards:**
|
| 47 |
+
* Ticket created: **+0.25**
|
| 48 |
+
* Correct label/priority/team: **+0.20 each**
|
| 49 |
+
* Correct channel notification: **+0.15**
|
| 50 |
+
* **Penalties:**
|
| 51 |
+
* Wrong channel message: **-0.05 each** (prevents "spray and pray" messaging)
|
| 52 |
+
* Duplicate tickets: **-0.10 each**
|
| 53 |
+
* Zero valid actions (inaction): **-1.0**
|
| 54 |
+
|
| 55 |
+
### Training with GRPO
|
| 56 |
+
We use **GRPO (Group Relative Policy Optimization)** to train a 1.7B parameter model.
|
| 57 |
+
* **SFT Warmup:** We begin with a short Supervised Fine-Tuning phase on baseline demonstrations to establish the JSON output format. Without this, the model outputs prose, resulting in a zero gradient for GRPO.
|
| 58 |
+
* **Optimization:** After SFT, GRPO optimizes directly against episode rewards. Each training step is a full PM-Ops episode (up to 12 turns). Weights are updated based on relative performance across the generation group.
|
| 59 |
+
|
| 60 |
+
The result is an agent that doesn't just "know" that reading conventions is good practice—it has learned through thousands of varied configurations that skipping the runbook reliably leads to penalties. The behavior becomes **instinctual, not just instructed.**
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+

|
| 64 |
+
|
| 65 |
+
Link of image1: https://drive.google.com/uc?export=view&id=1JdwYukKrEaMTaOwc1W4Q8bBRjxwGZ2be
|
| 66 |
+
|
| 67 |
+
Link of image2: https://drive.google.com/file/d/1Wb8G0WEPPvAFppBNSVMZ7EMFBPFjsP8n/view?usp=sharing
|
| 68 |
+
|
client.py
CHANGED
|
@@ -1,6 +1,6 @@
|
|
| 1 |
"""OpenEnv client wrapper for PM-Ops environment."""
|
| 2 |
import os
|
| 3 |
-
from openenv.core.
|
| 4 |
from models import PMOpsAction, PMOpsObservation
|
| 5 |
|
| 6 |
|
|
@@ -8,7 +8,6 @@ class PMOpsEnv(EnvClient):
|
|
| 8 |
action_class = PMOpsAction
|
| 9 |
observation_class = PMOpsObservation
|
| 10 |
|
| 11 |
-
def __init__(self, base_url: str = None, token: str = None):
|
| 12 |
url = base_url or os.getenv("API_BASE_URL", "https://adityaguntur-pm-ops.hf.space")
|
| 13 |
-
|
| 14 |
-
super().__init__(base_url=url, token=tok)
|
|
|
|
| 1 |
"""OpenEnv client wrapper for PM-Ops environment."""
|
| 2 |
import os
|
| 3 |
+
from openenv.core.env_client import EnvClient
|
| 4 |
from models import PMOpsAction, PMOpsObservation
|
| 5 |
|
| 6 |
|
|
|
|
| 8 |
action_class = PMOpsAction
|
| 9 |
observation_class = PMOpsObservation
|
| 10 |
|
| 11 |
+
def __init__(self, base_url: str | None = None, token: str | None= None):
|
| 12 |
url = base_url or os.getenv("API_BASE_URL", "https://adityaguntur-pm-ops.hf.space")
|
| 13 |
+
super().__init__(base_url=url)
|
|
|
models.py
CHANGED
|
@@ -46,4 +46,4 @@ class PMOpsState(State):
|
|
| 46 |
codebase: Dict[str, Any] = Field(default_factory=dict)
|
| 47 |
step_count: int = Field(default=0)
|
| 48 |
finished: bool = Field(default=False)
|
| 49 |
-
episode_id: str = Field(default="")
|
|
|
|
| 46 |
codebase: Dict[str, Any] = Field(default_factory=dict)
|
| 47 |
step_count: int = Field(default=0)
|
| 48 |
finished: bool = Field(default=False)
|
| 49 |
+
episode_id: str | None = Field(default="")
|
server/pm_ops_environment.py
CHANGED
|
@@ -54,12 +54,18 @@ class PMOpsEnvironment(Environment[PMOpsAction, PMOpsObservation, PMOpsState]):
|
|
| 54 |
difficulty = rng.choice(_DIFFICULTY_POOL)
|
| 55 |
task_type = rng.choice(_TASK_TYPES)
|
| 56 |
|
|
|
|
|
|
|
|
|
|
| 57 |
for attempt in range(10):
|
| 58 |
org = generate_org_config(seed + attempt, difficulty)
|
| 59 |
scenario = generate_scenario(task_type, org, seed + attempt)
|
| 60 |
if _oracle_check(scenario):
|
| 61 |
break
|
| 62 |
|
|
|
|
|
|
|
|
|
|
| 63 |
channels = list(org["oncall_channels"].values())
|
| 64 |
noise = org.get("noise_channels", [])
|
| 65 |
|
|
@@ -87,14 +93,16 @@ class PMOpsEnvironment(Environment[PMOpsAction, PMOpsObservation, PMOpsState]):
|
|
| 87 |
)
|
| 88 |
|
| 89 |
def step(self, action: PMOpsAction, timeout_s: Optional[float] = None, **kwargs) -> PMOpsObservation:
|
| 90 |
-
if self._ticketing is None:
|
| 91 |
raise RuntimeError("Call reset() before step()")
|
| 92 |
|
|
|
|
|
|
|
| 93 |
if self._done:
|
| 94 |
return PMOpsObservation(
|
| 95 |
step=self._step_count,
|
| 96 |
max_steps=MAX_STEPS,
|
| 97 |
-
task_brief=
|
| 98 |
last_action_result={"ok": False, "error": "Episode already finished"},
|
| 99 |
app_state_deltas={"ticketing": [], "chat": [], "codebase": []},
|
| 100 |
steps_remaining=0,
|
|
@@ -122,7 +130,7 @@ class PMOpsEnvironment(Environment[PMOpsAction, PMOpsObservation, PMOpsState]):
|
|
| 122 |
return PMOpsObservation(
|
| 123 |
step=self._step_count,
|
| 124 |
max_steps=MAX_STEPS,
|
| 125 |
-
task_brief=
|
| 126 |
last_action_result=result,
|
| 127 |
app_state_deltas=deltas,
|
| 128 |
steps_remaining=max(0, MAX_STEPS - self._step_count),
|
|
@@ -147,6 +155,12 @@ class PMOpsEnvironment(Environment[PMOpsAction, PMOpsObservation, PMOpsState]):
|
|
| 147 |
at = action.action_type
|
| 148 |
args = action.args or {}
|
| 149 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 150 |
if at == "meta.noop":
|
| 151 |
return {"ok": True, "data": "No operation."}
|
| 152 |
|
|
@@ -170,13 +184,13 @@ class PMOpsEnvironment(Environment[PMOpsAction, PMOpsObservation, PMOpsState]):
|
|
| 170 |
if at.startswith("ticketing."):
|
| 171 |
op = at.split(".", 1)[1]
|
| 172 |
handlers = {
|
| 173 |
-
"create_ticket":
|
| 174 |
-
"update_ticket":
|
| 175 |
-
"get_ticket":
|
| 176 |
-
"list_tickets":
|
| 177 |
-
"assign_ticket":
|
| 178 |
-
"comment_ticket":
|
| 179 |
-
"transition_ticket":
|
| 180 |
}
|
| 181 |
if op not in handlers:
|
| 182 |
return {"ok": False, "error": f"Unknown ticketing action: {op}"}
|
|
@@ -185,9 +199,9 @@ class PMOpsEnvironment(Environment[PMOpsAction, PMOpsObservation, PMOpsState]):
|
|
| 185 |
if at.startswith("codebase."):
|
| 186 |
op = at.split(".", 1)[1]
|
| 187 |
handlers = {
|
| 188 |
-
"list_commits":
|
| 189 |
-
"get_commit":
|
| 190 |
-
"list_prs":
|
| 191 |
}
|
| 192 |
if op not in handlers:
|
| 193 |
return {"ok": False, "error": f"Unknown codebase action: {op}"}
|
|
@@ -196,10 +210,10 @@ class PMOpsEnvironment(Environment[PMOpsAction, PMOpsObservation, PMOpsState]):
|
|
| 196 |
if at.startswith("chat."):
|
| 197 |
op = at.split(".", 1)[1]
|
| 198 |
handlers = {
|
| 199 |
-
"post_message":
|
| 200 |
-
"read_channel":
|
| 201 |
-
"list_channels":
|
| 202 |
-
"search":
|
| 203 |
}
|
| 204 |
if op not in handlers:
|
| 205 |
return {"ok": False, "error": f"Unknown chat action: {op}"}
|
|
@@ -208,6 +222,9 @@ class PMOpsEnvironment(Environment[PMOpsAction, PMOpsObservation, PMOpsState]):
|
|
| 208 |
return {"ok": False, "error": f"Unknown action_type: {at}"}
|
| 209 |
|
| 210 |
def _grade(self) -> float:
|
|
|
|
|
|
|
|
|
|
| 211 |
final_state = {
|
| 212 |
"ticketing": self._ticketing.snapshot(),
|
| 213 |
"chat": self._chat.snapshot(),
|
|
|
|
| 54 |
difficulty = rng.choice(_DIFFICULTY_POOL)
|
| 55 |
task_type = rng.choice(_TASK_TYPES)
|
| 56 |
|
| 57 |
+
org: Optional[Dict[str, Any]] = None
|
| 58 |
+
scenario: Optional[Dict[str, Any]] = None
|
| 59 |
+
|
| 60 |
for attempt in range(10):
|
| 61 |
org = generate_org_config(seed + attempt, difficulty)
|
| 62 |
scenario = generate_scenario(task_type, org, seed + attempt)
|
| 63 |
if _oracle_check(scenario):
|
| 64 |
break
|
| 65 |
|
| 66 |
+
if org is None or scenario is None:
|
| 67 |
+
raise RuntimeError("Failed to initialize episode state")
|
| 68 |
+
|
| 69 |
channels = list(org["oncall_channels"].values())
|
| 70 |
noise = org.get("noise_channels", [])
|
| 71 |
|
|
|
|
| 93 |
)
|
| 94 |
|
| 95 |
def step(self, action: PMOpsAction, timeout_s: Optional[float] = None, **kwargs) -> PMOpsObservation:
|
| 96 |
+
if self._ticketing is None or self._codebase is None or self._chat is None or self._scenario is None:
|
| 97 |
raise RuntimeError("Call reset() before step()")
|
| 98 |
|
| 99 |
+
scenario = self._scenario
|
| 100 |
+
|
| 101 |
if self._done:
|
| 102 |
return PMOpsObservation(
|
| 103 |
step=self._step_count,
|
| 104 |
max_steps=MAX_STEPS,
|
| 105 |
+
task_brief=scenario["brief"],
|
| 106 |
last_action_result={"ok": False, "error": "Episode already finished"},
|
| 107 |
app_state_deltas={"ticketing": [], "chat": [], "codebase": []},
|
| 108 |
steps_remaining=0,
|
|
|
|
| 130 |
return PMOpsObservation(
|
| 131 |
step=self._step_count,
|
| 132 |
max_steps=MAX_STEPS,
|
| 133 |
+
task_brief=scenario["brief"],
|
| 134 |
last_action_result=result,
|
| 135 |
app_state_deltas=deltas,
|
| 136 |
steps_remaining=max(0, MAX_STEPS - self._step_count),
|
|
|
|
| 155 |
at = action.action_type
|
| 156 |
args = action.args or {}
|
| 157 |
|
| 158 |
+
ticketing = self._ticketing
|
| 159 |
+
codebase = self._codebase
|
| 160 |
+
chat = self._chat
|
| 161 |
+
if ticketing is None or codebase is None or chat is None:
|
| 162 |
+
return {"ok": False, "error": "Environment not initialized. Call reset() before step()."}
|
| 163 |
+
|
| 164 |
if at == "meta.noop":
|
| 165 |
return {"ok": True, "data": "No operation."}
|
| 166 |
|
|
|
|
| 184 |
if at.startswith("ticketing."):
|
| 185 |
op = at.split(".", 1)[1]
|
| 186 |
handlers = {
|
| 187 |
+
"create_ticket": ticketing.create_ticket,
|
| 188 |
+
"update_ticket": ticketing.update_ticket,
|
| 189 |
+
"get_ticket": ticketing.get_ticket,
|
| 190 |
+
"list_tickets": ticketing.list_tickets,
|
| 191 |
+
"assign_ticket": ticketing.assign_ticket,
|
| 192 |
+
"comment_ticket": ticketing.comment_ticket,
|
| 193 |
+
"transition_ticket": ticketing.transition_ticket,
|
| 194 |
}
|
| 195 |
if op not in handlers:
|
| 196 |
return {"ok": False, "error": f"Unknown ticketing action: {op}"}
|
|
|
|
| 199 |
if at.startswith("codebase."):
|
| 200 |
op = at.split(".", 1)[1]
|
| 201 |
handlers = {
|
| 202 |
+
"list_commits": codebase.list_commits,
|
| 203 |
+
"get_commit": codebase.get_commit,
|
| 204 |
+
"list_prs": codebase.list_prs,
|
| 205 |
}
|
| 206 |
if op not in handlers:
|
| 207 |
return {"ok": False, "error": f"Unknown codebase action: {op}"}
|
|
|
|
| 210 |
if at.startswith("chat."):
|
| 211 |
op = at.split(".", 1)[1]
|
| 212 |
handlers = {
|
| 213 |
+
"post_message": chat.post_message,
|
| 214 |
+
"read_channel": chat.read_channel,
|
| 215 |
+
"list_channels": chat.list_channels,
|
| 216 |
+
"search": chat.search,
|
| 217 |
}
|
| 218 |
if op not in handlers:
|
| 219 |
return {"ok": False, "error": f"Unknown chat action: {op}"}
|
|
|
|
| 222 |
return {"ok": False, "error": f"Unknown action_type: {at}"}
|
| 223 |
|
| 224 |
def _grade(self) -> float:
|
| 225 |
+
if self._ticketing is None or self._chat is None or self._org_config is None or self._scenario is None:
|
| 226 |
+
return 0.0
|
| 227 |
+
|
| 228 |
final_state = {
|
| 229 |
"ticketing": self._ticketing.snapshot(),
|
| 230 |
"chat": self._chat.snapshot(),
|
training/Competent Pink Train (1).ipynb
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
training/__init__.py
ADDED
|
File without changes
|
training/dataset.py
ADDED
|
@@ -0,0 +1,110 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Generate fixed-seed training dataset for PM-Ops triage task.
|
| 2 |
+
|
| 3 |
+
Each seed is validated against the env's own RNG so the task type, difficulty,
|
| 4 |
+
org_config, and scenario the env generates at reset(seed=S) EXACTLY matches the
|
| 5 |
+
brief embedded in the prompt. Previously, the dataset generated triage briefs but
|
| 6 |
+
the env silently ran a different task type (release_notes, dep_update, etc.) for
|
| 7 |
+
the same seed — causing env_score=0 for every episode.
|
| 8 |
+
"""
|
| 9 |
+
import json
|
| 10 |
+
import os
|
| 11 |
+
import random
|
| 12 |
+
import sys
|
| 13 |
+
|
| 14 |
+
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
| 15 |
+
|
| 16 |
+
from server.world.org_generator import generate_org_config
|
| 17 |
+
from server.world.scenario_gen import generate_scenario
|
| 18 |
+
|
| 19 |
+
SEED_PREFIX = "SEED:"
|
| 20 |
+
|
| 21 |
+
# Must match pm_ops_environment.py constants exactly
|
| 22 |
+
_ENV_DIFFICULTY_POOL = ["easy", "medium", "medium", "hard"]
|
| 23 |
+
_ENV_TASK_TYPES = ["triage", "incident_routing", "release_notes", "dep_update"]
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
def _env_params(seed: int) -> tuple[str, str]:
|
| 27 |
+
"""Predict difficulty + task_type the env will choose for this seed.
|
| 28 |
+
|
| 29 |
+
Replicates the first two RNG calls in PMOpsEnvironment.reset() so we can
|
| 30 |
+
filter seeds to those that produce the task type we want.
|
| 31 |
+
"""
|
| 32 |
+
rng = random.Random(seed)
|
| 33 |
+
difficulty = rng.choice(_ENV_DIFFICULTY_POOL)
|
| 34 |
+
task_type = rng.choice(_ENV_TASK_TYPES)
|
| 35 |
+
return difficulty, task_type
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def _valid_triage(scenario: dict) -> bool:
|
| 39 |
+
exp = scenario.get("expected", {})
|
| 40 |
+
return bool(exp.get("channel")) and bool(exp.get("team"))
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def generate_triage_dataset(n_episodes: int = 150, base_seed: int = 42) -> list[dict]:
|
| 44 |
+
"""Return dataset rows where the env WILL run a triage episode for the embedded seed.
|
| 45 |
+
|
| 46 |
+
We pre-simulate the env's RNG to only include seeds where:
|
| 47 |
+
1. env.reset(seed) picks task_type="triage"
|
| 48 |
+
2. The resulting org + scenario are valid (have expected channel + team)
|
| 49 |
+
3. The difficulty, org_config, and brief exactly match what the env will use
|
| 50 |
+
|
| 51 |
+
This eliminates the mismatch where the dataset had triage briefs but the env
|
| 52 |
+
graded as release_notes → guaranteed env_score=0.
|
| 53 |
+
"""
|
| 54 |
+
rng = random.Random(base_seed)
|
| 55 |
+
rows = []
|
| 56 |
+
|
| 57 |
+
while len(rows) < n_episodes:
|
| 58 |
+
seed = rng.randint(0, 2**31)
|
| 59 |
+
|
| 60 |
+
# Only use seeds the env will run as triage
|
| 61 |
+
difficulty, task_type = _env_params(seed)
|
| 62 |
+
if task_type != "triage":
|
| 63 |
+
continue
|
| 64 |
+
|
| 65 |
+
# Generate org + scenario using the SAME difficulty + seed the env will use
|
| 66 |
+
org, scenario = None, None
|
| 67 |
+
for attempt in range(10):
|
| 68 |
+
org = generate_org_config(seed + attempt, difficulty)
|
| 69 |
+
scenario = generate_scenario("triage", org, seed + attempt)
|
| 70 |
+
if _valid_triage(scenario):
|
| 71 |
+
break
|
| 72 |
+
|
| 73 |
+
if not _valid_triage(scenario):
|
| 74 |
+
continue
|
| 75 |
+
|
| 76 |
+
rows.append({
|
| 77 |
+
"prompt": f"{SEED_PREFIX}{seed} | {scenario['brief']}",
|
| 78 |
+
"seed": seed,
|
| 79 |
+
"difficulty": difficulty,
|
| 80 |
+
})
|
| 81 |
+
|
| 82 |
+
return rows
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def parse_seed_from_prompt(prompt: str) -> int | None:
|
| 86 |
+
"""Extract seed embedded by generate_triage_dataset."""
|
| 87 |
+
if not prompt.startswith(SEED_PREFIX):
|
| 88 |
+
return None
|
| 89 |
+
try:
|
| 90 |
+
return int(prompt[len(SEED_PREFIX):].split(" | ")[0])
|
| 91 |
+
except ValueError:
|
| 92 |
+
return None
|
| 93 |
+
|
| 94 |
+
|
| 95 |
+
def save_dataset(rows: list[dict], path: str) -> None:
|
| 96 |
+
with open(path, "w") as f:
|
| 97 |
+
for row in rows:
|
| 98 |
+
f.write(json.dumps(row) + "\n")
|
| 99 |
+
print(f"Saved {len(rows)} episodes → {path}")
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
def load_dataset(path: str) -> list[dict]:
|
| 103 |
+
with open(path) as f:
|
| 104 |
+
return [json.loads(line) for line in f]
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
if __name__ == "__main__":
|
| 108 |
+
out = os.path.join(os.path.dirname(__file__), "triage_dataset.jsonl")
|
| 109 |
+
rows = generate_triage_dataset(n_episodes=150)
|
| 110 |
+
save_dataset(rows, out)
|
training/pm_ops_trainer.py
ADDED
|
@@ -0,0 +1,187 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""PMOpsGRPOTrainer — GRPOTrainer subclass with direct reward injection.
|
| 2 |
+
|
| 3 |
+
Problem being solved
|
| 4 |
+
--------------------
|
| 5 |
+
TRL 1.2.0 + Unsloth PatchFastRL strips custom rollout_func output keys
|
| 6 |
+
before calling reward_funcs. The 'reward' key never arrives in kwargs, so
|
| 7 |
+
reward_func always returns 0.0 and GRPO sees constant reward → zero gradient.
|
| 8 |
+
|
| 9 |
+
Root cause: TRL builds reward_kwargs from the *dataset* batch columns (only
|
| 10 |
+
'prompt'). rollout_func extra keys are not stored in the dataset, so they
|
| 11 |
+
never reach _calculate_rewards via inputs.
|
| 12 |
+
|
| 13 |
+
Fix
|
| 14 |
+
---
|
| 15 |
+
Wrap rollout_func inside __init__ using a closure that captures self directly
|
| 16 |
+
(standard Python closure semantics — no forward reference tricks needed).
|
| 17 |
+
Each time rollout_func is called during training, the wrapper stores the
|
| 18 |
+
'reward' list into self._rollout_reward_cache. _calculate_rewards then injects
|
| 19 |
+
those cached values as a tensor, bypassing the broken kwargs path entirely.
|
| 20 |
+
|
| 21 |
+
Usage
|
| 22 |
+
-----
|
| 23 |
+
from training.pm_ops_trainer import PMOpsGRPOTrainer
|
| 24 |
+
|
| 25 |
+
trainer = PMOpsGRPOTrainer(
|
| 26 |
+
model=model,
|
| 27 |
+
processing_class=tokenizer,
|
| 28 |
+
reward_funcs=reward_func, # kept as fallback; not called on inject path
|
| 29 |
+
train_dataset=dataset,
|
| 30 |
+
args=grpo_config,
|
| 31 |
+
rollout_func=rollout_func, # must return 'reward': list[float] in output
|
| 32 |
+
)
|
| 33 |
+
|
| 34 |
+
rollout_func contract
|
| 35 |
+
---------------------
|
| 36 |
+
rollout_func(prompts, trainer=None) must return a dict containing at minimum:
|
| 37 |
+
{
|
| 38 |
+
"prompt_ids": list[list[int]],
|
| 39 |
+
"completion_ids": list[list[int]],
|
| 40 |
+
"logprobs": list[list[float]],
|
| 41 |
+
"reward": list[float], # one combined float per episode
|
| 42 |
+
}
|
| 43 |
+
"""
|
| 44 |
+
|
| 45 |
+
import torch
|
| 46 |
+
from trl.trainer.grpo_trainer import GRPOTrainer
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
class PMOpsGRPOTrainer(GRPOTrainer):
|
| 50 |
+
"""Drop-in GRPOTrainer with closure-based reward caching.
|
| 51 |
+
|
| 52 |
+
Wraps rollout_func at construction time to intercept the 'reward' list
|
| 53 |
+
on every call. _calculate_rewards injects these rewards directly as a
|
| 54 |
+
tensor — no kwargs, no inputs lookup.
|
| 55 |
+
Falls back to standard GRPOTrainer behaviour when cache is empty.
|
| 56 |
+
"""
|
| 57 |
+
|
| 58 |
+
def __init__(self, *args, rollout_func=None, **kwargs):
|
| 59 |
+
# Reward cache populated by the rollout wrapper, consumed by _calculate_rewards
|
| 60 |
+
self._rollout_reward_cache: list[float] = []
|
| 61 |
+
self._rollout_capture_calls = 0
|
| 62 |
+
self._warned_rollout_bypass = False
|
| 63 |
+
|
| 64 |
+
wrapped_rollout_func = None
|
| 65 |
+
if rollout_func is not None:
|
| 66 |
+
original_rollout = rollout_func
|
| 67 |
+
|
| 68 |
+
# self is captured directly by closure; accept flexible call signatures
|
| 69 |
+
# because cloud runtimes may invoke rollout_func with positional/keyword variations.
|
| 70 |
+
def _capturing_rollout(prompts, trainer=None, *rollout_args, **rollout_kwargs):
|
| 71 |
+
rollout_kwargs.setdefault("trainer", trainer)
|
| 72 |
+
result = original_rollout(prompts, *rollout_args, **rollout_kwargs)
|
| 73 |
+
|
| 74 |
+
raw_rewards = result.get("reward", result.get("rewards", []))
|
| 75 |
+
if raw_rewards is None:
|
| 76 |
+
rewards = []
|
| 77 |
+
elif isinstance(raw_rewards, torch.Tensor):
|
| 78 |
+
rewards = [float(r) for r in raw_rewards.detach().cpu().flatten().tolist()]
|
| 79 |
+
elif isinstance(raw_rewards, (int, float)):
|
| 80 |
+
rewards = [float(raw_rewards)]
|
| 81 |
+
else:
|
| 82 |
+
rewards = [float(r) for r in list(raw_rewards)]
|
| 83 |
+
|
| 84 |
+
self._rollout_reward_cache = rewards
|
| 85 |
+
self._rollout_capture_calls += 1
|
| 86 |
+
print(f"[PMOpsGRPOTrainer] captured {len(rewards)} rewards "
|
| 87 |
+
f"(mean={sum(rewards)/len(rewards):.3f})" if rewards else
|
| 88 |
+
"[PMOpsGRPOTrainer] WARNING: rollout returned no rewards")
|
| 89 |
+
return result
|
| 90 |
+
|
| 91 |
+
wrapped_rollout_func = _capturing_rollout
|
| 92 |
+
rollout_func = wrapped_rollout_func
|
| 93 |
+
|
| 94 |
+
super().__init__(*args, rollout_func=rollout_func, **kwargs)
|
| 95 |
+
|
| 96 |
+
# Keep wrapper bound explicitly in case an upstream patch reassigns rollout_func.
|
| 97 |
+
if wrapped_rollout_func is not None:
|
| 98 |
+
self.rollout_func = wrapped_rollout_func
|
| 99 |
+
|
| 100 |
+
def _calculate_rewards(
|
| 101 |
+
self,
|
| 102 |
+
inputs,
|
| 103 |
+
prompts,
|
| 104 |
+
completions,
|
| 105 |
+
completion_ids_list,
|
| 106 |
+
):
|
| 107 |
+
"""Inject cached rewards when available; fall back to reward_funcs otherwise."""
|
| 108 |
+
def _coerce_rewards(raw):
|
| 109 |
+
if raw is None:
|
| 110 |
+
return []
|
| 111 |
+
if isinstance(raw, torch.Tensor):
|
| 112 |
+
return [float(r) for r in raw.detach().cpu().flatten().tolist()]
|
| 113 |
+
if isinstance(raw, (int, float)):
|
| 114 |
+
return [float(raw)]
|
| 115 |
+
return [float(r) for r in list(raw)]
|
| 116 |
+
|
| 117 |
+
cache = self._rollout_reward_cache
|
| 118 |
+
n = len(completions)
|
| 119 |
+
|
| 120 |
+
# Some TRL variants forward rollout extra_fields into `inputs` directly.
|
| 121 |
+
if not cache:
|
| 122 |
+
input_rewards: list[float] = []
|
| 123 |
+
for row in inputs:
|
| 124 |
+
if isinstance(row, dict) and "reward" in row:
|
| 125 |
+
input_rewards.append(float(row["reward"]))
|
| 126 |
+
else:
|
| 127 |
+
input_rewards = []
|
| 128 |
+
break
|
| 129 |
+
if input_rewards:
|
| 130 |
+
cache = input_rewards
|
| 131 |
+
|
| 132 |
+
# Cloud fallback: some patched runtimes skip the normal rollout capture path
|
| 133 |
+
# before calling _calculate_rewards. Actively invoke rollout_func here.
|
| 134 |
+
# Pass ALL n prompts (including repeated ones for num_generations > 1) so
|
| 135 |
+
# the rollout_func can generate distinct rewards per generation — NOT tile-mod.
|
| 136 |
+
if not cache and self.rollout_func is not None and prompts:
|
| 137 |
+
try:
|
| 138 |
+
print(f"[PMOpsGRPOTrainer] cache empty — probing rollout_func for {n} rewards")
|
| 139 |
+
out = self.rollout_func(list(prompts), trainer=self)
|
| 140 |
+
cache = self._rollout_reward_cache
|
| 141 |
+
if not cache and isinstance(out, dict):
|
| 142 |
+
cache = _coerce_rewards(out.get("reward", out.get("rewards", [])))
|
| 143 |
+
self._rollout_reward_cache = cache
|
| 144 |
+
except Exception as exc:
|
| 145 |
+
print(f"[PMOpsGRPOTrainer] rollout probe failed: {exc!r}")
|
| 146 |
+
|
| 147 |
+
if cache:
|
| 148 |
+
if len(cache) == n:
|
| 149 |
+
rewards_list = cache
|
| 150 |
+
elif len(cache) > n:
|
| 151 |
+
rewards_list = cache[:n]
|
| 152 |
+
else:
|
| 153 |
+
# Still short — extend with mean rather than tile-mod so we don't
|
| 154 |
+
# duplicate rewards for the same prompt (tile-mod → zero advantage).
|
| 155 |
+
mean_r = sum(cache) / len(cache)
|
| 156 |
+
rewards_list = list(cache) + [mean_r] * (n - len(cache))
|
| 157 |
+
print(f"[PMOpsGRPOTrainer] WARNING: cache has {len(cache)} rewards for n={n}; "
|
| 158 |
+
f"padding with mean={mean_r:.3f}. Consider matching num_generations.")
|
| 159 |
+
|
| 160 |
+
device = self.accelerator.device
|
| 161 |
+
rewards = torch.tensor(
|
| 162 |
+
[float(r) for r in rewards_list],
|
| 163 |
+
dtype=torch.float32,
|
| 164 |
+
device=device,
|
| 165 |
+
).unsqueeze(1) # [batch_size, 1]
|
| 166 |
+
|
| 167 |
+
std = rewards.std().item() if n > 1 else 0.0
|
| 168 |
+
self.log({"reward/injected_mean": rewards.mean().item(),
|
| 169 |
+
"reward/injected_std": std})
|
| 170 |
+
self._rollout_reward_cache = [] # consume cache
|
| 171 |
+
print(f"[PMOpsGRPOTrainer] injected {n} rewards "
|
| 172 |
+
f"mean={rewards.mean().item():.3f} std={std:.3f}")
|
| 173 |
+
return rewards
|
| 174 |
+
|
| 175 |
+
if self.rollout_func is not None and self._rollout_capture_calls == 0 and not self._warned_rollout_bypass:
|
| 176 |
+
print(
|
| 177 |
+
"[PMOpsGRPOTrainer] WARNING: rollout_func was never called before reward calculation. "
|
| 178 |
+
"This cloud runtime is likely bypassing rollout_func (TRL/Unsloth mismatch), so "
|
| 179 |
+
"reward_funcs only receive prompts/completion_ids/trainer_state and no 'reward' key."
|
| 180 |
+
)
|
| 181 |
+
self._warned_rollout_bypass = True
|
| 182 |
+
|
| 183 |
+
# Fallback: standard TRL reward_funcs path
|
| 184 |
+
print("[PMOpsGRPOTrainer] cache empty — falling back to reward_funcs")
|
| 185 |
+
return super()._calculate_rewards(
|
| 186 |
+
inputs, prompts, completions, completion_ids_list
|
| 187 |
+
)
|
training/prompts.py
ADDED
|
@@ -0,0 +1,94 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""System prompt and observation formatting for PM-Ops GRPO training."""
|
| 2 |
+
|
| 3 |
+
SYSTEM_PROMPT = """You are an expert PM Operations agent inside a software organization.
|
| 4 |
+
|
| 5 |
+
## Your Tools
|
| 6 |
+
- **Ticketing** (Jira-like): create_ticket, update_ticket, assign_ticket, transition_ticket, get_ticket, list_tickets, comment_ticket
|
| 7 |
+
- **Chat** (Slack-like): post_message, read_channel, list_channels, search
|
| 8 |
+
- **Codebase** (GitHub-like): list_commits, get_commit, list_prs
|
| 9 |
+
- **Meta**: read_runbook, finish, noop
|
| 10 |
+
|
| 11 |
+
## CRITICAL — Read the Runbook First
|
| 12 |
+
Every organization uses different conventions. Your FIRST action MUST be `meta.read_runbook`.
|
| 13 |
+
The runbook tells you the EXACT values to use:
|
| 14 |
+
- `label_taxonomy`: valid ticket labels (e.g. "defect", NOT "bug")
|
| 15 |
+
- `priority_levels`: valid priorities (e.g. "critical", NOT "P1")
|
| 16 |
+
- `team_map`: service → owning team
|
| 17 |
+
- `oncall_channels`: service → channel to notify
|
| 18 |
+
|
| 19 |
+
## Action Format
|
| 20 |
+
Reason through the problem, then emit exactly ONE JSON code block as your final output:
|
| 21 |
+
|
| 22 |
+
```json
|
| 23 |
+
{"action_type": "meta.read_runbook", "args": {}}
|
| 24 |
+
```
|
| 25 |
+
|
| 26 |
+
Full list of valid action_type values:
|
| 27 |
+
meta.read_runbook | meta.finish | meta.noop
|
| 28 |
+
ticketing.create_ticket | ticketing.update_ticket | ticketing.get_ticket
|
| 29 |
+
ticketing.list_tickets | ticketing.assign_ticket | ticketing.comment_ticket
|
| 30 |
+
ticketing.transition_ticket
|
| 31 |
+
codebase.list_commits | codebase.get_commit | codebase.list_prs
|
| 32 |
+
chat.post_message | chat.read_channel | chat.list_channels | chat.search
|
| 33 |
+
|
| 34 |
+
## Triage Strategy (follow this order)
|
| 35 |
+
1. `meta.read_runbook` — learn this org's label, priority, team, and channel conventions
|
| 36 |
+
2. `ticketing.create_ticket` — use EXACT label from label_taxonomy, EXACT priority from priority_levels
|
| 37 |
+
3. `ticketing.assign_ticket` — assign to the team that owns the affected service (team_map)
|
| 38 |
+
4. `chat.post_message` — post to the oncall channel for that service (oncall_channels)
|
| 39 |
+
5. `meta.finish` — end the episode
|
| 40 |
+
|
| 41 |
+
## Rules
|
| 42 |
+
- NEVER guess label or priority values — use only what the runbook tells you
|
| 43 |
+
- NEVER post to noise channels (#random, #general, #water-cooler, etc.)
|
| 44 |
+
- ONLY post to the oncall channel for the affected service
|
| 45 |
+
- Call `meta.finish` when done — do not exceed steps unnecessarily
|
| 46 |
+
"""
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def format_observation(obs, step: int, last_action_result: dict | None = None) -> str:
|
| 50 |
+
"""Convert a PMOpsObservation into a user-facing message string."""
|
| 51 |
+
if hasattr(obs, "task_brief"):
|
| 52 |
+
task_brief = obs.task_brief
|
| 53 |
+
steps_remaining = obs.steps_remaining
|
| 54 |
+
result = obs.last_action_result
|
| 55 |
+
else:
|
| 56 |
+
task_brief = obs.get("task_brief", "")
|
| 57 |
+
steps_remaining = obs.get("steps_remaining", 40)
|
| 58 |
+
result = obs.get("last_action_result", {})
|
| 59 |
+
|
| 60 |
+
if last_action_result is not None:
|
| 61 |
+
result = last_action_result
|
| 62 |
+
|
| 63 |
+
parts = [f"## Task\n{task_brief}", f"\n**Step {step}** | {steps_remaining} steps remaining"]
|
| 64 |
+
|
| 65 |
+
if result:
|
| 66 |
+
ok = result.get("ok", False)
|
| 67 |
+
if ok:
|
| 68 |
+
data = result.get("data", "")
|
| 69 |
+
if isinstance(data, dict) and "org_config" in data:
|
| 70 |
+
org = data["org_config"]
|
| 71 |
+
labels = list(org.get("label_taxonomy", {}).values())
|
| 72 |
+
priorities = org.get("priority_levels", [])
|
| 73 |
+
team_map = org.get("team_map", {})
|
| 74 |
+
channels = org.get("oncall_channels", {})
|
| 75 |
+
noise = org.get("noise_channels", [])
|
| 76 |
+
lines = [
|
| 77 |
+
"=== RUNBOOK ===",
|
| 78 |
+
f"Valid labels : {', '.join(labels)}",
|
| 79 |
+
f"Valid priorities : {', '.join(priorities)}",
|
| 80 |
+
f"Team map : {' | '.join(f'{s} -> {t}' for s, t in team_map.items())}",
|
| 81 |
+
f"Oncall channels : {' | '.join(f'{s} -> {c}' for s, c in channels.items())}",
|
| 82 |
+
f"Noise channels : {', '.join(noise)} <-- NEVER post here",
|
| 83 |
+
"=== END RUNBOOK ===",
|
| 84 |
+
]
|
| 85 |
+
data_str = "\n".join(lines)
|
| 86 |
+
else:
|
| 87 |
+
data_str = str(data)[:2000]
|
| 88 |
+
parts.append(f"\n**Result (success):**\n{data_str}")
|
| 89 |
+
else:
|
| 90 |
+
error = result.get("error", "Unknown error")
|
| 91 |
+
parts.append(f"\n**Result (error):** {error}")
|
| 92 |
+
|
| 93 |
+
parts.append("\nWhat is your next action? Think it through, then output a JSON code block.")
|
| 94 |
+
return "\n".join(parts)
|
training/rewards.py
ADDED
|
@@ -0,0 +1,100 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Runbook-compliance reward for PM-Ops GRPO training.
|
| 2 |
+
|
| 3 |
+
Design philosophy (v4)
|
| 4 |
+
----------------------
|
| 5 |
+
Previous reward formulas (json_ratio, ok_ratio, diversity_bonus) were all constant
|
| 6 |
+
once the model learned to output valid JSON — giving zero GRPO advantage.
|
| 7 |
+
|
| 8 |
+
Root problem: env_score was always 0 because:
|
| 9 |
+
1. The model used hardcoded label='sec-issue' / priority='P1', which fail
|
| 10 |
+
TicketingApp validation for most org configs → no ticket created → grader
|
| 11 |
+
returns 0.0 immediately.
|
| 12 |
+
2. Dataset seeds produced triage briefs but the env ran a DIFFERENT task type
|
| 13 |
+
(release_notes, dep_update) for the same seed → guaranteed mismatch.
|
| 14 |
+
|
| 15 |
+
Fix: reward is computed from the agent's ACTIONS compared to RUNBOOK DATA.
|
| 16 |
+
The org_config (returned by meta.read_runbook) varies by seed — different orgs
|
| 17 |
+
have different valid labels, priorities, teams, and oncall channels. The model's
|
| 18 |
+
fixed template (sec-issue, P1, infra) scores well for some orgs and badly for
|
| 19 |
+
others, creating the reward VARIANCE that GRPO needs.
|
| 20 |
+
|
| 21 |
+
Components (sum = 1.0 when all correct):
|
| 22 |
+
read_runbook 0.10 — process: did agent read runbook first?
|
| 23 |
+
valid_label 0.20 — used a label from label_taxonomy? (+0.20 / -0.10)
|
| 24 |
+
valid_priority 0.15 — used a priority from priority_levels? (+0.15 / -0.10)
|
| 25 |
+
valid_team 0.20 — assigned to a team from team_map? (+0.20 / -0.10)
|
| 26 |
+
right_channel 0.25 — posted to an oncall channel? (+0.25 / -0.10 per wrong)
|
| 27 |
+
env_bonus 0.10 — env grader bonus when everything lines up correctly
|
| 28 |
+
|
| 29 |
+
This is computed IN the rollout (not by TRL reward_funcs) because it needs
|
| 30 |
+
access to the per-episode runbook data.
|
| 31 |
+
"""
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def compute_rollout_reward(
|
| 35 |
+
*,
|
| 36 |
+
read_runbook_done: bool,
|
| 37 |
+
valid_labels: set, # from org_config.label_taxonomy.values()
|
| 38 |
+
valid_priorities: set, # from org_config.priority_levels
|
| 39 |
+
valid_teams: set, # from org_config.team_map.values()
|
| 40 |
+
oncall_channels: set, # from org_config.oncall_channels.values()
|
| 41 |
+
ticket_label: str | None, # label used in create_ticket (None if no ticket)
|
| 42 |
+
ticket_priority: str | None, # priority used in create_ticket
|
| 43 |
+
assigned_team: str | None, # team from assign_ticket (None if not called)
|
| 44 |
+
posted_channels: list[str], # all channels from chat.post_message
|
| 45 |
+
env_score: float, # final env grader score (0–1)
|
| 46 |
+
valid_json_count: int, # steps with parseable JSON output
|
| 47 |
+
) -> float:
|
| 48 |
+
"""Compute a single reward scalar for one episode.
|
| 49 |
+
|
| 50 |
+
Returns a value in [-1.0, 1.0].
|
| 51 |
+
"""
|
| 52 |
+
if valid_json_count == 0:
|
| 53 |
+
return -1.0
|
| 54 |
+
|
| 55 |
+
reward = 0.0
|
| 56 |
+
|
| 57 |
+
# 1. Runbook read (+0.10 process bonus)
|
| 58 |
+
if read_runbook_done:
|
| 59 |
+
reward += 0.10
|
| 60 |
+
|
| 61 |
+
# 2. Ticket label valid (only scored if a ticket was created)
|
| 62 |
+
if ticket_label is not None:
|
| 63 |
+
if valid_labels:
|
| 64 |
+
if ticket_label in valid_labels:
|
| 65 |
+
reward += 0.20
|
| 66 |
+
else:
|
| 67 |
+
reward -= 0.10 # wrong label — validation would have rejected it
|
| 68 |
+
|
| 69 |
+
# 3. Ticket priority valid
|
| 70 |
+
if ticket_priority is not None:
|
| 71 |
+
if valid_priorities:
|
| 72 |
+
if ticket_priority in valid_priorities:
|
| 73 |
+
reward += 0.15
|
| 74 |
+
else:
|
| 75 |
+
reward -= 0.10
|
| 76 |
+
|
| 77 |
+
# 4. Ticket assigned to a valid team
|
| 78 |
+
if assigned_team is not None:
|
| 79 |
+
if valid_teams:
|
| 80 |
+
if assigned_team in valid_teams:
|
| 81 |
+
reward += 0.20
|
| 82 |
+
else:
|
| 83 |
+
reward -= 0.10
|
| 84 |
+
|
| 85 |
+
# 5. Posted to the right oncall channel
|
| 86 |
+
if posted_channels:
|
| 87 |
+
if oncall_channels:
|
| 88 |
+
correct_posts = [ch for ch in posted_channels if ch in oncall_channels]
|
| 89 |
+
wrong_posts = [ch for ch in posted_channels if ch not in oncall_channels]
|
| 90 |
+
if correct_posts:
|
| 91 |
+
reward += 0.25
|
| 92 |
+
reward -= 0.10 * len(wrong_posts) # -0.10 per channel-spray post
|
| 93 |
+
else:
|
| 94 |
+
# Posted without reading runbook — can't verify, mild penalty
|
| 95 |
+
reward -= 0.05 * len(posted_channels)
|
| 96 |
+
|
| 97 |
+
# 6. Env grader bonus — full env score adds on top when everything is correct
|
| 98 |
+
reward += env_score * 0.10
|
| 99 |
+
|
| 100 |
+
return max(-1.0, min(1.0, reward))
|
training/rollout.py
ADDED
|
@@ -0,0 +1,430 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Multi-turn rollout function for PM-Ops GRPO training.
|
| 2 |
+
|
| 3 |
+
History design (no duplicates):
|
| 4 |
+
Each turn record stores the obs that TRIGGERED the completion, not the result.
|
| 5 |
+
build_messages reconstructs:
|
| 6 |
+
[sys] [user:obs_0] [asst:comp_0] [user:obs_1] [asst:comp_1] ... [user:current_obs]
|
| 7 |
+
current_obs is never in history — it becomes the final user message.
|
| 8 |
+
|
| 9 |
+
Other design decisions:
|
| 10 |
+
- Action format B: chain-of-thought reasoning + ```json block
|
| 11 |
+
- Truncation: runbook-pinned sliding window (runbook pair always kept)
|
| 12 |
+
- Fallback cascade: read_runbook (early) → noop (mid) → finish (late)
|
| 13 |
+
- Three-pass JSON extraction: code block → raw JSON → regex
|
| 14 |
+
"""
|
| 15 |
+
import json
|
| 16 |
+
import re
|
| 17 |
+
from typing import Any
|
| 18 |
+
|
| 19 |
+
import torch
|
| 20 |
+
import torch.nn.functional as F
|
| 21 |
+
|
| 22 |
+
from training.dataset import parse_seed_from_prompt
|
| 23 |
+
from training.prompts import SYSTEM_PROMPT, format_observation
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
# ---------------------------------------------------------------------------
|
| 27 |
+
# HF model.generate() — replaces generate_rollout_completions (vLLM-only)
|
| 28 |
+
# ---------------------------------------------------------------------------
|
| 29 |
+
|
| 30 |
+
def _get_model_for_generation(trainer):
|
| 31 |
+
"""Unwrap the model safely regardless of accelerate/PEFT/DDP wrapping.
|
| 32 |
+
|
| 33 |
+
Priority:
|
| 34 |
+
1. accelerator.unwrap_model — handles DDP + PEFT + DeepSpeed
|
| 35 |
+
2. trainer.model.module — plain DDP wrapping
|
| 36 |
+
3. trainer.model — unwrapped (local or single-GPU)
|
| 37 |
+
"""
|
| 38 |
+
if hasattr(trainer, "accelerator"):
|
| 39 |
+
return trainer.accelerator.unwrap_model(trainer.model)
|
| 40 |
+
if hasattr(trainer.model, "module"):
|
| 41 |
+
return trainer.model.module
|
| 42 |
+
return trainer.model
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def _generate_no_vllm(trainer, prompt_text: str, tokenizer,
|
| 46 |
+
max_new_tokens: int = 512, temperature: float = 1.1) -> dict:
|
| 47 |
+
"""Generate one completion using HF model.generate() without vLLM.
|
| 48 |
+
|
| 49 |
+
Returns the same dict shape as generate_rollout_completions so the rest
|
| 50 |
+
of rollout_once is unchanged:
|
| 51 |
+
prompt_ids: list[int]
|
| 52 |
+
completion_ids: list[int]
|
| 53 |
+
logprobs: list[float] (per-token log-prob under current policy)
|
| 54 |
+
text: str
|
| 55 |
+
"""
|
| 56 |
+
model = _get_model_for_generation(trainer)
|
| 57 |
+
|
| 58 |
+
# Device: prefer accelerator.device, fall back to first param device
|
| 59 |
+
if hasattr(trainer, "accelerator"):
|
| 60 |
+
device = trainer.accelerator.device
|
| 61 |
+
else:
|
| 62 |
+
device = next(model.parameters()).device
|
| 63 |
+
|
| 64 |
+
enc = tokenizer(prompt_text, return_tensors="pt").to(device)
|
| 65 |
+
prompt_len = enc["input_ids"].shape[1]
|
| 66 |
+
|
| 67 |
+
with torch.no_grad():
|
| 68 |
+
out = model.generate(
|
| 69 |
+
**enc,
|
| 70 |
+
max_new_tokens=max_new_tokens,
|
| 71 |
+
do_sample=True,
|
| 72 |
+
temperature=temperature,
|
| 73 |
+
top_p=0.95,
|
| 74 |
+
top_k=50,
|
| 75 |
+
pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
|
| 76 |
+
output_scores=True,
|
| 77 |
+
return_dict_in_generate=True,
|
| 78 |
+
)
|
| 79 |
+
|
| 80 |
+
prompt_ids = enc["input_ids"][0].tolist()
|
| 81 |
+
completion_ids = out.sequences[0][prompt_len:].tolist()
|
| 82 |
+
|
| 83 |
+
# Per-token log-probs from output.scores (one score tensor per new token)
|
| 84 |
+
logprobs = [
|
| 85 |
+
F.log_softmax(score[0], dim=-1)[tok_id].item()
|
| 86 |
+
for score, tok_id in zip(out.scores, completion_ids)
|
| 87 |
+
]
|
| 88 |
+
|
| 89 |
+
text = tokenizer.decode(completion_ids, skip_special_tokens=True)
|
| 90 |
+
return {
|
| 91 |
+
"prompt_ids": prompt_ids,
|
| 92 |
+
"completion_ids": completion_ids,
|
| 93 |
+
"logprobs": logprobs,
|
| 94 |
+
"text": text,
|
| 95 |
+
}
|
| 96 |
+
|
| 97 |
+
MAX_STEPS = 40
|
| 98 |
+
# ~3000 tokens at 4 chars/token; leaves room for completion tokens
|
| 99 |
+
MAX_PROMPT_CHARS = 12_000
|
| 100 |
+
RUNBOOK_RESPONSE_MAX_CHARS = 3_000
|
| 101 |
+
HISTORY_PAIRS = 6 # max (user+asst) pairs kept from non-runbook history
|
| 102 |
+
|
| 103 |
+
|
| 104 |
+
# ---------------------------------------------------------------------------
|
| 105 |
+
# JSON extraction — three-pass, most-specific first
|
| 106 |
+
# ---------------------------------------------------------------------------
|
| 107 |
+
|
| 108 |
+
def extract_json_action(text: str) -> dict | None:
|
| 109 |
+
# Pass 1: last ```json ... ``` block
|
| 110 |
+
matches = re.findall(r"```json\s*(\{.*?\})\s*```", text, re.DOTALL)
|
| 111 |
+
if matches:
|
| 112 |
+
try:
|
| 113 |
+
return json.loads(matches[-1])
|
| 114 |
+
except json.JSONDecodeError:
|
| 115 |
+
pass
|
| 116 |
+
|
| 117 |
+
# Pass 2: entire output is JSON
|
| 118 |
+
stripped = text.strip()
|
| 119 |
+
if stripped.startswith("{"):
|
| 120 |
+
try:
|
| 121 |
+
return json.loads(stripped)
|
| 122 |
+
except json.JSONDecodeError:
|
| 123 |
+
pass
|
| 124 |
+
|
| 125 |
+
# Pass 3: any {...action_type...} pattern — allow one level of nested {} (e.g. "args": {})
|
| 126 |
+
matches = re.findall(
|
| 127 |
+
r'\{(?:[^{}]|\{[^{}]*\})*"action_type"\s*:\s*"[^"]*"(?:[^{}]|\{[^{}]*\})*\}',
|
| 128 |
+
text, re.DOTALL,
|
| 129 |
+
)
|
| 130 |
+
if matches:
|
| 131 |
+
try:
|
| 132 |
+
return json.loads(matches[-1])
|
| 133 |
+
except json.JSONDecodeError:
|
| 134 |
+
pass
|
| 135 |
+
|
| 136 |
+
return None
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
def step_aware_fallback(step: int, max_steps: int = MAX_STEPS) -> dict:
|
| 140 |
+
"""Safe fallback that degrades gracefully across the episode."""
|
| 141 |
+
if step <= 1:
|
| 142 |
+
return {"action_type": "meta.read_runbook", "args": {}}
|
| 143 |
+
elif step >= max_steps - 3:
|
| 144 |
+
return {"action_type": "meta.finish", "args": {}}
|
| 145 |
+
return {"action_type": "meta.noop", "args": {}}
|
| 146 |
+
|
| 147 |
+
|
| 148 |
+
# ---------------------------------------------------------------------------
|
| 149 |
+
# Context builder — runbook-pinned sliding window, no duplicate user turns
|
| 150 |
+
# ---------------------------------------------------------------------------
|
| 151 |
+
|
| 152 |
+
def _chars(messages: list[dict]) -> int:
|
| 153 |
+
return sum(len(m["content"]) for m in messages)
|
| 154 |
+
|
| 155 |
+
|
| 156 |
+
def build_messages(
|
| 157 |
+
turn_history: list[dict],
|
| 158 |
+
current_obs_text: str,
|
| 159 |
+
) -> list[dict]:
|
| 160 |
+
"""Build prompt messages with runbook-pinned sliding window truncation.
|
| 161 |
+
|
| 162 |
+
turn_history entries: {"obs_text": str, "completion": str, "is_runbook": bool}
|
| 163 |
+
obs_text = observation that triggered this completion (user side)
|
| 164 |
+
completion = model output for that step (assistant side)
|
| 165 |
+
|
| 166 |
+
Final conversation shape:
|
| 167 |
+
[sys] [user:obs_0][asst:comp_0] ... [user:obs_k][asst:comp_k] [user:current_obs]
|
| 168 |
+
No entry in turn_history represents current_obs — it's only the final user turn.
|
| 169 |
+
"""
|
| 170 |
+
# Task brief is always prepended to current_obs so it stays visible even
|
| 171 |
+
# when old context is truncated.
|
| 172 |
+
messages: list[dict] = [{"role": "system", "content": SYSTEM_PROMPT}]
|
| 173 |
+
|
| 174 |
+
# Separate the runbook exchange from the rest
|
| 175 |
+
runbook_turn: dict | None = None
|
| 176 |
+
general: list[dict] = []
|
| 177 |
+
|
| 178 |
+
for turn in turn_history:
|
| 179 |
+
if turn["is_runbook"] and runbook_turn is None:
|
| 180 |
+
runbook_turn = turn
|
| 181 |
+
else:
|
| 182 |
+
general.append(turn)
|
| 183 |
+
|
| 184 |
+
# Pin runbook exchange (truncate only if absurdly large)
|
| 185 |
+
if runbook_turn is not None:
|
| 186 |
+
rb_comp = runbook_turn["completion"][:RUNBOOK_RESPONSE_MAX_CHARS]
|
| 187 |
+
messages.append({"role": "user", "content": runbook_turn["obs_text"]})
|
| 188 |
+
messages.append({"role": "assistant", "content": rb_comp})
|
| 189 |
+
|
| 190 |
+
# Fill budget with most-recent general turns (newest first, then reverse)
|
| 191 |
+
fixed_chars = _chars(messages) + len(current_obs_text)
|
| 192 |
+
budget = MAX_PROMPT_CHARS - fixed_chars
|
| 193 |
+
window: list[dict] = []
|
| 194 |
+
|
| 195 |
+
for turn in reversed(general[-HISTORY_PAIRS:]):
|
| 196 |
+
pair_chars = len(turn["obs_text"]) + len(turn["completion"])
|
| 197 |
+
if budget - pair_chars < 0:
|
| 198 |
+
break
|
| 199 |
+
window.append(turn)
|
| 200 |
+
budget -= pair_chars
|
| 201 |
+
|
| 202 |
+
for turn in reversed(window):
|
| 203 |
+
messages.append({"role": "user", "content": turn["obs_text"]})
|
| 204 |
+
messages.append({"role": "assistant", "content": turn["completion"]})
|
| 205 |
+
|
| 206 |
+
# Current observation is always the final user turn (never stored in history)
|
| 207 |
+
messages.append({"role": "user", "content": current_obs_text})
|
| 208 |
+
return messages
|
| 209 |
+
|
| 210 |
+
|
| 211 |
+
# ---------------------------------------------------------------------------
|
| 212 |
+
# Observation normaliser — handles both object and dict forms
|
| 213 |
+
# ---------------------------------------------------------------------------
|
| 214 |
+
|
| 215 |
+
def _obs_to_dict(obs: Any) -> dict:
|
| 216 |
+
if isinstance(obs, dict):
|
| 217 |
+
return obs
|
| 218 |
+
fields = ("task_brief", "last_action_result", "step", "steps_remaining",
|
| 219 |
+
"reward", "done", "token_budget_remaining")
|
| 220 |
+
return {f: getattr(obs, f, None) for f in fields}
|
| 221 |
+
|
| 222 |
+
|
| 223 |
+
def _current_obs_text(obs_dict: dict, step: int, task_brief: str) -> str:
|
| 224 |
+
"""Format observation, prepending a task reminder so it survives truncation."""
|
| 225 |
+
reminder = f"**Task reminder:** {task_brief[:200]}\n\n"
|
| 226 |
+
return reminder + format_observation(obs_dict, step, obs_dict.get("last_action_result"))
|
| 227 |
+
|
| 228 |
+
|
| 229 |
+
# ---------------------------------------------------------------------------
|
| 230 |
+
# Single-episode rollout
|
| 231 |
+
# ---------------------------------------------------------------------------
|
| 232 |
+
|
| 233 |
+
def rollout_once(
|
| 234 |
+
trainer,
|
| 235 |
+
sync_env,
|
| 236 |
+
tokenizer,
|
| 237 |
+
dataset_prompt: str,
|
| 238 |
+
max_steps: int = 15,
|
| 239 |
+
gen_offset: int = 0,
|
| 240 |
+
) -> dict:
|
| 241 |
+
"""Play one full PM-Ops episode. Returns trajectory + reward signals.
|
| 242 |
+
|
| 243 |
+
gen_offset: added to seed so each GRPO generation explores a different env
|
| 244 |
+
episode even when receiving the same prompt (same base seed).
|
| 245 |
+
"""
|
| 246 |
+
seed = parse_seed_from_prompt(dataset_prompt)
|
| 247 |
+
if seed is not None:
|
| 248 |
+
result = sync_env.reset(seed=seed + gen_offset)
|
| 249 |
+
else:
|
| 250 |
+
result = sync_env.reset()
|
| 251 |
+
|
| 252 |
+
obs = result.observation if hasattr(result, "observation") else result
|
| 253 |
+
obs_dict = _obs_to_dict(obs)
|
| 254 |
+
task_brief: str = obs_dict.get("task_brief") or dataset_prompt
|
| 255 |
+
|
| 256 |
+
# Flat trajectory buffers (TRL expects flat lists across all steps)
|
| 257 |
+
prompt_ids: list = []
|
| 258 |
+
completion_ids: list = []
|
| 259 |
+
logprobs: list = []
|
| 260 |
+
|
| 261 |
+
# Turn history for context building
|
| 262 |
+
# Each entry: {"obs_text": str, "completion": str, "is_runbook": bool}
|
| 263 |
+
turn_history: list[dict] = []
|
| 264 |
+
|
| 265 |
+
# Rollout accumulators
|
| 266 |
+
valid_action_count = 0
|
| 267 |
+
final_score = 0.0
|
| 268 |
+
step = 0
|
| 269 |
+
done = False
|
| 270 |
+
|
| 271 |
+
# Runbook-compliance reward tracking
|
| 272 |
+
read_runbook_done = False
|
| 273 |
+
valid_labels: set[str] = set() # org label_taxonomy values
|
| 274 |
+
valid_priorities: set[str] = set() # org priority_levels
|
| 275 |
+
valid_teams: set[str] = set() # org team_map values
|
| 276 |
+
oncall_channels: set[str] = set() # org oncall_channels values
|
| 277 |
+
ticket_label: str | None = None # label used in create_ticket
|
| 278 |
+
ticket_priority: str | None = None # priority used in create_ticket
|
| 279 |
+
assigned_team: str | None = None # team from assign_ticket
|
| 280 |
+
posted_channels: list[str] = [] # every channel posted to
|
| 281 |
+
|
| 282 |
+
while not done and step < max_steps:
|
| 283 |
+
# obs_text for THIS step — stored in history BEFORE stepping
|
| 284 |
+
obs_text = _current_obs_text(obs_dict, step, task_brief)
|
| 285 |
+
|
| 286 |
+
messages = build_messages(turn_history, obs_text)
|
| 287 |
+
prompt_text = tokenizer.apply_chat_template(
|
| 288 |
+
messages,
|
| 289 |
+
add_generation_prompt=True,
|
| 290 |
+
tokenize=False,
|
| 291 |
+
enable_thinking=False,
|
| 292 |
+
)
|
| 293 |
+
|
| 294 |
+
rollout_out = _generate_no_vllm(trainer, prompt_text, tokenizer)
|
| 295 |
+
prompt_ids.extend(rollout_out["prompt_ids"])
|
| 296 |
+
completion_ids.extend(rollout_out["completion_ids"])
|
| 297 |
+
logprobs.extend(rollout_out["logprobs"])
|
| 298 |
+
|
| 299 |
+
completion_text = rollout_out["text"]
|
| 300 |
+
|
| 301 |
+
# Parse action; fall back gracefully on parse failure
|
| 302 |
+
parsed = extract_json_action(completion_text)
|
| 303 |
+
is_valid_json = parsed is not None
|
| 304 |
+
|
| 305 |
+
if not is_valid_json:
|
| 306 |
+
# Log the raw output on step 0 to diagnose format failures
|
| 307 |
+
if step == 0 and valid_action_count == 0:
|
| 308 |
+
snippet = repr(completion_text[:300])
|
| 309 |
+
print(f"[rollout] step=0 NO JSON — raw output: {snippet}")
|
| 310 |
+
parsed = step_aware_fallback(step, max_steps)
|
| 311 |
+
else:
|
| 312 |
+
valid_action_count += 1
|
| 313 |
+
|
| 314 |
+
action_type: str = parsed.get("action_type", "meta.noop")
|
| 315 |
+
args: dict = parsed.get("args", {})
|
| 316 |
+
|
| 317 |
+
if action_type == "meta.read_runbook" and is_valid_json:
|
| 318 |
+
read_runbook_done = True
|
| 319 |
+
|
| 320 |
+
if action_type == "ticketing.create_ticket" and is_valid_json and ticket_label is None:
|
| 321 |
+
ticket_label = args.get("label")
|
| 322 |
+
ticket_priority = args.get("priority")
|
| 323 |
+
|
| 324 |
+
if action_type == "ticketing.assign_ticket" and is_valid_json and assigned_team is None:
|
| 325 |
+
assigned_team = args.get("team")
|
| 326 |
+
|
| 327 |
+
if action_type == "chat.post_message" and is_valid_json:
|
| 328 |
+
ch = args.get("channel", "")
|
| 329 |
+
if ch:
|
| 330 |
+
posted_channels.append(ch)
|
| 331 |
+
|
| 332 |
+
# Store the (obs_text, completion) pair BEFORE stepping the env
|
| 333 |
+
# is_runbook marks this turn for pinning in future context windows
|
| 334 |
+
turn_history.append({
|
| 335 |
+
"obs_text": obs_text,
|
| 336 |
+
"completion": completion_text,
|
| 337 |
+
"is_runbook": (action_type == "meta.read_runbook" and is_valid_json),
|
| 338 |
+
})
|
| 339 |
+
|
| 340 |
+
# Step the environment — obs_dict now holds the NEXT state
|
| 341 |
+
result = sync_env.step({"action_type": action_type, "args": args})
|
| 342 |
+
new_obs = result.observation if hasattr(result, "observation") else result
|
| 343 |
+
obs_dict = _obs_to_dict(new_obs)
|
| 344 |
+
|
| 345 |
+
# Extract full org_config from runbook response (one step after the call)
|
| 346 |
+
last_result = obs_dict.get("last_action_result") or {}
|
| 347 |
+
if action_type == "meta.read_runbook" and last_result.get("ok"):
|
| 348 |
+
data = last_result.get("data") or {}
|
| 349 |
+
if isinstance(data, dict):
|
| 350 |
+
org = data.get("org_config") or {}
|
| 351 |
+
valid_labels = set(org.get("label_taxonomy", {}).values())
|
| 352 |
+
valid_priorities = set(org.get("priority_levels", []))
|
| 353 |
+
valid_teams = set(org.get("team_map", {}).values())
|
| 354 |
+
oncall_channels = set(org.get("oncall_channels", {}).values())
|
| 355 |
+
|
| 356 |
+
done = bool(getattr(result, "done", obs_dict.get("done", False)))
|
| 357 |
+
final_score = float(getattr(result, "reward", obs_dict.get("reward", 0.0)))
|
| 358 |
+
step += 1
|
| 359 |
+
|
| 360 |
+
from training.rewards import compute_rollout_reward
|
| 361 |
+
combined = compute_rollout_reward(
|
| 362 |
+
read_runbook_done = read_runbook_done,
|
| 363 |
+
valid_labels = valid_labels,
|
| 364 |
+
valid_priorities = valid_priorities,
|
| 365 |
+
valid_teams = valid_teams,
|
| 366 |
+
oncall_channels = oncall_channels,
|
| 367 |
+
ticket_label = ticket_label,
|
| 368 |
+
ticket_priority = ticket_priority,
|
| 369 |
+
assigned_team = assigned_team,
|
| 370 |
+
posted_channels = posted_channels,
|
| 371 |
+
env_score = final_score,
|
| 372 |
+
valid_json_count = valid_action_count,
|
| 373 |
+
)
|
| 374 |
+
|
| 375 |
+
print(
|
| 376 |
+
f"[rollout] steps={step} env={final_score:.3f} "
|
| 377 |
+
f"label={'✓' if ticket_label and ticket_label in valid_labels else '✗' if ticket_label else '-'} "
|
| 378 |
+
f"priority={'✓' if ticket_priority and ticket_priority in valid_priorities else '✗' if ticket_priority else '-'} "
|
| 379 |
+
f"team={'✓' if assigned_team and assigned_team in valid_teams else '✗' if assigned_team else '-'} "
|
| 380 |
+
f"channel={'✓' if any(ch in oncall_channels for ch in posted_channels) else '✗' if posted_channels else '-'} "
|
| 381 |
+
f"→ reward={combined:.3f}"
|
| 382 |
+
)
|
| 383 |
+
return {
|
| 384 |
+
"prompt_ids": prompt_ids,
|
| 385 |
+
"completion_ids": completion_ids,
|
| 386 |
+
"logprobs": logprobs,
|
| 387 |
+
"reward": combined,
|
| 388 |
+
}
|
| 389 |
+
|
| 390 |
+
|
| 391 |
+
# ---------------------------------------------------------------------------
|
| 392 |
+
# GRPOTrainer-compatible rollout function (factory)
|
| 393 |
+
# ---------------------------------------------------------------------------
|
| 394 |
+
|
| 395 |
+
def make_rollout_func(sync_env, tokenizer, max_steps: int = 15):
|
| 396 |
+
"""Bind env + tokenizer; return the function GRPOTrainer calls each batch.
|
| 397 |
+
|
| 398 |
+
max_steps per task:
|
| 399 |
+
triage → 15 (solvable in 5, cap gives room for exploration)
|
| 400 |
+
incident_routing → 20
|
| 401 |
+
release_notes → 30
|
| 402 |
+
dep_update → 30
|
| 403 |
+
"""
|
| 404 |
+
def rollout_func(prompts: list[str], trainer=None) -> dict:
|
| 405 |
+
out: dict[str, list] = {
|
| 406 |
+
"prompt_ids": [],
|
| 407 |
+
"completion_ids": [],
|
| 408 |
+
"logprobs": [],
|
| 409 |
+
"reward": [],
|
| 410 |
+
}
|
| 411 |
+
# Track how many times each unique prompt has appeared so we can pass
|
| 412 |
+
# a gen_offset — ensures repeated prompts (num_generations > 1) hit
|
| 413 |
+
# different env seeds and produce different rollouts.
|
| 414 |
+
prompt_seen: dict[str, int] = {}
|
| 415 |
+
for prompt_text in prompts:
|
| 416 |
+
gen_offset = prompt_seen.get(prompt_text, 0)
|
| 417 |
+
prompt_seen[prompt_text] = gen_offset + 1
|
| 418 |
+
episode = rollout_once(
|
| 419 |
+
trainer=trainer,
|
| 420 |
+
sync_env=sync_env,
|
| 421 |
+
tokenizer=tokenizer,
|
| 422 |
+
dataset_prompt=prompt_text,
|
| 423 |
+
max_steps=max_steps,
|
| 424 |
+
gen_offset=gen_offset,
|
| 425 |
+
)
|
| 426 |
+
for k in out:
|
| 427 |
+
out[k].append(episode[k])
|
| 428 |
+
return out
|
| 429 |
+
|
| 430 |
+
return rollout_func
|
training/smoke_test.py
ADDED
|
@@ -0,0 +1,267 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Smoke test — runs locally without TRL/vLLM.
|
| 2 |
+
|
| 3 |
+
Tests in order:
|
| 4 |
+
1. Dataset generation
|
| 5 |
+
2. JSON extraction (edge cases)
|
| 6 |
+
3. build_messages structure (no duplicates, correct alternation)
|
| 7 |
+
4. Server startup + env connection
|
| 8 |
+
5. One full heuristic episode (env API end-to-end)
|
| 9 |
+
|
| 10 |
+
Run from repo root: python training/smoke_test.py
|
| 11 |
+
"""
|
| 12 |
+
import json
|
| 13 |
+
import os
|
| 14 |
+
import subprocess
|
| 15 |
+
import sys
|
| 16 |
+
import time
|
| 17 |
+
|
| 18 |
+
import requests
|
| 19 |
+
|
| 20 |
+
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
| 21 |
+
|
| 22 |
+
PASS = "\033[92m[PASS]\033[0m"
|
| 23 |
+
FAIL = "\033[91m[FAIL]\033[0m"
|
| 24 |
+
INFO = "\033[94m[INFO]\033[0m"
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def check(label: str, cond: bool, detail: str = "") -> bool:
|
| 28 |
+
if cond:
|
| 29 |
+
print(f"{PASS} {label}")
|
| 30 |
+
else:
|
| 31 |
+
print(f"{FAIL} {label}" + (f" — {detail}" if detail else ""))
|
| 32 |
+
return cond
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
# ---------------------------------------------------------------------------
|
| 36 |
+
# 1. Dataset
|
| 37 |
+
# ---------------------------------------------------------------------------
|
| 38 |
+
print("\n── 1. Dataset ──────────────────────────────────────────────────────")
|
| 39 |
+
from training.dataset import generate_triage_dataset, parse_seed_from_prompt
|
| 40 |
+
|
| 41 |
+
rows = generate_triage_dataset(n_episodes=5, base_seed=99)
|
| 42 |
+
check("generates 5 rows", len(rows) == 5)
|
| 43 |
+
check("each row has prompt+seed+difficulty", all(
|
| 44 |
+
"prompt" in r and "seed" in r and "difficulty" in r for r in rows
|
| 45 |
+
))
|
| 46 |
+
seed_back = parse_seed_from_prompt(rows[0]["prompt"])
|
| 47 |
+
check("seed round-trips through prompt string", seed_back == rows[0]["seed"],
|
| 48 |
+
f"got {seed_back}, expected {rows[0]['seed']}")
|
| 49 |
+
print(f" sample prompt: {rows[0]['prompt'][:100]}...")
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
# ---------------------------------------------------------------------------
|
| 53 |
+
# 2. JSON extraction
|
| 54 |
+
# ---------------------------------------------------------------------------
|
| 55 |
+
print("\n── 2. JSON extraction ──────────────────────────────────────────────")
|
| 56 |
+
from training.rollout import extract_json_action, step_aware_fallback
|
| 57 |
+
|
| 58 |
+
cases = [
|
| 59 |
+
# (description, input_text, expected_action_type)
|
| 60 |
+
("code block", 'Thinking...\n```json\n{"action_type": "meta.read_runbook", "args": {}}\n```', "meta.read_runbook"),
|
| 61 |
+
("raw JSON", '{"action_type": "meta.finish", "args": {}}', "meta.finish"),
|
| 62 |
+
("regex hit", 'I will do this: {"action_type": "meta.noop", "args": {}} done.', "meta.noop"),
|
| 63 |
+
("no JSON", "I cannot decide.", None),
|
| 64 |
+
("bad block", '```json\n{broken\n```', None),
|
| 65 |
+
]
|
| 66 |
+
for desc, text, expected in cases:
|
| 67 |
+
result = extract_json_action(text)
|
| 68 |
+
got = result.get("action_type") if result else None
|
| 69 |
+
check(f"extract: {desc}", got == expected, f"got {got!r}, expected {expected!r}")
|
| 70 |
+
|
| 71 |
+
fb = step_aware_fallback(0)
|
| 72 |
+
check("fallback step 0 → read_runbook", fb["action_type"] == "meta.read_runbook")
|
| 73 |
+
fb = step_aware_fallback(20)
|
| 74 |
+
check("fallback step 20 → noop", fb["action_type"] == "meta.noop")
|
| 75 |
+
fb = step_aware_fallback(38)
|
| 76 |
+
check("fallback step 38 → finish", fb["action_type"] == "meta.finish")
|
| 77 |
+
|
| 78 |
+
# reward_no_wrong_channels logic
|
| 79 |
+
from training.rewards import (
|
| 80 |
+
WEIGHT_FINAL_SCORE, WEIGHT_NO_WRONG_CHANNELS,
|
| 81 |
+
WEIGHT_VALID_JSON, WEIGHT_READ_RUNBOOK, WEIGHT_EFFICIENCY,
|
| 82 |
+
)
|
| 83 |
+
weights_sum = (WEIGHT_FINAL_SCORE + WEIGHT_NO_WRONG_CHANNELS +
|
| 84 |
+
WEIGHT_VALID_JSON + WEIGHT_READ_RUNBOOK + WEIGHT_EFFICIENCY)
|
| 85 |
+
check("reward weights sum to 1.0", abs(weights_sum - 1.0) < 1e-9, f"sum={weights_sum}")
|
| 86 |
+
|
| 87 |
+
|
| 88 |
+
# ---------------------------------------------------------------------------
|
| 89 |
+
# 3. build_messages structure
|
| 90 |
+
# ---------------------------------------------------------------------------
|
| 91 |
+
print("\n── 3. build_messages ───────────────────────────────────────────────")
|
| 92 |
+
from training.rollout import build_messages
|
| 93 |
+
|
| 94 |
+
# Empty history (step 0)
|
| 95 |
+
msgs = build_messages([], "First observation")
|
| 96 |
+
check("step 0: starts with system", msgs[0]["role"] == "system")
|
| 97 |
+
check("step 0: ends with user", msgs[-1]["role"] == "user")
|
| 98 |
+
check("step 0: final user is current obs", msgs[-1]["content"] == "First observation")
|
| 99 |
+
|
| 100 |
+
# Simulate 3 turns
|
| 101 |
+
history = [
|
| 102 |
+
{"obs_text": "obs_0", "completion": "comp_0", "is_runbook": True},
|
| 103 |
+
{"obs_text": "obs_1", "completion": "comp_1", "is_runbook": False},
|
| 104 |
+
{"obs_text": "obs_2", "completion": "comp_2", "is_runbook": False},
|
| 105 |
+
]
|
| 106 |
+
msgs = build_messages(history, "current_obs")
|
| 107 |
+
|
| 108 |
+
# Validate alternation: after system, must be user/asst/user/asst/.../user
|
| 109 |
+
roles = [m["role"] for m in msgs]
|
| 110 |
+
check("roles start with system", roles[0] == "system")
|
| 111 |
+
check("roles end with user", roles[-1] == "user")
|
| 112 |
+
pairs_ok = all(
|
| 113 |
+
roles[i] == "user" and roles[i+1] == "assistant"
|
| 114 |
+
for i in range(1, len(roles) - 2, 2)
|
| 115 |
+
)
|
| 116 |
+
check("strict user/asst alternation throughout", pairs_ok, str(roles))
|
| 117 |
+
|
| 118 |
+
# Verify current_obs appears exactly once as the last message
|
| 119 |
+
final_user_content = msgs[-1]["content"]
|
| 120 |
+
check("current_obs is last message content", final_user_content == "current_obs")
|
| 121 |
+
all_contents = [m["content"] for m in msgs]
|
| 122 |
+
check("current_obs not duplicated", all_contents.count("current_obs") == 1)
|
| 123 |
+
|
| 124 |
+
# Verify runbook is pinned (obs_0/comp_0 should appear even with many turns)
|
| 125 |
+
all_content_str = " ".join(all_contents)
|
| 126 |
+
check("runbook obs pinned", "obs_0" in all_content_str)
|
| 127 |
+
check("runbook completion pinned", "comp_0" in all_content_str)
|
| 128 |
+
|
| 129 |
+
|
| 130 |
+
# ---------------------------------------------------------------------------
|
| 131 |
+
# 4. Server startup
|
| 132 |
+
# ---------------------------------------------------------------------------
|
| 133 |
+
print("\n── 4. Server startup ───────────────────────────────────────────────")
|
| 134 |
+
REPO_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
| 135 |
+
ENV_URL = "http://localhost:8765" # non-standard port to avoid conflicts
|
| 136 |
+
|
| 137 |
+
proc = subprocess.Popen(
|
| 138 |
+
[sys.executable, "-m", "uvicorn", "server.app:app",
|
| 139 |
+
"--host", "0.0.0.0", "--port", "8765"],
|
| 140 |
+
cwd=REPO_DIR,
|
| 141 |
+
stdout=subprocess.DEVNULL,
|
| 142 |
+
stderr=subprocess.DEVNULL,
|
| 143 |
+
)
|
| 144 |
+
|
| 145 |
+
ready = False
|
| 146 |
+
for _ in range(20):
|
| 147 |
+
try:
|
| 148 |
+
r = requests.get(f"{ENV_URL}/", timeout=1)
|
| 149 |
+
if r.status_code == 200:
|
| 150 |
+
ready = True
|
| 151 |
+
break
|
| 152 |
+
except Exception:
|
| 153 |
+
pass
|
| 154 |
+
time.sleep(1)
|
| 155 |
+
|
| 156 |
+
check("server started on :8765", ready)
|
| 157 |
+
if not ready:
|
| 158 |
+
proc.terminate()
|
| 159 |
+
print(f"{FAIL} Cannot continue without server.")
|
| 160 |
+
sys.exit(1)
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
# ---------------------------------------------------------------------------
|
| 164 |
+
# 5. Full heuristic episode (env API end-to-end)
|
| 165 |
+
# ---------------------------------------------------------------------------
|
| 166 |
+
print("\n── 5. End-to-end episode ───────────────────────────────────────────")
|
| 167 |
+
from openenv.core import GenericEnvClient
|
| 168 |
+
|
| 169 |
+
client = GenericEnvClient(base_url=ENV_URL).sync()
|
| 170 |
+
client.connect()
|
| 171 |
+
|
| 172 |
+
try:
|
| 173 |
+
# reset with explicit seed (tests seed passthrough)
|
| 174 |
+
res = client.reset(seed=rows[0]["seed"])
|
| 175 |
+
obs = res.observation if hasattr(res, "observation") else res
|
| 176 |
+
|
| 177 |
+
def _get(o, k):
|
| 178 |
+
return getattr(o, k, None) if not isinstance(o, dict) else o.get(k)
|
| 179 |
+
|
| 180 |
+
task_brief = _get(obs, "task_brief")
|
| 181 |
+
check("reset returns task_brief", bool(task_brief), repr(task_brief)[:80])
|
| 182 |
+
|
| 183 |
+
# Step 1: read runbook
|
| 184 |
+
res = client.step({"action_type": "meta.read_runbook", "args": {}})
|
| 185 |
+
obs = res.observation if hasattr(res, "observation") else res
|
| 186 |
+
last = _get(obs, "last_action_result") or {}
|
| 187 |
+
check("read_runbook ok=True", last.get("ok") is True)
|
| 188 |
+
check("runbook contains org_config", "org_config" in str(last.get("data", "")))
|
| 189 |
+
|
| 190 |
+
org = last.get("data", {}).get("org_config", {}) if isinstance(last.get("data"), dict) else {}
|
| 191 |
+
labels = list(org.get("label_taxonomy", {}).values())
|
| 192 |
+
priorities = org.get("priority_levels", [])
|
| 193 |
+
teams = list(org.get("team_map", {}).values())
|
| 194 |
+
channels = list(org.get("oncall_channels", {}).values())
|
| 195 |
+
|
| 196 |
+
check("org has labels", len(labels) > 0)
|
| 197 |
+
check("org has priorities", len(priorities) > 0)
|
| 198 |
+
|
| 199 |
+
# Step 2: create ticket with correct label/priority
|
| 200 |
+
res = client.step({
|
| 201 |
+
"action_type": "ticketing.create_ticket",
|
| 202 |
+
"args": {
|
| 203 |
+
"summary": f"Triage: {task_brief[:60]}",
|
| 204 |
+
"description": task_brief,
|
| 205 |
+
"label": labels[0],
|
| 206 |
+
"priority": priorities[1] if len(priorities) > 1 else priorities[0],
|
| 207 |
+
"assignee": teams[0] if teams else "backend",
|
| 208 |
+
},
|
| 209 |
+
})
|
| 210 |
+
obs = res.observation if hasattr(res, "observation") else res
|
| 211 |
+
last = _get(obs, "last_action_result") or {}
|
| 212 |
+
check("create_ticket ok=True", last.get("ok") is True, str(last.get("error")))
|
| 213 |
+
ticket_id = (last.get("data") or {}).get("id") if isinstance(last.get("data"), dict) else None
|
| 214 |
+
check("ticket_id returned", bool(ticket_id), repr(ticket_id))
|
| 215 |
+
|
| 216 |
+
# Step 3: assign ticket
|
| 217 |
+
if ticket_id and teams:
|
| 218 |
+
res = client.step({
|
| 219 |
+
"action_type": "ticketing.assign_ticket",
|
| 220 |
+
"args": {"ticket_id": ticket_id, "team": teams[0]},
|
| 221 |
+
})
|
| 222 |
+
obs = res.observation if hasattr(res, "observation") else res
|
| 223 |
+
last = _get(obs, "last_action_result") or {}
|
| 224 |
+
check("assign_ticket ok=True", last.get("ok") is True, str(last.get("error")))
|
| 225 |
+
|
| 226 |
+
# Step 4: post to oncall channel
|
| 227 |
+
if channels:
|
| 228 |
+
res = client.step({
|
| 229 |
+
"action_type": "chat.post_message",
|
| 230 |
+
"args": {"channel": channels[0], "text": f"Triage smoke test: {task_brief[:80]}"},
|
| 231 |
+
})
|
| 232 |
+
obs = res.observation if hasattr(res, "observation") else res
|
| 233 |
+
last = _get(obs, "last_action_result") or {}
|
| 234 |
+
check("post_message ok=True", last.get("ok") is True, str(last.get("error")))
|
| 235 |
+
|
| 236 |
+
# Step 5: finish + check reward
|
| 237 |
+
res = client.step({"action_type": "meta.finish", "args": {}})
|
| 238 |
+
done = getattr(res, "done", None)
|
| 239 |
+
reward = getattr(res, "reward", 0.0)
|
| 240 |
+
check("episode done after finish", done is True, f"done={done}")
|
| 241 |
+
check("reward is float in [0,1]", isinstance(reward, float) and 0.0 <= reward <= 1.0,
|
| 242 |
+
f"reward={reward}")
|
| 243 |
+
print(f" episode reward: {reward:.3f}")
|
| 244 |
+
|
| 245 |
+
# Verify make_rollout_func accepts max_steps param
|
| 246 |
+
from training.rollout import make_rollout_func
|
| 247 |
+
import inspect
|
| 248 |
+
sig = inspect.signature(make_rollout_func)
|
| 249 |
+
check("make_rollout_func has max_steps param", "max_steps" in sig.parameters)
|
| 250 |
+
check("max_steps default is 15", sig.parameters["max_steps"].default == 15)
|
| 251 |
+
|
| 252 |
+
# Verify no_wrong_channels_reward key in rollout output keys
|
| 253 |
+
from training.rollout import rollout_once as _rollout_once
|
| 254 |
+
sig2 = inspect.signature(_rollout_once)
|
| 255 |
+
check("rollout_once has max_steps param", "max_steps" in sig2.parameters)
|
| 256 |
+
|
| 257 |
+
finally:
|
| 258 |
+
client.close()
|
| 259 |
+
|
| 260 |
+
proc.terminate()
|
| 261 |
+
|
| 262 |
+
# ---------------------------------------------------------------------------
|
| 263 |
+
# Summary
|
| 264 |
+
# ---------------------------------------------------------------------------
|
| 265 |
+
print("\n────────────────────────────────────────────────────────────────────")
|
| 266 |
+
print("Smoke test complete. If all checks passed, training pipeline is ready.")
|
| 267 |
+
print("Next: push repo to HF and run training/train.ipynb on A100 GPU Space.")
|
training/train.ipynb
ADDED
|
@@ -0,0 +1,669 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"cells": [
|
| 3 |
+
{
|
| 4 |
+
"cell_type": "markdown",
|
| 5 |
+
"metadata": {},
|
| 6 |
+
"source": [
|
| 7 |
+
"# PM-Ops: Train a Project Management Agent with GRPO + Unsloth\n",
|
| 8 |
+
"\n",
|
| 9 |
+
"Fine-tune **Qwen3-1.7B** on PM-Ops triage tasks using GRPO (TRL) + Unsloth for memory efficiency.\n",
|
| 10 |
+
"\n",
|
| 11 |
+
"**GPU:** A100 40GB (HF Jupyter Space — uses HF credits) \n",
|
| 12 |
+
"**Time:** ~90 min (150 triage episodes, 1 epoch) \n",
|
| 13 |
+
"**Stack:** OpenEnv + TRL + Unsloth (hackathon-recommended stack)\n",
|
| 14 |
+
"\n",
|
| 15 |
+
"### Architecture\n",
|
| 16 |
+
"```\n",
|
| 17 |
+
"[HF A100 GPU Space]\n",
|
| 18 |
+
" ├── Unsloth + GRPOTrainer ← Qwen3-1.7B with LoRA, 4-bit quantised\n",
|
| 19 |
+
" └── PM-Ops FastAPI (localhost:8000) ← subprocess, <1ms per step\n",
|
| 20 |
+
"```\n",
|
| 21 |
+
"\n",
|
| 22 |
+
"### Reward design (5 independent signals — anti-hack per hackathon guide)\n",
|
| 23 |
+
"| Signal | Weight | Purpose |\n",
|
| 24 |
+
"|---|---|---|\n",
|
| 25 |
+
"| `reward_final_score` | 0.45 | Correctness: label + priority + team + channel |\n",
|
| 26 |
+
"| `reward_no_wrong_channels` | 0.15 | Anti-hack: penalise channel-spray |\n",
|
| 27 |
+
"| `reward_valid_json` | 0.15 | Format discipline |\n",
|
| 28 |
+
"| `reward_read_runbook` | 0.15 | Process: read before acting |\n",
|
| 29 |
+
"| `reward_efficiency` | 0.10 | Speed when correct |"
|
| 30 |
+
]
|
| 31 |
+
},
|
| 32 |
+
{
|
| 33 |
+
"cell_type": "markdown",
|
| 34 |
+
"metadata": {},
|
| 35 |
+
"source": [
|
| 36 |
+
"## 0. Install Dependencies"
|
| 37 |
+
]
|
| 38 |
+
},
|
| 39 |
+
{
|
| 40 |
+
"cell_type": "code",
|
| 41 |
+
"execution_count": null,
|
| 42 |
+
"metadata": {},
|
| 43 |
+
"outputs": [],
|
| 44 |
+
"source": [
|
| 45 |
+
"# Unsloth first — it pins compatible versions of torch/transformers\n",
|
| 46 |
+
"!pip install -Uq unsloth\n",
|
| 47 |
+
"!pip install -Uq \"trl>=0.17.0\" openenv-core datasets trackio\n",
|
| 48 |
+
"print('Dependencies installed.')"
|
| 49 |
+
]
|
| 50 |
+
},
|
| 51 |
+
{
|
| 52 |
+
"cell_type": "markdown",
|
| 53 |
+
"metadata": {},
|
| 54 |
+
"source": [
|
| 55 |
+
"## 1. Clone PM-Ops Repo"
|
| 56 |
+
]
|
| 57 |
+
},
|
| 58 |
+
{
|
| 59 |
+
"cell_type": "code",
|
| 60 |
+
"execution_count": null,
|
| 61 |
+
"metadata": {},
|
| 62 |
+
"outputs": [],
|
| 63 |
+
"source": [
|
| 64 |
+
"import os, sys\n",
|
| 65 |
+
"\n",
|
| 66 |
+
"REPO_URL = 'https://huggingface.co/spaces/adityaguntur/pm-ops'\n",
|
| 67 |
+
"REPO_DIR = '/content/pm_ops'\n",
|
| 68 |
+
"\n",
|
| 69 |
+
"if not os.path.exists(REPO_DIR):\n",
|
| 70 |
+
" !git clone --depth=1 -q {REPO_URL} {REPO_DIR}\n",
|
| 71 |
+
" print(f'Cloned → {REPO_DIR}')\n",
|
| 72 |
+
"else:\n",
|
| 73 |
+
" print(f'Already exists: {REPO_DIR}')\n",
|
| 74 |
+
"\n",
|
| 75 |
+
"for p in [REPO_DIR, os.path.join(REPO_DIR, 'training')]:\n",
|
| 76 |
+
" if p not in sys.path:\n",
|
| 77 |
+
" sys.path.insert(0, p)\n",
|
| 78 |
+
"\n",
|
| 79 |
+
"os.chdir(REPO_DIR)\n",
|
| 80 |
+
"print(f'Working directory: {os.getcwd()}')"
|
| 81 |
+
]
|
| 82 |
+
},
|
| 83 |
+
{
|
| 84 |
+
"cell_type": "markdown",
|
| 85 |
+
"metadata": {},
|
| 86 |
+
"source": [
|
| 87 |
+
"## 2. HuggingFace Login"
|
| 88 |
+
]
|
| 89 |
+
},
|
| 90 |
+
{
|
| 91 |
+
"cell_type": "code",
|
| 92 |
+
"execution_count": null,
|
| 93 |
+
"metadata": {},
|
| 94 |
+
"outputs": [],
|
| 95 |
+
"source": [
|
| 96 |
+
"from huggingface_hub import notebook_login\n",
|
| 97 |
+
"notebook_login()"
|
| 98 |
+
]
|
| 99 |
+
},
|
| 100 |
+
{
|
| 101 |
+
"cell_type": "markdown",
|
| 102 |
+
"metadata": {},
|
| 103 |
+
"source": [
|
| 104 |
+
"## 3. Start PM-Ops Server Locally\n",
|
| 105 |
+
"\n",
|
| 106 |
+
"Running server on localhost eliminates the ~200ms/step network cost of calling the HF Space.\n",
|
| 107 |
+
"With 15 steps × 150 episodes × 2 generations that's 4,500 saved round-trips."
|
| 108 |
+
]
|
| 109 |
+
},
|
| 110 |
+
{
|
| 111 |
+
"cell_type": "code",
|
| 112 |
+
"execution_count": null,
|
| 113 |
+
"metadata": {},
|
| 114 |
+
"outputs": [],
|
| 115 |
+
"source": [
|
| 116 |
+
"import subprocess, time, requests\n",
|
| 117 |
+
"\n",
|
| 118 |
+
"server_proc = subprocess.Popen(\n",
|
| 119 |
+
" [sys.executable, '-m', 'uvicorn', 'server.app:app',\n",
|
| 120 |
+
" '--host', '0.0.0.0', '--port', '8000'],\n",
|
| 121 |
+
" cwd=REPO_DIR,\n",
|
| 122 |
+
" stdout=subprocess.DEVNULL,\n",
|
| 123 |
+
" stderr=subprocess.DEVNULL,\n",
|
| 124 |
+
")\n",
|
| 125 |
+
"\n",
|
| 126 |
+
"ENV_URL = 'http://localhost:8000'\n",
|
| 127 |
+
"\n",
|
| 128 |
+
"for i in range(30):\n",
|
| 129 |
+
" try:\n",
|
| 130 |
+
" if requests.get(f'{ENV_URL}/', timeout=2).status_code == 200:\n",
|
| 131 |
+
" print(f'PM-Ops server ready at {ENV_URL} (pid={server_proc.pid})')\n",
|
| 132 |
+
" break\n",
|
| 133 |
+
" except Exception:\n",
|
| 134 |
+
" pass\n",
|
| 135 |
+
" time.sleep(1)\n",
|
| 136 |
+
"else:\n",
|
| 137 |
+
" raise RuntimeError('Server did not start in 30 s — check uvicorn install.')"
|
| 138 |
+
]
|
| 139 |
+
},
|
| 140 |
+
{
|
| 141 |
+
"cell_type": "markdown",
|
| 142 |
+
"metadata": {},
|
| 143 |
+
"source": [
|
| 144 |
+
"## 4. Verify Environment"
|
| 145 |
+
]
|
| 146 |
+
},
|
| 147 |
+
{
|
| 148 |
+
"cell_type": "code",
|
| 149 |
+
"execution_count": null,
|
| 150 |
+
"metadata": {},
|
| 151 |
+
"outputs": [],
|
| 152 |
+
"source": [
|
| 153 |
+
"from openenv.core import GenericEnvClient\n",
|
| 154 |
+
"\n",
|
| 155 |
+
"with GenericEnvClient(base_url=ENV_URL).sync() as _check:\n",
|
| 156 |
+
" res = _check.reset()\n",
|
| 157 |
+
" obs = res.observation if hasattr(res, 'observation') else res\n",
|
| 158 |
+
" brief = getattr(obs, 'task_brief', '') or obs.get('task_brief', '')\n",
|
| 159 |
+
" print(f'Task brief: {brief[:120]}...')\n",
|
| 160 |
+
" res2 = _check.step({'action_type': 'meta.read_runbook', 'args': {}})\n",
|
| 161 |
+
" print('Step OK — environment is working.')"
|
| 162 |
+
]
|
| 163 |
+
},
|
| 164 |
+
{
|
| 165 |
+
"cell_type": "markdown",
|
| 166 |
+
"metadata": {},
|
| 167 |
+
"source": [
|
| 168 |
+
"## 5. Load Model with Unsloth\n",
|
| 169 |
+
"\n",
|
| 170 |
+
"Unsloth applies:\n",
|
| 171 |
+
"- **4-bit quantisation** — halves GPU memory vs bf16\n",
|
| 172 |
+
"- **LoRA adapters** — only trains ~1% of parameters, much faster\n",
|
| 173 |
+
"- **Unsloth kernel optimisations** — 2× faster rollout generation\n",
|
| 174 |
+
"- **`use_gradient_checkpointing=\"unsloth\"`** — 30% less activation memory\n",
|
| 175 |
+
"\n",
|
| 176 |
+
"> **Save warning (hackathon guide point 16):** never merge LoRA into a 4-bit model directly.\n",
|
| 177 |
+
"> Use `model.save_pretrained_merged(..., save_method=\"merged_16bit\")` in cell 14."
|
| 178 |
+
]
|
| 179 |
+
},
|
| 180 |
+
{
|
| 181 |
+
"cell_type": "code",
|
| 182 |
+
"execution_count": null,
|
| 183 |
+
"metadata": {},
|
| 184 |
+
"outputs": [],
|
| 185 |
+
"source": [
|
| 186 |
+
"from unsloth import FastLanguageModel, PatchFastRL\n",
|
| 187 |
+
"from trl import GRPOTrainer # import before patching so patch takes effect\n",
|
| 188 |
+
"\n",
|
| 189 |
+
"PatchFastRL('GRPO', FastLanguageModel)\n",
|
| 190 |
+
"\n",
|
| 191 |
+
"MODEL_NAME = 'Qwen/Qwen3-1.7B'\n",
|
| 192 |
+
"MAX_SEQ_LEN = 4096 + 512 # max_prompt_length + max_completion_length\n",
|
| 193 |
+
"LORA_RANK = 16\n",
|
| 194 |
+
"\n",
|
| 195 |
+
"model, tokenizer = FastLanguageModel.from_pretrained(\n",
|
| 196 |
+
" model_name=MODEL_NAME,\n",
|
| 197 |
+
" max_seq_length=MAX_SEQ_LEN,\n",
|
| 198 |
+
" load_in_4bit=True,\n",
|
| 199 |
+
" fast_inference=True, # enables Unsloth's vLLM-compatible fast path\n",
|
| 200 |
+
" max_lora_rank=LORA_RANK,\n",
|
| 201 |
+
" gpu_memory_utilization=0.6, # ~24 GB on A100 40GB for KV cache\n",
|
| 202 |
+
")\n",
|
| 203 |
+
"\n",
|
| 204 |
+
"model = FastLanguageModel.get_peft_model(\n",
|
| 205 |
+
" model,\n",
|
| 206 |
+
" r=LORA_RANK,\n",
|
| 207 |
+
" target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj',\n",
|
| 208 |
+
" 'gate_proj', 'up_proj', 'down_proj'],\n",
|
| 209 |
+
" lora_alpha=LORA_RANK,\n",
|
| 210 |
+
" use_gradient_checkpointing='unsloth',\n",
|
| 211 |
+
" random_state=42,\n",
|
| 212 |
+
")\n",
|
| 213 |
+
"tokenizer.pad_token = tokenizer.eos_token\n",
|
| 214 |
+
"print(f'Model ready: {MODEL_NAME} (4-bit + LoRA r={LORA_RANK})')"
|
| 215 |
+
]
|
| 216 |
+
},
|
| 217 |
+
{
|
| 218 |
+
"cell_type": "markdown",
|
| 219 |
+
"metadata": {},
|
| 220 |
+
"source": [
|
| 221 |
+
"## 6. Generate Training Dataset\n",
|
| 222 |
+
"\n",
|
| 223 |
+
"150 fixed-seed triage episodes. Seed is embedded in the prompt string so `env.reset(seed=...)` \n",
|
| 224 |
+
"reproduces the exact same org config — training is fully reproducible."
|
| 225 |
+
]
|
| 226 |
+
},
|
| 227 |
+
{
|
| 228 |
+
"cell_type": "code",
|
| 229 |
+
"execution_count": null,
|
| 230 |
+
"metadata": {},
|
| 231 |
+
"outputs": [],
|
| 232 |
+
"source": [
|
| 233 |
+
"from datasets import Dataset\n",
|
| 234 |
+
"from training.dataset import generate_triage_dataset\n",
|
| 235 |
+
"\n",
|
| 236 |
+
"N_EPISODES = 150\n",
|
| 237 |
+
"\n",
|
| 238 |
+
"rows = generate_triage_dataset(n_episodes=N_EPISODES, base_seed=42)\n",
|
| 239 |
+
"dataset = Dataset.from_list([{'prompt': r['prompt']} for r in rows])\n",
|
| 240 |
+
"print(f'Dataset: {len(dataset)} triage episodes')\n",
|
| 241 |
+
"print(f'Difficulties: {set(r[\"difficulty\"] for r in rows)}')\n",
|
| 242 |
+
"print(f'Sample: {dataset[0][\"prompt\"][:120]}...')"
|
| 243 |
+
]
|
| 244 |
+
},
|
| 245 |
+
{
|
| 246 |
+
"cell_type": "markdown",
|
| 247 |
+
"metadata": {},
|
| 248 |
+
"source": [
|
| 249 |
+
"## 7. Create Persistent Environment Client"
|
| 250 |
+
]
|
| 251 |
+
},
|
| 252 |
+
{
|
| 253 |
+
"cell_type": "code",
|
| 254 |
+
"execution_count": null,
|
| 255 |
+
"metadata": {},
|
| 256 |
+
"outputs": [],
|
| 257 |
+
"source": [
|
| 258 |
+
"from openenv.core import GenericEnvClient\n",
|
| 259 |
+
"\n",
|
| 260 |
+
"sync_env = GenericEnvClient(base_url=ENV_URL).sync()\n",
|
| 261 |
+
"sync_env.connect()\n",
|
| 262 |
+
"print('Persistent training connection established.')"
|
| 263 |
+
]
|
| 264 |
+
},
|
| 265 |
+
{
|
| 266 |
+
"cell_type": "markdown",
|
| 267 |
+
"metadata": {},
|
| 268 |
+
"source": [
|
| 269 |
+
"## 8. Build Rollout Function\n",
|
| 270 |
+
"\n",
|
| 271 |
+
"- `max_steps=15` — triage can be solved in 5 steps; 15 gives exploration room without burning compute \n",
|
| 272 |
+
"- Runbook-pinned truncation — org config stays in context even when older turns scroll out \n",
|
| 273 |
+
"- Step-aware fallback — `read_runbook` (early) → `noop` (mid) → `finish` (late) \n",
|
| 274 |
+
"- Channel tracking — feeds `reward_no_wrong_channels` to catch spray behaviour"
|
| 275 |
+
]
|
| 276 |
+
},
|
| 277 |
+
{
|
| 278 |
+
"cell_type": "code",
|
| 279 |
+
"execution_count": null,
|
| 280 |
+
"metadata": {},
|
| 281 |
+
"outputs": [],
|
| 282 |
+
"source": [
|
| 283 |
+
"from training.rollout import make_rollout_func\n",
|
| 284 |
+
"\n",
|
| 285 |
+
"TRAIN_MAX_STEPS = 15 # triage; increase for other tasks\n",
|
| 286 |
+
"\n",
|
| 287 |
+
"rollout_func = make_rollout_func(\n",
|
| 288 |
+
" sync_env=sync_env,\n",
|
| 289 |
+
" tokenizer=tokenizer,\n",
|
| 290 |
+
" max_steps=TRAIN_MAX_STEPS,\n",
|
| 291 |
+
")\n",
|
| 292 |
+
"print(f'Rollout ready (max_steps={TRAIN_MAX_STEPS})')"
|
| 293 |
+
]
|
| 294 |
+
},
|
| 295 |
+
{
|
| 296 |
+
"cell_type": "markdown",
|
| 297 |
+
"metadata": {},
|
| 298 |
+
"source": [
|
| 299 |
+
"## 9. Define Reward Functions"
|
| 300 |
+
]
|
| 301 |
+
},
|
| 302 |
+
{
|
| 303 |
+
"cell_type": "code",
|
| 304 |
+
"execution_count": null,
|
| 305 |
+
"metadata": {},
|
| 306 |
+
"outputs": [],
|
| 307 |
+
"source": [
|
| 308 |
+
"from training.rewards import ALL_REWARD_FUNCS, WEIGHT_FINAL_SCORE, WEIGHT_NO_WRONG_CHANNELS\n",
|
| 309 |
+
"print(f'Reward functions ({len(ALL_REWARD_FUNCS)}):')\n",
|
| 310 |
+
"for f in ALL_REWARD_FUNCS:\n",
|
| 311 |
+
" print(f' {f.__name__}')"
|
| 312 |
+
]
|
| 313 |
+
},
|
| 314 |
+
{
|
| 315 |
+
"cell_type": "markdown",
|
| 316 |
+
"metadata": {},
|
| 317 |
+
"source": [
|
| 318 |
+
"## 10. Configure GRPO Training"
|
| 319 |
+
]
|
| 320 |
+
},
|
| 321 |
+
{
|
| 322 |
+
"cell_type": "code",
|
| 323 |
+
"execution_count": null,
|
| 324 |
+
"metadata": {},
|
| 325 |
+
"outputs": [],
|
| 326 |
+
"source": [
|
| 327 |
+
"from trl import GRPOConfig\n",
|
| 328 |
+
"\n",
|
| 329 |
+
"OUTPUT_DIR = 'pm-ops-grpo-Qwen3-1.7B-triage'\n",
|
| 330 |
+
"HF_REPO_ID = f'your-hf-username/{OUTPUT_DIR}' # ← update before running\n",
|
| 331 |
+
"\n",
|
| 332 |
+
"grpo_config = GRPOConfig(\n",
|
| 333 |
+
" # Training\n",
|
| 334 |
+
" num_train_epochs=1,\n",
|
| 335 |
+
" learning_rate=5e-6,\n",
|
| 336 |
+
" gradient_accumulation_steps=64,\n",
|
| 337 |
+
" per_device_train_batch_size=1,\n",
|
| 338 |
+
" warmup_steps=10,\n",
|
| 339 |
+
" num_generations=2,\n",
|
| 340 |
+
"\n",
|
| 341 |
+
" # Sequence lengths — longer than Wordle due to JSON + reasoning\n",
|
| 342 |
+
" max_completion_length=512,\n",
|
| 343 |
+
" max_prompt_length=4096,\n",
|
| 344 |
+
"\n",
|
| 345 |
+
" # vLLM — Unsloth handles colocate mode via fast_inference=True above\n",
|
| 346 |
+
" use_vllm=True,\n",
|
| 347 |
+
"\n",
|
| 348 |
+
" # Output + logging\n",
|
| 349 |
+
" output_dir=OUTPUT_DIR,\n",
|
| 350 |
+
" report_to='trackio',\n",
|
| 351 |
+
" trackio_space_id=OUTPUT_DIR,\n",
|
| 352 |
+
" logging_steps=1,\n",
|
| 353 |
+
" save_steps=25,\n",
|
| 354 |
+
" gradient_checkpointing=False, # Unsloth handles this via get_peft_model\n",
|
| 355 |
+
")\n",
|
| 356 |
+
"\n",
|
| 357 |
+
"print(f'Output: {OUTPUT_DIR}')\n",
|
| 358 |
+
"print(f'Effective batch size: {grpo_config.per_device_train_batch_size * grpo_config.gradient_accumulation_steps}')"
|
| 359 |
+
]
|
| 360 |
+
},
|
| 361 |
+
{
|
| 362 |
+
"cell_type": "markdown",
|
| 363 |
+
"metadata": {},
|
| 364 |
+
"source": [
|
| 365 |
+
"## 11. Create Trainer"
|
| 366 |
+
]
|
| 367 |
+
},
|
| 368 |
+
{
|
| 369 |
+
"cell_type": "code",
|
| 370 |
+
"execution_count": null,
|
| 371 |
+
"metadata": {},
|
| 372 |
+
"outputs": [],
|
| 373 |
+
"source": [
|
| 374 |
+
"from trl import GRPOTrainer\n",
|
| 375 |
+
"\n",
|
| 376 |
+
"trainer = GRPOTrainer(\n",
|
| 377 |
+
" model=model, # Unsloth model object (not string)\n",
|
| 378 |
+
" processing_class=tokenizer,\n",
|
| 379 |
+
" reward_funcs=ALL_REWARD_FUNCS,\n",
|
| 380 |
+
" train_dataset=dataset,\n",
|
| 381 |
+
" args=grpo_config,\n",
|
| 382 |
+
" rollout_func=rollout_func,\n",
|
| 383 |
+
")\n",
|
| 384 |
+
"print('GRPOTrainer ready.')"
|
| 385 |
+
]
|
| 386 |
+
},
|
| 387 |
+
{
|
| 388 |
+
"cell_type": "markdown",
|
| 389 |
+
"metadata": {},
|
| 390 |
+
"source": [
|
| 391 |
+
"## 12. GPU Check"
|
| 392 |
+
]
|
| 393 |
+
},
|
| 394 |
+
{
|
| 395 |
+
"cell_type": "code",
|
| 396 |
+
"execution_count": null,
|
| 397 |
+
"metadata": {},
|
| 398 |
+
"outputs": [],
|
| 399 |
+
"source": [
|
| 400 |
+
"import torch\n",
|
| 401 |
+
"gpu = torch.cuda.get_device_properties(0)\n",
|
| 402 |
+
"reserved_before = round(torch.cuda.max_memory_reserved() / 1024**3, 2)\n",
|
| 403 |
+
"total_gb = round(gpu.total_memory / 1024**3, 2)\n",
|
| 404 |
+
"print(f'GPU: {gpu.name}')\n",
|
| 405 |
+
"print(f'Memory: {total_gb} GB total, {reserved_before} GB reserved')\n",
|
| 406 |
+
"assert total_gb >= 38, f'Need A100 40GB, got {total_gb:.1f} GB — switch runtime.'"
|
| 407 |
+
]
|
| 408 |
+
},
|
| 409 |
+
{
|
| 410 |
+
"cell_type": "markdown",
|
| 411 |
+
"metadata": {},
|
| 412 |
+
"source": [
|
| 413 |
+
"## 13. Train (~90 min on A100)\n",
|
| 414 |
+
"\n",
|
| 415 |
+
"Watch **trackio** for per-reward-signal curves. Key signals to monitor:\n",
|
| 416 |
+
"- `reward_final_score` — should trend upward over training\n",
|
| 417 |
+
"- `reward_valid_json` — should stay > 0.7 (model reliably outputs JSON)\n",
|
| 418 |
+
"- `reward_read_runbook` — should hit 1.0 quickly and stay there\n",
|
| 419 |
+
"- `reward_no_wrong_channels` — should increase (less channel spray)\n",
|
| 420 |
+
"\n",
|
| 421 |
+
"If `reward_valid_json` < 0.3 for first 20 steps → model struggling with format. Stop and reduce `max_completion_length` to 256."
|
| 422 |
+
]
|
| 423 |
+
},
|
| 424 |
+
{
|
| 425 |
+
"cell_type": "code",
|
| 426 |
+
"execution_count": null,
|
| 427 |
+
"metadata": {},
|
| 428 |
+
"outputs": [],
|
| 429 |
+
"source": [
|
| 430 |
+
"trainer_stats = trainer.train()"
|
| 431 |
+
]
|
| 432 |
+
},
|
| 433 |
+
{
|
| 434 |
+
"cell_type": "code",
|
| 435 |
+
"execution_count": null,
|
| 436 |
+
"metadata": {},
|
| 437 |
+
"outputs": [],
|
| 438 |
+
"source": [
|
| 439 |
+
"used_gb = round(torch.cuda.max_memory_reserved() / 1024**3, 2)\n",
|
| 440 |
+
"train_mins = round(trainer_stats.metrics.get('train_runtime', 0) / 60, 1)\n",
|
| 441 |
+
"print(f'Training time : {train_mins} min')\n",
|
| 442 |
+
"print(f'Peak GPU usage: {used_gb} GB / {total_gb} GB ({round(used_gb/total_gb*100, 1)}%)')"
|
| 443 |
+
]
|
| 444 |
+
},
|
| 445 |
+
{
|
| 446 |
+
"cell_type": "markdown",
|
| 447 |
+
"metadata": {},
|
| 448 |
+
"source": [
|
| 449 |
+
"## 14. Save Model\n",
|
| 450 |
+
"\n",
|
| 451 |
+
"**Important:** Unsloth uses LoRA + 4-bit quantisation. Do NOT use `trainer.save_model()` directly —\n",
|
| 452 |
+
"merging LoRA into a 4-bit base corrupts weights (hackathon guide point 16).\n",
|
| 453 |
+
"Use `save_pretrained_merged` which dequantises first, then merges cleanly into bf16."
|
| 454 |
+
]
|
| 455 |
+
},
|
| 456 |
+
{
|
| 457 |
+
"cell_type": "code",
|
| 458 |
+
"execution_count": null,
|
| 459 |
+
"metadata": {},
|
| 460 |
+
"outputs": [],
|
| 461 |
+
"source": [
|
| 462 |
+
"sync_env.close() # close training connection before saving\n",
|
| 463 |
+
"\n",
|
| 464 |
+
"# Merge LoRA into bf16 weights and save — safe Unsloth path\n",
|
| 465 |
+
"model.save_pretrained_merged(OUTPUT_DIR, tokenizer, save_method='merged_16bit')\n",
|
| 466 |
+
"print(f'Model saved → {OUTPUT_DIR}')\n",
|
| 467 |
+
"\n",
|
| 468 |
+
"# Push merged model to HF Hub\n",
|
| 469 |
+
"model.push_to_hub_merged(HF_REPO_ID, tokenizer, save_method='merged_16bit')\n",
|
| 470 |
+
"print(f'Pushed → {HF_REPO_ID}')"
|
| 471 |
+
]
|
| 472 |
+
},
|
| 473 |
+
{
|
| 474 |
+
"cell_type": "markdown",
|
| 475 |
+
"metadata": {},
|
| 476 |
+
"source": [
|
| 477 |
+
"## 15. Evaluate: Baseline vs Trained\n",
|
| 478 |
+
"\n",
|
| 479 |
+
"10 fresh triage episodes on the remote HF Space (not localhost) — clean eval setup."
|
| 480 |
+
]
|
| 481 |
+
},
|
| 482 |
+
{
|
| 483 |
+
"cell_type": "code",
|
| 484 |
+
"execution_count": null,
|
| 485 |
+
"metadata": {},
|
| 486 |
+
"outputs": [],
|
| 487 |
+
"source": [
|
| 488 |
+
"from transformers import AutoModelForCausalLM\n",
|
| 489 |
+
"from openenv.core import GenericEnvClient\n",
|
| 490 |
+
"from training.rollout import extract_json_action, step_aware_fallback, build_messages, _obs_to_dict\n",
|
| 491 |
+
"from training.prompts import format_observation\n",
|
| 492 |
+
"from inference import baseline_agent\n",
|
| 493 |
+
"\n",
|
| 494 |
+
"EVAL_URL = 'https://adityaguntur-pm-ops.hf.space'\n",
|
| 495 |
+
"N_EVAL = 10\n",
|
| 496 |
+
"EVAL_MAX_STEPS = 15\n",
|
| 497 |
+
"\n",
|
| 498 |
+
"eval_model = AutoModelForCausalLM.from_pretrained(\n",
|
| 499 |
+
" OUTPUT_DIR, torch_dtype='auto', device_map='auto'\n",
|
| 500 |
+
")\n",
|
| 501 |
+
"\n",
|
| 502 |
+
"\n",
|
| 503 |
+
"def _get(o, k, default=None):\n",
|
| 504 |
+
" return getattr(o, k, default) if not isinstance(o, dict) else o.get(k, default)\n",
|
| 505 |
+
"\n",
|
| 506 |
+
"\n",
|
| 507 |
+
"def eval_trained(sync_env, model, tokenizer, n=N_EVAL):\n",
|
| 508 |
+
" scores = []\n",
|
| 509 |
+
" for i in range(n):\n",
|
| 510 |
+
" result = sync_env.reset()\n",
|
| 511 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 512 |
+
" task_brief = obs_dict.get('task_brief', '')\n",
|
| 513 |
+
" history, step, score = [], 0, 0.0\n",
|
| 514 |
+
" done = False\n",
|
| 515 |
+
"\n",
|
| 516 |
+
" while not done and step < EVAL_MAX_STEPS:\n",
|
| 517 |
+
" from training.rollout import _current_obs_text\n",
|
| 518 |
+
" obs_text = _current_obs_text(obs_dict, step, task_brief)\n",
|
| 519 |
+
" msgs = build_messages(history, obs_text)\n",
|
| 520 |
+
" prompt_text = tokenizer.apply_chat_template(\n",
|
| 521 |
+
" msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n",
|
| 522 |
+
" )\n",
|
| 523 |
+
" inputs = tokenizer([prompt_text], return_tensors='pt').to(model.device)\n",
|
| 524 |
+
" out_ids = model.generate(**inputs, max_new_tokens=512)\n",
|
| 525 |
+
" completion = tokenizer.decode(out_ids[0][len(inputs.input_ids[0]):], skip_special_tokens=True)\n",
|
| 526 |
+
"\n",
|
| 527 |
+
" parsed = extract_json_action(completion) or step_aware_fallback(step)\n",
|
| 528 |
+
" result = sync_env.step({'action_type': parsed['action_type'], 'args': parsed.get('args', {})})\n",
|
| 529 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 530 |
+
" done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n",
|
| 531 |
+
" score = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n",
|
| 532 |
+
" history.append({'obs_text': obs_text, 'completion': completion, 'is_runbook': False})\n",
|
| 533 |
+
" step += 1\n",
|
| 534 |
+
"\n",
|
| 535 |
+
" scores.append(score)\n",
|
| 536 |
+
" print(f' Trained ep {i+1}/{n}: score={score:.3f}')\n",
|
| 537 |
+
" return scores\n",
|
| 538 |
+
"\n",
|
| 539 |
+
"\n",
|
| 540 |
+
"def eval_baseline(sync_env, n=N_EVAL):\n",
|
| 541 |
+
" scores = []\n",
|
| 542 |
+
" for i in range(n):\n",
|
| 543 |
+
" result = sync_env.reset()\n",
|
| 544 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 545 |
+
" org_config, done, step, score = {}, False, 0, 0.0\n",
|
| 546 |
+
" while not done and step < EVAL_MAX_STEPS:\n",
|
| 547 |
+
" at, args = baseline_agent(obs_dict, org_config)\n",
|
| 548 |
+
" result = sync_env.step({'action_type': at, 'args': args})\n",
|
| 549 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 550 |
+
" done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n",
|
| 551 |
+
" score = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n",
|
| 552 |
+
" step += 1\n",
|
| 553 |
+
" scores.append(score)\n",
|
| 554 |
+
" print(f' Baseline ep {i+1}/{n}: score={score:.3f}')\n",
|
| 555 |
+
" return scores\n",
|
| 556 |
+
"\n",
|
| 557 |
+
"\n",
|
| 558 |
+
"print('--- Baseline ---')\n",
|
| 559 |
+
"with GenericEnvClient(base_url=EVAL_URL).sync() as env_eval:\n",
|
| 560 |
+
" baseline_scores = eval_baseline(env_eval)\n",
|
| 561 |
+
"\n",
|
| 562 |
+
"print('--- Trained ---')\n",
|
| 563 |
+
"with GenericEnvClient(base_url=EVAL_URL).sync() as env_eval:\n",
|
| 564 |
+
" trained_scores = eval_trained(env_eval, eval_model, tokenizer)\n",
|
| 565 |
+
"\n",
|
| 566 |
+
"print(f'\\nBaseline avg: {sum(baseline_scores)/N_EVAL:.3f}')\n",
|
| 567 |
+
"print(f'Trained avg: {sum(trained_scores)/N_EVAL:.3f}')\n",
|
| 568 |
+
"print(f'Improvement : +{(sum(trained_scores)-sum(baseline_scores))/N_EVAL:.3f}')"
|
| 569 |
+
]
|
| 570 |
+
},
|
| 571 |
+
{
|
| 572 |
+
"cell_type": "markdown",
|
| 573 |
+
"metadata": {},
|
| 574 |
+
"source": [
|
| 575 |
+
"## 16. Plot Results"
|
| 576 |
+
]
|
| 577 |
+
},
|
| 578 |
+
{
|
| 579 |
+
"cell_type": "code",
|
| 580 |
+
"execution_count": null,
|
| 581 |
+
"metadata": {},
|
| 582 |
+
"outputs": [],
|
| 583 |
+
"source": [
|
| 584 |
+
"import matplotlib.pyplot as plt\n",
|
| 585 |
+
"import numpy as np\n",
|
| 586 |
+
"\n",
|
| 587 |
+
"fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n",
|
| 588 |
+
"\n",
|
| 589 |
+
"ax = axes[0]\n",
|
| 590 |
+
"x, w = np.arange(N_EVAL), 0.35\n",
|
| 591 |
+
"ax.bar(x - w/2, baseline_scores, w, label='Baseline (heuristic)', color='steelblue', alpha=0.8)\n",
|
| 592 |
+
"ax.bar(x + w/2, trained_scores, w, label='Trained (GRPO+Unsloth)', color='coral', alpha=0.8)\n",
|
| 593 |
+
"ax.axhline(sum(baseline_scores)/N_EVAL, color='steelblue', linestyle='--', alpha=0.5, linewidth=1)\n",
|
| 594 |
+
"ax.axhline(sum(trained_scores)/N_EVAL, color='coral', linestyle='--', alpha=0.5, linewidth=1)\n",
|
| 595 |
+
"ax.set_xlabel('Eval episode')\n",
|
| 596 |
+
"ax.set_ylabel('Episode reward (0–1)')\n",
|
| 597 |
+
"ax.set_title('Per-episode: Baseline vs GRPO-trained')\n",
|
| 598 |
+
"ax.set_xticks(x)\n",
|
| 599 |
+
"ax.set_ylim(0, 1.05)\n",
|
| 600 |
+
"ax.legend()\n",
|
| 601 |
+
"\n",
|
| 602 |
+
"ax2 = axes[1]\n",
|
| 603 |
+
"avgs = [sum(baseline_scores)/N_EVAL, sum(trained_scores)/N_EVAL]\n",
|
| 604 |
+
"bars = ax2.bar(['Baseline', 'GRPO+Unsloth'], avgs, color=['steelblue', 'coral'], alpha=0.85, width=0.5)\n",
|
| 605 |
+
"for bar, val in zip(bars, avgs):\n",
|
| 606 |
+
" ax2.text(bar.get_x() + bar.get_width()/2, val + 0.01, f'{val:.3f}',\n",
|
| 607 |
+
" ha='center', fontsize=13, fontweight='bold')\n",
|
| 608 |
+
"ax2.set_ylabel('Average reward (0–1)')\n",
|
| 609 |
+
"ax2.set_title(f'Average over {N_EVAL} triage episodes')\n",
|
| 610 |
+
"ax2.set_ylim(0, 1.05)\n",
|
| 611 |
+
"\n",
|
| 612 |
+
"plt.tight_layout()\n",
|
| 613 |
+
"plt.savefig('eval_results.png', dpi=150, bbox_inches='tight')\n",
|
| 614 |
+
"plt.show()\n",
|
| 615 |
+
"print('Saved: eval_results.png ← embed this in README for hackathon submission')"
|
| 616 |
+
]
|
| 617 |
+
},
|
| 618 |
+
{
|
| 619 |
+
"cell_type": "markdown",
|
| 620 |
+
"metadata": {},
|
| 621 |
+
"source": [
|
| 622 |
+
"## 17. Teardown"
|
| 623 |
+
]
|
| 624 |
+
},
|
| 625 |
+
{
|
| 626 |
+
"cell_type": "code",
|
| 627 |
+
"execution_count": null,
|
| 628 |
+
"metadata": {},
|
| 629 |
+
"outputs": [],
|
| 630 |
+
"source": [
|
| 631 |
+
"server_proc.terminate()\n",
|
| 632 |
+
"print('Local PM-Ops server stopped.')"
|
| 633 |
+
]
|
| 634 |
+
},
|
| 635 |
+
{
|
| 636 |
+
"cell_type": "markdown",
|
| 637 |
+
"metadata": {},
|
| 638 |
+
"source": [
|
| 639 |
+
"---\n",
|
| 640 |
+
"## What to try next\n",
|
| 641 |
+
"\n",
|
| 642 |
+
"| Change | Where | Expected effect |\n",
|
| 643 |
+
"|---|---|---|\n",
|
| 644 |
+
"| All 4 task types | `generate_triage_dataset` → all-tasks generator | Tests generalization |\n",
|
| 645 |
+
"| 500 episodes | `N_EPISODES = 500` | Better coverage |\n",
|
| 646 |
+
"| Equal weights (0.20 each) | `WEIGHT_*` in `rewards.py` | Compare reward-shaping approaches |\n",
|
| 647 |
+
"| Longer training | `num_train_epochs=3` | More improvement |\n",
|
| 648 |
+
"| Larger model | `Qwen/Qwen3-4B` | Better reasoning, more memory |\n",
|
| 649 |
+
"| Higher LoRA rank | `r=64, lora_alpha=64` | More capacity |\n",
|
| 650 |
+
"| More steps for other tasks | `TRAIN_MAX_STEPS = 20` | Needed for release_notes/dep_update |"
|
| 651 |
+
]
|
| 652 |
+
}
|
| 653 |
+
],
|
| 654 |
+
"metadata": {
|
| 655 |
+
"kernelspec": {
|
| 656 |
+
"display_name": "Python 3",
|
| 657 |
+
"language": "python",
|
| 658 |
+
"name": "python3"
|
| 659 |
+
},
|
| 660 |
+
"language_info": {
|
| 661 |
+
"name": "python",
|
| 662 |
+
"version": "3.11.0"
|
| 663 |
+
},
|
| 664 |
+
"accelerator": "GPU",
|
| 665 |
+
"gpuClass": "premium"
|
| 666 |
+
},
|
| 667 |
+
"nbformat": 4,
|
| 668 |
+
"nbformat_minor": 4
|
| 669 |
+
}
|
training/train_v2.ipynb
ADDED
|
@@ -0,0 +1,715 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"cells": [
|
| 3 |
+
{
|
| 4 |
+
"cell_type": "markdown",
|
| 5 |
+
"metadata": {},
|
| 6 |
+
"source": [
|
| 7 |
+
"# PM-Ops GRPO Training v2\n",
|
| 8 |
+
"\n",
|
| 9 |
+
"### Fixes vs v1\n",
|
| 10 |
+
"| Bug | Root cause | Fix |\n",
|
| 11 |
+
"|-----|-----------|-----|\n",
|
| 12 |
+
"| All rewards = 0 | `rollout_func` kwargs broken in TRL 1.3.0.dev0 | TRL **1.2.0** stable + single `reward` key |\n",
|
| 13 |
+
"| Dep hell | Unsloth 2026.4.8 requires torch<2.11 / trl≤0.24, conflicts with Colab's base image | **Drop Unsloth**, use bitsandbytes + PEFT directly |\n",
|
| 14 |
+
"| Only 4 steps | T4 + grad_accum=64 | GPU auto-detect, adaptive config |\n",
|
| 15 |
+
"| Silent reward failure | No logging | Rollout prints every episode, reward_func warns on zero |\n",
|
| 16 |
+
"\n",
|
| 17 |
+
"**Stack:** TRL 1.2.0 · bitsandbytes (4-bit) · PEFT LoRA · OpenEnv \n",
|
| 18 |
+
"**GPU:** A100 → ~150 steps / 60 min · T4 → ~18 steps (smoke-test / short run)"
|
| 19 |
+
]
|
| 20 |
+
},
|
| 21 |
+
{
|
| 22 |
+
"cell_type": "markdown",
|
| 23 |
+
"metadata": {},
|
| 24 |
+
"source": [
|
| 25 |
+
"## 0. Install Dependencies"
|
| 26 |
+
]
|
| 27 |
+
},
|
| 28 |
+
{
|
| 29 |
+
"cell_type": "code",
|
| 30 |
+
"execution_count": null,
|
| 31 |
+
"metadata": {},
|
| 32 |
+
"outputs": [],
|
| 33 |
+
"source": [
|
| 34 |
+
"# Keep Colab's native torch (2.11.0) — don't touch it.\n",
|
| 35 |
+
"# Install everything that is compatible with torch 2.11.0.\n",
|
| 36 |
+
"\n",
|
| 37 |
+
"!pip install -q \"trl==1.2.0\"\n",
|
| 38 |
+
"!pip install -q \"bitsandbytes>=0.44.0\" # 4-bit quantisation\n",
|
| 39 |
+
"!pip install -q \"peft>=0.13.0\" # LoRA\n",
|
| 40 |
+
"!pip install -q \"accelerate>=1.0.0\"\n",
|
| 41 |
+
"!pip install -q \"datasets>=4.7.0\"\n",
|
| 42 |
+
"!pip install -q \"transformers>=5.0.0\"\n",
|
| 43 |
+
"!pip install -q \"openenv-core>=0.2.2\" fastapi \"uvicorn[standard]\" \"pydantic>=2.0.0\"\n",
|
| 44 |
+
"!pip install -q trackio\n",
|
| 45 |
+
"\n",
|
| 46 |
+
"print('Done — restart kernel, then run all cells from top.')"
|
| 47 |
+
]
|
| 48 |
+
},
|
| 49 |
+
{
|
| 50 |
+
"cell_type": "markdown",
|
| 51 |
+
"metadata": {},
|
| 52 |
+
"source": [
|
| 53 |
+
"## 1. Version Check + GPU Detect"
|
| 54 |
+
]
|
| 55 |
+
},
|
| 56 |
+
{
|
| 57 |
+
"cell_type": "code",
|
| 58 |
+
"execution_count": null,
|
| 59 |
+
"metadata": {},
|
| 60 |
+
"outputs": [],
|
| 61 |
+
"source": [
|
| 62 |
+
"import torch, trl\n",
|
| 63 |
+
"import trl.experimental.openenv # must exist for rollout_func\n",
|
| 64 |
+
"\n",
|
| 65 |
+
"print(f'torch : {torch.__version__}')\n",
|
| 66 |
+
"print(f'TRL : {trl.__version__} (need 1.2.x)')\n",
|
| 67 |
+
"assert trl.__version__.startswith('1.2'), \\\n",
|
| 68 |
+
" f'Wrong TRL: {trl.__version__}. Re-run install cell and restart kernel.'\n",
|
| 69 |
+
"print('trl.experimental.openenv : OK')\n",
|
| 70 |
+
"\n",
|
| 71 |
+
"gpu = torch.cuda.get_device_properties(0)\n",
|
| 72 |
+
"TOTAL_GB = round(gpu.total_memory / 1024**3, 1)\n",
|
| 73 |
+
"IS_A100 = TOTAL_GB >= 35\n",
|
| 74 |
+
"print(f'\\nGPU : {gpu.name} ({TOTAL_GB} GB)')\n",
|
| 75 |
+
"print(f'Mode : {\"A100 — full training\" if IS_A100 else \"T4 — reduced config (smoke-test quality)\"}')\n",
|
| 76 |
+
"\n",
|
| 77 |
+
"# Adaptive config\n",
|
| 78 |
+
"GRAD_ACCUM = 32 if IS_A100 else 8\n",
|
| 79 |
+
"NUM_GEN = 4 if IS_A100 else 2\n",
|
| 80 |
+
"MAX_COMP_LEN = 512 if IS_A100 else 256\n",
|
| 81 |
+
"\n",
|
| 82 |
+
"print(f'\\ngrad_accum={GRAD_ACCUM} num_gen={NUM_GEN} max_comp={MAX_COMP_LEN}')"
|
| 83 |
+
]
|
| 84 |
+
},
|
| 85 |
+
{
|
| 86 |
+
"cell_type": "markdown",
|
| 87 |
+
"metadata": {},
|
| 88 |
+
"source": [
|
| 89 |
+
"## 2. Clone PM-Ops Repo"
|
| 90 |
+
]
|
| 91 |
+
},
|
| 92 |
+
{
|
| 93 |
+
"cell_type": "code",
|
| 94 |
+
"execution_count": null,
|
| 95 |
+
"metadata": {},
|
| 96 |
+
"outputs": [],
|
| 97 |
+
"source": [
|
| 98 |
+
"import os, sys\n",
|
| 99 |
+
"\n",
|
| 100 |
+
"REPO_URL = 'https://huggingface.co/spaces/TheCrustaceans/Pm-ops'\n",
|
| 101 |
+
"REPO_DIR = '/content/Pm_ops'\n",
|
| 102 |
+
"\n",
|
| 103 |
+
"if not os.path.exists(REPO_DIR):\n",
|
| 104 |
+
" !git clone --depth=1 -q {REPO_URL} {REPO_DIR}\n",
|
| 105 |
+
" print(f'Cloned → {REPO_DIR}')\n",
|
| 106 |
+
"else:\n",
|
| 107 |
+
" !git -C {REPO_DIR} pull -q origin main\n",
|
| 108 |
+
" print(f'Pulled latest → {REPO_DIR}')\n",
|
| 109 |
+
"\n",
|
| 110 |
+
"for p in [REPO_DIR, os.path.join(REPO_DIR, 'training')]:\n",
|
| 111 |
+
" if p not in sys.path:\n",
|
| 112 |
+
" sys.path.insert(0, p)\n",
|
| 113 |
+
"os.chdir(REPO_DIR)\n",
|
| 114 |
+
"print(f'CWD: {os.getcwd()}')"
|
| 115 |
+
]
|
| 116 |
+
},
|
| 117 |
+
{
|
| 118 |
+
"cell_type": "markdown",
|
| 119 |
+
"metadata": {},
|
| 120 |
+
"source": [
|
| 121 |
+
"## 3. HuggingFace Login"
|
| 122 |
+
]
|
| 123 |
+
},
|
| 124 |
+
{
|
| 125 |
+
"cell_type": "code",
|
| 126 |
+
"execution_count": null,
|
| 127 |
+
"metadata": {},
|
| 128 |
+
"outputs": [],
|
| 129 |
+
"source": [
|
| 130 |
+
"from huggingface_hub import notebook_login\n",
|
| 131 |
+
"notebook_login()"
|
| 132 |
+
]
|
| 133 |
+
},
|
| 134 |
+
{
|
| 135 |
+
"cell_type": "markdown",
|
| 136 |
+
"metadata": {},
|
| 137 |
+
"source": [
|
| 138 |
+
"## 4. Start Local PM-Ops Server\n",
|
| 139 |
+
"\n",
|
| 140 |
+
"Localhost removes ~200 ms/step network round-trip vs calling the HF Space."
|
| 141 |
+
]
|
| 142 |
+
},
|
| 143 |
+
{
|
| 144 |
+
"cell_type": "code",
|
| 145 |
+
"execution_count": null,
|
| 146 |
+
"metadata": {},
|
| 147 |
+
"outputs": [],
|
| 148 |
+
"source": [
|
| 149 |
+
"import subprocess, time, requests\n",
|
| 150 |
+
"\n",
|
| 151 |
+
"server_proc = subprocess.Popen(\n",
|
| 152 |
+
" [sys.executable, '-m', 'uvicorn', 'server.app:app', '--host', '0.0.0.0', '--port', '8000'],\n",
|
| 153 |
+
" cwd=REPO_DIR, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,\n",
|
| 154 |
+
")\n",
|
| 155 |
+
"ENV_URL = 'http://localhost:8000'\n",
|
| 156 |
+
"\n",
|
| 157 |
+
"for _ in range(30):\n",
|
| 158 |
+
" try:\n",
|
| 159 |
+
" if requests.get(f'{ENV_URL}/', timeout=2).status_code == 200:\n",
|
| 160 |
+
" print(f'PM-Ops server ready pid={server_proc.pid}')\n",
|
| 161 |
+
" break\n",
|
| 162 |
+
" except Exception:\n",
|
| 163 |
+
" pass\n",
|
| 164 |
+
" time.sleep(1)\n",
|
| 165 |
+
"else:\n",
|
| 166 |
+
" raise RuntimeError('Server did not start in 30 s')"
|
| 167 |
+
]
|
| 168 |
+
},
|
| 169 |
+
{
|
| 170 |
+
"cell_type": "markdown",
|
| 171 |
+
"metadata": {},
|
| 172 |
+
"source": [
|
| 173 |
+
"## 5. Verify Environment"
|
| 174 |
+
]
|
| 175 |
+
},
|
| 176 |
+
{
|
| 177 |
+
"cell_type": "code",
|
| 178 |
+
"execution_count": null,
|
| 179 |
+
"metadata": {},
|
| 180 |
+
"outputs": [],
|
| 181 |
+
"source": [
|
| 182 |
+
"from openenv.core import GenericEnvClient\n",
|
| 183 |
+
"from training.rollout import _obs_to_dict\n",
|
| 184 |
+
"\n",
|
| 185 |
+
"with GenericEnvClient(base_url=ENV_URL).sync() as _env:\n",
|
| 186 |
+
" r = _env.reset()\n",
|
| 187 |
+
" obs = _obs_to_dict(r.observation if hasattr(r, 'observation') else r)\n",
|
| 188 |
+
" print(f'task_brief : {obs.get(\"task_brief\", \"?\")[:80]}...')\n",
|
| 189 |
+
" _env.step({'action_type': 'meta.read_runbook', 'args': {}})\n",
|
| 190 |
+
" print('Env step : OK')"
|
| 191 |
+
]
|
| 192 |
+
},
|
| 193 |
+
{
|
| 194 |
+
"cell_type": "markdown",
|
| 195 |
+
"metadata": {},
|
| 196 |
+
"source": [
|
| 197 |
+
"## 6. Load Model — bitsandbytes 4-bit + PEFT LoRA\n",
|
| 198 |
+
"\n",
|
| 199 |
+
"No Unsloth — uses stock HuggingFace `AutoModelForCausalLM` with `BitsAndBytesConfig`. \n",
|
| 200 |
+
"Same 4-bit + LoRA memory profile, ~2× slower generation than Unsloth's kernels."
|
| 201 |
+
]
|
| 202 |
+
},
|
| 203 |
+
{
|
| 204 |
+
"cell_type": "code",
|
| 205 |
+
"execution_count": null,
|
| 206 |
+
"metadata": {},
|
| 207 |
+
"outputs": [],
|
| 208 |
+
"source": [
|
| 209 |
+
"from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\n",
|
| 210 |
+
"from peft import LoraConfig, get_peft_model, TaskType\n",
|
| 211 |
+
"\n",
|
| 212 |
+
"MODEL_NAME = 'Qwen/Qwen3-1.7B'\n",
|
| 213 |
+
"LORA_RANK = 16\n",
|
| 214 |
+
"\n",
|
| 215 |
+
"bnb_cfg = BitsAndBytesConfig(\n",
|
| 216 |
+
" load_in_4bit=True,\n",
|
| 217 |
+
" bnb_4bit_compute_dtype=torch.bfloat16,\n",
|
| 218 |
+
" bnb_4bit_use_double_quant=True,\n",
|
| 219 |
+
" bnb_4bit_quant_type='nf4',\n",
|
| 220 |
+
")\n",
|
| 221 |
+
"\n",
|
| 222 |
+
"model = AutoModelForCausalLM.from_pretrained(\n",
|
| 223 |
+
" MODEL_NAME,\n",
|
| 224 |
+
" quantization_config=bnb_cfg,\n",
|
| 225 |
+
" device_map='auto',\n",
|
| 226 |
+
" torch_dtype=torch.bfloat16,\n",
|
| 227 |
+
" attn_implementation='eager', # safe default; use 'flash_attention_2' if flash-attn installed\n",
|
| 228 |
+
")\n",
|
| 229 |
+
"\n",
|
| 230 |
+
"tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
|
| 231 |
+
"tokenizer.pad_token = tokenizer.eos_token\n",
|
| 232 |
+
"tokenizer.padding_side = 'left'\n",
|
| 233 |
+
"\n",
|
| 234 |
+
"lora_cfg = LoraConfig(\n",
|
| 235 |
+
" r=LORA_RANK,\n",
|
| 236 |
+
" lora_alpha=LORA_RANK,\n",
|
| 237 |
+
" target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj'],\n",
|
| 238 |
+
" lora_dropout=0.0,\n",
|
| 239 |
+
" bias='none',\n",
|
| 240 |
+
" task_type=TaskType.CAUSAL_LM,\n",
|
| 241 |
+
")\n",
|
| 242 |
+
"model = get_peft_model(model, lora_cfg)\n",
|
| 243 |
+
"model.print_trainable_parameters()\n",
|
| 244 |
+
"\n",
|
| 245 |
+
"reserved = round(torch.cuda.max_memory_reserved() / 1024**3, 2)\n",
|
| 246 |
+
"print(f'GPU memory after load: {reserved} GB / {TOTAL_GB} GB')"
|
| 247 |
+
]
|
| 248 |
+
},
|
| 249 |
+
{
|
| 250 |
+
"cell_type": "markdown",
|
| 251 |
+
"metadata": {},
|
| 252 |
+
"source": [
|
| 253 |
+
"## 7. Generate Training Dataset"
|
| 254 |
+
]
|
| 255 |
+
},
|
| 256 |
+
{
|
| 257 |
+
"cell_type": "code",
|
| 258 |
+
"execution_count": null,
|
| 259 |
+
"metadata": {},
|
| 260 |
+
"outputs": [],
|
| 261 |
+
"source": [
|
| 262 |
+
"from datasets import Dataset\n",
|
| 263 |
+
"from training.dataset import generate_triage_dataset\n",
|
| 264 |
+
"\n",
|
| 265 |
+
"N_EPISODES = 150\n",
|
| 266 |
+
"rows = generate_triage_dataset(n_episodes=N_EPISODES, base_seed=42)\n",
|
| 267 |
+
"dataset = Dataset.from_list([{'prompt': r['prompt']} for r in rows])\n",
|
| 268 |
+
"print(f'Dataset : {len(dataset)} triage episodes')\n",
|
| 269 |
+
"print(f'Sample : {dataset[0][\"prompt\"][:100]}...')"
|
| 270 |
+
]
|
| 271 |
+
},
|
| 272 |
+
{
|
| 273 |
+
"cell_type": "markdown",
|
| 274 |
+
"metadata": {},
|
| 275 |
+
"source": [
|
| 276 |
+
"## 8. Create Persistent Env Client + Rollout Function\n",
|
| 277 |
+
"\n",
|
| 278 |
+
"**Key change from v1:** single `\"reward\"` key in the rollout output. \n",
|
| 279 |
+
"This removes all kwargs-plumbing complexity — the reward_func is a trivial passthrough."
|
| 280 |
+
]
|
| 281 |
+
},
|
| 282 |
+
{
|
| 283 |
+
"cell_type": "code",
|
| 284 |
+
"execution_count": null,
|
| 285 |
+
"metadata": {},
|
| 286 |
+
"outputs": [],
|
| 287 |
+
"source": [
|
| 288 |
+
"from openenv.core import GenericEnvClient\n",
|
| 289 |
+
"from training.rollout import rollout_once\n",
|
| 290 |
+
"\n",
|
| 291 |
+
"sync_env = GenericEnvClient(base_url=ENV_URL).sync()\n",
|
| 292 |
+
"sync_env.connect()\n",
|
| 293 |
+
"print('Training env connected')\n",
|
| 294 |
+
"\n",
|
| 295 |
+
"TRAIN_MAX_STEPS = 15\n",
|
| 296 |
+
"\n",
|
| 297 |
+
"\n",
|
| 298 |
+
"def make_rollout_func(env, tok, max_steps=TRAIN_MAX_STEPS):\n",
|
| 299 |
+
" def rollout_func(prompts, trainer=None):\n",
|
| 300 |
+
" out = {'prompt_ids': [], 'completion_ids': [], 'logprobs': [], 'reward': []}\n",
|
| 301 |
+
" for prompt in prompts:\n",
|
| 302 |
+
" ep = rollout_once(\n",
|
| 303 |
+
" trainer=trainer, sync_env=env, tokenizer=tok,\n",
|
| 304 |
+
" dataset_prompt=prompt, max_steps=max_steps,\n",
|
| 305 |
+
" )\n",
|
| 306 |
+
" combined = (\n",
|
| 307 |
+
" ep['final_score_reward'] * 0.45 +\n",
|
| 308 |
+
" ep['no_wrong_channels_reward'] * 0.15 +\n",
|
| 309 |
+
" ep['valid_json_reward'] * 0.15 +\n",
|
| 310 |
+
" ep['read_runbook_reward'] * 0.15 +\n",
|
| 311 |
+
" ep['efficiency_reward'] * 0.10\n",
|
| 312 |
+
" )\n",
|
| 313 |
+
" out['prompt_ids'].append(ep['prompt_ids'])\n",
|
| 314 |
+
" out['completion_ids'].append(ep['completion_ids'])\n",
|
| 315 |
+
" out['logprobs'].append(ep['logprobs'])\n",
|
| 316 |
+
" out['reward'].append(combined)\n",
|
| 317 |
+
" return out\n",
|
| 318 |
+
" return rollout_func\n",
|
| 319 |
+
"\n",
|
| 320 |
+
"\n",
|
| 321 |
+
"rollout_func = make_rollout_func(sync_env, tokenizer)\n",
|
| 322 |
+
"\n",
|
| 323 |
+
"\n",
|
| 324 |
+
"def reward_func(completions, **kwargs):\n",
|
| 325 |
+
" \"\"\"Trivial passthrough — reads pre-computed reward from rollout kwargs.\"\"\"\n",
|
| 326 |
+
" n = len(completions)\n",
|
| 327 |
+
" rewards = kwargs.get('reward', [])\n",
|
| 328 |
+
" if not rewards:\n",
|
| 329 |
+
" print(\n",
|
| 330 |
+
" f'[ERROR] reward_func: empty kwargs! '\n",
|
| 331 |
+
" f'TRL {trl.__version__} is not passing rollout keys. '\n",
|
| 332 |
+
" f'kwargs keys present: {list(kwargs.keys())}'\n",
|
| 333 |
+
" )\n",
|
| 334 |
+
" return [0.0] * n\n",
|
| 335 |
+
" return [float(r) for r in rewards]\n",
|
| 336 |
+
"\n",
|
| 337 |
+
"\n",
|
| 338 |
+
"print('rollout_func + reward_func defined')"
|
| 339 |
+
]
|
| 340 |
+
},
|
| 341 |
+
{
|
| 342 |
+
"cell_type": "markdown",
|
| 343 |
+
"metadata": {},
|
| 344 |
+
"source": [
|
| 345 |
+
"## 9. Smoke Test\n",
|
| 346 |
+
"\n",
|
| 347 |
+
"Verify the env + reward pipeline works *before* committing to a 60-minute training run."
|
| 348 |
+
]
|
| 349 |
+
},
|
| 350 |
+
{
|
| 351 |
+
"cell_type": "code",
|
| 352 |
+
"execution_count": null,
|
| 353 |
+
"metadata": {},
|
| 354 |
+
"outputs": [],
|
| 355 |
+
"source": [
|
| 356 |
+
"print('─── Smoke test ──────────────────────────────────')\n",
|
| 357 |
+
"\n",
|
| 358 |
+
"# 1. Env resets + steps correctly\n",
|
| 359 |
+
"with GenericEnvClient(base_url=ENV_URL).sync() as _t:\n",
|
| 360 |
+
" r1 = _t.reset(seed=42)\n",
|
| 361 |
+
" obs = _obs_to_dict(r1.observation if hasattr(r1, 'observation') else r1)\n",
|
| 362 |
+
" assert obs.get('task_brief'), 'FAIL: no task_brief in observation'\n",
|
| 363 |
+
" print(f'[OK] env.reset() task_brief present')\n",
|
| 364 |
+
"\n",
|
| 365 |
+
" r2 = _t.step({'action_type': 'meta.read_runbook', 'args': {}})\n",
|
| 366 |
+
" print(f'[OK] env.step(read_runbook)')\n",
|
| 367 |
+
"\n",
|
| 368 |
+
" r3 = _t.step({'action_type': 'meta.finish', 'args': {}})\n",
|
| 369 |
+
" obs3 = _obs_to_dict(r3.observation if hasattr(r3, 'observation') else r3)\n",
|
| 370 |
+
" final = float(getattr(r3, 'reward', obs3.get('reward', -1)))\n",
|
| 371 |
+
" print(f'[OK] env.step(finish) reward={final:.3f}')\n",
|
| 372 |
+
"\n",
|
| 373 |
+
"# 2. reward_func passthrough\n",
|
| 374 |
+
"rf = reward_func(['c1', 'c2'], reward=[0.3, 0.7])\n",
|
| 375 |
+
"assert rf == [0.3, 0.7], f'FAIL: passthrough broken: {rf}'\n",
|
| 376 |
+
"print(f'[OK] reward_func passthrough {rf}')\n",
|
| 377 |
+
"\n",
|
| 378 |
+
"# 3. reward_func warns correctly on empty kwargs\n",
|
| 379 |
+
"import io, warnings\n",
|
| 380 |
+
"rf_empty = reward_func(['c1'])\n",
|
| 381 |
+
"assert rf_empty == [0.0]\n",
|
| 382 |
+
"print(f'[OK] reward_func empty → [0.0] (error printed above is expected)')\n",
|
| 383 |
+
"\n",
|
| 384 |
+
"print('\\nSmoke test PASSED')"
|
| 385 |
+
]
|
| 386 |
+
},
|
| 387 |
+
{
|
| 388 |
+
"cell_type": "markdown",
|
| 389 |
+
"metadata": {},
|
| 390 |
+
"source": [
|
| 391 |
+
"## 10. Configure GRPO Training"
|
| 392 |
+
]
|
| 393 |
+
},
|
| 394 |
+
{
|
| 395 |
+
"cell_type": "code",
|
| 396 |
+
"execution_count": null,
|
| 397 |
+
"metadata": {},
|
| 398 |
+
"outputs": [],
|
| 399 |
+
"source": [
|
| 400 |
+
"from trl import GRPOConfig\n",
|
| 401 |
+
"\n",
|
| 402 |
+
"OUTPUT_DIR = 'pm-ops-grpo-Qwen3-1.7B-triage-v2'\n",
|
| 403 |
+
"HF_REPO_ID = f'Saurav1/{OUTPUT_DIR}'\n",
|
| 404 |
+
"\n",
|
| 405 |
+
"grpo_config = GRPOConfig(\n",
|
| 406 |
+
" # Training\n",
|
| 407 |
+
" num_train_epochs=1,\n",
|
| 408 |
+
" learning_rate=5e-6,\n",
|
| 409 |
+
" gradient_accumulation_steps=GRAD_ACCUM,\n",
|
| 410 |
+
" per_device_train_batch_size=1,\n",
|
| 411 |
+
" warmup_steps=5,\n",
|
| 412 |
+
" num_generations=NUM_GEN,\n",
|
| 413 |
+
"\n",
|
| 414 |
+
" # Sequence lengths\n",
|
| 415 |
+
" max_completion_length=MAX_COMP_LEN,\n",
|
| 416 |
+
" max_prompt_length=4096,\n",
|
| 417 |
+
"\n",
|
| 418 |
+
" # No vLLM — standard HF generate (compatible with bitsandbytes + PEFT)\n",
|
| 419 |
+
" use_vllm=False,\n",
|
| 420 |
+
"\n",
|
| 421 |
+
" # Output\n",
|
| 422 |
+
" output_dir=OUTPUT_DIR,\n",
|
| 423 |
+
" report_to='trackio',\n",
|
| 424 |
+
" trackio_space_id=OUTPUT_DIR,\n",
|
| 425 |
+
" logging_steps=1,\n",
|
| 426 |
+
" save_steps=25,\n",
|
| 427 |
+
" gradient_checkpointing=True, # saves ~30% activation memory\n",
|
| 428 |
+
")\n",
|
| 429 |
+
"\n",
|
| 430 |
+
"eff_batch = grpo_config.per_device_train_batch_size * GRAD_ACCUM\n",
|
| 431 |
+
"total_steps = len(dataset) * NUM_GEN // eff_batch\n",
|
| 432 |
+
"print(f'Output dir : {OUTPUT_DIR}')\n",
|
| 433 |
+
"print(f'Eff. batch : {eff_batch}')\n",
|
| 434 |
+
"print(f'Total steps : ~{total_steps}')"
|
| 435 |
+
]
|
| 436 |
+
},
|
| 437 |
+
{
|
| 438 |
+
"cell_type": "markdown",
|
| 439 |
+
"metadata": {},
|
| 440 |
+
"source": [
|
| 441 |
+
"## 11. Create Trainer"
|
| 442 |
+
]
|
| 443 |
+
},
|
| 444 |
+
{
|
| 445 |
+
"cell_type": "code",
|
| 446 |
+
"execution_count": null,
|
| 447 |
+
"metadata": {},
|
| 448 |
+
"outputs": [],
|
| 449 |
+
"source": [
|
| 450 |
+
"from training.pm_ops_trainer import PMOpsGRPOTrainer\n",
|
| 451 |
+
"\n",
|
| 452 |
+
"trainer = PMOpsGRPOTrainer(\n",
|
| 453 |
+
" model=model,\n",
|
| 454 |
+
" processing_class=tokenizer,\n",
|
| 455 |
+
" reward_funcs=reward_func, # fallback only — injected rewards take priority\n",
|
| 456 |
+
" train_dataset=dataset,\n",
|
| 457 |
+
" args=grpo_config,\n",
|
| 458 |
+
" rollout_func=rollout_func,\n",
|
| 459 |
+
")\n",
|
| 460 |
+
"print(\"PMOpsGRPOTrainer ready — direct reward injection active\")"
|
| 461 |
+
]
|
| 462 |
+
},
|
| 463 |
+
{
|
| 464 |
+
"cell_type": "markdown",
|
| 465 |
+
"metadata": {},
|
| 466 |
+
"source": [
|
| 467 |
+
"## 12. Train\n",
|
| 468 |
+
"\n",
|
| 469 |
+
"**Watch Colab stdout for rollout lines:**\n",
|
| 470 |
+
"```\n",
|
| 471 |
+
"[rollout] steps=5 final=0.720 json=0.80 runbook=1 no_wrong=1.00 eff=0.67 → combined=0.674\n",
|
| 472 |
+
"```\n",
|
| 473 |
+
"- `combined > 0` in stdout **and** `train/reward > 0` in trackio → working correctly \n",
|
| 474 |
+
"- `combined > 0` in stdout **but** `train/reward = 0` in trackio → kwargs still not flowing \n",
|
| 475 |
+
" → Add `print(kwargs.keys())` inside `reward_func` to debug \n",
|
| 476 |
+
"- `combined = 0` always → check env connection / reward logic"
|
| 477 |
+
]
|
| 478 |
+
},
|
| 479 |
+
{
|
| 480 |
+
"cell_type": "code",
|
| 481 |
+
"execution_count": null,
|
| 482 |
+
"metadata": {},
|
| 483 |
+
"outputs": [],
|
| 484 |
+
"source": [
|
| 485 |
+
"trainer_stats = trainer.train()"
|
| 486 |
+
]
|
| 487 |
+
},
|
| 488 |
+
{
|
| 489 |
+
"cell_type": "code",
|
| 490 |
+
"execution_count": null,
|
| 491 |
+
"metadata": {},
|
| 492 |
+
"outputs": [],
|
| 493 |
+
"source": [
|
| 494 |
+
"used_gb = round(torch.cuda.max_memory_reserved() / 1024**3, 2)\n",
|
| 495 |
+
"train_mins = round(trainer_stats.metrics.get('train_runtime', 0) / 60, 1)\n",
|
| 496 |
+
"print(f'Training time : {train_mins} min')\n",
|
| 497 |
+
"print(f'Peak GPU : {used_gb} GB / {TOTAL_GB} GB ({round(used_gb/TOTAL_GB*100,1)}%)')\n",
|
| 498 |
+
"print(f'Final loss : {trainer_stats.metrics.get(\"train_loss\", \"n/a\")}')"
|
| 499 |
+
]
|
| 500 |
+
},
|
| 501 |
+
{
|
| 502 |
+
"cell_type": "markdown",
|
| 503 |
+
"metadata": {},
|
| 504 |
+
"source": [
|
| 505 |
+
"## 13. Save + Push"
|
| 506 |
+
]
|
| 507 |
+
},
|
| 508 |
+
{
|
| 509 |
+
"cell_type": "code",
|
| 510 |
+
"execution_count": null,
|
| 511 |
+
"metadata": {},
|
| 512 |
+
"outputs": [],
|
| 513 |
+
"source": [
|
| 514 |
+
"sync_env.close()\n",
|
| 515 |
+
"\n",
|
| 516 |
+
"# Save LoRA adapter (not merged — push separately if needed)\n",
|
| 517 |
+
"trainer.save_model(OUTPUT_DIR)\n",
|
| 518 |
+
"tokenizer.save_pretrained(OUTPUT_DIR)\n",
|
| 519 |
+
"print(f'Saved LoRA adapter → {OUTPUT_DIR}')\n",
|
| 520 |
+
"\n",
|
| 521 |
+
"# Push to Hub\n",
|
| 522 |
+
"trainer.model.push_to_hub(HF_REPO_ID)\n",
|
| 523 |
+
"tokenizer.push_to_hub(HF_REPO_ID)\n",
|
| 524 |
+
"print(f'Pushed → https://huggingface.co/{HF_REPO_ID}')"
|
| 525 |
+
]
|
| 526 |
+
},
|
| 527 |
+
{
|
| 528 |
+
"cell_type": "markdown",
|
| 529 |
+
"metadata": {},
|
| 530 |
+
"source": [
|
| 531 |
+
"## 14. Merge LoRA → bf16 (Optional — for full-weight inference)"
|
| 532 |
+
]
|
| 533 |
+
},
|
| 534 |
+
{
|
| 535 |
+
"cell_type": "code",
|
| 536 |
+
"execution_count": null,
|
| 537 |
+
"metadata": {},
|
| 538 |
+
"outputs": [],
|
| 539 |
+
"source": [
|
| 540 |
+
"# Only run this if you have enough CPU RAM (~7 GB free) to hold the merged model.\n",
|
| 541 |
+
"# On T4 Colab this often OOMs — skip and load with PEFT for eval instead.\n",
|
| 542 |
+
"from peft import PeftModel\n",
|
| 543 |
+
"from transformers import AutoModelForCausalLM\n",
|
| 544 |
+
"\n",
|
| 545 |
+
"base = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16, device_map='cpu')\n",
|
| 546 |
+
"merged = PeftModel.from_pretrained(base, OUTPUT_DIR)\n",
|
| 547 |
+
"merged = merged.merge_and_unload()\n",
|
| 548 |
+
"merged.save_pretrained(f'{OUTPUT_DIR}-merged')\n",
|
| 549 |
+
"tokenizer.save_pretrained(f'{OUTPUT_DIR}-merged')\n",
|
| 550 |
+
"print(f'Merged model saved → {OUTPUT_DIR}-merged')"
|
| 551 |
+
]
|
| 552 |
+
},
|
| 553 |
+
{
|
| 554 |
+
"cell_type": "markdown",
|
| 555 |
+
"metadata": {},
|
| 556 |
+
"source": [
|
| 557 |
+
"## 15. Evaluate: Baseline vs Trained"
|
| 558 |
+
]
|
| 559 |
+
},
|
| 560 |
+
{
|
| 561 |
+
"cell_type": "code",
|
| 562 |
+
"execution_count": null,
|
| 563 |
+
"metadata": {},
|
| 564 |
+
"outputs": [],
|
| 565 |
+
"source": [
|
| 566 |
+
"from training.rollout import extract_json_action, step_aware_fallback, build_messages\n",
|
| 567 |
+
"from training.rollout import _current_obs_text, _obs_to_dict\n",
|
| 568 |
+
"from inference import baseline_agent\n",
|
| 569 |
+
"\n",
|
| 570 |
+
"EVAL_URL = 'https://adityaguntur-pm-ops.hf.space'\n",
|
| 571 |
+
"N_EVAL = 10\n",
|
| 572 |
+
"EVAL_MAX_STEPS = 15\n",
|
| 573 |
+
"\n",
|
| 574 |
+
"# Put model in eval/inference mode\n",
|
| 575 |
+
"model.eval()\n",
|
| 576 |
+
"\n",
|
| 577 |
+
"\n",
|
| 578 |
+
"def eval_trained(env, mdl, tok, n=N_EVAL):\n",
|
| 579 |
+
" scores = []\n",
|
| 580 |
+
" for i in range(n):\n",
|
| 581 |
+
" result = env.reset()\n",
|
| 582 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 583 |
+
" task_brief = obs_dict.get('task_brief', '')\n",
|
| 584 |
+
" history, step, score, done = [], 0, 0.0, False\n",
|
| 585 |
+
"\n",
|
| 586 |
+
" while not done and step < EVAL_MAX_STEPS:\n",
|
| 587 |
+
" obs_text = _current_obs_text(obs_dict, step, task_brief)\n",
|
| 588 |
+
" msgs = build_messages(history, obs_text)\n",
|
| 589 |
+
" prompt_t = tok.apply_chat_template(\n",
|
| 590 |
+
" msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n",
|
| 591 |
+
" )\n",
|
| 592 |
+
" inputs = tok([prompt_t], return_tensors='pt', truncation=True, max_length=4096)\n",
|
| 593 |
+
" inputs = {k: v.to(mdl.device) for k, v in inputs.items()}\n",
|
| 594 |
+
" with torch.no_grad():\n",
|
| 595 |
+
" out_ids = mdl.generate(**inputs, max_new_tokens=512,\n",
|
| 596 |
+
" do_sample=False, pad_token_id=tok.eos_token_id)\n",
|
| 597 |
+
" completion = tok.decode(out_ids[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)\n",
|
| 598 |
+
"\n",
|
| 599 |
+
" parsed = extract_json_action(completion) or step_aware_fallback(step)\n",
|
| 600 |
+
" result = env.step({'action_type': parsed['action_type'], 'args': parsed.get('args', {})})\n",
|
| 601 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 602 |
+
" done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n",
|
| 603 |
+
" score = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n",
|
| 604 |
+
" history.append({'obs_text': obs_text, 'completion': completion, 'is_runbook': False})\n",
|
| 605 |
+
" step += 1\n",
|
| 606 |
+
"\n",
|
| 607 |
+
" scores.append(score)\n",
|
| 608 |
+
" print(f' Trained ep {i+1}/{n}: score={score:.3f}')\n",
|
| 609 |
+
" return scores\n",
|
| 610 |
+
"\n",
|
| 611 |
+
"\n",
|
| 612 |
+
"def eval_baseline(env, n=N_EVAL):\n",
|
| 613 |
+
" scores = []\n",
|
| 614 |
+
" for i in range(n):\n",
|
| 615 |
+
" result = env.reset()\n",
|
| 616 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 617 |
+
" org_cfg, done, step, score = {}, False, 0, 0.0\n",
|
| 618 |
+
" while not done and step < EVAL_MAX_STEPS:\n",
|
| 619 |
+
" at, args = baseline_agent(obs_dict, org_cfg)\n",
|
| 620 |
+
" result = env.step({'action_type': at, 'args': args})\n",
|
| 621 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 622 |
+
" done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n",
|
| 623 |
+
" score = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n",
|
| 624 |
+
" step += 1\n",
|
| 625 |
+
" scores.append(score)\n",
|
| 626 |
+
" print(f' Baseline ep {i+1}/{n}: score={score:.3f}')\n",
|
| 627 |
+
" return scores\n",
|
| 628 |
+
"\n",
|
| 629 |
+
"\n",
|
| 630 |
+
"print('─── Baseline ───')\n",
|
| 631 |
+
"with GenericEnvClient(base_url=EVAL_URL).sync() as env_eval:\n",
|
| 632 |
+
" baseline_scores = eval_baseline(env_eval)\n",
|
| 633 |
+
"\n",
|
| 634 |
+
"print('─── Trained ────')\n",
|
| 635 |
+
"with GenericEnvClient(base_url=EVAL_URL).sync() as env_eval:\n",
|
| 636 |
+
" trained_scores = eval_trained(env_eval, model, tokenizer)\n",
|
| 637 |
+
"\n",
|
| 638 |
+
"print(f'\\nBaseline avg : {sum(baseline_scores)/N_EVAL:.3f}')\n",
|
| 639 |
+
"print(f'Trained avg : {sum(trained_scores)/N_EVAL:.3f}')\n",
|
| 640 |
+
"print(f'Improvement : +{(sum(trained_scores)-sum(baseline_scores))/N_EVAL:.3f}')"
|
| 641 |
+
]
|
| 642 |
+
},
|
| 643 |
+
{
|
| 644 |
+
"cell_type": "markdown",
|
| 645 |
+
"metadata": {},
|
| 646 |
+
"source": [
|
| 647 |
+
"## 16. Plot"
|
| 648 |
+
]
|
| 649 |
+
},
|
| 650 |
+
{
|
| 651 |
+
"cell_type": "code",
|
| 652 |
+
"execution_count": null,
|
| 653 |
+
"metadata": {},
|
| 654 |
+
"outputs": [],
|
| 655 |
+
"source": [
|
| 656 |
+
"import matplotlib.pyplot as plt\n",
|
| 657 |
+
"import numpy as np\n",
|
| 658 |
+
"\n",
|
| 659 |
+
"fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n",
|
| 660 |
+
"\n",
|
| 661 |
+
"ax, x, w = axes[0], np.arange(N_EVAL), 0.35\n",
|
| 662 |
+
"ax.bar(x - w/2, baseline_scores, w, label='Baseline', color='steelblue', alpha=0.8)\n",
|
| 663 |
+
"ax.bar(x + w/2, trained_scores, w, label='GRPO v2', color='coral', alpha=0.8)\n",
|
| 664 |
+
"ax.axhline(sum(baseline_scores)/N_EVAL, color='steelblue', ls='--', alpha=0.5, lw=1)\n",
|
| 665 |
+
"ax.axhline(sum(trained_scores)/N_EVAL, color='coral', ls='--', alpha=0.5, lw=1)\n",
|
| 666 |
+
"ax.set(xlabel='Episode', ylabel='Reward', title='Per-episode: Baseline vs GRPO v2',\n",
|
| 667 |
+
" xticks=x, ylim=(0, 1.05))\n",
|
| 668 |
+
"ax.legend()\n",
|
| 669 |
+
"\n",
|
| 670 |
+
"ax2 = axes[1]\n",
|
| 671 |
+
"avgs = [sum(baseline_scores)/N_EVAL, sum(trained_scores)/N_EVAL]\n",
|
| 672 |
+
"bars = ax2.bar(['Baseline', 'GRPO v2'], avgs, color=['steelblue', 'coral'], alpha=0.85, width=0.5)\n",
|
| 673 |
+
"for bar, val in zip(bars, avgs):\n",
|
| 674 |
+
" ax2.text(bar.get_x() + bar.get_width()/2, val + 0.01, f'{val:.3f}',\n",
|
| 675 |
+
" ha='center', fontsize=13, fontweight='bold')\n",
|
| 676 |
+
"ax2.set(ylabel='Average reward', title=f'Avg over {N_EVAL} episodes', ylim=(0, 1.05))\n",
|
| 677 |
+
"\n",
|
| 678 |
+
"plt.tight_layout()\n",
|
| 679 |
+
"plt.savefig('eval_results_v2.png', dpi=150, bbox_inches='tight')\n",
|
| 680 |
+
"plt.show()\n",
|
| 681 |
+
"print('Saved: eval_results_v2.png')"
|
| 682 |
+
]
|
| 683 |
+
},
|
| 684 |
+
{
|
| 685 |
+
"cell_type": "markdown",
|
| 686 |
+
"metadata": {},
|
| 687 |
+
"source": [
|
| 688 |
+
"## 17. Teardown"
|
| 689 |
+
]
|
| 690 |
+
},
|
| 691 |
+
{
|
| 692 |
+
"cell_type": "code",
|
| 693 |
+
"execution_count": null,
|
| 694 |
+
"metadata": {},
|
| 695 |
+
"outputs": [],
|
| 696 |
+
"source": [
|
| 697 |
+
"server_proc.terminate()\n",
|
| 698 |
+
"print('Local PM-Ops server stopped')"
|
| 699 |
+
]
|
| 700 |
+
}
|
| 701 |
+
],
|
| 702 |
+
"metadata": {
|
| 703 |
+
"kernelspec": {
|
| 704 |
+
"display_name": "Python 3",
|
| 705 |
+
"language": "python",
|
| 706 |
+
"name": "python3"
|
| 707 |
+
},
|
| 708 |
+
"language_info": {
|
| 709 |
+
"name": "python",
|
| 710 |
+
"version": "3.12.0"
|
| 711 |
+
}
|
| 712 |
+
},
|
| 713 |
+
"nbformat": 4,
|
| 714 |
+
"nbformat_minor": 4
|
| 715 |
+
}
|
training/train_v3.ipynb
ADDED
|
@@ -0,0 +1,814 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"cells": [
|
| 3 |
+
{
|
| 4 |
+
"cell_type": "markdown",
|
| 5 |
+
"id": "cell-0",
|
| 6 |
+
"metadata": {},
|
| 7 |
+
"source": [
|
| 8 |
+
"# PM-Ops GRPO Training v3 — SFT Warmup + GRPO\n",
|
| 9 |
+
"\n",
|
| 10 |
+
"## What changed from v2\n",
|
| 11 |
+
"| Problem in v2 | Fix in v3 |\n",
|
| 12 |
+
"|---|---|\n",
|
| 13 |
+
"| `json=0.00` — model never outputs JSON, GRPO gradient = 0 | **SFT warmup on 60 baseline traces first** |\n",
|
| 14 |
+
"| Reward = 0.300 every episode (no variance) | **env_score varies 0.25–1.0 after SFT** |\n",
|
| 15 |
+
"| Slow — no vLLM, eager attention | **Unsloth + `use_vllm=True`** |\n",
|
| 16 |
+
"| 5 reward signals, all from fallback | **2 signals: env_score × 0.85 + json_ratio × 0.15** |\n",
|
| 17 |
+
"| `meta.finish` never called in 15-step cap | **Fixed fallback respects `max_steps`** |\n",
|
| 18 |
+
"\n",
|
| 19 |
+
"## Why SFT before GRPO?\n",
|
| 20 |
+
"GRPO learns by comparing rewards across a *group* of generations. If all generations get the\n",
|
| 21 |
+
"same reward (because the model outputs garbage on every step), the advantage is 0 and weights\n",
|
| 22 |
+
"don't move. SFT warmup costs ~15 min and unlocks the full GRPO gradient.\n",
|
| 23 |
+
"\n",
|
| 24 |
+
"**Stack**: Unsloth + TRL 1.2.0 + OpenEnv · **GPU**: A100 → ~75 min total"
|
| 25 |
+
]
|
| 26 |
+
},
|
| 27 |
+
{
|
| 28 |
+
"cell_type": "markdown",
|
| 29 |
+
"id": "cell-1-md",
|
| 30 |
+
"metadata": {},
|
| 31 |
+
"source": [
|
| 32 |
+
"## 0. Install"
|
| 33 |
+
]
|
| 34 |
+
},
|
| 35 |
+
{
|
| 36 |
+
"cell_type": "code",
|
| 37 |
+
"execution_count": null,
|
| 38 |
+
"id": "cell-1",
|
| 39 |
+
"metadata": {},
|
| 40 |
+
"outputs": [],
|
| 41 |
+
"source": [
|
| 42 |
+
"# Unsloth + vLLM first — let Unsloth resolve torch compat\n",
|
| 43 |
+
"!pip install -q unsloth vllm\n",
|
| 44 |
+
"# TRL training stack\n",
|
| 45 |
+
"!pip install -q \"trl==1.2.0\" accelerate datasets\n",
|
| 46 |
+
"# PM-Ops server runtime\n",
|
| 47 |
+
"!pip install -q \"openenv-core>=0.2.2\" \"fastapi>=0.110.0\" \"uvicorn[standard]>=0.29.0\" \"pydantic>=2.0.0\"\n",
|
| 48 |
+
"# Experiment tracking\n",
|
| 49 |
+
"!pip install -q trackio\n",
|
| 50 |
+
"print('Done - restart kernel, then run all cells from top.')"
|
| 51 |
+
]
|
| 52 |
+
},
|
| 53 |
+
{
|
| 54 |
+
"cell_type": "markdown",
|
| 55 |
+
"id": "cell-2-md",
|
| 56 |
+
"metadata": {},
|
| 57 |
+
"source": [
|
| 58 |
+
"## 1. Imports + GPU Config"
|
| 59 |
+
]
|
| 60 |
+
},
|
| 61 |
+
{
|
| 62 |
+
"cell_type": "code",
|
| 63 |
+
"execution_count": null,
|
| 64 |
+
"id": "cell-2",
|
| 65 |
+
"metadata": {},
|
| 66 |
+
"outputs": [],
|
| 67 |
+
"source": [
|
| 68 |
+
"import torch\n",
|
| 69 |
+
"\n",
|
| 70 |
+
"# PatchFastRL can bypass custom rollout_func on some cloud runtimes.\n",
|
| 71 |
+
"# Keep it disabled for PMOpsGRPOTrainer rollout compatibility.\n",
|
| 72 |
+
"from unsloth import FastLanguageModel, PatchFastRL\n",
|
| 73 |
+
"ENABLE_FAST_RL_PATCH = False\n",
|
| 74 |
+
"if ENABLE_FAST_RL_PATCH:\n",
|
| 75 |
+
" PatchFastRL('GRPO', FastLanguageModel)\n",
|
| 76 |
+
" print('PatchFastRL enabled')\n",
|
| 77 |
+
"else:\n",
|
| 78 |
+
" print('PatchFastRL disabled for rollout_func compatibility')\n",
|
| 79 |
+
"\n",
|
| 80 |
+
"import trl\n",
|
| 81 |
+
"print(f'torch : {torch.__version__}')\n",
|
| 82 |
+
"print(f'TRL : {trl.__version__}')\n",
|
| 83 |
+
"\n",
|
| 84 |
+
"gpu = torch.cuda.get_device_properties(0)\n",
|
| 85 |
+
"TOTAL_GB = round(gpu.total_memory / 1024**3, 1)\n",
|
| 86 |
+
"IS_A100 = TOTAL_GB >= 35\n",
|
| 87 |
+
"print(f'GPU : {gpu.name} ({TOTAL_GB} GB)')\n",
|
| 88 |
+
"\n",
|
| 89 |
+
"# Adaptive config — T4 uses minimal settings for smoke-testing\n",
|
| 90 |
+
"NUM_GEN = 6 if IS_A100 else 2\n",
|
| 91 |
+
"GRAD_ACCUM = 32 if IS_A100 else 8\n",
|
| 92 |
+
"MAX_COMP_LEN = 384\n",
|
| 93 |
+
"N_SFT_EPISODES = 60 if IS_A100 else 15\n",
|
| 94 |
+
"N_GRPO_EPISODES = 150 if IS_A100 else 30\n",
|
| 95 |
+
"TRAIN_MAX_STEPS = 12 # triage solvable in 5; 12 gives exploration room\n",
|
| 96 |
+
"\n",
|
| 97 |
+
"print(f'num_gen={NUM_GEN} grad_accum={GRAD_ACCUM} '\n",
|
| 98 |
+
" f'sft_eps={N_SFT_EPISODES} grpo_eps={N_GRPO_EPISODES}')"
|
| 99 |
+
]
|
| 100 |
+
},
|
| 101 |
+
{
|
| 102 |
+
"cell_type": "markdown",
|
| 103 |
+
"id": "cell-3-md",
|
| 104 |
+
"metadata": {},
|
| 105 |
+
"source": [
|
| 106 |
+
"## 2. Clone PM-Ops Repo"
|
| 107 |
+
]
|
| 108 |
+
},
|
| 109 |
+
{
|
| 110 |
+
"cell_type": "code",
|
| 111 |
+
"execution_count": null,
|
| 112 |
+
"id": "cell-3",
|
| 113 |
+
"metadata": {},
|
| 114 |
+
"outputs": [],
|
| 115 |
+
"source": [
|
| 116 |
+
"import os, sys\n",
|
| 117 |
+
"\n",
|
| 118 |
+
"REPO_URL = 'https://huggingface.co/spaces/TheCrustaceans/Pm-ops'\n",
|
| 119 |
+
"REPO_DIR = '/content/Pm_ops'\n",
|
| 120 |
+
"\n",
|
| 121 |
+
"if not os.path.exists(REPO_DIR):\n",
|
| 122 |
+
" !git clone --depth=1 -q {REPO_URL} {REPO_DIR}\n",
|
| 123 |
+
" print(f'Cloned -> {REPO_DIR}')\n",
|
| 124 |
+
"else:\n",
|
| 125 |
+
" !git -C {REPO_DIR} pull -q origin main\n",
|
| 126 |
+
" print(f'Pulled -> {REPO_DIR}')\n",
|
| 127 |
+
"\n",
|
| 128 |
+
"for p in [REPO_DIR, os.path.join(REPO_DIR, 'training')]:\n",
|
| 129 |
+
" if p not in sys.path:\n",
|
| 130 |
+
" sys.path.insert(0, p)\n",
|
| 131 |
+
"os.chdir(REPO_DIR)\n",
|
| 132 |
+
"print(f'CWD: {os.getcwd()}')"
|
| 133 |
+
]
|
| 134 |
+
},
|
| 135 |
+
{
|
| 136 |
+
"cell_type": "markdown",
|
| 137 |
+
"id": "cell-4-md",
|
| 138 |
+
"metadata": {},
|
| 139 |
+
"source": [
|
| 140 |
+
"## 3. HuggingFace Login"
|
| 141 |
+
]
|
| 142 |
+
},
|
| 143 |
+
{
|
| 144 |
+
"cell_type": "code",
|
| 145 |
+
"execution_count": null,
|
| 146 |
+
"id": "cell-4",
|
| 147 |
+
"metadata": {},
|
| 148 |
+
"outputs": [],
|
| 149 |
+
"source": [
|
| 150 |
+
"from huggingface_hub import notebook_login\n",
|
| 151 |
+
"notebook_login()"
|
| 152 |
+
]
|
| 153 |
+
},
|
| 154 |
+
{
|
| 155 |
+
"cell_type": "markdown",
|
| 156 |
+
"id": "cell-5-md",
|
| 157 |
+
"metadata": {},
|
| 158 |
+
"source": [
|
| 159 |
+
"## 4. Start Local PM-Ops Server"
|
| 160 |
+
]
|
| 161 |
+
},
|
| 162 |
+
{
|
| 163 |
+
"cell_type": "code",
|
| 164 |
+
"execution_count": null,
|
| 165 |
+
"id": "cell-5",
|
| 166 |
+
"metadata": {},
|
| 167 |
+
"outputs": [],
|
| 168 |
+
"source": [
|
| 169 |
+
"import subprocess, time, requests\n",
|
| 170 |
+
"\n",
|
| 171 |
+
"server_proc = subprocess.Popen(\n",
|
| 172 |
+
" [sys.executable, '-m', 'uvicorn', 'server.app:app',\n",
|
| 173 |
+
" '--host', '0.0.0.0', '--port', '8000'],\n",
|
| 174 |
+
" cwd=REPO_DIR, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,\n",
|
| 175 |
+
")\n",
|
| 176 |
+
"ENV_URL = 'http://localhost:8000'\n",
|
| 177 |
+
"\n",
|
| 178 |
+
"for _ in range(30):\n",
|
| 179 |
+
" try:\n",
|
| 180 |
+
" if requests.get(f'{ENV_URL}/', timeout=2).status_code == 200:\n",
|
| 181 |
+
" print(f'PM-Ops server ready pid={server_proc.pid}')\n",
|
| 182 |
+
" break\n",
|
| 183 |
+
" except Exception:\n",
|
| 184 |
+
" pass\n",
|
| 185 |
+
" time.sleep(1)\n",
|
| 186 |
+
"else:\n",
|
| 187 |
+
" raise RuntimeError('Server did not start in 30 s')"
|
| 188 |
+
]
|
| 189 |
+
},
|
| 190 |
+
{
|
| 191 |
+
"cell_type": "markdown",
|
| 192 |
+
"id": "cell-6-md",
|
| 193 |
+
"metadata": {},
|
| 194 |
+
"source": [
|
| 195 |
+
"## 5. Verify Env"
|
| 196 |
+
]
|
| 197 |
+
},
|
| 198 |
+
{
|
| 199 |
+
"cell_type": "code",
|
| 200 |
+
"execution_count": null,
|
| 201 |
+
"id": "cell-6",
|
| 202 |
+
"metadata": {},
|
| 203 |
+
"outputs": [],
|
| 204 |
+
"source": [
|
| 205 |
+
"import trl.experimental.openenv # must be importable\n",
|
| 206 |
+
"from openenv.core import GenericEnvClient\n",
|
| 207 |
+
"from training.rollout import _obs_to_dict\n",
|
| 208 |
+
"\n",
|
| 209 |
+
"with GenericEnvClient(base_url=ENV_URL).sync() as _env:\n",
|
| 210 |
+
" r = _env.reset()\n",
|
| 211 |
+
" obs = _obs_to_dict(r.observation if hasattr(r, 'observation') else r)\n",
|
| 212 |
+
" print(f'task_brief : {obs.get(\"task_brief\", \"?\")[:80]}...')\n",
|
| 213 |
+
" _env.step({'action_type': 'meta.read_runbook', 'args': {}})\n",
|
| 214 |
+
" print('Env step : OK')"
|
| 215 |
+
]
|
| 216 |
+
},
|
| 217 |
+
{
|
| 218 |
+
"cell_type": "markdown",
|
| 219 |
+
"id": "cell-7-md",
|
| 220 |
+
"metadata": {},
|
| 221 |
+
"source": [
|
| 222 |
+
"## 6. Load Model — Unsloth 4-bit + LoRA"
|
| 223 |
+
]
|
| 224 |
+
},
|
| 225 |
+
{
|
| 226 |
+
"cell_type": "code",
|
| 227 |
+
"execution_count": null,
|
| 228 |
+
"id": "cell-7",
|
| 229 |
+
"metadata": {},
|
| 230 |
+
"outputs": [],
|
| 231 |
+
"source": [
|
| 232 |
+
"MODEL_NAME = 'Qwen/Qwen3-1.7B'\n",
|
| 233 |
+
"LORA_RANK = 16\n",
|
| 234 |
+
"\n",
|
| 235 |
+
"model, tokenizer = FastLanguageModel.from_pretrained(\n",
|
| 236 |
+
" model_name = MODEL_NAME,\n",
|
| 237 |
+
" max_seq_length = 4096 + MAX_COMP_LEN,\n",
|
| 238 |
+
" load_in_4bit = True,\n",
|
| 239 |
+
" fast_inference = False, # disable fast RL path to preserve rollout_func behaviour\n",
|
| 240 |
+
" max_lora_rank = LORA_RANK,\n",
|
| 241 |
+
" gpu_memory_utilization = 0.50, # leave headroom for SFT activations\n",
|
| 242 |
+
")\n",
|
| 243 |
+
"model = FastLanguageModel.get_peft_model(\n",
|
| 244 |
+
" model,\n",
|
| 245 |
+
" r = LORA_RANK,\n",
|
| 246 |
+
" target_modules = ['q_proj','k_proj','v_proj','o_proj',\n",
|
| 247 |
+
" 'gate_proj','up_proj','down_proj'],\n",
|
| 248 |
+
" lora_alpha = LORA_RANK,\n",
|
| 249 |
+
" use_gradient_checkpointing = 'unsloth',\n",
|
| 250 |
+
" random_state = 42,\n",
|
| 251 |
+
" )\n",
|
| 252 |
+
"tokenizer.pad_token = tokenizer.eos_token\n",
|
| 253 |
+
"tokenizer.padding_side = 'left'\n",
|
| 254 |
+
"model.print_trainable_parameters()\n",
|
| 255 |
+
"reserved = round(torch.cuda.max_memory_reserved() / 1024**3, 2)\n",
|
| 256 |
+
"print(f'GPU after load: {reserved} GB / {TOTAL_GB} GB')"
|
| 257 |
+
]
|
| 258 |
+
},
|
| 259 |
+
{
|
| 260 |
+
"cell_type": "markdown",
|
| 261 |
+
"id": "cell-8-md",
|
| 262 |
+
"metadata": {},
|
| 263 |
+
"source": [
|
| 264 |
+
"---\n",
|
| 265 |
+
"## Phase 1 — SFT Warmup\n",
|
| 266 |
+
"\n",
|
| 267 |
+
"**Goal**: teach the model the JSON output format and PM-ops workflow before GRPO.\n",
|
| 268 |
+
"\n",
|
| 269 |
+
"We run `baseline_agent` (the deterministic heuristic) for 60 episodes and record every\n",
|
| 270 |
+
"(observation, action) pair as a supervised example. Each episode produces ~6 steps:\n",
|
| 271 |
+
"`read_runbook` → `create_ticket` → `assign_ticket` → `list_channels` → `post_message` → `finish`.\n",
|
| 272 |
+
"\n",
|
| 273 |
+
"After 2 SFT epochs (~15 min), the model reliably outputs `\\`\\`\\`json ... \\`\\`\\`` blocks.\n",
|
| 274 |
+
"Without this, GRPO reward variance ≈ 0 and nothing is learned."
|
| 275 |
+
]
|
| 276 |
+
},
|
| 277 |
+
{
|
| 278 |
+
"cell_type": "markdown",
|
| 279 |
+
"id": "cell-9-md",
|
| 280 |
+
"metadata": {},
|
| 281 |
+
"source": [
|
| 282 |
+
"## 7. Generate SFT Demonstration Dataset"
|
| 283 |
+
]
|
| 284 |
+
},
|
| 285 |
+
{
|
| 286 |
+
"cell_type": "code",
|
| 287 |
+
"execution_count": null,
|
| 288 |
+
"id": "cell-9",
|
| 289 |
+
"metadata": {},
|
| 290 |
+
"outputs": [],
|
| 291 |
+
"source": [
|
| 292 |
+
"import json as _json\n",
|
| 293 |
+
"from datasets import Dataset\n",
|
| 294 |
+
"from inference import baseline_agent\n",
|
| 295 |
+
"from training.rollout import _obs_to_dict, _current_obs_text, build_messages\n",
|
| 296 |
+
"\n",
|
| 297 |
+
"\n",
|
| 298 |
+
"def generate_sft_dataset(env_url, tok, n_episodes, seed_start=2000):\n",
|
| 299 |
+
" \"\"\"Run baseline_agent for each episode; record (prompt, completion) pairs.\"\"\"\n",
|
| 300 |
+
" examples = []\n",
|
| 301 |
+
" with GenericEnvClient(base_url=env_url).sync() as env:\n",
|
| 302 |
+
" for i in range(n_episodes):\n",
|
| 303 |
+
" seed = seed_start + i\n",
|
| 304 |
+
" result = env.reset(seed=seed)\n",
|
| 305 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 306 |
+
" task_brief = obs_dict.get('task_brief', '')\n",
|
| 307 |
+
" turn_history = []\n",
|
| 308 |
+
" org_config = {}\n",
|
| 309 |
+
" step, done = 0, False\n",
|
| 310 |
+
"\n",
|
| 311 |
+
" while not done and step < 8:\n",
|
| 312 |
+
" obs_text = _current_obs_text(obs_dict, step, task_brief)\n",
|
| 313 |
+
" action_type, args = baseline_agent(obs_dict, org_config)\n",
|
| 314 |
+
"\n",
|
| 315 |
+
" # Target completion: JSON code block (what we want the model to learn)\n",
|
| 316 |
+
" payload = {'action_type': action_type, 'args': args}\n",
|
| 317 |
+
" completion = '```json\\n' + _json.dumps(payload) + '\\n```'\n",
|
| 318 |
+
"\n",
|
| 319 |
+
" msgs = build_messages(turn_history, obs_text)\n",
|
| 320 |
+
" prompt = tok.apply_chat_template(\n",
|
| 321 |
+
" msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n",
|
| 322 |
+
" )\n",
|
| 323 |
+
" # Full SFT text = prompt + target completion + eos\n",
|
| 324 |
+
" examples.append({'text': prompt + completion + tok.eos_token})\n",
|
| 325 |
+
"\n",
|
| 326 |
+
" turn_history.append({\n",
|
| 327 |
+
" 'obs_text' : obs_text,\n",
|
| 328 |
+
" 'completion': completion,\n",
|
| 329 |
+
" 'is_runbook': (action_type == 'meta.read_runbook'),\n",
|
| 330 |
+
" })\n",
|
| 331 |
+
"\n",
|
| 332 |
+
" result = env.step({'action_type': action_type, 'args': args})\n",
|
| 333 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 334 |
+
"\n",
|
| 335 |
+
" # Sync org_config from runbook response\n",
|
| 336 |
+
" if action_type == 'meta.read_runbook':\n",
|
| 337 |
+
" last = obs_dict.get('last_action_result') or {}\n",
|
| 338 |
+
" if last.get('ok'):\n",
|
| 339 |
+
" data = last.get('data') or {}\n",
|
| 340 |
+
" if isinstance(data, dict) and 'org_config' in data:\n",
|
| 341 |
+
" org_config.update(data['org_config'])\n",
|
| 342 |
+
"\n",
|
| 343 |
+
" done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n",
|
| 344 |
+
" step += 1\n",
|
| 345 |
+
"\n",
|
| 346 |
+
" if (i + 1) % 10 == 0:\n",
|
| 347 |
+
" print(f' {i+1}/{n_episodes} episodes — {len(examples)} examples')\n",
|
| 348 |
+
"\n",
|
| 349 |
+
" return examples\n",
|
| 350 |
+
"\n",
|
| 351 |
+
"\n",
|
| 352 |
+
"print(f'Generating {N_SFT_EPISODES} SFT demonstration episodes...')\n",
|
| 353 |
+
"sft_raw = generate_sft_dataset(ENV_URL, tokenizer, n_episodes=N_SFT_EPISODES)\n",
|
| 354 |
+
"sft_dataset = Dataset.from_list(sft_raw)\n",
|
| 355 |
+
"print(f'\\nSFT dataset : {len(sft_dataset)} examples (~{len(sft_dataset)//6} eps x 6 steps)')\n",
|
| 356 |
+
"print(f'Sample (first 300 chars):\\n{sft_raw[0][\"text\"][:300]}')"
|
| 357 |
+
]
|
| 358 |
+
},
|
| 359 |
+
{
|
| 360 |
+
"cell_type": "markdown",
|
| 361 |
+
"id": "cell-10-md",
|
| 362 |
+
"metadata": {},
|
| 363 |
+
"source": [
|
| 364 |
+
"## 8. SFT Training (~15 min on A100)"
|
| 365 |
+
]
|
| 366 |
+
},
|
| 367 |
+
{
|
| 368 |
+
"cell_type": "code",
|
| 369 |
+
"execution_count": null,
|
| 370 |
+
"id": "cell-10",
|
| 371 |
+
"metadata": {},
|
| 372 |
+
"outputs": [],
|
| 373 |
+
"source": [
|
| 374 |
+
"from trl import SFTTrainer, SFTConfig\n",
|
| 375 |
+
"\n",
|
| 376 |
+
"sft_cfg = SFTConfig(\n",
|
| 377 |
+
" dataset_text_field = 'text',\n",
|
| 378 |
+
" max_seq_length = 2048,\n",
|
| 379 |
+
" num_train_epochs = 2,\n",
|
| 380 |
+
" per_device_train_batch_size = 4,\n",
|
| 381 |
+
" gradient_accumulation_steps = 4,\n",
|
| 382 |
+
" learning_rate = 2e-4,\n",
|
| 383 |
+
" warmup_steps = 10,\n",
|
| 384 |
+
" output_dir = 'pm-ops-sft-warmup',\n",
|
| 385 |
+
" report_to = 'none',\n",
|
| 386 |
+
" logging_steps = 5,\n",
|
| 387 |
+
" save_strategy = 'no',\n",
|
| 388 |
+
" dataloader_num_workers = 0,\n",
|
| 389 |
+
")\n",
|
| 390 |
+
"\n",
|
| 391 |
+
"sft_steps = (\n",
|
| 392 |
+
" len(sft_dataset)\n",
|
| 393 |
+
" // (sft_cfg.per_device_train_batch_size * sft_cfg.gradient_accumulation_steps)\n",
|
| 394 |
+
" * sft_cfg.num_train_epochs\n",
|
| 395 |
+
")\n",
|
| 396 |
+
"print(f'SFT: {len(sft_dataset)} examples x {sft_cfg.num_train_epochs} epochs -> ~{sft_steps} steps')\n",
|
| 397 |
+
"\n",
|
| 398 |
+
"sft_trainer = SFTTrainer(\n",
|
| 399 |
+
" model=model, tokenizer=tokenizer,\n",
|
| 400 |
+
" train_dataset=sft_dataset, args=sft_cfg,\n",
|
| 401 |
+
")\n",
|
| 402 |
+
"sft_stats = sft_trainer.train()\n",
|
| 403 |
+
"\n",
|
| 404 |
+
"runtime = sft_stats.metrics.get('train_runtime', 0)\n",
|
| 405 |
+
"loss = sft_stats.metrics.get('train_loss', 0)\n",
|
| 406 |
+
"print(f'SFT done: {round(runtime/60, 1)} min, loss={loss:.3f}')"
|
| 407 |
+
]
|
| 408 |
+
},
|
| 409 |
+
{
|
| 410 |
+
"cell_type": "markdown",
|
| 411 |
+
"id": "cell-11-md",
|
| 412 |
+
"metadata": {},
|
| 413 |
+
"source": [
|
| 414 |
+
"## 9. Verify SFT Output — Model Must Output Valid JSON"
|
| 415 |
+
]
|
| 416 |
+
},
|
| 417 |
+
{
|
| 418 |
+
"cell_type": "code",
|
| 419 |
+
"execution_count": null,
|
| 420 |
+
"id": "cell-11",
|
| 421 |
+
"metadata": {},
|
| 422 |
+
"outputs": [],
|
| 423 |
+
"source": [
|
| 424 |
+
"from training.rollout import extract_json_action, _obs_to_dict, _current_obs_text, build_messages\n",
|
| 425 |
+
"\n",
|
| 426 |
+
"model.eval()\n",
|
| 427 |
+
"with GenericEnvClient(base_url=ENV_URL).sync() as _env:\n",
|
| 428 |
+
" r = _env.reset(seed=99001)\n",
|
| 429 |
+
" obs_dict = _obs_to_dict(r.observation if hasattr(r, 'observation') else r)\n",
|
| 430 |
+
" obs_text = _current_obs_text(obs_dict, 0, obs_dict.get('task_brief', ''))\n",
|
| 431 |
+
" msgs = build_messages([], obs_text)\n",
|
| 432 |
+
" prompt = tokenizer.apply_chat_template(\n",
|
| 433 |
+
" msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n",
|
| 434 |
+
" )\n",
|
| 435 |
+
"\n",
|
| 436 |
+
"inputs = tokenizer([prompt], return_tensors='pt').to(model.device)\n",
|
| 437 |
+
"with torch.no_grad():\n",
|
| 438 |
+
" out = model.generate(\n",
|
| 439 |
+
" **inputs, max_new_tokens=128, do_sample=False,\n",
|
| 440 |
+
" pad_token_id=tokenizer.eos_token_id\n",
|
| 441 |
+
" )\n",
|
| 442 |
+
"completion = tokenizer.decode(out[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)\n",
|
| 443 |
+
"parsed = extract_json_action(completion)\n",
|
| 444 |
+
"\n",
|
| 445 |
+
"print(f'Output : {completion[:400]}')\n",
|
| 446 |
+
"print(f'Parsed : {parsed}')\n",
|
| 447 |
+
"\n",
|
| 448 |
+
"if parsed is not None:\n",
|
| 449 |
+
" print('PASS: model outputs valid JSON after SFT')\n",
|
| 450 |
+
"else:\n",
|
| 451 |
+
" print('FAIL: still no valid JSON — run SFT cell again with more epochs or more data')\n",
|
| 452 |
+
"\n",
|
| 453 |
+
"model.train()"
|
| 454 |
+
]
|
| 455 |
+
},
|
| 456 |
+
{
|
| 457 |
+
"cell_type": "markdown",
|
| 458 |
+
"id": "cell-12-md",
|
| 459 |
+
"metadata": {},
|
| 460 |
+
"source": [
|
| 461 |
+
"---\n",
|
| 462 |
+
"## Phase 2 — GRPO\n",
|
| 463 |
+
"\n",
|
| 464 |
+
"Now that the model outputs valid JSON, GRPO can optimize for *correctness*.\n",
|
| 465 |
+
"\n",
|
| 466 |
+
"**Reward** (2 components, sum = 1.0):\n",
|
| 467 |
+
"\n",
|
| 468 |
+
"| Component | Weight | Signal |\n",
|
| 469 |
+
"|---|---|---|\n",
|
| 470 |
+
"| `env_score` | 0.85 | Env grader: 0.25 (ticket) + 0.20 (label) + 0.20 (priority) + 0.20 (team) + 0.15 (channel) |\n",
|
| 471 |
+
"| `json_ratio` | 0.15 | Fraction of steps with parseable JSON — maintains format quality |\n",
|
| 472 |
+
"\n",
|
| 473 |
+
"`env_score` naturally varies 0.25–1.0 per episode (the model may get the ticket right\n",
|
| 474 |
+
"but pick the wrong label, or get the channel wrong). This is the learning signal.\n",
|
| 475 |
+
"Anti-hacking: org_config values differ every episode (seeded), so the model cannot memorize answers."
|
| 476 |
+
]
|
| 477 |
+
},
|
| 478 |
+
{
|
| 479 |
+
"cell_type": "markdown",
|
| 480 |
+
"id": "cell-13-md",
|
| 481 |
+
"metadata": {},
|
| 482 |
+
"source": [
|
| 483 |
+
"## 10. Generate GRPO Training Dataset"
|
| 484 |
+
]
|
| 485 |
+
},
|
| 486 |
+
{
|
| 487 |
+
"cell_type": "code",
|
| 488 |
+
"execution_count": null,
|
| 489 |
+
"id": "cell-13",
|
| 490 |
+
"metadata": {},
|
| 491 |
+
"outputs": [],
|
| 492 |
+
"source": "from training.dataset import generate_triage_dataset\n\n# generate_triage_dataset now pre-simulates the env's RNG to only include seeds\n# where env.reset(seed) will actually run a TRIAGE episode — previously the env\n# silently ran release_notes/dep_update for the same seed, guaranteeing env_score=0.\nrows = generate_triage_dataset(n_episodes=N_GRPO_EPISODES, base_seed=42)\ngrpo_dataset = Dataset.from_list([{'prompt': r['prompt']} for r in rows])\nprint(f'GRPO dataset: {len(grpo_dataset)} triage episodes')\nprint(f'Difficulties: {set(r[\"difficulty\"] for r in rows)}')\nprint(f'Sample prompt: {rows[0][\"prompt\"][:120]}')"
|
| 493 |
+
},
|
| 494 |
+
{
|
| 495 |
+
"cell_type": "markdown",
|
| 496 |
+
"id": "cell-14-md",
|
| 497 |
+
"metadata": {},
|
| 498 |
+
"source": [
|
| 499 |
+
"## 11. GRPO Rollout + Reward Functions"
|
| 500 |
+
]
|
| 501 |
+
},
|
| 502 |
+
{
|
| 503 |
+
"cell_type": "code",
|
| 504 |
+
"execution_count": null,
|
| 505 |
+
"id": "cell-14",
|
| 506 |
+
"metadata": {},
|
| 507 |
+
"outputs": [],
|
| 508 |
+
"source": "from training.rollout import (\n _obs_to_dict, _current_obs_text, build_messages,\n extract_json_action, step_aware_fallback,\n _generate_no_vllm,\n)\nfrom training.dataset import parse_seed_from_prompt\nfrom training.rewards import compute_rollout_reward\n\ngrpo_env = GenericEnvClient(base_url=ENV_URL).sync()\ngrpo_env.connect()\nprint('GRPO training env connected')\n\nROLLOUT_TEMPERATURE = 1.1 # must be > 1.0 for rollout diversity\n\n\ndef run_grpo_episode(trainer, env, tok, dataset_prompt, max_steps=TRAIN_MAX_STEPS,\n gen_offset=0):\n \"\"\"Run one PM-ops triage episode and return trajectory + reward.\n\n Reward is runbook-compliance based (see training/rewards.py):\n - Did the model read the runbook?\n - Did it use a valid label/priority/team from the runbook?\n - Did it notify a correct oncall channel?\n\n This creates genuine reward variance across rollouts because org configs\n vary by seed — the same hardcoded label/team/channel is correct for some\n orgs and wrong for others, giving GRPO a real gradient signal.\n \"\"\"\n seed = parse_seed_from_prompt(dataset_prompt)\n if seed is not None:\n result = env.reset(seed=seed + gen_offset)\n else:\n result = env.reset()\n\n obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n task_brief = obs_dict.get('task_brief') or dataset_prompt\n\n prompt_ids, completion_ids, logprobs = [], [], []\n turn_history = []\n valid_json_count = 0\n env_score = 0.0\n step, done = 0, False\n _sample_logged = False\n\n # Runbook-compliance tracking\n read_runbook_done = False\n valid_labels: set = set()\n valid_priorities: set = set()\n valid_teams: set = set()\n oncall_channels: set = set()\n ticket_label: str | None = None\n ticket_priority: str | None = None\n assigned_team: str | None = None\n posted_channels: list = []\n\n while not done and step < max_steps:\n obs_text = _current_obs_text(obs_dict, step, task_brief)\n msgs = build_messages(turn_history, obs_text)\n prompt_text = tok.apply_chat_template(\n msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n )\n\n rollout_out = _generate_no_vllm(\n trainer, prompt_text, tok,\n max_new_tokens=MAX_COMP_LEN,\n temperature=ROLLOUT_TEMPERATURE,\n )\n prompt_ids.extend(rollout_out['prompt_ids'])\n completion_ids.extend(rollout_out['completion_ids'])\n logprobs.extend(rollout_out['logprobs'])\n completion_text = rollout_out['text']\n\n if not _sample_logged:\n print(f' [sample] {repr(completion_text[:180])}')\n _sample_logged = True\n\n parsed = extract_json_action(completion_text)\n if parsed is not None:\n valid_json_count += 1\n else:\n parsed = step_aware_fallback(step, max_steps)\n\n action_type = parsed.get('action_type', 'meta.noop')\n args = parsed.get('args', {})\n\n if action_type == 'meta.read_runbook' and parsed is not None:\n read_runbook_done = True\n\n if action_type == 'ticketing.create_ticket' and parsed is not None and ticket_label is None:\n ticket_label = args.get('label')\n ticket_priority = args.get('priority')\n\n if action_type == 'ticketing.assign_ticket' and parsed is not None and assigned_team is None:\n assigned_team = args.get('team')\n\n if action_type == 'chat.post_message' and parsed is not None:\n ch = args.get('channel', '')\n if ch:\n posted_channels.append(ch)\n\n turn_history.append({\n 'obs_text' : obs_text,\n 'completion': completion_text,\n 'is_runbook': (action_type == 'meta.read_runbook' and parsed is not None),\n })\n\n try:\n result = env.step({'action_type': action_type, 'args': args})\n except RuntimeError as exc:\n if 'VALIDATION_ERROR' in str(exc):\n result = env.step({'action_type': 'meta.noop', 'args': {}})\n else:\n raise\n\n obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n\n # Extract org_config from runbook response (available the step AFTER read_runbook)\n last_res = obs_dict.get('last_action_result') or {}\n if action_type == 'meta.read_runbook' and last_res.get('ok'):\n data = last_res.get('data') or {}\n if isinstance(data, dict):\n org = data.get('org_config') or {}\n valid_labels = set(org.get('label_taxonomy', {}).values())\n valid_priorities = set(org.get('priority_levels', []))\n valid_teams = set(org.get('team_map', {}).values())\n oncall_channels = set(org.get('oncall_channels', {}).values())\n\n done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n env_score = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n step += 1\n\n reward = compute_rollout_reward(\n read_runbook_done = read_runbook_done,\n valid_labels = valid_labels,\n valid_priorities = valid_priorities,\n valid_teams = valid_teams,\n oncall_channels = oncall_channels,\n ticket_label = ticket_label,\n ticket_priority = ticket_priority,\n assigned_team = assigned_team,\n posted_channels = posted_channels,\n env_score = env_score,\n valid_json_count = valid_json_count,\n )\n\n lbl_ok = '✓' if ticket_label and ticket_label in valid_labels else ('✗' if ticket_label else '-')\n pri_ok = '✓' if ticket_priority and ticket_priority in valid_priorities else ('✗' if ticket_priority else '-')\n tm_ok = '✓' if assigned_team and assigned_team in valid_teams else ('✗' if assigned_team else '-')\n ch_ok = '✓' if any(ch in oncall_channels for ch in posted_channels) else ('✗' if posted_channels else '-')\n print(f' [rollout] steps={step} env={env_score:.3f} '\n f'label={lbl_ok} pri={pri_ok} team={tm_ok} ch={ch_ok} '\n f'offset={gen_offset} -> reward={reward:.3f}')\n\n return {\n 'prompt_ids' : prompt_ids,\n 'completion_ids': completion_ids,\n 'logprobs' : logprobs,\n 'reward' : reward,\n }\n\n\ndef grpo_rollout_func(prompts, trainer=None):\n out = {'prompt_ids': [], 'completion_ids': [], 'logprobs': [], 'reward': []}\n prompt_seen: dict = {}\n for prompt in prompts:\n gen_offset = prompt_seen.get(prompt, 0)\n prompt_seen[prompt] = gen_offset + 1\n ep = run_grpo_episode(trainer, grpo_env, tokenizer, prompt,\n TRAIN_MAX_STEPS, gen_offset=gen_offset)\n for k in out:\n out[k].append(ep[k])\n return out\n\n\n_warned_missing_reward = False\ndef grpo_reward_func(completions, **kwargs):\n \"\"\"Passthrough — reward is pre-computed in grpo_rollout_func.\"\"\"\n global _warned_missing_reward\n rewards = kwargs.get('reward', [])\n if not rewards:\n if not _warned_missing_reward:\n print(f\"[WARN] reward_func fallback, no reward key. kwargs: {list(kwargs.keys())}\")\n _warned_missing_reward = True\n return [0.0] * len(completions)\n return [float(r) for r in rewards]\n\n\nprint(f'GRPO rollout ready max_steps={TRAIN_MAX_STEPS} temperature={ROLLOUT_TEMPERATURE}')"
|
| 509 |
+
},
|
| 510 |
+
{
|
| 511 |
+
"cell_type": "markdown",
|
| 512 |
+
"id": "cell-15-md",
|
| 513 |
+
"metadata": {},
|
| 514 |
+
"source": [
|
| 515 |
+
"## 12. GRPO Config + Trainer"
|
| 516 |
+
]
|
| 517 |
+
},
|
| 518 |
+
{
|
| 519 |
+
"cell_type": "code",
|
| 520 |
+
"execution_count": null,
|
| 521 |
+
"id": "cell-15",
|
| 522 |
+
"metadata": {},
|
| 523 |
+
"outputs": [],
|
| 524 |
+
"source": [
|
| 525 |
+
"from trl import GRPOConfig\n",
|
| 526 |
+
"from training.pm_ops_trainer import PMOpsGRPOTrainer\n",
|
| 527 |
+
"\n",
|
| 528 |
+
"OUTPUT_DIR = 'pm-ops-grpo-Qwen3-1.7B-triage-v3'\n",
|
| 529 |
+
"HF_REPO_ID = f'Saurav1/{OUTPUT_DIR}'\n",
|
| 530 |
+
"\n",
|
| 531 |
+
"grpo_cfg = GRPOConfig(\n",
|
| 532 |
+
" # Training\n",
|
| 533 |
+
" num_train_epochs = 2,\n",
|
| 534 |
+
" learning_rate = 1e-6, # lower LR: model has SFT init, don't overwrite it\n",
|
| 535 |
+
" gradient_accumulation_steps = GRAD_ACCUM,\n",
|
| 536 |
+
" per_device_train_batch_size = 1,\n",
|
| 537 |
+
" warmup_steps = 5,\n",
|
| 538 |
+
" num_generations = NUM_GEN,\n",
|
| 539 |
+
" # Sequence lengths\n",
|
| 540 |
+
" max_completion_length = MAX_COMP_LEN,\n",
|
| 541 |
+
" max_prompt_length = 4096,\n",
|
| 542 |
+
" # Keep disabled for rollout_func compatibility on cloud runtimes\n",
|
| 543 |
+
" use_vllm = False,\n",
|
| 544 |
+
" # Output\n",
|
| 545 |
+
" output_dir = OUTPUT_DIR,\n",
|
| 546 |
+
" report_to = 'trackio',\n",
|
| 547 |
+
" trackio_space_id = OUTPUT_DIR,\n",
|
| 548 |
+
" logging_steps = 1,\n",
|
| 549 |
+
" save_steps = 20,\n",
|
| 550 |
+
" gradient_checkpointing = False, # Unsloth handles this\n",
|
| 551 |
+
")\n",
|
| 552 |
+
"\n",
|
| 553 |
+
"eff_batch = grpo_cfg.per_device_train_batch_size * GRAD_ACCUM\n",
|
| 554 |
+
"total_steps = len(grpo_dataset) * NUM_GEN * grpo_cfg.num_train_epochs // eff_batch\n",
|
| 555 |
+
"print(f'GRPO: {len(grpo_dataset)} eps x {NUM_GEN} gen x {grpo_cfg.num_train_epochs} epochs -> ~{total_steps} steps')\n",
|
| 556 |
+
"\n",
|
| 557 |
+
"# PMOpsGRPOTrainer overrides _calculate_rewards to read the pre-computed\n",
|
| 558 |
+
"# 'reward' key directly from the rollout batch — bypasses broken kwargs plumbing.\n",
|
| 559 |
+
"trainer = PMOpsGRPOTrainer(\n",
|
| 560 |
+
" model = model,\n",
|
| 561 |
+
" processing_class = tokenizer,\n",
|
| 562 |
+
" reward_funcs = grpo_reward_func, # kept as fallback only\n",
|
| 563 |
+
" train_dataset = grpo_dataset,\n",
|
| 564 |
+
" args = grpo_cfg,\n",
|
| 565 |
+
" rollout_func = grpo_rollout_func,\n",
|
| 566 |
+
")\n",
|
| 567 |
+
"print(f'PMOpsGRPOTrainer ready: {type(trainer).__name__}')\n",
|
| 568 |
+
"assert isinstance(trainer, PMOpsGRPOTrainer), 'trainer must be PMOpsGRPOTrainer'\n",
|
| 569 |
+
"assert grpo_cfg.use_vllm is False, 'use_vllm must be False for rollout compatibility'"
|
| 570 |
+
]
|
| 571 |
+
},
|
| 572 |
+
{
|
| 573 |
+
"cell_type": "code",
|
| 574 |
+
"execution_count": null,
|
| 575 |
+
"id": "9a7c99b4",
|
| 576 |
+
"metadata": {},
|
| 577 |
+
"outputs": [],
|
| 578 |
+
"source": [
|
| 579 |
+
"# Preflight: ensure rollout returns reward and trainer received rollout_func\n",
|
| 580 |
+
"probe = grpo_rollout_func([grpo_dataset[0]['prompt']], trainer=trainer)\n",
|
| 581 |
+
"print('probe keys:', list(probe.keys()))\n",
|
| 582 |
+
"print('probe reward sample:', probe['reward'][:1])\n",
|
| 583 |
+
"assert len(probe['reward']) == 1, 'rollout probe did not return reward values'"
|
| 584 |
+
]
|
| 585 |
+
},
|
| 586 |
+
{
|
| 587 |
+
"cell_type": "markdown",
|
| 588 |
+
"id": "cell-16-md",
|
| 589 |
+
"metadata": {},
|
| 590 |
+
"source": [
|
| 591 |
+
"## 13. Train\n",
|
| 592 |
+
"\n",
|
| 593 |
+
"Watch stdout for:\n",
|
| 594 |
+
"- `[sample] '```json...'` — should look like valid JSON code blocks\n",
|
| 595 |
+
"- `[rollout] env=X.XXX` — **should trend upward over steps** (this is the signal)\n",
|
| 596 |
+
"- `[rollout] json=X.XX` — should stay > 0.7 (SFT maintains format quality)\n",
|
| 597 |
+
"\n",
|
| 598 |
+
"Watch trackio for the reward curve."
|
| 599 |
+
]
|
| 600 |
+
},
|
| 601 |
+
{
|
| 602 |
+
"cell_type": "code",
|
| 603 |
+
"execution_count": null,
|
| 604 |
+
"id": "cell-16",
|
| 605 |
+
"metadata": {},
|
| 606 |
+
"outputs": [],
|
| 607 |
+
"source": [
|
| 608 |
+
"trainer_stats = trainer.train()\n",
|
| 609 |
+
"\n",
|
| 610 |
+
"used_gb = round(torch.cuda.max_memory_reserved() / 1024**3, 2)\n",
|
| 611 |
+
"train_mins = round(trainer_stats.metrics.get('train_runtime', 0) / 60, 1)\n",
|
| 612 |
+
"print(f'Training time : {train_mins} min')\n",
|
| 613 |
+
"print(f'Peak GPU : {used_gb} GB / {TOTAL_GB} GB ({round(used_gb/TOTAL_GB*100, 1)}%)')\n",
|
| 614 |
+
"final_reward = trainer_stats.metrics.get('train/reward', trainer_stats.metrics.get('train_loss', '?'))\n",
|
| 615 |
+
"print(f'Final reward : {final_reward}')"
|
| 616 |
+
]
|
| 617 |
+
},
|
| 618 |
+
{
|
| 619 |
+
"cell_type": "markdown",
|
| 620 |
+
"id": "cell-17-md",
|
| 621 |
+
"metadata": {},
|
| 622 |
+
"source": [
|
| 623 |
+
"## 14. Save Model"
|
| 624 |
+
]
|
| 625 |
+
},
|
| 626 |
+
{
|
| 627 |
+
"cell_type": "code",
|
| 628 |
+
"execution_count": null,
|
| 629 |
+
"id": "cell-17",
|
| 630 |
+
"metadata": {},
|
| 631 |
+
"outputs": [],
|
| 632 |
+
"source": [
|
| 633 |
+
"grpo_env.close()\n",
|
| 634 |
+
"\n",
|
| 635 |
+
"# Unsloth merged save — dequantises first, then merges LoRA cleanly into bf16.\n",
|
| 636 |
+
"# Do NOT use trainer.save_model() directly on a 4-bit + LoRA model.\n",
|
| 637 |
+
"model.save_pretrained_merged(OUTPUT_DIR, tokenizer, save_method='merged_16bit')\n",
|
| 638 |
+
"model.push_to_hub_merged(HF_REPO_ID, tokenizer, save_method='merged_16bit')\n",
|
| 639 |
+
"print(f'Pushed -> https://huggingface.co/{HF_REPO_ID}')"
|
| 640 |
+
]
|
| 641 |
+
},
|
| 642 |
+
{
|
| 643 |
+
"cell_type": "markdown",
|
| 644 |
+
"id": "cell-18-md",
|
| 645 |
+
"metadata": {},
|
| 646 |
+
"source": [
|
| 647 |
+
"## 15. Evaluate: Baseline vs Trained"
|
| 648 |
+
]
|
| 649 |
+
},
|
| 650 |
+
{
|
| 651 |
+
"cell_type": "code",
|
| 652 |
+
"execution_count": null,
|
| 653 |
+
"id": "cell-18",
|
| 654 |
+
"metadata": {},
|
| 655 |
+
"outputs": [],
|
| 656 |
+
"source": [
|
| 657 |
+
"from transformers import AutoModelForCausalLM\n",
|
| 658 |
+
"\n",
|
| 659 |
+
"N_EVAL = 15\n",
|
| 660 |
+
"EVAL_MAX_STEPS = 12\n",
|
| 661 |
+
"EVAL_SEED_BASE = 9000\n",
|
| 662 |
+
"\n",
|
| 663 |
+
"eval_model = AutoModelForCausalLM.from_pretrained(\n",
|
| 664 |
+
" OUTPUT_DIR, torch_dtype=torch.bfloat16, device_map='auto'\n",
|
| 665 |
+
")\n",
|
| 666 |
+
"eval_model.eval()\n",
|
| 667 |
+
"\n",
|
| 668 |
+
"\n",
|
| 669 |
+
"def eval_trained(n=N_EVAL):\n",
|
| 670 |
+
" scores = []\n",
|
| 671 |
+
" with GenericEnvClient(base_url=ENV_URL).sync() as env:\n",
|
| 672 |
+
" for i in range(n):\n",
|
| 673 |
+
" result = env.reset(seed=EVAL_SEED_BASE + i)\n",
|
| 674 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 675 |
+
" task_brief = obs_dict.get('task_brief', '')\n",
|
| 676 |
+
" history, step, score, done = [], 0, 0.0, False\n",
|
| 677 |
+
"\n",
|
| 678 |
+
" while not done and step < EVAL_MAX_STEPS:\n",
|
| 679 |
+
" obs_text = _current_obs_text(obs_dict, step, task_brief)\n",
|
| 680 |
+
" msgs = build_messages(history, obs_text)\n",
|
| 681 |
+
" prompt = tokenizer.apply_chat_template(\n",
|
| 682 |
+
" msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n",
|
| 683 |
+
" )\n",
|
| 684 |
+
" inputs = tokenizer([prompt], return_tensors='pt', truncation=True, max_length=4096)\n",
|
| 685 |
+
" inputs = {k: v.to(eval_model.device) for k, v in inputs.items()}\n",
|
| 686 |
+
" with torch.no_grad():\n",
|
| 687 |
+
" out_ids = eval_model.generate(\n",
|
| 688 |
+
" **inputs, max_new_tokens=256, do_sample=False,\n",
|
| 689 |
+
" pad_token_id=tokenizer.eos_token_id\n",
|
| 690 |
+
" )\n",
|
| 691 |
+
" completion = tokenizer.decode(\n",
|
| 692 |
+
" out_ids[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True\n",
|
| 693 |
+
" )\n",
|
| 694 |
+
" parsed = extract_json_action(completion) or step_aware_fallback(step, EVAL_MAX_STEPS)\n",
|
| 695 |
+
" result = env.step({'action_type': parsed['action_type'], 'args': parsed.get('args', {})})\n",
|
| 696 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 697 |
+
" done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n",
|
| 698 |
+
" score = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n",
|
| 699 |
+
" history.append({'obs_text': obs_text, 'completion': completion, 'is_runbook': False})\n",
|
| 700 |
+
" step += 1\n",
|
| 701 |
+
"\n",
|
| 702 |
+
" scores.append(score)\n",
|
| 703 |
+
" print(f' Trained ep {i+1}/{n}: score={score:.3f}')\n",
|
| 704 |
+
" return scores\n",
|
| 705 |
+
"\n",
|
| 706 |
+
"\n",
|
| 707 |
+
"def eval_baseline(n=N_EVAL):\n",
|
| 708 |
+
" from inference import baseline_agent\n",
|
| 709 |
+
" scores = []\n",
|
| 710 |
+
" with GenericEnvClient(base_url=ENV_URL).sync() as env:\n",
|
| 711 |
+
" for i in range(n):\n",
|
| 712 |
+
" result = env.reset(seed=EVAL_SEED_BASE + i)\n",
|
| 713 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 714 |
+
" org_config, step, score, done = {}, 0, 0.0, False\n",
|
| 715 |
+
" while not done and step < EVAL_MAX_STEPS:\n",
|
| 716 |
+
" at, args = baseline_agent(obs_dict, org_config)\n",
|
| 717 |
+
" result = env.step({'action_type': at, 'args': args})\n",
|
| 718 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 719 |
+
" done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n",
|
| 720 |
+
" score = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n",
|
| 721 |
+
" step += 1\n",
|
| 722 |
+
" scores.append(score)\n",
|
| 723 |
+
" print(f' Baseline ep {i+1}/{n}: score={score:.3f}')\n",
|
| 724 |
+
" return scores\n",
|
| 725 |
+
"\n",
|
| 726 |
+
"\n",
|
| 727 |
+
"print('--- Baseline ---')\n",
|
| 728 |
+
"baseline_scores = eval_baseline()\n",
|
| 729 |
+
"print('\\n--- Trained ---')\n",
|
| 730 |
+
"trained_scores = eval_trained()\n",
|
| 731 |
+
"\n",
|
| 732 |
+
"print(f'\\nBaseline avg : {sum(baseline_scores)/N_EVAL:.3f}')\n",
|
| 733 |
+
"print(f'Trained avg : {sum(trained_scores)/N_EVAL:.3f}')\n",
|
| 734 |
+
"print(f'Delta : {(sum(trained_scores)-sum(baseline_scores))/N_EVAL:+.3f}')"
|
| 735 |
+
]
|
| 736 |
+
},
|
| 737 |
+
{
|
| 738 |
+
"cell_type": "markdown",
|
| 739 |
+
"id": "cell-19-md",
|
| 740 |
+
"metadata": {},
|
| 741 |
+
"source": [
|
| 742 |
+
"## 16. Plot Results"
|
| 743 |
+
]
|
| 744 |
+
},
|
| 745 |
+
{
|
| 746 |
+
"cell_type": "code",
|
| 747 |
+
"execution_count": null,
|
| 748 |
+
"id": "cell-19",
|
| 749 |
+
"metadata": {},
|
| 750 |
+
"outputs": [],
|
| 751 |
+
"source": [
|
| 752 |
+
"import matplotlib.pyplot as plt\n",
|
| 753 |
+
"import numpy as np\n",
|
| 754 |
+
"\n",
|
| 755 |
+
"fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n",
|
| 756 |
+
"\n",
|
| 757 |
+
"ax, x, w = axes[0], np.arange(N_EVAL), 0.35\n",
|
| 758 |
+
"ax.bar(x - w/2, baseline_scores, w, label='Baseline', color='steelblue', alpha=0.8)\n",
|
| 759 |
+
"ax.bar(x + w/2, trained_scores, w, label='GRPO v3', color='coral', alpha=0.8)\n",
|
| 760 |
+
"ax.axhline(sum(baseline_scores)/N_EVAL, color='steelblue', ls='--', alpha=0.5, lw=1.5)\n",
|
| 761 |
+
"ax.axhline(sum(trained_scores)/N_EVAL, color='coral', ls='--', alpha=0.5, lw=1.5)\n",
|
| 762 |
+
"ax.set(xlabel='Eval episode', ylabel='Reward (0-1)',\n",
|
| 763 |
+
" title='Per-episode reward: Baseline vs GRPO v3', xticks=x, ylim=(0, 1.05))\n",
|
| 764 |
+
"ax.legend()\n",
|
| 765 |
+
"\n",
|
| 766 |
+
"ax2 = axes[1]\n",
|
| 767 |
+
"avgs = [sum(baseline_scores)/N_EVAL, sum(trained_scores)/N_EVAL]\n",
|
| 768 |
+
"bars = ax2.bar(['Baseline', 'GRPO v3'], avgs, color=['steelblue', 'coral'], alpha=0.85, width=0.5)\n",
|
| 769 |
+
"for bar, val in zip(bars, avgs):\n",
|
| 770 |
+
" ax2.text(bar.get_x() + bar.get_width()/2, val + 0.01, f'{val:.3f}',\n",
|
| 771 |
+
" ha='center', fontsize=13, fontweight='bold')\n",
|
| 772 |
+
"ax2.set(ylabel='Average reward (0-1)', title=f'Average over {N_EVAL} triage episodes',\n",
|
| 773 |
+
" ylim=(0, 1.05))\n",
|
| 774 |
+
"\n",
|
| 775 |
+
"plt.tight_layout()\n",
|
| 776 |
+
"plt.savefig('eval_results_v3.png', dpi=150, bbox_inches='tight')\n",
|
| 777 |
+
"plt.show()\n",
|
| 778 |
+
"print('Saved: eval_results_v3.png')"
|
| 779 |
+
]
|
| 780 |
+
},
|
| 781 |
+
{
|
| 782 |
+
"cell_type": "markdown",
|
| 783 |
+
"id": "cell-20-md",
|
| 784 |
+
"metadata": {},
|
| 785 |
+
"source": [
|
| 786 |
+
"## 17. Teardown"
|
| 787 |
+
]
|
| 788 |
+
},
|
| 789 |
+
{
|
| 790 |
+
"cell_type": "code",
|
| 791 |
+
"execution_count": null,
|
| 792 |
+
"id": "cell-20",
|
| 793 |
+
"metadata": {},
|
| 794 |
+
"outputs": [],
|
| 795 |
+
"source": [
|
| 796 |
+
"server_proc.terminate()\n",
|
| 797 |
+
"print('Local PM-Ops server stopped')"
|
| 798 |
+
]
|
| 799 |
+
}
|
| 800 |
+
],
|
| 801 |
+
"metadata": {
|
| 802 |
+
"kernelspec": {
|
| 803 |
+
"display_name": "Python 3",
|
| 804 |
+
"language": "python",
|
| 805 |
+
"name": "python3"
|
| 806 |
+
},
|
| 807 |
+
"language_info": {
|
| 808 |
+
"name": "python",
|
| 809 |
+
"version": "3.11.0"
|
| 810 |
+
}
|
| 811 |
+
},
|
| 812 |
+
"nbformat": 4,
|
| 813 |
+
"nbformat_minor": 5
|
| 814 |
+
}
|
training/train_v4.ipynb
ADDED
|
@@ -0,0 +1,847 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"cells": [
|
| 3 |
+
{
|
| 4 |
+
"cell_type": "markdown",
|
| 5 |
+
"id": "v4-title",
|
| 6 |
+
"metadata": {},
|
| 7 |
+
"source": [
|
| 8 |
+
"# PM-Ops GRPO Training v4\n",
|
| 9 |
+
"\n",
|
| 10 |
+
"Clean rewrite — all known v3 bugs fixed.\n",
|
| 11 |
+
"\n",
|
| 12 |
+
"| Fix | Details |\n",
|
| 13 |
+
"|---|---|\n",
|
| 14 |
+
"| Constant reward=0.350 | Replaced env_score formula with runbook-compliance scoring |\n",
|
| 15 |
+
"| `gen_slot=0` override bug | gen_slot now correctly offsets env seed per GRPO generation |\n",
|
| 16 |
+
"| Double `env.step` per step | Removed duplicate step call inside try/except |\n",
|
| 17 |
+
"| Dataset/env task-type mismatch | `dataset.py` pre-filters seeds to triage-only episodes |\n",
|
| 18 |
+
"| Near-greedy generation | Temperature 1.1 + top_k=50 for diverse rollouts |\n",
|
| 19 |
+
"| PatchFastRL called twice | Called once, unconditionally, before any trl imports |\n",
|
| 20 |
+
"\n",
|
| 21 |
+
"**Reward**: read_runbook (+0.10) + valid_label (±0.20/0.10) + valid_priority (±0.15/0.10) + valid_team (±0.20/0.10) + right_channel (±0.25/0.10) + env_bonus (×0.10). Varies per org config — different valid values per seed. \n",
|
| 22 |
+
"**Stack**: Unsloth 2026.x · TRL 0.22.2 · Qwen3-1.7B · PMOpsGRPOTrainer \n",
|
| 23 |
+
"**GPU**: A100-40GB → ~2 hrs (SFT 15 min + GRPO 90 min)"
|
| 24 |
+
]
|
| 25 |
+
},
|
| 26 |
+
{
|
| 27 |
+
"cell_type": "markdown",
|
| 28 |
+
"id": "v4-s0",
|
| 29 |
+
"metadata": {},
|
| 30 |
+
"source": [
|
| 31 |
+
"## 0. Install"
|
| 32 |
+
]
|
| 33 |
+
},
|
| 34 |
+
{
|
| 35 |
+
"cell_type": "code",
|
| 36 |
+
"execution_count": null,
|
| 37 |
+
"id": "v4-install",
|
| 38 |
+
"metadata": {},
|
| 39 |
+
"outputs": [],
|
| 40 |
+
"source": [
|
| 41 |
+
"%%capture\n",
|
| 42 |
+
"import os\n",
|
| 43 |
+
"!pip install --upgrade -qqq uv\n",
|
| 44 |
+
"if 'COLAB_' not in ''.join(os.environ.keys()):\n",
|
| 45 |
+
" !uv pip install unsloth vllm\n",
|
| 46 |
+
"else:\n",
|
| 47 |
+
" import subprocess\n",
|
| 48 |
+
" is_t4 = 'Tesla T4' in str(subprocess.check_output(['nvidia-smi']))\n",
|
| 49 |
+
" _vllm = 'vllm==0.9.2' if is_t4 else 'vllm==0.15.1'\n",
|
| 50 |
+
" _triton = 'triton==3.2.0' if is_t4 else 'triton'\n",
|
| 51 |
+
" !uv pip install -qqq --upgrade {_vllm} torchvision bitsandbytes xformers unsloth\n",
|
| 52 |
+
" !uv pip install -qqq {_triton}\n",
|
| 53 |
+
"!uv pip install transformers==4.56.2\n",
|
| 54 |
+
"!uv pip install --no-deps trl==0.22.2\n",
|
| 55 |
+
"!pip install 'numpy==1.26.4' --break-system-packages -q\n",
|
| 56 |
+
"!pip install openenv openenv-core -q"
|
| 57 |
+
]
|
| 58 |
+
},
|
| 59 |
+
{
|
| 60 |
+
"cell_type": "markdown",
|
| 61 |
+
"id": "v4-s1",
|
| 62 |
+
"metadata": {},
|
| 63 |
+
"source": [
|
| 64 |
+
"## 1. GPU Config + Patch"
|
| 65 |
+
]
|
| 66 |
+
},
|
| 67 |
+
{
|
| 68 |
+
"cell_type": "code",
|
| 69 |
+
"execution_count": null,
|
| 70 |
+
"id": "v4-config",
|
| 71 |
+
"metadata": {},
|
| 72 |
+
"outputs": [],
|
| 73 |
+
"source": "import torch\nfrom unsloth import FastLanguageModel, PatchFastRL\n\n# Must patch BEFORE any trl imports\nPatchFastRL('GRPO', FastLanguageModel)\n\nimport trl\n\ngpu = torch.cuda.get_device_properties(0)\nTOTAL_GB = round(gpu.total_memory / 1024**3, 1)\nIS_A100 = TOTAL_GB >= 35\n\nNUM_GEN = 6 if IS_A100 else 2\nGRAD_ACCUM = 32 if IS_A100 else 8\nMAX_COMP_LEN = 192\n# T4 needs more SFT to reach low enough loss before GRPO.\n# 20 eps → loss≈2.0 is too high; 50 eps → loss≈0.8 is a better floor.\nN_SFT_EPISODES = 120 if IS_A100 else 50\nN_GRPO_EPISODES = 150 if IS_A100 else 30\nTRAIN_MAX_STEPS = 12\n\nprint(f'torch={torch.__version__} trl={trl.__version__}')\nprint(f'GPU: {gpu.name} ({TOTAL_GB} GB) IS_A100={IS_A100}')\nprint(f'num_gen={NUM_GEN} grad_accum={GRAD_ACCUM} '\n f'sft_eps={N_SFT_EPISODES} grpo_eps={N_GRPO_EPISODES} '\n f'max_steps={TRAIN_MAX_STEPS} max_comp_len={MAX_COMP_LEN}')"
|
| 74 |
+
},
|
| 75 |
+
{
|
| 76 |
+
"cell_type": "markdown",
|
| 77 |
+
"id": "v4-s2",
|
| 78 |
+
"metadata": {},
|
| 79 |
+
"source": [
|
| 80 |
+
"## 2. Clone PM-Ops Repo"
|
| 81 |
+
]
|
| 82 |
+
},
|
| 83 |
+
{
|
| 84 |
+
"cell_type": "code",
|
| 85 |
+
"execution_count": null,
|
| 86 |
+
"id": "v4-clone",
|
| 87 |
+
"metadata": {},
|
| 88 |
+
"outputs": [],
|
| 89 |
+
"source": [
|
| 90 |
+
"import os, sys\n",
|
| 91 |
+
"\n",
|
| 92 |
+
"REPO_URL = 'https://huggingface.co/spaces/TheCrustaceans/Pm-ops'\n",
|
| 93 |
+
"REPO_DIR = '/content/Pm_ops'\n",
|
| 94 |
+
"\n",
|
| 95 |
+
"if not os.path.exists(REPO_DIR):\n",
|
| 96 |
+
" !git clone --depth=1 -q {REPO_URL} {REPO_DIR}\n",
|
| 97 |
+
" print(f'Cloned -> {REPO_DIR}')\n",
|
| 98 |
+
"else:\n",
|
| 99 |
+
" !git -C {REPO_DIR} pull -q origin main\n",
|
| 100 |
+
" print(f'Pulled -> {REPO_DIR}')\n",
|
| 101 |
+
"\n",
|
| 102 |
+
"for p in [REPO_DIR, os.path.join(REPO_DIR, 'training')]:\n",
|
| 103 |
+
" if p not in sys.path:\n",
|
| 104 |
+
" sys.path.insert(0, p)\n",
|
| 105 |
+
"os.chdir(REPO_DIR)\n",
|
| 106 |
+
"print(f'CWD: {os.getcwd()}')\n",
|
| 107 |
+
"!git -C {REPO_DIR} log --oneline -3"
|
| 108 |
+
]
|
| 109 |
+
},
|
| 110 |
+
{
|
| 111 |
+
"cell_type": "markdown",
|
| 112 |
+
"id": "v4-s3",
|
| 113 |
+
"metadata": {},
|
| 114 |
+
"source": [
|
| 115 |
+
"## 3. HuggingFace Login"
|
| 116 |
+
]
|
| 117 |
+
},
|
| 118 |
+
{
|
| 119 |
+
"cell_type": "code",
|
| 120 |
+
"execution_count": null,
|
| 121 |
+
"id": "v4-hf-login",
|
| 122 |
+
"metadata": {},
|
| 123 |
+
"outputs": [],
|
| 124 |
+
"source": [
|
| 125 |
+
"from huggingface_hub import notebook_login\n",
|
| 126 |
+
"notebook_login()"
|
| 127 |
+
]
|
| 128 |
+
},
|
| 129 |
+
{
|
| 130 |
+
"cell_type": "markdown",
|
| 131 |
+
"id": "v4-s4",
|
| 132 |
+
"metadata": {},
|
| 133 |
+
"source": [
|
| 134 |
+
"## 4. Start PM-Ops Server"
|
| 135 |
+
]
|
| 136 |
+
},
|
| 137 |
+
{
|
| 138 |
+
"cell_type": "code",
|
| 139 |
+
"execution_count": null,
|
| 140 |
+
"id": "v4-server",
|
| 141 |
+
"metadata": {},
|
| 142 |
+
"outputs": [],
|
| 143 |
+
"source": [
|
| 144 |
+
"import subprocess, time, requests\n",
|
| 145 |
+
"\n",
|
| 146 |
+
"# Kill any leftover server from a previous run\n",
|
| 147 |
+
"subprocess.run(['pkill', '-f', 'uvicorn'], capture_output=True)\n",
|
| 148 |
+
"time.sleep(1)\n",
|
| 149 |
+
"\n",
|
| 150 |
+
"server_proc = subprocess.Popen(\n",
|
| 151 |
+
" [sys.executable, '-m', 'uvicorn', 'server.app:app',\n",
|
| 152 |
+
" '--host', '0.0.0.0', '--port', '8000'],\n",
|
| 153 |
+
" cwd=REPO_DIR, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,\n",
|
| 154 |
+
")\n",
|
| 155 |
+
"ENV_URL = 'http://localhost:8000'\n",
|
| 156 |
+
"\n",
|
| 157 |
+
"for _ in range(30):\n",
|
| 158 |
+
" try:\n",
|
| 159 |
+
" if requests.get(f'{ENV_URL}/', timeout=2).status_code == 200:\n",
|
| 160 |
+
" print(f'PM-Ops server ready pid={server_proc.pid}')\n",
|
| 161 |
+
" break\n",
|
| 162 |
+
" except Exception:\n",
|
| 163 |
+
" pass\n",
|
| 164 |
+
" time.sleep(1)\n",
|
| 165 |
+
"else:\n",
|
| 166 |
+
" raise RuntimeError('Server did not start in 30 s')"
|
| 167 |
+
]
|
| 168 |
+
},
|
| 169 |
+
{
|
| 170 |
+
"cell_type": "markdown",
|
| 171 |
+
"id": "v4-s5",
|
| 172 |
+
"metadata": {},
|
| 173 |
+
"source": [
|
| 174 |
+
"## 5. Verify Env"
|
| 175 |
+
]
|
| 176 |
+
},
|
| 177 |
+
{
|
| 178 |
+
"cell_type": "code",
|
| 179 |
+
"execution_count": null,
|
| 180 |
+
"id": "v4-verify-env",
|
| 181 |
+
"metadata": {},
|
| 182 |
+
"outputs": [],
|
| 183 |
+
"source": [
|
| 184 |
+
"from openenv.core import GenericEnvClient\n",
|
| 185 |
+
"from training.rollout import _obs_to_dict\n",
|
| 186 |
+
"\n",
|
| 187 |
+
"with GenericEnvClient(base_url=ENV_URL).sync() as _env:\n",
|
| 188 |
+
" r = _env.reset(seed=42)\n",
|
| 189 |
+
" obs = _obs_to_dict(r.observation if hasattr(r, 'observation') else r)\n",
|
| 190 |
+
" print(f'task_brief : {obs.get(\"task_brief\", \"?\")[:80]}...')\n",
|
| 191 |
+
" step_r = _env.step({'action_type': 'meta.read_runbook', 'args': {}})\n",
|
| 192 |
+
" rb_obs = _obs_to_dict(step_r.observation if hasattr(step_r, 'observation') else step_r)\n",
|
| 193 |
+
" data = (rb_obs.get('last_action_result') or {}).get('data', {})\n",
|
| 194 |
+
" org = data.get('org_config', {}) if isinstance(data, dict) else {}\n",
|
| 195 |
+
" print(f'org labels : {list(org.get(\"label_taxonomy\", {}).values())}')\n",
|
| 196 |
+
" print(f'org priorities: {org.get(\"priority_levels\", [])}')\n",
|
| 197 |
+
" print(f'org channels : {list(org.get(\"oncall_channels\", {}).values())}')\n",
|
| 198 |
+
" print('Env verify: OK')"
|
| 199 |
+
]
|
| 200 |
+
},
|
| 201 |
+
{
|
| 202 |
+
"cell_type": "markdown",
|
| 203 |
+
"id": "v4-s6",
|
| 204 |
+
"metadata": {},
|
| 205 |
+
"source": [
|
| 206 |
+
"## 6. Load Model — Unsloth 4-bit + LoRA"
|
| 207 |
+
]
|
| 208 |
+
},
|
| 209 |
+
{
|
| 210 |
+
"cell_type": "code",
|
| 211 |
+
"execution_count": null,
|
| 212 |
+
"id": "v4-model",
|
| 213 |
+
"metadata": {},
|
| 214 |
+
"outputs": [],
|
| 215 |
+
"source": [
|
| 216 |
+
"MODEL_NAME = 'Qwen/Qwen3-1.7B'\n",
|
| 217 |
+
"LORA_RANK = 16\n",
|
| 218 |
+
"\n",
|
| 219 |
+
"model, tokenizer = FastLanguageModel.from_pretrained(\n",
|
| 220 |
+
" model_name = MODEL_NAME,\n",
|
| 221 |
+
" max_seq_length = 4096 + MAX_COMP_LEN,\n",
|
| 222 |
+
" load_in_4bit = True,\n",
|
| 223 |
+
" fast_inference = False,\n",
|
| 224 |
+
" max_lora_rank = LORA_RANK,\n",
|
| 225 |
+
" gpu_memory_utilization = 0.55,\n",
|
| 226 |
+
")\n",
|
| 227 |
+
"model = FastLanguageModel.get_peft_model(\n",
|
| 228 |
+
" model,\n",
|
| 229 |
+
" r = LORA_RANK,\n",
|
| 230 |
+
" target_modules = ['q_proj','k_proj','v_proj','o_proj',\n",
|
| 231 |
+
" 'gate_proj','up_proj','down_proj'],\n",
|
| 232 |
+
" lora_alpha = LORA_RANK,\n",
|
| 233 |
+
" use_gradient_checkpointing = 'unsloth',\n",
|
| 234 |
+
" random_state = 42,\n",
|
| 235 |
+
")\n",
|
| 236 |
+
"tokenizer.pad_token = tokenizer.eos_token\n",
|
| 237 |
+
"tokenizer.padding_side = 'left'\n",
|
| 238 |
+
"model.print_trainable_parameters()\n",
|
| 239 |
+
"print(f'GPU after load: {round(torch.cuda.max_memory_reserved()/1024**3, 2)} GB / {TOTAL_GB} GB')"
|
| 240 |
+
]
|
| 241 |
+
},
|
| 242 |
+
{
|
| 243 |
+
"cell_type": "markdown",
|
| 244 |
+
"id": "v4-phase1",
|
| 245 |
+
"metadata": {},
|
| 246 |
+
"source": [
|
| 247 |
+
"---\n",
|
| 248 |
+
"## Phase 1 — SFT Warmup\n",
|
| 249 |
+
"\n",
|
| 250 |
+
"Teach the model the JSON output format and PM-ops workflow before GRPO. \n",
|
| 251 |
+
"`baseline_agent` runs N deterministic episodes → ~6 steps each → supervised (prompt, completion) pairs. \n",
|
| 252 |
+
"2 SFT epochs (~15 min on A100). Without SFT, all rollouts output freeform text → zero GRPO gradient."
|
| 253 |
+
]
|
| 254 |
+
},
|
| 255 |
+
{
|
| 256 |
+
"cell_type": "markdown",
|
| 257 |
+
"id": "v4-s7",
|
| 258 |
+
"metadata": {},
|
| 259 |
+
"source": [
|
| 260 |
+
"## 7. Generate SFT Dataset"
|
| 261 |
+
]
|
| 262 |
+
},
|
| 263 |
+
{
|
| 264 |
+
"cell_type": "code",
|
| 265 |
+
"execution_count": null,
|
| 266 |
+
"id": "v4-sft-data",
|
| 267 |
+
"metadata": {},
|
| 268 |
+
"outputs": [],
|
| 269 |
+
"source": [
|
| 270 |
+
"import json as _json\n",
|
| 271 |
+
"from datasets import Dataset\n",
|
| 272 |
+
"from inference import baseline_agent\n",
|
| 273 |
+
"from training.rollout import _obs_to_dict, _current_obs_text, build_messages\n",
|
| 274 |
+
"\n",
|
| 275 |
+
"\n",
|
| 276 |
+
"def generate_sft_dataset(env_url, tok, n_episodes, seed_start=2000):\n",
|
| 277 |
+
" examples = []\n",
|
| 278 |
+
" with GenericEnvClient(base_url=env_url).sync() as env:\n",
|
| 279 |
+
" for i in range(n_episodes):\n",
|
| 280 |
+
" result = env.reset(seed=seed_start + i)\n",
|
| 281 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 282 |
+
" task_brief = obs_dict.get('task_brief', '')\n",
|
| 283 |
+
" turn_history, org_config, step, done = [], {}, 0, False\n",
|
| 284 |
+
"\n",
|
| 285 |
+
" while not done and step < 8:\n",
|
| 286 |
+
" obs_text = _current_obs_text(obs_dict, step, task_brief)\n",
|
| 287 |
+
" action_type, args = baseline_agent(obs_dict, org_config)\n",
|
| 288 |
+
" payload = {'action_type': action_type, 'args': args}\n",
|
| 289 |
+
" completion = '```json\\n' + _json.dumps(payload) + '\\n```'\n",
|
| 290 |
+
" msgs = build_messages(turn_history, obs_text)\n",
|
| 291 |
+
" prompt = tok.apply_chat_template(\n",
|
| 292 |
+
" msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n",
|
| 293 |
+
" )\n",
|
| 294 |
+
" examples.append({'text': prompt + completion + tok.eos_token})\n",
|
| 295 |
+
" turn_history.append({\n",
|
| 296 |
+
" 'obs_text': obs_text, 'completion': completion,\n",
|
| 297 |
+
" 'is_runbook': (action_type == 'meta.read_runbook'),\n",
|
| 298 |
+
" })\n",
|
| 299 |
+
" result = env.step({'action_type': action_type, 'args': args})\n",
|
| 300 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 301 |
+
" if action_type == 'meta.read_runbook':\n",
|
| 302 |
+
" last = obs_dict.get('last_action_result') or {}\n",
|
| 303 |
+
" if last.get('ok'):\n",
|
| 304 |
+
" data = last.get('data') or {}\n",
|
| 305 |
+
" if isinstance(data, dict) and 'org_config' in data:\n",
|
| 306 |
+
" org_config.update(data['org_config'])\n",
|
| 307 |
+
" done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n",
|
| 308 |
+
" step += 1\n",
|
| 309 |
+
"\n",
|
| 310 |
+
" if (i + 1) % 20 == 0:\n",
|
| 311 |
+
" print(f' {i+1}/{n_episodes} eps — {len(examples)} examples')\n",
|
| 312 |
+
"\n",
|
| 313 |
+
" return examples\n",
|
| 314 |
+
"\n",
|
| 315 |
+
"\n",
|
| 316 |
+
"print(f'Generating {N_SFT_EPISODES} SFT demonstration episodes...')\n",
|
| 317 |
+
"sft_raw = generate_sft_dataset(ENV_URL, tokenizer, n_episodes=N_SFT_EPISODES)\n",
|
| 318 |
+
"sft_dataset = Dataset.from_list(sft_raw)\n",
|
| 319 |
+
"print(f'SFT dataset: {len(sft_dataset)} examples (~{len(sft_dataset)//6} eps × 6 steps)')\n",
|
| 320 |
+
"print(f'Sample (first 300 chars):\\n{sft_raw[0][\"text\"][:300]}')"
|
| 321 |
+
]
|
| 322 |
+
},
|
| 323 |
+
{
|
| 324 |
+
"cell_type": "markdown",
|
| 325 |
+
"id": "v4-s8",
|
| 326 |
+
"metadata": {},
|
| 327 |
+
"source": [
|
| 328 |
+
"## 8. SFT Training"
|
| 329 |
+
]
|
| 330 |
+
},
|
| 331 |
+
{
|
| 332 |
+
"cell_type": "code",
|
| 333 |
+
"execution_count": null,
|
| 334 |
+
"id": "v4-sft-train",
|
| 335 |
+
"metadata": {},
|
| 336 |
+
"outputs": [],
|
| 337 |
+
"source": [
|
| 338 |
+
"from trl import SFTTrainer, SFTConfig\n",
|
| 339 |
+
"\n",
|
| 340 |
+
"sft_cfg = SFTConfig(\n",
|
| 341 |
+
" dataset_text_field = 'text',\n",
|
| 342 |
+
" max_seq_length = 2048,\n",
|
| 343 |
+
" num_train_epochs = 2,\n",
|
| 344 |
+
" per_device_train_batch_size = 4,\n",
|
| 345 |
+
" gradient_accumulation_steps = 4,\n",
|
| 346 |
+
" learning_rate = 2e-4,\n",
|
| 347 |
+
" warmup_steps = 10,\n",
|
| 348 |
+
" output_dir = 'pm-ops-sft-warmup',\n",
|
| 349 |
+
" report_to = 'none',\n",
|
| 350 |
+
" logging_steps = 5,\n",
|
| 351 |
+
" save_strategy = 'no',\n",
|
| 352 |
+
" dataloader_num_workers = 0,\n",
|
| 353 |
+
")\n",
|
| 354 |
+
"sft_steps = (\n",
|
| 355 |
+
" len(sft_dataset)\n",
|
| 356 |
+
" // (sft_cfg.per_device_train_batch_size * sft_cfg.gradient_accumulation_steps)\n",
|
| 357 |
+
" * sft_cfg.num_train_epochs\n",
|
| 358 |
+
")\n",
|
| 359 |
+
"print(f'SFT: {len(sft_dataset)} examples × {sft_cfg.num_train_epochs} epochs → ~{sft_steps} steps')\n",
|
| 360 |
+
"\n",
|
| 361 |
+
"sft_trainer = SFTTrainer(model=model, tokenizer=tokenizer,\n",
|
| 362 |
+
" train_dataset=sft_dataset, args=sft_cfg)\n",
|
| 363 |
+
"sft_stats = sft_trainer.train()\n",
|
| 364 |
+
"loss = sft_stats.metrics.get('train_loss', 0)\n",
|
| 365 |
+
"runtime = sft_stats.metrics.get('train_runtime', 0)\n",
|
| 366 |
+
"print(f'SFT done: {round(runtime/60, 1)} min loss={loss:.3f}')"
|
| 367 |
+
]
|
| 368 |
+
},
|
| 369 |
+
{
|
| 370 |
+
"cell_type": "markdown",
|
| 371 |
+
"id": "v4-s9",
|
| 372 |
+
"metadata": {},
|
| 373 |
+
"source": [
|
| 374 |
+
"## 9. Verify SFT — Model Must Output Valid JSON"
|
| 375 |
+
]
|
| 376 |
+
},
|
| 377 |
+
{
|
| 378 |
+
"cell_type": "code",
|
| 379 |
+
"execution_count": null,
|
| 380 |
+
"id": "v4-sft-verify",
|
| 381 |
+
"metadata": {},
|
| 382 |
+
"outputs": [],
|
| 383 |
+
"source": [
|
| 384 |
+
"from training.rollout import extract_json_action\n",
|
| 385 |
+
"\n",
|
| 386 |
+
"model.eval()\n",
|
| 387 |
+
"with GenericEnvClient(base_url=ENV_URL).sync() as _env:\n",
|
| 388 |
+
" r = _env.reset(seed=99001)\n",
|
| 389 |
+
" obs_dict = _obs_to_dict(r.observation if hasattr(r, 'observation') else r)\n",
|
| 390 |
+
" obs_text = _current_obs_text(obs_dict, 0, obs_dict.get('task_brief', ''))\n",
|
| 391 |
+
" msgs = build_messages([], obs_text)\n",
|
| 392 |
+
" prompt = tokenizer.apply_chat_template(\n",
|
| 393 |
+
" msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n",
|
| 394 |
+
" )\n",
|
| 395 |
+
"\n",
|
| 396 |
+
"inputs = tokenizer([prompt], return_tensors='pt').to(model.device)\n",
|
| 397 |
+
"with torch.no_grad():\n",
|
| 398 |
+
" out = model.generate(**inputs, max_new_tokens=128, do_sample=False,\n",
|
| 399 |
+
" pad_token_id=tokenizer.eos_token_id)\n",
|
| 400 |
+
"completion = tokenizer.decode(out[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)\n",
|
| 401 |
+
"parsed = extract_json_action(completion)\n",
|
| 402 |
+
"\n",
|
| 403 |
+
"print(f'Output: {completion[:400]}')\n",
|
| 404 |
+
"print(f'Parsed: {parsed}')\n",
|
| 405 |
+
"if parsed is not None:\n",
|
| 406 |
+
" print('PASS ✓ — model outputs valid JSON after SFT')\n",
|
| 407 |
+
"else:\n",
|
| 408 |
+
" print('FAIL ✗ — still no valid JSON. Run SFT again with more episodes or epochs.')\n",
|
| 409 |
+
"model.train()"
|
| 410 |
+
]
|
| 411 |
+
},
|
| 412 |
+
{
|
| 413 |
+
"cell_type": "markdown",
|
| 414 |
+
"id": "v4-phase2",
|
| 415 |
+
"metadata": {},
|
| 416 |
+
"source": [
|
| 417 |
+
"---\n",
|
| 418 |
+
"## Phase 2 — GRPO\n",
|
| 419 |
+
"\n",
|
| 420 |
+
"Reward components — sum ≈ 0.90 when all correct:\n",
|
| 421 |
+
"\n",
|
| 422 |
+
"| Signal | Weight | Description |\n",
|
| 423 |
+
"|---|---|---|\n",
|
| 424 |
+
"| `read_runbook` | +0.10 | Did agent read the runbook first? |\n",
|
| 425 |
+
"| `valid_label` | +0.20 / −0.10 | Ticket label in org's `label_taxonomy`? |\n",
|
| 426 |
+
"| `valid_priority` | +0.15 / −0.10 | Ticket priority in org's `priority_levels`? |\n",
|
| 427 |
+
"| `valid_team` | +0.20 / −0.10 | Assigned team in org's `team_map`? |\n",
|
| 428 |
+
"| `right_channel` | +0.25 / −0.10/ch | Posted to org's `oncall_channels`? |\n",
|
| 429 |
+
"| `env_bonus` | ×0.10 | Env grader confirmation |\n",
|
| 430 |
+
"\n",
|
| 431 |
+
"Reward **varies per org config** (different valid values per seed) → genuine GRPO advantage signal. \n",
|
| 432 |
+
"Each of the N GRPO generations gets a different env seed via `gen_slot` offset."
|
| 433 |
+
]
|
| 434 |
+
},
|
| 435 |
+
{
|
| 436 |
+
"cell_type": "markdown",
|
| 437 |
+
"id": "v4-s10",
|
| 438 |
+
"metadata": {},
|
| 439 |
+
"source": [
|
| 440 |
+
"## 10. GRPO Training Dataset"
|
| 441 |
+
]
|
| 442 |
+
},
|
| 443 |
+
{
|
| 444 |
+
"cell_type": "code",
|
| 445 |
+
"execution_count": null,
|
| 446 |
+
"id": "v4-grpo-data",
|
| 447 |
+
"metadata": {},
|
| 448 |
+
"outputs": [],
|
| 449 |
+
"source": [
|
| 450 |
+
"from training.dataset import generate_triage_dataset\n",
|
| 451 |
+
"\n",
|
| 452 |
+
"rows = generate_triage_dataset(n_episodes=N_GRPO_EPISODES, base_seed=42)\n",
|
| 453 |
+
"grpo_dataset = Dataset.from_list([{'prompt': r['prompt']} for r in rows])\n",
|
| 454 |
+
"print(f'GRPO dataset: {len(grpo_dataset)} triage episodes')\n",
|
| 455 |
+
"print(f'Difficulties: {set(r[\"difficulty\"] for r in rows)}')\n",
|
| 456 |
+
"print(f'Sample: {rows[0][\"prompt\"][:120]}')"
|
| 457 |
+
]
|
| 458 |
+
},
|
| 459 |
+
{
|
| 460 |
+
"cell_type": "markdown",
|
| 461 |
+
"id": "v4-s11",
|
| 462 |
+
"metadata": {},
|
| 463 |
+
"source": [
|
| 464 |
+
"## 11. GRPO Rollout + Reward"
|
| 465 |
+
]
|
| 466 |
+
},
|
| 467 |
+
{
|
| 468 |
+
"cell_type": "code",
|
| 469 |
+
"execution_count": null,
|
| 470 |
+
"id": "v4-rollout",
|
| 471 |
+
"metadata": {},
|
| 472 |
+
"outputs": [],
|
| 473 |
+
"source": [
|
| 474 |
+
"import torch.nn.functional as F\n",
|
| 475 |
+
"from training.rollout import (\n",
|
| 476 |
+
" _obs_to_dict, _current_obs_text, build_messages,\n",
|
| 477 |
+
" extract_json_action, step_aware_fallback,\n",
|
| 478 |
+
")\n",
|
| 479 |
+
"from training.dataset import parse_seed_from_prompt\n",
|
| 480 |
+
"from training.rewards import compute_rollout_reward\n",
|
| 481 |
+
"\n",
|
| 482 |
+
"grpo_env = GenericEnvClient(base_url=ENV_URL).sync()\n",
|
| 483 |
+
"grpo_env.connect()\n",
|
| 484 |
+
"print('GRPO env connected')\n",
|
| 485 |
+
"\n",
|
| 486 |
+
"\n",
|
| 487 |
+
"def run_grpo_episode(trainer, env, tok, dataset_prompt, max_steps=TRAIN_MAX_STEPS, gen_slot=0):\n",
|
| 488 |
+
" \"\"\"One full PM-ops episode. gen_slot offsets seed so parallel GRPO generations\n",
|
| 489 |
+
" explore distinct env episodes even for the same base prompt.\"\"\"\n",
|
| 490 |
+
" seed = parse_seed_from_prompt(dataset_prompt)\n",
|
| 491 |
+
" result = env.reset(seed=seed + gen_slot) if seed is not None else env.reset()\n",
|
| 492 |
+
"\n",
|
| 493 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 494 |
+
" task_brief = obs_dict.get('task_brief') or dataset_prompt\n",
|
| 495 |
+
"\n",
|
| 496 |
+
" prompt_ids, completion_ids, logprobs = [], [], []\n",
|
| 497 |
+
" turn_history = []\n",
|
| 498 |
+
" valid_json_count = 0\n",
|
| 499 |
+
" env_score = 0.0\n",
|
| 500 |
+
" step, done = 0, False\n",
|
| 501 |
+
"\n",
|
| 502 |
+
" # Runbook-compliance tracking\n",
|
| 503 |
+
" read_runbook_done = False\n",
|
| 504 |
+
" valid_labels: set = set()\n",
|
| 505 |
+
" valid_priorities: set = set()\n",
|
| 506 |
+
" valid_teams: set = set()\n",
|
| 507 |
+
" oncall_channels: set = set()\n",
|
| 508 |
+
" ticket_label: str | None = None\n",
|
| 509 |
+
" ticket_priority: str | None = None\n",
|
| 510 |
+
" assigned_team: str | None = None\n",
|
| 511 |
+
" posted_channels: list = []\n",
|
| 512 |
+
"\n",
|
| 513 |
+
" _model = (trainer.accelerator.unwrap_model(trainer.model)\n",
|
| 514 |
+
" if hasattr(trainer, 'accelerator') else trainer.model)\n",
|
| 515 |
+
" _device = (trainer.accelerator.device\n",
|
| 516 |
+
" if hasattr(trainer, 'accelerator') else next(_model.parameters()).device)\n",
|
| 517 |
+
"\n",
|
| 518 |
+
" while not done and step < max_steps:\n",
|
| 519 |
+
" obs_text = _current_obs_text(obs_dict, step, task_brief)\n",
|
| 520 |
+
" msgs = build_messages(turn_history, obs_text)\n",
|
| 521 |
+
" prompt_text = tok.apply_chat_template(\n",
|
| 522 |
+
" msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n",
|
| 523 |
+
" )\n",
|
| 524 |
+
"\n",
|
| 525 |
+
" enc = tok(prompt_text, return_tensors='pt', truncation=True, max_length=4096).to(_device)\n",
|
| 526 |
+
" plen = enc['input_ids'].shape[1]\n",
|
| 527 |
+
" with torch.no_grad():\n",
|
| 528 |
+
" out = _model.generate(\n",
|
| 529 |
+
" **enc,\n",
|
| 530 |
+
" max_new_tokens = MAX_COMP_LEN,\n",
|
| 531 |
+
" do_sample = True,\n",
|
| 532 |
+
" temperature = 1.1,\n",
|
| 533 |
+
" top_p = 0.95,\n",
|
| 534 |
+
" top_k = 50,\n",
|
| 535 |
+
" pad_token_id = tok.pad_token_id or tok.eos_token_id,\n",
|
| 536 |
+
" output_scores = True,\n",
|
| 537 |
+
" return_dict_in_generate = True,\n",
|
| 538 |
+
" )\n",
|
| 539 |
+
"\n",
|
| 540 |
+
" cids = out.sequences[0][plen:].tolist()\n",
|
| 541 |
+
" completion_text = tok.decode(cids, skip_special_tokens=True)\n",
|
| 542 |
+
"\n",
|
| 543 |
+
" prompt_ids.extend(enc['input_ids'][0].tolist())\n",
|
| 544 |
+
" completion_ids.extend(cids)\n",
|
| 545 |
+
" logprobs.extend([\n",
|
| 546 |
+
" F.log_softmax(s[0], dim=-1)[t].item()\n",
|
| 547 |
+
" for s, t in zip(out.scores, cids)\n",
|
| 548 |
+
" ])\n",
|
| 549 |
+
"\n",
|
| 550 |
+
" if step == 0:\n",
|
| 551 |
+
" print(f' [sample] {repr(completion_text[:180])}')\n",
|
| 552 |
+
"\n",
|
| 553 |
+
" parsed = extract_json_action(completion_text)\n",
|
| 554 |
+
" is_valid = parsed is not None\n",
|
| 555 |
+
" if not is_valid:\n",
|
| 556 |
+
" parsed = step_aware_fallback(step, max_steps)\n",
|
| 557 |
+
" else:\n",
|
| 558 |
+
" valid_json_count += 1\n",
|
| 559 |
+
"\n",
|
| 560 |
+
" action_type = parsed.get('action_type', 'meta.noop')\n",
|
| 561 |
+
" args = parsed.get('args', {})\n",
|
| 562 |
+
" print(f' [step {step}] {action_type}')\n",
|
| 563 |
+
"\n",
|
| 564 |
+
" # Compliance signal capture\n",
|
| 565 |
+
" if action_type == 'meta.read_runbook' and is_valid:\n",
|
| 566 |
+
" read_runbook_done = True\n",
|
| 567 |
+
" if action_type == 'ticketing.create_ticket' and is_valid and ticket_label is None:\n",
|
| 568 |
+
" ticket_label = args.get('label')\n",
|
| 569 |
+
" ticket_priority = args.get('priority')\n",
|
| 570 |
+
" if action_type == 'ticketing.assign_ticket' and is_valid and assigned_team is None:\n",
|
| 571 |
+
" assigned_team = args.get('team')\n",
|
| 572 |
+
" if action_type == 'chat.post_message' and is_valid:\n",
|
| 573 |
+
" ch = args.get('channel', '')\n",
|
| 574 |
+
" if ch:\n",
|
| 575 |
+
" posted_channels.append(ch)\n",
|
| 576 |
+
"\n",
|
| 577 |
+
" turn_history.append({\n",
|
| 578 |
+
" 'obs_text' : obs_text,\n",
|
| 579 |
+
" 'completion': completion_text,\n",
|
| 580 |
+
" 'is_runbook': (action_type == 'meta.read_runbook' and is_valid),\n",
|
| 581 |
+
" })\n",
|
| 582 |
+
"\n",
|
| 583 |
+
" result = env.step({'action_type': action_type, 'args': args})\n",
|
| 584 |
+
" obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
|
| 585 |
+
"\n",
|
| 586 |
+
" # Extract org config from runbook response (one step after the call)\n",
|
| 587 |
+
" last_result = obs_dict.get('last_action_result') or {}\n",
|
| 588 |
+
" if action_type == 'meta.read_runbook' and last_result.get('ok'):\n",
|
| 589 |
+
" data = last_result.get('data') or {}\n",
|
| 590 |
+
" if isinstance(data, dict):\n",
|
| 591 |
+
" org = data.get('org_config') or {}\n",
|
| 592 |
+
" valid_labels = set(org.get('label_taxonomy', {}).values())\n",
|
| 593 |
+
" valid_priorities = set(org.get('priority_levels', []))\n",
|
| 594 |
+
" valid_teams = set(org.get('team_map', {}).values())\n",
|
| 595 |
+
" oncall_channels = set(org.get('oncall_channels', {}).values())\n",
|
| 596 |
+
"\n",
|
| 597 |
+
" done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n",
|
| 598 |
+
" env_score = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n",
|
| 599 |
+
" step += 1\n",
|
| 600 |
+
"\n",
|
| 601 |
+
" reward = compute_rollout_reward(\n",
|
| 602 |
+
" read_runbook_done = read_runbook_done,\n",
|
| 603 |
+
" valid_labels = valid_labels,\n",
|
| 604 |
+
" valid_priorities = valid_priorities,\n",
|
| 605 |
+
" valid_teams = valid_teams,\n",
|
| 606 |
+
" oncall_channels = oncall_channels,\n",
|
| 607 |
+
" ticket_label = ticket_label,\n",
|
| 608 |
+
" ticket_priority = ticket_priority,\n",
|
| 609 |
+
" assigned_team = assigned_team,\n",
|
| 610 |
+
" posted_channels = posted_channels,\n",
|
| 611 |
+
" env_score = env_score,\n",
|
| 612 |
+
" valid_json_count = valid_json_count,\n",
|
| 613 |
+
" )\n",
|
| 614 |
+
"\n",
|
| 615 |
+
" lbl = '✓' if ticket_label and ticket_label in valid_labels else ('✗' if ticket_label else '-')\n",
|
| 616 |
+
" pri = '✓' if ticket_priority and ticket_priority in valid_priorities else ('✗' if ticket_priority else '-')\n",
|
| 617 |
+
" tm = '✓' if assigned_team and assigned_team in valid_teams else ('✗' if assigned_team else '-')\n",
|
| 618 |
+
" ch = '✓' if any(c in oncall_channels for c in posted_channels) else ('✗' if posted_channels else '-')\n",
|
| 619 |
+
" print(f' [rollout] steps={step} env={env_score:.3f} '\n",
|
| 620 |
+
" f'label={lbl} priority={pri} team={tm} channel={ch} → reward={reward:.3f}')\n",
|
| 621 |
+
"\n",
|
| 622 |
+
" return {\n",
|
| 623 |
+
" 'prompt_ids' : prompt_ids,\n",
|
| 624 |
+
" 'completion_ids': completion_ids,\n",
|
| 625 |
+
" 'logprobs' : logprobs,\n",
|
| 626 |
+
" 'reward' : reward,\n",
|
| 627 |
+
" }\n",
|
| 628 |
+
"\n",
|
| 629 |
+
"\n",
|
| 630 |
+
"def grpo_rollout_func(prompts, trainer=None):\n",
|
| 631 |
+
" out = {'prompt_ids': [], 'completion_ids': [], 'logprobs': [], 'reward': []}\n",
|
| 632 |
+
" prompt_seen: dict = {}\n",
|
| 633 |
+
" for prompt in prompts:\n",
|
| 634 |
+
" gen_offset = prompt_seen.get(prompt, 0)\n",
|
| 635 |
+
" prompt_seen[prompt] = gen_offset + 1\n",
|
| 636 |
+
" ep = run_grpo_episode(\n",
|
| 637 |
+
" trainer, grpo_env, tokenizer, prompt, TRAIN_MAX_STEPS, gen_slot=gen_offset\n",
|
| 638 |
+
" )\n",
|
| 639 |
+
" for k in out:\n",
|
| 640 |
+
" out[k].append(ep[k])\n",
|
| 641 |
+
" return out\n",
|
| 642 |
+
"\n",
|
| 643 |
+
"\n",
|
| 644 |
+
"def grpo_reward_func(completions, **kwargs):\n",
|
| 645 |
+
" \"\"\"Passthrough — reward pre-computed in grpo_rollout_func.\"\"\"\n",
|
| 646 |
+
" rewards = kwargs.get('reward', [])\n",
|
| 647 |
+
" if not rewards:\n",
|
| 648 |
+
" print(f'[grpo_reward_func] no reward in kwargs — keys: {list(kwargs.keys())}')\n",
|
| 649 |
+
" return [0.0] * len(completions)\n",
|
| 650 |
+
" return [float(r) for r in rewards]\n",
|
| 651 |
+
"\n",
|
| 652 |
+
"\n",
|
| 653 |
+
"print(f'GRPO rollout ready max_steps={TRAIN_MAX_STEPS}')"
|
| 654 |
+
]
|
| 655 |
+
},
|
| 656 |
+
{
|
| 657 |
+
"cell_type": "markdown",
|
| 658 |
+
"id": "v4-s12",
|
| 659 |
+
"metadata": {},
|
| 660 |
+
"source": [
|
| 661 |
+
"## 12. GRPO Config + Trainer"
|
| 662 |
+
]
|
| 663 |
+
},
|
| 664 |
+
{
|
| 665 |
+
"cell_type": "code",
|
| 666 |
+
"execution_count": null,
|
| 667 |
+
"id": "v4-trainer",
|
| 668 |
+
"metadata": {},
|
| 669 |
+
"outputs": [],
|
| 670 |
+
"source": [
|
| 671 |
+
"from trl import GRPOConfig\n",
|
| 672 |
+
"from training.pm_ops_trainer import PMOpsGRPOTrainer\n",
|
| 673 |
+
"\n",
|
| 674 |
+
"OUTPUT_DIR = 'pm-ops-grpo-Qwen3-1.7B-triage-v4'\n",
|
| 675 |
+
"HF_REPO_ID = f'Saurav1/{OUTPUT_DIR}'\n",
|
| 676 |
+
"\n",
|
| 677 |
+
"grpo_cfg = GRPOConfig(\n",
|
| 678 |
+
" num_train_epochs = 2,\n",
|
| 679 |
+
" learning_rate = 1e-6,\n",
|
| 680 |
+
" gradient_accumulation_steps = GRAD_ACCUM,\n",
|
| 681 |
+
" per_device_train_batch_size = 1,\n",
|
| 682 |
+
" warmup_steps = 5,\n",
|
| 683 |
+
" num_generations = NUM_GEN,\n",
|
| 684 |
+
" # max_completion_length must match MAX_COMP_LEN so TRL's internal\n",
|
| 685 |
+
" # single-turn generation doesn't overflow and zero the policy loss\n",
|
| 686 |
+
" # (clipped_ratio=1.0 means all gradient is masked → loss=0.000).\n",
|
| 687 |
+
" max_completion_length = MAX_COMP_LEN,\n",
|
| 688 |
+
" max_prompt_length = 4096,\n",
|
| 689 |
+
" use_vllm = False,\n",
|
| 690 |
+
" output_dir = OUTPUT_DIR,\n",
|
| 691 |
+
" report_to = 'none',\n",
|
| 692 |
+
" logging_steps = 1,\n",
|
| 693 |
+
" save_steps = 20,\n",
|
| 694 |
+
" gradient_checkpointing = False,\n",
|
| 695 |
+
")\n",
|
| 696 |
+
"\n",
|
| 697 |
+
"eff_batch = grpo_cfg.per_device_train_batch_size * GRAD_ACCUM\n",
|
| 698 |
+
"total_steps = len(grpo_dataset) * NUM_GEN * grpo_cfg.num_train_epochs // eff_batch\n",
|
| 699 |
+
"print(f'GRPO: {len(grpo_dataset)} eps × {NUM_GEN} gen × {grpo_cfg.num_train_epochs} epochs → ~{total_steps} steps')\n",
|
| 700 |
+
"\n",
|
| 701 |
+
"trainer = PMOpsGRPOTrainer(\n",
|
| 702 |
+
" model = model,\n",
|
| 703 |
+
" processing_class = tokenizer,\n",
|
| 704 |
+
" reward_funcs = grpo_reward_func,\n",
|
| 705 |
+
" train_dataset = grpo_dataset,\n",
|
| 706 |
+
" args = grpo_cfg,\n",
|
| 707 |
+
" rollout_func = grpo_rollout_func,\n",
|
| 708 |
+
")\n",
|
| 709 |
+
"assert grpo_cfg.use_vllm is False\n",
|
| 710 |
+
"print(f'Trainer: {type(trainer).__name__} ready')"
|
| 711 |
+
]
|
| 712 |
+
},
|
| 713 |
+
{
|
| 714 |
+
"cell_type": "markdown",
|
| 715 |
+
"id": "v4-s13",
|
| 716 |
+
"metadata": {},
|
| 717 |
+
"source": [
|
| 718 |
+
"## 13. Preflight Probe"
|
| 719 |
+
]
|
| 720 |
+
},
|
| 721 |
+
{
|
| 722 |
+
"cell_type": "code",
|
| 723 |
+
"execution_count": null,
|
| 724 |
+
"id": "v4-probe",
|
| 725 |
+
"metadata": {},
|
| 726 |
+
"outputs": [],
|
| 727 |
+
"source": [
|
| 728 |
+
"# Run 2 episodes and verify reward varies (not stuck at a constant)\n",
|
| 729 |
+
"probe = grpo_rollout_func(\n",
|
| 730 |
+
" [grpo_dataset[0]['prompt'], grpo_dataset[1]['prompt']],\n",
|
| 731 |
+
" trainer=trainer,\n",
|
| 732 |
+
")\n",
|
| 733 |
+
"print(f'probe rewards : {probe[\"reward\"]}')\n",
|
| 734 |
+
"assert len(probe['reward']) == 2, 'need one reward per prompt'\n",
|
| 735 |
+
"assert all(isinstance(r, float) for r in probe['reward']), 'rewards must be float'\n",
|
| 736 |
+
"print('Preflight PASS ✓')"
|
| 737 |
+
]
|
| 738 |
+
},
|
| 739 |
+
{
|
| 740 |
+
"cell_type": "markdown",
|
| 741 |
+
"id": "v4-s14",
|
| 742 |
+
"metadata": {},
|
| 743 |
+
"source": [
|
| 744 |
+
"## 14. Train\n",
|
| 745 |
+
"\n",
|
| 746 |
+
"Watch for:\n",
|
| 747 |
+
"- `[sample] '\\`\\`\\`json ...'` — model should output JSON code blocks\n",
|
| 748 |
+
"- `label=✓ priority=✓ team=✓ channel=✓` — compliance signals the model is getting right\n",
|
| 749 |
+
"- `reward/injected_std > 0` in logs — confirms GRPO has a non-zero gradient signal"
|
| 750 |
+
]
|
| 751 |
+
},
|
| 752 |
+
{
|
| 753 |
+
"cell_type": "code",
|
| 754 |
+
"execution_count": null,
|
| 755 |
+
"id": "v4-train",
|
| 756 |
+
"metadata": {},
|
| 757 |
+
"outputs": [],
|
| 758 |
+
"source": [
|
| 759 |
+
"trainer_stats = trainer.train()\n",
|
| 760 |
+
"\n",
|
| 761 |
+
"train_mins = round(trainer_stats.metrics.get('train_runtime', 0) / 60, 1)\n",
|
| 762 |
+
"used_gb = round(torch.cuda.max_memory_reserved() / 1024**3, 2)\n",
|
| 763 |
+
"print(f'Training time : {train_mins} min')\n",
|
| 764 |
+
"print(f'Peak GPU : {used_gb} GB / {TOTAL_GB} GB ({round(used_gb/TOTAL_GB*100, 1)}%)')"
|
| 765 |
+
]
|
| 766 |
+
},
|
| 767 |
+
{
|
| 768 |
+
"cell_type": "markdown",
|
| 769 |
+
"id": "v4-s15",
|
| 770 |
+
"metadata": {},
|
| 771 |
+
"source": [
|
| 772 |
+
"## 15. Save + Push to HF"
|
| 773 |
+
]
|
| 774 |
+
},
|
| 775 |
+
{
|
| 776 |
+
"cell_type": "code",
|
| 777 |
+
"execution_count": null,
|
| 778 |
+
"id": "v4-save",
|
| 779 |
+
"metadata": {},
|
| 780 |
+
"outputs": [],
|
| 781 |
+
"source": [
|
| 782 |
+
"grpo_env.close()\n",
|
| 783 |
+
"\n",
|
| 784 |
+
"model.save_pretrained_merged(OUTPUT_DIR, tokenizer, save_method='merged_16bit')\n",
|
| 785 |
+
"model.push_to_hub_merged(HF_REPO_ID, tokenizer, save_method='merged_16bit')\n",
|
| 786 |
+
"print(f'Pushed -> https://huggingface.co/{HF_REPO_ID}')"
|
| 787 |
+
]
|
| 788 |
+
},
|
| 789 |
+
{
|
| 790 |
+
"cell_type": "markdown",
|
| 791 |
+
"id": "v4-eval-hdr",
|
| 792 |
+
"metadata": {},
|
| 793 |
+
"source": [
|
| 794 |
+
"---\n",
|
| 795 |
+
"## Evaluation — Baseline vs Trained"
|
| 796 |
+
]
|
| 797 |
+
},
|
| 798 |
+
{
|
| 799 |
+
"cell_type": "markdown",
|
| 800 |
+
"id": "v4-s16",
|
| 801 |
+
"metadata": {},
|
| 802 |
+
"source": [
|
| 803 |
+
"## 16. Evaluate"
|
| 804 |
+
]
|
| 805 |
+
},
|
| 806 |
+
{
|
| 807 |
+
"cell_type": "code",
|
| 808 |
+
"execution_count": null,
|
| 809 |
+
"id": "v4-eval",
|
| 810 |
+
"metadata": {},
|
| 811 |
+
"outputs": [],
|
| 812 |
+
"source": "from training.rollout import extract_json_action, step_aware_fallback\nfrom inference import baseline_agent\n\nN_EVAL = 15\nEVAL_MAX_STEPS = 12\nEVAL_SEED_BASE = 9000\n\n# Required after training: switch Unsloth from training hooks to fast inference.\n# Without this, model.generate() is very slow and hangs on T4.\nFastLanguageModel.for_inference(model)\nmodel.eval()\nprint('Model switched to fast inference mode')\n\n\ndef _model_step(obs_dict, task_brief, history, step):\n \"\"\"One greedy decode step. Returns (action_type, args, json_ok).\"\"\"\n obs_text = _current_obs_text(obs_dict, step, task_brief)\n msgs = build_messages(history, obs_text)\n prompt = tokenizer.apply_chat_template(\n msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n )\n inputs = tokenizer([prompt], return_tensors='pt', truncation=True, max_length=4096)\n inputs = {k: v.to(model.device) for k, v in inputs.items()}\n with torch.no_grad():\n out_ids = model.generate(\n **inputs,\n max_new_tokens = 128,\n do_sample = False,\n pad_token_id = tokenizer.eos_token_id,\n )\n completion = tokenizer.decode(\n out_ids[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True\n )\n parsed = extract_json_action(completion)\n json_ok = parsed is not None\n if not json_ok:\n parsed = step_aware_fallback(step, EVAL_MAX_STEPS)\n history.append({'obs_text': obs_text, 'completion': completion, 'is_runbook': False})\n return parsed.get('action_type', 'meta.noop'), parsed.get('args', {}), json_ok, obs_text\n\n\ndef run_eval(n=N_EVAL, verbose=False):\n scores = []\n with GenericEnvClient(base_url=ENV_URL).sync() as env:\n for i in range(n):\n result = env.reset(seed=EVAL_SEED_BASE + i)\n obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n task_brief = obs_dict.get('task_brief', '')\n history, step, score, done = [], 0, 0.0, False\n actions = []\n\n while not done and step < EVAL_MAX_STEPS:\n action_type, args, json_ok, _ = _model_step(obs_dict, task_brief, history, step)\n actions.append(f'{action_type}{\"\" if json_ok else \"!\"}')\n result = env.step({'action_type': action_type, 'args': args})\n obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n score = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n step += 1\n\n scores.append(score)\n action_summary = ' → '.join(a.split('.')[-1] for a in actions)\n print(f' Trained ep {i+1}/{n}: score={score:.3f} [{action_summary}]')\n return scores\n\n\ndef run_baseline(n=N_EVAL):\n scores = []\n with GenericEnvClient(base_url=ENV_URL).sync() as env:\n for i in range(n):\n result = env.reset(seed=EVAL_SEED_BASE + i)\n obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n org_config, step, score, done = {}, 0, 0.0, False\n actions = []\n while not done and step < EVAL_MAX_STEPS:\n at, args = baseline_agent(obs_dict, org_config)\n actions.append(at.split('.')[-1])\n result = env.step({'action_type': at, 'args': args})\n obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n score = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n if at == 'meta.read_runbook':\n last = obs_dict.get('last_action_result') or {}\n if last.get('ok'):\n data = last.get('data') or {}\n if isinstance(data, dict) and 'org_config' in data:\n org_config.update(data['org_config'])\n step += 1\n scores.append(score)\n action_summary = ' → '.join(actions)\n print(f' Baseline ep {i+1}/{n}: score={score:.3f} [{action_summary}]')\n return scores\n\n\nprint('--- Baseline ---')\nbaseline_scores = run_baseline()\nprint('\\n--- Trained ---')\ntrained_scores = run_eval()\n\nb_avg = sum(baseline_scores) / N_EVAL\nt_avg = sum(trained_scores) / N_EVAL\nprint(f'\\nBaseline avg : {b_avg:.3f}')\nprint(f'Trained avg : {t_avg:.3f}')\nprint(f'Delta : {t_avg - b_avg:+.3f}')\nprint()\nprint('Action key: ! = fallback (no valid JSON from model)')"
|
| 813 |
+
},
|
| 814 |
+
{
|
| 815 |
+
"cell_type": "markdown",
|
| 816 |
+
"id": "v4-s17",
|
| 817 |
+
"metadata": {},
|
| 818 |
+
"source": [
|
| 819 |
+
"## 17. Teardown"
|
| 820 |
+
]
|
| 821 |
+
},
|
| 822 |
+
{
|
| 823 |
+
"cell_type": "code",
|
| 824 |
+
"execution_count": null,
|
| 825 |
+
"id": "v4-teardown",
|
| 826 |
+
"metadata": {},
|
| 827 |
+
"outputs": [],
|
| 828 |
+
"source": [
|
| 829 |
+
"server_proc.terminate()\n",
|
| 830 |
+
"print('PM-Ops server stopped')"
|
| 831 |
+
]
|
| 832 |
+
}
|
| 833 |
+
],
|
| 834 |
+
"metadata": {
|
| 835 |
+
"kernelspec": {
|
| 836 |
+
"display_name": "Python 3",
|
| 837 |
+
"language": "python",
|
| 838 |
+
"name": "python3"
|
| 839 |
+
},
|
| 840 |
+
"language_info": {
|
| 841 |
+
"name": "python",
|
| 842 |
+
"version": "3.12.0"
|
| 843 |
+
}
|
| 844 |
+
},
|
| 845 |
+
"nbformat": 4,
|
| 846 |
+
"nbformat_minor": 5
|
| 847 |
+
}
|
training/triage_dataset.jsonl
ADDED
|
@@ -0,0 +1,150 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{"prompt": "SEED:478163327 | Support escalation: payments is intermittently failing for EU users. Triage and route per org process.", "seed": 478163327, "difficulty": "easy"}
|
| 2 |
+
{"prompt": "SEED:1181241943 | A user reported that the notifications is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1181241943, "difficulty": "medium"}
|
| 3 |
+
{"prompt": "SEED:958682846 | Monitoring alert: elevated error rate on api-gateway. File a bug ticket and page the right team.", "seed": 958682846, "difficulty": "medium"}
|
| 4 |
+
{"prompt": "SEED:440213415 | A user reported that the inventory is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 440213415, "difficulty": "easy"}
|
| 5 |
+
{"prompt": "SEED:1812140441 | A user reported that the payments is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1812140441, "difficulty": "easy"}
|
| 6 |
+
{"prompt": "SEED:127978094 | A user reported that the auth is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 127978094, "difficulty": "easy"}
|
| 7 |
+
{"prompt": "SEED:939042955 | Monitoring alert: elevated error rate on payments. File a bug ticket and page the right team.", "seed": 939042955, "difficulty": "medium"}
|
| 8 |
+
{"prompt": "SEED:113971123 | Monitoring alert: elevated error rate on inventory. File a bug ticket and page the right team.", "seed": 113971123, "difficulty": "medium"}
|
| 9 |
+
{"prompt": "SEED:1801823908 | Support escalation: api-gateway is intermittently failing for EU users. Triage and route per org process.", "seed": 1801823908, "difficulty": "medium"}
|
| 10 |
+
{"prompt": "SEED:1929338154 | A user reported that the checkout is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1929338154, "difficulty": "medium"}
|
| 11 |
+
{"prompt": "SEED:27911967 | Monitoring alert: elevated error rate on api-gateway. File a bug ticket and page the right team.", "seed": 27911967, "difficulty": "medium"}
|
| 12 |
+
{"prompt": "SEED:1815115025 | Monitoring alert: elevated error rate on payments. File a bug ticket and page the right team.", "seed": 1815115025, "difficulty": "medium"}
|
| 13 |
+
{"prompt": "SEED:1193448329 | Support escalation: checkout is intermittently failing for EU users. Triage and route per org process.", "seed": 1193448329, "difficulty": "medium"}
|
| 14 |
+
{"prompt": "SEED:924765563 | Monitoring alert: elevated error rate on search. File a bug ticket and page the right team.", "seed": 924765563, "difficulty": "medium"}
|
| 15 |
+
{"prompt": "SEED:438989805 | Support escalation: notifications is intermittently failing for EU users. Triage and route per org process.", "seed": 438989805, "difficulty": "easy"}
|
| 16 |
+
{"prompt": "SEED:1631775357 | Monitoring alert: elevated error rate on auth. File a bug ticket and page the right team.", "seed": 1631775357, "difficulty": "easy"}
|
| 17 |
+
{"prompt": "SEED:1541804686 | Support escalation: api-gateway is intermittently failing for EU users. Triage and route per org process.", "seed": 1541804686, "difficulty": "medium"}
|
| 18 |
+
{"prompt": "SEED:1136108454 | Support escalation: payments is intermittently failing for EU users. Triage and route per org process.", "seed": 1136108454, "difficulty": "easy"}
|
| 19 |
+
{"prompt": "SEED:1973214822 | Support escalation: payments is intermittently failing for EU users. Triage and route per org process.", "seed": 1973214822, "difficulty": "easy"}
|
| 20 |
+
{"prompt": "SEED:1625792787 | A user reported that the checkout is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1625792787, "difficulty": "easy"}
|
| 21 |
+
{"prompt": "SEED:1259191105 | Monitoring alert: elevated error rate on checkout. File a bug ticket and page the right team.", "seed": 1259191105, "difficulty": "medium"}
|
| 22 |
+
{"prompt": "SEED:825873196 | A user reported that the api-gateway is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 825873196, "difficulty": "easy"}
|
| 23 |
+
{"prompt": "SEED:196814233 | Support escalation: api-gateway is intermittently failing for EU users. Triage and route per org process.", "seed": 196814233, "difficulty": "medium"}
|
| 24 |
+
{"prompt": "SEED:1242911821 | Monitoring alert: elevated error rate on payments. File a bug ticket and page the right team.", "seed": 1242911821, "difficulty": "easy"}
|
| 25 |
+
{"prompt": "SEED:999829240 | Monitoring alert: elevated error rate on checkout. File a bug ticket and page the right team.", "seed": 999829240, "difficulty": "easy"}
|
| 26 |
+
{"prompt": "SEED:1632629719 | Monitoring alert: elevated error rate on notifications. File a bug ticket and page the right team.", "seed": 1632629719, "difficulty": "medium"}
|
| 27 |
+
{"prompt": "SEED:1947382419 | Support escalation: inventory is intermittently failing for EU users. Triage and route per org process.", "seed": 1947382419, "difficulty": "medium"}
|
| 28 |
+
{"prompt": "SEED:698594025 | Monitoring alert: elevated error rate on inventory. File a bug ticket and page the right team.", "seed": 698594025, "difficulty": "medium"}
|
| 29 |
+
{"prompt": "SEED:1525876051 | A user reported that the checkout is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1525876051, "difficulty": "medium"}
|
| 30 |
+
{"prompt": "SEED:1146660997 | Monitoring alert: elevated error rate on auth. File a bug ticket and page the right team.", "seed": 1146660997, "difficulty": "easy"}
|
| 31 |
+
{"prompt": "SEED:735034881 | Support escalation: inventory is intermittently failing for EU users. Triage and route per org process.", "seed": 735034881, "difficulty": "medium"}
|
| 32 |
+
{"prompt": "SEED:701808367 | Monitoring alert: elevated error rate on auth. File a bug ticket and page the right team.", "seed": 701808367, "difficulty": "hard"}
|
| 33 |
+
{"prompt": "SEED:1629748727 | A user reported that the api-gateway is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1629748727, "difficulty": "medium"}
|
| 34 |
+
{"prompt": "SEED:943239974 | Support escalation: search is intermittently failing for EU users. Triage and route per org process.", "seed": 943239974, "difficulty": "medium"}
|
| 35 |
+
{"prompt": "SEED:240251661 | A user reported that the payments is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 240251661, "difficulty": "medium"}
|
| 36 |
+
{"prompt": "SEED:137869475 | Support escalation: api-gateway is intermittently failing for EU users. Triage and route per org process.", "seed": 137869475, "difficulty": "medium"}
|
| 37 |
+
{"prompt": "SEED:1722989659 | Monitoring alert: elevated error rate on search. File a bug ticket and page the right team.", "seed": 1722989659, "difficulty": "medium"}
|
| 38 |
+
{"prompt": "SEED:284277889 | Monitoring alert: elevated error rate on payments. File a bug ticket and page the right team.", "seed": 284277889, "difficulty": "medium"}
|
| 39 |
+
{"prompt": "SEED:1351531223 | Monitoring alert: elevated error rate on payments. File a bug ticket and page the right team.", "seed": 1351531223, "difficulty": "medium"}
|
| 40 |
+
{"prompt": "SEED:2144181937 | Monitoring alert: elevated error rate on checkout. File a bug ticket and page the right team.", "seed": 2144181937, "difficulty": "hard"}
|
| 41 |
+
{"prompt": "SEED:1970753705 | Monitoring alert: elevated error rate on auth. File a bug ticket and page the right team.", "seed": 1970753705, "difficulty": "medium"}
|
| 42 |
+
{"prompt": "SEED:1137651678 | Support escalation: payments is intermittently failing for EU users. Triage and route per org process.", "seed": 1137651678, "difficulty": "medium"}
|
| 43 |
+
{"prompt": "SEED:1059257080 | Monitoring alert: elevated error rate on payments. File a bug ticket and page the right team.", "seed": 1059257080, "difficulty": "medium"}
|
| 44 |
+
{"prompt": "SEED:1840109255 | Support escalation: auth is intermittently failing for EU users. Triage and route per org process.", "seed": 1840109255, "difficulty": "hard"}
|
| 45 |
+
{"prompt": "SEED:1554762903 | Monitoring alert: elevated error rate on payments. File a bug ticket and page the right team.", "seed": 1554762903, "difficulty": "medium"}
|
| 46 |
+
{"prompt": "SEED:594130308 | Support escalation: checkout is intermittently failing for EU users. Triage and route per org process.", "seed": 594130308, "difficulty": "hard"}
|
| 47 |
+
{"prompt": "SEED:390452952 | Support escalation: inventory is intermittently failing for EU users. Triage and route per org process.", "seed": 390452952, "difficulty": "easy"}
|
| 48 |
+
{"prompt": "SEED:470939445 | A user reported that the payments is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 470939445, "difficulty": "medium"}
|
| 49 |
+
{"prompt": "SEED:687117441 | Support escalation: notifications is intermittently failing for EU users. Triage and route per org process.", "seed": 687117441, "difficulty": "hard"}
|
| 50 |
+
{"prompt": "SEED:272849424 | A user reported that the search is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 272849424, "difficulty": "hard"}
|
| 51 |
+
{"prompt": "SEED:1639042338 | Support escalation: checkout is intermittently failing for EU users. Triage and route per org process.", "seed": 1639042338, "difficulty": "hard"}
|
| 52 |
+
{"prompt": "SEED:1079815404 | Support escalation: search is intermittently failing for EU users. Triage and route per org process.", "seed": 1079815404, "difficulty": "easy"}
|
| 53 |
+
{"prompt": "SEED:491995979 | Monitoring alert: elevated error rate on auth. File a bug ticket and page the right team.", "seed": 491995979, "difficulty": "medium"}
|
| 54 |
+
{"prompt": "SEED:1461042543 | A user reported that the inventory is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1461042543, "difficulty": "easy"}
|
| 55 |
+
{"prompt": "SEED:1260573448 | Support escalation: auth is intermittently failing for EU users. Triage and route per org process.", "seed": 1260573448, "difficulty": "hard"}
|
| 56 |
+
{"prompt": "SEED:679282378 | Monitoring alert: elevated error rate on search. File a bug ticket and page the right team.", "seed": 679282378, "difficulty": "hard"}
|
| 57 |
+
{"prompt": "SEED:13938521 | A user reported that the auth is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 13938521, "difficulty": "medium"}
|
| 58 |
+
{"prompt": "SEED:767303988 | Support escalation: api-gateway is intermittently failing for EU users. Triage and route per org process.", "seed": 767303988, "difficulty": "easy"}
|
| 59 |
+
{"prompt": "SEED:1281810600 | A user reported that the payments is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1281810600, "difficulty": "medium"}
|
| 60 |
+
{"prompt": "SEED:656439677 | Monitoring alert: elevated error rate on checkout. File a bug ticket and page the right team.", "seed": 656439677, "difficulty": "medium"}
|
| 61 |
+
{"prompt": "SEED:693847829 | Monitoring alert: elevated error rate on auth. File a bug ticket and page the right team.", "seed": 693847829, "difficulty": "easy"}
|
| 62 |
+
{"prompt": "SEED:1392239675 | A user reported that the notifications is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1392239675, "difficulty": "hard"}
|
| 63 |
+
{"prompt": "SEED:83651970 | A user reported that the search is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 83651970, "difficulty": "easy"}
|
| 64 |
+
{"prompt": "SEED:1558990517 | Support escalation: checkout is intermittently failing for EU users. Triage and route per org process.", "seed": 1558990517, "difficulty": "medium"}
|
| 65 |
+
{"prompt": "SEED:1028439863 | Monitoring alert: elevated error rate on payments. File a bug ticket and page the right team.", "seed": 1028439863, "difficulty": "easy"}
|
| 66 |
+
{"prompt": "SEED:1034535593 | Monitoring alert: elevated error rate on notifications. File a bug ticket and page the right team.", "seed": 1034535593, "difficulty": "easy"}
|
| 67 |
+
{"prompt": "SEED:367878761 | Monitoring alert: elevated error rate on auth. File a bug ticket and page the right team.", "seed": 367878761, "difficulty": "hard"}
|
| 68 |
+
{"prompt": "SEED:297265480 | A user reported that the notifications is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 297265480, "difficulty": "medium"}
|
| 69 |
+
{"prompt": "SEED:551437149 | Monitoring alert: elevated error rate on notifications. File a bug ticket and page the right team.", "seed": 551437149, "difficulty": "hard"}
|
| 70 |
+
{"prompt": "SEED:709214891 | A user reported that the checkout is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 709214891, "difficulty": "medium"}
|
| 71 |
+
{"prompt": "SEED:1817363615 | Monitoring alert: elevated error rate on payments. File a bug ticket and page the right team.", "seed": 1817363615, "difficulty": "medium"}
|
| 72 |
+
{"prompt": "SEED:863937247 | A user reported that the inventory is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 863937247, "difficulty": "medium"}
|
| 73 |
+
{"prompt": "SEED:1713658874 | Support escalation: auth is intermittently failing for EU users. Triage and route per org process.", "seed": 1713658874, "difficulty": "medium"}
|
| 74 |
+
{"prompt": "SEED:1881625505 | A user reported that the notifications is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1881625505, "difficulty": "hard"}
|
| 75 |
+
{"prompt": "SEED:519709079 | Monitoring alert: elevated error rate on notifications. File a bug ticket and page the right team.", "seed": 519709079, "difficulty": "medium"}
|
| 76 |
+
{"prompt": "SEED:965067727 | Monitoring alert: elevated error rate on inventory. File a bug ticket and page the right team.", "seed": 965067727, "difficulty": "easy"}
|
| 77 |
+
{"prompt": "SEED:1452066459 | A user reported that the notifications is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1452066459, "difficulty": "easy"}
|
| 78 |
+
{"prompt": "SEED:988335251 | A user reported that the auth is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 988335251, "difficulty": "medium"}
|
| 79 |
+
{"prompt": "SEED:30884438 | Monitoring alert: elevated error rate on auth. File a bug ticket and page the right team.", "seed": 30884438, "difficulty": "easy"}
|
| 80 |
+
{"prompt": "SEED:252860896 | Support escalation: api-gateway is intermittently failing for EU users. Triage and route per org process.", "seed": 252860896, "difficulty": "medium"}
|
| 81 |
+
{"prompt": "SEED:289482182 | A user reported that the auth is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 289482182, "difficulty": "easy"}
|
| 82 |
+
{"prompt": "SEED:1419179580 | Support escalation: payments is intermittently failing for EU users. Triage and route per org process.", "seed": 1419179580, "difficulty": "easy"}
|
| 83 |
+
{"prompt": "SEED:1022222125 | Support escalation: checkout is intermittently failing for EU users. Triage and route per org process.", "seed": 1022222125, "difficulty": "medium"}
|
| 84 |
+
{"prompt": "SEED:2084839399 | Monitoring alert: elevated error rate on auth. File a bug ticket and page the right team.", "seed": 2084839399, "difficulty": "medium"}
|
| 85 |
+
{"prompt": "SEED:568275055 | Support escalation: notifications is intermittently failing for EU users. Triage and route per org process.", "seed": 568275055, "difficulty": "hard"}
|
| 86 |
+
{"prompt": "SEED:1043665192 | Monitoring alert: elevated error rate on inventory. File a bug ticket and page the right team.", "seed": 1043665192, "difficulty": "hard"}
|
| 87 |
+
{"prompt": "SEED:1748309234 | A user reported that the search is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1748309234, "difficulty": "medium"}
|
| 88 |
+
{"prompt": "SEED:405126228 | A user reported that the inventory is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 405126228, "difficulty": "easy"}
|
| 89 |
+
{"prompt": "SEED:1851350739 | Monitoring alert: elevated error rate on checkout. File a bug ticket and page the right team.", "seed": 1851350739, "difficulty": "medium"}
|
| 90 |
+
{"prompt": "SEED:1819256337 | A user reported that the payments is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1819256337, "difficulty": "hard"}
|
| 91 |
+
{"prompt": "SEED:2005855667 | Support escalation: inventory is intermittently failing for EU users. Triage and route per org process.", "seed": 2005855667, "difficulty": "easy"}
|
| 92 |
+
{"prompt": "SEED:422701550 | A user reported that the auth is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 422701550, "difficulty": "easy"}
|
| 93 |
+
{"prompt": "SEED:1729245242 | A user reported that the auth is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1729245242, "difficulty": "medium"}
|
| 94 |
+
{"prompt": "SEED:469306919 | Monitoring alert: elevated error rate on inventory. File a bug ticket and page the right team.", "seed": 469306919, "difficulty": "medium"}
|
| 95 |
+
{"prompt": "SEED:822873088 | Monitoring alert: elevated error rate on inventory. File a bug ticket and page the right team.", "seed": 822873088, "difficulty": "medium"}
|
| 96 |
+
{"prompt": "SEED:1926780541 | Support escalation: checkout is intermittently failing for EU users. Triage and route per org process.", "seed": 1926780541, "difficulty": "medium"}
|
| 97 |
+
{"prompt": "SEED:1811967841 | A user reported that the api-gateway is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1811967841, "difficulty": "medium"}
|
| 98 |
+
{"prompt": "SEED:1196342297 | Support escalation: payments is intermittently failing for EU users. Triage and route per org process.", "seed": 1196342297, "difficulty": "hard"}
|
| 99 |
+
{"prompt": "SEED:1072910527 | Monitoring alert: elevated error rate on api-gateway. File a bug ticket and page the right team.", "seed": 1072910527, "difficulty": "easy"}
|
| 100 |
+
{"prompt": "SEED:1903232052 | Support escalation: inventory is intermittently failing for EU users. Triage and route per org process.", "seed": 1903232052, "difficulty": "easy"}
|
| 101 |
+
{"prompt": "SEED:217275224 | A user reported that the inventory is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 217275224, "difficulty": "easy"}
|
| 102 |
+
{"prompt": "SEED:400560567 | Monitoring alert: elevated error rate on checkout. File a bug ticket and page the right team.", "seed": 400560567, "difficulty": "medium"}
|
| 103 |
+
{"prompt": "SEED:714300770 | A user reported that the notifications is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 714300770, "difficulty": "hard"}
|
| 104 |
+
{"prompt": "SEED:2085812759 | Support escalation: notifications is intermittently failing for EU users. Triage and route per org process.", "seed": 2085812759, "difficulty": "hard"}
|
| 105 |
+
{"prompt": "SEED:918037633 | Monitoring alert: elevated error rate on payments. File a bug ticket and page the right team.", "seed": 918037633, "difficulty": "hard"}
|
| 106 |
+
{"prompt": "SEED:251837136 | A user reported that the payments is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 251837136, "difficulty": "medium"}
|
| 107 |
+
{"prompt": "SEED:1627677155 | Support escalation: api-gateway is intermittently failing for EU users. Triage and route per org process.", "seed": 1627677155, "difficulty": "easy"}
|
| 108 |
+
{"prompt": "SEED:1676848676 | Monitoring alert: elevated error rate on inventory. File a bug ticket and page the right team.", "seed": 1676848676, "difficulty": "medium"}
|
| 109 |
+
{"prompt": "SEED:1954246074 | A user reported that the api-gateway is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1954246074, "difficulty": "medium"}
|
| 110 |
+
{"prompt": "SEED:1816803306 | Monitoring alert: elevated error rate on search. File a bug ticket and page the right team.", "seed": 1816803306, "difficulty": "hard"}
|
| 111 |
+
{"prompt": "SEED:664847319 | A user reported that the auth is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 664847319, "difficulty": "medium"}
|
| 112 |
+
{"prompt": "SEED:1274350418 | Support escalation: search is intermittently failing for EU users. Triage and route per org process.", "seed": 1274350418, "difficulty": "medium"}
|
| 113 |
+
{"prompt": "SEED:251183830 | Support escalation: api-gateway is intermittently failing for EU users. Triage and route per org process.", "seed": 251183830, "difficulty": "easy"}
|
| 114 |
+
{"prompt": "SEED:1346922426 | Monitoring alert: elevated error rate on checkout. File a bug ticket and page the right team.", "seed": 1346922426, "difficulty": "easy"}
|
| 115 |
+
{"prompt": "SEED:215359682 | Monitoring alert: elevated error rate on api-gateway. File a bug ticket and page the right team.", "seed": 215359682, "difficulty": "hard"}
|
| 116 |
+
{"prompt": "SEED:676168421 | Support escalation: notifications is intermittently failing for EU users. Triage and route per org process.", "seed": 676168421, "difficulty": "easy"}
|
| 117 |
+
{"prompt": "SEED:344076115 | A user reported that the notifications is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 344076115, "difficulty": "medium"}
|
| 118 |
+
{"prompt": "SEED:294296873 | Support escalation: payments is intermittently failing for EU users. Triage and route per org process.", "seed": 294296873, "difficulty": "easy"}
|
| 119 |
+
{"prompt": "SEED:1010193046 | Monitoring alert: elevated error rate on api-gateway. File a bug ticket and page the right team.", "seed": 1010193046, "difficulty": "hard"}
|
| 120 |
+
{"prompt": "SEED:514909066 | A user reported that the search is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 514909066, "difficulty": "medium"}
|
| 121 |
+
{"prompt": "SEED:170689150 | Support escalation: notifications is intermittently failing for EU users. Triage and route per org process.", "seed": 170689150, "difficulty": "easy"}
|
| 122 |
+
{"prompt": "SEED:1800557283 | Monitoring alert: elevated error rate on payments. File a bug ticket and page the right team.", "seed": 1800557283, "difficulty": "medium"}
|
| 123 |
+
{"prompt": "SEED:1119980130 | A user reported that the auth is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1119980130, "difficulty": "medium"}
|
| 124 |
+
{"prompt": "SEED:1349409434 | A user reported that the checkout is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1349409434, "difficulty": "medium"}
|
| 125 |
+
{"prompt": "SEED:1140805649 | Support escalation: checkout is intermittently failing for EU users. Triage and route per org process.", "seed": 1140805649, "difficulty": "hard"}
|
| 126 |
+
{"prompt": "SEED:562117925 | A user reported that the payments is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 562117925, "difficulty": "medium"}
|
| 127 |
+
{"prompt": "SEED:1963764597 | A user reported that the api-gateway is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1963764597, "difficulty": "medium"}
|
| 128 |
+
{"prompt": "SEED:311570307 | Monitoring alert: elevated error rate on inventory. File a bug ticket and page the right team.", "seed": 311570307, "difficulty": "easy"}
|
| 129 |
+
{"prompt": "SEED:1968321319 | A user reported that the payments is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1968321319, "difficulty": "easy"}
|
| 130 |
+
{"prompt": "SEED:314652384 | A user reported that the checkout is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 314652384, "difficulty": "medium"}
|
| 131 |
+
{"prompt": "SEED:1139027119 | Monitoring alert: elevated error rate on search. File a bug ticket and page the right team.", "seed": 1139027119, "difficulty": "medium"}
|
| 132 |
+
{"prompt": "SEED:1498981529 | A user reported that the payments is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1498981529, "difficulty": "easy"}
|
| 133 |
+
{"prompt": "SEED:1049193572 | A user reported that the payments is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1049193572, "difficulty": "medium"}
|
| 134 |
+
{"prompt": "SEED:1224011538 | Support escalation: notifications is intermittently failing for EU users. Triage and route per org process.", "seed": 1224011538, "difficulty": "medium"}
|
| 135 |
+
{"prompt": "SEED:1881988384 | A user reported that the payments is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1881988384, "difficulty": "medium"}
|
| 136 |
+
{"prompt": "SEED:33599984 | Support escalation: search is intermittently failing for EU users. Triage and route per org process.", "seed": 33599984, "difficulty": "medium"}
|
| 137 |
+
{"prompt": "SEED:444900634 | A user reported that the checkout is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 444900634, "difficulty": "medium"}
|
| 138 |
+
{"prompt": "SEED:1135872495 | A user reported that the api-gateway is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 1135872495, "difficulty": "easy"}
|
| 139 |
+
{"prompt": "SEED:459716009 | A user reported that the checkout is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 459716009, "difficulty": "medium"}
|
| 140 |
+
{"prompt": "SEED:1169726681 | Monitoring alert: elevated error rate on search. File a bug ticket and page the right team.", "seed": 1169726681, "difficulty": "medium"}
|
| 141 |
+
{"prompt": "SEED:904647469 | Support escalation: inventory is intermittently failing for EU users. Triage and route per org process.", "seed": 904647469, "difficulty": "medium"}
|
| 142 |
+
{"prompt": "SEED:874443787 | Monitoring alert: elevated error rate on payments. File a bug ticket and page the right team.", "seed": 874443787, "difficulty": "medium"}
|
| 143 |
+
{"prompt": "SEED:2098228320 | Monitoring alert: elevated error rate on inventory. File a bug ticket and page the right team.", "seed": 2098228320, "difficulty": "medium"}
|
| 144 |
+
{"prompt": "SEED:218179599 | A user reported that the inventory is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 218179599, "difficulty": "easy"}
|
| 145 |
+
{"prompt": "SEED:1819244083 | Monitoring alert: elevated error rate on api-gateway. File a bug ticket and page the right team.", "seed": 1819244083, "difficulty": "medium"}
|
| 146 |
+
{"prompt": "SEED:189351213 | Support escalation: inventory is intermittently failing for EU users. Triage and route per org process.", "seed": 189351213, "difficulty": "easy"}
|
| 147 |
+
{"prompt": "SEED:1432614422 | Monitoring alert: elevated error rate on search. File a bug ticket and page the right team.", "seed": 1432614422, "difficulty": "medium"}
|
| 148 |
+
{"prompt": "SEED:1125089309 | Monitoring alert: elevated error rate on payments. File a bug ticket and page the right team.", "seed": 1125089309, "difficulty": "medium"}
|
| 149 |
+
{"prompt": "SEED:1897669065 | Support escalation: inventory is intermittently failing for EU users. Triage and route per org process.", "seed": 1897669065, "difficulty": "hard"}
|
| 150 |
+
{"prompt": "SEED:41531046 | A user reported that the inventory is returning 500 errors. The issue started 20 minutes ago. Triage this bug: file a ticket with the correct label and priority, assign it to the right team, and notify the correct oncall channel.", "seed": 41531046, "difficulty": "easy"}
|