129 lines
6.1 KiB
Python
129 lines
6.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate persisted isolated burst scenarios and mapped SCADA observations."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import os
|
|
import sys
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
|
|
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("--target", default="tjw_bench_tjwater2_20261105")
|
|
parser.add_argument("--state-dir", type=Path, default=Path("data/benchmark-generation/tjwater2-20261105"))
|
|
parser.add_argument("--apply", action="store_true")
|
|
parser.add_argument("--prepare-only", action="store_true")
|
|
args = parser.parse_args()
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
os.umask(0o077)
|
|
from psycopg import connect
|
|
from psycopg.conninfo import make_conninfo
|
|
from psycopg.rows import dict_row
|
|
from app.core.config import settings
|
|
from scripts.benchmark.calibration import save_json, digest
|
|
from scripts.benchmark.generation import (
|
|
apply_project_runtime_migrations,
|
|
create_empty_database,
|
|
copy_project_inputs,
|
|
make_schedule,
|
|
register_benchmark_project,
|
|
)
|
|
from scripts.benchmark.simulation_design import load_simulation_design
|
|
from scripts.check_burst_simulation_design import resolve_project, readonly, validate_database
|
|
|
|
design = load_simulation_design(args.design)
|
|
schedule = make_schedule(design)
|
|
if not args.apply:
|
|
print(f"dry run: {len(schedule)} cases, target={args.target}, dates={schedule[0]['history_start']}..{schedule[-1]['reserved_end']}")
|
|
return 0
|
|
folder = args.state_dir.resolve()
|
|
folder.mkdir(parents=True, exist_ok=True, mode=0o700)
|
|
state_path = folder / "state.json"
|
|
import json
|
|
state = json.loads(state_path.read_text()) if state_path.exists() else {
|
|
"version": "burst-dataset-generation@1", "created_at": datetime.now(timezone.utc).isoformat(),
|
|
"target": args.target, "design_sha256": digest(args.design), "schedule": schedule,
|
|
"stages": {}, "cases": {}, "source_effect": "read-only", "status": "preparing"}
|
|
if state["target"] != args.target or state["design_sha256"] != digest(args.design):
|
|
raise ValueError("state belongs to a different target or design")
|
|
save_json(state_path, state)
|
|
biz, iot = asyncio.run(resolve_project(design))
|
|
with readonly(biz) as connection:
|
|
validate_database(connection, design, template=False)
|
|
target_biz, target_iot = make_conninfo(biz, dbname=args.target), make_conninfo(iot, dbname=args.target)
|
|
for role, source, template in (("biz", biz, settings.WNDB_SCHEMA_TEMPLATE_DB_NAME),
|
|
("iot", iot, settings.TIMESCALEDB_SCHEMA_TEMPLATE_DB_NAME)):
|
|
if not state["stages"].get(f"{role}_created"):
|
|
create_empty_database(source, args.target, template)
|
|
state["stages"][f"{role}_created"] = True
|
|
save_json(state_path, state)
|
|
print(f"created isolated {role} database: {args.target}", flush=True)
|
|
if not state["stages"].get("model_copied"):
|
|
state["model_counts"] = copy_project_inputs(biz, target_biz)
|
|
state["stages"]["model_copied"] = True
|
|
save_json(state_path, state)
|
|
print("copied model and device bindings from a read-only source snapshot", flush=True)
|
|
required_migrations = {
|
|
"025_mcp_query_contexts.sql",
|
|
"026_calculation_jobs.sql",
|
|
"027_calculation_job_lifecycle.sql",
|
|
"030_remove_calculation_queue_expiry.sql",
|
|
}
|
|
if not required_migrations.issubset(state.get("runtime_migrations", [])):
|
|
state["runtime_migrations"] = apply_project_runtime_migrations(
|
|
target_biz, Path(__file__).resolve().parents[1]
|
|
)
|
|
state["stages"]["runtime_migrations_applied"] = True
|
|
save_json(state_path, state)
|
|
print("installed isolated MCP and calculation queue storage", flush=True)
|
|
if not state["stages"].get("model_template_created"):
|
|
from app.native.wndb.core.project_templates import (
|
|
create_project_model_template,
|
|
ensure_replication_worker_capacity,
|
|
)
|
|
from app.native.wndb.core.projects import have_project
|
|
|
|
template = f"{args.target}_template"
|
|
if not have_project(template):
|
|
try:
|
|
ensure_replication_worker_capacity()
|
|
except RuntimeError as exc:
|
|
if "no reserved logical-replication worker capacity" not in str(exc):
|
|
raise
|
|
state["model_template_mode"] = "snapshot_per_run"
|
|
else:
|
|
create_project_model_template(args.target)
|
|
if have_project(template):
|
|
state["model_template"] = template
|
|
state["model_template_mode"] = "logical_replication"
|
|
state["stages"]["model_template_created"] = True
|
|
save_json(state_path, state)
|
|
print(f"configured hydraulic model isolation: {state['model_template_mode']}", flush=True)
|
|
if not state["stages"].get("metadata_registered"):
|
|
project = asyncio.run(register_benchmark_project(
|
|
source_project_id=str(design.project_id),
|
|
target=args.target,
|
|
target_biz_dsn=target_biz,
|
|
target_iot_dsn=target_iot,
|
|
))
|
|
state["project_id"] = str(project.id)
|
|
state["stages"]["metadata_registered"] = True
|
|
save_json(state_path, state)
|
|
print(f"registered isolated benchmark project: {project.id}", flush=True)
|
|
# Inspect isolated schema readiness before invoking backend persistence.
|
|
with connect(target_biz, row_factory=dict_row) as connection:
|
|
print("isolated BizDB analysis schema", connection.execute(
|
|
"SELECT column_name FROM information_schema.columns WHERE table_schema='analysis' AND table_name='runs' ORDER BY ordinal_position").fetchall(), flush=True)
|
|
state["status"] = "prepared"
|
|
save_json(state_path, state)
|
|
print(f"isolated preparation complete: {state_path}", flush=True)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|