feat(mcp): queue long-running calculations

This commit is contained in:
2026-10-29 18:25:15 +08:00
parent 2865811f1b
commit 63af998045
32 changed files with 1588 additions and 72 deletions
+8
View File
@@ -3,6 +3,14 @@ import os
import pandas as pd
def calculation_temp_directory(kind: str) -> str:
"""Keep native outputs inside the supervisor-owned job directory."""
root = os.environ.get("TJWATER_CALCULATION_JOB_DIR", "temp")
path = os.path.abspath(os.path.join(root, kind))
os.makedirs(path, exist_ok=True)
return path
def fill_time_gaps(
data: pd.DataFrame,
time_col: str = "time",
@@ -8,15 +8,14 @@ import sys
import pandas as pd
import wntr
from app.algorithms._utils import _cleanup_temp_files
from app.algorithms._utils import _cleanup_temp_files, calculation_temp_directory
_PIPE2LEAKNODE = None
_SIGNATURE_WORKER_DATA = {}
def _make_temp_prefix(tag):
temp_dir = os.path.abspath(os.path.join("temp", "burst_location"))
os.makedirs(temp_dir, exist_ok=True)
temp_dir = calculation_temp_directory("burst_location")
safe_tag = str(tag).replace(os.sep, "_").replace(" ", "_")
return os.path.join(temp_dir, f"{safe_tag}_{os.getpid()}")
@@ -15,7 +15,7 @@ from pymoo.optimize import minimize as pymoo_minimize
from pymoo.termination.default import DefaultSingleObjectiveTermination
from app.algorithms._utils import _cleanup_temp_files
from app.algorithms._utils import _cleanup_temp_files, calculation_temp_directory
_worker_data: dict[str, Any] = {}
DEFAULT_N_WORKERS = max(1, min(cpu_count() - 1, 4))
@@ -92,7 +92,7 @@ def _worker_evaluate(raw_ratios: np.ndarray) -> float:
demand_obj.base_value = original_val + per_node_leak
modifications.append((demand_obj, original_val))
temp_dir = os.path.abspath(os.path.join("temp", "leakage"))
temp_dir = calculation_temp_directory("leakage")
os.makedirs(temp_dir, exist_ok=True)
prefix = os.path.join(temp_dir, f"temp_{os.getpid()}")
@@ -569,7 +569,7 @@ class LeakageProblem(Problem):
modifications.append((demand_obj, original_val))
# 结果保存在根目录的temp/leakage文件夹中
temp_dir = os.path.abspath(os.path.join("temp", "leakage"))
temp_dir = calculation_temp_directory("leakage")
os.makedirs(temp_dir, exist_ok=True)
prefix = os.path.join(temp_dir, f"temp_{os.getpid()}")
+103
View File
@@ -0,0 +1,103 @@
"""Durable project calculation queue. Every worker update is generation fenced."""
from uuid import uuid4
from fastapi import HTTPException
from psycopg.types.json import Jsonb
class CalculationJobRepository:
def __init__(self, connection, project_id):
self.connection = connection
self.project_id = str(project_id)
async def enqueue(self, *, owner_user_id, request_key, tool, parameters, binding, timeout_seconds):
cursor = await self.connection.execute(
"""INSERT INTO workbench.calculation_jobs
(job_id,project_id,owner_user_id,request_key,tool,parameters,binding,timeout_seconds)
VALUES (%s,%s,%s,%s,%s,%s,%s,%s)
ON CONFLICT(project_id,owner_user_id,request_key) DO NOTHING RETURNING *""",
(uuid4(), self.project_id, owner_user_id, request_key, tool,
Jsonb(parameters), Jsonb(binding), timeout_seconds),
)
row = await cursor.fetchone()
if row:
return dict(row)
cursor = await self.connection.execute(
"SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND owner_user_id=%s AND request_key=%s",
(self.project_id, owner_user_id, request_key),
)
row = dict(await cursor.fetchone())
if row['tool'] != tool or row['parameters'] != parameters or row['binding'] != binding:
raise HTTPException(409, 'Calculation request_key was used with different inputs')
return row
async def get(self, job_id, owner_user_id=None):
cursor = await self.connection.execute(
"""SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND job_id=%s
AND (%s::uuid IS NULL OR owner_user_id=%s)""",
(self.project_id, job_id, owner_user_id, owner_user_id),
)
row = await cursor.fetchone()
return dict(row) if row else None
async def list(self, *, owner_user_id, binding, limit, offset, status=None):
cursor = await self.connection.execute(
"""SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND owner_user_id=%s
AND binding @> %s AND (%s::text IS NULL OR status=%s) ORDER BY created_at DESC,job_id LIMIT %s OFFSET %s""",
(self.project_id, owner_user_id, Jsonb(binding), status, status, limit, offset),
)
return [dict(row) for row in await cursor.fetchall()]
async def cancel(self, job_id, owner_user_id):
await self.connection.execute(
"""UPDATE workbench.calculation_jobs SET status=CASE WHEN status='queued' THEN 'cancelled'
ELSE 'cancel_requested' END,updated_at=clock_timestamp()
WHERE project_id=%s AND job_id=%s AND owner_user_id=%s AND status IN ('queued','running')""",
(self.project_id, job_id, owner_user_id),
)
return await self.get(job_id, owner_user_id)
async def recover(self):
# A crashed write is never automatically replayed.
await self.connection.execute(
"""UPDATE workbench.calculation_jobs SET status='interrupted',
error='{"code":"WORKER_INTERRUPTED"}'::jsonb,updated_at=clock_timestamp()
WHERE project_id=%s AND status IN ('running','cancel_requested')
AND lease_expires_at < clock_timestamp()""", (self.project_id,),
)
async def claim(self, owner, lease_seconds):
cursor = await self.connection.execute(
"""UPDATE workbench.calculation_jobs j SET status='running',lease_owner=%s,
lease_generation=j.lease_generation+1,
lease_expires_at=clock_timestamp()+%s*interval '1 second',updated_at=clock_timestamp()
WHERE j.job_id=(SELECT job_id FROM workbench.calculation_jobs
WHERE project_id=%s AND status='queued' ORDER BY created_at,job_id
FOR UPDATE SKIP LOCKED LIMIT 1) RETURNING j.*""",
(owner, lease_seconds, self.project_id),
)
row = await cursor.fetchone()
return dict(row) if row else None
async def heartbeat(self, row, owner, lease_seconds):
cursor = await self.connection.execute(
"""UPDATE workbench.calculation_jobs SET lease_expires_at=clock_timestamp()+%s*interval '1 second',
updated_at=clock_timestamp() WHERE project_id=%s AND job_id=%s AND lease_owner=%s
AND lease_generation=%s AND lease_expires_at>clock_timestamp()
AND status IN ('running','cancel_requested') RETURNING status""",
(lease_seconds, self.project_id, row['job_id'], owner, row['lease_generation']),
)
result = await cursor.fetchone()
return result['status'] if result else None
async def finish(self, row, owner, status, *, result=None, error=None):
cursor = await self.connection.execute(
"""UPDATE workbench.calculation_jobs SET status=%s,result=%s,error=%s,updated_at=clock_timestamp()
WHERE project_id=%s AND job_id=%s AND lease_owner=%s AND lease_generation=%s
AND lease_expires_at>clock_timestamp() AND status IN ('running','cancel_requested')
AND (status!='cancel_requested' OR %s IN ('cancelled','interrupted')
OR (%s='failed' AND %s='CLEANUP_FAILED')) RETURNING job_id""",
(status, Jsonb(result) if result is not None else None, Jsonb(error) if error else None,
self.project_id, row['job_id'], owner, row['lease_generation'], status, status, (error or {}).get('code')),
)
return await cursor.fetchone() is not None
@@ -75,6 +75,20 @@ class OperationalTaskRepository:
status=409,
)
async def pending_calculation_jobs(
self, task_id: str, run_id: str, execution_request_id: str,
) -> list[str]:
"""Only the exact original run blocks final artifacts and completion."""
cursor = await self.connection.execute(
"""SELECT job_id FROM workbench.calculation_jobs
WHERE project_id=%s AND binding->>'task_id'=%s
AND binding->>'run_id'=%s AND binding->>'execution_request_id'=%s
AND status IN ('queued','running','cancel_requested')
ORDER BY created_at,job_id LIMIT 100""",
(self.project_id, task_id, run_id, execution_request_id),
)
return [str(row["job_id"]) for row in await cursor.fetchall()]
async def get_task(self, task_id: str) -> dict | None:
cursor = await self.connection.execute(
"""SELECT document FROM workbench.operational_tasks
+12 -7
View File
@@ -347,20 +347,25 @@ class McpQueryApplication:
)
result = QueryResult.model_validate(previous["result"])
else:
result = (
await self.executor(
scope, binding, tool, args, full_snapshot=True
from app.mcp.calculation_jobs import should_queue, submit_calculation
if should_queue(tool, args):
result = await submit_calculation(scope, binding, tool, args, call_id)
else:
result = (
await self.executor(
scope, binding, tool, args, full_snapshot=True
)
if include_full
else await self.executor(scope, binding, tool, args)
)
if include_full
else await self.executor(scope, binding, tool, args)
)
# Never serialize an unbounded response into a model message or database receipt.
if len(result.model_dump_json().encode()) > 8 * 1024 * 1024:
return tool_error(
"RESULT_TOO_LARGE",
"Narrow the query or request a smaller page",
)
if binding.task_id and name not in {
if binding.task_id and not should_queue(tool, args) and name not in {
"get_calculation_job", "get_calculation_result", "cancel_calculation_job", "list_calculation_jobs",
"get_task_context",
"get_task_evidence",
}:
+196
View File
@@ -0,0 +1,196 @@
"""Durable MCP calculation submission and owner-scoped result retrieval.
Only safe server-owned binding fields are persisted. OAuth tokens and context
handles are never used to resume a calculation.
"""
from datetime import datetime, timezone
from uuid import uuid5, NAMESPACE_URL
from fastapi import HTTPException
from app.mcp.models import BoundContext, QueryResult, PageInfo
from app.infra.db.postgresql.calculation_jobs import CalculationJobRepository
QUEUED_TOOLS = frozenset({
"run_burst_location", "run_leakage_identification", "run_burst_detection",
"run_simulation", "run_burst_simulation", "run_contaminant_simulation",
"analyze_valve_closure", "analyze_flushing", "analyze_pressure_regulation",
"optimize_sensor_placement", "save_simulation_scenario", "create_sensor_placement",
})
CONTROL_TOOLS = frozenset({"get_calculation_job", "get_calculation_result", "cancel_calculation_job", "list_calculation_jobs"})
BINDING_FIELDS = {"project_id", "subject", "client_id", "native_session_id", "purpose", "parent_session_id", "task_id", "run_id", "execution_request_id", "tool_cap", "permission_cap", "task_scope_cap"}
def durable_binding(binding: BoundContext) -> dict:
return binding.model_dump(mode="json", include=BINDING_FIELDS)
def restore_binding(document: dict) -> BoundContext:
# Expiration governs admission via a context handle; durable work has its own
# lifecycle and must revalidate identity, membership and run on every use.
return BoundContext.model_validate({**document, "context_handle": "durable-calculation", "expires_at": datetime.max.replace(tzinfo=timezone.utc)})
def should_queue(tool, args) -> bool:
return tool.name in QUEUED_TOOLS and (not tool.mutating or not args.dry_run)
def envelope(source: str, data, *, job_id="list", page=None) -> QueryResult:
return QueryResult(source=source, source_ref="calculation:" + str(job_id), queried_at=datetime.now(timezone.utc), semantics="task", data=data, page=page, truncated=bool(page and (page.offset or page.has_more)),
data_gaps=["Calculation status and queue receipts are not analytical evidence."])
def public_job(row: dict) -> dict:
fields = {"job_id", "status", "tool", "created_at", "started_at", "finished_at", "updated_at", "error"}
from fastapi.encoders import jsonable_encoder
return {**jsonable_encoder({key: row[key] for key in fields if key in row}), "poll_after_ms": 1500}
async def submit_calculation(scope, binding, tool, args, call_id: str) -> QueryResult:
from app.mcp.query_service import authorize_query
await authorize_query(scope, binding, tool, args)
if binding.purpose == "task_reader":
raise HTTPException(403, "Read-only task agents cannot start calculations")
if tool.mutating:
from app.mcp.mutation_service import (
preview_change, check_digest, mutation_arguments, retryable_capacity_rejection,
)
from app.infra.db.postgresql.mcp_queries import McpQueryRepository
mutation_repo = McpQueryRepository(scope.connection, scope.project.project_id)
mutation_key = "mutation:" + binding.subject + ":" + binding.native_session_id
approved_parameters = {"tool": tool.name, "arguments": mutation_arguments(args, preview=False)}
async with scope.connection.transaction():
await mutation_repo.lock_receipt(mutation_key, str(args.request_id))
previous = await mutation_repo.receipt(mutation_key, str(args.request_id))
if previous:
if previous["parameters"] != approved_parameters:
raise HTTPException(409, "request_id already has different parameters")
if previous["result"].get("status") == "completed":
return QueryResult.model_validate(previous["result"]["output"])
if not retryable_capacity_rejection(previous["result"]):
raise HTTPException(409, "Operation already started; inspect the existing calculation or mutation status")
check_digest(args, await preview_change(scope, binding, tool.name, args))
parameters = args.model_dump(mode="json", exclude={"context_handle"})
request_id = str(args.request_id) if tool.mutating else call_id
key = str(uuid5(NAMESPACE_URL, ":".join([binding.subject, binding.client_id, binding.native_session_id, str(binding.task_id), str(binding.run_id), request_id])))
repo = CalculationJobRepository(scope.connection, scope.project.project_id)
if binding.task_id:
from app.mcp.query_service import task_service
from app.mcp.contexts import verify_task
task_repo = task_service(scope).repo
# Share the operational transaction lock with finalization. Admission
# cannot race a pending-job check and then publish work into a closed run.
async with task_repo.transaction():
await task_repo.lock(f"task:{binding.task_id}")
await task_repo.lock(f"run:{binding.run_id}")
await verify_task(scope, binding)
row = await repo.enqueue(
owner_user_id=scope.project.user_id, request_key=key, tool=tool.name,
parameters=parameters, binding=durable_binding(binding), timeout_seconds=args.timeout_seconds,
)
else:
row = await repo.enqueue(
owner_user_id=scope.project.user_id, request_key=key, tool=tool.name,
parameters=parameters, binding=durable_binding(binding), timeout_seconds=args.timeout_seconds,
)
await scope.connection.commit()
return envelope(tool.name, public_job(row), job_id=row["job_id"])
def check_binding(binding, row):
original = row["binding"]
if any(str(original.get(key)) != str(getattr(binding, key)) for key in ("subject", "client_id", "native_session_id", "task_id", "run_id", "execution_request_id")):
raise HTTPException(403, "Calculation belongs to a different session or task run")
async def authorize_original(scope, row):
from app.mcp.catalog import TOOLS
from app.mcp.query_service import authorize_query
from app.mcp.contexts import verify_task
original = restore_binding(row["binding"])
tool = TOOLS[row["tool"]]
args = tool.model.model_validate({**row["parameters"], "context_handle": "durable-calculation-context-handle"})
if (str(scope.project.user_id) != str(row["owner_user_id"])
or str(scope.project.project_id) != str(original.project_id)):
raise HTTPException(403, "Calculation owner is no longer authorized")
if original.task_id:
await verify_task(scope, original)
await authorize_query(scope, original, tool, args)
return original, tool, args
async def read_calculation_job(scope, binding, tool, args) -> QueryResult:
repo = CalculationJobRepository(scope.connection, scope.project.project_id)
if tool.name == "list_calculation_jobs":
identity_scope = binding.model_dump(mode="json", include={
"subject", "client_id", "native_session_id", "task_id", "run_id", "execution_request_id",
})
rows = await repo.list(owner_user_id=scope.project.user_id, binding=identity_scope, limit=args.limit + 1, offset=args.offset, status=args.status)
return envelope(tool.name, [public_job(row) for row in rows[:args.limit]], page=PageInfo(limit=args.limit, offset=args.offset, has_more=len(rows) > args.limit))
row = await repo.get(args.job_id, owner_user_id=scope.project.user_id)
if row is None:
raise HTTPException(404, "Calculation was not found")
check_binding(binding, row)
await authorize_original(scope, row)
if tool.name == "cancel_calculation_job":
if binding.purpose == "task_reader":
raise HTTPException(403, "Read-only task agents cannot cancel calculations")
row = await repo.cancel(args.job_id, owner_user_id=scope.project.user_id)
await scope.connection.commit()
if tool.name != "get_calculation_result" or row["status"] != "succeeded":
return envelope(tool.name, public_job(row), job_id=row["job_id"])
result = QueryResult.model_validate(row["result"])
result = await attach_calculation_evidence(scope, row, result)
value = result.data
for part in args.data_path:
if isinstance(value, dict) and part in value:
value = value[part]
else:
raise HTTPException(422, "data_path must select an existing object field")
if isinstance(value, list):
result.data = value[args.offset:args.offset + args.limit]
result.page = PageInfo(limit=args.limit, offset=args.offset, total=len(value), has_more=args.offset + args.limit < len(value))
result.truncated = result.truncated or args.offset > 0 or result.page.has_more
elif args.offset:
raise HTTPException(422, "offset requires a collection data_path")
else:
result.data = value
if len(result.model_dump_json().encode()) > 11000:
raise HTTPException(413, "Select a narrower data_path or smaller result page")
return result
async def execute_calculation_job(scope, row) -> QueryResult:
"""Worker entry point; success is persisted by the fenced supervisor."""
from app.mcp.query_service import execute_query
original, tool, args = await authorize_original(scope, row)
result = await execute_query(scope, original, tool, args)
return result
async def attach_calculation_evidence(scope, row, result):
from app.mcp.query_service import task_service
from app.mcp.contexts import verify_task
original, tool, _ = await authorize_original(scope, row)
if original.task_id:
from app.domain.schemas.operational_task import TaskDataSnapshot
from app.infra.db.postgresql.mcp_queries import McpQueryRepository
repo = McpQueryRepository(scope.connection, scope.project.project_id)
key, call_id = "calculation:" + str(row["job_id"]), "result"
# Result lookup can leave an implicit SELECT transaction open. End it so
# evidence and its receipt commit together before returning a response.
await scope.connection.commit()
task_repo = task_service(scope).repo
# Match operational finalization's lock order: event lock, then task.
async with task_repo.transaction():
await repo.lock_receipt(key, call_id)
previous = await repo.receipt(key, call_id)
if previous:
return QueryResult.model_validate(previous["result"])
await task_repo.lock(f"task:{original.task_id}")
await verify_task(scope, original)
result.source_ref = key
evidence, _ = await task_service(scope).record_data_snapshot(original.task_id, TaskDataSnapshot(
source=tool.name, source_ref=key, queried_at=result.queried_at, origin=result.origin,
data=result.model_dump(mode="json"), truncated=result.truncated))
result.evidence_id = evidence.evidence_id
await repo.save_receipt(key, call_id, {"job_id": str(row["job_id"])}, result.model_dump(mode="json"))
return result
+21
View File
@@ -0,0 +1,21 @@
"""Bounded inputs for durable computation control."""
from typing import Literal
from uuid import UUID
from pydantic import Field
from app.mcp.models import StrictInput, PageInput, Id
CalculationStatus = Literal["queued", "running", "cancel_requested", "succeeded", "failed", "cancelled", "interrupted", "timed_out"]
class CalculationJobInput(StrictInput):
job_id: UUID
class CalculationResultInput(CalculationJobInput):
data_path: list[Id] = Field(default_factory=list, max_length=8)
limit: int = Field(default=20, ge=1, le=500)
offset: int = Field(default=0, ge=0, le=1000000)
class CalculationJobsInput(PageInput):
status: CalculationStatus | None = None
+13 -2
View File
@@ -4,6 +4,7 @@ from dataclasses import dataclass
from pydantic import BaseModel
from mcp import types
from app.mcp import models as m
from app.mcp import calculation_models as jobs
from app.mcp import extended_models as e
from app.mcp.analysis_models import BurstDetectionInput, BurstLocationInput, ValveIsolationInput, LeakageIdentificationInput
from app.mcp.analysis_models import ScenarioInput, ContaminantSimulationInput, ValveClosureInput, FlushingInput, PressureRegulationInput, SensorPlacementInput
@@ -36,11 +37,15 @@ class QueryTool:
)
return types.Tool(
name=self.name,
description=self.description + suffix,
description=self.description + suffix + (
"。执行返回持久job_id;排队回执不是分析证据,使用get_calculation_job查询状态、get_calculation_result分页读取成功结果,不重复提交。"
if self.name in {"run_burst_location", "run_leakage_identification", "run_burst_detection", "run_simulation", "run_burst_simulation", "run_contaminant_simulation", "analyze_valve_closure", "analyze_flushing", "analyze_pressure_regulation", "optimize_sensor_placement", "save_simulation_scenario", "create_sensor_placement"}
else ""
),
inputSchema=self.model.model_json_schema(),
outputSchema=m.QueryResult.model_json_schema(),
annotations=types.ToolAnnotations(
readOnlyHint=not (self.computation or self.mutating), destructiveHint=self.mutating or self.name == "run_simulation", openWorldHint=self.name in {"search_project_places", "reverse_geocode"}
readOnlyHint=not (self.computation or self.mutating or self.name == "cancel_calculation_job"), destructiveHint=self.mutating or self.name == "run_simulation", openWorldHint=self.name in {"search_project_places", "reverse_geocode"}
),
)
@@ -256,4 +261,10 @@ _specs += [
("retry_operational_task", "预览或重新入队已终止任务;不自动启动新Agent会话", e.TaskRetryInput, "operational_task.write", "task", False, False, True),
("cancel_task_run", "预览或取消指定任务运行,校验运行版本", e.TaskCancelInput, "operational_task.write", "task", False, False, True),
]
_specs += [
("get_calculation_job", "读取当前身份与会话的持久计算状态;排队回执不是分析证据", jobs.CalculationJobInput, "webgis.view", "task"),
("get_calculation_result", "读取成功计算的原始结果;data_path选择嵌套集合,limit/offset按需分页", jobs.CalculationResultInput, "webgis.view", "task"),
("cancel_calculation_job", "请求取消本人当前会话的计算;运行中取消由独立计算进程处理,须继续核实最终状态", jobs.CalculationJobInput, "webgis.view", "task"),
("list_calculation_jobs", "分页读取本人当前会话及任务范围的计算状态", jobs.CalculationJobsInput, "webgis.view", "task"),
]
TOOLS = {row[0]: QueryTool(*row) for row in _specs}
+1 -1
View File
@@ -70,7 +70,7 @@ async def bind_context(
subject=principal.subject,
client_id=principal.client_id,
tool_cap=sorted(name for name, tool in TOOLS.items()
if (command.purpose != "task_reader" or not (tool.computation or tool.mutating))
if (command.purpose != "task_reader" or not (tool.computation or tool.mutating or name == "cancel_calculation_job"))
and (not tool.mutating or command.purpose == "ordinary_chat")),
permission_cap=sorted(permissions_for_context(scope.project)),
native_session_id=command.native_session_id,
+3
View File
@@ -131,6 +131,9 @@ async def execute_query(
full_snapshot: bool = False,
) -> QueryResult:
await authorize_query(scope, binding, tool, args)
if tool.name in {"get_calculation_job", "get_calculation_result", "cancel_calculation_job", "list_calculation_jobs"}:
from app.mcp.calculation_jobs import read_calculation_job
return await read_calculation_job(scope, binding, tool, args)
if tool.mutating:
from app.mcp.mutation_service import execute_mutation
return await execute_mutation(scope, binding, tool, args)
+6
View File
@@ -201,6 +201,8 @@ def temporary_project_name(project: str, purpose: str) -> str:
def temporary_project_database(project: str, purpose: str):
"""Clone one project's runnable model into an isolated temporary database."""
temporary_name = temporary_project_name(project, purpose)
from app.services.calculation_cleanup import record_calculation_database
record_calculation_database(temporary_name)
try:
copy_project(
get_project_template_database_name(project),
@@ -211,12 +213,15 @@ def temporary_project_database(project: str, purpose: str):
finally:
if have_project(temporary_name):
delete_project(temporary_name)
record_calculation_database(temporary_name, removed=True)
@contextmanager
def temporary_template_database(name_hint: str, purpose: str):
"""Create an empty schema-only temporary database from the fixed template."""
temporary_name = temporary_project_name(name_hint, purpose)
from app.services.calculation_cleanup import record_calculation_database
record_calculation_database(temporary_name)
try:
copy_project(
get_schema_template_database_name(),
@@ -227,6 +232,7 @@ def temporary_template_database(name_hint: str, purpose: str):
finally:
if have_project(temporary_name):
delete_project(temporary_name)
record_calculation_database(temporary_name, removed=True)
def have_project(name: str) -> bool:
+74
View File
@@ -0,0 +1,74 @@
"""Record only job-owned native databases so abrupt cancellation is recoverable."""
import asyncio
import json
import os
from pathlib import Path
import re
from uuid import UUID
def _names(directory: Path) -> list[str]:
path = directory / "databases.json"
if not path.exists():
return []
if path.is_symlink() or path.stat().st_size > 65536:
raise ValueError("Invalid calculation database ledger")
names = json.loads(path.read_text())
if not isinstance(names, list) or len(names) > 100 or any(
not isinstance(name, str) or not re.fullmatch(r"tjw_tmp_[a-z0-9_]+_[a-f0-9]{32}", name)
or len(name) > 63 for name in names
):
raise ValueError("Invalid calculation database ledger")
return names
def _write(directory: Path, names: list[str]) -> None:
path = directory / "databases.json"
staging = directory / "databases.json.new"
with staging.open("w") as output:
json.dump(names, output)
output.flush()
os.fsync(output.fileno())
staging.replace(path)
def record_calculation_database(name: str, *, removed: bool = False) -> None:
"""Called before clone creation and after successful normal cleanup."""
root = os.environ.get("TJWATER_CALCULATION_JOB_DIR")
if not root:
return
directory = Path(root)
names = _names(directory)
if removed:
names = [value for value in names if value != name]
elif name not in names:
names.append(name)
_write(directory, names)
async def cleanup_calculation_databases(project_id: str, job_directory: str) -> None:
directory = Path(job_directory)
names = _names(directory)
if not names:
return
from app.infra.db.metadb.database import SessionLocal
from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository
from app.infra.db.project_routing import ActiveProjectRouting, activate_project_routing
from app.native.wndb.core.projects import have_project, delete_project
async with SessionLocal() as session:
metadata = MetadataRepository(session)
project = await metadata.get_project_by_id(UUID(str(project_id)))
route = await metadata.get_project_db_routing(UUID(str(project_id)), "biz_data")
if not project or not route or route.db_type != "postgresql":
raise ValueError("Calculation cleanup routing unavailable")
routing = ActiveProjectRouting(project.code, route.dsn, project_id=project.id)
await session.commit()
def cleanup():
with activate_project_routing(routing):
for name in names[:]:
if have_project(name):
delete_project(name)
names.remove(name)
_write(directory, names)
await asyncio.to_thread(cleanup)
+318
View File
@@ -0,0 +1,318 @@
"""Independent, singleton calculation supervisor and isolated job subprocesses.
Never imported by API lifespan. Child output contains the original QueryResult;
algorithm stdout/stderr is discarded rather than interpreted as a protocol.
"""
from __future__ import annotations
import asyncio
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass
from datetime import datetime, timezone
import json
import os
from pathlib import Path
import signal
import shutil
import sys
import tempfile
import time
from uuid import UUID, uuid4
@dataclass(frozen=True)
class WorkerSettings:
project_ids: tuple[str, ...]
max_concurrency: int = 2
poll_seconds: float = 2
lease_seconds: int = 30
heartbeat_seconds: float = 5
@classmethod
def from_environment(cls):
projects = tuple(dict.fromkeys(str(UUID(p.strip())) for p in
os.environ['CALCULATION_WORKER_PROJECT_IDS'].split(',')))
result = cls(projects, int(os.getenv('CALCULATION_WORKER_MAX_CONCURRENCY', '2')),
float(os.getenv('CALCULATION_WORKER_POLL_SECONDS', '2')),
int(os.getenv('CALCULATION_WORKER_LEASE_SECONDS', '30')),
float(os.getenv('CALCULATION_WORKER_HEARTBEAT_SECONDS', '5')))
if not projects or not 1 <= result.max_concurrency <= 32:
raise ValueError('Configure explicit project IDs and concurrency between 1 and 32')
if not 0 < result.heartbeat_seconds < result.lease_seconds / 3 or result.poll_seconds <= 0:
raise ValueError('Heartbeat must be positive and below one third of the lease')
return result
@asynccontextmanager
async def queue_repository(project_id):
"""Resolve current routing from metadata, never from a job payload."""
from psycopg import AsyncConnection
from psycopg.rows import dict_row
from app.infra.db.metadb.database import SessionLocal
from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository
from app.infra.db.postgresql.calculation_jobs import CalculationJobRepository
async with SessionLocal() as session:
metadata = MetadataRepository(session)
project = await metadata.get_project_by_id(UUID(str(project_id)))
if not project or project.status != 'active':
raise ValueError('Calculation project is not active')
routing = await metadata.get_project_db_routing(UUID(str(project_id)), 'biz_data')
if not routing or routing.db_type != 'postgresql':
raise ValueError('Project business PostgreSQL routing is unavailable')
await session.commit()
async with await AsyncConnection.connect(routing.dsn, row_factory=dict_row) as conn:
yield CalculationJobRepository(conn, project_id)
@asynccontextmanager
async def execution_scope(row):
from fastapi import HTTPException
from app.mcp.scope import QueryMetadataRepository, QueryScope
from app.auth.project_dependencies import (resolve_project_context,
resolve_project_business_routing, get_project_pg_connection)
from app.infra.db.metadb.database import SessionLocal
from app.infra.db.project_routing import activate_project_routing
async with SessionLocal() as session:
metadata = QueryMetadataRepository(session)
user = await metadata.get_user_by_id(UUID(str(row['owner_user_id'])))
if not user or not user.is_active:
raise HTTPException(403, 'Calculation owner is inactive')
project = await resolve_project_context(str(row['project_id']), user, metadata)
routing = await resolve_project_business_routing(project, metadata)
await session.commit()
async with asynccontextmanager(get_project_pg_connection)(project, metadata) as conn:
await session.commit()
with activate_project_routing(routing):
yield QueryScope(project, user, conn, routing, metadata)
async def execute_child(input_path: str, output_path: str):
from app.mcp.calculation_jobs import execute_calculation_job
row = json.loads(Path(input_path).read_text())
try:
async with execution_scope(row) as scope:
result = await execute_calculation_job(scope, row)
payload = {'status': 'succeeded', 'result': result.model_dump(mode='json')}
except Exception as exc:
# Do not persist exception text: DSNs, SQL, credentials or model data can occur there.
category = 'CALCULATION_FAILED'
if getattr(exc, 'status_code', None) in (401, 403):
category = 'AUTHORIZATION_DENIED'
elif getattr(exc, 'status_code', None) == 409:
category = 'EXECUTION_CONFLICT'
elif getattr(exc, 'status_code', None) in (404, 410):
category = 'EXECUTION_CONTEXT_UNAVAILABLE'
payload = {'status': 'failed', 'error': {'code': category}}
encoded = json.dumps(payload, allow_nan=False)
if len(encoded.encode()) > 8 * 1024 * 1024:
encoded = json.dumps({'status': 'failed', 'error': {'code': 'RESULT_TOO_LARGE'}})
Path(output_path).write_text(encoded)
async def stop_process_group(process):
"""Always kill the group, even if its original leader already exited."""
try:
os.killpg(process.pid, signal.SIGTERM)
except ProcessLookupError:
pass
try:
await asyncio.wait_for(process.wait(), 2)
except TimeoutError:
pass
# Native pools can ignore SIGTERM or survive the Python leader.
try:
os.killpg(process.pid, signal.SIGKILL)
except ProcessLookupError:
pass
await process.wait()
if process.stdin is not None:
process.stdin.close()
@contextmanager
def job_directory(row):
root = Path(os.getenv('CALCULATION_WORKER_STATE_DIR', 'temp/calculation-jobs'))
root.mkdir(parents=True, exist_ok=True)
directory = Path(tempfile.mkdtemp(prefix=f"{row.get('job_id', 'job')}-", dir=root))
directory.chmod(0o700)
(directory / 'scope.json').write_text(json.dumps({'project_id': str(row['project_id']), 'job_id': str(row.get('job_id', 'job')),
'lease_generation': row.get('lease_generation')}))
try:
yield str(directory.resolve())
finally:
# Failed external resource cleanup must remain operator-recoverable.
if not (directory / 'cleanup-failed').exists():
shutil.rmtree(directory)
async def cleanup_job_resources(row, directory):
from app.services.calculation_cleanup import cleanup_calculation_databases
try:
await cleanup_calculation_databases(row['project_id'], directory)
except BaseException:
(Path(directory) / 'cleanup-failed').touch()
raise
async def recover_job_resources(project_ids):
"""Reconcile only this worker's explicit scopes after obtaining the singleton."""
root = Path(os.getenv('CALCULATION_WORKER_STATE_DIR', 'temp/calculation-jobs'))
if not root.exists():
return
for directory in root.iterdir():
if directory.is_symlink() or not directory.is_dir():
continue
scope_file = directory / 'scope.json'
if not scope_file.exists():
continue
row = json.loads(scope_file.read_text())
if str(row.get('project_id')) not in project_ids:
continue
job_id = str(UUID(row['job_id']))
if not directory.name.startswith(job_id + '-'):
raise RuntimeError('Calculation state directory identity is invalid')
async with queue_repository(row['project_id']) as repository:
stored = await repository.get(job_id)
if not stored or stored['lease_generation'] != row['lease_generation']:
raise RuntimeError('Calculation state ownership cannot be verified')
expiry = stored.get('lease_expires_at')
if expiry and expiry > datetime.now(timezone.utc) and stored['status'] in ('running', 'cancel_requested'):
raise RuntimeError('Previous calculation lease has not expired')
await cleanup_job_resources(row, str(directory.resolve()))
async with queue_repository(row['project_id']) as repository:
await repository.recover()
shutil.rmtree(directory)
class CalculationWorker:
def __init__(self, settings, *, repository=queue_repository):
self.settings = settings
self.repository = repository
self.owner = str(uuid4())
async def run_job(self, row):
from fastapi.encoders import jsonable_encoder
with job_directory(row) as directory:
input_path = Path(directory) / 'input.json'
output_path = Path(directory) / 'result.json'
input_path.write_text(json.dumps(jsonable_encoder(row)))
process = await asyncio.create_subprocess_exec(
sys.executable, '-m', 'app.services.calculation_worker',
'--guard', str(input_path), str(output_path), start_new_session=True,
stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.DEVNULL,
stderr=asyncio.subprocess.DEVNULL,
env={**os.environ, 'TMPDIR': directory, 'TJWATER_CALCULATION_JOB_DIR': directory},
)
started = time.monotonic()
cleaned = False
try:
while True:
async with asyncio.timeout(self.settings.heartbeat_seconds):
async with self.repository(row['project_id']) as repository:
status = await repository.heartbeat(row, self.owner, self.settings.lease_seconds)
if status is None:
return # Finally stops all owned computation on a lost lease.
terminal = None
if status == 'cancel_requested':
terminal = {'status': 'cancelled'}
elif time.monotonic() - started >= row['timeout_seconds']:
terminal = {'status': 'timed_out', 'error': {'code': 'EXECUTION_TIMEOUT'}}
elif process.returncode is not None:
try:
if output_path.stat().st_size > 8 * 1024 * 1024:
raise ValueError('Result exceeds persistence limit')
terminal = json.loads(output_path.read_text())
except (OSError, ValueError):
terminal = {'status': 'failed', 'error': {'code': 'CHILD_PROCESS_FAILED'}}
if terminal:
await stop_process_group(process)
try:
await cleanup_job_resources(row, directory)
cleaned = True
except Exception:
terminal = {'status': 'failed', 'error': {'code': 'CLEANUP_FAILED'}}
cleaned = True
async with self.repository(row['project_id']) as repository:
finished = await repository.finish(row, self.owner, **terminal)
if not finished:
status = await repository.heartbeat(row, self.owner, self.settings.lease_seconds)
if status == 'cancel_requested':
await repository.finish(row, self.owner, 'cancelled')
return
await asyncio.sleep(self.settings.heartbeat_seconds)
finally:
await stop_process_group(process)
if not cleaned:
await cleanup_job_resources(row, directory)
async def run(self):
from sqlalchemy import text
from app.infra.db.metadb.database import engine
# Dedicated session-level metadata lock: independent API workers never
# participate; duplicate supervisors exit before claiming any jobs.
async with engine.connect() as lock:
acquired = await lock.scalar(text("SELECT pg_try_advisory_lock(1748531781, 2)"))
await lock.commit()
if not acquired:
raise RuntimeError('Another calculation supervisor is already active')
tasks = set()
project_cursor = 0
try:
# Allow any old supervisor to detect a broken lock connection and
# stop its children before this replacement claims capacity.
await asyncio.sleep(self.settings.lease_seconds)
await recover_job_resources(self.settings.project_ids)
while True:
async with asyncio.timeout(self.settings.heartbeat_seconds):
await lock.execute(text('SELECT 1'))
await lock.commit()
for task in tuple(tasks):
if task.done():
tasks.remove(task)
task.result() # Fail closed if heartbeat/persistence failed.
projects = self.settings.project_ids
ordered = projects[project_cursor:] + projects[:project_cursor]
project_cursor = (project_cursor + 1) % len(projects)
for project_id in ordered:
if len(tasks) >= self.settings.max_concurrency:
break
async with self.repository(project_id) as repository:
await repository.recover()
row = await repository.claim(self.owner, self.settings.lease_seconds)
if row:
tasks.add(asyncio.create_task(self.run_job(row)))
await asyncio.sleep(min(self.settings.poll_seconds, self.settings.heartbeat_seconds))
finally:
for task in tasks:
task.cancel()
await asyncio.gather(*tasks, return_exceptions=True)
await lock.execute(text('SELECT pg_advisory_unlock(1748531781, 2)'))
await lock.commit()
async def guard_child(input_path, output_path):
"""A parent-death pipe kills the whole group even when the supervisor is SIGKILLed."""
reader = asyncio.StreamReader()
protocol = asyncio.StreamReaderProtocol(reader)
transport, _ = await asyncio.get_running_loop().connect_read_pipe(lambda: protocol, sys.stdin)
process = await asyncio.create_subprocess_exec(
sys.executable, '-m', 'app.services.calculation_worker', input_path, output_path,
stdin=asyncio.subprocess.DEVNULL,
)
completed = asyncio.create_task(process.wait())
parent_gone = asyncio.create_task(reader.read(1))
try:
done, _ = await asyncio.wait((completed, parent_gone), return_when=asyncio.FIRST_COMPLETED)
if parent_gone in done:
os.killpg(os.getpgrp(), signal.SIGKILL)
await completed
finally:
parent_gone.cancel()
transport.close()
if __name__ == '__main__':
if sys.argv[1] == '--guard':
asyncio.run(guard_child(sys.argv[2], sys.argv[3]))
else:
asyncio.run(execute_child(sys.argv[1], sys.argv[2]))
+14
View File
@@ -778,6 +778,10 @@ class OperationalTaskService:
"Only a running run accepts Agent artifacts", status=409)
now = self._now()
self._require_active_run_lease(active_run, now)
if command.report and (
command.report.stage == "final" or command.report.outcome == "insufficient_data"
):
await self._require_no_pending_calculations(task_id, active_run)
previous_report = await self.repo.get_latest_report(str(task_id))
require(not previous_report or previous_report.get("execution_generation") != task["execution_generation"]
or previous_report["report"].get("outcome") != "insufficient_data",
@@ -1866,6 +1870,15 @@ class OperationalTaskService:
task_ids.append(task_id)
return DispatchDueSchedulesResult(task_ids=task_ids, dispatched_at=now)
async def _require_no_pending_calculations(self, task_id: UUID, run: dict) -> None:
if hasattr(self.repo, "pending_calculation_jobs"):
pending = await self.repo.pending_calculation_jobs(
str(task_id), str(run["run_id"]), str(run["execution_request_id"]),
)
require(not pending, "PENDING_CALCULATIONS",
"Wait for submitted calculations before publishing a final report or completing the run",
refs=pending, status=409, recovery="retry_same_request")
async def _finish_run(
self, task_id: UUID, run_id: UUID, command: AgentRunCommand,
*, outcome: str, failure_code: str | None = None,
@@ -1890,6 +1903,7 @@ class OperationalTaskService:
task = await self._task(task_id)
final_outcome = outcome
if outcome == "succeeded":
await self._require_no_pending_calculations(task_id, run)
if hasattr(self.repo, "pending_run_messages"):
pending_messages = await self.repo.pending_run_messages(str(run_id))
require(not pending_messages, "PENDING_AGENT_MESSAGES",
+5 -2
View File
@@ -2,6 +2,7 @@ from collections.abc import Iterator
from contextlib import contextmanager
from pathlib import Path
from tempfile import NamedTemporaryFile
import os
from app.native.wndb.inp.exporter import dump_inp
@@ -17,9 +18,11 @@ def temporary_project_inp(
version: str = "2",
) -> Iterator[Path]:
"""Export one project model to an isolated INP and remove it afterwards."""
PROJECT_INP_DIRECTORY.mkdir(parents=True, exist_ok=True)
job_directory = os.environ.get("TJWATER_CALCULATION_JOB_DIR")
directory = Path(job_directory) / "inp" if job_directory else PROJECT_INP_DIRECTORY
directory.mkdir(parents=True, exist_ok=True)
with NamedTemporaryFile(
dir=PROJECT_INP_DIRECTORY,
dir=directory,
prefix=f"{purpose}_",
suffix=".inp",
delete=False,