#!/usr/bin/env python3 """ train_grpo_job.py — Self-contained GRPO training job for HF Jobs. Loads dataset from HF Hub, runs GRPO training with custom reward functions, pushes model to Hub on completion via HfApi.upload_folder(). """ import logging import os import re import shutil import subprocess import tempfile from pathlib import Path import torch from datasets import load_dataset from trl import GRPOTrainer, GRPOConfig logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") logger = logging.getLogger(__name__) # ─── Config ─────────────────────────────────────────────────────────────────── MODEL_NAME = "Qwen/Qwen2.5-Coder-0.5B-Instruct" DATASET_ID = "oxdev/smart-contract-security-sft" OUTPUT_DIR = "/tmp/grpo_output" HUB_MODEL_ID = "oxdev/security-auditor-grpo" FORGE_AVAILABLE = shutil.which("forge") is not None # ─── Reward Functions ───────────────────────────────────────────────────────── def extract_finding_block(text: str) -> dict | None: pattern = re.compile( r'FINDING\s*\|\s*contract:\s*(\S+)\s*\|\s*function:\s*(\S+)\s*\|' r'\s*bug_class:\s*(\S+)\s*\|\s*confidence:\s*(\d+)', re.IGNORECASE ) match = pattern.search(text) if not match: return None return { "contract": match.group(1), "function": match.group(2), "bug_class": match.group(3), "confidence": int(match.group(4)), } def extract_solidity_poc(text: str) -> str | None: pattern = re.compile(r'```solidity\s*\n(.*?)```', re.DOTALL) matches = pattern.findall(text) if not matches: return None for code in matches: if "is Test" in code or "function test_" in code: return code.strip() return max(matches, key=len).strip() if matches else None def _check_solidity_syntax(code: str) -> bool: required = [r'pragma\s+solidity', r'contract\s+\w+', r'function\s+\w+'] return all(re.search(p, code) for p in required) def run_forge_test(poc_code: str, timeout: int = 30) -> dict: if not FORGE_AVAILABLE: return { "compiled": False, "test_passed": False, "syntax_valid": _check_solidity_syntax(poc_code), } tmpdir = tempfile.mkdtemp(prefix="forge_poc_") try: test_dir = Path(tmpdir) / "test" test_dir.mkdir() (Path(tmpdir) / "foundry.toml").write_text('[profile.default]\nsrc = "src"\nout = "out"\nlibs = ["lib"]\nsolc_version = "0.8.24"\n') (Path(tmpdir) / "src").mkdir() try: subprocess.run( ["forge", "install", "foundry-rs/forge-std", "--no-git", "--no-commit"], cwd=tmpdir, capture_output=True, timeout=60, ) except Exception: pass (Path(tmpdir) / "remappings.txt").write_text("forge-std/=lib/forge-std/src/\n") (test_dir / "PoC.t.sol").write_text(poc_code) build = subprocess.run(["forge", "build"], cwd=tmpdir, capture_output=True, text=True, timeout=timeout) if build.returncode != 0: return {"compiled": False, "test_passed": False} test = subprocess.run(["forge", "test", "-vv"], cwd=tmpdir, capture_output=True, text=True, timeout=timeout) return {"compiled": True, "test_passed": test.returncode == 0 and "PASS" in test.stdout} except Exception: return {"compiled": False, "test_passed": False} finally: shutil.rmtree(tmpdir, ignore_errors=True) def security_audit_reward(completions, **kwargs): """Primary reward: FINDING block + PoC compilation + exploit verification.""" rewards = [] finding_count = compile_count = pass_count = 0 for completion in completions: text = completion[0]["content"] if isinstance(completion, list) else str(completion) reward = -1.0 finding = extract_finding_block(text) if finding: finding_count += 1 reward = 0.0 poc = extract_solidity_poc(text) if poc: reward = 0.2 result = run_forge_test(poc) if result.get("compiled") or result.get("syntax_valid", False): compile_count += 1 reward = 0.5 if result.get("test_passed"): pass_count += 1 reward = 1.0 elif any(kw in text.lower() for kw in ["vulnerability", "exploit", "bug", "finding"]): reward = -0.5 rewards.append(reward) n = len(rewards) if rewards else 1 logger.info(f"[reward] finding_rate={finding_count/n:.2f} compile_rate={compile_count/n:.2f} exploit_rate={pass_count/n:.2f}") return rewards def format_reward(completions, **kwargs): """Secondary reward: structural format compliance.""" rewards = [] for completion in completions: text = completion[0]["content"] if isinstance(completion, list) else str(completion) reward = 0.0 if re.search(r'FINDING\s*\|', text): fields = sum(bool(re.search(p, text)) for p in [r'path:', r'proof:', r'description:', r'fix:']) reward = 0.3 + (0.05 * fields) if re.search(r'```solidity', text): reward += 0.1 rewards.append(reward) return rewards # ─── Main ───────────────────────────────────────────────────────────────────── def main(): logger.info("=" * 60) logger.info("GRPO Training — Smart Contract Security Auditor") logger.info(f"Model: {MODEL_NAME}") logger.info(f"Dataset: {DATASET_ID}") logger.info(f"Forge available: {FORGE_AVAILABLE}") logger.info(f"GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'}") logger.info(f"CUDA available: {torch.cuda.is_available()}") if torch.cuda.is_available(): logger.info(f"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB") logger.info("=" * 60) # Load dataset logger.info("Loading dataset from HF Hub...") dataset = load_dataset(DATASET_ID, split="train") logger.info(f"Dataset: {len(dataset)} samples, columns={dataset.column_names}") # Configure GRPO — NO hub_model_id, NO log_completions, NO push_to_hub # This prevents ANY Hub calls during __init__ or training config = GRPOConfig( output_dir=OUTPUT_DIR, num_train_epochs=2, per_device_train_batch_size=2, gradient_accumulation_steps=2, num_generations=2, max_completion_length=512, learning_rate=5e-7, beta=0.0, scale_rewards=True, reward_weights=[0.7, 0.3], gradient_checkpointing=True, bf16=True, logging_steps=5, logging_first_step=True, logging_strategy="steps", disable_tqdm=True, save_strategy="steps", save_steps=50, save_total_limit=2, # CRITICAL: all Hub-related settings OFF to prevent 401 at init push_to_hub=False, log_completions=False, report_to="none", seed=42, ) # Train logger.info("Initializing GRPOTrainer...") trainer = GRPOTrainer( model=MODEL_NAME, args=config, reward_funcs=[security_audit_reward, format_reward], train_dataset=dataset, ) logger.info("GRPOTrainer initialized successfully!") logger.info("Starting training...") trainer.train() logger.info("Training complete!") # Save locally logger.info(f"Saving model to {OUTPUT_DIR}...") trainer.save_model(OUTPUT_DIR) # Manual push to hub using HfApi — safer and more explicit hf_token = os.environ.get("HF_TOKEN") if hf_token: logger.info(f"Pushing to hub: {HUB_MODEL_ID}") try: from huggingface_hub import HfApi api = HfApi(token=hf_token) # Create repo if needed (ignore error if exists) try: api.create_repo(repo_id=HUB_MODEL_ID, exist_ok=True) except Exception as e: logger.warning(f"create_repo warning (may already exist): {e}") # Upload entire output folder api.upload_folder( folder_path=OUTPUT_DIR, repo_id=HUB_MODEL_ID, commit_message="GRPO training complete — smart contract security auditor", ) logger.info(f"✅ Model pushed to https://huggingface.co/{HUB_MODEL_ID}") except Exception as e: logger.error(f"Push failed: {e}") logger.info(f"Model saved locally at {OUTPUT_DIR}") else: logger.warning("No HF_TOKEN found — model saved locally only") logger.info(f"Model at: {OUTPUT_DIR}") logger.info("=" * 60) logger.info("DONE") logger.info("=" * 60) if __name__ == "__main__": main()