311 lines
11 KiB
Python
311 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import shutil
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import fcntl
|
|
|
|
from ...config import settings
|
|
|
|
READINESS_CACHE_SECONDS = 300
|
|
_readiness_cache: tuple[float, dict[str, Any]] | None = None
|
|
_readiness_thread_lock = threading.Lock()
|
|
|
|
|
|
@contextmanager
|
|
def readiness_probe_lock():
|
|
settings.project_root.joinpath("var").mkdir(parents=True, exist_ok=True)
|
|
lock_path = settings.project_root / "var" / "runtime_readiness.lock"
|
|
with _readiness_thread_lock, lock_path.open("w") as lock_file:
|
|
fcntl.flock(lock_file, fcntl.LOCK_EX)
|
|
try:
|
|
yield
|
|
finally:
|
|
fcntl.flock(lock_file, fcntl.LOCK_UN)
|
|
|
|
|
|
def runtime_environment_specs() -> list[dict[str, Any]]:
|
|
return [
|
|
{
|
|
"role": "backend_task",
|
|
"name": settings.task_conda_env,
|
|
"label": "Backend, SegModel, YOLO",
|
|
"env_file": "envs/seg_smp.yml",
|
|
"required_imports": [
|
|
{"module": "fastapi", "package": "fastapi"},
|
|
{"module": "uvicorn", "package": "uvicorn"},
|
|
{"module": "torch", "package": "torch"},
|
|
{"module": "cv2", "package": "opencv-python"},
|
|
{"module": "segmentation_models_pytorch", "package": "segmentation-models-pytorch"},
|
|
{"module": "ultralytics", "package": "ultralytics"},
|
|
{"module": "albumentations", "package": "albumentations"},
|
|
{"module": "mmengine", "package": "mmengine"},
|
|
{"module": "mmseg", "package": "mmsegmentation"},
|
|
],
|
|
},
|
|
{
|
|
"role": "mmseg_full",
|
|
"name": settings.mmseg_conda_env,
|
|
"label": "MMSeg full mmcv runtime",
|
|
"env_file": "envs/seg_mmcv.yml",
|
|
"required_imports": [
|
|
{"module": "torch", "package": "torch"},
|
|
{"module": "cv2", "package": "opencv-python"},
|
|
{"module": "mmcv", "package": "mmcv"},
|
|
{"module": "mmcv._ext", "package": "mmcv"},
|
|
{"module": "mmengine", "package": "mmengine"},
|
|
{"module": "mmseg", "package": "mmsegmentation"},
|
|
],
|
|
},
|
|
]
|
|
|
|
|
|
def parse_nvidia_smi_csv(output: str) -> list[dict]:
|
|
gpus: list[dict] = []
|
|
for line in output.splitlines():
|
|
if not line.strip():
|
|
continue
|
|
parts = [part.strip() for part in line.split(",")]
|
|
if len(parts) < 7:
|
|
continue
|
|
index, name, total, used, free, util, temp = parts[:7]
|
|
try:
|
|
gpus.append(
|
|
{
|
|
"index": int(index),
|
|
"name": name,
|
|
"memory_total_mb": int(total),
|
|
"memory_used_mb": int(used),
|
|
"memory_free_mb": int(free),
|
|
"utilization_gpu_percent": int(util),
|
|
"temperature_c": int(temp),
|
|
}
|
|
)
|
|
except ValueError:
|
|
continue
|
|
return gpus
|
|
|
|
|
|
def get_gpus() -> dict:
|
|
cmd = [
|
|
"nvidia-smi",
|
|
"--query-gpu=index,name,memory.total,memory.used,memory.free,utilization.gpu,temperature.gpu",
|
|
"--format=csv,noheader,nounits",
|
|
]
|
|
try:
|
|
result = subprocess.run(cmd, capture_output=True, text=True, check=True)
|
|
return {"available": True, "gpus": parse_nvidia_smi_csv(result.stdout)}
|
|
except Exception as exc:
|
|
return {"available": False, "gpus": [], "error": str(exc)}
|
|
|
|
|
|
def get_conda_envs() -> dict:
|
|
try:
|
|
result = subprocess.run(["conda", "env", "list"], capture_output=True, text=True, check=True)
|
|
except Exception as exc:
|
|
return {"available": False, "envs": [], "error": str(exc)}
|
|
envs = []
|
|
for line in result.stdout.splitlines():
|
|
raw = line.strip()
|
|
if not raw or raw.startswith("#"):
|
|
continue
|
|
marker = "*" in raw.split()
|
|
parts = raw.replace("*", " ").split()
|
|
if len(parts) >= 2:
|
|
envs.append({"name": parts[0], "path": parts[-1], "active": marker})
|
|
return {"available": True, "envs": envs, "task_default": settings.task_conda_env, "mmseg_default": settings.mmseg_conda_env}
|
|
|
|
|
|
def probe_code(required_imports: list[dict[str, str]]) -> str:
|
|
imports_json = json.dumps(required_imports, ensure_ascii=False)
|
|
return f"""
|
|
import importlib
|
|
import importlib.metadata as metadata
|
|
import json
|
|
import platform
|
|
import sys
|
|
|
|
required = {imports_json}
|
|
checks = []
|
|
for item in required:
|
|
module_name = item["module"]
|
|
package_name = item.get("package") or module_name
|
|
try:
|
|
module = importlib.import_module(module_name)
|
|
version = getattr(module, "__version__", None)
|
|
if version is None:
|
|
try:
|
|
version = metadata.version(package_name)
|
|
except Exception:
|
|
version = None
|
|
checks.append({{
|
|
"module": module_name,
|
|
"package": package_name,
|
|
"passed": True,
|
|
"version": str(version) if version is not None else None,
|
|
}})
|
|
except Exception as exc:
|
|
checks.append({{
|
|
"module": module_name,
|
|
"package": package_name,
|
|
"passed": False,
|
|
"error": f"{{type(exc).__name__}}: {{exc}}",
|
|
}})
|
|
|
|
extra = {{"python": sys.version.split()[0], "platform": platform.platform()}}
|
|
try:
|
|
import torch
|
|
extra["torch_cuda_available"] = bool(torch.cuda.is_available())
|
|
extra["torch_cuda"] = getattr(torch.version, "cuda", None)
|
|
except Exception:
|
|
pass
|
|
print(json.dumps({{"checks": checks, "extra": extra}}, ensure_ascii=False))
|
|
"""
|
|
|
|
|
|
def parse_probe_stdout(stdout: str) -> dict[str, Any]:
|
|
for line in reversed(stdout.splitlines()):
|
|
text = line.strip()
|
|
if not text or not text.startswith("{"):
|
|
continue
|
|
return json.loads(text)
|
|
raise ValueError("probe did not emit JSON")
|
|
|
|
|
|
def inspect_conda_env(env_name: str, required_imports: list[dict[str, str]], timeout: int = 45, retries: int = 1) -> dict[str, Any]:
|
|
command = ["conda", "run", "-n", env_name, "python", "-c", probe_code(required_imports)]
|
|
attempts = []
|
|
result: subprocess.CompletedProcess[str] | None = None
|
|
for attempt in range(retries + 1):
|
|
result = subprocess.run(command, capture_output=True, text=True, timeout=timeout)
|
|
attempts.append({"attempt": attempt + 1, "returncode": result.returncode})
|
|
if result.returncode == 0 or result.returncode not in {139, -11} or attempt >= retries:
|
|
break
|
|
time.sleep(2)
|
|
if result is None:
|
|
raise RuntimeError("probe did not run")
|
|
detail = {
|
|
"command": command[:4] + ["..."],
|
|
"returncode": result.returncode,
|
|
"attempts": attempts,
|
|
"stdout_tail": result.stdout[-2000:],
|
|
"stderr_tail": result.stderr[-2000:],
|
|
}
|
|
if result.returncode != 0:
|
|
return {"passed": False, "checks": [], "extra": {}, "detail": detail}
|
|
try:
|
|
parsed = parse_probe_stdout(result.stdout)
|
|
except Exception as exc:
|
|
return {"passed": False, "checks": [], "extra": {}, "detail": {**detail, "error": str(exc)}}
|
|
checks = parsed.get("checks", [])
|
|
return {
|
|
"passed": bool(checks) and all(item.get("passed") for item in checks),
|
|
"checks": checks,
|
|
"extra": parsed.get("extra", {}),
|
|
"detail": detail,
|
|
}
|
|
|
|
|
|
def get_runtime_readiness(force: bool = False) -> dict[str, Any]:
|
|
global _readiness_cache
|
|
now = time.time()
|
|
if not force and _readiness_cache and now - _readiness_cache[0] < READINESS_CACHE_SECONDS:
|
|
cached = dict(_readiness_cache[1])
|
|
cached["cached"] = True
|
|
return cached
|
|
|
|
with readiness_probe_lock():
|
|
now = time.time()
|
|
if not force and _readiness_cache and now - _readiness_cache[0] < READINESS_CACHE_SECONDS:
|
|
cached = dict(_readiness_cache[1])
|
|
cached["cached"] = True
|
|
return cached
|
|
|
|
conda = get_conda_envs()
|
|
env_paths = {item["name"]: item["path"] for item in conda.get("envs", [])}
|
|
envs: list[dict[str, Any]] = []
|
|
for spec in runtime_environment_specs():
|
|
env_name = spec["name"]
|
|
exists = env_name in env_paths
|
|
env_report: dict[str, Any] = {
|
|
"role": spec["role"],
|
|
"name": env_name,
|
|
"label": spec["label"],
|
|
"env_file": spec["env_file"],
|
|
"path": env_paths.get(env_name),
|
|
"exists": exists,
|
|
"passed": False,
|
|
"checks": [],
|
|
"extra": {},
|
|
}
|
|
if exists:
|
|
probe = inspect_conda_env(env_name, spec["required_imports"])
|
|
env_report.update(probe)
|
|
envs.append(env_report)
|
|
|
|
payload = {
|
|
"available": bool(conda.get("available")),
|
|
"passed": bool(conda.get("available")) and all(item["passed"] for item in envs),
|
|
"generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(now)),
|
|
"cache_seconds": READINESS_CACHE_SECONDS,
|
|
"cached": False,
|
|
"envs": envs,
|
|
"specs": {
|
|
"bootstrap_script": "scripts/bootstrap_conda_envs.sh",
|
|
"verify_script": "scripts/verify_runtime_envs.py",
|
|
"env_files": [spec["env_file"] for spec in runtime_environment_specs()],
|
|
"task_default": settings.task_conda_env,
|
|
"mmseg_default": settings.mmseg_conda_env,
|
|
},
|
|
}
|
|
_readiness_cache = (now, payload)
|
|
return payload
|
|
|
|
|
|
def disk_usage() -> dict:
|
|
usage = shutil.disk_usage(settings.source_root)
|
|
return {
|
|
"path": str(settings.source_root),
|
|
"total": usage.total,
|
|
"used": usage.used,
|
|
"free": usage.free,
|
|
}
|
|
|
|
|
|
def scan_results() -> list[dict]:
|
|
roots = [
|
|
settings.source_root / "DataSet_Public_outputs",
|
|
settings.source_root / "BestMode_Predict_Results_DataSet_Public",
|
|
settings.source_root / "Hardisk",
|
|
settings.source_root / "Seg_All_In_One_Analysis",
|
|
]
|
|
exts = {".csv", ".png", ".jpg", ".jpeg", ".svg", ".log", ".pth", ".pt"}
|
|
results: list[dict] = []
|
|
for root in roots:
|
|
if not root.exists():
|
|
continue
|
|
for path in root.rglob("*"):
|
|
if path.is_file() and path.suffix.lower() in exts:
|
|
try:
|
|
stat = path.stat()
|
|
results.append(
|
|
{
|
|
"name": path.name,
|
|
"path": str(path.resolve()),
|
|
"relative_path": str(path.resolve().relative_to(settings.source_root)),
|
|
"size": stat.st_size,
|
|
"modified": stat.st_mtime,
|
|
"kind": path.suffix.lower().lstrip("."),
|
|
}
|
|
)
|
|
except OSError:
|
|
continue
|
|
results.sort(key=lambda item: item["modified"], reverse=True)
|
|
return results[:1000]
|