mirror of
https://gitlvb.teallvbs.xyz/IQ.Lvbs/IQ.Pilot.git
synced 2026-08-05 08:15:41 +08:00
IQ.Lvbs CI [bot]
392 lines
14 KiB
Python
392 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Copyright (c) IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos
|
|
"""
|
|
import asyncio
|
|
import hashlib
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import aiohttp
|
|
from cereal import custom
|
|
from openpilot.common.realtime import Ratekeeper
|
|
from openpilot.common.time_helpers import system_time_valid
|
|
from openpilot.iqpilot._proprietary_loader import ProprietaryModuleMissing, load_private_module
|
|
from openpilot.common.swaglog import cloudlog
|
|
from openpilot.system.hardware.hw import Paths
|
|
|
|
_TIME_SYNC_WAIT_TIMEOUT_S = 30.0
|
|
_TIME_SYNC_POLL_S = 0.5
|
|
|
|
|
|
def _wait_for_valid_clock(timeout: float = _TIME_SYNC_WAIT_TIMEOUT_S) -> None:
|
|
if system_time_valid():
|
|
return
|
|
cloudlog.warning("models_manager: system clock not yet valid, waiting for NTP before fetching")
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if system_time_valid():
|
|
cloudlog.warning("models_manager: system clock is now valid, resuming")
|
|
return
|
|
time.sleep(_TIME_SYNC_POLL_S)
|
|
cloudlog.warning("models_manager: gave up waiting for a valid clock, proceeding anyway")
|
|
|
|
try:
|
|
load_private_module(__name__, "iqpilot_private.models.manager")
|
|
_BaseIQModelManager = IQModelManager # noqa: F821
|
|
except ProprietaryModuleMissing:
|
|
from iqpilot.models_private_src.manager import IQModelManager as _BaseIQModelManager
|
|
|
|
from openpilot.iqpilot.selfdrive.iqmodeld.models.git_auth import get_aiohttp_auth
|
|
from openpilot.iqpilot.selfdrive.iqmodeld.models.helpers import (
|
|
bundle_files_ready,
|
|
get_active_bundle,
|
|
get_runtime_bundle_upgrade,
|
|
is_default_bundle,
|
|
persist_active_bundle,
|
|
)
|
|
|
|
|
|
_ACTIVE_BUNDLE_KEY = "ModelManager_ActiveBundle"
|
|
_DOWNLOAD_INDEX_KEY = "ModelManager_DownloadIndex"
|
|
_RUNNER_CACHE_KEY = "ModelRunnerTypeCache"
|
|
|
|
|
|
class IQModelManager(_BaseIQModelManager):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self._validated_active_key: tuple[tuple[str, str], ...] | None = None
|
|
self._manifest_refresh_key: tuple[tuple[str, str], ...] | None = None
|
|
|
|
@staticmethod
|
|
def _bundle_index(bundle) -> int | None:
|
|
try:
|
|
return int(getattr(bundle, "index", -1))
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
@staticmethod
|
|
def _bundle_files(bundle) -> list[tuple[str, str]]:
|
|
files = []
|
|
for model in getattr(bundle, "models", []) or []:
|
|
for artifact in (getattr(model, "metadata", None), getattr(model, "artifact", None)):
|
|
filename = getattr(artifact, "fileName", "") if artifact is not None else ""
|
|
if not filename:
|
|
continue
|
|
download_uri = getattr(artifact, "downloadUri", None)
|
|
sha256 = getattr(download_uri, "sha256", "") if download_uri is not None else ""
|
|
files.append((filename, sha256 or ""))
|
|
return files
|
|
|
|
@staticmethod
|
|
def _safe_model_path(filename: str) -> Path | None:
|
|
if not filename or os.path.basename(filename) != filename:
|
|
cloudlog.warning(f"Ignoring unsafe model filename {filename!r}")
|
|
return None
|
|
|
|
root = Path(Paths.model_root()).resolve()
|
|
path = (root / filename).resolve()
|
|
try:
|
|
path.relative_to(root)
|
|
except ValueError:
|
|
cloudlog.warning(f"Ignoring model path outside model root {path}")
|
|
return None
|
|
return path
|
|
|
|
@staticmethod
|
|
def _verify_file_sync(path: Path, expected_hash: str) -> bool:
|
|
if not path.is_file():
|
|
return False
|
|
if not expected_hash:
|
|
return True
|
|
|
|
sha256_hash = hashlib.sha256()
|
|
with open(path, "rb") as f:
|
|
for chunk in iter(lambda: f.read(1024 * 1024), b""):
|
|
sha256_hash.update(chunk)
|
|
return sha256_hash.hexdigest().lower() == expected_hash.lower()
|
|
|
|
def _bundle_validation_key(self, bundle) -> tuple[tuple[str, str], ...]:
|
|
return tuple(self._bundle_files(bundle))
|
|
|
|
def _bundle_files_valid(self, bundle) -> bool:
|
|
for filename, expected_hash in self._bundle_files(bundle):
|
|
path = self._safe_model_path(filename)
|
|
if path is None or not self._verify_file_sync(path, expected_hash):
|
|
return False
|
|
return True
|
|
|
|
def _remove_bundle_files(self, bundle) -> None:
|
|
for filename, _expected_hash in self._bundle_files(bundle):
|
|
path = self._safe_model_path(filename)
|
|
if path is None:
|
|
continue
|
|
for candidate in (path, Path(f"{path}.download")):
|
|
try:
|
|
if candidate.is_file():
|
|
candidate.unlink()
|
|
except OSError as e:
|
|
cloudlog.exception(f"Failed to remove model artifact {candidate}: {e}")
|
|
|
|
def _find_available_bundle(self, target):
|
|
target_index = self._bundle_index(target)
|
|
target_ref = getattr(target, "ref", None)
|
|
target_internal = getattr(target, "internalName", None)
|
|
target_display = getattr(target, "displayName", None)
|
|
|
|
for bundle in self.available_models:
|
|
if target_index is not None and self._bundle_index(bundle) == target_index:
|
|
return bundle
|
|
if target_ref and getattr(bundle, "ref", None) == target_ref:
|
|
return bundle
|
|
if target_internal and getattr(bundle, "internalName", None) == target_internal:
|
|
return bundle
|
|
if target_display and getattr(bundle, "displayName", None) == target_display:
|
|
return bundle
|
|
return None
|
|
|
|
def _bundle_matches(self, left, right) -> bool:
|
|
if left is None or right is None:
|
|
return False
|
|
|
|
left_index = self._bundle_index(left)
|
|
right_index = self._bundle_index(right)
|
|
if left_index is not None and right_index is not None and left_index == right_index:
|
|
return True
|
|
|
|
for attr in ("ref", "internalName", "displayName"):
|
|
left_value = getattr(left, attr, None)
|
|
if left_value and left_value == getattr(right, attr, None):
|
|
return True
|
|
return False
|
|
|
|
def _clear_active_bundle(self) -> None:
|
|
self.params.remove(_ACTIVE_BUNDLE_KEY)
|
|
self.params.remove(_RUNNER_CACHE_KEY)
|
|
self.active_bundle = None
|
|
self._validated_active_key = None
|
|
|
|
def _download_request_matches(self, bundle) -> bool:
|
|
bundle_index = self._bundle_index(bundle)
|
|
return bundle_index is not None and self._download_index() == bundle_index
|
|
|
|
def _queue_active_redownload_if_invalid(self) -> None:
|
|
if self.active_bundle is None:
|
|
self._validated_active_key = None
|
|
return
|
|
|
|
validation_key = self._bundle_validation_key(self.active_bundle)
|
|
if validation_key == self._validated_active_key:
|
|
return
|
|
|
|
if self._bundle_files_valid(self.active_bundle):
|
|
self._validated_active_key = validation_key
|
|
return
|
|
|
|
bundle = self._find_available_bundle(self.active_bundle) or self.active_bundle
|
|
bundle_index = self._bundle_index(bundle)
|
|
cloudlog.warning(f"Active model {_display_bundle_name(self.active_bundle)} is missing or corrupt; queueing redownload")
|
|
self._remove_bundle_files(bundle)
|
|
self._clear_active_bundle()
|
|
if bundle_index is not None and self._download_index() is None:
|
|
self.params.put(_DOWNLOAD_INDEX_KEY, bundle_index)
|
|
|
|
def _find_manifest_counterpart(self, target):
|
|
# never match by index: indexes shift between manifest generations, and a
|
|
# positional match could redownload a different model than the user selected
|
|
for attr in ("ref", "internalName", "displayName"):
|
|
value = getattr(target, attr, None)
|
|
if not value:
|
|
continue
|
|
for bundle in self.available_models:
|
|
if getattr(bundle, attr, None) == value:
|
|
return bundle
|
|
return None
|
|
|
|
def _queue_active_manifest_refresh(self) -> None:
|
|
active = self.active_bundle
|
|
if active is None or is_default_bundle(active):
|
|
return
|
|
if self._download_index() is not None:
|
|
return
|
|
|
|
counterpart = self._find_manifest_counterpart(active)
|
|
if counterpart is None:
|
|
return
|
|
counterpart_index = self._bundle_index(counterpart)
|
|
if counterpart_index is None:
|
|
return
|
|
|
|
active_files = dict(self._bundle_files(active))
|
|
stale = False
|
|
for filename, sha in self._bundle_files(counterpart):
|
|
if not sha:
|
|
continue
|
|
active_sha = active_files.get(filename)
|
|
# an empty recorded hash can't prove a mismatch, so it never triggers a redownload
|
|
if active_sha is None or (active_sha and active_sha.lower() != sha.lower()):
|
|
stale = True
|
|
break
|
|
if not stale:
|
|
self._manifest_refresh_key = None
|
|
return
|
|
|
|
# the manifest may be an expired offline cache, so keep the active bundle and its
|
|
# files in place: the download flow replaces artifacts atomically and only persists
|
|
# the counterpart as active once everything landed. One attempt per bundle per run
|
|
# so a dead network doesn't turn the 1Hz loop into a download-retry storm.
|
|
key = self._bundle_validation_key(active)
|
|
if key == self._manifest_refresh_key:
|
|
return
|
|
self._manifest_refresh_key = key
|
|
|
|
cloudlog.warning(f"Active model {_display_bundle_name(active)} artifacts are stale vs current manifest; queueing redownload")
|
|
self.params.put(_DOWNLOAD_INDEX_KEY, counterpart_index)
|
|
|
|
async def _download_file(self, url: str, path: str, model) -> None:
|
|
temp_path = f"{path}.download"
|
|
self._download_start_times[model.fileName] = time.monotonic()
|
|
|
|
try:
|
|
if os.path.exists(temp_path):
|
|
os.remove(temp_path)
|
|
|
|
async with aiohttp.ClientSession(auth=get_aiohttp_auth()) as session:
|
|
async with session.get(url) as response:
|
|
response.raise_for_status()
|
|
total_size = int(response.headers.get("content-length", 0))
|
|
bytes_downloaded = 0
|
|
|
|
with open(temp_path, "wb") as f:
|
|
async for chunk in response.content.iter_chunked(self._chunk_size):
|
|
f.write(chunk)
|
|
bytes_downloaded += len(chunk)
|
|
|
|
if self._download_index() is None:
|
|
raise Exception("Download cancelled")
|
|
|
|
if total_size > 0:
|
|
progress = (bytes_downloaded / total_size) * 100
|
|
model.downloadProgress.status = custom.IQModelManager.DownloadStatus.downloading
|
|
model.downloadProgress.progress = progress
|
|
model.downloadProgress.eta = self._calculate_eta(model.fileName, progress)
|
|
self._report_status()
|
|
|
|
f.flush()
|
|
os.fsync(f.fileno())
|
|
|
|
os.replace(temp_path, path)
|
|
|
|
except Exception:
|
|
if os.path.exists(temp_path):
|
|
os.remove(temp_path)
|
|
raise
|
|
|
|
finally:
|
|
self._download_start_times.pop(model.fileName, None)
|
|
|
|
async def _download_bundle(self, model_bundle: custom.IQModelManager.ModelBundle, destination_path: str) -> None:
|
|
self.selected_bundle = model_bundle
|
|
self.selected_bundle.status = custom.IQModelManager.DownloadStatus.downloading
|
|
os.makedirs(destination_path, exist_ok=True)
|
|
|
|
try:
|
|
if not self._download_request_matches(model_bundle):
|
|
raise RuntimeError("Download cancelled")
|
|
|
|
tasks = [self._process_model(model, destination_path) for model in self.selected_bundle.models]
|
|
await asyncio.gather(*tasks)
|
|
|
|
if not self._download_request_matches(model_bundle):
|
|
raise RuntimeError("Download cancelled")
|
|
|
|
self.active_bundle = self.selected_bundle
|
|
self.active_bundle.status = custom.IQModelManager.DownloadStatus.downloaded
|
|
self.params.put(_ACTIVE_BUNDLE_KEY, self.active_bundle.to_dict())
|
|
self.params.remove(_RUNNER_CACHE_KEY)
|
|
self.selected_bundle = None
|
|
|
|
except Exception:
|
|
if self._download_request_matches(model_bundle) and self.selected_bundle is not None:
|
|
self.selected_bundle.status = custom.IQModelManager.DownloadStatus.failed
|
|
else:
|
|
self.selected_bundle = None
|
|
raise
|
|
|
|
finally:
|
|
self._report_status()
|
|
|
|
def download(self, model_bundle: custom.IQModelManager.ModelBundle, destination_path: str) -> None:
|
|
asyncio.run(self._download_bundle(model_bundle, destination_path))
|
|
|
|
def _queue_tinygrad_upgrade(self) -> None:
|
|
if self.active_bundle is None:
|
|
return
|
|
|
|
replacement = get_runtime_bundle_upgrade(self.active_bundle, self.params, self.available_models)
|
|
if replacement is None or replacement is self.active_bundle:
|
|
return
|
|
|
|
if bundle_files_ready(replacement):
|
|
persist_active_bundle(self.params, replacement)
|
|
self.active_bundle = replacement
|
|
return
|
|
|
|
if self._download_index() is None and getattr(replacement, "index", None) is not None:
|
|
self.params.put("ModelManager_DownloadIndex", int(replacement.index))
|
|
cloudlog.warning(f"Queued tinygrad upgrade for retired bundle {getattr(self.active_bundle, 'internalName', '<unknown>')}")
|
|
|
|
def main_thread(self) -> None:
|
|
_wait_for_valid_clock()
|
|
rk = Ratekeeper(1, print_delay_threshold=None)
|
|
|
|
while True:
|
|
try:
|
|
# before NTP the TLS cert reads "not yet valid" and every fetch SSL-fails; one line, not spam
|
|
if not system_time_valid():
|
|
if not getattr(self, "_ntp_wait_logged", False):
|
|
cloudlog.warning("models_manager: waiting for NTP before fetching (system clock not valid)")
|
|
self._ntp_wait_logged = True
|
|
rk.keep_time()
|
|
continue
|
|
self._ntp_wait_logged = False
|
|
|
|
self.available_models = self.model_fetcher.get_available_bundles()
|
|
self.active_bundle = get_active_bundle(self.params)
|
|
self._queue_active_redownload_if_invalid()
|
|
self._queue_tinygrad_upgrade()
|
|
self._queue_active_manifest_refresh()
|
|
|
|
if (index_to_download := self._download_index()) is not None:
|
|
if model_to_download := next((model for model in self.available_models if model.index == index_to_download), None):
|
|
try:
|
|
self.download(model_to_download, Paths.model_root())
|
|
except Exception as e:
|
|
cloudlog.exception(e)
|
|
finally:
|
|
self.params.remove("ModelManager_DownloadIndex")
|
|
self.selected_bundle = None
|
|
|
|
if self.params.get("ModelManager_ClearCache"):
|
|
self.clear_model_cache()
|
|
self.params.remove("ModelManager_ClearCache")
|
|
|
|
self._report_status()
|
|
rk.keep_time()
|
|
|
|
except Exception as e:
|
|
cloudlog.exception(f"Error in main thread: {str(e)}")
|
|
rk.keep_time()
|
|
|
|
|
|
def _display_bundle_name(bundle) -> str:
|
|
return getattr(bundle, "internalName", None) or getattr(bundle, "displayName", None) or "<unknown>"
|
|
|
|
|
|
def main():
|
|
IQModelManager().main_thread()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|