"""Per-group task queues and subprocess execution.""" from __future__ import annotations import os import json import signal import subprocess import threading from pathlib import Path from .config import Settings from .discovery import Group, Task, TaskCatalog from .store import Run, RunStore class RunManager: def __init__(self, settings: Settings, catalog: TaskCatalog, store: RunStore) -> None: self.settings = settings self.catalog = catalog self.store = store self._wake_events: dict[str, threading.Event] = {} self._workers: dict[str, threading.Thread] = {} self._processes: dict[int, subprocess.Popen[bytes]] = {} self._lock = threading.Lock() def run_task( self, task: Task, input_data: dict[str, str | int | list[str]] | None = None, ) -> tuple[Run, bool]: run, created = self.store.enqueue( target_kind="task", target_path=task.id, queue_group=task.group_id, input_data=input_data, ) if created: self._wake_worker(task.group_id) return run, created def run_group(self, group: Group) -> tuple[Run, bool]: run, created = self.store.enqueue( target_kind="group", target_path=group.id, queue_group=group.id, ) if created: self._wake_worker(group.id) return run, created def request_cancel(self, run_id: int) -> Run | None: run = self.store.request_cancel(run_id) if run is None: return None if run.parent_run_id is None and run.target_kind == "group": child = self.store.active_child(run.id) if child is not None: self.request_cancel(child.id) self._terminate_active_process(run.id) return self.store.get(run.id) def _wake_worker(self, group_id: str) -> None: with self._lock: event = self._wake_events.setdefault(group_id, threading.Event()) worker = self._workers.get(group_id) if worker is None or not worker.is_alive(): worker = threading.Thread( target=self._worker_loop, args=(group_id, event), name=f"task-group-{group_id}", daemon=True, ) self._workers[group_id] = worker worker.start() event.set() def _worker_loop(self, group_id: str, event: threading.Event) -> None: while True: event.clear() run = self.store.claim_next(group_id) if run is None: event.wait() continue self._execute_root_run(run) def _execute_root_run(self, run: Run) -> None: if run.target_kind == "task": task = self.catalog.task(run.target_path) if task is None: self.store.finish( run.id, status="failed", message="Task is no longer available.", ) return self._execute_task(run, task) return group = self.catalog.group(run.target_path) if group is None: self.store.finish( run.id, status="failed", message="Group is no longer available.", ) return self._execute_group(run, group) def _execute_group(self, parent_run: Run, group: Group) -> None: for task in group.tasks: if self.store.is_cancel_requested(parent_run.id): self.store.finish( parent_run.id, status="cancelled", message="Group cancellation requested.", ) return child = self.store.create_child(parent_run, task.id) completed_child = self._execute_task(child, task, parent_run_id=parent_run.id) if completed_child.status != "succeeded": parent_status = ( "cancelled" if completed_child.status == "cancelled" or self.store.is_cancel_requested(parent_run.id) else "failed" ) self.store.finish( parent_run.id, status=parent_status, message=f"Stopped after {task.id}: {completed_child.status}.", ) return self.store.finish(parent_run.id, status="succeeded") def _execute_task( self, run: Run, task: Task, *, parent_run_id: int | None = None, ) -> Run: if self.store.is_cancel_requested(run.id) or ( parent_run_id is not None and self.store.is_cancel_requested(parent_run_id) ): return self.store.finish(run.id, status="cancelled", message="Cancelled before execution.") task_state = self.settings.state_root / "tasks" / task.group_id task_state.mkdir(parents=True, exist_ok=True) log_path = self.settings.state_root / "runs" / f"{run.id}.log" log_path.parent.mkdir(parents=True, exist_ok=True) input_path = self.settings.state_root / "runs" / f"{run.id}.input.json" try: input_path.write_text(run.input_json + "\n", encoding="utf-8") input_path.chmod(0o600) except OSError as error: return self.store.finish( run.id, status="failed", message=f"Could not prepare task input: {error}", ) environment = { "PATH": "/usr/local/bin:/usr/bin:/bin", "PYTHONUNBUFFERED": "1", "SERVER_MAINTENANCE_CREDENTIALS": str(self.settings.credentials_root), "SERVER_MAINTENANCE_STATE": str(task_state), "SERVER_MAINTENANCE_INPUT": str(input_path), } if self.settings.timezone_name: environment["TZ"] = self.settings.timezone_name try: with log_path.open("wb") as log_file: process = subprocess.Popen( [self.settings.python_executable, str(task.path)], cwd=task.path.parent, stdin=subprocess.DEVNULL, stdout=log_file, stderr=subprocess.STDOUT, env=environment, start_new_session=True, ) with self._lock: self._processes[run.id] = process exit_code = process.wait() except OSError as error: return self.store.finish( run.id, status="failed", message=f"Could not start task: {error}", log_path=str(log_path), ) finally: with self._lock: self._processes.pop(run.id, None) if self.store.is_cancel_requested(run.id) or ( parent_run_id is not None and self.store.is_cancel_requested(parent_run_id) ): return self.store.finish( run.id, status="cancelled", exit_code=exit_code, log_path=str(log_path), ) return self.store.finish( run.id, status="succeeded" if exit_code == 0 else "failed", exit_code=exit_code, log_path=str(log_path), ) def _terminate_active_process(self, run_id: int) -> None: with self._lock: process = self._processes.get(run_id) if process is None or process.poll() is not None: return try: os.killpg(process.pid, signal.SIGTERM) except ProcessLookupError: return threading.Thread( target=self._kill_after_grace, args=(process,), name=f"task-cancel-{run_id}", daemon=True, ).start() def _kill_after_grace(self, process: subprocess.Popen[bytes]) -> None: try: process.wait(timeout=self.settings.cancel_grace_seconds) except subprocess.TimeoutExpired: try: os.killpg(process.pid, signal.SIGKILL) except ProcessLookupError: pass