156 lines
6.7 KiB
Python
156 lines
6.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Run one isolated online-model day through MCP and materialize noisy SCADA."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import sys
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--start", default="2026-05-01T00:00:00+08:00")
|
|
parser.add_argument(
|
|
"--generation-state",
|
|
type=Path,
|
|
default=Path("data/benchmark-generation/tjwater2-20261105/state.json"),
|
|
)
|
|
parser.add_argument(
|
|
"--state",
|
|
type=Path,
|
|
default=Path("data/benchmark-generation/tjwater2-20261105/normal-day-state.json"),
|
|
)
|
|
parser.add_argument(
|
|
"--design",
|
|
type=Path,
|
|
default=Path("resources/benchmarks/tjwater2-v2-burst-balanced.json"),
|
|
)
|
|
parser.add_argument("--agent-env", type=Path, default=Path("../next-tjwater-agent/.local.env"))
|
|
parser.add_argument("--noise-seed", default="tjwater2-normal-scada-v1")
|
|
args = parser.parse_args()
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
os.umask(0o077)
|
|
|
|
from dotenv import load_dotenv
|
|
from fastapi.testclient import TestClient
|
|
from psycopg import connect
|
|
from psycopg.conninfo import make_conninfo
|
|
from psycopg.rows import dict_row
|
|
|
|
from app.main import app
|
|
from scripts.benchmark.calibration import save_json
|
|
from scripts.benchmark.generation import write_realtime_scada_day
|
|
from scripts.benchmark.simulation_design import load_simulation_design
|
|
from scripts.check_burst_simulation_design import resolve_project
|
|
from scripts.run_burst_benchmark_mcp import exchange_token, execute_batch, mcp_call
|
|
|
|
load_dotenv(Path(__file__).resolve().parents[1] / ".env", override=False)
|
|
start = datetime.fromisoformat(args.start)
|
|
if start.tzinfo is None or start.minute % 15 or start.second or start.microsecond:
|
|
raise ValueError("start must be timezone-aware and aligned to 15 minutes")
|
|
generation = json.loads(args.generation_state.read_text(encoding="utf-8"))
|
|
if generation.get("status") != "prepared" or not generation.get("project_id"):
|
|
raise ValueError("isolated benchmark project is not prepared")
|
|
state_path = args.state.resolve()
|
|
state_path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
|
|
state = json.loads(state_path.read_text(encoding="utf-8")) if state_path.exists() else {
|
|
"version": "normal-scada-day@1",
|
|
"target": generation["target"],
|
|
"project_id": generation["project_id"],
|
|
"start": start.isoformat(),
|
|
"duration_seconds": 86400,
|
|
"noise_seed": args.noise_seed,
|
|
"status": "preparing",
|
|
}
|
|
expected = {
|
|
"target": generation["target"],
|
|
"project_id": generation["project_id"],
|
|
"start": start.isoformat(),
|
|
"noise_seed": args.noise_seed,
|
|
}
|
|
if any(state.get(key) != value for key, value in expected.items()):
|
|
raise ValueError("normal-day state belongs to different inputs")
|
|
save_json(state_path, state)
|
|
|
|
design = load_simulation_design(args.design)
|
|
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"])
|
|
with connect(business_dsn, row_factory=dict_row) as connection:
|
|
node = connection.execute(
|
|
"""SELECT node_id FROM asset.scada_devices
|
|
WHERE device_type='pressure' AND node_id IS NOT NULL ORDER BY device_id LIMIT 1"""
|
|
).fetchone()["node_id"]
|
|
link = connection.execute(
|
|
"""SELECT link_id FROM asset.scada_devices
|
|
WHERE device_type='pipe_flow' AND link_id IS NOT NULL ORDER BY device_id LIMIT 1"""
|
|
).fetchone()["link_id"]
|
|
|
|
if not state.get("job_id"):
|
|
token = exchange_token(args.agent_env)
|
|
with TestClient(app, base_url="http://127.0.0.1:8000") as client:
|
|
headers = {
|
|
"Authorization": "Bearer " + token,
|
|
"X-Project-Id": state["project_id"],
|
|
}
|
|
binding = client.post(
|
|
"/api/v1/mcp-query-contexts",
|
|
headers=headers,
|
|
json={"native_session_id": "benchmark-normal-scada-v1", "purpose": "ordinary_chat"},
|
|
)
|
|
binding.raise_for_status()
|
|
queued = mcp_call(client, token, state["project_id"], "run_simulation", {
|
|
"context_handle": binding.json()["context_handle"],
|
|
"timeout_seconds": 240,
|
|
"start_time": start.isoformat(),
|
|
"duration": 86400,
|
|
"output_node_ids": [node],
|
|
"output_link_ids": [link],
|
|
})
|
|
state.update({"job_id": queued["job_id"], "status": queued["status"]})
|
|
save_json(state_path, state)
|
|
print({"job_id": state["job_id"], "status": state["status"]}, flush=True)
|
|
|
|
async def stored_job(include_result: bool = False):
|
|
from app.services.calculation_worker import queue_repository
|
|
|
|
async with queue_repository(state["project_id"]) as repository:
|
|
return await repository.get(state["job_id"], include_result=include_result)
|
|
|
|
stored = asyncio.run(stored_job())
|
|
if stored["status"] == "queued":
|
|
asyncio.run(execute_batch(state["project_id"], [state["job_id"]], 1))
|
|
stored = asyncio.run(stored_job())
|
|
if stored["status"] != "succeeded":
|
|
raise RuntimeError(f"normal online simulation ended with {stored['status']}")
|
|
result = asyncio.run(stored_job(include_result=True))["result"]
|
|
calculation_data = (result or {}).get("data") or {}
|
|
scenario_result = calculation_data.get("result") or calculation_data
|
|
units = scenario_result.get("units") or {}
|
|
if str(units.get("pressure", "")).upper() not in {"MTR", "M", "METERS"}:
|
|
raise ValueError(f"unexpected model pressure unit: {units.get('pressure')}")
|
|
if str(units.get("flow", "")).upper() != "LPS":
|
|
raise ValueError(f"unexpected model flow unit: {units.get('flow')}")
|
|
if not state.get("scada"):
|
|
state["status"] = "simulated"
|
|
save_json(state_path, state)
|
|
state["scada"] = write_realtime_scada_day(
|
|
business_dsn=business_dsn,
|
|
timescale_dsn=timescale_dsn,
|
|
start=start,
|
|
noise_seed=args.noise_seed,
|
|
)
|
|
state["model_units"] = units
|
|
state["status"] = "scada_materialized"
|
|
save_json(state_path, state)
|
|
print(json.dumps({"status": state["status"], "job_id": state["job_id"], "scada": state["scada"]}, ensure_ascii=False), flush=True)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|