Fix HF Job bootstrap: transformers>=4.51 for trl 0.18, datasets<4; simplify to colab-style OpenEnv SQL reward.

ultimate_sota_training.py

@@ -10,13 +10,10 @@ Environment (control cost vs quality on HF Jobs / local GPU):
   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)
-  REASONING_XML_TAG         — XML tag name for chain-of-thought (default: think)
-  TRAIN_MAX_STEPS           — GRPO optimizer steps (default: 200; was 30 for smoke)
-  TRAIN_NUM_EPOCHS, TRAIN_LR, GRPO_NUM_GENERATIONS, GRPO_MAX_COMPLETION_LEN
-  PER_DEVICE_TRAIN_BS, GRAD_ACCUM
-  TRL_REPORT_TO             — none | wandb | tensorboard (auto: wandb if key else tensorboard)
-  BOOTSTRAP_*_VERSION       — pin transformers / accelerate / trl for HF Jobs (see bootstrap_deps)
-  Artifacts: artifacts/reward_components.jsonl, artifacts/trainer_on_log.jsonl, tensorboard/
+  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
@@ -33,12 +30,9 @@ Key stability choices:
 
 from __future__ import annotations
 
-import contextvars
 import json
-import math
 import os
 import random
-import re
 import subprocess
 import sys
 import time
@@ -47,10 +41,6 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Dict, List, Optional
 
-# Set by TrainerCallback so reward funcs can tag JSONL rows with the real global_step.
-CURRENT_GRPO_STEP: contextvars.ContextVar[int] = contextvars.ContextVar("CURRENT_GRPO_STEP", default=-1)
-
-
 def _run(cmd: List[str], *, check: bool = True) -> subprocess.CompletedProcess:
     return subprocess.run(cmd, check=check)
 
@@ -82,34 +72,29 @@ def bootstrap_deps() -> None:
     # 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.4.1,<4.4.0",
+            "datasets>=3.2.0,<4.0.0",
             "matplotlib",
             "tensorboard",
