155 lines
6.9 KiB
Python
155 lines
6.9 KiB
Python
"""Worker-local, bounded asynchronous inference. No operator-network dependencies."""
|
|
|
|
import hashlib
|
|
import json
|
|
import threading
|
|
import time
|
|
import traceback
|
|
from pathlib import Path
|
|
|
|
|
|
class LatestInference:
|
|
def __init__(self, factory, policy, directory: Path, clock=time.monotonic):
|
|
self.factory, self.policy, self.directory, self.clock = factory, policy, directory, clock
|
|
self.condition = threading.Condition()
|
|
self.enabled = self.closed = self.ready = False
|
|
self.ever_ready = False
|
|
self.epoch = self.count = self.dropped = 0
|
|
self.pending = self.result = None
|
|
self.error = None
|
|
self.thread = threading.Thread(target=self._run, name="polygon-inference", daemon=True)
|
|
|
|
def start(self):
|
|
self.thread.start()
|
|
|
|
def enable(self, enabled):
|
|
with self.condition:
|
|
if enabled != self.enabled:
|
|
self.enabled = enabled
|
|
self.epoch += 1
|
|
self.pending = self.result = None
|
|
self.condition.notify_all()
|
|
|
|
def submit(self, rgb, frame_id, captured_at, simulation_time_ns, observation=None):
|
|
with self.condition:
|
|
if not self.enabled or self.closed:
|
|
return
|
|
if self.pending is not None:
|
|
self.dropped += 1
|
|
self.pending = (rgb, frame_id, captured_at, simulation_time_ns, self.epoch, observation)
|
|
self.condition.notify_all()
|
|
|
|
def command(self, now, frame_deadline=0.5, command_deadline=0.5):
|
|
with self.condition:
|
|
result = self.result
|
|
if not self.enabled:
|
|
return 0.0, 0.0, "paused", result
|
|
if self.error:
|
|
return 0.0, 0.0, "inference-error", result
|
|
if result is None or now - result["captured_at"] > frame_deadline:
|
|
return 0.0, 0.0, "stale-camera", result
|
|
if now - result["completed_at"] > command_deadline:
|
|
return 0.0, 0.0, "stale-command", result
|
|
decision = result["decision"]
|
|
return decision["speed_mps"], decision["yaw_rate_rps"], "none", result
|
|
|
|
def _run(self):
|
|
model = None
|
|
policy_epoch = None
|
|
self.directory.mkdir(parents=True, exist_ok=True)
|
|
try:
|
|
with (self.directory / "decisions.jsonl").open("a", encoding="utf-8") as journal:
|
|
while True:
|
|
with self.condition:
|
|
self.condition.wait_for(
|
|
lambda: self.closed or (self.enabled and self.pending is not None)
|
|
)
|
|
if self.closed:
|
|
return
|
|
rgb, frame_id, captured_at, sim_ns, epoch, observation = self.pending
|
|
self.pending = None
|
|
try:
|
|
if model is None:
|
|
model = self.factory()
|
|
model.ready()
|
|
started = self.clock()
|
|
evidence = {}
|
|
if self.policy is None:
|
|
if policy_epoch is not None and policy_epoch != epoch:
|
|
model.reset()
|
|
policy_epoch = epoch
|
|
decision, boxes, evidence = model.infer_observation(
|
|
rgb, observation, frame_id
|
|
)
|
|
else:
|
|
road, boxes = model.infer(rgb)
|
|
with self.condition:
|
|
if epoch != self.epoch or not self.enabled:
|
|
continue
|
|
if self.policy is not None and policy_epoch != epoch:
|
|
self.policy.reset()
|
|
policy_epoch = epoch
|
|
if self.policy is not None:
|
|
decision = self.policy.decide(road, boxes).model_dump()
|
|
completed = self.clock()
|
|
record = dict(
|
|
frame_id=frame_id,
|
|
captured_at=captured_at,
|
|
simulation_time_ns=sim_ns,
|
|
completed_at=completed,
|
|
inference_ms=(completed - started) * 1000,
|
|
decision=decision,
|
|
)
|
|
with self.condition:
|
|
self.ready, self.ever_ready, self.error = True, True, None
|
|
if epoch != self.epoch or not self.enabled:
|
|
continue # A pause/resume fences every earlier in-flight decision.
|
|
self.result = record
|
|
self.count += 1
|
|
# Archive source RGB on Worker, independently of the video stream.
|
|
from PIL import Image
|
|
|
|
image_path = self.directory / f"{frame_id:08d}.jpg"
|
|
Image.fromarray(rgb).save(image_path, "JPEG", quality=90)
|
|
record = {
|
|
**record,
|
|
"image": image_path.name,
|
|
"image_sha256": hashlib.sha256(image_path.read_bytes()).hexdigest(),
|
|
"obstacles": boxes,
|
|
**evidence,
|
|
}
|
|
journal.write(json.dumps(record, allow_nan=False) + "\n")
|
|
journal.flush()
|
|
except Exception as exc:
|
|
with (self.directory / "errors.jsonl").open(
|
|
"a", encoding="utf-8"
|
|
) as errors:
|
|
errors.write(
|
|
json.dumps(
|
|
{
|
|
"frame_id": frame_id,
|
|
"monotonic": self.clock(),
|
|
"error": str(exc),
|
|
"type": type(exc).__name__,
|
|
"traceback": traceback.format_exc(),
|
|
}
|
|
)
|
|
+ "\n"
|
|
)
|
|
with self.condition:
|
|
self.error, self.ready, self.result = type(exc).__name__, False, None
|
|
if model is not None:
|
|
model.close()
|
|
model = None
|
|
with self.condition:
|
|
self.condition.wait(timeout=0.5)
|
|
finally:
|
|
if model is not None:
|
|
model.close()
|
|
|
|
def close(self):
|
|
with self.condition:
|
|
self.closed = True
|
|
self.condition.notify_all()
|
|
self.thread.join(timeout=12)
|