229 lines
8.1 KiB
Python
229 lines
8.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
VORTEX_GPU — Linux Host Node Agent (port of vortex-node-agent.ps1)
|
|
Target: Ubuntu (nightmare .128, RTX 4080 SUPER 16GB)
|
|
|
|
Responsibilities (mirrors the Windows agent):
|
|
- Register this GPU box + stream nvidia-smi telemetry every N sec.
|
|
- Poll the gateway for jobs:
|
|
* shell : run an arbitrary cmd, return stdout
|
|
* hashcat / comfyui : run the command against the local GPU
|
|
* provision_comfyui : launch an ISOLATED ComfyUI instance on a dedicated
|
|
port + dedicated user/output/input dirs (tenant sees
|
|
a clean private machine; the physical GPU is shared/hidden)
|
|
* destroy_instance : kill the ComfyUI process for an instance
|
|
|
|
Auth: X-Node-Secret header (matches server.ts nodeAuthorized).
|
|
Runs as a systemd service: vortex-node-agent.service
|
|
"""
|
|
import json
|
|
import os
|
|
import socket
|
|
import subprocess
|
|
import time
|
|
import urllib.request
|
|
|
|
GATEWAY = os.environ.get("VORTEX_GATEWAY", "http://10.30.20.127:3000")
|
|
SECRET = os.environ.get("VORTEX_NODE_SECRET", "99496a5bf30b5a7411d3a60bf096ca642815fcf2be361113")
|
|
INTERVAL = int(os.environ.get("VORTEX_INTERVAL", "5"))
|
|
HOSTNAME = socket.gethostname()
|
|
|
|
HOME = os.path.expanduser("~")
|
|
COMFY_DIR = os.path.join(HOME, "ComfyUI")
|
|
COMFY_PY = os.path.join(COMFY_DIR, "venv", "bin", "python")
|
|
INST_DIR = os.path.join(HOME, "vortex-agent", "instances")
|
|
os.makedirs(INST_DIR, exist_ok=True)
|
|
|
|
_instance_pids = {} # instanceId -> pid
|
|
|
|
|
|
def http(method, path, body=None):
|
|
req = urllib.request.Request(GATEWAY + path, method=method)
|
|
req.add_header("X-Node-Secret", SECRET)
|
|
data = None
|
|
if body is not None:
|
|
data = json.dumps(body).encode()
|
|
req.add_header("Content-Type", "application/json")
|
|
with urllib.request.urlopen(req, data=data, timeout=60) as r:
|
|
return json.loads(r.read().decode())
|
|
|
|
|
|
def gpu_info():
|
|
try:
|
|
out = subprocess.run(
|
|
["nvidia-smi",
|
|
"--query-gpu=name,driver_version,memory.total,memory.used,utilization.gpu,temperature.gpu",
|
|
"--format=csv,noheader,nounits"],
|
|
capture_output=True, text=True, timeout=10).stdout.strip()
|
|
except Exception:
|
|
return None
|
|
if not out:
|
|
return None
|
|
p = [x.strip() for x in out.split(",")]
|
|
if len(p) < 6:
|
|
return None
|
|
try:
|
|
return {"gpuModel": p[0], "driverVersion": p[1], "memTotalMb": int(p[2]),
|
|
"memUsedMb": int(p[3]), "gpuUtilPct": int(p[4]), "tempC": int(p[5])}
|
|
except ValueError:
|
|
return None
|
|
|
|
|
|
def _procstat():
|
|
with open("/proc/stat") as f:
|
|
fields = f.readline().split()[1:]
|
|
idle = int(fields[3]) + int(fields[4])
|
|
total = sum(int(x) for x in fields)
|
|
return idle, total
|
|
|
|
|
|
def cpu_util():
|
|
try:
|
|
i1, t1 = _procstat()
|
|
time.sleep(0.2)
|
|
i2, t2 = _procstat()
|
|
if t2 == t1:
|
|
return 0
|
|
return int(100 * (1 - (i2 - i1) / (t2 - t1)))
|
|
except Exception:
|
|
return 0
|
|
|
|
|
|
def ram_info():
|
|
mem = {}
|
|
with open("/proc/meminfo") as f:
|
|
for line in f:
|
|
if ":" in line:
|
|
k = line.split(":")[0]
|
|
v = line.split(":")[1].strip().split()[0]
|
|
mem[k] = int(v)
|
|
total_gb = round(mem["MemTotal"] / 1024 / 1024, 1)
|
|
avail_gb = round(mem.get("MemAvailable", 0) / 1024 / 1024, 1)
|
|
return total_gb, round(total_gb - avail_gb, 1)
|
|
|
|
|
|
def uptime_sec():
|
|
return int(float(open("/proc/uptime").read().split()[0]))
|
|
|
|
|
|
def run_shell(command):
|
|
try:
|
|
r = subprocess.run(["bash", "-c", command], capture_output=True, text=True, timeout=600)
|
|
return (r.returncode == 0), (r.stdout or "") + (r.stderr or "")
|
|
except Exception as e:
|
|
return False, f"error: {e}"
|
|
|
|
|
|
def provision_comfyui(instance_id, port, user_dir):
|
|
os.makedirs(user_dir, exist_ok=True)
|
|
for sub in ("user", "output", "input", "models", "custom_nodes"):
|
|
os.makedirs(os.path.join(user_dir, sub), exist_ok=True)
|
|
if not os.path.exists(os.path.join(COMFY_DIR, "main.py")):
|
|
return False, "ComfyUI base not found on this node"
|
|
if not os.path.exists(COMFY_PY):
|
|
return False, "ComfyUI venv not found"
|
|
args = [COMFY_PY, os.path.join(COMFY_DIR, "main.py"),
|
|
"--listen", "0.0.0.0", "--port", str(port),
|
|
"--user-directory", os.path.join(user_dir, "user"),
|
|
"--output-directory", os.path.join(user_dir, "output"),
|
|
"--input-directory", os.path.join(user_dir, "input")]
|
|
try:
|
|
proc = subprocess.Popen(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
|
|
start_new_session=True)
|
|
_instance_pids[instance_id] = proc.pid
|
|
with open(os.path.join(user_dir, "instance.pid"), "w") as f:
|
|
f.write(str(proc.pid))
|
|
return True, f"launched pid={proc.pid} port={port}"
|
|
except Exception as e:
|
|
return False, f"failed to launch: {e}"
|
|
|
|
|
|
def destroy_instance(instance_id):
|
|
pid = _instance_pids.pop(instance_id, None)
|
|
if pid is None:
|
|
pidfile = os.path.join(INST_DIR, instance_id, "instance.pid")
|
|
if os.path.exists(pidfile):
|
|
try:
|
|
pid = int(open(pidfile).read().strip())
|
|
except Exception:
|
|
pid = None
|
|
if pid:
|
|
try:
|
|
os.kill(pid, 15)
|
|
return True, f"killed pid={pid}"
|
|
except Exception as e:
|
|
return True, f"kill pid={pid}: {e}"
|
|
return True, "no pid found"
|
|
|
|
|
|
def handle_job(job):
|
|
kind = job.get("kind")
|
|
cmd = job.get("command", "")
|
|
payload = job.get("payload", {}) or {}
|
|
if kind in ("shell", "hashcat", "comfyui"):
|
|
ok, result = run_shell(cmd)
|
|
elif kind == "provision_comfyui":
|
|
ok, result = provision_comfyui(payload.get("instanceId", "inst"),
|
|
payload.get("port", 8189),
|
|
payload.get("userDir") or os.path.join(INST_DIR, payload.get("instanceId", "inst")))
|
|
elif kind == "destroy_instance":
|
|
ok, result = destroy_instance(payload.get("instanceId", ""))
|
|
else:
|
|
ok, result = False, "unknown job kind"
|
|
try:
|
|
http("POST", f"/api/node/jobs/{job['id']}/result", {"ok": ok, "result": result})
|
|
except Exception as e:
|
|
print(f"[vortex-agent] result post failed: {e}", flush=True)
|
|
|
|
|
|
def main():
|
|
g = gpu_info()
|
|
ram_t, _ = ram_info()
|
|
print(f"[vortex-agent] registering node '{HOSTNAME}' with {GATEWAY} ...", flush=True)
|
|
try:
|
|
http("POST", "/api/node/register", {
|
|
"hostname": HOSTNAME,
|
|
"gpuModel": g["gpuModel"] if g else "GPU",
|
|
"driverVersion": g["driverVersion"] if g else "",
|
|
"memTotalMb": g["memTotalMb"] if g else 0,
|
|
"ramTotalGb": ram_t,
|
|
})
|
|
print(f"[vortex-agent] registered. GPU: {g['gpuModel'] if g else '?'}", flush=True)
|
|
except Exception as e:
|
|
print(f"[vortex-agent] register failed: {e}", flush=True)
|
|
|
|
print("[vortex-agent] agent active.", flush=True)
|
|
while True:
|
|
g = gpu_info()
|
|
ram_t, ram_u = ram_info()
|
|
cpu = cpu_util()
|
|
up = uptime_sec()
|
|
try:
|
|
http("POST", "/api/node/report", {
|
|
"hostname": HOSTNAME,
|
|
"gpuModel": g["gpuModel"] if g else "GPU",
|
|
"driverVersion": g["driverVersion"] if g else "",
|
|
"memTotalMb": g["memTotalMb"] if g else 0,
|
|
"memUsedMb": g["memUsedMb"] if g else 0,
|
|
"gpuUtilPct": g["gpuUtilPct"] if g else 0,
|
|
"tempC": g["tempC"] if g else 0,
|
|
"cpuUtilPct": cpu, "ramTotalGb": ram_t, "ramUsedGb": ram_u,
|
|
"uptimeSec": up,
|
|
})
|
|
except Exception as e:
|
|
print(f"[vortex-agent] heartbeat dropped: {e}", flush=True)
|
|
|
|
try:
|
|
resp = http("GET", f"/api/node/jobs?hostname={HOSTNAME}")
|
|
for job in resp.get("jobs", []):
|
|
print(f"[vortex-agent] job {job.get('id')} kind={job.get('kind')}", flush=True)
|
|
handle_job(job)
|
|
except Exception:
|
|
pass
|
|
|
|
time.sleep(INTERVAL)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|