Integrate task runner into python tools
This commit is contained in:
ajp_anton
@@ -0,0 +1,219 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user