110 lines
3.1 KiB
Python
110 lines
3.1 KiB
Python
#!/usr/bin/env python3
|
|
"""尝试通过 NVIDIA Sync 连接 DGXA，直到 LLM3D capabilities 绿灯。"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from urllib.request import urlopen
|
|
|
|
DEFAULT_BASE = os.environ.get("LLM3D_BASE", "http://192.168.130.16:8008").rstrip("/")
|
|
NVSYNC = Path(
|
|
os.environ.get(
|
|
"NVSYNC_EXE",
|
|
str(
|
|
Path(os.environ.get("LOCALAPPDATA", ""))
|
|
/ "Programs"
|
|
/ "nvidia-sync"
|
|
/ "resources"
|
|
/ "bin"
|
|
/ "nvsync-amd64.exe"
|
|
),
|
|
)
|
|
)
|
|
|
|
|
|
def caps(base: str) -> dict:
|
|
with urlopen(f"{base}/api/capabilities", timeout=90) as r:
|
|
return json.loads(r.read().decode("utf-8"))
|
|
|
|
|
|
def is_ready(c: dict) -> bool:
|
|
dgx = c.get("dgx") or {}
|
|
return bool(dgx.get("tokenrig_ready") or dgx.get("unirig_ready")) and bool(
|
|
dgx.get("unimate_ready")
|
|
)
|
|
|
|
|
|
def try_nvsync_connect(alias: str = "DGXA") -> None:
|
|
if not NVSYNC.is_file():
|
|
print(f"nvsync missing: {NVSYNC}", flush=True)
|
|
return
|
|
print(f"nvsync connect {alias} --detach ...", flush=True)
|
|
try:
|
|
subprocess.run(
|
|
[str(NVSYNC), "connect", alias, "--detach"],
|
|
timeout=90,
|
|
check=False,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
except subprocess.TimeoutExpired:
|
|
print("nvsync connect timed out (device may be offline)", flush=True)
|
|
st = subprocess.run(
|
|
[str(NVSYNC), "status", alias],
|
|
timeout=30,
|
|
check=False,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
print(st.stdout or st.stderr or "(no status)", flush=True)
|
|
|
|
|
|
def main() -> int:
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--base", default=DEFAULT_BASE)
|
|
ap.add_argument("--alias", default="DGXA")
|
|
ap.add_argument("--retries", type=int, default=30)
|
|
ap.add_argument("--interval", type=float, default=20.0)
|
|
ap.add_argument("--no-connect", action="store_true")
|
|
args = ap.parse_args()
|
|
|
|
for i in range(1, args.retries + 1):
|
|
print(f"=== attempt {i}/{args.retries} ===", flush=True)
|
|
if not args.no_connect:
|
|
try_nvsync_connect(args.alias)
|
|
try:
|
|
c = caps(args.base)
|
|
except Exception as exc: # noqa: BLE001
|
|
print(f"capabilities error: {exc}", flush=True)
|
|
time.sleep(args.interval)
|
|
continue
|
|
dgx = c.get("dgx") or {}
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"tokenrig_ready": dgx.get("tokenrig_ready"),
|
|
"unimate_ready": dgx.get("unimate_ready"),
|
|
"hy3d_ready": dgx.get("hy3d_ready"),
|
|
"tokenrig_detail": dgx.get("tokenrig_detail"),
|
|
},
|
|
ensure_ascii=False,
|
|
),
|
|
flush=True,
|
|
)
|
|
if is_ready(c):
|
|
print("READY", flush=True)
|
|
return 0
|
|
time.sleep(args.interval)
|
|
|
|
print("NOT_READY", flush=True)
|
|
return 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|