253 lines
9.0 KiB
Python
253 lines
9.0 KiB
Python
#!/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()
|