models: publish the download speed and derive the eta from it (#2052)

* models: publish the download speed and derive the eta from it

Fills in DownloadProgress.speed (bytes per second, declared with the parallel
downloader). The reporter measures the rate between its 0.5 s reports, smooths
it, and derives the eta from the bytes remaining at that rate instead of
extrapolating elapsed time. Speed is zeroed when a download ends, fails or is
cancelled and for a cached artifact, and copied across artifacts that share a
file like the other progress fields.

The big UI's Models page shows "12.3 MB/s" next to the progress bar and the
mici models panel shows it beside the percentage. Both take the max across a
bundle's artifacts, not the sum, since artifacts sharing a file mirror the same
transfer and only one file is ever in flight.

A stall decays the smoothed speed toward 0 without ever reaching it, so below
0.05 MB/s, which would only display as "0.0 MB/s", it snaps to 0. Both UIs
then show "--.- MB/s" while downloading with no rate (before the first reading
or during a stall) instead of a stuck "0.0 MB/s" or nothing, and the speed
recovers straight to the real rate once data flows again. The snap also keeps
the eta in range: without it a stall of about 25 s drove remaining / speed past
the capnp UInt32 eta, and the capnp error failed the download outright instead
of leaving the stall to the read timeout and retries. A test holds a piece
stalled until the published speed reaches exactly 0, and errors on that
overflow without the snap.

Assisted-by: Claude Code (Claude Opus 4.8)
Assisted-by: Claude Code (Claude Opus 5.5)

* remove unneeded comments

Assisted-by: Claude Code (Claude Opus 5.5)

* models: take the speed baseline before the workers start

The reporter took its starting byte count after the piece workers were
submitted, and starting each pool thread lets earlier workers run, so
bytes landed in that gap never counted toward a speed sample. Take the
baseline in _download_file before any worker exists and pass it in.

The test runs every piece inline in submit, so the whole file lands
before the reporter starts; with the late baseline the speed stays 0.

Assisted-by: Claude Code (Claude Opus 5.5)

.github/workflows/test_models.yml

@@ -3,62 +3,9 @@ on:
   pull_request:
     paths:
       - 'tinygrad_repo'
-      - 'openpilot/sunnypilot/modeld_v2/**'
   workflow_dispatch:
 
-env:
-  CI: 1
-  GIT_CONFIG_COUNT: 1
-  GIT_CONFIG_KEY_0: lfs.fetchexclude
-  GIT_CONFIG_VALUE_0: openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl
-
 jobs:
-  plain-artifact:
-    name: Plain artifact round trip and pickle converter
-    runs-on: ubuntu-24.04
-    env:
-      PYTHONPATH: ${{ github.workspace }}/tinygrad_repo:${{ github.workspace }}
-      TINYGRAD_TESTS: 1
-    steps:
-      - uses: actions/checkout@v7
-      - run: ./tools/op.sh setup
-      - name: Build openpilot
-        run: scons -j$(nproc)
-      - name: Compile, load, cache and fuse a model with the pinned tinygrad
-        env:
-          DEV: CPU
-        run: tools/op.sh test -v openpilot/sunnypilot/modeld_v2/tests/test_plain_artifact.py
-      - name: Install LLVM 20 and ROCm comgr for the amd mock, as tinygrad_repo/.github/actions/setup-tinygrad does
-        run: |
-          . /etc/os-release
-          curl -fsSL https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc > /dev/null
-          echo "deb http://apt.llvm.org/$VERSION_CODENAME/ llvm-toolchain-$VERSION_CODENAME-20 main" | sudo tee /etc/apt/sources.list.d/llvm.list
-          sudo mkdir -p /etc/apt/keyrings
-          curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
-          echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/7.1 $VERSION_CODENAME main" | sudo tee /etc/apt/sources.list.d/rocm.list
-          printf 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600\n' | sudo tee /etc/apt/preferences.d/rocm-pin-600
-          sudo apt-get -qq update && sudo apt-get -qq install -y --no-install-recommends libllvm20 clang-20 lld-20 comgr
-          printf '/opt/rocm/lib\n/opt/rocm/lib64\n' | sudo tee /etc/ld.so.conf.d/rocm.conf && sudo ldconfig
-          echo "/usr/lib/llvm-20/bin" >> "$GITHUB_PATH"
-      - name: Convert a pickled USB+AMD jit and replay it under tinygrad's amd mock
-        env:
-          DEV: MOCKUSB+AMD
-          CHECK_OOB: 0
-          GMMU: 0
-        run: tools/op.sh test -v openpilot/sunnypilot/modeld_v2/tests/test_convert_upstream_pkl.py
-      - name: Where tinygrad moved, for the helper or failure point named above
-        if: failure() && github.event_name == 'pull_request'
-        run: |
-          git fetch -q --depth=1 origin "${{ github.base_ref }}"
-          old=$(git rev-parse "origin/${{ github.base_ref }}:tinygrad_repo") new=$(git rev-parse HEAD:tinygrad_repo)
-          git -C tinygrad_repo fetch -q origin "$old" || true
-          echo "tinygrad_repo $old -> $new, in what helpers.py and plain_artifact_compiler.py read:"
-          git -C tinygrad_repo diff "$old" "$new" -- tinygrad/uop/ops.py tinygrad/engine/jit.py tinygrad/engine/realize.py tinygrad/device.py \
-            tinygrad/helpers.py tinygrad/dtype.py tinygrad/renderer/__init__.py
-          echo "and in what convert_upstream_pkl.py decodes:"
-          git -C tinygrad_repo diff --stat "$old" "$new" -- tinygrad/runtime/autogen/am/pm4_nv.py tinygrad/runtime/autogen/amd_gpu.py
-          git -C tinygrad_repo diff "$old" "$new" -- tinygrad/runtime/support/hcq2.py tinygrad/runtime/support/usb.py tinygrad/runtime/ops_amd.py
-
   generate-matrix:
     runs-on: ubuntu-latest
     outputs:

openpilot/selfdrive/ui/sunnypilot/layouts/settings/models.py

@@ -243,7 +243,9 @@ class ModelsLayout(Widget):
     if ds.verifying in statuses:
       return {"name": name, "downloading": True, "progress": progress, "status_text": tr("verifying")}
     if ds.downloading in statuses:
-      return {"name": name, "downloading": True, "progress": progress}
+      speed = max((p.speed for p in progresses), default=0.0)
+      status_text = f"{speed / 1024 ** 2:.1f} MB/s" if speed > 0 else "--.- MB/s"
+      return {"name": name, "downloading": True, "progress": progress, "status_text": status_text}
     if statuses <= {ds.downloaded, ds.cached}:
       return {"name": name, "text_color": ON_COLOR, "icon": "icons/checkmark.png"}
     # circled_slash is authored grey; tinting it again only darkens it

openpilot/selfdrive/ui/sunnypilot/mici/layouts/models.py

@@ -257,6 +257,7 @@ class ModelsLayoutMici(NavScroller):
       count = 0
       verifying = False
       canceling = False
+      speed = 0.0
       for model in manager.selectedBundle.models:
         count += 1
         p = model.artifact.downloadProgress
@@ -265,6 +266,7 @@ class ModelsLayoutMici(NavScroller):
           progress += p.progress
           verifying = verifying or p.status == custom.ModelManagerSP.DownloadStatus.verifying
           canceling = canceling or p.canceling
+          speed = max(speed, p.speed)  # only one file downloads at a time, so the non-zero speed is the live one
         elif p.status in (custom.ModelManagerSP.DownloadStatus.downloaded,
                           custom.ModelManagerSP.DownloadStatus.cached):
           progress += 100.0
@@ -279,7 +281,10 @@ class ModelsLayoutMici(NavScroller):
       self.current_model_info.current_model_text.set_text(name_text)
       self.current_model_info.info_header.set_text(tr("progress") + self._download_progress)
       self.current_model_info.info_header._shimmer = True
-      self.current_model_info.info_text.set_text(f"{progress/count:.2f}%")
+      progress_text = f"{progress/count:.2f}%"
+      if not verifying and not canceling:
+        progress_text += f"  |  {speed / 1024 ** 2:.1f} MB/s" if speed > 0 else "  |  --.- MB/s"
+      self.current_model_info.info_text.set_text(progress_text)
 
     elif manager.selectedBundle and manager.selectedBundle.status == custom.ModelManagerSP.DownloadStatus.downloaded:
       self.current_model_info.info_header.set_text(tr("downloaded"))

openpilot/sunnypilot/modeld_v2/convert_upstream_pkl.py

@@ -7,7 +7,6 @@ See the LICENSE.md file in the root directory for more details.
 """
 import argparse
 import dataclasses
-import inspect
 import itertools
 import mmap
 import re
@@ -46,11 +45,6 @@ def _address(v):
   return _root(v.src[0]), add + start, size
 
 
-def _buffer(u):
-  from tinygrad.uop.ops import Ops
-  return u.buffer if u.op is Ops.BUFFER else getattr(u.arg, "buffer", None)
-
-
 def _patched(buf, topo) -> tuple[bytearray, dict]:
   from tinygrad.uop.ops import Ops
   blob, rows = bytearray(buf.max_numel() * buf.dtype.itemsize), {}
@@ -78,7 +72,7 @@ def _evaluate(u, env):
     table, index = _root(u.src[0].src[0]), _evaluate(u.src[0].src[1], env)
     return struct.unpack_from(u.dtype.fmt, table.arg, index * u.dtype.itemsize)[0] if table.op is Ops.BINARY and index is not None else None
   if u.op in GroupOp.ALU and None not in (values := [_evaluate(x, env) for x in u.src]):
-    return exec_alu(u.op, u.dtype, values, *([False] if "truncate_output" in inspect.signature(exec_alu).parameters else []))
+    return exec_alu(u.op, u.dtype, values, False)
   return None
 
 
@@ -118,7 +112,7 @@ def convert(jit, kwargs: dict, frame: dict, target) -> tuple[dict, memoryview]:
   from tinygrad.uop.ops import Ops
   submit = jit.captured._linear.src[0].without_after
   kernels, topo = submit.arg.aux.kernels, list(jit.captured._linear.toposort(enter_calls=True))
-  allocs = {str(name): u for u in topo if u.op is Ops.ALLOC and (name := getattr(u, "tag", None) or u.arg.name) is not None}
+  allocs = {str(u.tag): u for u in topo if u.op is Ops.ALLOC and u.tag is not None}
   (commands, patched), kernargs = _patched(allocs["cmdbuf_compute_0"], topo), allocs["kernargs_compute_0"]
   argument_words, pokes = _patched(kernargs, topo)[1], _input_pokes(submit, kernargs)
   sh = {name: getattr(amd_gpu, f"reg{name}") + amd_gpu.GC_BASE__INST0_SEG0 - pm4.PACKET3_SET_SH_REG_START
@@ -164,7 +158,7 @@ def convert(jit, kwargs: dict, frame: dict, target) -> tuple[dict, memoryview]:
         args.append({"input": pokes[slot][0]})
       elif slot in argument_words:
         base, at, nbytes = _address(argument_words[slot])
-        args.append({"buffer": buffers.setdefault(base, len(buffers)), "offset": at, "nbytes": nbytes or base.max_numel() * base.dtype.itemsize, "dtype": u8})
+        args.append({"buffer": buffers.setdefault(base, len(buffers)), "offset": at, "nbytes": nbytes or base.arg.buffer.nbytes, "dtype": u8})
     outs, ins = kernel[6]
     calls.append({"kind": "program", "name": re.sub(r"\x1b\[[0-9;]*m", "", kernel[1]), "lib": programs[image_alloc], "target": dataclasses.asdict(target),
                   "signature": [[f"data{i}", i, u8, [a.get("nbytes", 1)]] for i, a in enumerate(args)], "global_size": grid, "local_size": local,
@@ -172,13 +166,12 @@ def convert(jit, kwargs: dict, frame: dict, target) -> tuple[dict, memoryview]:
     needs_init |= {args[i]["buffer"] for i in ins if "buffer" in args[i] and args[i]["buffer"] not in written}
     written |= {args[i]["buffer"] for i in outs if "buffer" in args[i]}
 
-  assert all(_buffer(b) is not None for b, i in buffers.items() if i in needs_init), "a kernel reads a scratch before any kernel writes it"
   names = list(jit.captured.expected_names)
   inputs = [{"slot": slot, "name": names[slot] if slot < len(names) else None, "dtype": repr(u.dtype), "numel": int(u.max_numel()), "device": u.device}
             for slot, u in enumerate(_prepare_jit_inputs((), kwargs)[0])]
   header = {"input_names": names, "calls": calls, "returns": {"kind": "NoneType", "tensors": []}, "inputs": inputs,
-            "buffers": [{"nbytes": bb.nbytes if (bb := _buffer(b)) else b.max_numel() * b.dtype.itemsize, "device": bb.device if bb else b.device, "dtype": u8}
-                        | ({"init": blob.add(bb.ensure_allocated().as_memoryview())} if i in needs_init else {}) for b, i in buffers.items()],
+            "buffers": [{"nbytes": b.arg.buffer.nbytes, "device": b.arg.buffer.device, "dtype": u8}
+                        | ({"init": blob.add(b.arg.buffer.ensure_allocated().as_memoryview())} if i in needs_init else {}) for b, i in buffers.items()],
             "check": {"inputs": [[path, store(array), device] for path, array, device in frame["inputs"]],
                       "outputs": [store(array) for array in frame["outputs"]], "output_container": frame["output_container"]}}
   with sink:

openpilot/sunnypilot/modeld_v2/helpers.py

@@ -250,12 +250,6 @@ def _allocate_buffers(header: dict, blob: memoryview, device_map: dict | None, s
   return buffers
 
 
-def _fields(what, /, **values):  # a tinygrad dataclass built or replaced with the fields it has today: a dropped field is fine, a new required one is not
-  known = {field.name for field in dataclasses.fields(what)}
-  values = {name: value for name, value in values.items() if name in known}
-  return what(**values) if isinstance(what, type) else dataclasses.replace(what, **values)
-
-
 def _program_uop(call: dict, lib: bytes):
   from tinygrad.helpers import Target
   from tinygrad.uop.ops import KernelInfo, Ops, ProgramInfo, UOp
@@ -263,9 +257,9 @@ def _program_uop(call: dict, lib: bytes):
   for name, compact_slot, dtype, shape in call["signature"]:
     parameter = UOp.param(call["globals"][compact_slot], _dtype(dtype), math.prod(shape))
     image = (shape[0], shape[1]) if len(shape) == 3 and shape[2] == 4 else None
-    parameters.append(parameter.replace(arg=_fields(parameter.arg, name=name, image=image)))
-  info = _fields(ProgramInfo, global_size=tuple(call["global_size"]), local_size=tuple(call["local_size"]), globals=tuple(call["globals"]),
-                 outs=tuple(call["outs"]), ins=tuple(call["ins"]), target=_fields(Target, **call["target"]))
+    parameters.append(parameter.replace(arg=dataclasses.replace(parameter.arg, name=name, image=image)))
+  info = ProgramInfo(global_size=tuple(call["global_size"]), local_size=tuple(call["local_size"]), vars=(), globals=tuple(call["globals"]),
+                     outs=tuple(call["outs"]), ins=tuple(call["ins"]), target=Target(**call["target"]))
   return UOp(Ops.PROGRAM, src=(UOp(Ops.SINK, arg=KernelInfo(name=call["name"])), UOp(Ops.LINEAR, src=tuple(parameters)),
                                UOp(Ops.SOURCE, arg=""), UOp(Ops.BINARY, arg=lib)), arg=info)
 

openpilot/sunnypilot/modeld_v2/plain_artifact_compiler.py

@@ -23,7 +23,7 @@ def record_kernel_linear():
   recorded, hcq_compile, compile_linear = [], realize.hcq_compile, jit.compile_linear
 
   def recording_compile_linear(*args, **kwargs):
-    with mock.patch.object(realize, "hcq_compile", lambda linear, *a, **k: (linear.src and recorded.append(linear), hcq_compile(linear, *a, **k))[1]):
+    with mock.patch.object(realize, "hcq_compile", lambda linear, *a, **k: (recorded.append(linear), hcq_compile(linear, *a, **k))[1]):
       return compile_linear(*args, **kwargs)
 
   with mock.patch.object(jit, "compile_linear", recording_compile_linear):
@@ -38,18 +38,6 @@ def capture_frame(jit, kwargs: dict, output_container: str | None = None):
   return {"inputs": inputs, "outputs": [tensor.numpy().copy() for tensor in outputs], "output_container": output_container}, result
 
 
-def kernel_signature(ast, elf) -> list:
-  from tinygrad.dtype import AddrSpace
-  if hasattr(ast, "kernel_params"):
-    if any(p.arg.addrspace is AddrSpace.ALU for p in ast.kernel_params):
-      raise ValueError(f"{elf.name} takes a variable, which plain artifacts don't support")
-    compact = {slot: i for i, slot in enumerate(ast.arg.globals)}
-    return [[p.arg.name, compact[p.arg.slot], repr(p.dtype), [int(x) for x in p._shape]] for p in ast.kernel_params]
-  if not all(isinstance(slot, int) for _, slot, _, _ in elf.signature):
-    raise ValueError("TinyELF.signature lost its slots and this tinygrad has no UOp.kernel_params: update kernel_signature")
-  return [[name, slot, repr(dtype), [int(x) for x in shape]] for name, slot, dtype, shape in elf.signature]
-
-
 def serialize_jit(jit, kernel_linear, check: dict | None = None) -> tuple[dict, memoryview]:
   from tinygrad.uop.ops import Ops
   sink = tempfile.TemporaryFile(dir=".")
@@ -63,12 +51,6 @@ def serialize_jit(jit, kernel_linear, check: dict | None = None) -> tuple[dict,
     if uop.op is Ops.PARAM:
       input_specs.setdefault(uop.arg.slot, {"dtype": repr(uop.dtype), "numel": int(uop.max_numel()), "device": uop.device})
       return {"input": uop.arg.slot}
-    root, offset = uop, 0
-    while root.op in (Ops.AFTER, Ops.BITCAST, Ops.SHRINK):
-      offset, root = offset + (root.src[1].arg * root.src[0].dtype.itemsize if root.op is Ops.SHRINK else 0), root.src[0]
-    if root.op is Ops.ALLOC:
-      base = bases.setdefault(id(root), (len(bases), root, root.dtype))
-      return {"buffer": base[0], "offset": offset, "nbytes": int(uop.max_numel()) * uop.dtype.itemsize, "dtype": repr(uop.dtype)}
     buffer = uop.buffer
     root = buffer._base if buffer._base is not None else buffer
     base = bases.setdefault(id(root), (len(bases), root, uop.base.dtype if buffer._base is None else uop.dtype))
@@ -81,7 +63,7 @@ def serialize_jit(jit, kernel_linear, check: dict | None = None) -> tuple[dict,
       elf, info = ast.to_elf(), ast.arg
       global_size, local_size = info.launch_dims({})
       calls.append({"kind": "program", "name": elf.name, "lib": blob.add(bytes(elf.lib)), "target": dataclasses.asdict(elf.target),
-                    "signature": kernel_signature(ast, elf),
+                    "signature": [[name, slot, repr(dtype), [int(x) for x in shape]] for name, slot, dtype, shape in elf.signature],
                     "global_size": [int(x) for x in global_size], "local_size": [int(x) for x in local_size],
                     "globals": list(info.globals), "outs": list(info.outs), "ins": list(info.ins), "args": arguments})
       outs, ins = info.outs, info.ins
@@ -95,9 +77,7 @@ def serialize_jit(jit, kernel_linear, check: dict | None = None) -> tuple[dict,
 
   ret = jit.captured.ret
   returns = {"kind": type(ret).__name__, "tensors": [describe(t.uop.base) | {"shape": [int(x) for x in t.shape]} for t in flatten(ret)]}
-  from tinygrad.device import Buffer
-  assert not [base for index, base, _ in bases.values() if index in needs_init and not isinstance(base, Buffer)], "a kernel reads a scratch before writing it"
-  buffers = [{"nbytes": base.nbytes if isinstance(base, Buffer) else int(base.max_numel()) * dtype.itemsize, "device": base.device, "dtype": repr(dtype)}
+  buffers = [{"nbytes": base.nbytes, "device": base.device, "dtype": repr(dtype)}
              | ({"init": blob.add(base.as_memoryview())} if index in needs_init else {}) for index, base, dtype in bases.values()]
   names = list(jit.captured.expected_names)
   if check is not None:

openpilot/sunnypilot/modeld_v2/tests/test_convert_upstream_pkl.py

@@ -1,121 +0,0 @@
-"""
-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 re
-import sys
-import tempfile
-import unittest
-from unittest import mock
-import numpy as np
-
-os.environ.setdefault("CHECK_OOB", "0")
-os.environ.setdefault("GMMU", "0")
-if sys.platform == "darwin" and os.environ.get("DEV", "").startswith("MOCKUSB"):  # the mock's VRAM is a memfd, which macOS lacks: an unlinked file is the same
-  from tinygrad.runtime.autogen import libc
-
-  def _memfd_create(name, flags):
-    fd, path = tempfile.mkstemp()
-    os.unlink(path)
-    return fd
-  libc.memfd_create = _memfd_create
-
-from openpilot.common.test import OpenpilotTestCase
-from openpilot.sunnypilot.modeld_v2 import convert_upstream_pkl
-from openpilot.sunnypilot.modeld_v2.convert_upstream_pkl import convert
-from openpilot.sunnypilot.modeld_v2.helpers import PlainJit, load_model
-from openpilot.sunnypilot.modeld_v2.plain_artifact_compiler import capture_frame
-
-IMG, DESIRE, STATE, OUT = (6, 8, 8), (8,), (4, 16), (1, 32)
-
-
-class TestConvertUpstreamPkl(OpenpilotTestCase):
-  @classmethod
-  def setUpClass(cls):
-    super().setUpClass()
-    if not os.environ.get("TINYGRAD_TESTS"):
-      raise unittest.SkipTest("TINYGRAD_TESTS not set. skipping converter test")
-    from tinygrad import Device, Tensor
-    assert Device.DEFAULT == "AMD" and getattr(Device["AMD"], "is_usb", False), f"{Device.DEFAULT} is not an AMD device behind the USB bridge"
-    assert (Tensor([1, 2, 3]) * 2).tolist() == [2, 4, 6]
-    cls.tmp = tempfile.TemporaryDirectory()
-    cls.rng = np.random.default_rng(0)
-    cls.weights = Tensor(cls.rng.standard_normal((16, 32)).astype(np.float32)).realize()
-    cls.bias = Tensor(cls.rng.standard_normal((1, 32)).astype(np.float32)).realize()
-    cls.input_specs = {"input_imgs": (IMG, "uint8", Device.DEFAULT), "desire": (DESIRE, "float32", Device.DEFAULT),
-                       "features": (STATE, "float32", Device.DEFAULT)}
-    cls.output_specs = {"outputs": (OUT, "float32", Device.DEFAULT), "next_features": (STATE, "float32", Device.DEFAULT)}
-
-  @classmethod
-  def tearDownClass(cls):
-    cls.tmp.cleanup()
-    super().tearDownClass()
-
-  def model(self, input_imgs, desire, features, output_buffers):
-    from tinygrad import Tensor
-    feat = features.mean(axis=0).reshape(1, 16) + input_imgs.float().mean() + desire.sum()
-    Tensor.realize(output_buffers["outputs"].assign((feat @ self.weights + self.bias).relu()),
-                   output_buffers["next_features"].assign(features[1:].cat(feat * 2)))
-
-  def inputs(self, seed: int) -> dict:
-    from examples.openpilot.helpers import allocate_inputs
-    rng = np.random.default_rng(seed)
-
-    def initialize(arrays: dict):
-      for value in arrays.values():
-        value[...] = rng.standard_normal(value.shape) if value.dtype.kind == "f" else rng.integers(0, 256, value.shape)
-    return allocate_inputs(self.input_specs, initialize) | {"output_buffers": allocate_inputs(self.output_specs)}
-
-  def jit(self):
-    from tinygrad import TinyJit
-    run = TinyJit(self.model, prune=True)
-    kwargs = self.inputs(1)
-    run(**kwargs)
-    run(**kwargs)
-    return run
-
-  def assert_matches_eager(self, plain, seed: int):
-    fresh, expected = self.inputs(seed), self.inputs(seed)
-    plain(**fresh)
-    self.model(**expected)
-    for name in self.output_specs:
-      assert np.array_equal(fresh["output_buffers"][name].numpy(), expected["output_buffers"][name].numpy()), name
-
-  def convert(self, jit) -> tuple[dict, memoryview]:
-    from tinygrad import Device
-    kwargs = self.inputs(2)
-    frame, _ = capture_frame(jit, kwargs, "output_buffers")
-    header, blob = convert(jit, kwargs, frame, Device[Device.DEFAULT].renderer.target)
-    kernels = jit.captured._linear.src[0].without_after.arg.aux.kernels
-    assert [call["name"] for call in header["calls"]] == [re.sub(r"\x1b\[[0-9;]*m", "", kernel[1]) for kernel in kernels]
-    assert all(re.fullmatch(r"\w+", call["name"]) for call in header["calls"]), "kernel names: failure point 1"
-    assert [(tuple(call["outs"]), tuple(call["ins"])) for call in header["calls"]] == [tuple(kernel[6]) for kernel in kernels], "outs/ins: failure point 1"
-    bound = {argument["input"] for call in header["calls"] for argument in call["args"] if "input" in argument}
-    assert bound == {spec["slot"] for spec in header["inputs"]}, f"input slots bound by pokes {sorted(bound)}: failure point 4"
-    assert [spec["name"] for spec in header["inputs"]][:len(self.input_specs)] == sorted(self.input_specs)
-    return header, blob
-
-  def test_converted_jit_matches_eager(self):  # the submit decoded into plain calls, replayed against its own frame
-    jit = self.jit()
-    plain = PlainJit(*self.convert(jit))
-    self.assert_matches_eager(plain, 3)
-    self.assert_matches_eager(plain, 4)
-
-  def test_pickle_round_trip_through_main(self):
-    from examples.openpilot.helpers import dump_pickle, load_pickle
-    from tinygrad.engine.realize import lower_and_compile
-    pkl, out = os.path.join(self.tmp.name, "tiny_tinygrad.pkl"), os.path.join(self.tmp.name, "tiny_plain.pkl")
-    dump_pickle({"run": self.jit(), "input_specs": self.input_specs, "output_specs": self.output_specs, "metadata": {"name": "tiny"}}, pkl)
-
-    artifact = load_pickle(pkl, out_of_band=True)
-    artifact["run"].captured._linear = lower_and_compile(artifact["run"].captured._linear)
-    self.assert_matches_eager(PlainJit(*self.convert(artifact["run"])), 5)
-
-    with mock.patch.object(sys, "argv", ["convert_upstream_pkl", pkl, out]):
-      convert_upstream_pkl.main()
-    bundle = load_model(out, cache=False, verify=True)
-    assert bundle["metadata"] == {"name": "tiny"} and bundle["input_specs"] == self.input_specs and bundle["output_specs"] == self.output_specs
-    self.assert_matches_eager(bundle["run"], 6)

openpilot/sunnypilot/modeld_v2/tests/test_plain_artifact.py

@@ -1,162 +0,0 @@
-"""
-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 dataclasses
-import functools
-import inspect
-import os
-import sys
-import tempfile
-from unittest import mock
-import numpy as np
-import unittest
-
-import openpilot.sunnypilot.modeld_v2.helpers as helpers
-from openpilot.common.test import OpenpilotTestCase
-from openpilot.sunnypilot.modeld_v2 import compile_warp
-from openpilot.sunnypilot.modeld_v2.frame_fusion import frame_part, pack_layout
-from openpilot.sunnypilot.modeld_v2.helpers import PlainJit, cache_is_fresh, load_model, read_parts, write_bundle
-from openpilot.sunnypilot.modeld_v2.plain_artifact_compiler import capture_frame, record_kernel_linear, serialize_jit
-
-CAM_W, CAM_H, MODEL_W, MODEL_H = 32, 16, 16, 8
-FRAME = f"{CAM_W},{CAM_H},{CAM_W},{CAM_H},{CAM_H // 2},{CAM_W * CAM_H * 3 // 2}"
-FRAME_SIZE = CAM_W * CAM_H * 3 // 2
-HOST = {"desire": (8,), "traffic_convention": (2,)}
-STATE, OUT = (4, 16), (1, 32)
-
-
-def compile_jit(fn, inputs: dict, output_container: str | None = None) -> tuple[dict, memoryview]:
-  from tinygrad import TinyJit
-  run = TinyJit(fn, prune=True)
-  with record_kernel_linear() as recorded:
-    run(**inputs)
-    run(**inputs)
-    check, _ = capture_frame(run, inputs, output_container)
-  assert len(recorded) == 1, f"expected one captured kernel list, got {len(recorded)}"
-  return serialize_jit(run, recorded[0], check)
-
-
-class TestPlainArtifact(OpenpilotTestCase):
-  @classmethod
-  def setUpClass(cls):
-    super().setUpClass()
-    if not os.environ.get("TINYGRAD_TESTS"):
-      raise unittest.SkipTest("TINYGRAD_TESTS env var not set, skipping artifact test.")
-    from tinygrad import Tensor
-    cls.tmp = tempfile.TemporaryDirectory()
-    cls.rng = np.random.default_rng(0)
-    cls.weights = Tensor(cls.rng.standard_normal((16, 32)).astype(np.float32)).realize()
-    cls.bias = Tensor(cls.rng.standard_normal((1, 32)).astype(np.float32)).realize()
-
-  @classmethod
-  def tearDownClass(cls):
-    cls.tmp.cleanup()
-    super().tearDownClass()
-
-  def model(self, x, h):  # a two input, one weight model: the shape of a policy without the size
-    return ((x @ self.weights + self.bias).relu() + h.mean()).contiguous()
-
-  def tensors(self, seed: int):
-    from tinygrad import Tensor
-    rng = np.random.default_rng(seed)
-    return {"x": Tensor(rng.standard_normal((1, 16)).astype(np.float32)).realize(), "h": Tensor(rng.standard_normal((4, 16)).astype(np.float32)).realize()}
-
-  def artifact(self, name: str = "model") -> str:
-    path = os.path.join(self.tmp.name, f"{name}_tinygrad.pkl")
-    if not os.path.exists(path):
-      write_bundle(path, {"run": compile_jit(self.model, self.tensors(1))}, {"metadata": {"name": name}})
-    return path
-
-  def test_round_trip_matches_jit(self):  # compiler -> bundle -> PlainJit, replay verified, then fresh inputs bit exact with eager tinygrad
-    bundle = load_model(self.artifact(), cache=False, verify=True)
-    assert bundle["metadata"] == {"name": "model"} and isinstance(bundle["run"], PlainJit)
-    fresh = self.tensors(2)
-    assert np.array_equal(bundle["run"](**fresh).numpy(), self.model(**fresh).numpy())
-
-  def test_replay_check_catches_a_changed_kernel(self):
-    parts, extras = read_parts(self.artifact())
-    header, blob = parts["run"]
-    poisoned = dict(header, check=dict(header["check"], outputs=[header["check"]["inputs"][0][1]]))
-    with self.assertRaises(helpers.ArtifactMismatch):
-      PlainJit(poisoned, blob)
-
-  def test_cache_round_trip_holds_no_weights_and_is_rebuilt_when_stale(self):
-    path = self.artifact("cached")
-    expected = load_model(path, cache=False, verify=False)["run"](**self.tensors(3)).numpy()
-    load_model(path)
-    assert cache_is_fresh(path)
-    with open(f"{path}.cache", "rb") as file:
-      cache = file.read()
-    assert self.weights.numpy().tobytes() not in cache and len(cache) < os.path.getsize(path)
-    with mock.patch.object(helpers, "PlainJit", side_effect=AssertionError("a fresh cache must not rebuild")):
-      cached = load_model(path)
-    assert np.array_equal(cached["run"](**self.tensors(3)).numpy(), expected)
-    with mock.patch.object(helpers, "tinygrad_identity", return_value="another tinygrad"):
-      assert not cache_is_fresh(path)
-      assert isinstance(load_model(path)["run"], PlainJit)
-      assert cache_is_fresh(path)
-    assert not cache_is_fresh(path)  # keyed on the tinygrad that wrote it, so the next boot after a bump rebuilds once
-    os.utime(path)
-    with mock.patch.object(helpers, "tinygrad_identity", return_value="another tinygrad"):
-      assert not cache_is_fresh(path)
-
-  def test_fused_with_warp_matches_warp_then_model(self):
-    from tinygrad import Tensor
-    warp_path = os.path.join(self.tmp.name, "warp_tinygrad.pkl")
-    with mock.patch.object(sys, "argv", ["compile_warp", "--frame", FRAME, "--warp-to", f"{MODEL_W}x{MODEL_H}", "--output", warp_path]):
-      compile_warp.main()
-    warp = load_model(warp_path, cache=False)["run"]
-
-    def native(new_img, desire, traffic_convention, state_feat_q, output_buffers):
-      feat = state_feat_q.mean(axis=0).reshape(1, 16) + new_img.float().mean() + desire.sum() + traffic_convention[0]
-      Tensor.realize(output_buffers["outputs"].assign((feat @ self.weights + self.bias).relu()),
-                     output_buffers["next_state_feat_q"].assign(state_feat_q[1:].cat(feat * 2)))
-
-    def model_inputs(new_img):
-      return {"new_img": new_img} | {name: Tensor(np.ones(shape, np.float32) * 0.5).realize() for name, shape in HOST.items()} | \
-             {"state_feat_q": Tensor(self.rng.standard_normal(STATE).astype(np.float32)).realize(),
-              "output_buffers": {"outputs": Tensor.zeros(OUT).contiguous().realize(), "next_state_feat_q": Tensor.zeros(STATE).contiguous().realize()}}
-    frames = Tensor(self.rng.integers(0, 256, (2, FRAME_SIZE), dtype=np.uint8)).realize()
-    tfm = Tensor(np.tile(np.eye(3, dtype=np.float32), (2, 1, 1))).realize()
-    path = os.path.join(self.tmp.name, "native_tinygrad.pkl")
-    part = compile_jit(native, model_inputs(warp(input_frame=frames, M_inv=tfm).realize()), "output_buffers")
-    extras = {"input_specs": {"new_img": ((2, 6, MODEL_H // 2, MODEL_W // 2), "|u1", "CPU"), "state_feat_q": (STATE, "<f4", "CPU")}
-              | {name: (shape, "<f4", "CPU") for name, shape in HOST.items()},
-              "output_specs": {"outputs": (OUT, "<f4", "CPU"), "next_state_feat_q": (STATE, "<f4", "CPU")},
-              "metadata": {"output_shapes": {"outputs": OUT, "next_state_feat_q": STATE}}}
-    write_bundle(path, {"run": part}, extras)
-
-    fused = load_model(path, cache=False, extend=functools.partial(frame_part, warp_path))
-    assert fused["run"].input_names == ["host", "state_feat_q"]
-    offsets, packed = pack_layout(HOST, FRAME_SIZE)
-    host = np.zeros(packed, np.uint8)
-    inputs = model_inputs(None)
-    for name, shape in ({"tfm": (2, 3, 3)} | HOST).items():
-      np.ndarray(shape, np.float32, host, offsets[name])[...] = tfm.numpy() if name == "tfm" else inputs[name].numpy()
-    host[offsets["frames"]:] = frames.numpy().ravel()
-    fused["run"](host=Tensor(host, device="NPY"), state_feat_q=inputs["state_feat_q"], output_buffers=inputs["output_buffers"])
-
-    expected = model_inputs(warp(input_frame=frames, M_inv=tfm).realize()) | {k: inputs[k] for k in ("state_feat_q",)}
-    native(**expected)
-    for name in ("outputs", "next_state_feat_q"):
-      assert np.array_equal(inputs["output_buffers"][name].numpy(), expected["output_buffers"][name].numpy()), name
-
-  def test_tinygrad_api_the_loader_reads(self):
-    import tinygrad.engine.jit as jit
-    from tinygrad.device import Buffer, TinyELF
-    from tinygrad.helpers import Target
-    from tinygrad.uop.ops import ParamArg, ProgramInfo, UOp
-    fields = {cls: {f.name for f in dataclasses.fields(cls)} for cls in (ProgramInfo, ParamArg, Target)}
-    assert fields[ProgramInfo] >= {"global_size", "local_size", "globals", "outs", "ins", "target"}, "ProgramInfo changed: update helpers._program_uop"
-    assert fields[ParamArg] >= {"slot", "dtype", "name", "image"}, "ParamArg changed: update helpers._program_uop"
-    assert fields[Target] >= {"device", "renderer", "arch"}, "Target changed: update the 'target' key in plain_artifact_compiler.serialize_jit"
-    assert "signature" in getattr(TinyELF, "__annotations__", {}) or hasattr(UOp, "kernel_params"), \
-      "kernel signatures moved: update plain_artifact_compiler.serialize_jit"
-    assert callable(getattr(UOp, "param", None)) and callable(getattr(UOp, "from_buffer", None)), "UOp.param/from_buffer: update helpers.PlainJit"
-    assert callable(getattr(UOp, "store_call", None)) and callable(getattr(UOp, "call", None)), "UOp.call/store_call: update helpers.PlainJit"
-    assert callable(getattr(Buffer, "view", None)) and callable(getattr(Buffer, "ensure_allocated", None)), "Buffer.view: update helpers._upload"
-    assert callable(getattr(jit, "compile_linear", None)) and callable(getattr(jit, "_prepare_jit_inputs", None)), "engine.jit: update helpers._Jit"
-    assert "expected_names" in inspect.signature(jit.CapturedJit).parameters, "CapturedJit: update helpers.PlainJit and plain_artifact_compiler.serialize_jit"

openpilot/sunnypilot/models/manager.py

@@ -34,6 +34,8 @@ TRANSIENT_HTTP_STATUS = frozenset({408, 429, 500, 502, 503, 504})
 REPORT_INTERVAL = 0.5  # seconds between progress publications
 # byte ranges count encoded bytes but iter_content yields decoded ones, so ask for an unencoded body
 NO_COMPRESSION = {"Accept-Encoding": "identity"}
+SPEED_SMOOTHING = 0.7  # weight of the previous speed sample in the published speed
+SPEED_FLOOR = 0.05 * 1024 ** 2  # bytes/s; anything slower would only display as 0.0 MB/s, so it reads as no rate
 
 
 class DownloadCancelled(Exception):
@@ -154,6 +156,7 @@ class ModelManagerSP:
         artifact.downloadProgress.status = source_artifact.downloadProgress.status
         artifact.downloadProgress.progress = source_artifact.downloadProgress.progress
         artifact.downloadProgress.eta = source_artifact.downloadProgress.eta
+        artifact.downloadProgress.speed = source_artifact.downloadProgress.speed
 
   def _calculate_eta(self, filename: str, progress: float) -> int:
     if filename not in self._download_start_times or progress <= 0:
@@ -167,28 +170,43 @@ class ModelManagerSP:
     eta = total_estimated_time - elapsed_time
     return max(1, int(eta))  # Return at least 1 second if download is ongoing
 
-  def _set_progress(self, artifact, status, progress: float, eta: int = 0) -> None:
+  def _set_progress(self, artifact, status, progress: float, eta: int = 0, speed: float = 0.0) -> None:
     artifact.downloadProgress.status = status
     artifact.downloadProgress.progress = progress
     artifact.downloadProgress.eta = eta
+    artifact.downloadProgress.speed = speed
     self._sync_artifact_progress(artifact)
     self._report_status()
 
-  def _publish_progress(self, artifact, done_bytes: int, total: int) -> None:
+  def _publish_progress(self, artifact, done_bytes: int, total: int, speed: float) -> None:
     # held at 99 until the file passes its hash check
     progress = min(99.0, done_bytes / total * 100) if total > 0 else 0.0
-    eta = self._calculate_eta(artifact.fileName, progress)
-    self._set_progress(artifact, custom.ModelManagerSP.DownloadStatus.downloading, progress, eta)
+    if speed > 0 and total > 0:
+      eta = max(1, int((total - done_bytes) / speed))
+    else:
+      eta = self._calculate_eta(artifact.fileName, progress)
+    self._set_progress(artifact, custom.ModelManagerSP.DownloadStatus.downloading, progress, eta, speed)
 
-  async def _report_until_done(self, tasks: list[asyncio.Future], artifact, progress: _Progress, total: int) -> None:
+  async def _report_until_done(self, tasks: list[asyncio.Future], artifact, progress: _Progress, total: int,
+                               started: tuple[int, float]) -> None:
     pending = set(tasks)
+    last_bytes, last_time = started
+    speed = 0.0
     while pending:
       done, pending = await asyncio.wait(pending, timeout=REPORT_INTERVAL)
       for task in done:
         task.result()  # re-raises a worker failure
       if self._download_interrupted():
         raise DownloadCancelled("Download cancelled")
-      self._publish_progress(artifact, progress.snapshot(), total)
+      now = time.monotonic()
+      done_bytes = progress.snapshot()
+      if (dt := now - last_time) > 0:
+        instant = max(0.0, (done_bytes - last_bytes) / dt)  # bytes/s this tick; clamped since retried pieces can lower the byte count
+        speed = instant if speed == 0 else SPEED_SMOOTHING * speed + (1 - SPEED_SMOOTHING) * instant
+        if speed < SPEED_FLOOR:  # a stall decays toward 0 without reaching it, and a near-0 speed overflows the eta
+          speed = 0.0
+      last_time, last_bytes = now, done_bytes
+      self._publish_progress(artifact, done_bytes, total, speed)
 
   async def _download_file(self, url: str, path: str, artifact) -> None:
     self._download_start_times[artifact.fileName] = time.monotonic()
@@ -199,12 +217,13 @@ class ModelManagerSP:
     await loop.run_in_executor(None, _prepare_target, path, total)
     progress = _Progress()
     cancel = threading.Event()
+    started = progress.snapshot(), time.monotonic()  # before any worker can land bytes
 
     with ThreadPoolExecutor(max_workers=MAX_CONCURRENT_PIECES) as pool:
       tasks = [loop.run_in_executor(pool, _fetch_piece, url, path, i, start, end, progress, cancel, self._block_size)
                for i, (start, end) in enumerate(pieces)]
       try:
-        await self._report_until_done(tasks, artifact, progress, total)
+        await self._report_until_done(tasks, artifact, progress, total, started)
       except BaseException as e:
         cancel.set()
         if isinstance(e, DownloadCancelled):
@@ -247,6 +266,7 @@ class ModelManagerSP:
       self._download_start_times.pop(artifact.fileName, None)
       artifact.downloadProgress.status = status.failed
       artifact.downloadProgress.eta = 0
+      artifact.downloadProgress.speed = 0
       self._sync_artifact_progress(artifact)
       if self.selected_bundle:
         self.selected_bundle.status = status.failed
@@ -259,6 +279,7 @@ class ModelManagerSP:
         os.remove(full_path)
       artifact.downloadProgress.status = status.failed
       artifact.downloadProgress.eta = 0
+      artifact.downloadProgress.speed = 0
       self._sync_artifact_progress(artifact)
       if self.selected_bundle:
         self.selected_bundle.status = status.failed
@@ -301,6 +322,7 @@ class ModelManagerSP:
           artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.cached
           artifact.downloadProgress.progress = 100
           artifact.downloadProgress.eta = 0
+          artifact.downloadProgress.speed = 0
         else:
           seen_artifacts.add(artifact.fileName)
           await self._process_artifact(artifact, destination_path)

openpilot/sunnypilot/models/tests/test_manager_download.py

@@ -6,6 +6,7 @@ See the LICENSE.md file in the root directory for more details.
 """
 
 import asyncio
+import concurrent.futures
 import contextlib
 import dataclasses
 import hashlib
@@ -129,6 +130,13 @@ def threaded_server(handler):
     thread.join()
 
 
+class InlineExecutor(concurrent.futures.Executor):
+  def submit(self, fn, /, *args, **kwargs):
+    future: concurrent.futures.Future = concurrent.futures.Future()
+    future.set_result(fn(*args, **kwargs))  # the piece lands before the reporter starts
+    return future
+
+
 class ManagerDownloadTestBase(OpenpilotTestCase):
   def setUp(self):
     super().setUp()
@@ -139,7 +147,7 @@ class ManagerDownloadTestBase(OpenpilotTestCase):
     self.addCleanup(self._tmp.cleanup)
     self.dest = self._tmp.name
 
-    for name, value in (('PIECE_SIZE', PIECE_SIZE), ('REPORT_INTERVAL', 0.05), ('PIECE_BACKOFF', 0.01)):
+    for name, value in (('PIECE_SIZE', PIECE_SIZE), ('REPORT_INTERVAL', 0.05), ('PIECE_BACKOFF', 0.01), ('SPEED_FLOOR', 100)):
       self.enterContext(mock.patch.object(manager_module, name, value))
 
     self.reported: list[float] = []
@@ -229,6 +237,31 @@ class TestManagerDownload(ManagerDownloadTestBase):
     assert self.reported and self.reported == sorted(self.reported), "progress only rises"
     assert max(self.reported) <= 99.0, "progress holds below 100 until the hash check"
     assert artifact.fileName not in self.manager._download_start_times
+    assert artifact.downloadProgress.speed > 0, "a running download shows its speed"
+
+  def test_speed_counts_bytes_landed_before_the_first_report(self):
+    self.enterContext(mock.patch.object(manager_module, 'ThreadPoolExecutor', lambda max_workers: InlineExecutor()))
+    artifact = self.download_file()
+    assert artifact.downloadProgress.speed > 0
+
+  def test_eta_comes_from_speed(self):
+    artifact = self.make_artifact()
+    self.manager._publish_progress(artifact, done_bytes=400, total=1000, speed=100)
+    assert artifact.downloadProgress.eta == 6
+
+  def test_stalled_speed_drops_to_zero(self):
+    stall_event = self.stall(0)
+    seen_speed = threading.Event()
+
+    def release_once_zero(*args):
+      if self.artifact.downloadProgress.speed > 0:
+        seen_speed.set()
+      elif seen_speed.is_set():
+        stall_event.set()
+    self.pm.send.side_effect = release_once_zero
+    self.download_file()
+    assert self.server.stall_released_by_event is True, "the stalled speed never reached exactly 0"
+    assert self.artifact.downloadProgress.speed > 0, "the speed must recover once data flows again"
 
   def test_pieces_download_in_parallel(self):
     stall_event = self.stall(0)
@@ -303,7 +336,8 @@ class TestProcessArtifact(ManagerDownloadTestBase):
         self.download_bundle(source)
         assert ACTIVE_BUNDLE_KEYS[source] in store and ACTIVE_BUNDLE_KEYS[other] not in store
         assert self.manager.selected_bundle.status == ds.downloaded
-        assert self.artifact.downloadProgress.progress == 100 and self.artifact.downloadProgress.eta == 0
+        assert self.artifact.downloadProgress.progress == 100
+        assert self.artifact.downloadProgress.eta == 0 and self.artifact.downloadProgress.speed == 0
         assert self.read_path() == WHOLE_BODY
         os.remove(self.path)