modeld_v2: add missing dynamic unpickler helpers

openpilot/sunnypilot/modeld_v2/helpers.py

@@ -0,0 +1,117 @@
+"""
+Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors.
+
+This file is part of sunnypilot and is licensed under the MIT License.
+See the LICENSE.md file in the root directory for more details.
+"""
+
+import io
+import struct
+import pickle
+import inspect
+import importlib
+import tempfile
+import shutil
+import enum
+
+
+def _patch_system_flock_acquire():
+  try:
+    from tinygrad.runtime.support.system import System
+    original_flock_acquire = System.flock_acquire
+    acquired_locks: dict[str, int] = {}
+
+    def flock_acquire(name: str) -> int:
+      if name in acquired_locks:
+        return acquired_locks[name]
+      lock_file_descriptor = original_flock_acquire(name)
+      acquired_locks[name] = lock_file_descriptor
+      return lock_file_descriptor
+    System.flock_acquire = flock_acquire
+  except (ImportError, AttributeError):
+    pass
+_patch_system_flock_acquire()
+
+def _pad_args(func, args, kwargs):
+  try:
+    sig = inspect.signature(func)
+  except Exception:
+    return args, kwargs
+  params = list(sig.parameters.values())
+  if inspect.isfunction(func) and params and params[0].name in ('cls', 'self'):
+    params = params[1:]
+
+  new_args = list(args)
+  has_varargs = any(p.kind == inspect.Parameter.VAR_POSITIONAL for p in params)
+  if len(new_args) > len(params) and not has_varargs:
+    new_args = new_args[:len(params)]
+
+  for i in range(len(new_args), len(params)):
+    param = params[i]
+    if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
+      continue
+    val = param.default if param.default is not inspect.Parameter.empty else None
+    new_args.append(val)
+  return new_args, kwargs
+
+
+def _dynamic_factory(real_class):
+  if isinstance(real_class, type) and issubclass(real_class, enum.Enum):
+    return real_class
+
+  def factory(*args, **kwargs):
+    try:
+      return real_class(*args, **kwargs)
+    except TypeError:
+      new_args, new_kwargs = _pad_args(real_class, args, kwargs)
+      return real_class(*new_args, **new_kwargs)
+
+  class DynamicMeta(type(real_class)):
+    def __call__(cls, *args, **kwargs):
+      return factory(*args, **kwargs)
+
+  class DynamicProxy(real_class, metaclass=DynamicMeta):
+    __slots__ = ()
+
+    def __new__(cls, *args, **kwargs):
+      return factory(*args, **kwargs)
+
+  DynamicProxy.__name__ = real_class.__name__
+  DynamicProxy.__module__ = real_class.__module__
+  return DynamicProxy
+
+
+class DynamicTinygradUnpickler(pickle.Unpickler):
+  def find_class(self, module, name):
+    real_class = getattr(importlib.import_module(module), name)
+    if module.startswith("tinygrad"):
+      return _dynamic_factory(real_class)
+    return real_class
+
+
+def load_oob(f):
+  opcodes = f.read(struct.unpack('<q', f.read(8))[0])
+  def buffers():
+    while (h := f.read(8)):
+      pb = pickle.PickleBuffer(bytearray(struct.unpack('<q', h)[0]))
+      f.readinto(pb)
+      yield pb
+  return DynamicTinygradUnpickler(io.BytesIO(opcodes), buffers=buffers()).load()
+
+
+def dump_oob(obj, f):
+  with tempfile.TemporaryFile(dir=".") as tmp:
+    def buffer_cb(buffer):
+      data = buffer.raw()
+      tmp.write(struct.pack('<q', len(data)))
+      tmp.write(data)
+      buffer.release()
+      return False
+
+    stream = io.BytesIO()
+    pickle.Pickler(stream, protocol=5, buffer_callback=buffer_cb).dump(obj)
+    opcodes = stream.getvalue()
+    f.write(struct.pack('<q', len(opcodes)))
+    f.write(opcodes)
+    tmp.seek(0)
+    shutil.copyfileobj(tmp, f)

