385 lines
14 KiB
Python
385 lines
14 KiB
Python
"""Per-group task queues and subprocess execution."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import signal
|
|
import shutil
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from collections.abc import Mapping
|
|
from pathlib import Path
|
|
|
|
from .config import Settings
|
|
from .discovery import BackgroundTask, Group, Task, TaskCatalog
|
|
from .inputs import PendingUpload
|
|
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,
|
|
uploads: Mapping[str, PendingUpload] | 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:
|
|
try:
|
|
uploaded_inputs = self._save_uploads(run.id, uploads or {})
|
|
except (OSError, ValueError) as error:
|
|
self._delete_uploads(run.id)
|
|
return self.store.finish(
|
|
run.id,
|
|
status="failed",
|
|
message=f"Could not save uploaded file: {error}",
|
|
), True
|
|
if uploaded_inputs:
|
|
run = self.store.set_input(run.id, {**(input_data or {}), **uploaded_inputs})
|
|
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)
|
|
completed = self.store.get(run.id)
|
|
if completed is not None and completed.status == "cancelled":
|
|
self._delete_uploads(run.id)
|
|
return completed
|
|
|
|
def artifacts(self, run_id: int) -> list[str]:
|
|
root = self._artifact_root(run_id)
|
|
if not root.is_dir():
|
|
return []
|
|
return [
|
|
path.name
|
|
for path in sorted(root.iterdir(), key=lambda item: item.name)
|
|
if path.is_file() and not path.is_symlink()
|
|
]
|
|
|
|
def artifact_path(self, run_id: int, filename: str) -> Path | None:
|
|
path = self._artifact_root(run_id) / filename
|
|
if Path(filename).name != filename or path.is_symlink() or not path.is_file():
|
|
return None
|
|
return path
|
|
|
|
def delete_artifact(self, run_id: int, filename: str) -> None:
|
|
path = self.artifact_path(run_id, filename)
|
|
if path is None:
|
|
return
|
|
path.unlink(missing_ok=True)
|
|
try:
|
|
path.parent.rmdir()
|
|
except OSError:
|
|
pass
|
|
|
|
def cleanup_transient_files(self, *, artifact_max_age_seconds: int = 86_400) -> None:
|
|
"""Remove uploads from interrupted runs and expired undownloaded artifacts."""
|
|
|
|
runs_root = self.settings.state_root / "runs"
|
|
if not runs_root.is_dir():
|
|
return
|
|
cutoff = time.time() - artifact_max_age_seconds
|
|
for path in runs_root.iterdir():
|
|
if not path.is_dir():
|
|
continue
|
|
run_id_text, separator, kind = path.name.partition(".")
|
|
if not separator or not run_id_text.isdigit() or kind not in {"uploads", "artifacts"}:
|
|
continue
|
|
run = self.store.get(int(run_id_text))
|
|
if kind == "uploads" and (run is None or run.status not in {"queued", "running"}):
|
|
shutil.rmtree(path, ignore_errors=True)
|
|
elif kind == "artifacts" and path.stat().st_mtime < cutoff:
|
|
shutil.rmtree(path, ignore_errors=True)
|
|
|
|
def run_background_task(self, task: BackgroundTask) -> None:
|
|
"""Run a declared internal task without creating high-volume run history."""
|
|
|
|
log_path = self.settings.state_root / "background" / f"{task.id}.log"
|
|
log_path.parent.mkdir(parents=True, exist_ok=True)
|
|
environment = self._environment(task.group_id)
|
|
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,
|
|
)
|
|
exit_code = process.wait()
|
|
if exit_code:
|
|
log_file.write(f"Background task exited with status {exit_code}.\n".encode())
|
|
except OSError as error:
|
|
log_path.write_text(f"Could not start background task: {error}\n", encoding="utf-8")
|
|
|
|
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.")
|
|
|
|
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"
|
|
artifact_root = self._artifact_root(run.id)
|
|
artifact_root.mkdir(parents=True, exist_ok=True)
|
|
try:
|
|
input_path.write_text(run.input_json + "\n", encoding="utf-8")
|
|
input_path.chmod(0o600)
|
|
except OSError as error:
|
|
self._delete_uploads(run.id)
|
|
shutil.rmtree(artifact_root, ignore_errors=True)
|
|
return self.store.finish(
|
|
run.id,
|
|
status="failed",
|
|
message=f"Could not prepare task input: {error}",
|
|
)
|
|
environment = self._environment(task.group_id)
|
|
environment["SERVER_MAINTENANCE_INPUT"] = str(input_path)
|
|
environment["SERVER_MAINTENANCE_ARTIFACTS"] = str(artifact_root)
|
|
|
|
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:
|
|
self._delete_uploads(run.id)
|
|
shutil.rmtree(artifact_root, ignore_errors=True)
|
|
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)
|
|
self._delete_uploads(run.id)
|
|
|
|
if self.store.is_cancel_requested(run.id) or (
|
|
parent_run_id is not None and self.store.is_cancel_requested(parent_run_id)
|
|
):
|
|
shutil.rmtree(artifact_root, ignore_errors=True)
|
|
return self.store.finish(
|
|
run.id,
|
|
status="cancelled",
|
|
exit_code=exit_code,
|
|
log_path=str(log_path),
|
|
)
|
|
completed = self.store.finish(
|
|
run.id,
|
|
status="succeeded" if exit_code == 0 else "failed",
|
|
exit_code=exit_code,
|
|
log_path=str(log_path),
|
|
)
|
|
if completed.status != "succeeded":
|
|
shutil.rmtree(artifact_root, ignore_errors=True)
|
|
else:
|
|
try:
|
|
artifact_root.rmdir()
|
|
except OSError:
|
|
pass
|
|
self.cleanup_transient_files()
|
|
return completed
|
|
|
|
def _environment(self, group_id: str) -> dict[str, str]:
|
|
task_state = self.settings.state_root / "tasks" / group_id
|
|
task_state.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
|
|
return environment
|
|
|
|
def _upload_root(self, run_id: int) -> Path:
|
|
return self.settings.state_root / "runs" / f"{run_id}.uploads"
|
|
|
|
def _artifact_root(self, run_id: int) -> Path:
|
|
return self.settings.state_root / "runs" / f"{run_id}.artifacts"
|
|
|
|
def _save_uploads(
|
|
self,
|
|
run_id: int,
|
|
uploads: Mapping[str, PendingUpload],
|
|
) -> dict[str, str]:
|
|
if not uploads:
|
|
return {}
|
|
root = self._upload_root(run_id)
|
|
root.mkdir(parents=True, exist_ok=False)
|
|
saved = {}
|
|
try:
|
|
for field_name, upload in uploads.items():
|
|
suffix = upload.extension or ".upload"
|
|
destination = root / f"{field_name}{suffix}"
|
|
temporary = destination.with_suffix(destination.suffix + ".part")
|
|
total = 0
|
|
with temporary.open("xb") as output:
|
|
while chunk := upload.stream.read(64 * 1024):
|
|
total += len(chunk)
|
|
if upload.maximum_bytes is not None and total > upload.maximum_bytes:
|
|
raise ValueError(f"{field_name} exceeds its size limit")
|
|
output.write(chunk)
|
|
temporary.replace(destination)
|
|
destination.chmod(0o600)
|
|
saved[field_name] = str(destination)
|
|
except Exception:
|
|
shutil.rmtree(root, ignore_errors=True)
|
|
raise
|
|
return saved
|
|
|
|
def _delete_uploads(self, run_id: int) -> None:
|
|
shutil.rmtree(self._upload_root(run_id), ignore_errors=True)
|
|
|
|
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
|