models: test tinygrad concurrency

.github/workflows/test_models.yml

@@ -0,0 +1,86 @@
+name: Test Models Compatibility With New Tinygrad
+on:
+  push:
+    paths:
+      - 'tinygrad_repo/**'
+      - '.github/workflows/test_models.yml'
+      - 'openpilot/sunnypilot/modeld_v2/**'
+  pull_request:
+    paths:
+      - 'tinygrad_repo/**'
+      - '.github/workflows/test_models.yml'
+      - 'openpilot/sunnypilot/modeld_v2/**'
+  workflow_dispatch:
+
+jobs:
+  generate-matrix:
+    runs-on: ubuntu-latest
+    outputs:
+      models: ${{ steps.set-matrix.outputs.models }}
+    steps:
+      - uses: actions/checkout@v4
+      - name: Fetch and Parse json
+        id: set-matrix
+        run: |
+          python3 -c '
+          import json, urllib.request, sys, os, re
+
+          with open("openpilot/sunnypilot/models/fetcher.py", "r") as f:
+            content = f.read()
+            match = re.search(r"MODEL_URL\s*=\s*[\"'"'"']([^\"'"'"']+)[\"'"'"']", content)
+            if not match:
+              raise ValueError("MODEL_URL not found in fetcher.py")
+            url = match.group(1)
+
+          req = urllib.request.urlopen(url)
+          data = json.loads(req.read())
+
+          artifacts = []
+          for bundle in data.get("bundles", []):
+            for model in bundle.get("models", []):
+              if "artifact" in model:
+                artifacts.append(model["artifact"])
+
+          with open(os.environ["GITHUB_OUTPUT"], "a") as f:
+            f.write(f"models={json.dumps(artifacts)}\n")
+          '
+
+  test-model:
+    name: Test ${{ matrix.artifact.file_name }}
+    needs: generate-matrix
+    runs-on: [self-hosted, tici]
+    strategy:
+      fail-fast: false
+      matrix:
+        artifact: ${{ fromJson(needs.generate-matrix.outputs.models) }}
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          submodules: true
+
+      - name: Download Model Chunks in Parallel
+        run: |
+          mkdir -p /tmp/model_chunks
+          echo '${{ toJson(matrix.artifact.chunks) }}' > chunks.json
+
+          BASE_URL="${{ matrix.artifact.download_uri.url }}"
+          BASE_DIR=$(dirname "$BASE_URL")
+
+          python3 -c '
+          import json, os
+          with open("chunks.json") as f:
+            chunks = json.load(f)
+          manifest_path = f"/tmp/model_chunks/${{ matrix.artifact.file_name }}.chunkmanifest"
+          with open(manifest_path, "w") as f:
+            f.write(str(len(chunks)))
+          for c in chunks:
+            print(c["file_name"])
+          ' | xargs -I {} -P 10 bash -c "echo 'Downloading {}...' && curl -s -L -o /tmp/model_chunks/{} $BASE_DIR/{}"
+
+      - name: Run Model Compatibility Test
+        env:
+          MODEL_BASE_NAME: ${{ matrix.artifact.file_name }}
+          MODEL_CHUNK_DIR: "/tmp/model_chunks"
+          PYTHONPATH: ".:./tinygrad_repo"
+        run: |
+          python3 -m pytest openpilot/sunnypilot/modeld_v2/tests/test_models.py

openpilot/sunnypilot/modeld_v2/helpers.py

@@ -0,0 +1,93 @@
+import io
+import struct
+import pickle
+import inspect
+import importlib
+import enum
+
+
+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 _enum_factory(enum_class):
+  def factory(*args, **kwargs):
+    try:
+      return enum_class(*args, **kwargs)
+    # OptOps and UOp objects in the .pkl are left over from the compilation phase,
+    # reassignment does nothing because they aren't tied to the execution graph
+    # It never executes or evaluates the UOp nodes again.
+    except ValueError:
+      return list(enum_class)[0]
+  factory.__name__ = enum_class.__name__
+  factory.__module__ = enum_class.__module__
+  return factory
+
+
+def _dynamic_factory(real_class):
+  if isinstance(real_class, type) and issubclass(real_class, enum.Enum):
+    return _enum_factory(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):
+    if module == "tinygrad.ops":
+      try:
+        importlib.import_module("tinygrad.uops")
+        module = "tinygrad.uops"
+      except ImportError:
+        pass
+    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()

openpilot/sunnypilot/modeld_v2/modeld.py

@@ -17,7 +17,7 @@ from tinygrad.tensor import Tensor
 
 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
@@ -52,6 +52,7 @@ from openpilot.sunnypilot.modeld_v2.compile_modeld import (derive_frame_skip, ma
                                                            WARP_INPUTS, POLICY_INPUTS)
 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.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,22 @@
+import os
+import pytest
+from openpilot.common.file_chunker import open_file_chunked
+from openpilot.sunnypilot.modeld_v2.helpers import load_oob
+
+
+class TestLegacyModels:
+  def test_legacy_model_load(self):
+    base_name = os.environ.get("MODEL_BASE_NAME")
+    if not base_name:
+      pytest.skip("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 e:
+      pytest.fail(f"Failed to open chunked file {base_path}: {e}")
+
+    obj = load_oob(f)
+    assert isinstance(obj, dict), "Parsed object is not a dictionary"
+    assert 'metadata' in obj, "Metadata key is missing"