IQ.Pilot Prebuilt Release @ 27f668a

iqpilot/selfdrive/iqmodeld/tests/test_model_bundle_downloader.py

@@ -0,0 +1,133 @@
+"""
+Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
+"""
+import hashlib
+import http.server
+import os
+import threading
+
+import pytest
+
+from iqpilot.selfdrive.iqmodeld import model_bundle_downloader as dl
+
+
+class _RangeHandler(http.server.BaseHTTPRequestHandler):
+  store: dict[str, bytes] = {}
+  cut_first: dict[str, int] = {}
+  hits: list[tuple[str, str | None]] = []
+
+  def log_message(self, *a):
+    pass
+
+  def do_GET(self):
+    oid = self.path.rsplit("/", 1)[-1]
+    data = self.store[oid]
+    rng = self.headers.get("Range")
+    self.hits.append((oid, rng))
+    start = int(rng.split("=")[1].rstrip("-")) if rng else 0
+    body = data[start:]
+    cut = self.cut_first.pop(oid, None)
+    if cut is not None:
+      body = body[:cut]
+    self.send_response(206 if rng else 200)
+    self.send_header("Content-Length", str(len(body)))
+    if rng:
+      self.send_header("Content-Range", f"bytes {start}-{start + len(body) - 1}/{len(data)}")
+    self.end_headers()
+    self.wfile.write(body)
+
+
+@pytest.fixture
+def server():
+  srv = http.server.ThreadingHTTPServer(("127.0.0.1", 0), _RangeHandler)
+  t = threading.Thread(target=srv.serve_forever, daemon=True)
+  t.start()
+  yield srv
+  srv.shutdown()
+  srv.server_close()
+
+
+def _objects(parts):
+  return [{"oid": hashlib.sha256(p).hexdigest(), "size": len(p)} for p in parts]
+
+
+def test_resume_continues_a_cut_part_and_reuses_finished_parts(server, tmp_path, monkeypatch):
+  parts = [os.urandom(300_000), os.urandom(300_000), os.urandom(120_000)]
+  objs = _objects(parts)
+  _RangeHandler.store = {o["oid"]: p for o, p in zip(objs, parts, strict=True)}
+  _RangeHandler.hits = []
+  _RangeHandler.cut_first = {objs[1]["oid"]: 100_000}
+  port = server.server_address[1]
+  monkeypatch.setattr(dl, "_requests_auth", lambda: None)
+  monkeypatch.setattr(dl, "_resolve_oid", lambda session, base, oid, size, auth: (f"http://127.0.0.1:{port}/o/{oid}", {}))
+  monkeypatch.setattr(dl, "MODELS_BASE_URLS", ("http://unused",))
+  monkeypatch.setattr(dl, "STREAM_RETRIES", 3)
+  monkeypatch.setattr(dl, "CHUNK", 64 * 1024)
+  whole = b"".join(parts)
+  dst = str(tmp_path / "model.pkl")
+  out = dl.download_lfs_bundle(objs, dst, hashlib.sha256(whole).hexdigest(), len(whole))
+  with open(dst, "rb") as f:
+    assert out == dst and f.read() == whole
+  assert not os.path.exists(dst + ".parts")
+  ranges = [r for o, r in _RangeHandler.hits if o == objs[1]["oid"]]
+  assert ranges[0] is None and ranges[1] == "bytes=100000-"
+  assert sum(1 for o, _ in _RangeHandler.hits if o == objs[0]["oid"]) == 1
+
+
+def test_corrupt_finished_part_is_refetched(server, tmp_path, monkeypatch):
+  parts = [os.urandom(200_000), os.urandom(50_000)]
+  objs = _objects(parts)
+  _RangeHandler.store = {o["oid"]: p for o, p in zip(objs, parts, strict=True)}
+  _RangeHandler.hits = []
+  _RangeHandler.cut_first = {}
+  port = server.server_address[1]
+  monkeypatch.setattr(dl, "_requests_auth", lambda: None)
+  monkeypatch.setattr(dl, "_resolve_oid", lambda session, base, oid, size, auth: (f"http://127.0.0.1:{port}/o/{oid}", {}))
+  monkeypatch.setattr(dl, "MODELS_BASE_URLS", ("http://unused",))
+  dst = str(tmp_path / "model.pkl")
+  os.makedirs(dst + ".parts")
+  with open(dl._part_path(dst, objs[0]["oid"]), "wb") as f:
+    f.write(os.urandom(200_000))
+  whole = b"".join(parts)
+  dl.download_lfs_bundle(objs, dst, hashlib.sha256(whole).hexdigest(), len(whole))
+  with open(dst, "rb") as f:
+    assert f.read() == whole
+
+
+def test_hf_single_file_resumes_after_cut(server, tmp_path, monkeypatch):
+  data = os.urandom(700_000)
+  oid = hashlib.sha256(data).hexdigest()
+  _RangeHandler.store = {oid: data}
+  _RangeHandler.hits = []
+  _RangeHandler.cut_first = {oid: 250_000}
+  port = server.server_address[1]
+  monkeypatch.setattr(dl, "_hf", lambda: ({"Authorization": "Bearer test"}, lambda p: f"http://127.0.0.1:{port}/o/{oid}"))
+  monkeypatch.setattr(dl, "STREAM_RETRIES", 3)
+  monkeypatch.setattr(dl, "CHUNK", 64 * 1024)
+  dst = str(tmp_path / "policy.pkl")
+  out = dl.download_hf_file("egpu/policy/x.pkl", dst, oid, len(data))
+  with open(dst, "rb") as f:
+    assert out == dst and f.read() == data
+  ranges = [r for o, r in _RangeHandler.hits if o == oid]
+  assert ranges[0] is None and ranges[1] == "bytes=250000-"
+  assert not os.path.exists(dst + ".hfpart")
+
+
+def test_download_onnx_prefers_hf_then_falls_back(tmp_path, monkeypatch):
+  from iqpilot.selfdrive.iqmodeld import egpu_helpers as eh
+  meta = {"key": "m", "sha256": "ab" * 32, "download": {"kind": "comma_lfs", "size": 5}}
+  monkeypatch.setattr(eh, "onnx_cache_path", lambda m: str(tmp_path / "m.onnx"))
+  monkeypatch.setattr("iqpilot.selfdrive.iqmodeld.egpu_model.download_descriptor", lambda m: ("commalfs:" + m["sha256"], 5), raising=False)
+  calls = []
+  import iqpilot.selfdrive.iqmodeld.model_bundle_downloader as dlm
+  monkeypatch.setattr(dlm, "download_hf_file", lambda path, dst, sha, size, progress_cb=None: (calls.append(("hf", path)), open(dst, "wb").close(), dst)[2])
+  monkeypatch.setattr(eh, "resolve_download_url", lambda *a, **k: (calls.append(("lfs",)), "http://unused")[1])
+  out = eh.download_onnx(meta)
+  assert calls == [("hf", "onnx/" + "ab" * 32 + ".onnx")] and out == str(tmp_path / "m.onnx")
+  calls.clear()
+  def boom(*a, **k):
+    calls.append(("hf-fail",)); raise RuntimeError("hf down")
+  monkeypatch.setattr(dlm, "download_hf_file", boom)
+  with pytest.raises(Exception):
+    eh.download_onnx(meta)
+  assert calls[:2] == [("hf-fail",), ("lfs",)]