Add runtime environment readiness checks
This commit is contained in:
@@ -1,11 +1,70 @@
|
||||
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] = []
|
||||
@@ -63,6 +122,152 @@ def get_conda_envs() -> dict:
|
||||
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 {
|
||||
|
||||
Reference in New Issue
Block a user