feat(simulation): add Worker AI polygon runtime and terrain navigation

This commit is contained in:
DCCONSTRUCTIONS
2026-09-25 16:40:45 +03:00
parent a7c64e009d
commit f01bd39037
88 changed files with 9918 additions and 108 deletions
+398
View File
@@ -0,0 +1,398 @@
"""Synthetic protocol checks; these do not qualify rendering or AI model quality."""
import base64
import hashlib
import io
import json
import numpy as np
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from PIL import Image
from pydantic import ValidationError
from test_observatory_recorded_jobs import _definitions
from k1link.observatory.recorded_jobs import ObservatoryRecordedJobQueue
from k1link.simulation.ai_polygon.contracts import (
Decision,
RunApplied,
RunCreate,
RunSample,
WorkerHello,
WorkerWorldCreate,
WorldCreate,
WorldSettings,
)
from k1link.simulation.ai_polygon.policy import RoadPolicy
from k1link.simulation.ai_polygon.runs import RunStore
from k1link.simulation.ai_polygon.worlds import WorldStore, inspect_gaussian_ply
from k1link.web.ai_polygon_api import build_ai_polygon_router
def ply(value=0.0):
names = [
"x",
"y",
"z",
"opacity",
"f_dc_0",
"f_dc_1",
"f_dc_2",
"scale_0",
"scale_1",
"scale_2",
"rot_0",
"rot_1",
"rot_2",
"rot_3",
]
header = "ply\nformat binary_little_endian 1.0\nelement vertex 1\n"
header += "".join(f"property float {name}\n" for name in names) + "end_header\n"
return header.encode() + np.full(14, value, dtype="<f4").tobytes()
def make_world(worlds):
data = ply()
world = worlds.create(
WorldCreate(
name="Synthetic protocol fixture",
filename="fixture.PLY",
byte_length=len(data),
author="test",
license="CC0",
)
)
worlds.append(world["world_id"], 0, data)
worlds.complete(world["world_id"])
return worlds.configure(world["world_id"], WorldSettings(prepared=True))
@pytest.fixture
def runs(tmp_path):
queue = ObservatoryRecordedJobQueue(tmp_path, definitions=_definitions())
worlds = WorldStore(tmp_path)
store = RunStore(worlds, queue)
hello = WorkerHello(
worker_id="worker-006",
instance_id="a" * 32,
runtime="isaac-sim-6.1",
model_ids=["ddrnet-goose-pytorch-reference", "rf_detr_large"],
profile_sha256="b" * 64,
runtime_sources={key: "c" * 64 for key in ("worker", "scene", "models", "robot")},
)
store.register(hello)
return store, hello, make_world(worlds)
def sample(sequence=0, **changes):
stream = io.BytesIO()
Image.new("RGB", (800, 600)).save(stream, "JPEG")
row = dict(
sequence=sequence,
simulation_time_ns=sequence * 100_000_000,
inference_ms=5,
pose_xy=(0, 0),
decision=Decision(
speed_mps=0, yaw_rate_rps=0, reason="no-road", road_fraction=0, obstacle_count=0
),
image_jpeg_base64=base64.b64encode(stream.getvalue()).decode(),
)
return RunSample(**{**row, **changes})
def test_resume_and_atomic_publication_recovery(tmp_path):
worlds = WorldStore(tmp_path)
data = ply()
row = worlds.create(
WorldCreate(
name=" Test ", filename="scan.ply", byte_length=len(data), author="Test", license="CC0"
)
)
key = row["world_id"]
worlds.append(key, 0, data[:20])
assert worlds.prefix_hashes(key)["chunks"] == [
{"byte_length": 20, "sha256": hashlib.sha256(data[:20]).hexdigest()}
]
with pytest.raises(RuntimeError):
worlds.append(key, 0, data[:20])
with pytest.raises(RuntimeError):
worlds.complete(key)
worlds.append(key, 20, data[20:])
directory = worlds.directory(key)
(directory / "source.part").replace(directory / "source.ply")
recovered = WorldStore(tmp_path).complete(key)
assert recovered["status"] == "available"
assert recovered["sha256"] == hashlib.sha256(data).hexdigest()
assert recovered["name"] == "Test"
assert worlds.complete(key) == recovered
@pytest.mark.parametrize(
"data", [b"not ply", ply()[:-1], ply() + b"xxxx", ply(float("nan")), ply(float("inf"))]
)
def test_invalid_ply_rejected(tmp_path, data):
path = tmp_path / "bad.ply"
path.write_bytes(data)
with pytest.raises(ValueError):
inspect_gaussian_ply(path)
@pytest.mark.parametrize(
"settings",
[{"spawn_xy": [float("nan"), 0]}, {"rotation_degrees": [0, 0, 900]}, {"max_speed_mps": 2}],
)
def test_nonfinite_or_out_of_bounds_settings_rejected(settings):
with pytest.raises(ValidationError):
WorldSettings(**settings)
def test_start_snapshots_and_idempotency(runs):
store, hello, world = runs
request = RunCreate(world_id=world["world_id"], max_steps=2)
row = store.start(request, "request-001")
assert store.start(request, "request-001") == row
store.worlds.configure(world["world_id"], WorldSettings(max_speed_mps=0.5, prepared=True))
assert store.get(row["run_id"])["world"]["settings"]["max_speed_mps"] == 0.3
with pytest.raises(RuntimeError):
store.start(RunCreate(world_id=world["world_id"], max_steps=3), "request-001")
with pytest.raises(RuntimeError):
store.start(request, "request-002")
assert store.poll(hello.instance_id, None)["action"] == "load"
def test_sequence_clock_pause_step_stop_and_resource_release(runs):
store, hello, world = runs
key = store.start(RunCreate(world_id=world["world_id"], max_steps=2), "request-001")["run_id"]
with pytest.raises(RuntimeError):
store.control(key, "step")
store.poll(hello.instance_id, key)
store.sample(key, hello.instance_id, sample())
with pytest.raises(RuntimeError):
store.sample(key, hello.instance_id, sample())
store.applied(
key,
hello.instance_id,
RunApplied(sequence=0, simulation_time_ns=100_000_000, physics_steps=6, pose_xy=(0, 0)),
)
with pytest.raises(ValueError):
store.sample(key, hello.instance_id, sample(1, simulation_time_ns=999))
store.control(key, "pause")
assert store.get(key)["state"] == "running"
assert store.poll(hello.instance_id, key)["run"]["state"] == "paused"
store.control(key, "step")
assert store.poll(hello.instance_id, key)["action"] == "step"
assert store.poll(hello.instance_id, key)["action"] == "pause"
store.control(key, "stop")
with pytest.raises(RuntimeError):
store.sample(key, hello.instance_id, sample(1))
store.finish(key, hello.instance_id, "stopped", "")
store.queue.reserve_simulation("airun-" + "b" * 32)
def test_expired_worker_stops_run_but_does_not_claim_gpu_released(runs):
store, hello, world = runs
key = store.start(RunCreate(world_id=world["world_id"]), "request-001")["run_id"]
store.seen -= 21
assert not store.status()["available"]
assert store.get(key)["state"] == "failed"
with pytest.raises(RuntimeError):
store.sample(key, hello.instance_id, sample())
with pytest.raises(RuntimeError):
store.queue.reserve_simulation("airun-" + "b" * 32)
store.register(hello)
assert store.finish(key, hello.instance_id, "failed", "")["state"] == "failed"
store.queue.reserve_simulation("airun-" + "b" * 32)
def test_restart_fences_worker_and_preserves_journal(runs):
store, hello, world = runs
key = store.start(RunCreate(world_id=world["world_id"], max_steps=2), "request-001")["run_id"]
store.poll(hello.instance_id, key)
store.sample(key, hello.instance_id, sample())
restarted = RunStore(store.worlds, store.queue)
assert restarted.get(key)["state"] == "failed"
with pytest.raises(RuntimeError):
restarted.sample(key, hello.instance_id, sample(1))
saved = json.loads((store.directory(key) / "decisions.jsonl").read_text())
assert (
saved["image_sha256"]
== hashlib.sha256((store.directory(key) / "frames/000000.jpg").read_bytes()).hexdigest()
)
assert "image_jpeg_base64" not in saved
def test_completion_requires_expected_steps(runs):
store, hello, world = runs
key = store.start(RunCreate(world_id=world["world_id"], max_steps=1), "request-001")["run_id"]
with pytest.raises(RuntimeError):
store.finish(key, hello.instance_id, "completed", "")
store.poll(hello.instance_id, key)
store.sample(key, hello.instance_id, sample())
with pytest.raises(RuntimeError):
store.finish(key, hello.instance_id, "completed", "")
store.applied(
key,
hello.instance_id,
RunApplied(sequence=0, simulation_time_ns=100_000_000, physics_steps=6, pose_xy=(0, 0)),
)
assert store.finish(key, hello.instance_id, "completed", "")["state"] == "completed"
def test_api_auth_and_bounded_upload(tmp_path):
app = FastAPI()
app.include_router(build_ai_polygon_router(tmp_path))
with TestClient(app) as client:
root = "/api/v1/ai-polygon"
assert client.get(root + "/catalog").json()["runtime"]["available"] is False
assert client.post(root + "/worker/poll", json={"instance_id": "a" * 32}).status_code == 401
assert client.post(root + "/worlds", json={"filename": "../bad.ply"}).status_code == 422
world = make_world(WorldStore(tmp_path))
response = client.get(root + f"/worlds/{world['world_id']}/source.ply")
assert response.status_code == 200 and response.content == ply()
assert (
client.post(
root + "/runs",
headers={"Idempotency-Key": "test-run-001"},
json={"world_id": world["world_id"]},
).status_code
== 409
)
assert not any((tmp_path / "ai-polygon/runs").iterdir())
def test_camera_policy_brakes_and_waits_before_resume():
policy = RoadPolicy(0.3)
road = np.ones((512, 512), dtype=bool)
assert policy.decide(road, []).speed_mps == 0
policy.decide(road, [])
assert policy.decide(road, []).speed_mps == 0.3
obstacle = policy.decide(road, [(0.4, 0.2, 0.6, 0.8)])
assert obstacle.reason == "obstacle" and obstacle.speed_mps == 0
assert policy.decide(road, []).speed_mps == 0
assert policy.decide(np.zeros_like(road), []).reason == "no-road"
with pytest.raises(ValueError):
policy.decide(road, [(float("nan"), 0, 1, 1)])
def worker_asset():
return WorkerWorldCreate(
name="Synthetic paired scene",
filename="fixture.ply",
byte_length=len(ply()),
author="test",
license="CC0",
sha256="a" * 64,
collider_sha256="b" * 64,
splat_count=1,
settings=WorldSettings(prepared=True),
)
def test_worker_asset_is_manifest_only_and_retry_preserves_settings(tmp_path):
worlds = WorldStore(tmp_path)
request = worker_asset()
row = worlds.register_worker_asset(request, "worker-006")
assert {p.name for p in worlds.directory(row["world_id"]).iterdir()} == {"world.json"}
changed = worlds.configure(row["world_id"], WorldSettings(spawn_xy=(3, 4), prepared=True))
assert worlds.register_worker_asset(request, "worker-006") == changed
with pytest.raises(RuntimeError):
worlds.register_worker_asset(
request.model_copy(update={"collider_sha256": "c" * 64}), "worker-006"
)
with pytest.raises(RuntimeError):
worlds.prefix_hashes(row["world_id"])
def test_worker_asset_api_requires_current_authenticated_instance(tmp_path):
app = FastAPI()
app.include_router(build_ai_polygon_router(tmp_path))
base = "/api/v1/ai-polygon"
store = RunStore(WorldStore(tmp_path))
hello = WorkerHello(
worker_id="worker-006",
instance_id="a" * 32,
runtime="isaac-sim-6.1",
model_ids=["test"],
runtime_sources={k: "c" * 64 for k in ("worker", "scene", "models", "robot")},
profile_sha256="b" * 64,
)
auth = {"Authorization": "Bearer " + store.token, "Worker-Instance": hello.instance_id}
body = worker_asset().model_dump(mode="json")
with TestClient(app) as client:
assert (
client.post(
base + "/worker/worlds", json=body, headers={"Worker-Instance": hello.instance_id}
).status_code
== 401
)
assert client.post(base + "/worker/worlds", json=body, headers=auth).status_code == 409
assert (
client.post(
base + "/worker/register", json=hello.model_dump(), headers=auth
).status_code
== 200
)
response = client.post(base + "/worker/worlds", json=body, headers=auth)
assert response.status_code == 201
row = response.json()
assert row["storage"] == {"kind": "worker", "worker_id": "worker-006"}
assert client.get(base + f"/worlds/{row['world_id']}/source.ply").status_code == 409
assert (
client.post(
base + "/worker/worlds", json=body, headers={**auth, "Worker-Instance": "d" * 32}
).status_code
== 409
)
def test_camera_policy_turns_towards_connected_road():
policy = RoadPolicy(0.3)
road = np.zeros((512, 512), dtype=bool)
road[:, 100:260] = True
policy.decide(road, [])
policy.decide(road, [])
decision = policy.decide(road, [])
assert decision.yaw_rate_rps > 0 and 0 < decision.speed_mps < 0.3
def test_progress_does_not_claim_camera_ready(runs):
store, hello, world = runs
run = store.start(RunCreate(world_id=world["world_id"]), "progress-case-001")
store.progress(run["run_id"], hello.instance_id, "scene")
assert store.get(run["run_id"])["phase"] == "scene"
assert store.get(run["run_id"])["state"] == "starting"
store.control(run["run_id"], "stop")
assert store.progress(run["run_id"], hello.instance_id, "models")["control"] == "stop"
assert store.get(run["run_id"])["state"] == "stopping"
def test_worker_stop_race_accepts_already_exited_process(monkeypatch):
import importlib.util
from pathlib import Path
runtime = Path(__file__).resolve().parents[1] / "simulation/ai-polygon"
monkeypatch.syspath_prepend(str(runtime))
spec = importlib.util.spec_from_file_location("polygon_worker_stop_test", runtime / "worker.py")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
calls = []
class Child:
pid = 123
def poll(self):
return None
def wait(self, timeout):
calls.append(("wait", timeout))
return 0
def taskkill(*args, **kwargs):
# Windows reports a non-zero taskkill when the process exits in the race.
assert kwargs["check"] is False
calls.append(("taskkill", 255))
monkeypatch.setattr(module.subprocess, "run", taskkill)
module.terminate_episode(Child())
assert calls == [("taskkill", 255), ("wait", 30)]