feat(mcp): queue long-running calculations
This commit is contained in:
@@ -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()}")
|
||||
|
||||
|
||||
@@ -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
@@ -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",
|
||||
}:
|
||||
|
||||
@@ -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
|
||||
@@ -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
@@ -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
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
@@ -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]))
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user