IQ.Pilot Release Commit @ 461be14

iqpilot/selfdrive/iqmodeld/model_bundle_downloader.py

@@ -27,6 +27,63 @@ def _requests_auth():
   return None
 
 
+def _hf():
+  import importlib
+  for mod in ("iqpilot_private.models.git_auth", "iqpilot.selfdrive.iqmodeld.models.git_auth"):
+    try:
+      m = importlib.import_module(mod)
+      return m.get_hf_headers(), m.hf_resolve_url
+    except Exception:
+      continue
+  return None, None
+
+
+def download_hf_file(hf_path: str, dst: str, sha256: str, size: int, progress_cb=None) -> str:
+  import requests
+  headers, resolve = _hf()
+  if resolve is None:
+    raise RuntimeError("no HF credentials available")
+  url = resolve(hf_path)
+  os.makedirs(os.path.dirname(dst), exist_ok=True)
+  tmp = dst + ".hfpart"
+  last_error: Exception | None = None
+  for _attempt in range(STREAM_RETRIES):
+    try:
+      have = os.path.getsize(tmp) if os.path.isfile(tmp) else 0
+      if size and have > size:
+        os.remove(tmp)
+        have = 0
+      if not size or have < size:
+        req_headers = dict(headers)
+        if have:
+          req_headers["Range"] = f"bytes={have}-"
+        with requests.get(url, headers=req_headers, stream=True, timeout=HTTP_TIMEOUT_S, allow_redirects=True) as r:
+          r.raise_for_status()
+          if have and r.status_code != 206:
+            have = 0
+          with open(tmp, "ab" if have else "wb") as f:
+            got = have
+            for chunk in r.iter_content(CHUNK):
+              f.write(chunk)
+              got += len(chunk)
+              if progress_cb is not None and size:
+                progress_cb(min(1.0, got / size))
+      digest = hashlib.sha256()
+      with open(tmp, "rb") as f:
+        for chunk in iter(lambda: f.read(CHUNK), b""):
+          digest.update(chunk)
+      if size and os.path.getsize(tmp) != size:
+        raise RuntimeError(f"size mismatch: {os.path.getsize(tmp)}/{size} bytes")
+      if sha256 and digest.hexdigest() != sha256:
+        os.remove(tmp)
+        raise RuntimeError("sha256 mismatch")
+      os.replace(tmp, dst)
+      return dst
+    except Exception as e:
+      last_error = e
+  raise RuntimeError(f"HF download failed: {last_error}")
+
+
 def _lfs_endpoint(base_url: str) -> str:
   return base_url.split("/raw/", 1)[0] + ".git/info/lfs"
 
@@ -47,41 +104,108 @@ def _resolve_oid(session, base_url: str, oid: str, size: int, auth):
   return action["href"], action.get("header", {})
 
 
+def _part_path(dst: str, oid: str) -> str:
+  return os.path.join(dst + ".parts", oid)
+
+
+def _part_complete(path: str, oid: str, size: int) -> bool:
+  if not os.path.isfile(path) or os.path.getsize(path) != size:
+    return False
+  digest = hashlib.sha256()
+  with open(path, "rb") as f:
+    for chunk in iter(lambda: f.read(CHUNK), b""):
+      digest.update(chunk)
+  return digest.hexdigest() == oid
+
+
+def _fetch_part(session, base_url: str, obj: dict, path: str, auth, progress) -> None:
+  size = int(obj["size"])
+  have = os.path.getsize(path) if os.path.isfile(path) else 0
+  if have > size:
+    os.remove(path)
+    have = 0
+  href, headers = _resolve_oid(session, base_url, obj["oid"], size, auth)
+  obj_auth = None if headers.get("Authorization") else auth
+  # LFS parts are content-addressed (oid == sha256), so a half-written part can be resumed with a
+  # Range request and verified afterwards instead of being thrown away on every restart.
+  if have:
+    headers = {**headers, "Range": f"bytes={have}-"}
+  with session.get(href, headers=headers, stream=True, timeout=HTTP_TIMEOUT_S, auth=obj_auth) as r:
+    r.raise_for_status()
+    if have and r.status_code != 206:
+      have = 0
+    with open(path, "ab" if have else "wb") as f:
+      for chunk in r.iter_content(CHUNK):
+        f.write(chunk)
+        progress(len(chunk))
+
+
 def download_lfs_bundle(objects: list, dst: str, sha256: str, size: int, progress_cb=None) -> str:
   import requests
   auth = _requests_auth()
   session = requests.Session()
