Optimize FLM job encoding

openpilot/starpilot/flm/controller_replay.py

@@ -1513,22 +1513,36 @@ def pack_primitive_snapshot(value):
   import pickle
   import zlib
 
-  def primitive(item):
-    if item is None or type(item) in (bool, int, float, str, bytes):
-      if type(item) is float and not math.isfinite(item):
+  # Pickle preserves shared primitive containers. Validate each actual
+  # container once rather than expanding a shared source graph repeatedly.
+  # This memo is local to this call; it never caches evidence or permissions.
+  checked = set()
+  active = set()
+
+  def primitive(item, recurse):
+    kind = type(item)
+    if item is None or kind in (bool, int, float, str, bytes):
+      if kind is float and not math.isfinite(item):
         raise ValueError('nonfinite job primitive')
       return
-    if type(item) in (tuple, list):
-      for child in item:
-        primitive(child)
+    if kind not in (tuple, list, dict):
+      raise ValueError('fit job requires primitives')
+    identity = id(item)
+    if identity in active:
+      raise ValueError('cyclic fit job primitive')
+    if identity in checked:
       return
-    if type(item) is dict and all(type(key) is str for key in item):
-      for child in item.values():
-        primitive(child)
-      return
-    raise ValueError('fit job requires primitives')
+    if kind is dict and any(type(key) is not str for key in item):
+      raise ValueError('fit job requires primitives')
+    active.add(identity)
+    try:
+      for child in (item.values() if kind is dict else item):
+        recurse(child, recurse)
+    finally:
+      active.remove(identity)
+    checked.add(identity)
 
-  primitive(value)
+  primitive(value, primitive)
   raw = pickle.dumps(value, protocol=5)
   if len(raw) > MAX_FIT_RAW_BYTES:
     raise ValueError('fit job raw byte bound')

openpilot/starpilot/flm/tests/test_fit_wire_sharing.py

@@ -0,0 +1,54 @@
+"""Bound primitive validation work without changing the worker wire."""
+import math
+import pickle
+import zlib
+from types import SimpleNamespace
+
+import pytest
+
+from openpilot.starpilot.flm import controller_replay as wire
+
+
+def test_shared_primitives_keep_exact_wire_types_and_signed_zero():
+  row = {'negative_zero': -0., 'integer': 1, 'boolean': True, 'bytes': b'value'}
+  value = {'rows': [row] * 8192}
+  encoded = wire.pack_primitive_snapshot(value)
+  assert encoded == zlib.compress(pickle.dumps(value, protocol=5), 1)
+  restored = wire.unpack_primitive_snapshot(encoded)['rows']
+  assert len(restored) == 8192 and restored[0] is restored[-1]
+  assert math.copysign(1., restored[0]['negative_zero']) == -1.
+  assert type(restored[0]['integer']) is int and type(restored[0]['boolean']) is bool
+  assert restored[0]['bytes'] == b'value'
+
+
+def test_shared_container_validation_is_bounded(monkeypatch):
+  checked = []
+
+  def finite(value):
+    checked.append(value)
+    return math.isfinite(value)
+
+  monkeypatch.setattr(wire, 'math', SimpleNamespace(isfinite=finite))
+  row = {'value': .25}
+  wire.pack_primitive_snapshot({'rows': [row] * 8192})
+  assert checked == [.25]
+
+
+@pytest.mark.parametrize('bad', (object(), {'value': float('nan')}, {'value': float('inf')}, {1: True}))
+def test_invalid_shared_data_still_rejects(bad):
+  with pytest.raises(ValueError):
+    wire.pack_primitive_snapshot([bad, bad])
+
+
+def test_cycle_rejects_without_recursing_to_the_interpreter_limit():
+  value = []
+  value.append(value)
+  with pytest.raises(ValueError, match='cyclic fit job primitive'):
+    wire.pack_primitive_snapshot(value)
+
+
+@pytest.mark.parametrize('limit,message', (('MAX_FIT_RAW_BYTES', 'raw byte bound'), ('MAX_FIT_SNAPSHOT_BYTES', 'IPC byte bound')))
+def test_original_wire_size_limits_remain_enforced(monkeypatch, limit, message):
+  monkeypatch.setattr(wire, limit, 1)
+  with pytest.raises(ValueError, match=message):
+    wire.pack_primitive_snapshot({'value': b'payload'})

prebuilt.json

@@ -1 +1 @@
-{"protected_inventory_sha256": "a30e5882f8ffbce74d36b82d4eee17cfa2ff5e7f37ec36a29ce066183d581765", "schema_version": 1}
+{"protected_inventory_sha256": "9a6c99695cbbbfc91dcd112a64a47937f35f40f690f4ec1cd34073d8cdbd2db6", "schema_version": 1}