StarPilot/openpilot/starpilot/models/tests/test_shared_camera_warp.py

from dataclasses import replace
from pathlib import Path
from types import SimpleNamespace
import unittest
from unittest.mock import Mock, patch

import numpy as np
from tinygrad import Context, dtypes

from openpilot.starpilot.models.runner import CatalogModelState
from openpilot.starpilot.models.tests.test_runner import artifact_fixture


def fake_warp(shape=(2, 6, 2, 2), dtype=dtypes.uint8, device="CPU"):
  return SimpleNamespace(shape=shape, dtype=dtype, device=device)


def make_runner(artifact=None, camera=(8, 8)):
  artifact = artifact_fixture() if artifact is None else artifact
  artifact[camera] = Mock(return_value=fake_warp())
  with Context(DEV="CPU:LLVM"), patch("openpilot.starpilot.models.runner.load_oob", return_value=artifact):
    runner = CatalogModelState(camera[0], camera[1], Path("verified.pkl"), "v15", False)
  return runner


def frame_inputs(runner):
  frames = {name: np.zeros(runner.frame_buf_size, dtype=np.uint8) for name in runner.vision_input_names}
  transforms = {name: np.eye(3, dtype=np.float32) for name in runner.vision_input_names}
  inputs = {"desire_pulse": np.array([0, 1, 0, 0, 0, 0, 0, 0], dtype=np.float32),
            "prev_action": np.array([2, 3], dtype=np.float32)}
  return frames, transforms, inputs


