"""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]))