Update Python tools
This commit is contained in:
ajp_anton
+124
-3
@@ -2,15 +2,19 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
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
|
||||
|
||||
|
||||
@@ -28,6 +32,7 @@ class RunManager:
|
||||
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",
|
||||
@@ -36,6 +41,17 @@ class RunManager:
|
||||
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
|
||||
|
||||
@@ -58,7 +74,55 @@ class RunManager:
|
||||
if child is not None:
|
||||
self.request_cancel(child.id)
|
||||
self._terminate_active_process(run.id)
|
||||
return self.store.get(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."""
|
||||
@@ -171,10 +235,14 @@ class RunManager:
|
||||
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",
|
||||
@@ -182,6 +250,7 @@ class RunManager:
|
||||
)
|
||||
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:
|
||||
@@ -198,6 +267,8 @@ class RunManager:
|
||||
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",
|
||||
@@ -207,22 +278,33 @@ class RunManager:
|
||||
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),
|
||||
)
|
||||
return self.store.finish(
|
||||
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
|
||||
@@ -237,6 +319,45 @@ class RunManager:
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user