220 lines
7.7 KiB
Python
220 lines
7.7 KiB
Python
"""Per-group task queues and subprocess execution."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
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) -> tuple[Run, bool]:
|
|
run, created = self.store.enqueue(
|
|
target_kind="task",
|
|
target_path=task.id,
|
|
queue_group=task.group_id,
|
|
)
|
|
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)
|
|
environment = {
|
|
"PATH": "/usr/local/bin:/usr/bin:/bin",
|
|
"PYTHONUNBUFFERED": "1",
|
|
"SERVER_MAINTENANCE_CREDENTIALS": str(self.settings.credentials_root),
|
|
"SERVER_MAINTENANCE_STATE": str(task_state),
|
|
}
|
|
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
|