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