class TestSharedCameraWarp(unittest.TestCase):
  def test_shared_frame_runs_one_warp_and_two_independent_policies(self):
    lateral, longitudinal, ordinary = (make_runner() for _ in range(3))
    frames, transforms, inputs = frame_inputs(lateral)
    lat_output = lateral.run(frames, transforms, inputs)
    packet = lateral.last_warp
    callback = Mock()
    long_output = longitudinal.run(frames, transforms, inputs, callback, shared_warp=packet)
    normal_output = ordinary.run(frames, transforms, inputs)
    lateral.warp_enqueue.assert_called_once()
    longitudinal.warp_enqueue.assert_not_called()
    ordinary.warp_enqueue.assert_called_once()
    callback.assert_called_once_with()
    self.assertIs(longitudinal.run_policy.call_args.kwargs["warped"], packet.tensor)
    self.assertIs(longitudinal.last_warp, packet)
    for key in lat_output:
      np.testing.assert_array_equal(long_output[key], normal_output[key])
    for key in lateral.input_queues:
      self.assertIsNot(lateral.input_queues[key], longitudinal.input_queues[key])
    for key in lateral.npy:
      self.assertFalse(np.shares_memory(lateral.npy[key], longitudinal.npy[key]))
    longitudinal.run(frames, transforms, {**inputs, "desire_pulse": np.zeros(8)})
    np.testing.assert_array_equal(lateral.prev_desire, inputs["desire_pulse"])
    np.testing.assert_array_equal(longitudinal.prev_desire, np.zeros(8))
    self.assertIs(lateral.last_warp, packet)

  def test_history_length_and_sampling_remain_private(self):
    artifact = artifact_fixture()
    artifact["metadata"]["model"]["input_shapes"]["img"] = (1, 6, 2, 2)
    artifact["metadata"]["model"]["input_shapes"]["big_img"] = (1, 6, 2, 2)
    artifact["frame_skip"] = 1
    lateral, longitudinal = make_runner(), make_runner(artifact)
    self.assertEqual(lateral.camera_warp_descriptor, longitudinal.camera_warp_descriptor)
    self.assertNotEqual(lateral.input_queues["img_q"].shape, longitudinal.input_queues["img_q"].shape)
    args = frame_inputs(lateral)
    lateral.run(*args)
    longitudinal.run(*args, shared_warp=lateral.last_warp)
    longitudinal.warp_enqueue.assert_not_called()

  def test_descriptor_and_actual_tensor_mismatches_rejected_before_state_update(self):
    lateral, longitudinal = make_runner(), make_runner()
    args = frame_inputs(lateral)
    lateral.run(*args)
    packet = lateral.last_warp
    descriptor = packet.descriptor
    mismatches = [replace(packet, descriptor=replace(descriptor, camera_size=(16, 8))),
                  replace(packet, descriptor=replace(descriptor, nv12_layout=(1, 2, 3, 4))),
                  replace(packet, descriptor=replace(descriptor, image_layouts=((4, 4), (4, 4)))),
                  replace(packet, descriptor=replace(descriptor, device="AMD")),
                  replace(packet, tensor=fake_warp(shape=(2, 6, 4, 4))),
                  replace(packet, tensor=fake_warp(dtype=dtypes.float32)),
                  replace(packet, tensor=fake_warp(device="AMD"))]
    for mismatch in mismatches:
      with self.subTest(packet=mismatch), self.assertRaises(ValueError):
        longitudinal.run(*args, shared_warp=mismatch)
    longitudinal.run_policy.assert_not_called()
    np.testing.assert_array_equal(longitudinal.prev_desire, np.zeros(8))

  def test_new_buffers_or_transforms_cannot_reuse_camera_warp(self):
    lateral, longitudinal = make_runner(), make_runner()
    frames, transforms, inputs = frame_inputs(lateral)
    lateral.run(frames, transforms, inputs)
    packet = lateral.last_warp
    other_frames = {key: value.copy() for key, value in frames.items()}
    with self.assertRaisesRegex(ValueError, "frame/transforms"):
      longitudinal.run(other_frames, transforms, inputs, shared_warp=packet)
    changed = {key: value.copy() for key, value in transforms.items()}
    changed["big_img"][0, 2] = 1
    with self.assertRaisesRegex(ValueError, "frame/transforms"):
      longitudinal.run(frames, changed, inputs, shared_warp=packet)
    longitudinal.run_policy.assert_not_called()

  def test_old_packet_and_repeat_consumption_rejected(self):
    lateral, longitudinal = make_runner(), make_runner()
    args = frame_inputs(lateral)
    lateral.run(*args)
    packet = lateral.last_warp
    longitudinal.run(*args, shared_warp=packet)
    with self.assertRaisesRegex(ValueError, "already consumed"):
      longitudinal.run(*args, shared_warp=packet)
    lateral.run(*args)
    with self.assertRaisesRegex(ValueError, "stale"):
      make_runner().run(*args, shared_warp=packet)

  def test_failed_output_and_warmup_invalidate_exposed_packet(self):
    producer = make_runner()
    args = frame_inputs(producer)
    producer.run(*args)
    packet = producer.last_warp
    producer.run_policy.return_value = (Mock(numpy=lambda: np.array([np.nan])),)
    with self.assertRaisesRegex(ValueError, "not finite"):
      producer.run(*args)
    self.assertIsNone(producer.last_warp)
    with self.assertRaisesRegex(ValueError, "stale"):
      make_runner().run(*args, shared_warp=packet)
    producer = make_runner()
    producer.run(*frame_inputs(producer))
    producer.warmup()
    self.assertIsNone(producer.last_warp)
    self.assertIsNone(producer._last_shared_warp)

  def test_borrowed_packet_relay_keeps_original_producer_lifetime(self):
    for invalidate in ("advance", "reset"):
      producer, relay, consumer = (make_runner() for _ in range(3))
      args = frame_inputs(producer)
      producer.run(*args)
      packet = producer.last_warp
      relay.run(*args, shared_warp=packet)
      self.assertIs(relay.last_warp, packet)
      self.assertIs(relay.last_warp.producer(), producer)
      if invalidate == "advance":
        producer.run(*args)
      else:
        producer._reset_state()
      with self.subTest(invalidate=invalidate), self.assertRaisesRegex(ValueError, "stale"):
        consumer.run(*args, shared_warp=relay.last_warp)
      consumer.run_policy.assert_not_called()
      consumer.warp_enqueue.assert_not_called()

  def test_history_in_warp_cannot_share_but_default_path_is_preserved(self):
    artifact = artifact_fixture()
    artifact["image_history_pipeline"] = "warp"
    consumer = make_runner(artifact)
    producer = make_runner()
    args = frame_inputs(producer)
    producer.run(*args)
    with self.assertRaisesRegex(ValueError, "history"):
      consumer.run(*args, shared_warp=producer.last_warp)
    consumer.warp_enqueue.return_value = ("road-history", "wide-history")
    consumer.run(*args)
    consumer.warp_enqueue.assert_called_once()
    self.assertEqual(consumer.run_policy.call_args.kwargs["img"], "road-history")
    self.assertEqual(consumer.run_policy.call_args.kwargs["big_img"], "wide-history")
    self.assertIsNone(consumer.last_warp)