261 lines
11 KiB
Python
261 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Submit, execute, and materialize isolated burst benchmark cases through MCP."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
from uuid import uuid4
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--design", type=Path, default=Path("resources/benchmarks/tjwater2-v2-burst-balanced.json"))
|
|
parser.add_argument("--state-dir", type=Path, default=Path("data/benchmark-generation/tjwater2-20261105"))
|
|
parser.add_argument("--agent-env", type=Path, default=Path("../next-tjwater-agent/.local.env"))
|
|
parser.add_argument("--limit", type=int, default=1, help="Maximum new cases to complete in this invocation")
|
|
parser.add_argument("--batches", type=int, default=1, help="Run isolated child processes for this many batches")
|
|
parser.add_argument("--local-executor-concurrency", type=int, default=1, choices=(1, 2))
|
|
args = parser.parse_args()
|
|
if args.limit < 1 or args.limit > args.local_executor_concurrency:
|
|
parser.error("--limit must be between 1 and local executor concurrency")
|
|
if args.batches < 1:
|
|
parser.error("--batches must be positive")
|
|
if args.batches > 1:
|
|
command = [
|
|
sys.executable,
|
|
str(Path(__file__).resolve()),
|
|
"--design", str(args.design),
|
|
"--state-dir", str(args.state_dir),
|
|
"--agent-env", str(args.agent_env),
|
|
"--limit", str(args.limit),
|
|
"--local-executor-concurrency", str(args.local_executor_concurrency),
|
|
"--batches", "1",
|
|
]
|
|
for _ in range(args.batches):
|
|
subprocess.run(command, check=True)
|
|
return 0
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
os.umask(0o077)
|
|
|
|
from dotenv import load_dotenv
|
|
from psycopg.conninfo import make_conninfo
|
|
from scripts.benchmark.calibration import save_json
|
|
from scripts.benchmark.generation import write_analysis_scada
|
|
from scripts.benchmark.simulation_design import load_simulation_design
|
|
from scripts.check_burst_simulation_design import resolve_project
|
|
|
|
load_dotenv(Path(__file__).resolve().parents[1] / ".env", override=False)
|
|
state_path = args.state_dir.resolve() / "state.json"
|
|
state = json.loads(state_path.read_text(encoding="utf-8"))
|
|
project_id = state.get("project_id")
|
|
if not project_id or state.get("status") != "prepared":
|
|
raise ValueError("run isolated preparation before MCP generation")
|
|
design = load_simulation_design(args.design)
|
|
pipes = {pipe.pipe_id: pipe for pipe in design.pipes}
|
|
source_biz, source_iot = asyncio.run(resolve_project(design))
|
|
business_dsn = make_conninfo(source_biz, dbname=state["target"])
|
|
timescale_dsn = make_conninfo(source_iot, dbname=state["target"])
|
|
|
|
pending = [
|
|
row for row in state["schedule"]
|
|
if state["cases"].get(row["case_key"], {}).get("status") != "scada_materialized"
|
|
][:args.limit]
|
|
for offset in range(0, len(pending), args.local_executor_concurrency):
|
|
batch = pending[offset:offset + args.local_executor_concurrency]
|
|
job_ids = submit_batch(
|
|
batch=batch,
|
|
pipes=pipes,
|
|
project_id=project_id,
|
|
agent_env=args.agent_env,
|
|
state=state,
|
|
state_path=state_path,
|
|
)
|
|
asyncio.run(execute_batch(project_id, job_ids, args.local_executor_concurrency))
|
|
for row in batch:
|
|
case = state["cases"][row["case_key"]]
|
|
run = find_analysis_run(business_dsn, case["scenario_name"])
|
|
summary = write_analysis_scada(
|
|
business_dsn=business_dsn,
|
|
timescale_dsn=timescale_dsn,
|
|
run_id=str(run["run_id"]),
|
|
)
|
|
case.update({
|
|
"analysis_run_id": str(run["run_id"]),
|
|
"status": "scada_materialized",
|
|
"scada": summary,
|
|
})
|
|
save_json(state_path, state)
|
|
print({"case_key": row["case_key"], "job_id": case["job_id"], "scada_rows": summary["rows"]}, flush=True)
|
|
return 0
|
|
|
|
|
|
def exchange_token(agent_env: Path) -> str:
|
|
import httpx
|
|
from dotenv import dotenv_values
|
|
from app.mcp.config import McpSettings
|
|
from scripts.run_operational_worker import CredentialProvider, WorkerSettings
|
|
|
|
values = dotenv_values(agent_env)
|
|
mcp = McpSettings()
|
|
|
|
async def exchange() -> str:
|
|
issuer = (values.get("TJWATER_MCP_DEV_HTTP_ISSUER_URI") or values["TJWATER_MCP_ISSUER_URI"]).rstrip("/")
|
|
async with httpx.AsyncClient(timeout=15) as client:
|
|
product = await CredentialProvider(WorkerSettings.from_environment()).headers(client)
|
|
response = await client.post(issuer + "/protocol/openid-connect/token", data={
|
|
"grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
|
|
"client_id": values["TJWATER_MCP_CLIENT_ID"],
|
|
"client_secret": values["TJWATER_MCP_CLIENT_SECRET"],
|
|
"subject_token": product["Authorization"].removeprefix("Bearer "),
|
|
"subject_token_type": "urn:ietf:params:oauth:token-type:access_token",
|
|
"requested_token_type": "urn:ietf:params:oauth:token-type:access_token",
|
|
"audience": mcp.audience,
|
|
"resource": mcp.resource_uri,
|
|
"scope": "mcp.catalog.read mcp.query mcp.context.bind",
|
|
})
|
|
response.raise_for_status()
|
|
return response.json()["access_token"]
|
|
|
|
return asyncio.run(exchange())
|
|
|
|
|
|
def mcp_call(client, token: str, project_id: str, name: str, arguments: dict) -> dict:
|
|
from mcp import types
|
|
|
|
method = "tools/call"
|
|
call_id = str(uuid4())
|
|
meta = {
|
|
types.PROTOCOL_VERSION_META_KEY: "2026-07-28",
|
|
types.CLIENT_CAPABILITIES_META_KEY: {},
|
|
types.CLIENT_INFO_META_KEY: {"name": "benchmark-generator", "version": "1"},
|
|
"io.tjwater/callId": call_id,
|
|
}
|
|
response = client.post("/mcp", headers={
|
|
"Authorization": "Bearer " + token,
|
|
"X-Project-Id": project_id,
|
|
"MCP-Protocol-Version": "2026-07-28",
|
|
"MCP-Method": method,
|
|
"MCP-Name": name,
|
|
"Accept": "application/json, text/event-stream",
|
|
}, json={
|
|
"jsonrpc": "2.0",
|
|
"id": call_id,
|
|
"method": method,
|
|
"params": {"_meta": meta, "name": name, "arguments": arguments},
|
|
})
|
|
response.raise_for_status()
|
|
payload = response.json()
|
|
result = payload.get("result") or {}
|
|
if payload.get("error") or result.get("isError"):
|
|
raise RuntimeError("MCP tool call failed")
|
|
return result["structuredContent"]["data"]
|
|
|
|
|
|
def submit_batch(*, batch, pipes, project_id, agent_env, state, state_path) -> list[str]:
|
|
from fastapi.testclient import TestClient
|
|
from app.main import app
|
|
from scripts.benchmark.calibration import save_json
|
|
|
|
token = exchange_token(agent_env)
|
|
session_id = "benchmark-tjwater2-burst-v1"
|
|
jobs = []
|
|
with TestClient(app, base_url="http://127.0.0.1:8000") as client:
|
|
common = {"Authorization": "Bearer " + token, "X-Project-Id": project_id}
|
|
for row in batch:
|
|
case = state["cases"].setdefault(row["case_key"], {})
|
|
request_id = case.setdefault("request_id", str(uuid4()))
|
|
scenario_name = case.setdefault(
|
|
"scenario_name",
|
|
f"bench-{row['case_key']}-{row['event_at'][0:16].replace('-', '').replace(':', '').replace('T', '')}",
|
|
)
|
|
binding = client.post(
|
|
"/api/v1/mcp-query-contexts",
|
|
headers=common,
|
|
json={"native_session_id": session_id, "purpose": "ordinary_chat"},
|
|
)
|
|
binding.raise_for_status()
|
|
context_handle = binding.json()["context_handle"]
|
|
pipe = pipes[row["pipe_id"]]
|
|
parameters = {
|
|
"start_time": row["event_at"],
|
|
"duration": 18000,
|
|
"output_node_ids": pipe.output_node_ids,
|
|
"output_link_ids": pipe.output_link_ids,
|
|
"burst_pipe_ids": [row["pipe_id"]],
|
|
"orifice_areas_cm2": [row["area_cm2"]],
|
|
}
|
|
preview = mcp_call(client, token, project_id, "save_simulation_scenario", {
|
|
"context_handle": context_handle,
|
|
"timeout_seconds": 240,
|
|
"request_id": request_id,
|
|
"dry_run": True,
|
|
"name": scenario_name,
|
|
"scenario": "burst",
|
|
"parameters": parameters,
|
|
})
|
|
case["expected_digest"] = preview["expected_digest"]
|
|
save_json(state_path, state)
|
|
queued = mcp_call(client, token, project_id, "save_simulation_scenario", {
|
|
"context_handle": context_handle,
|
|
"timeout_seconds": 240,
|
|
"request_id": request_id,
|
|
"dry_run": False,
|
|
"expected_digest": preview["expected_digest"],
|
|
"name": scenario_name,
|
|
"scenario": "burst",
|
|
"parameters": parameters,
|
|
})
|
|
case.update({"job_id": queued["job_id"], "status": queued["status"]})
|
|
jobs.append(queued["job_id"])
|
|
save_json(state_path, state)
|
|
return jobs
|
|
|
|
|
|
async def execute_batch(project_id: str, expected_job_ids: list[str], concurrency: int) -> None:
|
|
from app.services.calculation_worker import CalculationWorker, WorkerSettings, queue_repository
|
|
|
|
worker = CalculationWorker(WorkerSettings(
|
|
max_concurrency=concurrency,
|
|
poll_seconds=1,
|
|
lease_seconds=30,
|
|
heartbeat_seconds=5,
|
|
))
|
|
rows = []
|
|
for expected in expected_job_ids:
|
|
async with queue_repository(project_id) as repository:
|
|
await repository.recover()
|
|
row = await repository.claim(worker.owner, worker.settings.lease_seconds)
|
|
if row is None or str(row["job_id"]) != expected:
|
|
raise RuntimeError("isolated benchmark queue order changed unexpectedly")
|
|
rows.append(row)
|
|
await asyncio.gather(*(worker.run_job(row) for row in rows))
|
|
for row in rows:
|
|
async with queue_repository(project_id) as repository:
|
|
stored = await repository.get(row["job_id"])
|
|
if stored["status"] != "succeeded":
|
|
raise RuntimeError(f"benchmark calculation failed with {stored['status']}")
|
|
|
|
|
|
def find_analysis_run(business_dsn: str, scenario_name: str) -> dict:
|
|
from psycopg import connect
|
|
from psycopg.rows import dict_row
|
|
|
|
with connect(business_dsn, row_factory=dict_row) as connection:
|
|
rows = connection.execute(
|
|
"""SELECT run_id,name,status,started_at FROM analysis.runs
|
|
WHERE name=%s ORDER BY created_at DESC LIMIT 2""",
|
|
(scenario_name,),
|
|
).fetchall()
|
|
if len(rows) != 1 or rows[0]["status"] != "completed":
|
|
raise ValueError("persisted scenario run is missing, duplicated, or incomplete")
|
|
return dict(rows[0])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|