deveshshetty/sysadmin-game
0
1"""Gradio app for Sysadmin Game GRPO training on HF Spaces."""2 3import os4import re5import json6import threading7import time8from dataclasses import dataclass, field9from typing import Optional10from collections import defaultdict11 12import gradio as gr13import torch14import matplotlib.pyplot as plt15from transformers import AutoModelForCausalLM, AutoTokenizer16 17 18# ============== Configuration ==============19 20@dataclass21class Config:22 model_name: str = "Qwen/Qwen2.5-Coder-3B-Instruct"23 num_steps: int = 5024 episodes_per_step: int = 425 group_size: int = 226 max_turns: int = 827 learning_rate: float = 1e-528 temperature: float = 0.729 max_seq_length: int = 204830 31 32SYSTEM_PROMPT = """You are a Linux sysadmin agent. You diagnose and fix system issues by running shell commands.33 34RULES:35- Run exactly ONE command per response36- Wrap your command in <bash> and </bash> tags37- You may optionally think first in <think> tags38- Do NOT explain, just run the command39 40Example response:41<think>Check what's using port 80</think>42<bash>ss -tlnp | grep :80</bash>43 44Another example:45<bash>systemctl status nginx</bash>46 47ALWAYS use <bash>command</bash> format. Never use markdown code blocks."""48 49 50# ============== Training State ==============51 52class TrainingState:53 def __init__(self):54 self.is_training = False55 self.current_step = 056 self.total_steps = 057 self.history = {"steps": [], "rewards": [], "fix_rates": [], "losses": []}58 self.logs = []59 self.model = None60 self.tokenizer = None61 62 def log(self, msg: str):63 self.logs.append(f"[{time.strftime('%H:%M:%S')}] {msg}")64 if len(self.logs) > 500:65 self.logs = self.logs[-500:]66 67state = TrainingState()68 69 70# ============== Model Loading ==============71 72def load_model(model_name: str, progress=gr.Progress()):73 """Load model for training."""74 try:75 progress(0.1, desc="Loading tokenizer...")76 state.log(f"Loading model: {model_name}")77 78 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)79 if tokenizer.pad_token is None:80 tokenizer.pad_token = tokenizer.eos_token81 tokenizer.padding_side = "left"82 state.log(f"Tokenizer loaded, vocab size: {len(tokenizer)}")83 84 progress(0.3, desc="Loading model...")85 86 device = "cuda" if torch.cuda.is_available() else "cpu"87 # Use bfloat16 for training stability (no GradScaler needed unlike fp16)88 if device == "cuda" and torch.cuda.is_bf16_supported():89 dtype = torch.bfloat1690 elif device == "cuda":91 dtype = torch.float32 # fall back to fp32 if bf16 not supported92 else:93 dtype = torch.float3294 state.log(f"Using device: {device}, dtype: {dtype}")95 96 model = AutoModelForCausalLM.from_pretrained(97 model_name,98 torch_dtype=dtype,99 device_map="auto" if device == "cuda" else None,100 low_cpu_mem_usage=True,101 trust_remote_code=True,102 )103 104 # Don't resize embeddings — Qwen models already match tokenizer size105 state.log(f"Model vocab: {model.config.vocab_size}, tokenizer vocab: {len(tokenizer)}")106 107 # Enable gradient checkpointing to reduce VRAM usage during training108 model.gradient_checkpointing_enable()109 state.log("Gradient checkpointing enabled (saves ~40% VRAM during training)")110 111 if device == "cpu":112 model = model.to(device)113 114 state.model = model115 state.tokenizer = tokenizer116 state.log(f"Model loaded successfully on {device}")117 118 progress(1.0, desc="Done!")119 return f"✅ Model loaded: {model_name} on {device}"120 121 except Exception as e:122 import traceback123 err_msg = f"{type(e).__name__}: {str(e)}"124 state.log(f"Model loading FAILED: {err_msg}")125 state.log(f"Traceback: {traceback.format_exc()[-500:]}")126 return f"❌ Failed to load model: {err_msg}"127 128 129# ============== Real HTTP Environment ==============130 131class RealEnvHTTP:132 """HTTP client for the real Sysadmin Game environment server.133 134 Mirrors the SimulatedEnv dict-based interface so run_episode works135 unchanged whether we're talking to a real Docker sandbox or the fake sim.136 """137 138 def __init__(self, base_url: str):139 self.base_url = base_url.rstrip("/")140 self.episode_id: Optional[str] = None141 142 def _post(self, endpoint: str, payload: dict) -> dict:143 from urllib import request as urlrequest144 from urllib.error import URLError, HTTPError145 import traceback146 147 url = f"{self.base_url}{endpoint}"148 data = json.dumps(payload).encode("utf-8")149 headers = {"Content-Type": "application/json", "Accept": "application/json"}150 req = urlrequest.Request(url, data=data, headers=headers, method="POST")151 152 try:153 with urlrequest.urlopen(req, timeout=90.0) as resp:154 return json.loads(resp.read().decode("utf-8"))155 except HTTPError as e:156 # Read error body for details157 error_body = ""158 try:159 error_body = e.read().decode("utf-8")160 except:161 pass162 state.log(f"HTTP {e.code} from {endpoint}: {error_body[:200]}")163 raise RuntimeError(f"HTTP {e.code} from {endpoint}: {error_body[:100]}") from e164 except URLError as e:165 state.log(f"Connection failed to {url}: {e.reason}")166 raise RuntimeError(f"Cannot connect to {url}: {e.reason}") from e167 except json.JSONDecodeError as e:168 state.log(f"Invalid JSON from {endpoint}: {str(e)}")169 raise RuntimeError(f"Invalid JSON response from {endpoint}") from e170 except Exception as e:171 state.log(f"Request failed {endpoint}: {type(e).__name__}: {str(e)}")172 raise173 174 def reset(self, scenario_id: Optional[str] = None) -> dict:175 try:176 resp = self._post("/reset", {"scenario_id": scenario_id})177 self.episode_id = resp["metadata"]["episode_id"]178 return {179 "output": resp["output"],180 "done": resp["done"],181 "reward": resp["reward"],182 "metadata": resp["metadata"],183 }184 except Exception as e:185 state.log(f"Reset failed: {type(e).__name__}: {str(e)[:100]}")186 raise187 188 def step(self, command: str) -> dict:189 try:190 resp = self._post("/step", {"command": command, "episode_id": self.episode_id})191 return {192 "output": resp["output"],193 "done": resp["done"],194 "reward": resp["reward"],195 "metadata": resp["metadata"],196 }197 except Exception as e:198 state.log(f"Step failed for '{command[:30]}': {type(e).__name__}: {str(e)[:80]}")199 raise200 201 202def check_env_server(env_url: str) -> str:203 """Ping the env server and return a status string."""204 if not env_url.strip():205 return "⚠️ No server URL — will use simulated environment"206 try:207 from urllib import request as urlrequest208 url = env_url.rstrip("/") + "/health"209 with urlrequest.urlopen(url, timeout=5.0) as resp:210 data = json.loads(resp.read().decode("utf-8"))211 scenarios_url = env_url.rstrip("/") + "/scenarios"212 with urlrequest.urlopen(scenarios_url, timeout=5.0) as resp:213 sc = json.loads(resp.read().decode("utf-8"))214 return (215 f"✅ Connected to real environment server\n"216 f"Status: {data.get('status')} Active episodes: {data.get('active_episodes', 0)}\n"217 f"Train scenarios: {', '.join(sc.get('train', []))}"218 )219 except Exception as e:220 return f"❌ Cannot reach server at {env_url}\nError: {e}"221 222 223# ============== Simulated Environment (fallback) ==============224 225SIMULATED_SCENARIOS = [226 {227 "id": "nginx_syntax",228 "complaint": "nginx won't start, getting config errors",229 "solution_pattern": r"nginx -t|vim.*nginx|nano.*nginx|systemctl restart nginx",230 "diagnostics": ["systemctl status nginx", "nginx -t", "cat /etc/nginx"],231 },232 {233 "id": "disk_full",234 "complaint": "Can't save files, disk full errors everywhere",235 "solution_pattern": r"rm |find.*-delete|truncate|df -h",236 "diagnostics": ["df -h", "du -sh", "ls -la /var/log"],237 },238 {239 "id": "port_bound",240 "complaint": "nginx says address already in use on port 80",241 "solution_pattern": r"kill|systemctl stop|fuser -k",242 "diagnostics": ["ss -tlnp", "netstat -tlnp", "lsof -i :80"],243 },244]245 246 247class SimulatedEnv:248 """Fallback when no real server URL is provided."""249 250 def __init__(self):251 self.scenario = None252 self.commands_run = []253 self.fixed = False254 255 def reset(self, scenario_id=None) -> dict:256 import random257 if scenario_id:258 self.scenario = next(259 (s for s in SIMULATED_SCENARIOS if s["id"] == scenario_id),260 random.choice(SIMULATED_SCENARIOS),261 )262 else:263 self.scenario = random.choice(SIMULATED_SCENARIOS)264 self.commands_run = []265 self.fixed = False266 return {267 "output": self.scenario["complaint"],268 "done": False,269 "reward": 0.0,270 "metadata": {"scenario_id": self.scenario["id"]},271 }272 273 def step(self, command: str) -> dict:274 self.commands_run.append(command)275 reward = -0.01276 277 for diag in self.scenario["diagnostics"]:278 if diag.split()[0] in command:279 reward += 0.1280 break281 282 if re.search(self.scenario["solution_pattern"], command, re.IGNORECASE):283 self.fixed = True284 reward += 1.0285 286 done = self.fixed or len(self.commands_run) >= 15287 return {288 "output": f"[Simulated output for: {command}]",289 "done": done,290 "reward": reward,291 "metadata": {"scenario_id": self.scenario["id"], "fixed": self.fixed},292 }293 294 295# ============== Episode Runner ==============296 297def parse_response(response: str) -> Optional[str]:298 """Extract command from model response. Handles multiple formats."""299 # Try <bash>...</bash> tags first (preferred)300 bash_match = re.search(r"<bash>(.*?)</bash>", response, re.DOTALL)301 if bash_match:302 return bash_match.group(1).strip()303 304 # Try ```bash or ```sh code blocks305 code_match = re.search(r"```(?:bash|sh|shell)?\n?(.*?)```", response, re.DOTALL)306 if code_match:307 cmd = code_match.group(1).strip()308 # Take only first line if multiple commands309 return cmd.split("\n")[0].strip()310 311 # Try single backtick `command`312 tick_match = re.search(r"`([^`]+)`", response)313 if tick_match:314 cmd = tick_match.group(1).strip()315 # Only accept if it looks like a command (starts with common commands)316 cmd_starters = ("ls", "cat", "grep", "find", "ps", "ss", "netstat", "df",317 "du", "systemctl", "service", "nginx", "kill", "rm", "mv",318 "cp", "chmod", "chown", "apt", "yum", "pip", "docker",319 "journalctl", "tail", "head", "less", "more", "lsof",320 "free", "top", "htop", "mount", "umount", "fdisk",321 "curl", "wget", "ssh", "scp", "tar", "gzip", "fuser",322 "truncate", "echo", "sudo", "id", "whoami", "stat")323 if any(cmd.startswith(s) for s in cmd_starters):324 return cmd325 326 # Last resort: look for lines starting with $ or # (shell prompts)327 prompt_match = re.search(r"^[\$#]\s*(.+)$", response, re.MULTILINE)328 if prompt_match:329 return prompt_match.group(1).strip()330 331 return None332 333 334def run_episode(env, model, tokenizer, config: Config, episode_num: int = 0) -> dict:335 """Run single episode against env (real or simulated)."""336 import traceback337 338 model.eval() # Disable gradient checkpointing for faster generation339 340 try:341 obs = env.reset()342 except Exception as e:343 state.log(f" Episode {episode_num}: RESET FAILED - {type(e).__name__}: {str(e)[:100]}")344 raise RuntimeError(f"Episode reset failed: {e}") from e345 346 scenario_id = obs["metadata"]["scenario_id"]347 state.log(f" Episode {episode_num}: scenario={scenario_id}")348 349 messages = [350 {"role": "system", "content": SYSTEM_PROMPT},351 {"role": "user", "content": obs["output"]},352 ]353 354 trajectory = {355 "scenario_id": scenario_id,356 "prompts": [],357 "responses": [],358 "rewards": [],359 "commands": [],360 }361 362 total_reward = 0.0363 364 for turn in range(config.max_turns):365 try:366 prompt = tokenizer.apply_chat_template(367 messages, tokenize=False, add_generation_prompt=True368 )369 inputs = tokenizer(370 prompt,371 return_tensors="pt",372 truncation=True,373 max_length=config.max_seq_length - 200,374 ).to(model.device)375 376 with torch.no_grad():377 outputs = model.generate(378 **inputs,379 max_new_tokens=200,380 temperature=config.temperature,381 do_sample=True,382 pad_token_id=tokenizer.pad_token_id,383 eos_token_id=tokenizer.eos_token_id,384 )385 386 response = tokenizer.decode(387 outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True388 )389 command = parse_response(response)390 391 trajectory["prompts"].append(prompt)392 trajectory["responses"].append(response)393 trajectory["commands"].append(command)394 395 if not command:396 state.log(f" Turn {turn}: no command extracted, penalty -0.1")397 trajectory["rewards"].append(-0.1)398 total_reward -= 0.1399 break400 401 obs = env.step(command)402 reward = obs["reward"]403 trajectory["rewards"].append(reward)404 total_reward += reward405 406 # Log command and reward407 cmd_short = command[:40] + "..." if len(command) > 40 else command408 state.log(f" Turn {turn}: `{cmd_short}` → reward={reward:+.2f}")409 410 messages.append({"role": "assistant", "content": response})411 messages.append(412 {"role": "tool", "content": f"<output>\n{obs['output']}\n</output>"}413 )414 415 if obs["done"]:416 fixed = obs["metadata"].get("fixed", False)417 reason = obs["metadata"].get("termination_reason", "unknown")418 state.log(f" Done: fixed={fixed}, reason={reason}, total={total_reward:+.2f}")419 break420 421 except Exception as e:422 import traceback423 err_type = type(e).__name__424 err_msg = str(e)[:100]425 state.log(f" Turn {turn} ERROR: {err_type}: {err_msg}")426 # Log full traceback for debugging427 tb_lines = traceback.format_exc().split('\n')[-4:-1]428 for line in tb_lines:429 if line.strip():430 state.log(f" {line.strip()}")431 trajectory["rewards"].append(-0.1)432 total_reward -= 0.1433 break434 435 trajectory["total_reward"] = total_reward436 trajectory["fixed"] = obs["metadata"].get("fixed", False)437 trajectory["num_commands"] = len([c for c in trajectory["commands"] if c])438 return trajectory439 440 441# ============== GRPO Training ==============442 443def compute_advantages(trajectories: list, group_size: int) -> list:444 """Compute GRPO advantages (normalize within group)."""445 processed = []446 for i in range(0, len(trajectories) - group_size + 1, group_size):447 group = trajectories[i : i + group_size]448 rewards = [t["total_reward"] for t in group]449 mean_r = sum(rewards) / len(rewards)450 std_r = max((sum((r - mean_r) ** 2 for r in rewards) / len(rewards)) ** 0.5, 1e-6)451 452 for traj in group:453 advantage = (traj["total_reward"] - mean_r) / std_r454 for prompt, response, reward in zip(455 traj["prompts"], traj["responses"], traj["rewards"]456 ):457 processed.append(458 {"prompt": prompt, "response": response, "advantage": advantage}459 )460 return processed461 462 463def _check_model_health(model) -> bool:464 """Return True if model weights are healthy (no NaN/Inf)."""465 for name, param in model.named_parameters():466 if torch.isnan(param).any() or torch.isinf(param).any():467 state.log(f" ⚠️ NaN/Inf detected in {name} — model is corrupted!")468 return False469 return True470 471 472def train_step(model, tokenizer, env, config: Config, optimizer, step_num: int = 0) -> tuple:473 """Single GRPO training step."""474 state.log(f"Step {step_num}: Collecting {config.episodes_per_step} episodes...")475 476 if torch.cuda.is_available():477 torch.cuda.empty_cache()478 479 # Check model health before running episodes480 if not _check_model_health(model):481 state.log(f"Step {step_num}: SKIPPING — model weights are corrupted (NaN/Inf). Stop training and reload model.")482 return None, None, None483 484 trajectories = []485 consecutive_failures = 0486 for ep_idx in range(config.episodes_per_step):487 try:488 traj = run_episode(env, model, tokenizer, config, episode_num=ep_idx)489 trajectories.append(traj)490 consecutive_failures = 0 # Reset on success491 except Exception as e:492 consecutive_failures += 1493 err_type = type(e).__name__494 state.log(f" Episode {ep_idx} FAILED: {err_type}: {str(e)[:100]}")495 if torch.cuda.is_available():496 torch.cuda.empty_cache()497 # If too many consecutive failures, likely server is down498 if consecutive_failures >= 3:499 state.log(f" ⚠️ {consecutive_failures} consecutive failures - check environment server!")500 break501 502 if not trajectories:503 state.log(f"Step {step_num}: No trajectories collected!")504 return None, None, None505 506 # Summarize episodes507 rewards = [t["total_reward"] for t in trajectories]508 fix_count = sum(1 for t in trajectories if t["fixed"])509 fix_rate = fix_count / len(trajectories)510 avg_reward = sum(rewards) / len(rewards)511 total_cmds = sum(t.get("num_commands", 0) for t in trajectories)512 513 state.log(f"Step {step_num}: Episodes done - avg_reward={avg_reward:.3f}, fixed={fix_count}/{len(trajectories)}, commands={total_cmds}")514 515 training_data = compute_advantages(trajectories, config.group_size)516 if not training_data:517 state.log(f"Step {step_num}: No training data after advantage computation")518 return avg_reward, fix_rate, 0.0519 520 state.log(f"Step {step_num}: Training on {len(training_data)} examples...")521 522 # Free inference VRAM before training523 if torch.cuda.is_available():524 torch.cuda.empty_cache()525 526 model.train()527 total_loss = 0.0528 valid_examples = 0529 skipped_nan = 0530 skipped_error = 0531 532 # Process one example at a time: forward → backward → step → zero533 # This avoids holding multiple computation graphs in VRAM534 for idx, ex in enumerate(training_data):535 try:536 inputs = tokenizer(537 ex["prompt"] + ex["response"],538 return_tensors="pt",539 truncation=True,540 max_length=config.max_seq_length,541 ).to(model.device)542 543 outputs = model(**inputs, labels=inputs.input_ids)544 loss = outputs.loss * ex["advantage"]545 546 if torch.isnan(loss) or torch.isinf(loss):547 skipped_nan += 1548 del inputs, outputs, loss549 continue550 551 loss.backward()552 total_loss += loss.item()553 valid_examples += 1554 555 # Step after each example to free graph memory immediately556 grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)557 if not (torch.isnan(grad_norm) or torch.isinf(grad_norm)):558 optimizer.step()559 optimizer.zero_grad()560 561 # Free memory562 del inputs, outputs, loss563 if torch.cuda.is_available():564 torch.cuda.empty_cache()565 566 except Exception as e:567 skipped_error += 1568 if skipped_error <= 2:569 state.log(f" Training example {idx} error: {type(e).__name__}: {str(e)[:60]}")570 optimizer.zero_grad()571 if torch.cuda.is_available():572 torch.cuda.empty_cache()573 574 575 avg_loss = total_loss / max(valid_examples, 1)576 577 # Log summary with any issues578 issues = []579 if skipped_nan > 0:580 issues.append(f"{skipped_nan} NaN")581 if skipped_error > 0:582 issues.append(f"{skipped_error} errors")583 issue_str = f" (skipped: {', '.join(issues)})" if issues else ""584 585 state.log(f"Step {step_num}: Loss={avg_loss:.4f} (from {valid_examples} examples){issue_str}")586 587 return avg_reward, fix_rate, avg_loss588 589 590def run_training(591 num_steps: int,592 episodes_per_step: int,593 learning_rate: float,594 env_url: str,595):596 """Main training loop — yields live updates after each step.597 598 Uses real HTTP env if env_url is set, else simulated.599 This is a generator so Gradio streams status/plot/logs in real-time.600 """601 if state.model is None:602 yield "❌ Load a model first!", None, ""603 return604 605 config = Config(606 num_steps=int(num_steps),607 episodes_per_step=int(episodes_per_step),608 learning_rate=float(learning_rate),609 )610 611 state.is_training = True612 state.total_steps = config.num_steps613 state.history = {"steps": [], "rewards": [], "fix_rates": [], "losses": []}614 615 # Pick environment: real if URL provided, fake otherwise616 if env_url.strip():617 env = RealEnvHTTP(env_url.strip())618 env_label = f"Real server: {env_url.strip()}"619 else:620 env = SimulatedEnv()621 env_label = "Simulated (no server URL)"622 623 optimizer = torch.optim.AdamW(state.model.parameters(), lr=config.learning_rate)624 625 state.log(f"Starting training: {config.num_steps} steps, {config.episodes_per_step} eps/step")626 state.log(f"Environment: {env_label}")627 628 # Yield initial state so user sees logs immediately629 yield f"⏳ Starting training...", None, "\n".join(state.logs[-50:])630 631 try:632 for step in range(config.num_steps):633 if not state.is_training:634 state.log("Training stopped by user")635 break636 637 state.current_step = step + 1638 639 try:640 reward, fix_rate, loss = train_step(641 state.model, state.tokenizer, env, config, optimizer, step_num=step + 1642 )643 644 if reward is not None:645 state.history["steps"].append(step + 1)646 state.history["rewards"].append(reward)647 state.history["fix_rates"].append(fix_rate)648 state.history["losses"].append(loss)649 state.log(650 f"═══ Step {step+1} Summary: reward={reward:.3f}, fix={fix_rate:.1%}, loss={loss:.4f} ═══"651 )652 except Exception as e:653 import traceback654 state.log(f"Step {step+1} FAILED: {type(e).__name__}: {str(e)[:100]}")655 state.log(f" Traceback: {traceback.format_exc()[-300:]}")656 # Continue to next step instead of crashing657 658 # Yield after every step so UI updates live659 fig = create_training_plot()660 yield (661 f"⏳ Step {step+1}/{config.num_steps}",662 fig,663 "\n".join(state.logs[-50:]),664 )665 plt.close(fig) if fig else None666 667 state.is_training = False668 state.log("Training complete!")669 fig = create_training_plot()670 final_reward = state.history["rewards"][-1] if state.history["rewards"] else 0.0671 yield f"✅ Training complete! Final reward: {final_reward:.3f}", fig, "\n".join(state.logs[-50:])672 673 except Exception as e:674 import traceback675 state.is_training = False676 err_msg = f"{type(e).__name__}: {str(e)}"677 state.log(f"Training CRASHED: {err_msg}")678 state.log(f"Full traceback:\n{traceback.format_exc()}")679 fig = create_training_plot()680 yield f"❌ Training failed: {err_msg}", fig, "\n".join(state.logs[-50:])681 682 683def stop_training():684 state.is_training = False685 state.log("Training stopped by user")686 return "Training stopped"687 688 689def create_training_plot():690 """Create training curves plot."""691 if not state.history["steps"]:692 return None693 694 fig, axes = plt.subplots(1, 3, figsize=(15, 4))695 696 axes[0].plot(state.history["steps"], state.history["rewards"], "b-", lw=2)697 axes[0].set_xlabel("Step")698 axes[0].set_ylabel("Reward")699 axes[0].set_title("Average Reward")700 axes[0].grid(True, alpha=0.3)701 702 axes[1].plot(state.history["steps"], state.history["fix_rates"], "g-", lw=2)703 axes[1].set_xlabel("Step")704 axes[1].set_ylabel("Fix Rate")705 axes[1].set_title("Success Rate")706 axes[1].set_ylim(0, 1)707 axes[1].grid(True, alpha=0.3)708 709 axes[2].plot(state.history["steps"], state.history["losses"], "r-", lw=2)710 axes[2].set_xlabel("Step")711 axes[2].set_ylabel("Loss")712 axes[2].set_title("Policy Loss")713 axes[2].grid(True, alpha=0.3)714 715 plt.tight_layout()716 return fig717 718 719# ============== Demo Tab ==============720 721def run_demo(complaint: str):722 """Run a single demo episode."""723 if state.model is None:724 return "Load a model first!"725 726 messages = [727 {"role": "system", "content": SYSTEM_PROMPT},728 {"role": "user", "content": complaint},729 ]730 731 prompt = state.tokenizer.apply_chat_template(732 messages, tokenize=False, add_generation_prompt=True733 )734 inputs = state.tokenizer(prompt, return_tensors="pt").to(state.model.device)735 736 with torch.no_grad():737 outputs = state.model.generate(738 **inputs,739 max_new_tokens=300,740 temperature=0.7,741 do_sample=True,742 pad_token_id=state.tokenizer.eos_token_id,743 )744 745 response = state.tokenizer.decode(746 outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True747 )748 return response749 750 751# ============== Gradio UI ==============752 753with gr.Blocks(title="Sysadmin Game - GRPO Training", theme=gr.themes.Soft()) as demo:754 gr.Markdown("# 🛠️ Sysadmin Game: GRPO Training")755 gr.Markdown(756 "Train LLMs to diagnose and fix Linux systems using reinforcement learning.\n\n"757 "**With real environment:** Start your environment server locally, expose it via "758 "ngrok, paste the URL below. **Without:** falls back to a simulated environment."759 )760 761 with gr.Tabs():762 # ── Setup Tab ──────────────────────────────────────────────────────────763 with gr.Tab("1. Setup"):764 gr.Markdown("### Load Model")765 with gr.Row():766 model_dropdown = gr.Dropdown(767 choices=[768 "Qwen/Qwen2.5-Coder-0.5B-Instruct",769 "Qwen/Qwen2.5-Coder-1.5B-Instruct",770 "Qwen/Qwen2.5-Coder-3B-Instruct",771 ],772 value="Qwen/Qwen2.5-Coder-0.5B-Instruct",773 label="Model",774 )775 load_btn = gr.Button("Load Model", variant="primary")776 load_status = gr.Textbox(label="Status", interactive=False)777 778 gr.Markdown("### Environment Server (optional — for real Docker sandbox)")779 with gr.Row():780 env_url_input = gr.Textbox(781 label="Environment Server URL",782 placeholder="https://abc123.ngrok-free.app (leave empty for simulated env)",783 scale=4,784 )785 check_btn = gr.Button("Check Connection", scale=1)786 env_status = gr.Textbox(label="Environment Status", interactive=False, lines=3)787 788 load_btn.click(load_model, inputs=[model_dropdown], outputs=[load_status])789 check_btn.click(check_env_server, inputs=[env_url_input], outputs=[env_status])790 791 # ── Training Tab ───────────────────────────────────────────────────────792 with gr.Tab("2. Training"):793 gr.Markdown("### GRPO Training Configuration")794 with gr.Row():795 num_steps = gr.Slider(10, 200, value=50, step=10, label="Training Steps")796 episodes = gr.Slider(2, 16, value=4, step=2, label="Episodes per Step")797 lr = gr.Number(value=1e-5, label="Learning Rate")798 799 env_url_train = gr.Textbox(800 label="Environment Server URL (copy from Setup tab)",801 placeholder="https://abc123.ngrok-free.app or leave empty for simulated",802 )803 804 with gr.Row():805 train_btn = gr.Button("Start Training", variant="primary")806 stop_btn = gr.Button("Stop", variant="stop")807 808 train_status = gr.Textbox(label="Status", interactive=False)809 with gr.Row():810 train_logs = gr.Textbox(811 label="Logs", lines=15, interactive=False,812 autoscroll=True, scale=3,813 )814 train_plot = gr.Plot(label="Training Curves", scale=2)815 816 train_btn.click(817 run_training,818 inputs=[num_steps, episodes, lr, env_url_train],819 outputs=[train_status, train_plot, train_logs],820 )821 stop_btn.click(stop_training, outputs=[train_status])822 823 # ── Demo Tab ───────────────────────────────────────────────────────────824 with gr.Tab("3. Demo"):825 gr.Markdown("### Test the Model")826 complaint_input = gr.Textbox(827 label="User Complaint",828 placeholder="e.g., nginx won't start, getting config errors",829 lines=2,830 )831 demo_btn = gr.Button("Get Diagnosis", variant="primary")832 demo_output = gr.Textbox(label="Model Response", lines=10)833 834 demo_btn.click(run_demo, inputs=[complaint_input], outputs=[demo_output])835 836 gr.Examples(837 examples=[838 ["nginx won't start, says something about address already in use"],839 ["Can't write any files, getting 'No space left on device' errors"],840 ["Getting permission denied when trying to read /var/log/syslog"],841 ],842 inputs=[complaint_input],843 )844 845demo.launch(server_name="0.0.0.0", server_port=7860)846