213 lines
8.7 KiB
Python
213 lines
8.7 KiB
Python
"""Single-owner Windows coordinator; no GPU work while idle.
|
|
|
|
Core owns admission. Each episode gets a fresh native Isaac process and the
|
|
additive model stack. A missing Core heartbeat terminates this episode; it never
|
|
resumes an old command. A leftover active.json requires explicit reconciliation.
|
|
"""
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
from core_client import CoreClient
|
|
from model_stack import ROOT, ModelStack, sha256
|
|
|
|
|
|
def terminate_episode(child):
|
|
"""A concurrent normal exit is success; only a still-running child is failure."""
|
|
if child.poll() is None:
|
|
subprocess.run(
|
|
["taskkill", "/PID", str(child.pid), "/T", "/F"],
|
|
capture_output=True,
|
|
check=False,
|
|
timeout=30,
|
|
)
|
|
child.wait(timeout=30)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--core", default="http://127.0.0.1:18081")
|
|
parser.add_argument("--token-file", type=Path, required=True)
|
|
parser.add_argument("--state", type=Path, required=True)
|
|
parser.add_argument("--isaac", type=Path, required=True)
|
|
parser.add_argument("--once", action="store_true")
|
|
args = parser.parse_args()
|
|
args.state.mkdir(parents=True, exist_ok=True)
|
|
active_file = args.state / "active.json"
|
|
if active_file.exists():
|
|
raise RuntimeError("Reconcile the previous simulation before starting a new worker")
|
|
lock = args.state / "worker.lock"
|
|
instance = uuid4().hex
|
|
client = CoreClient(args.core, args.token_file, instance)
|
|
stack = ModelStack()
|
|
stack.preflight()
|
|
robot_root = args.isaac.parent / "assets/jetbot-6.1-v1"
|
|
robot_manifest = robot_root / "asset-manifest.json"
|
|
assets = json.loads(robot_manifest.read_text(encoding="utf-8-sig"))
|
|
for asset in assets["files"]:
|
|
if sha256(robot_root / asset["path"]) != asset["sha256"]:
|
|
raise RuntimeError("Prepared Jetbot asset changed")
|
|
sources = {
|
|
"worker": sha256(ROOT / "worker.py"),
|
|
"scene": sha256(ROOT / "run_scene.py"),
|
|
"models": sha256(stack.profile_path),
|
|
"robot": sha256(robot_manifest),
|
|
}
|
|
hello = {
|
|
"worker_id": "worker-006-ai-polygon",
|
|
"instance_id": instance,
|
|
"runtime": "isaac-sim-6.1",
|
|
"model_ids": [m["id"] for m in stack.profile["models"]],
|
|
"profile_sha256": hashlib.sha256(json.dumps(sources, sort_keys=True).encode()).hexdigest(),
|
|
"runtime_sources": sources,
|
|
}
|
|
lost = threading.Event()
|
|
finished = threading.Event()
|
|
child = None
|
|
|
|
def heartbeat():
|
|
while not finished.wait(3):
|
|
try:
|
|
client.request("/worker/heartbeat", {"instance_id": instance})
|
|
except Exception:
|
|
lost.set()
|
|
return
|
|
|
|
fd = os.open(lock, os.O_WRONLY | os.O_CREAT | os.O_EXCL)
|
|
os.close(fd)
|
|
try:
|
|
client.request("/worker/register", hello)
|
|
threading.Thread(target=heartbeat, daemon=True).start()
|
|
while not lost.is_set():
|
|
polled = client.request("/worker/poll", {"instance_id": instance})
|
|
if polled["action"] == "idle":
|
|
time.sleep(0.5)
|
|
continue
|
|
run = polled["run"]
|
|
if not re.fullmatch(r"airun-[a-f0-9]{32}", run["run_id"]):
|
|
raise ValueError("Invalid run identity")
|
|
episode = args.state / run["run_id"]
|
|
episode.mkdir()
|
|
run_file = episode / "run.json"
|
|
run_file.write_text(json.dumps(run))
|
|
active_file.write_text(
|
|
json.dumps(
|
|
{
|
|
"run_id": run["run_id"],
|
|
"instance_id": instance,
|
|
"profile_sha256": hello["profile_sha256"],
|
|
}
|
|
)
|
|
)
|
|
result = {"outcome": "failed", "message": "Worker прервал прогон."}
|
|
released = False
|
|
try:
|
|
cache = args.state / "worlds"
|
|
cache.mkdir(exist_ok=True)
|
|
source = cache / (run["world"]["sha256"] + ".ply")
|
|
client.download(run["world"], source)
|
|
if lost.is_set():
|
|
raise RuntimeError("Core connection was lost during scene preparation")
|
|
control = client.request(
|
|
"/worker/runs/" + run["run_id"] + "/progress", {"phase": "models"}
|
|
)
|
|
if control["control"] == "stop":
|
|
raise InterruptedError("Stopped before model startup")
|
|
stack.start(
|
|
cancelled=lambda run_id=run["run_id"]: (
|
|
lost.is_set()
|
|
or client.request("/runs/" + run_id)["control"] == "stop"
|
|
)
|
|
)
|
|
if lost.is_set():
|
|
raise RuntimeError("Core connection was lost during model startup")
|
|
control = client.request(
|
|
"/worker/runs/" + run["run_id"] + "/progress", {"phase": "scene"}
|
|
)
|
|
if control["control"] == "stop":
|
|
raise InterruptedError("Stopped before scene startup")
|
|
command = [
|
|
str(args.isaac / "python.bat"),
|
|
str(ROOT / "run_scene.py"),
|
|
"--run",
|
|
str(run_file),
|
|
"--source",
|
|
str(source),
|
|
"--core",
|
|
args.core,
|
|
"--token-file",
|
|
str(args.token_file),
|
|
"--instance",
|
|
instance,
|
|
]
|
|
with (episode / "isaac.log").open("wb") as output:
|
|
child = subprocess.Popen(command, stdout=output, stderr=subprocess.STDOUT)
|
|
stopped = False
|
|
next_control_check = 0.0
|
|
deadline = time.monotonic() + 300 + run["request"]["max_steps"] * 15
|
|
while (
|
|
child.poll() is None and not lost.is_set() and time.monotonic() < deadline
|
|
):
|
|
if time.monotonic() >= next_control_check:
|
|
state = client.request("/runs/" + run["run_id"])
|
|
if state["control"] == "stop" or state["state"] == "failed":
|
|
stopped = state["control"] == "stop"
|
|
break
|
|
next_control_check = time.monotonic() + 1
|
|
time.sleep(0.2)
|
|
if child.poll() is None:
|
|
terminate_episode(child)
|
|
child.wait(timeout=30)
|
|
child = None
|
|
result_file = episode / "result.json"
|
|
if stopped:
|
|
result = {"outcome": "stopped", "message": "Прогон остановлен."}
|
|
elif result_file.exists() and not lost.is_set():
|
|
result = json.loads(result_file.read_text())
|
|
except Exception as exc:
|
|
(episode / "worker-error.txt").write_text(type(exc).__name__ + ": " + str(exc))
|
|
finally:
|
|
if child is not None and child.poll() is None:
|
|
terminate_episode(child)
|
|
child.wait(timeout=30)
|
|
stack.stop()
|
|
released = True
|
|
# A Stop may race with the final sample/normal native exit. It remains a
|
|
# successful cancellation only after native and model cleanup above.
|
|
if not lost.is_set():
|
|
state = client.request("/runs/" + run["run_id"])
|
|
if state["control"] == "stop" and state["state"] == "stopping":
|
|
result = {"outcome": "stopped", "message": "Движение и inference остановлены."}
|
|
# Do not release the reservation until native process AND GPU containers are gone.
|
|
for attempt in range(10):
|
|
try:
|
|
if lost.is_set():
|
|
client.request("/worker/register", hello)
|
|
client.request(
|
|
"/worker/runs/" + run["run_id"] + "/finish",
|
|
{**result, "instance_id": instance, "resources_released": released},
|
|
)
|
|
active_file.unlink()
|
|
break
|
|
except Exception:
|
|
if attempt == 9:
|
|
raise
|
|
time.sleep(1)
|
|
if args.once or lost.is_set():
|
|
break
|
|
finally:
|
|
finished.set()
|
|
lock.unlink(missing_ok=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|