Add Docker-free multi-turn SQLite text-to-SQL workflow (BIRD)

examples/grpo_sqlite_sql/README.md

@@ -0,0 +1,70 @@
+# Example: Multi-turn GRPO for Text-to-SQL on SQLite (BIRD)
+
+This example trains an agent that explores a real SQLite database with read-only queries and then
+submits one SQL query, using the [BIRD](https://bird-bench.github.io/) benchmark. It complements the
+[InterCode-SQL example](../grpo_intercode_sql/README.md): it needs **no Docker or MySQL server**,
+because each rollout opens the `.sqlite` file directly.
+
+## How the workflow works
+
+`sqlite_sql_workflow` (`trinity/common/workflows/envs/sqlite_sql/`) gives the model the database
+schema, the question and BIRD's external knowledge. At each turn the model replies with exactly one
+of:
+
+- `<action>SQL</action>`: run a read-only query; the first `max_obs_rows` rows (or the error) come back
+  as the next user message
+- `<answer>SQL</answer>`: submit the final query and end the episode
+
+The reward is BIRD / Spider execution accuracy: 1.0 if the submitted query returns the same set of
+rows as the reference query, else 0.0. Episodes that run out of turns score 0 unless
+`auto_submit: true`, which grades the last query that ran successfully.
+
+Because the reward comes from executing model-written SQL, every query runs in a sandbox: the
+database is opened read-only behind an SQLite authorizer that only allows reads (no writes, `ATTACH`,
+`PRAGMA` or extensions), one statement is executed per query, and a query is interrupted after
+`timeout` seconds. Concurrent rollouts therefore cannot change the data other rollouts are scored
+on. Experience `info` also records `gold_error` (the reference query fails) and `gold_empty` (the
+reference returns no rows, so any query returning nothing would match), so these tasks can be
+filtered.
+
+## Prepare data
+
+1. Download the BIRD train set and databases from <https://bird-bench.github.io/>, and BIRD
+   Mini-Dev (500 SQLite questions, disjoint from train) from
+   <https://bird-bench.oss-cn-beijing.aliyuncs.com/minidev.zip>.
+2. Create the taskset:
+
+```bash
+python examples/grpo_sqlite_sql/get_bird_data.py \
+    --train_json train/train.json \
+    --test_json minidev/MINIDEV/mini_dev_sqlite.json
+```
+
+Each line of `bird_data/train.jsonl` and `bird_data/test.jsonl` has `db_id`, `question`, `evidence`
+and `SQL`. Any dataset in this format with `<db_root>/<db_id>/<db_id>.sqlite` databases (e.g. Spider)
+works the same way.
+
+## Run
+
+```bash
+export BIRD_TRAIN_DB_ROOT=/path/to/train/train_databases
+export BIRD_DEV_DB_ROOT=/path/to/minidev/MINIDEV/dev_databases
+trinity run --config examples/grpo_sqlite_sql/sqlite_sql.yaml
+```
+
+Workflow arguments (`workflow_args`):
+
+| Argument | Default | Description |
+|---|---|---|
+| `db_root` | required | Directory with `<db_id>/<db_id>.sqlite` |
+| `max_env_steps` | 5 | Maximum number of model turns |
+| `max_obs_rows` | 10 | Rows of a query result shown to the model |
+| `timeout` | 30 | Per-query timeout in seconds (the official BIRD evaluator uses 30) |
+| `auto_submit` | false | Grade the last successful query if the model never submits |
+
+Notes on BIRD Mini-Dev as a reward source:
+
+- 16 of the 500 reference queries return the single value `1`, so a constant `SELECT 1` earns the
+  full reward on 3.2% of the tasks. Watch rollouts for constant answers on hard tasks.
+- Two reference queries (question ids 518 and 701) take longer than 30 s on a typical machine,
+  so they score 0 for every model at the default timeout, as they do in the official evaluator.

examples/grpo_sqlite_sql/get_bird_data.py

@@ -0,0 +1,53 @@
+"""
+Create the Trinity taskset for `sqlite_sql_workflow` from BIRD.
+
+Download the BIRD train set (https://bird-bench.github.io/) and BIRD Mini-Dev
+(https://bird-bench.oss-cn-beijing.aliyuncs.com/minidev.zip), then run:
+
+    python examples/grpo_sqlite_sql/get_bird_data.py \
+        --train_json train/train.json \
+        --test_json minidev/MINIDEV/mini_dev_sqlite.json
+
+This writes `examples/grpo_sqlite_sql/bird_data/train.jsonl` and `test.jsonl` with the fields
+the workflow reads: `db_id`, `question`, `evidence` and `SQL`.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+FIELDS = ("db_id", "question", "evidence", "SQL")
+
+
+def load(path: Path) -> list[dict]:
+    text = path.read_text(encoding="utf-8")
+    if path.suffix == ".jsonl":
+        return [json.loads(line) for line in text.splitlines() if line.strip()]
+    return json.loads(text)
+
+
+def write(records: list[dict], path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as f:
+        for record in records:
+            row = {key: record.get(key) or "" for key in FIELDS}
+            f.write(json.dumps(row, ensure_ascii=False) + "\n")
+    print(f"Wrote {len(records)} tasks to {path}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--train_json", type=Path, required=True)
+    parser.add_argument("--test_json", type=Path, required=True)
+    parser.add_argument(
+        "--output_dir", type=Path, default=Path("examples/grpo_sqlite_sql/bird_data")
+    )
+    args = parser.parse_args()
+    write(load(args.train_json), args.output_dir / "train.jsonl")
+    write(load(args.test_json), args.output_dir / "test.jsonl")
+
+
+if __name__ == "__main__":
+    main()

examples/grpo_sqlite_sql/sqlite_sql.yaml

@@ -0,0 +1,81 @@
+project: "SQLITE_SQL"
+name: "BIRD_SQLITE_SQL_RFT"
+checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints}
+algorithm:
+  algorithm_type: grpo
+  repeat_times: 8
+  optimizer:
+    lr: 1e-6
+model:
+  model_path: ${oc.env:TRINITY_MODEL_PATH,Qwen/Qwen2.5-7B-Instruct}
+  max_prompt_tokens: 12288
+  max_response_tokens: 1024
+  max_model_len: 13312
+cluster:
+  node_num: 1
+  gpu_per_node: 8
+buffer:
+  total_epochs: 1
+  batch_size: 16
+  explorer_input:
+    taskset:
+      name: bird_sqlite_sql
+      storage_type: file
+      path: 'examples/grpo_sqlite_sql/bird_data'
+      split: train
+      format:
+        prompt_key: 'question'
+      rollout_args:
+        temperature: 1.0
+        logprobs: 0
+      workflow_args:
+        db_root: ${oc.env:BIRD_TRAIN_DB_ROOT,train/train_databases}
+        max_env_steps: 5
+        max_obs_rows: 10
+        timeout: 30
+    eval_tasksets:
+      - name: bird_minidev_eval
+        storage_type: file
+        path: 'examples/grpo_sqlite_sql/bird_data'
+        split: test
+        format:
+          prompt_key: 'question'
+        rollout_args:
+          temperature: 0.0
+          logprobs: 0
+        workflow_args:
+          db_root: ${oc.env:BIRD_DEV_DB_ROOT,minidev/MINIDEV/dev_databases}
+          max_env_steps: 5
+          max_obs_rows: 10
+          timeout: 30
+        repeat_times: 1
+    default_workflow_type: 'sqlite_sql_workflow'
+  trainer_input:
+    experience_buffer:
+      name: sqlite_sql_buffer
+      storage_type: queue
+      path: 'sqlite:///sqlite_sql.db'
+explorer:
+  eval_on_startup: true
+  eval_interval: 20
+  runner_per_model: 16
+  max_timeout: 3600
+  rollout_model:
+    engine_num: 2
+    tensor_parallel_size: 2
+    enable_prefix_caching: true
+    enforce_eager: true
+    dtype: bfloat16
+    seed: 42
+    gpu_memory_utilization: 0.7
+    enable_chunked_prefill: true
+synchronizer:
+  sync_method: 'nccl'
+  sync_interval: 8
+  sync_timeout: 1200
+trainer:
+  save_interval: 100
+  grad_clip: 1.0
+  use_dynamic_bsz: true
+  max_token_len_per_gpu: 13312
+  ulysses_sequence_parallel_size: 1

tests/common/sqlite_sql_env_test.py

@@ -0,0 +1,109 @@
+import sqlite3
+
+import pytest
+
+from trinity.common.workflows.envs.sqlite_sql.sqlite_sql_env import (
+    SQLiteSQLEnv,
+    db_file,
+    execute_sql,
+    parse_response,
+    schema_text,
+)
+
+GOLD = "SELECT name FROM customers WHERE city = 'Hangzhou'"
+
+
+@pytest.fixture()
+def db_path(tmp_path):
+    path = db_file(str(tmp_path), "shop")
+    path.parent.mkdir()
+    conn = sqlite3.connect(path)
+    conn.executescript(
+        """
+        CREATE TABLE customers(id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT, city TEXT);
+        INSERT INTO customers(name, city) VALUES ('Ann','Hangzhou'),('Bob','Shanghai'),('Cat','Hangzhou');
+        """
+    )
+    conn.commit()
+    conn.close()
+    return path
+
+
+def test_parse_response():
+    assert parse_response("<think>x</think><action>SELECT 1;</action>") == ("action", "SELECT 1")
+    assert parse_response("<answer>\n```sql\nSELECT 2\n```\n</answer>") == ("answer", "SELECT 2")
+    assert parse_response("<action>SELECT 1</action> then <answer>SELECT 3</answer>") == (
+        "answer",
+        "SELECT 3",
+    )
+    # tags inside the thinking section are ignored
+    assert parse_response("<think><answer>SELECT 9</answer></think>") == (None, "")
+    assert parse_response("SELECT 1") == (None, "")
+
+
+def test_schema_text_skips_internal_tables(db_path):
+    schema = schema_text(db_path)
+    assert "CREATE TABLE customers" in schema
+    assert "sqlite_sequence" not in schema
+
+
+def test_explore_then_answer(db_path):
+    env = SQLiteSQLEnv(db_path, GOLD, max_obs_rows=2)
+
+    obs = env.step("<action>SELECT name, city FROM customers ORDER BY id</action>")
+    assert not obs.done and obs.error is None
+    assert obs.observation.splitlines()[0] == "name | city"
+    assert "Ann | Hangzhou" in obs.observation and "Cat" not in obs.observation
+    assert "(3 rows in total)" in obs.observation
+
+    result = env.step("<answer>SELECT name FROM customers WHERE city LIKE 'Hang%'</answer>")
+    assert result.done and result.reward == 1.0
+    assert env.submitted and env.num_queries == 1
+
+
+def test_wrong_answer(db_path):
+    result = SQLiteSQLEnv(db_path, GOLD).step("<answer>SELECT name FROM customers</answer>")
+    assert result.done and result.reward == 0.0
+
+
+@pytest.mark.parametrize(
+    "sql",
+    [
+        "DELETE FROM customers",
+        "UPDATE customers SET city = 'Hangzhou'",
+        "DROP TABLE customers",
+        "ATTACH DATABASE ':memory:' AS other",
+        "PRAGMA table_info(customers)",
+        "SELECT 1; DELETE FROM customers",
+    ],
+)
+def test_writes_are_rejected_and_do_not_reach_the_database(db_path, sql):
+    env = SQLiteSQLEnv(db_path, GOLD)
+    obs = env.step(f"<action>{sql}</action>")
+    assert obs.observation.startswith("Error:") and not obs.done
+    assert env.num_errors == 1
+    rows, _, _ = execute_sql(db_path, "SELECT COUNT(*) FROM customers")
+    assert rows == [(3,)]
+
+
+def test_runaway_query_is_interrupted(db_path):
+    env = SQLiteSQLEnv(db_path, GOLD, timeout=0.5)
+    obs = env.step(
+        "<action>WITH RECURSIVE r(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM r) "
+        "SELECT COUNT(*) FROM r</action>"
+    )
+    assert obs.observation == "Error: timeout"
+
+
+def test_invalid_format_keeps_episode_open(db_path):
+    obs = SQLiteSQLEnv(db_path, GOLD).step("I think the answer is Ann.")
+    assert not obs.done and obs.error == "no action"
+
+
+def test_broken_or_empty_reference(db_path):
+    broken = SQLiteSQLEnv(db_path, "SELECT nope FROM customers")
+    assert broken.gold_error is not None
+    assert broken.step(f"<answer>{GOLD}</answer>").reward == 0.0
+
+    empty = SQLiteSQLEnv(db_path, "SELECT name FROM customers WHERE city = 'Beijing'")
+    assert empty.gold_empty

trinity/common/workflows/__init__.py

@@ -36,6 +36,7 @@
         "sciworld_workflow": "trinity.common.workflows.envs.sciworld.sciworld_workflow.SciWorldWorkflow",
         "webshop_workflow": "trinity.common.workflows.envs.webshop.webshop_workflow.WebShopWorkflow",
         "intercode_sql_workflow": "trinity.common.workflows.envs.intercode_sql.intercode_sql_workflow.InterCodeSQLWorkflow",
+        "sqlite_sql_workflow": "trinity.common.workflows.envs.sqlite_sql.sqlite_sql_workflow.SQLiteSQLWorkflow",
         # general multi-turn
         "step_wise_alfworld_workflow": "trinity.common.workflows.envs.alfworld.alfworld_workflow.StepWiseAlfworldWorkflow",
         # ruler/llm_as_a_judge