openpilot/sunnypilot/modeld_v2/modeld.py

@@ -16,7 +16,7 @@ from setproctitle import setproctitle
 
 import openpilot.cereal.messaging as messaging
 from openpilot.common.hardware import COMMA_HARDWARE
-from openpilot.selfdrive.modeld.helpers import chestnut_present, load_oob
+from openpilot.selfdrive.modeld.helpers import chestnut_present
 from openpilot.cereal import log
 from opendbc.car.structs import car
 from openpilot.cereal.services import SERVICE_LIST
@@ -41,6 +41,7 @@ from openpilot.sunnypilot.modeld_v2.meta_helper import load_meta_constants
 from openpilot.sunnypilot.modeld_v2.camera_offset_helper import CameraOffsetHelper
 from openpilot.sunnypilot.livedelay.helpers import get_lat_delay
 from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase
+from openpilot.sunnypilot.modeld_v2.helpers import load_oob
 from openpilot.sunnypilot.modeld_v2.model_adapters import get_model_adapter
 from openpilot.sunnypilot.models.helpers import get_active_bundle
 from openpilot.sunnypilot.selfdrive.controls.lib.relc import RoadEdgeLaneChangeController

openpilot/sunnypilot/modeld_v2/tests/test_models.py

@@ -0,0 +1,77 @@
+"""
+Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors.
+
+This file is part of sunnypilot and is licensed under the MIT License.
+See the LICENSE.md file in the root directory for more details.
+"""
+
+import os
+import unittest
+from unittest.mock import patch
+from openpilot.common.file_chunker import open_file_chunked
+from openpilot.sunnypilot.modeld_v2.helpers import load_oob
+from tinygrad.device import Device
+
+
+class TestLegacyModels(unittest.TestCase):
+  def test_legacy_model_load(self):
+    base_name = os.environ.get("MODEL_BASE_NAME")
+    if not base_name:
+      raise unittest.SkipTest("MODEL_BASE_NAME env var not set, skipping integration test.")
+    chunk_dir = os.environ.get("MODEL_CHUNK_DIR", "/tmp/model_chunks")
+    base_path = os.path.join(chunk_dir, base_name)
+
+    try:
+      f = open_file_chunked(base_path)
+    except Exception as error:
+      self.fail(f"Failed to open chunked file {base_path}: {error}")
+    self.addCleanup(f.close)
+
+    real_getitem = Device.__class__.__getitem__
+
+    def safe_getitem(device_self, ix):
+      if ix == "QCOM" and not os.path.exists("/dev/kgsl-3d0"):
+        return real_getitem(device_self, "CPU")
+      if ix == "AMD" and not os.path.exists("/dev/kfd"):
+        return real_getitem(device_self, "CPU")
+      return real_getitem(device_self, ix)
+
+    with patch.object(Device.__class__, "__getitem__", safe_getitem):
+      obj = load_oob(f)
+
+      assert isinstance(obj, dict), "Parsed object is not a dictionary"
+      assert "metadata" in obj, "Metadata key is missing"
+
+      def traverse_jits(d):
+        if hasattr(d, "captured") and d.captured is not None:
+          linear_attr = getattr(d.captured, "_linear", None)
+          if linear_attr is None:
+            try:
+              linear_attr = getattr(d.captured, "linear", None)
+            except AssertionError as e:
+              if "alloc size must be positive" in str(e):
+                return
+              raise
+
+          if linear_attr is not None:
+            stack = [linear_attr]
+            visited = set()
+            while stack:
+              node = stack.pop()
+              if id(node) in visited:
+                continue
+              visited.add(id(node))
+              if hasattr(node, "src") and node.src is not None:
+                for child in node.src:
+                  stack.append(child)
+        elif isinstance(d, dict):
+          for v in d.values():
+            traverse_jits(v)
+        elif isinstance(d, (list, tuple)):
+          for v in d:
+            traverse_jits(v)
+
+      if "run_model" in obj:
+        traverse_jits(obj["run_model"])
+      if "run_policy" in obj:
+        traverse_jits(obj["run_policy"])