Aditya Guntur commited on
Commit
3d2e8d7
·
2 Parent(s): 36068f1818155d

fix: resolve merge conflict, restore HF frontmatter + full README

Browse files
.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
+ ![PM-Ops Demo](https://drive.google.com/uc?export=view&id=1JdwYukKrEaMTaOwc1W4Q8bBRjxwGZ2be)
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.client import EnvClient
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
- tok = token or os.getenv("HF_TOKEN")
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=self._scenario["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=self._scenario["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": self._ticketing.create_ticket,
174
- "update_ticket": self._ticketing.update_ticket,
175
- "get_ticket": self._ticketing.get_ticket,
176
- "list_tickets": self._ticketing.list_tickets,
177
- "assign_ticket": self._ticketing.assign_ticket,
178
- "comment_ticket": self._ticketing.comment_ticket,
179
- "transition_ticket": self._ticketing.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": self._codebase.list_commits,
189
- "get_commit": self._codebase.get_commit,
190
- "list_prs": self._codebase.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": self._chat.post_message,
200
- "read_channel": self._chat.read_channel,
201
- "list_channels": self._chat.list_channels,
202
- "search": self._chat.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"}