feat(simulation): add Worker AI polygon runtime and terrain navigation
This commit is contained in:
DCCONSTRUCTIONS
@@ -0,0 +1,212 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user