122 lines
3.9 KiB
Python
122 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
import threading
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from . import db
|
|
from .commands import CommandSpec
|
|
from .config import settings
|
|
from .modules import build_module_task
|
|
from .schemas import JobCreate
|
|
|
|
_running: dict[str, subprocess.Popen] = {}
|
|
_lock = threading.Lock()
|
|
|
|
|
|
def default_conda_env_for_job(job_type: str) -> str:
|
|
if job_type.startswith("mmseg."):
|
|
return settings.mmseg_conda_env
|
|
return settings.task_conda_env
|
|
|
|
|
|
def _build_task(request: JobCreate) -> CommandSpec:
|
|
conda_env = request.conda_env or default_conda_env_for_job(request.type)
|
|
spec = build_module_task(request.type, request.params, conda_env)
|
|
if spec is None:
|
|
raise ValueError(f"unsupported job type: {request.type}")
|
|
env = dict(spec.env)
|
|
if request.gpus:
|
|
env["CUDA_VISIBLE_DEVICES"] = ",".join(str(gpu) for gpu in request.gpus)
|
|
return CommandSpec(
|
|
command=spec.command,
|
|
cwd=spec.cwd,
|
|
description=spec.description,
|
|
env=env,
|
|
stdin_text=spec.stdin_text,
|
|
)
|
|
|
|
|
|
def create_job(request: JobCreate) -> dict:
|
|
spec = _build_task(request)
|
|
job_id = uuid.uuid4().hex
|
|
settings.log_dir.mkdir(parents=True, exist_ok=True)
|
|
log_path = settings.log_dir / f"{job_id}.log"
|
|
job = {
|
|
"id": job_id,
|
|
"type": request.type,
|
|
"status": "queued",
|
|
"params": request.params,
|
|
"command": spec.command,
|
|
"cwd": str(spec.cwd),
|
|
"description": spec.description,
|
|
"pid": None,
|
|
"exit_code": None,
|
|
"created_at": db.utc_now(),
|
|
"started_at": None,
|
|
"finished_at": None,
|
|
"log_path": str(log_path),
|
|
"error": None,
|
|
}
|
|
db.insert_job(job)
|
|
thread = threading.Thread(target=_run_job, args=(job_id, spec, log_path), daemon=True)
|
|
thread.start()
|
|
return db.get_job(job_id)
|
|
|
|
|
|
def _run_job(job_id: str, spec: CommandSpec, log_path: Path) -> None:
|
|
env = os.environ.copy()
|
|
env.update(spec.env)
|
|
db.update_job(job_id, status="running", started_at=db.utc_now())
|
|
try:
|
|
with log_path.open("ab") as log_file:
|
|
log_file.write(("COMMAND: " + " ".join(spec.command) + "\n").encode("utf-8"))
|
|
log_file.flush()
|
|
process = subprocess.Popen(
|
|
spec.command,
|
|
cwd=str(spec.cwd),
|
|
env=env,
|
|
stdin=subprocess.PIPE if spec.stdin_text is not None else None,
|
|
stdout=log_file,
|
|
stderr=subprocess.STDOUT,
|
|
start_new_session=True,
|
|
)
|
|
with _lock:
|
|
_running[job_id] = process
|
|
db.update_job(job_id, pid=process.pid)
|
|
if spec.stdin_text is not None and process.stdin is not None:
|
|
process.stdin.write(spec.stdin_text.encode("utf-8"))
|
|
process.stdin.close()
|
|
exit_code = process.wait()
|
|
with _lock:
|
|
_running.pop(job_id, None)
|
|
current = db.get_job(job_id)
|
|
if current and current["status"] == "cancelled":
|
|
db.update_job(job_id, exit_code=exit_code, finished_at=db.utc_now())
|
|
elif exit_code == 0:
|
|
db.update_job(job_id, status="success", exit_code=exit_code, finished_at=db.utc_now())
|
|
else:
|
|
db.update_job(job_id, status="failed", exit_code=exit_code, finished_at=db.utc_now())
|
|
except Exception as exc:
|
|
with _lock:
|
|
_running.pop(job_id, None)
|
|
db.update_job(job_id, status="failed", error=str(exc), finished_at=db.utc_now())
|
|
|
|
|
|
def cancel_job(job_id: str) -> dict | None:
|
|
job = db.get_job(job_id)
|
|
if not job:
|
|
return None
|
|
db.update_job(job_id, status="cancelled", finished_at=db.utc_now())
|
|
with _lock:
|
|
process = _running.get(job_id)
|
|
if process and process.poll() is None:
|
|
try:
|
|
os.killpg(process.pid, signal.SIGTERM)
|
|
except ProcessLookupError:
|
|
pass
|
|
return db.get_job(job_id)
|