-            "wandb",
-        ]
-    )
-
-    _tf = os.environ.get("BOOTSTRAP_TRANSFORMERS_VERSION", "4.48.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",
             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",
@@ -180,7 +165,6 @@ import transformers.utils.hub
 if not hasattr(transformers.utils.hub, "TRANSFORMERS_CACHE"):
     transformers.utils.hub.TRANSFORMERS_CACHE = "/tmp"
 
-from transformers import TrainerCallback
 from trl import GRPOConfig, GRPOTrainer
 from unsloth import FastLanguageModel
 
@@ -199,24 +183,6 @@ def get_request_timeout() -> float:
     return float(os.environ.get("OPENENV_REQUEST_TIMEOUT_SEC", "120"))
 
 
-def build_system_prompt() -> str:
-    """Single prompt template for every task (easy → expert); tag name is configurable."""
-    tag = os.environ.get("REASONING_XML_TAG", "think")
-    return f"""You are an elite SQL engineer. You fix broken SQLite analytics queries using the task description and the broken query.
-You MUST output your reasoning process inside <{tag}> tags.
-After you have finished thinking, you MUST output the exact fixed SQL query inside <sql> tags.
-Do not output any markdown blocks like ```sql.
-
-Example:
-<{tag}>
-I will check joins, filters, and aggregation, then write a corrected SELECT or WITH query.
-</{tag}>
-<sql>
-WITH OrderTotals AS (SELECT order_id, SUM(amount) AS total FROM line_items GROUP BY order_id)
-SELECT o.id, ot.total FROM orders o JOIN OrderTotals ot ON o.id = ot.order_id;
-</sql>"""
-
-
 def _fetch_task_ids(client: httpx.Client) -> List[str]:
     raw = os.environ.get("OPENENV_TASK_IDS", "").strip()
     if raw:
@@ -232,10 +198,11 @@ def _fetch_task_ids(client: httpx.Client) -> List[str]:
 
 
 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")))
-    system = build_system_prompt()
+    marker = os.environ.get("COMPLETION_SQL_MARKER", "Fixed SQL:")
 
     print(f"Connecting to OpenEnv at {bridge} (timeout={timeout}s)...")
     rows: List[Dict[str, Any]] = []
@@ -254,10 +221,10 @@ def make_real_dataset() -> Dataset:
             obs = resp.json()["observation"]
 
             prompt = (
-                f"{system}\n\n"
+                "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\n"
-                f"Provide your <{os.environ.get('REASONING_XML_TAG', 'think')}> and <sql> output:"
+                f"Broken Query: {obs['original_query']}\n"
+                f"{marker}"
             )
             for _ in range(rows_per_task):
                 rows.append({"prompt": prompt, "task_id": t_id})
@@ -267,147 +234,42 @@ def make_real_dataset() -> Dataset:
     print(f"Dataset: {len(rows)} prompts ({rows_per_task} per task).")
     return Dataset.from_list(rows)
 
-# --- 3. MULTI-REWARD SHAPING + JSONL logging (per-component batch stats) ---
-
-_REWARD_COMPONENTS_JSONL: Optional[Path] = None
-
-
-def extract_xml_tag(text, tag):
-    pattern = f"<{tag}>(.*?)</{tag}>"
-    match = re.search(pattern, text, re.DOTALL)
-    return match.group(1).strip() if match else None
 
+# --- 3. One live OpenEnv reward (colab_real_world style) ---
 
-def _reward_batch_stats(values: List[float]) -> Dict[str, float]:
-    if not values:
-        return {"mean": 0.0, "std": 0.0, "min": 0.0, "max": 0.0}
-    n = len(values)
-    mean = sum(values) / n
-    var = sum((x - mean) ** 2 for x in values) / max(n - 1, 1)
-    return {"mean": mean, "std": math.sqrt(var), "min": min(values), "max": max(values)}
 
-
-def _append_jsonl(path: Path, row: Dict[str, Any]) -> None:
-    path.parent.mkdir(parents=True, exist_ok=True)
-    with path.open("a", encoding="utf-8") as f:
-        f.write(json.dumps(row, ensure_ascii=False, default=str) + "\n")
-
-
-def _log_reward_component(name: str, values: List[float]) -> None:
-    if _REWARD_COMPONENTS_JSONL is None:
-        return
-    _append_jsonl(
-        _REWARD_COMPONENTS_JSONL,
-        {
-            "time_epoch_s": time.time(),
-            "global_step": CURRENT_GRPO_STEP.get(),
-            "reward_component": name,
-            "n": len(values),
-            **_reward_batch_stats(values),
-        },
-    )
-
-
-def format_reward_func(completions, **kwargs):
-    """Reward 1: CoT + sql XML tags (+0.1). Tag name follows REASONING_XML_TAG."""
-    tag = os.environ.get("REASONING_XML_TAG", "think")
-    rewards = []
-    for comp in completions:
-        has_think = extract_xml_tag(comp, tag) is not None
-        has_sql = extract_xml_tag(comp, "sql") is not None
-        rewards.append(0.1 if (has_think and has_sql) else 0.0)
-    _log_reward_component("format_xml", rewards)
-    return rewards
-
-
-def syntax_reward_func(completions, **kwargs):
-    """Reward 2: Does the SQL look like valid code? (+0.2)"""
-    rewards = []
-    for comp in completions:
-        sql = extract_xml_tag(comp, "sql")
-        if sql and (sql.upper().startswith("SELECT") or sql.upper().startswith("WITH")):
-            rewards.append(0.2)
-        else:
-            rewards.append(0.0)
-    _log_reward_component("syntax_select_with", rewards)
-    return rewards
-
-
-def execution_reward_func(completions, task_id, **kwargs):
-    """Reward 3: live OpenEnv submit_query against the real Space/API (not a stub)."""
-    rewards: List[float] = []
+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 query, t_id in zip(completions, task_id):
-            sql = extract_xml_tag(query, "sql")
+        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
-
-            session_headers = {"X-Session-Id": str(uuid.uuid4())}
+            hdr = {"X-Session-Id": str(uuid.uuid4())}
             try:
-                r0 = client.post("/reset", json={"task_id": t_id}, headers=session_headers)
-                r0.raise_for_status()
+                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=session_headers,
+                    headers=hdr,
                 )
                 resp.raise_for_status()
-                reward = float(resp.json().get("reward", 0.0))
+                r = float(resp.json().get("reward", 0.0))
             except Exception:
-                reward = 0.0
-
-            reward += random.uniform(-1e-6, 1e-6)
-            rewards.append(reward)
-    _log_reward_component("openenv_execution", rewards)
-    return rewards
-
-
-def length_shape_reward_func(completions, **kwargs):
-    """Reward 4: soft preference for shorter completions (bounded; does not replace execution reward)."""
-    cap = float(os.environ.get("COMPLETION_SOFT_CHAR_CAP", "3500"))
-    bonus_max = float(os.environ.get("LENGTH_BONUS_MAX", "0.05"))
-    rewards: List[float] = []
-    for comp in completions:
-        L = len(comp) if comp else 0
-        if L <= 0:
-            rewards.append(0.0)
-        else:
-            rewards.append(bonus_max * max(0.0, 1.0 - min(L, cap) / cap))
-    _log_reward_component("length_shape", rewards)
+                r = 0.0
+            r += random.uniform(-1e-6, 1e-6)
+            rewards.append(r)
     return rewards
 
 
-class GrpoStepContextCallback(TrainerCallback):
-    """Expose true global_step to reward funcs for JSONL alignment."""
-
-    def on_step_begin(self, args, state, control, **kwargs):
-        CURRENT_GRPO_STEP.set(int(state.global_step))
-
-
-class JsonlOnLogCallback(TrainerCallback):
-    """Mirror every trainer `logs` dict to JSONL (loss, learning_rate, reward keys, etc.)."""
-
-    def __init__(self, path: Path):
-        self.path = path
-        self.path.parent.mkdir(parents=True, exist_ok=True)
-        self._fp = path.open("w", encoding="utf-8")
-
-    def on_log(self, args, state, control, logs=None, **kwargs):
-        if not logs:
-            return
-        row: Dict[str, Any] = {"global_step": int(state.global_step), **dict(logs)}
-        self._fp.write(json.dumps(row, ensure_ascii=False, default=str) + "\n")
-        self._fp.flush()
-
-    def on_train_end(self, args, state, control, **kwargs):
-        try:
-            self._fp.close()
-        except Exception:
-            pass
-
 # --- 3b. ARTIFACTS / PLOTS (REAL, FROM LOGS) ---
 
 @dataclass(frozen=True)
