""" OpenEnv GRPO training (production-oriented, simple stack). Produces real training artifacts (trainer log_history, metrics JSON, reward plots) and optional Hub push of LoRA weights. Every execution reward calls your live Space (or local server) at OPENENV_BASE_URL — not a mock. Environment (control cost vs quality on HF Jobs / local GPU): OPENENV_BASE_URL — OpenEnv HTTP root (default: Space URL from openenv.yaml) OPENENV_TASK_IDS — Comma list; if unset, uses GET /tasks from the server ROWS_PER_TASK — GRPO rows per task_id (default: 48) OPENENV_REQUEST_TIMEOUT_SEC — HTTP timeout for reset/step (default: 120) TRAIN_MAX_STEPS — GRPO steps (default 200) TRL_REPORT_TO — none | wandb | tensorboard (auto: wandb if key else none) BOOTSTRAP_*_VERSION — pin transformers / accelerate / trl (defaults satisfy trl>=4.50) Artifacts: artifacts/train_log_history.jsonl, metrics, plots HF_HUB_REPO_ID — push target (default md896/sota-sql-agent-7b) SKIP_HUB_PUSH=1 — do not push after train HF_TOKEN / HUGGING_FACE_HUB_TOKEN — Hub auth for push Designed for Hugging Face Jobs / Spaces where: - system Python may be externally managed (PEP-668) → uses --break-system-packages - preinstalled CUDA/PyTorch stacks can conflict with optional vision packages Key stability choices: - Avoid importing torchvision in text-only runs (it can break when torch/torchvision versions are mismatched by dependency resolution). - Produce plots and metrics from the *actual* GRPO run (no hard-coded scores). """ from __future__ import annotations import json import os import random import subprocess import sys import time import uuid from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, List, Optional def _run(cmd: List[str], *, check: bool = True) -> subprocess.CompletedProcess: return subprocess.run(cmd, check=check) def _pip(args: List[str], *, check: bool = True) -> subprocess.CompletedProcess: return _run([sys.executable, "-m", "pip", *args], check=check) def bootstrap_deps() -> None: """ Best-effort dependency bootstrap for ephemeral HF containers. Set SKIP_BOOTSTRAP=1 to disable. Pins: BOOTSTRAP_TRANSFORMERS_VERSION, BOOTSTRAP_ACCELERATE_VERSION, BOOTSTRAP_TRL_VERSION. """ if os.environ.get("SKIP_BOOTSTRAP") == "1": return # Ensure text-only transformers runs never hard-import torchvision even if it # is present in the base image. os.environ.setdefault("TRANSFORMERS_NO_TORCHVISION", "1") # Ubuntu 24.04+ images may mark system Python as "externally managed" # (PEP-668). Prefer an explicit opt-out for all pip ops in ephemeral jobs. os.environ.setdefault("PIP_BREAK_SYSTEM_PACKAGES", "1") print("Bootstrapping dependencies...") # Text-only run: torchvision/torchaudio are not required and are a common source # of crashes when torch versions shift in container images. _pip(["uninstall", "--break-system-packages", "-y", "torchvision", "torchaudio"], check=False) _pip(["uninstall", "-y", "torchao"], check=False) # trl 0.18.x needs transformers>=4.50. datasets 4.x pulls huggingface-hub 1.x which breaks 4.5x. _tf = os.environ.get("BOOTSTRAP_TRANSFORMERS_VERSION", "4.51.3") _acc = os.environ.get("BOOTSTRAP_ACCELERATE_VERSION", "0.34.2") _trl = os.environ.get("BOOTSTRAP_TRL_VERSION", "0.18.2") _pip( [ "install", "--break-system-packages", "httpx>=0.27.0", "datasets>=3.2.0,<4.0.0", "matplotlib", "tensorboard", f"transformers=={_tf}", f"accelerate=={_acc}", f"trl=={_trl}", ] ) if os.environ.get("WANDB_API_KEY"): _pip(["install", "--break-system-packages", "wandb"], check=False) _pip( [ "install", "--break-system-packages", "--force-reinstall", "--no-deps", f"transformers=={_tf}", f"accelerate=={_acc}", ] ) _pip(["install", "--break-system-packages", "--no-deps", f"trl=={_trl}"]) _pip(["uninstall", "-y", "torchao"], check=False) _pip(["uninstall", "--break-system-packages", "-y", "torchvision", "torchaudio"], check=False) # Keep bootstrap import-free; training imports happen below. bootstrap_deps() import httpx import torch from datasets import Dataset from huggingface_hub import HfApi from transformers import AutoModelForCausalLM, AutoTokenizer from trl import GRPOConfig, GRPOTrainer # --- 1. CONFIGURATION (env-first; defaults match openenv.yaml) --- _DEFAULT_OPENENV_BASE = "https://md896-sql-debug-env.hf.space" BYPASS_HEADERS: Dict[str, str] = {} MODEL_NAME = os.environ.get("TRAIN_MODEL_NAME", "Qwen/Qwen2.5-Coder-0.5B-Instruct") def get_bridge_url() -> str: return os.environ.get("OPENENV_BASE_URL", _DEFAULT_OPENENV_BASE).rstrip("/") def get_request_timeout() -> float: return float(os.environ.get("OPENENV_REQUEST_TIMEOUT_SEC", "120")) def _fetch_task_ids(client: httpx.Client) -> List[str]: raw = os.environ.get("OPENENV_TASK_IDS", "").strip() if raw: return [x.strip() for x in raw.split(",") if x.strip()] r = client.get("/tasks", timeout=get_request_timeout()) r.raise_for_status() body = r.json() tasks = body.get("tasks") or [] ids = [t["task_id"] for t in tasks if t.get("task_id")] if not ids: raise RuntimeError("/tasks returned no task_id entries") return ids def make_real_dataset() -> Dataset: """Plain prompts + live /tasks (same spirit as colab_real_world.py, HF Space instead of loca.lt).""" bridge = get_bridge_url() timeout = get_request_timeout() rows_per_task = max(1, int(os.environ.get("ROWS_PER_TASK", "48"))) marker = os.environ.get("COMPLETION_SQL_MARKER", "Fixed SQL:") print(f"Connecting to OpenEnv at {bridge} (timeout={timeout}s)...") rows: List[Dict[str, Any]] = [] with httpx.Client(base_url=bridge, headers=BYPASS_HEADERS, timeout=timeout) as client: h = client.get("/health", timeout=min(30.0, timeout)) h.raise_for_status() print(f"OpenEnv health: {h.json()}") task_ids = _fetch_task_ids(client) print(f"Training task_ids ({len(task_ids)}): {task_ids}") for t_id in task_ids: resp = client.post("/reset", json={"task_id": t_id}) resp.raise_for_status() obs = resp.json()["observation"] prompt = ( "Fix the following SQL query and provide only the fixed SQL.\n" f"Task: {obs['task_description']}\n" f"Broken Query: {obs['original_query']}\n" f"{marker}" ) for _ in range(rows_per_task): rows.append({"prompt": prompt, "task_id": t_id}) if not rows: raise RuntimeError("Failed to build dataset (no rows).") print(f"Dataset: {len(rows)} prompts ({rows_per_task} per task).") return Dataset.from_list(rows) def make_task_dataset(task_id: str, rows_per_task: int) -> Dataset: bridge = get_bridge_url() timeout = get_request_timeout() marker = os.environ.get("COMPLETION_SQL_MARKER", "Fixed SQL:") with httpx.Client(base_url=bridge, headers=BYPASS_HEADERS, timeout=timeout) as client: resp = client.post("/reset", json={"task_id": task_id}) resp.raise_for_status() obs = resp.json()["observation"] prompt = ( "Fix the following SQL query and provide only the fixed SQL.\n" f"Task: {obs['task_description']}\n" f"Broken Query: {obs['original_query']}\n" f"{marker}" ) rows = [{"prompt": prompt, "task_id": task_id} for _ in range(max(1, rows_per_task))] return Dataset.from_list(rows) # --- 3. One live OpenEnv reward (colab_real_world style) --- def openenv_sql_reward_func(completions, task_id, **kwargs): """Score completions by executing extracted SQL against the real OpenEnv HTTP API.""" base = get_bridge_url() timeout = get_request_timeout() marker = os.environ.get("COMPLETION_SQL_MARKER", "Fixed SQL:") rewards: List[float] = [] with httpx.Client(base_url=base, headers=BYPASS_HEADERS, timeout=timeout) as client: for completion, t_id in zip(completions, task_id): if marker in completion: sql = completion.split(marker, 1)[-1].strip() else: sql = completion.strip() if not sql: rewards.append(0.0) continue hdr = {"X-Session-Id": str(uuid.uuid4())} try: client.post("/reset", json={"task_id": t_id}, headers=hdr).raise_for_status() resp = client.post( "/step", json={"action": {"action_type": "submit_query", "query": sql}}, headers=hdr, ) resp.raise_for_status() r = float(resp.json().get("reward", 0.0)) except Exception: r = 0.0 r += random.uniform(-1e-6, 1e-6) rewards.append(r) return rewards def eval_model_reward( model: AutoModelForCausalLM, tokenizer: AutoTokenizer, dataset: Dataset, *, max_items: int, ) -> float: subset = dataset.select(range(min(max_items, len(dataset)))) prompts = subset["prompt"] task_ids = subset["task_id"] completions: List[str] = [] for prompt in prompts: inputs = tokenizer(prompt, return_tensors="pt").to(model.device) with torch.no_grad(): out = model.generate( **inputs, max_new_tokens=256, do_sample=True, temperature=float(os.environ.get("EVAL_TEMPERATURE", "0.7")), top_p=float(os.environ.get("EVAL_TOP_P", "0.9")), renormalize_logits=True, remove_invalid_values=True, pad_token_id=tokenizer.eos_token_id, ) completions.append(tokenizer.decode(out[0], skip_special_tokens=True)) rewards = openenv_sql_reward_func(completions, task_ids) return float(sum(rewards) / max(len(rewards), 1)) # --- 3b. ARTIFACTS / PLOTS (REAL, FROM LOGS) --- @dataclass(frozen=True) class ArtifactPaths: root: Path @property def logs_jsonl(self) -> Path: return self.root / "train_log_history.jsonl" @property def metrics_json(self) -> Path: return self.root / "train_metrics.json" @property def reward_curve_png(self) -> Path: return self.root / "reward_curve.png" def _ensure_dir(path: Path) -> None: path.mkdir(parents=True, exist_ok=True) def save_log_history(log_history: List[Dict[str, Any]], paths: ArtifactPaths) -> None: _ensure_dir(paths.root) with paths.logs_jsonl.open("w", encoding="utf-8") as f: for row in log_history: f.write(json.dumps(row, ensure_ascii=False) + "\n") def extract_reward_series(log_history: List[Dict[str, Any]]) -> List[tuple[float, float]]: """ Returns [(step, reward_like_value)] extracted from trainer log_history. TRL log keys vary; this is resilient and will pick the most relevant. """ candidates = [ "reward", "rewards/mean", "rewards", "train/reward", "train/rewards", "objective/mean_reward", "mean_reward", ] series: List[tuple[float, float]] = [] for row in log_history: step = row.get("step") or row.get("global_step") or row.get("epoch") if step is None: continue value = None for key in candidates: if key in row and isinstance(row[key], (int, float)): value = float(row[key]) break if value is None: # fallback: pick any numeric key containing "reward" for k, v in row.items(): if "reward" in str(k).lower() and isinstance(v, (int, float)): value = float(v) break if value is None: continue series.append((float(step), value)) # de-dup by step while preserving order seen = set() deduped: List[tuple[float, float]] = [] for s, v in series: if s in seen: continue seen.add(s) deduped.append((s, v)) return deduped def write_metrics(log_history: List[Dict[str, Any]], reward_series: List[tuple[float, float]], paths: ArtifactPaths) -> None: metrics = { "generated_at_epoch_s": time.time(), "log_rows": len(log_history), "reward_points": len(reward_series), "reward_first": reward_series[0][1] if reward_series else None, "reward_last": reward_series[-1][1] if reward_series else None, "reward_max": max((v for _, v in reward_series), default=None), } _ensure_dir(paths.root) paths.metrics_json.write_text(json.dumps(metrics, indent=2), encoding="utf-8") def plot_reward_curve(reward_series: List[tuple[float, float]], paths: ArtifactPaths) -> None: if not reward_series: print("⚠️ No reward series found in log history; skipping plot.") return import matplotlib.pyplot as plt xs = [s for s, _ in reward_series] ys = [v for _, v in reward_series] plt.figure(figsize=(9, 4)) plt.plot(xs, ys, linewidth=2) plt.title("GRPO Reward Over Time (from run logs)") plt.xlabel("step") plt.ylabel("reward (extracted)") plt.grid(True, linestyle="--", alpha=0.4) _ensure_dir(paths.root) plt.tight_layout() plt.savefig(paths.reward_curve_png, dpi=200) print(f"Saved {paths.reward_curve_png}") def _resolve_report_to() -> str: raw = os.environ.get("TRL_REPORT_TO", "").strip().lower() if raw in ("", "auto"): return "wandb" if os.environ.get("WANDB_API_KEY") else "none" if raw in ("false", "no", "off", "none"): return "none" return raw # --- 4. Simple GRPO training loop (live OpenEnv rewards) --- def run_sota_train(): max_steps = int(os.environ.get("TRAIN_MAX_STEPS", "200")) out_dir = os.environ.get("OUTPUT_DIR", "./sota_results") print(f"Starting GRPO on {MODEL_NAME}...") print( f"OpenEnv={get_bridge_url()} | max_steps={max_steps} | " f"rows_per_task={os.environ.get('ROWS_PER_TASK', '48')} | " f"report_to={_resolve_report_to()}" ) tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) tokenizer.pad_token = tokenizer.eos_token use_cuda = torch.cuda.is_available() # L4/A10/A100 are typically more numerically stable with bf16 than fp16 for RL-style sampling. torch_dtype = torch.bfloat16 if use_cuda else torch.float32 model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, torch_dtype=torch_dtype, device_map="auto", attn_implementation=os.environ.get("ATTN_IMPLEMENTATION", "eager"), ) # Runtime generation safety defaults (used by both eval and GRPO generate path). model.generation_config.remove_invalid_values = True model.generation_config.renormalize_logits = True model.generation_config.top_p = float(os.environ.get("GRPO_TOP_P", "0.9")) model.generation_config.temperature = float(os.environ.get("GRPO_TEMPERATURE", "0.7")) train_dataset = make_real_dataset() print("Quick baseline eval (pre-train)...") baseline_avg_reward = eval_model_reward(model, tokenizer, train_dataset, max_items=8) hard_eval_n = int(os.environ.get("HARD_EVAL_SAMPLES", "16")) hard_dataset = make_task_dataset("hard_finance_explosion", rows_per_task=hard_eval_n) base_hard_reward = eval_model_reward(model, tokenizer, hard_dataset, max_items=hard_eval_n) report_to = _resolve_report_to() tb_dir = Path(out_dir) / "tensorboard" if report_to == "tensorboard": _ensure_dir(tb_dir) per_device_bs = int(os.environ.get("PER_DEVICE_TRAIN_BS", "1")) grad_accum = int(os.environ.get("GRAD_ACCUM", "2")) requested_num_gen = int(os.environ.get("GRPO_NUM_GENERATIONS", "8")) effective_bs = max(1, per_device_bs * grad_accum) if effective_bs % requested_num_gen != 0: valid = [d for d in range(2, effective_bs + 1) if effective_bs % d == 0] num_gen = valid[-1] if valid else 2 print( f"Adjusting GRPO_NUM_GENERATIONS from {requested_num_gen} to {num_gen} " f"for effective batch size {effective_bs}." ) else: num_gen = requested_num_gen _cfg: Dict[str, Any] = dict( output_dir=out_dir, learning_rate=float(os.environ.get("TRAIN_LR", "5e-6")), per_device_train_batch_size=per_device_bs, gradient_accumulation_steps=grad_accum, num_generations=num_gen, max_completion_length=int(os.environ.get("GRPO_MAX_COMPLETION_LEN", "256")), temperature=float(os.environ.get("GRPO_TEMPERATURE", "0.7")), top_p=float(os.environ.get("GRPO_TOP_P", "0.9")), bf16=bool(use_cuda), fp16=False, num_train_epochs=int(os.environ.get("TRAIN_NUM_EPOCHS", "1")), max_steps=max_steps, logging_steps=int(os.environ.get("LOGGING_STEPS", "1")), logging_first_step=True, report_to=report_to, ) if report_to == "tensorboard": _cfg["logging_dir"] = str(tb_dir) training_args = GRPOConfig(**_cfg) trainer = GRPOTrainer( model=model, reward_funcs=[openenv_sql_reward_func], args=training_args, train_dataset=train_dataset, processing_class=tokenizer, ) print("Training with live execution rewards against OpenEnv...") trainer.train() print("Quick eval (post-train)...") post_avg_reward = eval_model_reward(model, tokenizer, train_dataset, max_items=8) trained_hard_reward = eval_model_reward(model, tokenizer, hard_dataset, max_items=hard_eval_n) # --- Save artifacts (real logs/plots) --- artifacts = ArtifactPaths(root=Path(out_dir) / "artifacts") log_history = getattr(trainer.state, "log_history", []) or [] save_log_history(log_history, artifacts) reward_series = extract_reward_series(log_history) write_metrics(log_history, reward_series, artifacts) # augment metrics with before/after metrics_path = artifacts.metrics_json try: metrics = json.loads(metrics_path.read_text(encoding="utf-8")) except Exception: metrics = {} metrics.update( { "openenv_base_url": get_bridge_url(), "train_max_steps": max_steps, "model_name": MODEL_NAME, "baseline_avg_reward": baseline_avg_reward, "post_avg_reward": post_avg_reward, "delta_avg_reward": post_avg_reward - baseline_avg_reward, "base_hard_reward": base_hard_reward, "trained_hard_reward": trained_hard_reward, "delta_hard_reward": trained_hard_reward - base_hard_reward, "tensorboard_dir": str(tb_dir) if report_to == "tensorboard" else None, "report_to": report_to, } ) metrics_path.write_text(json.dumps(metrics, indent=2), encoding="utf-8") plot_reward_curve(reward_series, artifacts) try: import matplotlib.pyplot as plt labels = ["baseline", "post-train"] values = [baseline_avg_reward, post_avg_reward] plt.figure(figsize=(5, 4)) plt.bar(labels, values, color=["#94a3b8", "#22c55e"]) plt.ylim(0, max(1.0, max(values) * 1.1)) plt.title("Avg execution reward (sampled)") plt.ylabel("avg reward") out_path = artifacts.root / "before_after_avg_reward.png" plt.tight_layout() plt.savefig(out_path, dpi=200) print(f"Saved {out_path}") except Exception as e: print(f"Could not generate before/after plot: {e}") model_dir = os.environ.get("MODEL_SAVE_DIR", "./sota_sql_agent_full") print("\nSaving trained model locally...") model.save_pretrained(model_dir) hub_id = os.environ.get("MODEL_HUB_REPO_ID", os.environ.get("HF_HUB_REPO_ID", "md896/sql-debug-agent-qwen05b-grpo")) token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") if os.environ.get("SKIP_HUB_PUSH", "").strip() in ("1", "true", "yes"): print("SKIP_HUB_PUSH set — not pushing to Hub.") else: try: model.push_to_hub(hub_id, token=token) tokenizer.push_to_hub(hub_id, token=token) print(f"Pushed trained model to https://huggingface.co/{hub_id}") except Exception as e: print(f"Hub push failed (set HF_TOKEN / MODEL_HUB_REPO_ID or SKIP_HUB_PUSH=1): {e}") # Upload run artifacts back to the Space repo so you can download/view them. artifact_space = os.environ.get("ARTIFACT_SPACE_ID", "md896/sql-debug-env") run_tag = time.strftime("%Y%m%d-%H%M%S") try: if token: api = HfApi(token=token) api.upload_folder( repo_id=artifact_space, repo_type="space", folder_path=str(artifacts.root), path_in_repo=f"artifacts/runs/{run_tag}", commit_message=f"Add training artifacts {run_tag}", ) print(f"Uploaded artifacts to https://huggingface.co/spaces/{artifact_space}/tree/main/artifacts/runs/{run_tag}") else: print("No HF token in job env; skipping artifact upload.") except Exception as e: print(f"Artifact upload failed: {e}") print(f"\nTraining artifacts under {artifacts.root}") if __name__ == "__main__": run_sota_train()