-  os.makedirs(os.path.dirname(dst), exist_ok=True)
-  tmp = dst + ".part"
+  os.makedirs(dst + ".parts", exist_ok=True)
   total = int(size) or sum(int(o["size"]) for o in objects)
+  done_bytes = sum(int(o["size"]) for o in objects if _part_complete(_part_path(dst, o["oid"]), o["oid"], int(o["size"])))
+  got = [done_bytes]
+
+  def progress(n: int) -> None:
+    got[0] += n
+    if progress_cb is not None and total:
+      progress_cb(min(1.0, got[0] / total))
+
   last_error: Exception | None = None
   for base_url in MODELS_BASE_URLS:
-    for attempt in range(STREAM_RETRIES):
+    for _attempt in range(STREAM_RETRIES):
       try:
-        digest = hashlib.sha256()
-        got = 0
-        with open(tmp, "wb") as f:
-          for obj in objects:
-            href, headers = _resolve_oid(session, base_url, obj["oid"], int(obj["size"]), auth)
-            obj_auth = None if headers.get("Authorization") else auth
-            with session.get(href, headers=headers, stream=True, timeout=120, auth=obj_auth) as r:
-              r.raise_for_status()
-              for chunk in r.iter_content(CHUNK):
-                f.write(chunk)
-                digest.update(chunk)
-                got += len(chunk)
-                if progress_cb is not None and total:
-                  progress_cb(min(1.0, got / total))
-        if total and got != total:
-          raise RuntimeError(f"size mismatch: {got}/{total} bytes")
-        if sha256 and digest.hexdigest() != sha256:
-          raise RuntimeError("sha256 mismatch")
-        os.replace(tmp, dst)
-        return dst
+        for obj in objects:
+          path = _part_path(dst, obj["oid"])
+          if _part_complete(path, obj["oid"], int(obj["size"])):
+            continue
+          got[0] = done_bytes
+          _fetch_part(session, base_url, obj, path, auth, progress)
+          if not _part_complete(path, obj["oid"], int(obj["size"])):
+            if os.path.getsize(path) >= int(obj["size"]):
+              os.remove(path)
+            raise RuntimeError(f"part {obj['oid'][:12]} incomplete or failed verification")
+          done_bytes += int(obj["size"])
+          got[0] = done_bytes
+        break
       except Exception as e:
         last_error = e
-        try:
-          os.remove(tmp)
-        except OSError:
-          pass
-  raise RuntimeError(f"model bundle download failed: {last_error}")
+    else:
+      continue
+    break
+  else:
+    raise RuntimeError(f"model bundle download failed: {last_error}")
+
+  tmp = dst + ".part"
+  digest = hashlib.sha256()
+  with open(tmp, "wb") as out:
+    for obj in objects:
+      with open(_part_path(dst, obj["oid"]), "rb") as f:
+        for chunk in iter(lambda: f.read(CHUNK), b""):
+          out.write(chunk)
+          digest.update(chunk)
+  if total and os.path.getsize(tmp) != total:
+    os.remove(tmp)
+    raise RuntimeError(f"size mismatch: {os.path.getsize(tmp) if os.path.exists(tmp) else 0}/{total} bytes")
+  if sha256 and digest.hexdigest() != sha256:
+    os.remove(tmp)
+    for obj in objects:
+      try:
+        os.remove(_part_path(dst, obj["oid"]))
+      except OSError:
+        pass
+    raise RuntimeError("sha256 mismatch")
+  os.replace(tmp, dst)
+  for obj in objects:
+    try:
+      os.remove(_part_path(dst, obj["oid"]))
+    except OSError:
+      pass
+  try:
+    os.rmdir(dst + ".parts")
+  except OSError:
+    pass
+  return dst