300 lines
14 KiB
Python
300 lines
14 KiB
Python
"""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 app.infra.db.calculation_worker import queue_repository as open_repository
|
|
|
|
async with open_repository(project_id) as repository:
|
|
yield repository
|
|
|
|
|
|
@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 app.infra.db.calculation_worker import calculation_supervisor_lease
|
|
|
|
# Dedicated session-level metadata lock: independent API workers never
|
|
# participate; duplicate supervisors exit before claiming any jobs.
|
|
async with calculation_supervisor_lease() as lease:
|
|
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 lease.ping()
|
|
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)
|
|
|
|
|
|
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]))
|