565 lines
20 KiB
Python
565 lines
20 KiB
Python
|
"""Lightweight resource gate and task ledger for local stage execution.
|
|||
|
|
|
|||
|
|
The project runs on one workstation, so a small file-backed ledger is enough to
|
|||
|
|
coordinate FastAPI and Celery worker processes without adding another service.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import json
|
|||
|
|
import os
|
|||
|
|
import time
|
|||
|
|
from contextlib import contextmanager
|
|||
|
|
from dataclasses import dataclass
|
|||
|
|
from datetime import datetime
|
|||
|
|
from pathlib import Path
|
|||
|
|
from typing import Any
|
|||
|
|
|
|||
|
|
import fcntl
|
|||
|
|
|
|||
|
|
|
|||
|
|
TERMINAL_STATES = {"success", "completed", "failed", "failure", "cancelled", "revoked"}
|
|||
|
|
ACTIVE_STATES = {"waiting_resource", "running"}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _project_root() -> Path:
|
|||
|
|
return Path(os.environ.get("PROJECT_ROOT", Path(__file__).resolve().parents[2])).resolve()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _state_dir() -> Path:
|
|||
|
|
path = Path(os.environ.get("BBAV5_RESOURCE_STATE_DIR", _project_root() / "workspace" / "resource_governor"))
|
|||
|
|
path.mkdir(parents=True, exist_ok=True)
|
|||
|
|
return path
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _state_path() -> Path:
|
|||
|
|
return _state_dir() / "state.json"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _config_path() -> Path:
|
|||
|
|
return _state_dir() / "config.json"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _lock_path() -> Path:
|
|||
|
|
return _state_dir() / "state.lock"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _now() -> str:
|
|||
|
|
return datetime.now().isoformat()
|
|||
|
|
|
|||
|
|
|
|||
|
|
DEFAULT_RESOURCE_CONFIG = {
|
|||
|
|
"total_cpu": 20,
|
|||
|
|
"max_running_tasks": 11,
|
|||
|
|
"gene_family_project_max_cpu": 20,
|
|||
|
|
"rnaseq_project_max_cpu": 20,
|
|||
|
|
"rnaseq_qc_sample_parallel_jobs": 2,
|
|||
|
|
"rnaseq_align_sample_parallel_jobs": 2,
|
|||
|
|
"stage_cpu_overrides": {},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _sanitize_stage_cpu_overrides(value: Any) -> dict[str, dict[str, int]]:
|
|||
|
|
if not isinstance(value, dict):
|
|||
|
|
return {}
|
|||
|
|
sanitized: dict[str, dict[str, int]] = {}
|
|||
|
|
for stage, item in value.items():
|
|||
|
|
if not isinstance(stage, str) or not isinstance(item, dict):
|
|||
|
|
continue
|
|||
|
|
try:
|
|||
|
|
max_cpu = int(item.get("max_cpu", item.get("default_cpu", 1)))
|
|||
|
|
default_cpu = int(item.get("default_cpu", max_cpu))
|
|||
|
|
except (TypeError, ValueError):
|
|||
|
|
continue
|
|||
|
|
max_cpu = max(4, max_cpu)
|
|||
|
|
default_cpu = max(4, default_cpu)
|
|||
|
|
sanitized[stage] = {"default_cpu": default_cpu, "max_cpu": max_cpu}
|
|||
|
|
return sanitized
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _read_config_unlocked() -> dict[str, Any]:
|
|||
|
|
config = dict(DEFAULT_RESOURCE_CONFIG)
|
|||
|
|
config["total_cpu"] = int(os.environ.get("BBAV5_RESOURCE_TOTAL_CPU", config["total_cpu"]))
|
|||
|
|
config["max_running_tasks"] = int(os.environ.get("BBAV5_RESOURCE_MAX_RUNNING_TASKS", config["max_running_tasks"]))
|
|||
|
|
config["rnaseq_project_max_cpu"] = int(os.environ.get("BBAV5_RNASEQ_PROJECT_MAX_CPU", config["rnaseq_project_max_cpu"]))
|
|||
|
|
path = _config_path()
|
|||
|
|
if path.exists():
|
|||
|
|
try:
|
|||
|
|
saved = json.loads(path.read_text(encoding="utf-8"))
|
|||
|
|
if isinstance(saved, dict):
|
|||
|
|
config.update({k: v for k, v in saved.items() if k in config})
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
for key, default in DEFAULT_RESOURCE_CONFIG.items():
|
|||
|
|
if key == "stage_cpu_overrides":
|
|||
|
|
config[key] = _sanitize_stage_cpu_overrides(config.get(key, default))
|
|||
|
|
continue
|
|||
|
|
if isinstance(default, str):
|
|||
|
|
config[key] = str(config.get(key, default) or "").strip()
|
|||
|
|
continue
|
|||
|
|
try:
|
|||
|
|
config[key] = max(1, int(config.get(key, default)))
|
|||
|
|
except (TypeError, ValueError):
|
|||
|
|
config[key] = default
|
|||
|
|
total = max(1, int(config.get("total_cpu", DEFAULT_RESOURCE_CONFIG["total_cpu"])))
|
|||
|
|
config["total_cpu"] = total
|
|||
|
|
config["gene_family_project_max_cpu"] = max(1, min(total, int(config.get("gene_family_project_max_cpu", total))))
|
|||
|
|
config["rnaseq_project_max_cpu"] = max(1, min(total, int(config.get("rnaseq_project_max_cpu", total))))
|
|||
|
|
return config
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_resource_config() -> dict[str, Any]:
|
|||
|
|
return _read_config_unlocked()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def update_resource_config(updates: dict[str, Any]) -> dict[str, Any]:
|
|||
|
|
allowed = set(DEFAULT_RESOURCE_CONFIG)
|
|||
|
|
current = _read_config_unlocked()
|
|||
|
|
for key, value in (updates or {}).items():
|
|||
|
|
if key not in allowed:
|
|||
|
|
continue
|
|||
|
|
if value is None or value == "":
|
|||
|
|
continue
|
|||
|
|
if key == "stage_cpu_overrides":
|
|||
|
|
current[key] = _sanitize_stage_cpu_overrides(value)
|
|||
|
|
elif isinstance(DEFAULT_RESOURCE_CONFIG.get(key), str):
|
|||
|
|
current[key] = str(value or "").strip()
|
|||
|
|
else:
|
|||
|
|
current[key] = int(value)
|
|||
|
|
total = max(1, int(current.get("total_cpu", DEFAULT_RESOURCE_CONFIG["total_cpu"])))
|
|||
|
|
current["total_cpu"] = total
|
|||
|
|
current["max_running_tasks"] = max(1, int(current.get("max_running_tasks", DEFAULT_RESOURCE_CONFIG["max_running_tasks"])))
|
|||
|
|
current["gene_family_project_max_cpu"] = max(1, min(total, int(current.get("gene_family_project_max_cpu", total))))
|
|||
|
|
current["rnaseq_project_max_cpu"] = max(1, min(total, int(current.get("rnaseq_project_max_cpu", total))))
|
|||
|
|
current["rnaseq_qc_sample_parallel_jobs"] = max(1, int(current.get("rnaseq_qc_sample_parallel_jobs", 2)))
|
|||
|
|
current["rnaseq_align_sample_parallel_jobs"] = max(1, int(current.get("rnaseq_align_sample_parallel_jobs", 2)))
|
|||
|
|
current["stage_cpu_overrides"] = _sanitize_stage_cpu_overrides(current.get("stage_cpu_overrides", {}))
|
|||
|
|
path = _config_path()
|
|||
|
|
path.write_text(json.dumps(current, ensure_ascii=False, indent=2), encoding="utf-8")
|
|||
|
|
return current
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _total_cpu() -> int:
|
|||
|
|
return max(1, int(_read_config_unlocked().get("total_cpu", 20)))
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _max_running_tasks() -> int:
|
|||
|
|
return max(1, int(_read_config_unlocked().get("max_running_tasks", 11)))
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _pid_alive(pid: Any) -> bool:
|
|||
|
|
try:
|
|||
|
|
value = int(pid)
|
|||
|
|
except (TypeError, ValueError):
|
|||
|
|
return False
|
|||
|
|
return value > 0 and Path(f"/proc/{value}").exists()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _total_memory_gb() -> float:
|
|||
|
|
try:
|
|||
|
|
with open("/proc/meminfo", "r", encoding="utf-8") as handle:
|
|||
|
|
for line in handle:
|
|||
|
|
if line.startswith("MemTotal:"):
|
|||
|
|
kb = int(line.split()[1])
|
|||
|
|
return round(kb / 1024 / 1024, 2)
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
return 0.0
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _gpu_summary() -> list[dict[str, Any]]:
|
|||
|
|
try:
|
|||
|
|
import subprocess
|
|||
|
|
result = subprocess.run(
|
|||
|
|
[
|
|||
|
|
"nvidia-smi",
|
|||
|
|
"--query-gpu=name,memory.total",
|
|||
|
|
"--format=csv,noheader,nounits",
|
|||
|
|
],
|
|||
|
|
capture_output=True,
|
|||
|
|
text=True,
|
|||
|
|
check=False,
|
|||
|
|
timeout=5,
|
|||
|
|
)
|
|||
|
|
if result.returncode != 0:
|
|||
|
|
return []
|
|||
|
|
gpus = []
|
|||
|
|
for idx, line in enumerate(result.stdout.splitlines()):
|
|||
|
|
if not line.strip():
|
|||
|
|
continue
|
|||
|
|
parts = [part.strip() for part in line.split(",")]
|
|||
|
|
gpus.append({
|
|||
|
|
"index": idx,
|
|||
|
|
"name": parts[0] if parts else "unknown",
|
|||
|
|
"memory_mb": int(float(parts[1])) if len(parts) > 1 and parts[1] else 0,
|
|||
|
|
})
|
|||
|
|
return gpus
|
|||
|
|
except Exception:
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
|
|||
|
|
def detect_hardware() -> dict[str, Any]:
|
|||
|
|
"""Return local hardware facts used for RNA-seq resource planning."""
|
|||
|
|
total_cpu = os.cpu_count() or _total_cpu()
|
|||
|
|
configured_cpu = _total_cpu()
|
|||
|
|
memory_gb = _total_memory_gb()
|
|||
|
|
return {
|
|||
|
|
"cpu_count": int(total_cpu),
|
|||
|
|
"configured_total_cpu": int(configured_cpu),
|
|||
|
|
"memory_gb": memory_gb,
|
|||
|
|
"gpus": _gpu_summary(),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def rnaseq_project_cpu_limit(requested_cpu: int | None = None) -> int:
|
|||
|
|
"""Return the CPU budget for one RNA-seq project within the global governor."""
|
|||
|
|
config = _read_config_unlocked()
|
|||
|
|
configured_total = int(config.get("total_cpu", _total_cpu()))
|
|||
|
|
project_limit = int(config.get("rnaseq_project_max_cpu", configured_total))
|
|||
|
|
if requested_cpu:
|
|||
|
|
try:
|
|||
|
|
project_limit = min(project_limit, int(requested_cpu))
|
|||
|
|
except (TypeError, ValueError):
|
|||
|
|
pass
|
|||
|
|
return max(1, min(project_limit, configured_total))
|
|||
|
|
|
|||
|
|
|
|||
|
|
def recommend_rnaseq_parallelism(
|
|||
|
|
total_cpu: int | None = None,
|
|||
|
|
sample_parallel_jobs: int | None = None,
|
|||
|
|
threads_per_sample: int | None = None,
|
|||
|
|
sample_count: int | None = None,
|
|||
|
|
stage: str = "rnaseq_align",
|
|||
|
|
) -> dict[str, int]:
|
|||
|
|
"""Compute per-project RNA-seq sample/thread settings."""
|
|||
|
|
cpu_total = max(1, int(total_cpu or rnaseq_project_cpu_limit()))
|
|||
|
|
if sample_parallel_jobs is None or int(sample_parallel_jobs or 0) < 1:
|
|||
|
|
jobs = int(_read_config_unlocked().get("rnaseq_qc_sample_parallel_jobs" if stage == "rnaseq_qc" else "rnaseq_align_sample_parallel_jobs", 2)) if stage in {"rnaseq_qc", "rnaseq_align"} else 1
|
|||
|
|
else:
|
|||
|
|
jobs = max(1, int(sample_parallel_jobs or 1))
|
|||
|
|
if sample_count:
|
|||
|
|
jobs = min(jobs, max(1, int(sample_count)))
|
|||
|
|
if threads_per_sample is None or int(threads_per_sample) < 1:
|
|||
|
|
threads = max(1, cpu_total // jobs)
|
|||
|
|
else:
|
|||
|
|
threads = max(1, int(threads_per_sample))
|
|||
|
|
if jobs * threads > cpu_total:
|
|||
|
|
threads = max(1, cpu_total // jobs)
|
|||
|
|
effective_cpu = max(1, jobs * threads)
|
|||
|
|
return {
|
|||
|
|
"total_cpu": effective_cpu,
|
|||
|
|
"project_cpu_limit": cpu_total,
|
|||
|
|
"sample_parallel_jobs": jobs,
|
|||
|
|
"threads_per_sample": threads,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _empty_state() -> dict[str, Any]:
|
|||
|
|
return {
|
|||
|
|
"version": 1,
|
|||
|
|
"updated_at": _now(),
|
|||
|
|
"limits": {
|
|||
|
|
"total_cpu": _total_cpu(),
|
|||
|
|
"max_running_tasks": _max_running_tasks(),
|
|||
|
|
"hardware": detect_hardware(),
|
|||
|
|
"config": get_resource_config(),
|
|||
|
|
},
|
|||
|
|
"tasks": {},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
@contextmanager
|
|||
|
|
def _locked_state():
|
|||
|
|
lock_file = _lock_path()
|
|||
|
|
state_file = _state_path()
|
|||
|
|
with lock_file.open("a+", encoding="utf-8") as lock:
|
|||
|
|
fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
|
|||
|
|
try:
|
|||
|
|
if state_file.exists():
|
|||
|
|
try:
|
|||
|
|
state = json.loads(state_file.read_text(encoding="utf-8"))
|
|||
|
|
except Exception:
|
|||
|
|
state = _empty_state()
|
|||
|
|
else:
|
|||
|
|
state = _empty_state()
|
|||
|
|
if not isinstance(state, dict):
|
|||
|
|
state = _empty_state()
|
|||
|
|
state.setdefault("tasks", {})
|
|||
|
|
state["limits"] = {
|
|||
|
|
"total_cpu": _total_cpu(),
|
|||
|
|
"max_running_tasks": _max_running_tasks(),
|
|||
|
|
"hardware": detect_hardware(),
|
|||
|
|
"config": get_resource_config(),
|
|||
|
|
}
|
|||
|
|
_cleanup_stale_locked(state)
|
|||
|
|
yield state
|
|||
|
|
state["updated_at"] = _now()
|
|||
|
|
tmp = state_file.with_suffix(".json.tmp")
|
|||
|
|
tmp.write_text(json.dumps(state, ensure_ascii=False, indent=2), encoding="utf-8")
|
|||
|
|
tmp.replace(state_file)
|
|||
|
|
finally:
|
|||
|
|
fcntl.flock(lock.fileno(), fcntl.LOCK_UN)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _cleanup_stale_locked(state: dict[str, Any]) -> None:
|
|||
|
|
tasks = state.setdefault("tasks", {})
|
|||
|
|
for task_id, task in list(tasks.items()):
|
|||
|
|
status = str(task.get("status") or "").lower()
|
|||
|
|
if status in TERMINAL_STATES:
|
|||
|
|
continue
|
|||
|
|
if status == "running" and task.get("pid") and not _pid_alive(task.get("pid")):
|
|||
|
|
task["status"] = "failed"
|
|||
|
|
task["ended_at"] = task.get("ended_at") or _now()
|
|||
|
|
task["message"] = "worker process disappeared; resource lease released"
|
|||
|
|
task["cpu"] = 0
|
|||
|
|
elif status == "waiting_resource" and task.get("pid") and not _pid_alive(task.get("pid")):
|
|||
|
|
task["status"] = "cancelled"
|
|||
|
|
task["ended_at"] = task.get("ended_at") or _now()
|
|||
|
|
task["message"] = "waiting worker process disappeared"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _usage_locked(state: dict[str, Any]) -> dict[str, Any]:
|
|||
|
|
used_cpu = 0
|
|||
|
|
running = 0
|
|||
|
|
waiting = 0
|
|||
|
|
by_class: dict[str, int] = {}
|
|||
|
|
for task in state.get("tasks", {}).values():
|
|||
|
|
status = str(task.get("status") or "").lower()
|
|||
|
|
if status == "running":
|
|||
|
|
running += 1
|
|||
|
|
cpu = int(task.get("cpu") or 0)
|
|||
|
|
used_cpu += cpu
|
|||
|
|
cls = str(task.get("resource_class") or "unknown")
|
|||
|
|
by_class[cls] = by_class.get(cls, 0) + cpu
|
|||
|
|
elif status == "waiting_resource":
|
|||
|
|
waiting += 1
|
|||
|
|
limits = state.get("limits", {})
|
|||
|
|
total_cpu = int(limits.get("total_cpu") or _total_cpu())
|
|||
|
|
return {
|
|||
|
|
"used_cpu": used_cpu,
|
|||
|
|
"available_cpu": max(0, total_cpu - used_cpu),
|
|||
|
|
"running_tasks": running,
|
|||
|
|
"waiting_tasks": waiting,
|
|||
|
|
"by_resource_class_cpu": by_class,
|
|||
|
|
"limits": limits,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def record_task(
|
|||
|
|
*,
|
|||
|
|
task_id: str,
|
|||
|
|
stage: str,
|
|||
|
|
job_id: str,
|
|||
|
|
prefix: str = "",
|
|||
|
|
status: str = "queued",
|
|||
|
|
cpu: int = 0,
|
|||
|
|
queue: str = "",
|
|||
|
|
resource_class: str = "",
|
|||
|
|
concurrency_group: str = "",
|
|||
|
|
celery_task_id: str = "",
|
|||
|
|
message: str = "",
|
|||
|
|
) -> dict[str, Any]:
|
|||
|
|
with _locked_state() as state:
|
|||
|
|
task = state["tasks"].setdefault(task_id, {})
|
|||
|
|
task.update(
|
|||
|
|
{
|
|||
|
|
"task_id": task_id,
|
|||
|
|
"stage": stage,
|
|||
|
|
"job_id": job_id,
|
|||
|
|
"prefix": prefix,
|
|||
|
|
"status": status,
|
|||
|
|
"cpu": int(cpu or 0) if status == "running" else 0,
|
|||
|
|
"requested_cpu": int(cpu or 0),
|
|||
|
|
"queue": queue,
|
|||
|
|
"resource_class": resource_class,
|
|||
|
|
"concurrency_group": concurrency_group,
|
|||
|
|
"celery_task_id": celery_task_id or task.get("celery_task_id", ""),
|
|||
|
|
"pid": os.getpid() if status in ACTIVE_STATES else task.get("pid"),
|
|||
|
|
"message": message,
|
|||
|
|
"updated_at": _now(),
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
task.setdefault("created_at", _now())
|
|||
|
|
if status == "running":
|
|||
|
|
task["started_at"] = task.get("started_at") or _now()
|
|||
|
|
if status in TERMINAL_STATES:
|
|||
|
|
task["ended_at"] = task.get("ended_at") or _now()
|
|||
|
|
task["cpu"] = 0
|
|||
|
|
return {"task": dict(task), "usage": _usage_locked(state)}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def try_acquire(
|
|||
|
|
*,
|
|||
|
|
task_id: str,
|
|||
|
|
stage: str,
|
|||
|
|
job_id: str,
|
|||
|
|
prefix: str,
|
|||
|
|
cpu: int,
|
|||
|
|
queue: str = "",
|
|||
|
|
resource_class: str = "",
|
|||
|
|
concurrency_group: str = "",
|
|||
|
|
celery_task_id: str = "",
|
|||
|
|
) -> dict[str, Any]:
|
|||
|
|
cpu = max(1, int(cpu or 1))
|
|||
|
|
with _locked_state() as state:
|
|||
|
|
tasks = state.setdefault("tasks", {})
|
|||
|
|
existing = tasks.get(task_id, {})
|
|||
|
|
if str(existing.get("status") or "").lower() == "running":
|
|||
|
|
return {"granted": True, "task": dict(existing), "usage": _usage_locked(state)}
|
|||
|
|
|
|||
|
|
usage = _usage_locked(state)
|
|||
|
|
limits = state.get("limits", {})
|
|||
|
|
total_cpu = int(limits.get("total_cpu") or _total_cpu())
|
|||
|
|
max_tasks = int(limits.get("max_running_tasks") or _max_running_tasks())
|
|||
|
|
config = get_resource_config()
|
|||
|
|
project_limit_key = "rnaseq_project_max_cpu" if resource_class == "rnaseq" else "gene_family_project_max_cpu"
|
|||
|
|
project_limit = max(1, int(config.get(project_limit_key, total_cpu)))
|
|||
|
|
project_used_cpu = 0
|
|||
|
|
group_running = False
|
|||
|
|
for existing_task in tasks.values():
|
|||
|
|
if str(existing_task.get("status") or "").lower() != "running":
|
|||
|
|
continue
|
|||
|
|
if (
|
|||
|
|
concurrency_group
|
|||
|
|
and str(existing_task.get("concurrency_group") or "") == concurrency_group
|
|||
|
|
and str(existing_task.get("task_id") or "") != task_id
|
|||
|
|
):
|
|||
|
|
group_running = True
|
|||
|
|
if str(existing_task.get("job_id") or "") != str(job_id):
|
|||
|
|
continue
|
|||
|
|
existing_class = str(existing_task.get("resource_class") or "")
|
|||
|
|
existing_key = "rnaseq_project_max_cpu" if existing_class == "rnaseq" else "gene_family_project_max_cpu"
|
|||
|
|
if existing_key == project_limit_key:
|
|||
|
|
project_used_cpu += int(existing_task.get("cpu") or 0)
|
|||
|
|
project_has_cpu = True
|
|||
|
|
# CPU token caps are disabled; max_running_tasks still prevents unbounded task fan-out.
|
|||
|
|
if usage["running_tasks"] < max_tasks and project_has_cpu and not group_running:
|
|||
|
|
task = tasks.setdefault(task_id, {})
|
|||
|
|
task.update(
|
|||
|
|
{
|
|||
|
|
"task_id": task_id,
|
|||
|
|
"stage": stage,
|
|||
|
|
"job_id": job_id,
|
|||
|
|
"prefix": prefix,
|
|||
|
|
"status": "running",
|
|||
|
|
"cpu": cpu,
|
|||
|
|
"requested_cpu": cpu,
|
|||
|
|
"queue": queue,
|
|||
|
|
"resource_class": resource_class,
|
|||
|
|
"concurrency_group": concurrency_group,
|
|||
|
|
"celery_task_id": celery_task_id or task.get("celery_task_id", ""),
|
|||
|
|
"pid": os.getpid(),
|
|||
|
|
"started_at": task.get("started_at") or _now(),
|
|||
|
|
"updated_at": _now(),
|
|||
|
|
"message": "resource lease granted",
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
return {"granted": True, "task": dict(task), "usage": _usage_locked(state)}
|
|||
|
|
|
|||
|
|
task = tasks.setdefault(task_id, {})
|
|||
|
|
task.update(
|
|||
|
|
{
|
|||
|
|
"task_id": task_id,
|
|||
|
|
"stage": stage,
|
|||
|
|
"job_id": job_id,
|
|||
|
|
"prefix": prefix,
|
|||
|
|
"status": "waiting_resource",
|
|||
|
|
"cpu": 0,
|
|||
|
|
"requested_cpu": cpu,
|
|||
|
|
"queue": queue,
|
|||
|
|
"resource_class": resource_class,
|
|||
|
|
"concurrency_group": concurrency_group,
|
|||
|
|
"celery_task_id": celery_task_id or task.get("celery_task_id", ""),
|
|||
|
|
"pid": os.getpid(),
|
|||
|
|
"updated_at": _now(),
|
|||
|
|
"message": f"waiting for {cpu} CPU token(s)",
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
usage = _usage_locked(state)
|
|||
|
|
return {
|
|||
|
|
"granted": False,
|
|||
|
|
"task": dict(task),
|
|||
|
|
"usage": usage,
|
|||
|
|
"reason": (
|
|||
|
|
f"等待互斥资源组 {concurrency_group}（容量 1）"
|
|||
|
|
if group_running else (
|
|||
|
|
f"资源不足：需要 CPU {cpu}，当前可用 CPU {usage['available_cpu']}，"
|
|||
|
|
f"运行任务 {usage['running_tasks']}/{max_tasks}，"
|
|||
|
|
f"项目已用 CPU {project_used_cpu}/{project_limit}"
|
|||
|
|
)
|
|||
|
|
),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def release(task_id: str, status: str = "completed", message: str = "") -> dict[str, Any]:
|
|||
|
|
with _locked_state() as state:
|
|||
|
|
task = state.setdefault("tasks", {}).setdefault(task_id, {"task_id": task_id})
|
|||
|
|
task["status"] = status
|
|||
|
|
task["cpu"] = 0
|
|||
|
|
task["ended_at"] = task.get("ended_at") or _now()
|
|||
|
|
task["updated_at"] = _now()
|
|||
|
|
if message:
|
|||
|
|
task["message"] = message
|
|||
|
|
return {"task": dict(task), "usage": _usage_locked(state)}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def snapshot(include_terminal: bool = False) -> dict[str, Any]:
|
|||
|
|
with _locked_state() as state:
|
|||
|
|
tasks = []
|
|||
|
|
for task in state.get("tasks", {}).values():
|
|||
|
|
status = str(task.get("status") or "").lower()
|
|||
|
|
if include_terminal or status not in TERMINAL_STATES:
|
|||
|
|
tasks.append(dict(task))
|
|||
|
|
tasks.sort(key=lambda item: item.get("updated_at", ""), reverse=True)
|
|||
|
|
return {
|
|||
|
|
"state_path": str(_state_path()),
|
|||
|
|
"updated_at": state.get("updated_at"),
|
|||
|
|
"limits": state.get("limits", {}),
|
|||
|
|
"usage": _usage_locked(state),
|
|||
|
|
"tasks": tasks,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_task(task_id: str) -> dict[str, Any] | None:
|
|||
|
|
with _locked_state() as state:
|
|||
|
|
task = state.get("tasks", {}).get(task_id)
|
|||
|
|
if not task:
|
|||
|
|
return None
|
|||
|
|
return {"task": dict(task), "usage": _usage_locked(state)}
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class ResourceLease:
|
|||
|
|
task_id: str
|
|||
|
|
stage: str
|
|||
|
|
job_id: str
|
|||
|
|
prefix: str
|
|||
|
|
cpu: int
|
|||
|
|
queue: str = ""
|
|||
|
|
resource_class: str = ""
|
|||
|
|
concurrency_group: str = ""
|
|||
|
|
celery_task_id: str = ""
|
|||
|
|
poll_seconds: float = 5.0
|
|||
|
|
|
|||
|
|
def wait(self, update_state=None) -> dict[str, Any]:
|
|||
|
|
while True:
|
|||
|
|
result = try_acquire(
|
|||
|
|
task_id=self.task_id,
|
|||
|
|
stage=self.stage,
|
|||
|
|
job_id=self.job_id,
|
|||
|
|
prefix=self.prefix,
|
|||
|
|
cpu=self.cpu,
|
|||
|
|
queue=self.queue,
|
|||
|
|
resource_class=self.resource_class,
|
|||
|
|
concurrency_group=self.concurrency_group,
|
|||
|
|
celery_task_id=self.celery_task_id,
|
|||
|
|
)
|
|||
|
|
if result.get("granted"):
|
|||
|
|
return result
|
|||
|
|
if update_state:
|
|||
|
|
update_state(result)
|
|||
|
|
time.sleep(self.poll_seconds)
|
|||
|
|
|
|||
|
|
def release(self, status: str = "completed", message: str = "") -> dict[str, Any]:
|
|||
|
|
return release(self.task_id, status=status, message=message)
|