"""Bold full-stack training pipeline: SFT -> DPO -> GRPO -> Adv DPO -> Eval. ONE script. Two modes via --smoke flag: --smoke : 5 steps per stage, ~10 min total. Validates code path. Free Colab. (default) : full step counts, ~3-4h on H200. Production. Usage in Colab T4 (smoke): !python scripts/bold_pipeline.py --smoke Usage on HF Jobs H200 (production): hf jobs uv run --flavor h200 \\ https://huggingface.co/spaces/sai1906/opsguard/resolve/main/scripts/bold_pipeline.py """ from __future__ import annotations import argparse import gc import inspect import json import os import random import sys from pathlib import Path def _wait_for_cuda(): import time as _t print("[cuda] sleeping 30s for driver init...", flush=True) _t.sleep(30) try: import torch n = torch.cuda.device_count() if torch.cuda.is_available() else 0 if n > 0: print(f"[cuda] ready: {n} device(s), {torch.cuda.get_device_name(0)}", flush=True) else: print("[cuda] device_count=0 after sleep — letting model load trigger lazy init", flush=True) except Exception as e: print(f"[cuda] probe error (non-fatal, model load will retry): {e}", flush=True) def _setup_workdir(): import subprocess work = Path("/tmp/opsguard") if not (work / "data" / "sft_traces.jsonl").exists(): if work.exists(): subprocess.run(["rm", "-rf", str(work)], check=True) token = os.environ.get("HF_TOKEN", "") url = (f"https://user:{token}@huggingface.co/spaces/sai1906/opsguard" if token else "https://huggingface.co/spaces/sai1906/opsguard") subprocess.run(["git", "clone", "--depth", "1", url, str(work)], check=True) sys.path.insert(0, str(work)) os.chdir(work) needed = work / "data" / "sft_traces.jsonl" if not needed.exists(): raise FileNotFoundError( f"{needed} missing after clone. cwd={os.getcwd()} " f"contents={list(work.iterdir()) if work.exists() else 'workdir-missing'}" ) print(f"[workdir] cwd={os.getcwd()} data files OK", flush=True) return work _wait_for_cuda() WORK = _setup_workdir() def _safe_kwargs(cls, kw): sig = inspect.signature(cls).parameters return {k: v for k, v in kw.items() if k in sig} def _free(): import torch gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() def _push_lora(folder, repo_id): from huggingface_hub import HfApi api = HfApi(token=os.environ["HF_TOKEN"]) try: api.create_repo(repo_id, repo_type="model", exist_ok=True) except Exception as e: print(f" create_repo({repo_id}): {e}", flush=True) api.upload_folder(folder_path=folder, repo_id=repo_id, repo_type="model") print(f" pushed -> https://huggingface.co/{repo_id}", flush=True) def stage_sft(args, model_name, hub_repo): print("\n" + "=" * 70 + "\n=== STAGE 1: SFT WARMSTART ===\n" + "=" * 70, flush=True) import torch from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training from datasets import Dataset from trl import SFTConfig, SFTTrainer print(f"loading {model_name} (4bit={args.use_4bit})...", flush=True) model_kwargs = {"torch_dtype": torch.bfloat16, "device_map": "auto", "attn_implementation": "sdpa"} if args.use_4bit: bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4") model_kwargs["quantization_config"] = bnb tok = AutoTokenizer.from_pretrained(model_name) if tok.pad_token_id is None: tok.pad_token_id = tok.eos_token_id model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs) model.config.use_cache = False if args.use_4bit: model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True) else: model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) model.enable_input_require_grads() lc = LoraConfig(r=args.lora_r, lora_alpha=args.lora_alpha, lora_dropout=0.05, bias="none", target_modules=["q_proj","k_proj","v_proj","o_proj", "gate_proj","up_proj","down_proj"], task_type="CAUSAL_LM") model = get_peft_model(model, lc) for n, p in model.named_parameters(): if p.requires_grad and ("lora_" in n or "lora_A" in n or "lora_B" in n): p.data = p.data.to(torch.float32) model.print_trainable_parameters() rows = [json.loads(l) for l in open("data/sft_traces.jsonl")] print(f"loaded {len(rows)} SFT traces", flush=True) texts = [r["prompt"] + "\n\nACTION:\n" + r["completion"] + tok.eos_token for r in rows] if args.smoke: texts = texts[:64] ds = Dataset.from_list([{"text": t} for t in texts]) n_epochs = 1 if args.smoke else args.sft_epochs raw = dict( output_dir="/tmp/opsguard-sft", per_device_train_batch_size=args.batch_size, gradient_accumulation_steps=args.grad_accum, num_train_epochs=n_epochs, learning_rate=2e-5, max_grad_norm=1.0, weight_decay=0.01, warmup_ratio=0.1, adam_epsilon=1e-7, optim="adamw_torch", gradient_checkpointing=True, gradient_checkpointing_kwargs={"use_reentrant": False}, logging_steps=2, save_strategy="epoch", save_total_limit=1, bf16=True, max_seq_length=2048, dataset_text_field="text", report_to="none", push_to_hub=False, max_steps=5 if args.smoke else -1, ) cfg = SFTConfig(**_safe_kwargs(SFTConfig, raw)) trainer = SFTTrainer(model=model, train_dataset=ds, args=cfg, processing_class=tok) trainer.train() out_path = "/tmp/opsguard-sft-lora" model.save_pretrained(out_path) tok.save_pretrained(out_path) if hub_repo and not args.smoke: _push_lora(out_path, hub_repo + "-sft") del trainer, model _free() return out_path def stage_dpo(args, model_name, sft_lora_path, hub_repo): print("\n" + "=" * 70 + "\n=== STAGE 2: DPO (R2Vul-style preference pairs) ===\n" + "=" * 70, flush=True) import torch from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig from peft import PeftModel from datasets import Dataset from trl import DPOConfig, DPOTrainer pref_path = "data/preference_pairs.jsonl" if not Path(pref_path).exists(): print(f"WARN: {pref_path} missing — skipping DPO", flush=True) return sft_lora_path rows = [json.loads(l) for l in open(pref_path)] print(f"loaded {len(rows)} preference pairs", flush=True) if args.smoke: rows = rows[:32] ds = Dataset.from_list([ {"prompt": r["prompt"], "chosen": r["chosen"], "rejected": r["rejected"]} for r in rows ]) model_kwargs = {"torch_dtype": torch.bfloat16, "device_map": "auto", "attn_implementation": "sdpa"} if args.use_4bit: bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4") model_kwargs["quantization_config"] = bnb tok = AutoTokenizer.from_pretrained(sft_lora_path) base = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs) base.config.use_cache = False if not args.use_4bit: base.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) base.enable_input_require_grads() model = PeftModel.from_pretrained(base, sft_lora_path, is_trainable=True) for n, p in model.named_parameters(): if p.requires_grad and "lora_" in n: p.data = p.data.to(torch.float32) n_epochs = 1 if args.smoke else args.dpo_epochs raw = dict( output_dir="/tmp/opsguard-dpo", per_device_train_batch_size=1, gradient_accumulation_steps=8, num_train_epochs=n_epochs, learning_rate=5e-6, max_grad_norm=1.0, weight_decay=0.01, adam_epsilon=1e-7, optim="adamw_torch", gradient_checkpointing=True, gradient_checkpointing_kwargs={"use_reentrant": False}, beta=0.1, max_length=2048, max_prompt_length=1500, logging_steps=2, save_strategy="epoch", save_total_limit=1, bf16=True, report_to="none", push_to_hub=False, remove_unused_columns=False, max_steps=5 if args.smoke else -1, ) cfg = DPOConfig(**_safe_kwargs(DPOConfig, raw)) try: trainer = DPOTrainer(model=model, args=cfg, train_dataset=ds, processing_class=tok) except TypeError: trainer = DPOTrainer(model=model, args=cfg, train_dataset=ds, tokenizer=tok) trainer.train() out_path = "/tmp/opsguard-dpo-lora" model.save_pretrained(out_path) tok.save_pretrained(out_path) if hub_repo and not args.smoke: _push_lora(out_path, hub_repo + "-dpo") del trainer, model, base _free() return out_path def stage_grpo(args, model_name, dpo_lora_path, hub_repo): print("\n" + "=" * 70 + "\n=== STAGE 3: GRPO via custom reward fn (no TRL OpenEnv tools) ===\n" + "=" * 70, flush=True) import torch from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig from peft import PeftModel from datasets import Dataset from trl import GRPOConfig, GRPOTrainer from server.opsguard_environment import OpsguardEnvironment from models import OpsguardAction, ActionType from scripts.system_prompt import SYSTEM_PROMPT, format_observation, parse_action model_kwargs = {"torch_dtype": torch.bfloat16, "device_map": "auto", "attn_implementation": "sdpa"} if args.use_4bit: bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4") model_kwargs["quantization_config"] = bnb tok = AutoTokenizer.from_pretrained(dpo_lora_path) if tok.pad_token_id is None: tok.pad_token_id = tok.eos_token_id base = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs) base.config.use_cache = False if not args.use_4bit: base.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) base.enable_input_require_grads() model = PeftModel.from_pretrained(base, dpo_lora_path, is_trainable=True) for n, p in model.named_parameters(): if p.requires_grad and "lora_" in n: p.data = p.data.to(torch.float32) print("building prompt dataset (sampled env observations)...", flush=True) n_prompts = 24 if args.smoke else args.grpo_n_prompts scenarios = ["E2_social_eng_buildup", "E3_compromised_maintainer", "E4_multi_vector"] prompt_rows = [] env = OpsguardEnvironment() for i in range(n_prompts): sid = scenarios[i % len(scenarios)] obs = env.reset(scenario_id=sid, seed=i) for _ in range(random.randint(0, 5)): obs = env.step(OpsguardAction(action_type=ActionType.WAIT)) if obs.done: obs = env.reset(scenario_id=sid, seed=i) break prompt_text = SYSTEM_PROMPT + "\n\nOBSERVATION:\n" + format_observation(obs) + "\n\nACTION:\n" prompt_rows.append({"prompt": prompt_text, "scenario": sid, "seed": i}) ds = Dataset.from_list(prompt_rows) def reward_fn(completions, prompts=None, scenario=None, seed=None, **kwargs): rewards = [] scenarios_b = scenario if scenario is not None else ["E2_social_eng_buildup"] * len(completions) seeds_b = seed if seed is not None else [0] * len(completions) for comp, sid, sd in zip(completions, scenarios_b, seeds_b): try: action = parse_action(comp) env_local = OpsguardEnvironment() env_local.reset(scenario_id=sid, seed=int(sd)) obs = env_local.step(action) rewards.append(float(obs.reward) if obs.reward is not None else 0.0) except Exception: rewards.append(-1.0) return rewards n_steps = 5 if args.smoke else args.grpo_steps raw = dict( output_dir="/tmp/opsguard-grpo", per_device_train_batch_size=1, gradient_accumulation_steps=4, num_generations=args.grpo_num_generations, max_steps=n_steps, max_completion_length=256, max_prompt_length=1500, beta=0.001, learning_rate=5e-6, max_grad_norm=1.0, weight_decay=0.01, adam_epsilon=1e-7, optim="adamw_torch", gradient_checkpointing=True, gradient_checkpointing_kwargs={"use_reentrant": False}, logging_steps=1, save_strategy="no", bf16=True, report_to="none", push_to_hub=False, use_vllm=False, ) cfg = GRPOConfig(**_safe_kwargs(GRPOConfig, raw)) try: trainer = GRPOTrainer(model=model, args=cfg, reward_funcs=reward_fn, train_dataset=ds, processing_class=tok) except TypeError: trainer = GRPOTrainer(model=model, args=cfg, reward_funcs=reward_fn, train_dataset=ds, tokenizer=tok) trainer.train() out_path = "/tmp/opsguard-grpo-lora" model.save_pretrained(out_path) tok.save_pretrained(out_path) if hub_repo and not args.smoke: _push_lora(out_path, hub_repo + "-grpo") del trainer, model, base _free() return out_path def stage_eval(args, model_name, final_lora_path): print("\n" + "=" * 70 + "\n=== STAGE 4: EVAL — trained vs base on Datadog held-out ===\n" + "=" * 70, flush=True) import torch from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig from peft import PeftModel from models import OpsguardAction, ActionType from scripts.system_prompt import SYSTEM_PROMPT, format_observation, parse_action cve_rows = [json.loads(l) for l in open("data/datadog_extracted.jsonl")] if args.smoke: cve_rows = cve_rows[:20] print(f"evaluating on {len(cve_rows)} real Datadog samples...", flush=True) model_kwargs = {"torch_dtype": torch.bfloat16, "device_map": "auto"} if args.use_4bit: bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4") model_kwargs["quantization_config"] = bnb tok = AutoTokenizer.from_pretrained(final_lora_path) base = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs) model = PeftModel.from_pretrained(base, final_lora_path) model.eval() @torch.inference_mode() def predict(text_payload): prompt = SYSTEM_PROMPT + "\n\nOBSERVATION:\n" + text_payload + "\n\nACTION:\n" inp = tok(prompt, return_tensors="pt", truncation=True, max_length=1800).to(model.device) out = model.generate(**inp, max_new_tokens=200, do_sample=False, pad_token_id=tok.eos_token_id) return tok.decode(out[0][inp["input_ids"].shape[1]:], skip_special_tokens=True) REJECT = {"reject_pr", "close_spam", "flag_security"} n_caught = 0 for r in cve_rows: synth_obs = json.dumps({ "current_issue": { "title": r.get("package", "") + " " + r.get("filename", ""), "body": r.get("diff_preview", "")[:1000], "is_pr": True, "author_login": "unknown", "pr_diff_preview": r.get("diff_preview", "")[:800], } }) try: text = predict(synth_obs) action = parse_action(text) at = action.action_type.value if hasattr(action.action_type, "value") else str(action.action_type) if at in REJECT: n_caught += 1 except Exception: pass catch_rate = n_caught / max(1, len(cve_rows)) print(f"\nFINAL: trained LoRA caught {n_caught}/{len(cve_rows)} = {catch_rate:.3f} catch rate on Datadog held-out", flush=True) Path("/tmp/eval_post_train").mkdir(exist_ok=True) Path("/tmp/eval_post_train/cve_results.json").write_text(json.dumps({ "n_samples": len(cve_rows), "n_caught": n_caught, "catch_rate": round(catch_rate, 3), })) return catch_rate def main(): ap = argparse.ArgumentParser() ap.add_argument("--smoke", action="store_true", help="5 steps per stage, ~10 min on Colab T4. Validates code path.") ap.add_argument("--model", type=str, default="Qwen/Qwen2.5-7B-Instruct") ap.add_argument("--use-4bit", action="store_true", help="Use bnb 4-bit (default false = bf16 full). Toggle for VRAM-constrained env.") ap.add_argument("--lora-r", type=int, default=32, help="LoRA rank (default 32 for stability). 128 for bold.") ap.add_argument("--lora-alpha", type=int, default=64, help="LoRA alpha (default 64 = 2*r, mild scaling).") ap.add_argument("--sft-epochs", type=int, default=5) ap.add_argument("--dpo-epochs", type=int, default=3) ap.add_argument("--grpo-steps", type=int, default=200) ap.add_argument("--grpo-num-generations", type=int, default=4) ap.add_argument("--grpo-n-prompts", type=int, default=256) ap.add_argument("--batch-size", type=int, default=2) ap.add_argument("--grad-accum", type=int, default=4) ap.add_argument("--hub-repo", type=str, default="sai1906/opsguard", help="Hub repo prefix; suffixes -sft, -dpo, -grpo appended.") ap.add_argument("--skip", type=str, default="", help="Comma-separated stages to skip: sft,dpo,grpo,eval") args = ap.parse_args() if args.smoke and not args.use_4bit: print("[smoke] forcing --use-4bit on Colab/free GPUs", flush=True) args.use_4bit = True skipped = set(s.strip() for s in args.skip.split(",") if s.strip()) print(f"=== BOLD PIPELINE smoke={args.smoke} 4bit={args.use_4bit} " f"r={args.lora_r} alpha={args.lora_alpha} ===", flush=True) print(f" SFT epochs={args.sft_epochs} DPO epochs={args.dpo_epochs} GRPO steps={args.grpo_steps}", flush=True) sft_lora = stage_sft(args, args.model, args.hub_repo) if "sft" not in skipped else None dpo_lora = stage_dpo(args, args.model, sft_lora or args.model, args.hub_repo) if "dpo" not in skipped else sft_lora grpo_lora = stage_grpo(args, args.model, dpo_lora or sft_lora or args.model, args.hub_repo) if "grpo" not in skipped else dpo_lora final = grpo_lora or dpo_lora or sft_lora if final and "eval" not in skipped: stage_eval(args, args.model, final) print("\n=== DONE ===", flush=True) if __name__ == "__main__": main()