feat: run fixed research pipelines within durable budgets

2026-09-08 22:16:05 +08:00
parent eb4850a003
commit 7860434b04
25 changed files with 2357 additions and 37 deletions
@@ -0,0 +1,121 @@
+"""Stage-three PostgreSQL acceptance in dedicated databases only."""
+
+import asyncio
+import os
+import subprocess
+from pathlib import Path
+
+from alembic import command
+from alembic.config import Config
+from cryptography.fernet import Fernet
+
+NAME = "wq_research_stage3_test"
+RESTORE = "wq_research_restore_stage3"
+os.environ.update(
+    DATABASE_URL=f"postgresql+asyncpg://postgres:research-test-only@127.0.0.1:18436/{NAME}",
+    ADMIN_PASSWORD="research-acceptance-only",
+    ENCRYPTION_KEY=Fernet.generate_key().decode(),
+)
+
+
+def docker(*args, **kwargs):
+    return subprocess.run(["docker", "exec", "-i", "wq-research-acceptance-pg", *args], check=True, **kwargs)
+
+
+async def acceptance():
+    from unittest.mock import patch
+
+    import httpx
+    from sqlalchemy import select
+
+    from app.config import Settings
+    from app.main import create_app
+    from app.models import TemplateInput
+    from app.research.workspace_contracts import TemplateSpec
+    from tests.test_ai import configure
+    from tests.test_backtests import setup
+    from tests.test_research_flows import begin, get, test_fixed_two_rounds_and_idempotent_start
+    from tests.test_research_workspace import template
+
+    app = create_app(Settings(_env_file=None, enable_runner=False, public_origin="http://testserver"))
+    async with app.router.lifespan_context(app):
+        async with httpx.AsyncClient(
+            transport=httpx.ASGITransport(app), base_url="http://testserver", headers={"X-WQ-Request": "1"}
+        ) as client:
+            assert (
+                await client.post(
+                    "/api/v1/auth/login", json={"username": "admin", "password": "research-acceptance-only"}
+                )
+            ).status_code == 200
+            # Configure a deterministic model; no provider or real platform network.
+            from tests.ai_fake import fake_model
+
+            app.state.ai.model_factory = fake_model
+            await configure(app, client)
+            platform, lane = await setup(app)
+            calls = []
+
+            async def model(ai, context, output_type, revision):
+                calls.append(context)
+                value = template()
+                value["expression"] = f"rank({{field}}) + {len(calls)}"
+                return TemplateSpec.model_validate(value), {
+                    "model": "fixture",
+                    "revision": revision,
+                    "usage": {"requests": 1},
+                }
+
+            async with app.state.sessions() as db:
+                fixed = await db.scalar(select(TemplateInput))
+            body = {
+                "request_id": "finite-run",
+                "name": "PG 有限研究",
+                "input_ids": [fixed.id],
+                "hypothesis": "排名稳定性",
+                "settings": {"region": "USA", "universe": "TOP3000", "delay": 1},
+                "budget": {"max_rounds": 2, "max_simulations": 4, "max_model_calls": 3},
+                "batch_candidates": 2,
+            }
+            with patch("app.research.runtime.request_model", model):
+                await test_fixed_two_rounds_and_idempotent_start(app, client, (body, platform, lane, calls))
+                short = {
+                    **body,
+                    "request_id": "concurrent-budget",
+                    "budget": {"max_rounds": 1, "max_simulations": 1, "max_model_calls": 2},
+                }
+                first, second = await asyncio.gather(begin(client, short), begin(client, short))
+                assert first["id"] == second["id"]
+                from app.research.runtime import ResearchRuntime
+
+                another = ResearchRuntime(app.state.sessions, app.state.ai, app.state.runner)
+                for _ in range(7):
+                    await asyncio.gather(
+                        app.state.research.advance(first["id"]), another.advance(first["id"])
+                    )
+                result = await get(client, first["id"])
+                assert result["status"] == "budget_exhausted" and result["simulations_used"] == 0
+                assert result["model_calls_used"] == 1
+    print("PASS PostgreSQL: two-round execution, idempotent starts, concurrent reservations and budget gate")
+
+
+if __name__ == "__main__":
+    docker("createdb", "-U", "postgres", NAME)
+    with Path("/tmp/wq-research-stage2.dump").open("rb") as source:
+        docker("pg_restore", "-U", "postgres", "-d", NAME, stdin=source)
+    config = Config("alembic.ini")
+    command.upgrade(config, "0008")
+    command.check(config)
+    asyncio.run(acceptance())
+    dump = Path("/tmp/wq-research-stage3.dump")
+    with dump.open("wb") as output:
+        docker("pg_dump", "-U", "postgres", "-Fc", NAME, stdout=output)
+    docker("createdb", "-U", "postgres", RESTORE)
+    with dump.open("rb") as source:
+        docker("pg_restore", "-U", "postgres", "-d", RESTORE, stdin=source)
+    query = "SELECT (SELECT count(*) FROM research_revisions),(SELECT count(*) FROM research_flow_runs),(SELECT count(*) FROM research_step_runs),(SELECT note FROM research WHERE alpha_id='OLD_RESEARCH')"
+    a = docker("psql", "-U", "postgres", "-d", NAME, "-Atc", query, capture_output=True).stdout
+    b = docker("psql", "-U", "postgres", "-d", RESTORE, "-Atc", query, capture_output=True).stdout
+    assert a == b
+    print(
+        "PASS PostgreSQL 17: 0007 → 0008 and pg_dump/pg_restore preserve flow budgets, steps and old research notes"
+    )