@@ -520,9 +382,7 @@ def plot_reward_curve(reward_series: List[tuple[float, float]], paths: ArtifactP
 def _resolve_report_to() -> str:
     raw = os.environ.get("TRL_REPORT_TO", "").strip().lower()
     if raw in ("", "auto"):
-        if os.environ.get("WANDB_API_KEY"):
-            return "wandb"
-        return "tensorboard"
+        return "wandb" if os.environ.get("WANDB_API_KEY") else "none"
     if raw in ("false", "no", "off", "none"):
         return "none"
     return raw
@@ -530,14 +390,8 @@ def _resolve_report_to() -> str:
 
 # --- 4. Unsloth GRPO training loop (live OpenEnv rewards) ---
 def run_sota_train():
-    global _REWARD_COMPONENTS_JSONL
-
     max_steps = int(os.environ.get("TRAIN_MAX_STEPS", "200"))
     out_dir = os.environ.get("OUTPUT_DIR", "./sota_results")
-    artifacts_early = Path(out_dir) / "artifacts"
-    _ensure_dir(artifacts_early)
-    _REWARD_COMPONENTS_JSONL = artifacts_early / "reward_components.jsonl"
-    _REWARD_COMPONENTS_JSONL.write_text("", encoding="utf-8")
 
     print(f"Starting Unsloth GRPO on {MODEL_NAME}...")
     print(
@@ -566,9 +420,7 @@ def run_sota_train():
     train_dataset = make_real_dataset()
 
     def quick_exec_eval(max_items: int = 8) -> float:
-        """
-        Quick before/after check: sample prompts, generate CoT + sql, score via live OpenEnv.
-        """
+        """Sample prompts, generate completions, score with the same OpenEnv SQL reward."""
         subset = train_dataset.select(range(min(max_items, len(train_dataset))))
         prompts = subset["prompt"]
         task_ids = subset["task_id"]
@@ -586,7 +438,7 @@ def run_sota_train():
                 )
             completions.append(tokenizer.decode(out[0], skip_special_tokens=True))
 
-        rewards = execution_reward_func(completions, task_ids)
+        rewards = openenv_sql_reward_func(completions, task_ids)
         return float(sum(rewards) / max(len(rewards), 1))
 
     print("Quick baseline eval (pre-train)...")
@@ -603,7 +455,7 @@ def run_sota_train():
         per_device_train_batch_size=int(os.environ.get("PER_DEVICE_TRAIN_BS", "1")),
         gradient_accumulation_steps=int(os.environ.get("GRAD_ACCUM", "2")),
         num_generations=int(os.environ.get("GRPO_NUM_GENERATIONS", "8")),
-        max_completion_length=int(os.environ.get("GRPO_MAX_COMPLETION_LEN", "512")),
+        max_completion_length=int(os.environ.get("GRPO_MAX_COMPLETION_LEN", "256")),
         temperature=float(os.environ.get("GRPO_TEMPERATURE", "0.9")),
         num_train_epochs=int(os.environ.get("TRAIN_NUM_EPOCHS", "1")),
         max_steps=max_steps,
@@ -615,24 +467,12 @@ def run_sota_train():
         _cfg["logging_dir"] = str(tb_dir)
     training_args = GRPOConfig(**_cfg)
 
-    trainer_logs_path = artifacts_early / "trainer_on_log.jsonl"
-    trainer_logs_path.write_text("", encoding="utf-8")
-
     trainer = GRPOTrainer(
         model=model,
-        reward_funcs=[
-            format_reward_func,
-            syntax_reward_func,
-            execution_reward_func,
-            length_shape_reward_func,
-        ],
+        reward_funcs=[openenv_sql_reward_func],
         args=training_args,
         train_dataset=train_dataset,
         processing_class=tokenizer,
-        callbacks=[
-            GrpoStepContextCallback(),
-            JsonlOnLogCallback(trainer_logs_path),
-        ],
     )
 
     print("Training with live execution rewards against OpenEnv...")
@@ -661,8 +501,6 @@ def run_sota_train():
             "baseline_avg_reward": baseline_avg_reward,
             "post_avg_reward": post_avg_reward,
             "delta_avg_reward": post_avg_reward - baseline_avg_reward,
-            "reward_components_jsonl": str(artifacts_early / "reward_components.jsonl"),
-            "trainer_on_log_jsonl": str(artifacts_early / "trainer_on_log.jsonl"),
             "tensorboard_dir": str(tb_dir) if report_to == "tensorboard" else None,
             "report_to": report_to,
         }