565 lines
20 KiB
Python
Executable File
565 lines
20 KiB
Python
Executable File
"""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)
|