RonyForAI/Mirage_DB_RL

1"""2inference.py — Baseline inference script for the Mirage_RL OpenEnv environment.3Place this file in the root directory of the project (Mirage_RL/).4 5MANDATORY ENVIRONMENT VARIABLES:6    HF_TOKEN         Your Hugging Face / API key.7    API_BASE_URL     The API endpoint for the LLM.8                     Default: https://router.huggingface.co/v19    MODEL_NAME       The model identifier to use for inference.10                     Default: Qwen/Qwen2.5-72B-Instruct11 12STDOUT FORMAT (exact — one line per tag, no newlines within a line):13    [START] task=<task_name> env=<benchmark> model=<model_name>14    [STEP]  step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null>15    [END]   success=<true|false> steps=<n> score=<score> rewards=<r1,r2,...,rn>16 17Usage:18    HF_TOKEN=<key> API_BASE_URL=<url> MODEL_NAME=<model> python inference.py19"""20 21from __future__ import annotations22 23import os24import sys25import time26import textwrap27from typing import List, Optional28 29# ── Path setup: works when run locally or inside Docker ───────────────────────30_HERE = os.path.dirname(os.path.abspath(__file__))31if _HERE not in sys.path:32    sys.path.insert(0, _HERE)33 34# ── OpenAI client ─────────────────────────────────────────────────────────────35try:36    from openai import OpenAI37except ImportError:38    sys.exit("ERROR: 'openai' package not found. Run: pip install openai>=1.0.0")39 40# ── Environment + task imports ────────────────────────────────────────────────41try:42    from Mirage_RL.server.Mirage_RL_environment import QueryEnv43    from Mirage_RL.server.tasks import TASKS, grade44    from Mirage_RL.models import QueryAction45except ImportError:46    try:47        from server.Mirage_RL_environment import QueryEnv          # type: ignore48        from server.tasks import TASKS, grade                      # type: ignore49        from models import QueryAction                             # type: ignore50    except ImportError as exc:51        sys.exit(52            f"ERROR: Cannot import Mirage_RL modules.\n"53            f"Run inference.py from the Mirage_RL/ directory or install the package.\n"54            f"Details: {exc}"55        )56 57# ─────────────────────────────────────────────────────────────────────────────58# Configuration59# ─────────────────────────────────────────────────────────────────────────────60 61API_KEY      = os.getenv("HF_TOKEN") or os.getenv("API_KEY") or os.getenv("OPENAI_API_KEY", "")62API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"63MODEL_NAME   = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct"64BENCHMARK    = "mirage_rl"65 66TASK_ORDER              = ["easy", "medium", "hard"]67MAX_STEPS               = 10      # hard upper bound per episode (tasks finish naturally)68MAX_RETRIES             = 369RETRY_DELAY             = 1.070SUCCESS_SCORE_THRESHOLD = 0.5     # score >= 0.5 counts as success71 72 73# ─────────────────────────────────────────────────────────────────────────────74# Mandatory log helpers  (exact format — do not modify field names or order)75# ─────────────────────────────────────────────────────────────────────────────76 77def log_start(task: str, env: str, model: str) -> None:78    print(f"[START] task={task} env={env} model={model}", flush=True)79 80 81def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:82    error_val = error if error else "null"83    done_val  = str(done).lower()84    print(85        f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}",86        flush=True,87    )88 89 90def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:91    rewards_str = ",".join(f"{r:.2f}" for r in rewards)92    print(93        f"[END] success={str(success).lower()} steps={steps} score={score:.2f} rewards={rewards_str}",94        flush=True,95    )96 97 98# ─────────────────────────────────────────────────────────────────────────────99# Action string formatter  (used in [STEP] log)100# ─────────────────────────────────────────────────────────────────────────────101 102_JOIN_NAMES  = {0: "hash", 1: "nested_loop", 2: "merge_sort"}103_INDEX_NAMES = {0: "no_index", 1: "use_index"}104 105def format_action(action: QueryAction, tables: list) -> str:106    table_name = tables[action.next_table] if action.next_table < len(tables) else str(action.next_table)107    join_name  = _JOIN_NAMES.get(action.join_type, str(action.join_type))108    idx_name   = _INDEX_NAMES.get(action.use_index, str(action.use_index))109    return f"join({table_name},{join_name},{idx_name})"110 111 112# ─────────────────────────────────────────────────────────────────────────────113# Prompt builder114# ─────────────────────────────────────────────────────────────────────────────115 116SYSTEM_PROMPT = textwrap.dedent("""117    You are a database query optimizer. Your job is to build an efficient join plan118    by deciding, one table at a time, which table to join next and how.119 120    You will see table statistics including estimated row counts, join selectivities,121    and index availability for each table. These estimates may not be exact — real122    database planners always operate on statistics that can differ from ground truth.123 124    TWO things determine your score:125 126    1. JOIN METHOD quality (60% of score): which algorithm and index you choose.127       Different join algorithms have very different cost profiles.128       Using an index when one is available significantly reduces scan cost.129 130    2. JOIN ORDER quality (40% of score): which table you pick at each step.131       Joining a table with a large output early causes the intermediate result132       to explode — every subsequent join must then probe against that larger set.133       To minimise total cost: prefer joining tables with the smallest134       (estimated_rows × selectivity) output first, saving large tables for later.135 136    The observation shows 'intermediate_size': the estimated size of the accumulated137    intermediate result from all joins so far. Keep this number small by joining138    highly selective, small-output tables early in the sequence.139 140    Respond with ONLY a valid JSON object — no markdown, no explanation:141    {"next_table": <int>, "join_type": <0|1|2>, "use_index": <0|1>}142""").strip()143 144 145def build_user_prompt(obs, task_config, step: int) -> str:146    lines = []147    for i in range(len(obs.tables)):148        status   = "JOINED" if i in obs.chosen_order else "available"149        idx_info = "indexed" if obs.has_index[i] else "no_index"150        est_out  = obs.table_rows[i] * obs.selectivities[i]   # estimated output rows151        lines.append(152            f"  [{i}] {obs.tables[i]:15s}  est_rows={obs.table_rows[i]:>12,}  "153            f"sel={obs.selectivities[i]:.3f}  est_output={est_out:>12,.0f}  "154            f"{idx_info:10s}  [{status}]"155        )156 157    remaining_display = [f"[{i}]{obs.tables[i]}" for i in obs.remaining_tables]158    joined_display    = [obs.tables[i] for i in obs.chosen_order]159 160    return textwrap.dedent(f"""161        Step {step} of episode  |  Task: {task_config.name}  [{task_config.difficulty.upper()}]162 163        Query:164        {obs.query_context}165 166        Table statistics (est_rows = planner estimates, may differ from true cardinality):167        {chr(10).join(lines)}168 169        Join progress:170          Already joined       : {joined_display if joined_display else "(none — first step)"}171          Remaining            : {remaining_display}172          Accumulated cost     : {obs.current_cost:.2f}173          Intermediate size    : {obs.intermediate_size:,.0f}  ← keep this small; explodes if you join large-output tables early174 175        Decide the next join. Choose next_table from {obs.remaining_tables}.176        Reason about join ORDER (intermediate blowup) AND join METHOD (algorithm + index).177        Output the JSON.178    """).strip()179 180 181 182# ─────────────────────────────────────────────────────────────────────────────183# LLM call + greedy fallback184# ─────────────────────────────────────────────────────────────────────────────185 186import random as _random187 188def random_fallback(obs) -> dict:189    """190    Random fallback when LLM fails — picks a uniformly random remaining table,191    random join type, random index usage. This ensures LLM failures are penalised192    rather than silently rescued by an optimal greedy choice.193    """194    table  = _random.choice(obs.remaining_tables)195    j_type = _random.randint(0, 2)196    use_ix = _random.randint(0, 1)197    return {"next_table": table, "join_type": j_type, "use_index": use_ix}198 199 200def get_action(client: OpenAI, obs, task_config, step: int) -> tuple[dict, Optional[str]]:201    """202    Ask LLM for next action. Returns (action_dict, error_or_None).203    Falls back to greedy after MAX_RETRIES failures.204    """205    user_prompt = build_user_prompt(obs, task_config, step)206    last_error: Optional[str] = None207 208    for attempt in range(MAX_RETRIES):209        try:210            completion = client.chat.completions.create(211                model=MODEL_NAME,212                messages=[213                    {"role": "system", "content": SYSTEM_PROMPT},214                    {"role": "user",   "content": user_prompt},215                ],216                temperature=0.0,217                max_tokens=128,218                stream=False,219            )220            text = (completion.choices[0].message.content or "").strip()221 222            # Strip markdown fences if present223            if text.startswith("```"):224                parts = text.split("```")225                text = parts[1].lstrip("json").strip() if len(parts) > 1 else text226 227            import json228            action = json.loads(text)229 230            # Validate231            if action.get("next_table") not in obs.remaining_tables:232                raise ValueError(f"next_table={action.get('next_table')} not in {obs.remaining_tables}")233            action["join_type"] = int(action.get("join_type", 2))234            action["use_index"] = int(action.get("use_index", 1))235            if action["join_type"] not in (0, 1, 2):236                action["join_type"] = 2237            if action["use_index"] not in (0, 1):238                action["use_index"] = 1239 240            return action, None241 242        except Exception as exc:243            last_error = str(exc)244            if attempt < MAX_RETRIES - 1:245                time.sleep(RETRY_DELAY)246 247    # All retries exhausted → random fallback (penalises LLM failures)248    return random_fallback(obs), f"llm_failed:{last_error}"249 250 251# ─────────────────────────────────────────────────────────────────────────────252# Single task runner253# ─────────────────────────────────────────────────────────────────────────────254 255def run_task(task_id: str, client: OpenAI) -> float:256    """Run one task episode. Returns final score in [0.0, 1.0]."""257    task_config = TASKS[task_id]258    env = QueryEnv()259 260    rewards:     List[float]  = []261    steps_taken: int          = 0262    score:       float        = 0.0263    success:     bool         = False264 265    log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME)266 267    try:268        obs = env.reset(task_id=task_id, seed=42)  # seed=42 → reproducible scenario + noise269 270        for step in range(1, MAX_STEPS + 1):271            if obs.done:272                break273 274            action_dict, error = get_action(client, obs, task_config, step)275 276            action = QueryAction(277                next_table=int(action_dict["next_table"]),278                join_type=int(action_dict["join_type"]),279                use_index=int(action_dict["use_index"]),280            )281 282            obs = env.step(action)283 284            reward       = obs.reward285            done         = obs.done286            steps_taken  = step287 288            rewards.append(reward)289 290            log_step(291                step=step,292                action=format_action(action, list(obs.tables)),293                reward=reward,294                done=done,295                error=error,296            )297 298            if done:299                break300 301        # Final grader score — the environment returns the episode score natively on done302        score   = rewards[-1] if rewards else 0.0303        score   = min(max(score, 0.0), 1.0)304        success = score >= SUCCESS_SCORE_THRESHOLD305 306    except Exception as exc:307        print(f"[DEBUG] Task {task_id} exception: {exc}", flush=True)308        score   = 0.0309        success = False310 311    finally:312        try:313            if hasattr(env, "close"):314                env.close()315        except Exception:316            pass317        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)318 319    return score320 321 322# ─────────────────────────────────────────────────────────────────────────────323# Entry point324# ─────────────────────────────────────────────────────────────────────────────325 326def main() -> None:327    if not API_KEY:328        print(329            "ERROR: No API key found.\n"330            "Set HF_TOKEN (or API_KEY / OPENAI_API_KEY) before running.",331            file=sys.stderr,332        )333        sys.exit(1)334 335    client = OpenAI(api_key=API_KEY, base_url=API_BASE_URL)336 337    results: list[dict] = []338    for task_id in TASK_ORDER:339        score = run_task(task_id, client)340        results.append({"task_id": task_id, "score": score})341 342    # Human-readable summary (goes to stdout after all [END] lines)343    avg = sum(r["score"] for r in results) / len(results)344    sep = "=" * 52345    print(f"\n{sep}", flush=True)346    print("  MIRAGE_RL BASELINE RESULTS", flush=True)347    print(sep, flush=True)348    for r in results:349        bar = "█" * int(r["score"] * 20)350        print(f"  {r['task_id']:8s} | {r['score']:.2f} | {bar}", flush=True)351    print(sep, flush=True)352    print(f"  {'AVERAGE':8s} | {avg:.2f}", flush=True)353    print(sep, flush=True)354 355 356if __name__ == "__main__":357    main()358