mirror of
https://github.com/sunnypilot/sunnypilot.git
synced 2026-07-26 03:02:06 +08:00
Compare commits
45 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| d9d656e1bf | |||
| 528fcda7bc | |||
| 09df48a316 | |||
| 1e1d23c81c | |||
| 6c535e2be9 | |||
| b640d8afb0 | |||
| 0d3d510494 | |||
| 79964f65f1 | |||
| 7e40ca0442 | |||
| 4aea44a4bf | |||
| f70170220d | |||
| d3319c3042 | |||
| 955f5cb73c | |||
| fd22de1c9a | |||
| ee3583df33 | |||
| 70424bd661 | |||
| d215eab1d4 | |||
| 91316c8cb5 | |||
| fa284be7e6 | |||
| 6e7d9e5e52 | |||
| a4a7c2335d | |||
| ae573c7c3f | |||
| 7d7b6ee306 | |||
| 6a4c59c3e0 | |||
| fb5cb7a1cc | |||
| 049dfd2eaa | |||
| be20848487 | |||
| cdd232b606 | |||
| b21c70b1ba | |||
| 67e5bd3c1e | |||
| 74692d0b5f | |||
| dd35c27981 | |||
| 159140e64e | |||
| f1ab6c8dfb | |||
| e1fe30fd3e | |||
| fba521dcff | |||
| a8ef55bfaa | |||
| a232f54e2d | |||
| 2697008aa7 | |||
| ad5abd242a | |||
| 6c1e0f370b | |||
| 1083f5bf21 | |||
| dc5116c718 | |||
| 8611e08dc6 | |||
| dc0f73c63b |
@@ -12,11 +12,11 @@ on:
|
||||
required: false
|
||||
type: string
|
||||
recompiled_dir:
|
||||
description: 'Existing recompiled directory number (e.g. 3 for recompiled3)'
|
||||
description: 'Existing recompiled directory number (e.g. 1 for recompiled1)'
|
||||
required: true
|
||||
type: string
|
||||
json_version:
|
||||
description: 'driving_models version number to update (e.g. 5 for driving_models_v5.json)'
|
||||
description: 'driving_models version number to update (e.g. 18 for driving_models_v18.json)'
|
||||
required: true
|
||||
type: string
|
||||
artifact_suffix:
|
||||
@@ -63,12 +63,11 @@ on:
|
||||
default: 'None'
|
||||
options:
|
||||
- None
|
||||
- Simple Plan Models
|
||||
- Space Lab Models
|
||||
- TR Models
|
||||
- DTR Models
|
||||
- Master Models
|
||||
- Release Models
|
||||
- 2025 World Models
|
||||
- 2026 World Models
|
||||
- Custom Merge Models
|
||||
- FOF series models
|
||||
- Other
|
||||
custom_model_folder:
|
||||
description: 'Custom model folder name (if "Other" selected)'
|
||||
|
||||
@@ -30,6 +30,11 @@ on:
|
||||
required: false
|
||||
type: string
|
||||
default: ''
|
||||
target_hardware:
|
||||
description: 'Hardware target to compile for (qcom or usbgpu)'
|
||||
required: false
|
||||
type: string
|
||||
default: 'qcom'
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
upstream_branch:
|
||||
@@ -46,6 +51,14 @@ on:
|
||||
required: false
|
||||
type: boolean
|
||||
default: true
|
||||
target_hardware:
|
||||
description: 'Hardware target to compile for'
|
||||
required: true
|
||||
type: choice
|
||||
options:
|
||||
- qcom
|
||||
- usbgpu
|
||||
default: 'qcom'
|
||||
|
||||
|
||||
run-name: Build model [${{ inputs.custom_name || inputs.upstream_branch }}] from ref [${{ inputs.upstream_branch }}]
|
||||
@@ -173,7 +186,17 @@ jobs:
|
||||
COMPILE_MODELD="${{ github.workspace }}/openpilot/sunnypilot/modeld_v2/compile_modeld.py"
|
||||
MODEL_SIZE=$(python3 -c "from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE as s; print(f'{s[0]}x{s[1]}')")
|
||||
CAMERA_RES=$(python3 -c "from openpilot.common.transformations.camera import _ar_ox_fisheye as a, _os_fisheye as o; print(f'{a.width}x{a.height} {o.width}x{o.height}')")
|
||||
TG_FLAGS="DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1"
|
||||
|
||||
if [ "${{ inputs.target_hardware }}" == "usbgpu" ]; then
|
||||
echo "USBGPU build"
|
||||
export USBGPU=1
|
||||
TG_FLAGS="DEV=AMD USBGPU=1 IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1"
|
||||
OUTPUT_PKL="${{ env.MODELS_DIR }}/big_driving_tinygrad.pkl"
|
||||
else
|
||||
echo "QCOM build"
|
||||
TG_FLAGS="DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1"
|
||||
OUTPUT_PKL="${{ env.MODELS_DIR }}/driving_tinygrad.pkl"
|
||||
fi
|
||||
|
||||
# Generate metadata for all ONNX files
|
||||
find "${{ env.MODELS_DIR }}" -maxdepth 1 -name '*.onnx' | while IFS= read -r onnx_file; do
|
||||
@@ -187,6 +210,7 @@ jobs:
|
||||
OFF_POLICY_ONNX="${{ env.MODELS_DIR }}/driving_off_policy.onnx"
|
||||
ON_POLICY_ONNX="${{ env.MODELS_DIR }}/driving_on_policy.onnx"
|
||||
SUPERCOMBO_ONNX="${{ env.MODELS_DIR }}/supercombo.onnx"
|
||||
[ -f "$SUPERCOMBO_ONNX" ] || SUPERCOMBO_ONNX="${{ env.MODELS_DIR }}/driving_supercombo.onnx"
|
||||
|
||||
MODEL_TYPE="" ONNX_ARGS="" OUTPUT_NAME=""
|
||||
if [ -f "$VISION_ONNX" ]; then
|
||||
@@ -207,24 +231,15 @@ jobs:
|
||||
fi
|
||||
|
||||
if [ -n "$MODEL_TYPE" ]; then
|
||||
echo "Detected: $MODEL_TYPE -> driving_tinygrad.pkl"
|
||||
echo "Detected: $MODEL_TYPE -> $OUTPUT_PKL"
|
||||
env ${TG_FLAGS} python3 "$COMPILE_MODELD" \
|
||||
--model-type $MODEL_TYPE \
|
||||
--model-size $MODEL_SIZE \
|
||||
--camera-resolutions $CAMERA_RES \
|
||||
$ONNX_ARGS \
|
||||
--output "${{ env.MODELS_DIR }}/driving_tinygrad.pkl"
|
||||
--output "$OUTPUT_PKL"
|
||||
fi
|
||||
|
||||
- name: Validate Model Outputs
|
||||
run: |
|
||||
source /etc/profile
|
||||
export UV_PROJECT_ENVIRONMENT=${HOME}/venv
|
||||
export VIRTUAL_ENV=$UV_PROJECT_ENVIRONMENT
|
||||
python3 "${{ github.workspace }}/release/ci/model_generator.py" \
|
||||
--validate-only \
|
||||
--model-dir "${{ env.MODELS_DIR }}"
|
||||
|
||||
- name: Prepare Output
|
||||
run: |
|
||||
sudo rm -rf ${{ env.OUTPUT_DIR }}
|
||||
|
||||
+1
-1
Submodule opendbc_repo updated: b129bf1d71...d6b9c1adaa
@@ -137,10 +137,16 @@ struct ModelManagerSP @0xaedffd8f31e7b55d {
|
||||
eta @2 :UInt32;
|
||||
}
|
||||
|
||||
struct Chunk {
|
||||
fileName @0 :Text;
|
||||
sha256 @1 :Text;
|
||||
}
|
||||
|
||||
struct Artifact {
|
||||
fileName @0 :Text;
|
||||
downloadUri @1 :DownloadUri;
|
||||
downloadProgress @2 :DownloadProgress;
|
||||
chunks @3 :List(Chunk);
|
||||
}
|
||||
|
||||
struct Model {
|
||||
@@ -155,6 +161,7 @@ struct ModelManagerSP @0xaedffd8f31e7b55d {
|
||||
policy @3;
|
||||
offPolicy @4;
|
||||
onPolicy @5;
|
||||
chunked @6;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -132,6 +132,7 @@ struct OnroadEvent @0xc4fa6047f024e718 {
|
||||
userBookmark @95;
|
||||
excessiveActuation @96;
|
||||
audioFeedback @97;
|
||||
|
||||
soundsUnavailableDEPRECATED @47;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -134,9 +134,6 @@ inline static std::unordered_map<std::string, ParamKeyAttributes> keys = {
|
||||
{"UsbGpuPresent", {CLEAR_ON_MANAGER_START | CLEAR_ON_OFFROAD_TRANSITION, BOOL}},
|
||||
{"UsbGpuCompiled", {CLEAR_ON_MANAGER_START | CLEAR_ON_OFFROAD_TRANSITION, BOOL}},
|
||||
{"Version", {PERSISTENT, STRING}},
|
||||
{"WgpuEnabled", {CLEAR_ON_MANAGER_START | DEVELOPMENT_ONLY, BOOL}},
|
||||
{"WgpuModelName", {CLEAR_ON_MANAGER_START | CLEAR_ON_OFFROAD_TRANSITION | DEVELOPMENT_ONLY, STRING}},
|
||||
{"WgpuReady", {CLEAR_ON_MANAGER_START | DEVELOPMENT_ONLY, BOOL}},
|
||||
|
||||
// --- sunnypilot params --- //
|
||||
{"ApiCache_DriveStats", {PERSISTENT, JSON}},
|
||||
|
||||
@@ -37,9 +37,6 @@ available = probe_devices()
|
||||
if 'CUDA' in available:
|
||||
tg_backend = 'CUDA'
|
||||
tg_flags = f'DEV={tg_backend}'
|
||||
elif 'METAL' in available:
|
||||
tg_backend = 'METAL'
|
||||
tg_flags = f'DEV={tg_backend} FLOAT16=1'
|
||||
elif 'QCOM' in available:
|
||||
tg_backend = 'QCOM'
|
||||
tg_flags = f'DEV={tg_backend} IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1'
|
||||
@@ -58,7 +55,6 @@ tg_devices = { # which device to put jit inputs to at runtime
|
||||
}
|
||||
|
||||
USBGPU = usbgpu_present() # or release # TODO always build big model on release
|
||||
WGPU = os.getenv('WGPU') == '1'
|
||||
if USBGPU:
|
||||
usbgpu_tg_flags = f'DEBUG=2 DEV=USB+AMD:LLVM WARP_DEV={tg_backend} FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0'
|
||||
# the USB+AMD GPU takes an exclusive flock; serialize all targets that touch it
|
||||
@@ -88,13 +84,13 @@ compile_modeld_script = [
|
||||
model_w, model_h = MEDMODEL_INPUT_SIZE
|
||||
frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
|
||||
|
||||
for usbgpu in [False, True] if USBGPU or WGPU else [False]:
|
||||
for usbgpu in [False, True] if USBGPU else [False]:
|
||||
target_pkl_path = File(modeld_pkl_path(usbgpu)).abspath
|
||||
# BIG_INTO_SMALL=1 builds the default target from the big model, e.g. to test it without a USB GPU
|
||||
file_prefix, cmd_flags = ('big_', usbgpu_tg_flags if USBGPU else tg_flags) if usbgpu else ('big_' if os.getenv('BIG_INTO_SMALL') else '', tg_flags)
|
||||
file_prefix, cmd_flags = ('big_', usbgpu_tg_flags) if usbgpu else ('big_' if os.getenv('BIG_INTO_SMALL') else '', tg_flags)
|
||||
driving_onnx_deps = get_existing_chunks(File(f"models/{file_prefix}driving_supercombo.onnx").abspath)
|
||||
camera_res_args = ' '.join(f'{cw}x{ch}' for cw, ch in CAMERA_CONFIGS)
|
||||
cmd = (f'{cmd_flags} {mac_brew_string} {sys.executable} {modeld_dir}/compile_modeld.py '
|
||||
cmd = (f'{cmd_flags} {mac_brew_string} python3 {modeld_dir}/compile_modeld.py '
|
||||
f'--model-size {model_w}x{model_h} '
|
||||
f'--camera-resolutions {camera_res_args} '
|
||||
f'--onnx {File(f"models/{file_prefix}driving_supercombo.onnx").abspath} '
|
||||
@@ -108,7 +104,7 @@ for usbgpu in [False, True] if USBGPU or WGPU else [False]:
|
||||
tinygrad_files + compile_modeld_script + driving_onnx_deps + [Value(chunk_targets), chunker_file],
|
||||
[cmd, Action(do_chunk, " [CHUNK] $TARGET")],
|
||||
)
|
||||
if usbgpu and USBGPU:
|
||||
if usbgpu:
|
||||
lenv.SideEffect(usbgpu_lock, node)
|
||||
|
||||
# get model metadata
|
||||
|
||||
@@ -26,7 +26,6 @@ from openpilot.common.file_chunker import open_file_chunked, get_manifest_path
|
||||
from openpilot.selfdrive.modeld.constants import ModelConstants, Plan
|
||||
from openpilot.selfdrive.modeld.helpers import usbgpu_present, modeld_pkl_path, get_tg_input_devices, load_oob
|
||||
from openpilot.selfdrive.modeld.usbgpu_link import wait_usbgpu_link
|
||||
from openpilot.tools.wgpu.zmq import WGPU_CAR_PARAMS, ZmqPubMaster, ZmqSubMaster, ZmqSubSocket
|
||||
|
||||
from openpilot.sunnypilot.livedelay.helpers import get_lat_delay
|
||||
from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase
|
||||
@@ -77,44 +76,26 @@ class FrameMeta:
|
||||
self.frame_id, self.timestamp_sof, self.timestamp_eof = vipc.frame_id, vipc.timestamp_sof, vipc.timestamp_eof
|
||||
|
||||
|
||||
def copy_nv12_to_venus(buf: VisionBuf, dst: np.ndarray, nv12: tuple[int, int, int, int]) -> None:
|
||||
stride, y_height, uv_height, _ = nv12
|
||||
if buf.stride < buf.width:
|
||||
raise ValueError(f"invalid VisionIPC stride {buf.stride} for width {buf.width}")
|
||||
|
||||
src = np.frombuffer(buf.data, dtype=np.uint8)
|
||||
src_size = buf.uv_offset + buf.stride * (buf.height // 2)
|
||||
if src.size < src_size:
|
||||
raise ValueError(f"VisionIPC buffer has {src.size} bytes, expected at least {src_size}")
|
||||
|
||||
dst[:stride * y_height].reshape(y_height, stride)[:buf.height, :buf.width] = \
|
||||
src[:buf.stride * buf.height].reshape(buf.height, buf.stride)[:, :buf.width]
|
||||
dst[stride * y_height:stride * (y_height + uv_height)].reshape(uv_height, stride)[:buf.height // 2, :buf.width] = \
|
||||
src[buf.uv_offset:src_size].reshape(buf.height // 2, buf.stride)[:, :buf.width]
|
||||
|
||||
|
||||
class ModelState(ModelStateBase):
|
||||
prev_desire: np.ndarray # for tracking the rising edge of the pulse
|
||||
|
||||
def __init__(self, cam_w: int, cam_h: int, usbgpu: bool, big_model: bool = False):
|
||||
def __init__(self, cam_w: int, cam_h: int, usbgpu: bool):
|
||||
ModelStateBase.__init__(self)
|
||||
self.LAT_SMOOTH_SECONDS = LAT_SMOOTH_SECONDS
|
||||
input_devices = get_tg_input_devices(PROCESS_NAME, usbgpu)
|
||||
self.WARP_DEV, self.QUEUE_DEV = input_devices['WARP_DEV'], input_devices['QUEUE_DEV']
|
||||
jits = load_oob(open_file_chunked(modeld_pkl_path(usbgpu or big_model)))
|
||||
jits = load_oob(open_file_chunked(modeld_pkl_path(usbgpu)))
|
||||
metadata = jits['metadata']
|
||||
self.input_shapes = metadata['input_shapes']
|
||||
self.vision_input_names = [k for k in self.input_shapes if 'img' in k]
|
||||
self.output_slices = metadata['output_slices']
|
||||
|
||||
self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32)
|
||||
self.copy_vision_buffers = self.WARP_DEV.split(":")[0] == "METAL"
|
||||
|
||||
self.frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
|
||||
self.input_queues, self.npy = make_input_queues(self.input_shapes, self.frame_skip, device=self.QUEUE_DEV)
|
||||
self.full_frames: dict[str, Tensor] = {}
|
||||
self._blob_cache: dict[tuple[str, int], Tensor] = {}
|
||||
self._vision_staging: dict[str, np.ndarray] = {}
|
||||
self.parser = Parser()
|
||||
self.frame_buf_params = {k: get_nv12_info(cam_w, cam_h) for k in ('img', 'big_img')}
|
||||
self.run_policy = jits['run_policy']
|
||||
@@ -127,22 +108,13 @@ class ModelState(ModelStateBase):
|
||||
def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray],
|
||||
inputs: dict[str, np.ndarray]) -> dict[str, np.ndarray] | None:
|
||||
for key in bufs.keys():
|
||||
nv12 = self.frame_buf_params[key]
|
||||
yuv_size = nv12[3]
|
||||
if self.copy_vision_buffers:
|
||||
# VisionIPC supplies a CPU pointer. Metal's external_ptr expects an MTLBuffer object,
|
||||
# so wrap its NV12 pixels in the Venus layout expected by the compiled warp, then copy.
|
||||
if key not in self._vision_staging:
|
||||
self._vision_staging[key] = np.zeros(yuv_size, dtype=np.uint8)
|
||||
copy_nv12_to_venus(bufs[key], self._vision_staging[key], nv12)
|
||||
self.full_frames[key] = Tensor(self._vision_staging[key], device=self.WARP_DEV).realize()
|
||||
else:
|
||||
# There is a ringbuffer of imgs, just cache tensors pointing to all of them.
|
||||
frame = np.frombuffer(bufs[key].data, dtype=np.uint8, count=yuv_size)
|
||||
cache_key = (key, frame.ctypes.data)
|
||||
if cache_key not in self._blob_cache:
|
||||
self._blob_cache[cache_key] = Tensor.from_blob(frame.ctypes.data, (yuv_size,), dtype='uint8', device=self.WARP_DEV)
|
||||
self.full_frames[key] = self._blob_cache[cache_key]
|
||||
ptr = np.frombuffer(bufs[key].data, dtype=np.uint8).ctypes.data
|
||||
yuv_size = self.frame_buf_params[key][3]
|
||||
# There is a ringbuffer of imgs, just cache tensors pointing to all of them
|
||||
cache_key = (key, ptr)
|
||||
if cache_key not in self._blob_cache:
|
||||
self._blob_cache[cache_key] = Tensor.from_blob(ptr, (yuv_size,), dtype='uint8', device=self.WARP_DEV)
|
||||
self.full_frames[key] = self._blob_cache[cache_key]
|
||||
|
||||
# Model decides when action is completed, so desire input is just a pulse triggered on rising edge
|
||||
inputs['desire_pulse'][0] = 0
|
||||
@@ -167,32 +139,18 @@ class ModelState(ModelStateBase):
|
||||
return outputs_dict
|
||||
|
||||
|
||||
def main(demo=False, remote_addr: str | None = None, big_model: bool = False):
|
||||
def main(demo=False):
|
||||
cloudlog.warning("modeld init")
|
||||
|
||||
_present = usbgpu_present()
|
||||
_compiled = os.path.isfile(get_manifest_path(modeld_pkl_path(usbgpu=True)))
|
||||
USBGPU = _present and _compiled
|
||||
if big_model and not _compiled:
|
||||
raise FileNotFoundError(f"big model is not compiled: {modeld_pkl_path(usbgpu=True)}")
|
||||
params = Params()
|
||||
params.put_bool("UsbGpuPresent", _present)
|
||||
params.put_bool("UsbGpuCompiled", _compiled)
|
||||
|
||||
config_realtime_process(7, 54)
|
||||
|
||||
remote_CP = None
|
||||
if remote_addr is not None:
|
||||
# Do not attach to VisionIPC until all startup prerequisites are available.
|
||||
# Otherwise its notification queue grows while waiting for the infrequent
|
||||
# bridged carParams message and reconnect starts tens of seconds behind.
|
||||
cloudlog.warning("waiting for remote carParams")
|
||||
car_params_socket = ZmqSubSocket(WGPU_CAR_PARAMS, remote_addr, conflate=True)
|
||||
raw_car_params = car_params_socket.receive()
|
||||
assert raw_car_params is not None
|
||||
remote_CP = messaging.log_from_bytes(raw_car_params, car.CarParams)
|
||||
cloudlog.info("modeld got remote CarParams: %s", remote_CP.brand)
|
||||
|
||||
# visionipc clients
|
||||
while True:
|
||||
available_streams = VisionIpcClient.available_streams("camerad", block=False)
|
||||
@@ -220,14 +178,12 @@ def main(demo=False, remote_addr: str | None = None, big_model: bool = False):
|
||||
wait_usbgpu_link()
|
||||
st = time.monotonic()
|
||||
cloudlog.warning("loading model")
|
||||
model = ModelState(vipc_client_main.width, vipc_client_main.height, USBGPU, big_model=big_model)
|
||||
model = ModelState(vipc_client_main.width, vipc_client_main.height, USBGPU)
|
||||
cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting")
|
||||
|
||||
# messaging
|
||||
output_services = ["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"]
|
||||
pm = ZmqPubMaster(output_services) if remote_addr is not None else PubMaster(output_services)
|
||||
services = ["deviceState", "carState", "roadCameraState", "liveCalibration", "driverMonitoringState", "carControl", "liveDelay"]
|
||||
sm = ZmqSubMaster(services, remote_addr) if remote_addr is not None else SubMaster(services)
|
||||
pm = PubMaster(["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"])
|
||||
sm = SubMaster(["deviceState", "carState", "roadCameraState", "liveCalibration", "driverMonitoringState", "carControl", "liveDelay"])
|
||||
|
||||
publish_state = PublishState()
|
||||
params = Params()
|
||||
@@ -247,9 +203,6 @@ def main(demo=False, remote_addr: str | None = None, big_model: bool = False):
|
||||
|
||||
if demo:
|
||||
CP = get_demo_car_params()
|
||||
elif remote_addr is not None:
|
||||
assert remote_CP is not None
|
||||
CP = remote_CP
|
||||
else:
|
||||
CP = messaging.log_from_bytes(params.get("CarParams", block=True), car.CarParams)
|
||||
cloudlog.info("modeld got CarParams: %s", CP.brand)
|
||||
@@ -379,9 +332,7 @@ if __name__ == "__main__":
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--demo', action='store_true', help='A boolean for demo mode.')
|
||||
parser.add_argument('--remote', metavar='ADDRESS', help='Run against a remote device over the cereal ZMQ bridge.')
|
||||
parser.add_argument('--big-model', action='store_true', help='Use the locally compiled big driving model.')
|
||||
args = parser.parse_args()
|
||||
main(demo=args.demo, remote_addr=args.remote, big_model=args.big_model)
|
||||
main(demo=args.demo)
|
||||
except KeyboardInterrupt:
|
||||
cloudlog.warning("got SIGINT")
|
||||
|
||||
@@ -722,6 +722,7 @@ EVENTS: dict[int, dict[str, Alert | AlertCallbackType]] = {
|
||||
ET.NO_ENTRY: NoEntryAlert("Driving Model Lagging"),
|
||||
ET.PERMANENT: modeld_lagging_alert,
|
||||
},
|
||||
|
||||
# Besides predicting the path, lane lines and lead car data the model also
|
||||
# predicts the current velocity and rotation speed of the car. If the model is
|
||||
# very uncertain about the current velocity while the car is moving, this
|
||||
|
||||
@@ -111,7 +111,6 @@ class SelfdriveD(CruiseHelper):
|
||||
self.is_metric = self.params.get_bool("IsMetric")
|
||||
self.is_ldw_enabled = self.params.get_bool("IsLdwEnabled")
|
||||
self.disengage_on_accelerator = self.params.get_bool("DisengageOnAccelerator")
|
||||
self.wgpu_enabled = self.params.get_bool("WgpuEnabled")
|
||||
|
||||
car_recognized = self.CP.brand != 'mock'
|
||||
|
||||
@@ -399,12 +398,12 @@ class SelfdriveD(CruiseHelper):
|
||||
has_disable_events = self.events.contains(ET.NO_ENTRY) and (self.events.contains(ET.SOFT_DISABLE) or self.events.contains(ET.IMMEDIATE_DISABLE))
|
||||
no_system_errors = (not has_disable_events) or (len(self.events) == num_events)
|
||||
if not self.sm.all_checks() and no_system_errors:
|
||||
# if not self.sm.all_alive():
|
||||
# self.events.add(EventName.commIssue)
|
||||
# elif not self.sm.all_freq_ok():
|
||||
# self.events.add(EventName.commIssueAvgFreq)
|
||||
# else:
|
||||
# self.events.add(EventName.commIssue)
|
||||
if not self.sm.all_alive():
|
||||
self.events.add(EventName.commIssue)
|
||||
elif not self.sm.all_freq_ok():
|
||||
self.events.add(EventName.commIssueAvgFreq)
|
||||
else:
|
||||
self.events.add(EventName.commIssue)
|
||||
|
||||
logs = {
|
||||
'invalid': [s for s, valid in self.sm.valid.items() if not valid],
|
||||
@@ -417,13 +416,13 @@ class SelfdriveD(CruiseHelper):
|
||||
else:
|
||||
self.logged_comm_issue = None
|
||||
|
||||
# if not self.CP.notCar:
|
||||
# if not self.sm['livePose'].posenetOK:
|
||||
# self.events.add(EventName.posenetInvalid)
|
||||
# if not self.sm['livePose'].inputsOK:
|
||||
# self.events.add(EventName.locationdTemporaryError)
|
||||
# if not self.sm['liveParameters'].valid and cal_status == log.LiveCalibrationData.Status.calibrated and not TESTING_CLOSET and (not SIMULATION or REPLAY):
|
||||
# self.events.add(EventName.paramsdTemporaryError)
|
||||
if not self.CP.notCar:
|
||||
if not self.sm['livePose'].posenetOK:
|
||||
self.events.add(EventName.posenetInvalid)
|
||||
if not self.sm['livePose'].inputsOK:
|
||||
self.events.add(EventName.locationdTemporaryError)
|
||||
if not self.sm['liveParameters'].valid and cal_status == log.LiveCalibrationData.Status.calibrated and not TESTING_CLOSET and (not SIMULATION or REPLAY):
|
||||
self.events.add(EventName.paramsdTemporaryError)
|
||||
|
||||
# conservative HW alert. if the data or frequency are off, locationd will throw an error
|
||||
if any((self.sm.frame - self.sm.recv_frame[s])*DT_CTRL > 10. for s in self.sensor_packets):
|
||||
@@ -468,9 +467,9 @@ class SelfdriveD(CruiseHelper):
|
||||
self.distance_traveled += abs(CS.vEgo) * DT_CTRL
|
||||
|
||||
# TODO: fix simulator
|
||||
# if not SIMULATION or REPLAY:
|
||||
# if self.sm['modelV2'].frameDropPerc > 1 and not self.wgpu_enabled:
|
||||
# self.events.add(EventName.modeldLagging)
|
||||
if not SIMULATION or REPLAY:
|
||||
if self.sm['modelV2'].frameDropPerc > 1:
|
||||
self.events.add(EventName.modeldLagging)
|
||||
|
||||
# mute canBusMissing event if in Park, as it sometimes may trigger a false alarm with MADS in Paused state
|
||||
if CS.gearShifter == car.CarState.GearShifter.park and self.mads.enabled:
|
||||
@@ -626,7 +625,6 @@ class SelfdriveD(CruiseHelper):
|
||||
self.is_metric = self.params.get_bool("IsMetric")
|
||||
self.is_ldw_enabled = self.params.get_bool("IsLdwEnabled")
|
||||
self.disengage_on_accelerator = self.params.get_bool("DisengageOnAccelerator")
|
||||
self.wgpu_enabled = self.params.get_bool("WgpuEnabled")
|
||||
self.experimental_mode = self.params.get_bool("ExperimentalMode") and self.CP.openpilotLongitudinalControl
|
||||
self.personality = self.params.get("LongitudinalPersonality", return_default=True)
|
||||
|
||||
|
||||
@@ -248,10 +248,8 @@ class MiciHomeLayout(Widget):
|
||||
|
||||
# ***** Center-aligned bottom section icons *****
|
||||
self._experimental_icon.set_visible(ui_state.experimental_mode)
|
||||
wgpu_running = ui_state.wgpu_enabled and ui_state.sm.alive["modelV2"] and ui_state.sm.valid["modelV2"]
|
||||
self._egpu_icon.set_visible((ui_state.usbgpu and ui_state.usbgpu_compiled) or wgpu_running)
|
||||
self._egpu_icon_gray.set_visible((ui_state.usbgpu and not ui_state.usbgpu_compiled) or
|
||||
(ui_state.wgpu_ready and not wgpu_running))
|
||||
self._egpu_icon.set_visible(ui_state.usbgpu and ui_state.usbgpu_compiled)
|
||||
self._egpu_icon_gray.set_visible(ui_state.usbgpu and not ui_state.usbgpu_compiled)
|
||||
self._mic_icon.set_visible(ui_state.recording_audio)
|
||||
self._body_icon.set_visible(bool(ui_state.is_body))
|
||||
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import pyray as rl
|
||||
from openpilot.cereal import log
|
||||
@@ -162,20 +160,6 @@ class AugmentedRoadView(CameraView):
|
||||
text_color=rl.Color(255, 255, 255, int(255 * 0.9)),
|
||||
alignment=rl.GuiTextAlignment.TEXT_ALIGN_CENTER,
|
||||
alignment_vertical=rl.GuiTextAlignmentVertical.TEXT_ALIGN_MIDDLE)
|
||||
self._model_status_key_labels = [
|
||||
UnifiedLabel("", 21, FontWeight.ROMAN, text_color=rl.Color(210, 210, 210, 220),
|
||||
alignment=rl.GuiTextAlignment.TEXT_ALIGN_LEFT,
|
||||
alignment_vertical=rl.GuiTextAlignmentVertical.TEXT_ALIGN_MIDDLE,
|
||||
wrap_text=False)
|
||||
for _ in range(6)
|
||||
]
|
||||
self._model_status_value_labels = [
|
||||
UnifiedLabel("", 22, FontWeight.SEMI_BOLD,
|
||||
alignment=rl.GuiTextAlignment.TEXT_ALIGN_RIGHT,
|
||||
alignment_vertical=rl.GuiTextAlignmentVertical.TEXT_ALIGN_MIDDLE,
|
||||
wrap_text=False)
|
||||
for _ in range(6)
|
||||
]
|
||||
|
||||
self._fade_texture = gui_app.texture("icons_mici/onroad/onroad_fade.png")
|
||||
|
||||
@@ -246,8 +230,6 @@ class AugmentedRoadView(CameraView):
|
||||
self._driver_state_renderer.set_position(self._rect.x + 16, self._rect.y + 10)
|
||||
self._driver_state_renderer.render()
|
||||
|
||||
self._render_model_status()
|
||||
|
||||
self._hud_renderer.set_can_draw_top_icons(alert_to_render is None)
|
||||
self._hud_renderer.set_wheel_critical_icon(alert_to_render is not None and not not_animating_out and
|
||||
alert_to_render.visual_alert == car.CarControl.HUDControl.VisualAlert.steerRequired)
|
||||
@@ -266,58 +248,6 @@ class AugmentedRoadView(CameraView):
|
||||
|
||||
self._bookmark_icon.render(self.rect)
|
||||
|
||||
def _render_model_status(self):
|
||||
model = ui_state.sm["modelV2"]
|
||||
if ui_state.sm.seen["modelV2"] and model.timestampEof:
|
||||
model_age_ms = max(0., (time.monotonic_ns() - model.timestampEof) / 1e6)
|
||||
execution_ms = max(0., model.modelExecutionTime * 1e3)
|
||||
io_queue_ms = max(0., model_age_ms - execution_ms)
|
||||
frame_drop = model.frameDropPerc
|
||||
age_text = f"{model_age_ms:.0f} ms"
|
||||
execution_text = f"{execution_ms:.0f} ms"
|
||||
io_queue_text = f"{io_queue_ms:.0f} ms"
|
||||
frame_drop_text = f"{frame_drop:.1f}%"
|
||||
else:
|
||||
model_age_ms = float("inf")
|
||||
age_text = execution_text = io_queue_text = "-- ms"
|
||||
frame_drop_text = "--%"
|
||||
|
||||
if ui_state.wgpu_enabled:
|
||||
source, model_name = "WGPU", ui_state.wgpu_model_name
|
||||
else:
|
||||
source, model_name = "LOCAL", "DEVICE"
|
||||
|
||||
if model_age_ms < 200:
|
||||
color = rl.Color(100, 255, 120, 230)
|
||||
elif model_age_ms < 500:
|
||||
color = rl.Color(255, 210, 80, 230)
|
||||
else:
|
||||
color = rl.Color(255, 120, 80, 230)
|
||||
|
||||
panel_w, panel_h = 250, 178
|
||||
status_rect = rl.Rectangle(self._content_rect.x + self._content_rect.width - panel_w - 18,
|
||||
self._content_rect.y + 18, panel_w, panel_h)
|
||||
rl.draw_rectangle_rounded(status_rect, 0.14, 8, rl.Color(0, 0, 0, 175))
|
||||
|
||||
rows = (
|
||||
("SOURCE", source),
|
||||
("MODEL", model_name),
|
||||
("AGE", age_text),
|
||||
("EXEC", execution_text),
|
||||
("IO/QUEUE", io_queue_text),
|
||||
("DROPPED", frame_drop_text),
|
||||
)
|
||||
row_h = 26
|
||||
for i, ((key, value), key_label, value_label) in enumerate(
|
||||
zip(rows, self._model_status_key_labels, self._model_status_value_labels, strict=True)
|
||||
):
|
||||
row_rect = rl.Rectangle(status_rect.x + 14, status_rect.y + 11 + i * row_h, status_rect.width - 28, row_h)
|
||||
key_label.set_text(key)
|
||||
key_label.render(rl.Rectangle(row_rect.x, row_rect.y, 105, row_rect.height))
|
||||
value_label.set_text(value)
|
||||
value_label.set_text_color(color)
|
||||
value_label.render(rl.Rectangle(row_rect.x + 108, row_rect.y, row_rect.width - 108, row_rect.height))
|
||||
|
||||
def _switch_stream_if_needed(self, sm):
|
||||
if sm['selfdriveState'].experimentalMode and WIDE_CAM in self.available_streams:
|
||||
v_ego = sm['carState'].vEgo
|
||||
|
||||
@@ -43,7 +43,7 @@ class ModelsLayout(Widget):
|
||||
self._initialize_items()
|
||||
|
||||
self.clear_cache_item.action_item.set_value(f"{self.calculate_cache_size():.2f} MB")
|
||||
for ctrl, key in [(self.lane_turn_value_control, "LaneTurnValue"), (self.delay_control, "LagdToggleDelay")]:
|
||||
for ctrl, key in [(self.lane_turn_value_control, "LaneTurnValue"), (self.delay_control, "LagdToggleDelay"), (self.camera_offset, "CameraOffset")]:
|
||||
ctrl.action_item.set_value(int(float(ui_state.params.get(key, return_default=True)) * 100))
|
||||
|
||||
self._scroller = Scroller(self.items, line_separator=True, spacing=0)
|
||||
@@ -93,9 +93,14 @@ class ModelsLayout(Widget):
|
||||
|
||||
self.lagd_toggle = toggle_item_sp(tr("Live Learning Steer Delay"), "", param="LagdToggle")
|
||||
|
||||
self.camera_offset = option_item_sp(tr("Adjust Camera Offset"), "CameraOffset", -35, 35,
|
||||
tr("Virtually shift camera's perspective to move model's center to Left(+ values) or Right (- values)"),
|
||||
1, None, True, "", style.BUTTON_ACTION_WIDTH, None, True,
|
||||
lambda v: f"{v / 100:.2f} m")
|
||||
|
||||
self.items = [self.current_model_item, self.cancel_download_item, self.supercombo_label, self.vision_label,
|
||||
self.policy_label, self.off_policy_label, self.on_policy_label, self.refresh_item, self.clear_cache_item, self.lane_turn_desire_toggle,
|
||||
self.lane_turn_value_control, self.lagd_toggle, self.delay_control]
|
||||
self.policy_label, self.off_policy_label, self.on_policy_label, self.refresh_item, self.clear_cache_item,
|
||||
self.lane_turn_desire_toggle, self.lane_turn_value_control, self.lagd_toggle, self.delay_control, self.camera_offset]
|
||||
|
||||
def _update_lagd_description(self, lagd_toggle: bool):
|
||||
desc = tr("Enable this for the car to learn and adapt its steering response time. Disable to use a fixed steering response time. " +
|
||||
@@ -232,6 +237,7 @@ class ModelsLayout(Widget):
|
||||
advanced_controls: bool = ui_state.params.get_bool("ShowAdvancedControls")
|
||||
turn_desire: bool = ui_state.params.get_bool("LaneTurnDesire")
|
||||
live_delay: bool = ui_state.params.get_bool("LagdToggle")
|
||||
camera_offset: bool = ui_state.params.get("ModelManager_ActiveBundle") is not None
|
||||
|
||||
self.lane_turn_desire_toggle.action_item.set_state(turn_desire)
|
||||
self.lane_turn_value_control.set_visible(turn_desire and advanced_controls)
|
||||
@@ -240,6 +246,7 @@ class ModelsLayout(Widget):
|
||||
new_step = int(round(100 / CV.MPH_TO_KPH)) if ui_state.is_metric else 100
|
||||
if self.lane_turn_value_control.action_item is not None and self.lane_turn_value_control.action_item.value_change_step != new_step:
|
||||
self.lane_turn_value_control.action_item.value_change_step = new_step
|
||||
self.camera_offset.set_visible(camera_offset)
|
||||
|
||||
self._update_lagd_description(live_delay)
|
||||
self.model_manager = ui_state.sm["modelManagerSP"]
|
||||
|
||||
@@ -82,9 +82,6 @@ class UIState(UIStateSP):
|
||||
self.experimental_mode: bool = self.params.get_bool("ExperimentalMode")
|
||||
self.usbgpu: bool = self.params.get_bool("UsbGpuPresent")
|
||||
self.usbgpu_compiled: bool = self.params.get_bool("UsbGpuCompiled")
|
||||
self.wgpu_enabled: bool = self.params.get_bool("WgpuEnabled")
|
||||
self.wgpu_model_name: str = self.params.get("WgpuModelName") or "UNKNOWN"
|
||||
self.wgpu_ready: bool = self.params.get_bool("WgpuReady")
|
||||
self.started: bool = False
|
||||
self.ignition: bool = False
|
||||
self.recording_audio: bool = False
|
||||
@@ -216,9 +213,6 @@ class UIState(UIStateSP):
|
||||
self.experimental_mode = self.params.get_bool("ExperimentalMode")
|
||||
self.usbgpu = self.params.get_bool("UsbGpuPresent")
|
||||
self.usbgpu_compiled = self.params.get_bool("UsbGpuCompiled")
|
||||
self.wgpu_enabled = self.params.get_bool("WgpuEnabled")
|
||||
self.wgpu_model_name = self.params.get("WgpuModelName") or "UNKNOWN"
|
||||
self.wgpu_ready = self.params.get_bool("WgpuReady")
|
||||
|
||||
UIStateSP.update_params(self)
|
||||
|
||||
|
||||
@@ -10,471 +10,291 @@ import argparse
|
||||
import os
|
||||
import pickle
|
||||
import time
|
||||
from functools import partial
|
||||
from collections import defaultdict
|
||||
|
||||
from functools import partial
|
||||
import numpy as np
|
||||
from tinygrad.tensor import Tensor
|
||||
os.environ['GMMU'] = '0'
|
||||
|
||||
def _patch_tinygrad_fetch_fw():
|
||||
import hashlib
|
||||
import pathlib
|
||||
import zstandard
|
||||
from tinygrad import helpers
|
||||
_orig_fetch_fw = helpers.fetch_fw
|
||||
def fetch_fw(path, name, sha256):
|
||||
p = pathlib.Path(f"/lib/firmware/{path}/{name}.zst")
|
||||
if p.is_file():
|
||||
blob = zstandard.ZstdDecompressor().stream_reader(p.read_bytes()).read()
|
||||
if hashlib.sha256(blob).hexdigest() == sha256:
|
||||
return blob
|
||||
return _orig_fetch_fw(path, name, sha256)
|
||||
helpers.fetch_fw = fetch_fw
|
||||
_patch_tinygrad_fetch_fw()
|
||||
|
||||
from openpilot.selfdrive.modeld.compile_modeld import NV12Frame, make_frame_prepare, sample_desire, sample_skip, shift_and_sample
|
||||
from tinygrad import dtypes
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.engine.jit import TinyJit
|
||||
|
||||
from openpilot.selfdrive.modeld.compile_modeld import (
|
||||
NV12Frame, make_frame_prepare,
|
||||
shift_and_sample, sample_skip, sample_desire,
|
||||
)
|
||||
from tinygrad.tensor import Tensor
|
||||
|
||||
MODEL_TYPES = ('vision_policy', 'supercombo', 'vision_multi_policy')
|
||||
|
||||
|
||||
def _detect_desire_key(policy_input_shapes):
|
||||
for k in policy_input_shapes:
|
||||
if k.startswith('desire'):
|
||||
return k
|
||||
return None
|
||||
def _detect_desire_key(shapes: dict) -> str | None:
|
||||
return next((key for key in shapes if key.startswith('desire')), None)
|
||||
|
||||
|
||||
def _detect_vision_keys(vision_input_shapes):
|
||||
img_keys = sorted([k for k in vision_input_shapes if 'img' in k])
|
||||
road_key = next((k for k in img_keys if 'big' not in k), None)
|
||||
wide_key = next((k for k in img_keys if 'big' in k), None)
|
||||
if road_key is None or wide_key is None:
|
||||
raise ValueError(f"Cannot determine road/wide image keys from {list(vision_input_shapes.keys())}")
|
||||
return road_key, wide_key
|
||||
def _detect_vision_keys(shapes: dict) -> tuple[str | None, str | None]:
|
||||
img_keys = sorted(key for key in shapes if 'img' in key)
|
||||
return (
|
||||
next((key for key in img_keys if 'big' not in key), None),
|
||||
next((key for key in img_keys if 'big' in key), None)
|
||||
)
|
||||
|
||||
|
||||
def make_split_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, device):
|
||||
road_key, _ = _detect_vision_keys(vision_input_shapes)
|
||||
img = vision_input_shapes[road_key]
|
||||
n_frames = img[1] // 6
|
||||
img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img[2], img[3])
|
||||
|
||||
fb = policy_input_shapes['features_buffer']
|
||||
desire_key = _detect_desire_key(policy_input_shapes)
|
||||
dp = policy_input_shapes[desire_key]
|
||||
tc = policy_input_shapes.get('traffic_convention', (1, 2))
|
||||
|
||||
npy = {
|
||||
'desire': np.zeros(dp[2], dtype=np.float32),
|
||||
'traffic_convention': np.zeros(tc, dtype=np.float32),
|
||||
'tfm': np.zeros((3, 3), dtype=np.float32),
|
||||
'big_tfm': np.zeros((3, 3), dtype=np.float32),
|
||||
}
|
||||
|
||||
handled = {'features_buffer', desire_key, 'traffic_convention'}
|
||||
for key, shape in policy_input_shapes.items():
|
||||
if key in handled:
|
||||
continue
|
||||
npy[key] = np.zeros(shape, dtype=np.float32)
|
||||
|
||||
input_queues = {
|
||||
'img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
|
||||
'big_img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
|
||||
'feat_q': Tensor(np.zeros((frame_skip * (fb[1] - 1) + 1, fb[0], fb[2]), dtype=np.float32), device=device).contiguous().realize(),
|
||||
'desire_q': Tensor(np.zeros((frame_skip * dp[1], dp[0], dp[2]), dtype=np.float32), device=device).contiguous().realize(),
|
||||
**{k: Tensor(v, device='NPY').realize() for k, v in npy.items()},
|
||||
}
|
||||
return input_queues, npy
|
||||
def derive_frame_skip(vision_input_shapes: dict, policy_input_shapes: dict) -> int:
|
||||
features_buffer = policy_input_shapes.get('features_buffer')
|
||||
return 1 if not features_buffer or features_buffer[1] >= 99 else 4
|
||||
|
||||
|
||||
def make_run_split_policy(vision_runner, policy_runner, nv12: NV12Frame, model_w, model_h,
|
||||
vision_features_slice, frame_skip, desire_key, extra_policy_keys,
|
||||
vision_road_key, vision_wide_key, prepare_only=False):
|
||||
frame_prepare = make_frame_prepare(nv12, model_w, model_h)
|
||||
sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
|
||||
sample_desire_fn = partial(sample_desire, frame_skip=frame_skip)
|
||||
|
||||
def run_policy(img_q, big_img_q, feat_q, desire_q, desire, traffic_convention, tfm, big_tfm, frame, big_frame, **extra):
|
||||
npy_tensors = [tfm.to(Device.DEFAULT), big_tfm.to(Device.DEFAULT),
|
||||
desire.to(Device.DEFAULT), traffic_convention.to(Device.DEFAULT)]
|
||||
extra_device = {k: extra[k].to(Device.DEFAULT) for k in extra_policy_keys}
|
||||
Tensor.realize(*npy_tensors, *extra_device.values())
|
||||
tfm, big_tfm, desire, traffic_convention = npy_tensors
|
||||
|
||||
img = shift_and_sample(img_q, frame_prepare(frame, tfm).unsqueeze(0), sample_skip_fn)
|
||||
big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm).unsqueeze(0), sample_skip_fn)
|
||||
|
||||
if prepare_only:
|
||||
return img, big_img
|
||||
|
||||
vision_out = next(iter(vision_runner({vision_road_key: img, vision_wide_key: big_img}).values())).cast('float32')
|
||||
|
||||
new_feat = vision_out[:, vision_features_slice].reshape(1, -1).unsqueeze(0)
|
||||
feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn)
|
||||
desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn)
|
||||
|
||||
inputs = {'features_buffer': feat_buf, desire_key: desire_buf, 'traffic_convention': traffic_convention, **extra_device}
|
||||
policy_out = next(iter(policy_runner(inputs).values())).cast('float32')
|
||||
|
||||
return vision_out, policy_out
|
||||
return run_policy
|
||||
|
||||
|
||||
def compile_split_policy(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip,
|
||||
vision_runner, policy_runner, vision_metadata, policy_metadata):
|
||||
print(f"Compiling combined policy JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...")
|
||||
|
||||
vision_features_slice = vision_metadata['output_slices']['hidden_state']
|
||||
vision_input_shapes = vision_metadata['input_shapes']
|
||||
policy_input_shapes = policy_metadata['input_shapes']
|
||||
desire_key = _detect_desire_key(policy_input_shapes)
|
||||
extra_policy_keys = [k for k in policy_input_shapes if k not in ('features_buffer', desire_key, 'traffic_convention')]
|
||||
vision_road_key, vision_wide_key = _detect_vision_keys(vision_input_shapes)
|
||||
|
||||
_run = make_run_split_policy(vision_runner, policy_runner, nv12, model_w, model_h,
|
||||
vision_features_slice, frame_skip, desire_key, extra_policy_keys,
|
||||
vision_road_key, vision_wide_key, prepare_only)
|
||||
run_policy_jit = TinyJit(_run, prune=True)
|
||||
|
||||
SEED = 42
|
||||
|
||||
def random_inputs_run_fn(fn, seed, test_val=None, test_buffers=None, expect_match=True):
|
||||
input_queues, npy = make_split_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, Device.DEFAULT)
|
||||
rng = np.random.default_rng(seed)
|
||||
Tensor.manual_seed(seed)
|
||||
|
||||
testing = test_val is not None or test_buffers is not None
|
||||
n_runs = 1 if testing else 3
|
||||
|
||||
for i in range(n_runs):
|
||||
frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize()
|
||||
big_frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize()
|
||||
for v in npy.values():
|
||||
v[:] = rng.standard_normal(v.shape).astype(v.dtype)
|
||||
Device.default.synchronize()
|
||||
st = time.perf_counter()
|
||||
outs = fn(**input_queues, frame=frame, big_frame=big_frame)
|
||||
mt = time.perf_counter()
|
||||
Device.default.synchronize()
|
||||
et = time.perf_counter()
|
||||
print(f" [{i+1}/{n_runs}] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms")
|
||||
|
||||
if i == 0:
|
||||
val = [np.copy(v.numpy()) for v in outs]
|
||||
buffers = [np.copy(v.numpy().copy()) for v in input_queues.values()]
|
||||
|
||||
if test_val is not None:
|
||||
match = all(np.array_equal(a, b) for a, b in zip(val, test_val, strict=True))
|
||||
assert match == expect_match, f"outputs {'differ from' if expect_match else 'match'} baseline (seed={seed})"
|
||||
if test_buffers is not None:
|
||||
match = all(np.array_equal(a, b) for a, b in zip(buffers, test_buffers, strict=True))
|
||||
assert match == expect_match, f"buffers {'differ from' if expect_match else 'match'} baseline (seed={seed})"
|
||||
return fn, val, buffers
|
||||
|
||||
print('capture + replay')
|
||||
run_policy_jit, test_val, test_buffers = random_inputs_run_fn(run_policy_jit, SEED)
|
||||
|
||||
print('pickle round trip')
|
||||
run_policy_jit = pickle.loads(pickle.dumps(run_policy_jit))
|
||||
random_inputs_run_fn(run_policy_jit, SEED, test_val, test_buffers, expect_match=True)
|
||||
random_inputs_run_fn(run_policy_jit, SEED+1, test_val, test_buffers, expect_match=False)
|
||||
return run_policy_jit
|
||||
|
||||
|
||||
def derive_frame_skip(vision_input_shapes, policy_input_shapes):
|
||||
fb = policy_input_shapes.get('features_buffer')
|
||||
if fb is None:
|
||||
return 1
|
||||
fb_history = fb[1]
|
||||
if fb_history >= 99:
|
||||
return 1
|
||||
return 4
|
||||
|
||||
|
||||
def make_supercombo_input_queues(input_shapes, frame_skip, device):
|
||||
img_shape = input_shapes.get('img', input_shapes.get('input_imgs'))
|
||||
if img_shape is None:
|
||||
raise ValueError("No img input found in model shapes")
|
||||
def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = Device.DEFAULT) -> tuple[dict, dict]:
|
||||
road_key, _ = _detect_vision_keys(input_shapes)
|
||||
if not road_key:
|
||||
raise ValueError("Vision road key missing from input shapes.")
|
||||
|
||||
img_shape = input_shapes[road_key]
|
||||
n_frames = img_shape[1] // 6
|
||||
img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img_shape[2], img_shape[3])
|
||||
|
||||
numpy_keys = {}
|
||||
queue_keys = {}
|
||||
desire_key = _detect_desire_key(input_shapes)
|
||||
if not desire_key:
|
||||
raise ValueError("Desire key missing from input shapes.")
|
||||
|
||||
desire_shape = input_shapes[desire_key]
|
||||
features_buffer = input_shapes.get('features_buffer')
|
||||
|
||||
npy_arrays = {
|
||||
'desire': np.zeros(desire_shape[2], dtype=np.float32),
|
||||
'tfm': np.zeros((3, 3), dtype=np.float32),
|
||||
'big_tfm': np.zeros((3, 3), dtype=np.float32)
|
||||
}
|
||||
|
||||
for key, shape in input_shapes.items():
|
||||
if 'img' in key:
|
||||
continue
|
||||
if len(shape) == 3 and shape[1] > 1:
|
||||
if key.startswith('desire'):
|
||||
numpy_keys[key] = np.zeros(shape[2], dtype=np.float32)
|
||||
queue_keys[f'{key}_q'] = Tensor(
|
||||
np.zeros((frame_skip * shape[1], shape[0], shape[2]), dtype=np.float32),
|
||||
device=device).contiguous().realize()
|
||||
elif key == 'features_buffer':
|
||||
queue_keys['feat_q'] = Tensor(
|
||||
np.zeros((frame_skip * (shape[1] - 1) + 1, shape[0], shape[2]), dtype=np.float32),
|
||||
device=device).contiguous().realize()
|
||||
else:
|
||||
numpy_keys[key] = np.zeros(shape, dtype=np.float32)
|
||||
elif len(shape) == 2:
|
||||
numpy_keys[key] = np.zeros(shape, dtype=np.float32)
|
||||
if key not in npy_arrays and 'img' not in key and key not in ('features_buffer', desire_key):
|
||||
npy_arrays[key] = np.zeros(shape, dtype=np.float32)
|
||||
|
||||
if 'traffic_convention' not in numpy_keys:
|
||||
tc_shape = input_shapes.get('traffic_convention', (1, 2))
|
||||
numpy_keys['traffic_convention'] = np.zeros(tc_shape, dtype=np.float32)
|
||||
|
||||
numpy_keys['tfm'] = np.zeros((3, 3), dtype=np.float32)
|
||||
numpy_keys['big_tfm'] = np.zeros((3, 3), dtype=np.float32)
|
||||
|
||||
input_queues = {
|
||||
queues = {
|
||||
'img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
|
||||
'big_img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
|
||||
**queue_keys,
|
||||
**{k: Tensor(v, device='NPY').realize() for k, v in numpy_keys.items()},
|
||||
'desire_q': Tensor(np.zeros((frame_skip * desire_shape[1], desire_shape[0], desire_shape[2]),
|
||||
dtype=np.float32), device=device).contiguous().realize()
|
||||
}
|
||||
return input_queues, numpy_keys
|
||||
|
||||
if features_buffer:
|
||||
queues['feat_q'] = Tensor(np.zeros((frame_skip * (features_buffer[1] - 1) + 1, features_buffer[0], features_buffer[2]),
|
||||
dtype=np.float32), device=device).contiguous().realize()
|
||||
|
||||
queues.update({key: Tensor(value, device='NPY').realize() for key, value in npy_arrays.items()})
|
||||
return queues, npy_arrays
|
||||
|
||||
|
||||
def make_run_supercombo(model_runner, nv12: NV12Frame, model_w, model_h,
|
||||
features_slice, frame_skip, input_shapes, prepare_only=False):
|
||||
frame_prepare = make_frame_prepare(nv12, model_w, model_h)
|
||||
def make_split_input_queues(vision_input_shapes: dict, policy_input_shapes: dict, frame_skip: int, device: str = Device.DEFAULT) -> tuple[dict, dict]:
|
||||
return generate_queues_and_npy({**vision_input_shapes, **policy_input_shapes}, frame_skip, device)
|
||||
|
||||
|
||||
def make_supercombo_input_queues(input_shapes: dict, frame_skip: int, device: str = Device.DEFAULT) -> tuple[dict, dict]:
|
||||
return generate_queues_and_npy(input_shapes, frame_skip, device)
|
||||
|
||||
|
||||
def create_jit_runner(vision_runner, policy_runners: list, nv12: NV12Frame, model_size: tuple[int, int],
|
||||
features_slice: slice, frame_skip: int, input_shapes: dict, prepare_only: bool):
|
||||
frame_prepare = make_frame_prepare(nv12, *model_size)
|
||||
sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
|
||||
sample_desire_fn = partial(sample_desire, frame_skip=frame_skip)
|
||||
|
||||
desire_key = _detect_desire_key(input_shapes)
|
||||
if desire_key is None:
|
||||
raise ValueError(f"No desire* key found in input_shapes: {list(input_shapes.keys())}")
|
||||
road_img_key, wide_img_key = _detect_vision_keys(input_shapes)
|
||||
extra_policy_keys = [k for k in input_shapes
|
||||
if k not in (desire_key, 'features_buffer', 'traffic_convention')
|
||||
and 'img' not in k]
|
||||
road_key, wide_key = _detect_vision_keys(input_shapes)
|
||||
|
||||
def run_supercombo(img_q, big_img_q, feat_q, desire_q,
|
||||
frame, big_frame, **kwargs):
|
||||
desire = kwargs.get(desire_key)
|
||||
if not desire_key or not road_key or not wide_key:
|
||||
raise ValueError("Missing required vision or desire keys in input shapes.")
|
||||
|
||||
extra_keys = [key for key in input_shapes if key not in (desire_key, 'features_buffer', 'traffic_convention') and 'img' not in key]
|
||||
|
||||
def runner(img_q, big_img_q, feat_q, frame, big_frame, tfm, big_tfm, **kwargs):
|
||||
desire_q = kwargs['desire_q']
|
||||
desire = kwargs['desire']
|
||||
traffic_convention = kwargs.get('traffic_convention')
|
||||
tfm = kwargs['tfm']
|
||||
big_tfm = kwargs['big_tfm']
|
||||
|
||||
tfm = tfm.to(Device.DEFAULT)
|
||||
big_tfm = big_tfm.to(Device.DEFAULT)
|
||||
desire = desire.to(Device.DEFAULT)
|
||||
traffic_convention = traffic_convention.to(Device.DEFAULT)
|
||||
Tensor.realize(tfm, big_tfm, desire, traffic_convention)
|
||||
npys = [tfm.to(Device.DEFAULT), big_tfm.to(Device.DEFAULT), desire.to(Device.DEFAULT)]
|
||||
if traffic_convention is not None:
|
||||
npys.append(traffic_convention.to(Device.DEFAULT))
|
||||
|
||||
img = shift_and_sample(img_q, frame_prepare(frame, tfm).unsqueeze(0), sample_skip_fn)
|
||||
big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm).unsqueeze(0), sample_skip_fn)
|
||||
extra_tensors = {key: kwargs[key].to(Device.DEFAULT) for key in extra_keys if key in kwargs}
|
||||
Tensor.realize(*npys, *extra_tensors.values())
|
||||
|
||||
tfm_dev, big_tfm_dev, desire_dev = npys[:3]
|
||||
traffic_conv_dev = npys[3] if traffic_convention is not None else None
|
||||
|
||||
img = shift_and_sample(img_q, frame_prepare(frame, tfm_dev).unsqueeze(0), sample_skip_fn)
|
||||
big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm_dev).unsqueeze(0), sample_skip_fn)
|
||||
|
||||
if prepare_only:
|
||||
return img, big_img
|
||||
|
||||
desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn)
|
||||
feat_buf = sample_skip_fn(feat_q)
|
||||
desire_buf = shift_and_sample(desire_q, desire_dev.reshape(1, 1, -1), sample_desire_fn)
|
||||
inputs = {desire_key: desire_buf, **extra_tensors}
|
||||
|
||||
inputs = {road_img_key: img, wide_img_key: big_img,
|
||||
desire_key: desire_buf, 'features_buffer': feat_buf,
|
||||
'traffic_convention': traffic_convention}
|
||||
for k in extra_policy_keys:
|
||||
if k in kwargs:
|
||||
inputs[k] = kwargs[k].to(Device.DEFAULT)
|
||||
if traffic_conv_dev is not None:
|
||||
inputs['traffic_convention'] = traffic_conv_dev
|
||||
|
||||
model_out = next(iter(model_runner(inputs).values())).cast('float32')
|
||||
if vision_runner:
|
||||
vision_out = next(iter(vision_runner({road_key: img, wide_key: big_img}).values()))
|
||||
vision_out_cast = vision_out.cast('float32')
|
||||
new_feat = vision_out_cast[:, features_slice].reshape(1, -1).unsqueeze(0)
|
||||
inputs['features_buffer'] = shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
|
||||
policy_outs = [next(iter(pol_runner(inputs).values())).cast('float32') for pol_runner in policy_runners]
|
||||
return (vision_out_cast, *policy_outs) if len(policy_outs) > 1 else (vision_out_cast, policy_outs[0])
|
||||
inputs.update({road_key: img, wide_key: big_img, 'features_buffer': sample_skip_fn(feat_q)})
|
||||
policy_out = next(iter(policy_runners[0](inputs).values())).cast('float32')
|
||||
new_feat = policy_out[:, features_slice].reshape(1, -1).unsqueeze(0)
|
||||
shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
|
||||
return policy_out
|
||||
|
||||
new_feat = model_out[:, features_slice].reshape(1, -1).unsqueeze(0)
|
||||
shift_and_sample(feat_q, new_feat, sample_skip_fn)
|
||||
|
||||
return model_out
|
||||
|
||||
return run_supercombo
|
||||
return runner
|
||||
|
||||
|
||||
def make_run_vision_multi_policy(vision_runner, policy_runners, nv12: NV12Frame, model_w, model_h,
|
||||
vision_features_slice, frame_skip, desire_key, extra_policy_keys,
|
||||
vision_road_key, vision_wide_key, prepare_only=False):
|
||||
frame_prepare = make_frame_prepare(nv12, model_w, model_h)
|
||||
sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
|
||||
sample_desire_fn = partial(sample_desire, frame_skip=frame_skip)
|
||||
def compile_and_warmup(nv12: NV12Frame, model_size: tuple[int, int], prepare_only: bool, frame_skip: int, vision_runner, policy_runners: list, metadata: dict):
|
||||
print(f"Compiling combined JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...")
|
||||
|
||||
def run_multi_policy(img_q, big_img_q, feat_q, desire_q, desire,
|
||||
traffic_convention, tfm, big_tfm, frame, big_frame, **extra):
|
||||
npy_tensors = [tfm.to(Device.DEFAULT), big_tfm.to(Device.DEFAULT),
|
||||
desire.to(Device.DEFAULT), traffic_convention.to(Device.DEFAULT)]
|
||||
extra_device = {k: extra[k].to(Device.DEFAULT) for k in extra_policy_keys}
|
||||
Tensor.realize(*npy_tensors, *extra_device.values())
|
||||
tfm, big_tfm, desire, traffic_convention = npy_tensors
|
||||
all_shapes = {key: value for meta in metadata.values() for key, value in meta['input_shapes'].items()}
|
||||
|
||||
img = shift_and_sample(img_q, frame_prepare(frame, tfm).unsqueeze(0), sample_skip_fn)
|
||||
big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm).unsqueeze(0), sample_skip_fn)
|
||||
feat_meta = metadata.get('vision') or metadata.get('model') or metadata.get('policy')
|
||||
if not feat_meta:
|
||||
raise ValueError("Could not find vision, model, or policy metadata.")
|
||||
|
||||
if prepare_only:
|
||||
return img, big_img
|
||||
features_slice = feat_meta['output_slices']['hidden_state']
|
||||
WARP_DEV = 'CPU' if "USBGPU" in os.environ else Device.DEFAULT
|
||||
|
||||
vision_out = next(iter(vision_runner({vision_road_key: img, vision_wide_key: big_img}).values())).cast('float32')
|
||||
run_func = create_jit_runner(vision_runner, policy_runners, nv12, model_size, features_slice, frame_skip, all_shapes, prepare_only)
|
||||
run_jit = TinyJit(run_func, prune=True)
|
||||
queues, npy_arrays = generate_queues_and_npy(all_shapes, frame_skip, Device.DEFAULT)
|
||||
|
||||
new_feat = vision_out[:, vision_features_slice].reshape(1, -1).unsqueeze(0)
|
||||
feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn)
|
||||
desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn)
|
||||
|
||||
inputs = {'features_buffer': feat_buf, desire_key: desire_buf, 'traffic_convention': traffic_convention, **extra_device}
|
||||
|
||||
policy_outputs = []
|
||||
for runner in policy_runners:
|
||||
policy_out = next(iter(runner(inputs).values())).cast('float32')
|
||||
policy_outputs.append(policy_out)
|
||||
|
||||
return (vision_out, *policy_outputs)
|
||||
|
||||
return run_multi_policy
|
||||
|
||||
|
||||
def _warmup_and_serialize(run_jit, input_queues, npy, nv12):
|
||||
for i in range(3):
|
||||
rng = np.random.default_rng(42 + i)
|
||||
frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize()
|
||||
big_frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize()
|
||||
for v in npy.values():
|
||||
v[:] = rng.standard_normal(v.shape).astype(v.dtype)
|
||||
np.random.seed(42 + i)
|
||||
frame = Tensor.randint(nv12.size, low=0, high=256, dtype=dtypes.uint8, device=WARP_DEV).realize()
|
||||
big_frame = Tensor.randint(nv12.size, low=0, high=256, dtype=dtypes.uint8, device=WARP_DEV).realize()
|
||||
for arr in npy_arrays.values():
|
||||
arr[:] = np.random.randn(*arr.shape).astype(arr.dtype)
|
||||
|
||||
Device.default.synchronize()
|
||||
st = time.perf_counter()
|
||||
run_jit(**input_queues, frame=frame, big_frame=big_frame)
|
||||
mt = time.perf_counter()
|
||||
start_time = time.perf_counter()
|
||||
run_jit(**queues, frame=frame, big_frame=big_frame)
|
||||
mid_time = time.perf_counter()
|
||||
Device.default.synchronize()
|
||||
et = time.perf_counter()
|
||||
print(f" [{i + 1}/3] enqueue {(mt - st) * 1e3:6.2f} ms -- total {(et - st) * 1e3:6.2f} ms")
|
||||
return pickle.loads(pickle.dumps(run_jit))
|
||||
print(f" [{i + 1}/3] enqueue {(mid_time - start_time) * 1e3:6.2f} ms -- total {(time.perf_counter() - start_time) * 1e3:6.2f} ms")
|
||||
|
||||
return pickle.loads(pickle.dumps(run_jit)) if not prepare_only else run_jit
|
||||
|
||||
|
||||
def compile_supercombo(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip,
|
||||
model_runner, metadata):
|
||||
print(f"Compiling combined supercombo JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...")
|
||||
|
||||
features_slice = metadata['output_slices']['hidden_state']
|
||||
input_shapes = metadata['input_shapes']
|
||||
|
||||
_run = make_run_supercombo(model_runner, nv12, model_w, model_h,
|
||||
features_slice, frame_skip, input_shapes, prepare_only)
|
||||
run_jit = TinyJit(_run, prune=True)
|
||||
|
||||
input_queues, npy = make_supercombo_input_queues(input_shapes, frame_skip, Device.DEFAULT)
|
||||
|
||||
run_jit = _warmup_and_serialize(run_jit, input_queues, npy, nv12)
|
||||
return run_jit
|
||||
def _parse_size(size_str: str) -> tuple[int, int]:
|
||||
width, height = size_str.lower().split('x')
|
||||
return int(width), int(height)
|
||||
|
||||
|
||||
def compile_multi_policy(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip,
|
||||
vision_runner, policy_runners, vision_metadata, policy_metadata):
|
||||
print(f"Compiling combined multi-policy JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...")
|
||||
|
||||
vision_features_slice = vision_metadata['output_slices']['hidden_state']
|
||||
vision_input_shapes = vision_metadata['input_shapes']
|
||||
policy_input_shapes = policy_metadata['input_shapes']
|
||||
desire_key = _detect_desire_key(policy_input_shapes)
|
||||
extra_policy_keys = [k for k in policy_input_shapes if k not in ('features_buffer', desire_key, 'traffic_convention')]
|
||||
vision_road_key, vision_wide_key = _detect_vision_keys(vision_input_shapes)
|
||||
|
||||
_run = make_run_vision_multi_policy(vision_runner, policy_runners, nv12, model_w, model_h,
|
||||
vision_features_slice, frame_skip, desire_key, extra_policy_keys,
|
||||
vision_road_key, vision_wide_key, prepare_only)
|
||||
run_jit = TinyJit(_run, prune=True)
|
||||
|
||||
input_queues, npy = make_split_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, Device.DEFAULT)
|
||||
|
||||
run_jit = _warmup_and_serialize(run_jit, input_queues, npy, nv12)
|
||||
return run_jit
|
||||
def read_file_chunked_to_shm(path):
|
||||
if not path:
|
||||
return None
|
||||
import atexit
|
||||
from openpilot.common.file_chunker import read_file_chunked
|
||||
from openpilot.system.hardware.hw import Paths
|
||||
shm_path = os.path.join(Paths.shm_path(), os.path.basename(path))
|
||||
atexit.register(lambda: os.path.exists(shm_path) and os.remove(shm_path))
|
||||
with open(shm_path, 'wb') as f:
|
||||
f.write(read_file_chunked(path))
|
||||
return shm_path
|
||||
|
||||
|
||||
def _parse_size(s):
|
||||
w, h = s.lower().split('x')
|
||||
return int(w), int(h)
|
||||
def _compile_for_resolutions(camera_resolutions: list, model_size: tuple[int, int], frame_skip: int,
|
||||
vision_runner, policy_runners: list, metadata: dict) -> dict:
|
||||
from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
||||
return {
|
||||
(cam_w, cam_h): {
|
||||
name: compile_and_warmup(NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)), model_size, prepare_only,
|
||||
frame_skip, vision_runner, policy_runners, metadata)
|
||||
for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)]
|
||||
}
|
||||
for cam_w, cam_h in camera_resolutions
|
||||
}
|
||||
|
||||
|
||||
def _load_policy_runners(args: argparse.Namespace) -> tuple[list, list]:
|
||||
runners, keys = [], []
|
||||
for name, onnx_arg in [('policy', args.policy_onnx), ('off_policy', args.off_policy_onnx), ('on_policy', args.on_policy_onnx)]:
|
||||
if onnx_arg:
|
||||
runners.append(OnnxRunner(onnx_arg))
|
||||
keys.append(name)
|
||||
return runners, keys
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
||||
from openpilot.selfdrive.modeld.get_model_metadata import make_metadata_dict
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
|
||||
p = argparse.ArgumentParser(description="Compile combined JIT pkl for sunnypilot modeld_v2")
|
||||
p.add_argument('--model-type', choices=MODEL_TYPES, required=True)
|
||||
p.add_argument('--model-size', type=_parse_size, required=True, help='model input WxH')
|
||||
p.add_argument('--camera-resolutions', type=_parse_size, nargs='+', required=True)
|
||||
p.add_argument('--frame-skip', type=int, default=None, help='frame skip value (auto-derived if not provided)')
|
||||
p.add_argument('--output', required=True)
|
||||
parser = argparse.ArgumentParser(description="Compile combined JIT pkl for sunnypilot modeld_v2")
|
||||
parser.add_argument('--model-type', choices=MODEL_TYPES, required=True)
|
||||
parser.add_argument('--model-size', type=_parse_size, required=True, help='model input WxH')
|
||||
parser.add_argument('--camera-resolutions', type=_parse_size, nargs='+', required=True)
|
||||
parser.add_argument('--frame-skip', type=int, default=None, help='frame skip value (auto-derived if not provided)')
|
||||
parser.add_argument('--output', required=True)
|
||||
|
||||
p.add_argument('--vision-onnx', help='vision ONNX (for split models)')
|
||||
p.add_argument('--policy-onnx', help='policy ONNX (for vision_policy)')
|
||||
p.add_argument('--off-policy-onnx', help='off-policy ONNX (for vision_multi_policy)')
|
||||
p.add_argument('--on-policy-onnx', help='on-policy ONNX (for vision_multi_policy)')
|
||||
p.add_argument('--supercombo-onnx', help='supercombo ONNX (for supercombo)')
|
||||
parser.add_argument('--vision-onnx', help='vision ONNX (for split models)')
|
||||
parser.add_argument('--policy-onnx', help='policy ONNX (for vision_policy)')
|
||||
parser.add_argument('--off-policy-onnx', help='off-policy ONNX (for vision_multi_policy)')
|
||||
parser.add_argument('--on-policy-onnx', help='on-policy ONNX (for vision_multi_policy)')
|
||||
parser.add_argument('--supercombo-onnx', help='supercombo ONNX (for supercombo)')
|
||||
|
||||
args = p.parse_args()
|
||||
out = defaultdict(dict)
|
||||
args = parser.parse_args()
|
||||
output_data = defaultdict(dict)
|
||||
|
||||
args.vision_onnx = read_file_chunked_to_shm(args.vision_onnx)
|
||||
args.policy_onnx = read_file_chunked_to_shm(args.policy_onnx)
|
||||
args.off_policy_onnx = read_file_chunked_to_shm(args.off_policy_onnx)
|
||||
args.on_policy_onnx = read_file_chunked_to_shm(args.on_policy_onnx)
|
||||
args.supercombo_onnx = read_file_chunked_to_shm(args.supercombo_onnx)
|
||||
|
||||
vision_runner = OnnxRunner(args.vision_onnx) if args.vision_onnx else None
|
||||
|
||||
if args.model_type == 'vision_policy':
|
||||
assert args.vision_onnx and args.policy_onnx
|
||||
vision_runner = OnnxRunner(args.vision_onnx)
|
||||
policy_runner = OnnxRunner(args.policy_onnx)
|
||||
out['metadata']['vision'] = make_metadata_dict(args.vision_onnx)
|
||||
out['metadata']['policy'] = make_metadata_dict(args.policy_onnx)
|
||||
|
||||
frame_skip = args.frame_skip if args.frame_skip is not None else derive_frame_skip(out['metadata']['vision']['input_shapes'],
|
||||
out['metadata']['policy']['input_shapes'])
|
||||
|
||||
for cam_w, cam_h in args.camera_resolutions:
|
||||
nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
|
||||
model_w, model_h = args.model_size
|
||||
out[(cam_w, cam_h)] = {
|
||||
name: compile_split_policy(nv12, model_w, model_h, prepare_only, frame_skip,
|
||||
vision_runner, policy_runner,
|
||||
out['metadata']['vision'], out['metadata']['policy'])
|
||||
for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)]
|
||||
}
|
||||
|
||||
assert vision_runner and args.policy_onnx
|
||||
policy_runners = [OnnxRunner(args.policy_onnx)]
|
||||
output_data['metadata'] = {'vision': make_metadata_dict(args.vision_onnx), 'policy': make_metadata_dict(args.policy_onnx)}
|
||||
elif args.model_type == 'supercombo':
|
||||
assert args.supercombo_onnx
|
||||
model_runner = OnnxRunner(args.supercombo_onnx)
|
||||
out['metadata']['model'] = make_metadata_dict(args.supercombo_onnx)
|
||||
|
||||
frame_skip = args.frame_skip if args.frame_skip is not None else derive_frame_skip({}, out['metadata']['model']['input_shapes'])
|
||||
|
||||
for cam_w, cam_h in args.camera_resolutions:
|
||||
nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
|
||||
model_w, model_h = args.model_size
|
||||
out[(cam_w, cam_h)] = {
|
||||
name: compile_supercombo(nv12, model_w, model_h, prepare_only, frame_skip,
|
||||
model_runner, out['metadata']['model'])
|
||||
for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)]
|
||||
}
|
||||
|
||||
policy_runners = [OnnxRunner(args.supercombo_onnx)]
|
||||
output_data['metadata'] = {'model': make_metadata_dict(args.supercombo_onnx)}
|
||||
elif args.model_type == 'vision_multi_policy':
|
||||
assert args.vision_onnx
|
||||
vision_runner = OnnxRunner(args.vision_onnx)
|
||||
out['metadata']['vision'] = make_metadata_dict(args.vision_onnx)
|
||||
assert vision_runner
|
||||
policy_runners, policy_names = _load_policy_runners(args)
|
||||
output_data['metadata'] = {'vision': make_metadata_dict(args.vision_onnx)}
|
||||
for name in policy_names:
|
||||
runner_arg = getattr(args, f"{name}_onnx")
|
||||
output_data['metadata'][name] = make_metadata_dict(runner_arg)
|
||||
|
||||
policy_runners = []
|
||||
policy_onnxes = []
|
||||
if args.policy_onnx:
|
||||
policy_onnxes.append(('policy', args.policy_onnx))
|
||||
if args.off_policy_onnx:
|
||||
policy_onnxes.append(('off_policy', args.off_policy_onnx))
|
||||
if args.on_policy_onnx:
|
||||
policy_onnxes.append(('on_policy', args.on_policy_onnx))
|
||||
policy_keys = [key for key in output_data['metadata'].keys() if key != 'vision']
|
||||
first_policy_meta = output_data['metadata'][policy_keys[0]] if policy_keys else {}
|
||||
vision_meta = output_data['metadata'].get('vision', {})
|
||||
|
||||
for name, onnx_path in policy_onnxes:
|
||||
runner = OnnxRunner(onnx_path)
|
||||
policy_runners.append(runner)
|
||||
out['metadata'][name] = make_metadata_dict(onnx_path)
|
||||
derived_frame_skip = args.frame_skip or derive_frame_skip(vision_meta.get('input_shapes', {}), first_policy_meta.get('input_shapes', {}))
|
||||
output_data.update(_compile_for_resolutions(args.camera_resolutions, args.model_size, derived_frame_skip,
|
||||
vision_runner, policy_runners, output_data['metadata']))
|
||||
|
||||
first_policy_key = policy_onnxes[0][0]
|
||||
frame_skip = args.frame_skip if args.frame_skip is not None else derive_frame_skip(out['metadata']['vision']['input_shapes'],
|
||||
out['metadata'][first_policy_key]['input_shapes'])
|
||||
with open(args.output, "wb") as file:
|
||||
pickle.dump(output_data, file)
|
||||
|
||||
for cam_w, cam_h in args.camera_resolutions:
|
||||
nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
|
||||
model_w, model_h = args.model_size
|
||||
out[(cam_w, cam_h)] = {
|
||||
name: compile_multi_policy(nv12, model_w, model_h, prepare_only, frame_skip,
|
||||
vision_runner, policy_runners,
|
||||
out['metadata']['vision'], out['metadata'][first_policy_key])
|
||||
for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)]
|
||||
}
|
||||
|
||||
with open(args.output, "wb") as f:
|
||||
pickle.dump(out, f)
|
||||
pkl_size = os.path.getsize(args.output)
|
||||
print(f"Saved combined JIT to {args.output} ({pkl_size / 1e6:.2f} MB)")
|
||||
|
||||
from openpilot.common.file_chunker import chunk_file, get_chunk_targets
|
||||
chunk_targets = get_chunk_targets(args.output, pkl_size)
|
||||
chunk_file(args.output, chunk_targets)
|
||||
num_chunks = len(chunk_targets) - 1
|
||||
print(f"Chunked into {num_chunks} file(s)")
|
||||
print(f"Chunked into {len(chunk_targets) - 1} file(s)")
|
||||
|
||||
@@ -7,7 +7,8 @@ See the LICENSE.md file in the root directory for more details.
|
||||
"""
|
||||
|
||||
import os
|
||||
from openpilot.common.hardware import TICI
|
||||
os.environ['GMMU'] = '0'
|
||||
from openpilot.system.hardware import TICI
|
||||
os.environ['DEV'] = 'QCOM' if TICI else 'CPU'
|
||||
USBGPU = "USBGPU" in os.environ
|
||||
if USBGPU:
|
||||
@@ -16,11 +17,10 @@ if USBGPU:
|
||||
import pickle
|
||||
import time
|
||||
import numpy as np
|
||||
import openpilot.cereal.messaging as messaging
|
||||
from openpilot.cereal import log
|
||||
from opendbc.car.structs import car
|
||||
import cereal.messaging as messaging
|
||||
from cereal import car, log
|
||||
from setproctitle import setproctitle
|
||||
from openpilot.cereal.messaging import PubMaster, SubMaster
|
||||
from cereal.messaging import PubMaster, SubMaster
|
||||
from msgq.visionipc import VisionIpcClient, VisionStreamType, VisionBuf
|
||||
from opendbc.car.car_helpers import get_demo_car_params
|
||||
from openpilot.common.swaglog import cloudlog
|
||||
@@ -42,7 +42,7 @@ from openpilot.sunnypilot.livedelay.helpers import get_lat_delay
|
||||
from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase
|
||||
from openpilot.sunnypilot.models.helpers import get_active_bundle
|
||||
|
||||
PROCESS_NAME = "openpilot.selfdrive.modeld.modeld_tinygrad"
|
||||
PROCESS_NAME = "selfdrive.modeld.modeld_tinygrad"
|
||||
|
||||
|
||||
def _pkl_exists(path):
|
||||
@@ -55,7 +55,7 @@ def _find_driving_pkl(bundle):
|
||||
return override
|
||||
if bundle is None or not bundle.models:
|
||||
return None
|
||||
from openpilot.common.hardware.hw import Paths
|
||||
from openpilot.system.hardware.hw import Paths
|
||||
model_root = Paths.model_root()
|
||||
|
||||
pkl_name = bundle.models[0].artifact.fileName
|
||||
@@ -104,12 +104,14 @@ class ModelState(ModelStateBase):
|
||||
from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, make_split_input_queues
|
||||
from tinygrad.device import Device
|
||||
|
||||
from openpilot.common.file_chunker import open_file_chunked
|
||||
from openpilot.common.file_chunker import read_file_chunked
|
||||
|
||||
cloudlog.warning(f"loading combined pkl: {pkl_path}")
|
||||
jits = pickle.load(open_file_chunked(pkl_path))
|
||||
jits = pickle.loads(read_file_chunked(pkl_path))
|
||||
|
||||
self.DEV = Device.DEFAULT
|
||||
self.WARP_DEV = 'CPU' if USBGPU else self.DEV
|
||||
self.QUEUE_DEV = self.DEV
|
||||
|
||||
metadata = jits['metadata']
|
||||
if 'model' in metadata:
|
||||
@@ -121,7 +123,7 @@ class ModelState(ModelStateBase):
|
||||
self._vision_input_names = [k for k in model_metadata['input_shapes'] if 'img' in k]
|
||||
from openpilot.sunnypilot.modeld_v2.compile_modeld import make_supercombo_input_queues
|
||||
frame_skip = derive_frame_skip({}, model_metadata['input_shapes'])
|
||||
self.input_queues, self.numpy_inputs = make_supercombo_input_queues(model_metadata['input_shapes'], frame_skip, device=self.DEV)
|
||||
self.input_queues, self.numpy_inputs = make_supercombo_input_queues(model_metadata['input_shapes'], frame_skip, device=self.QUEUE_DEV)
|
||||
else:
|
||||
vision_metadata = metadata['vision']
|
||||
policy_keys = [k for k in metadata if k != 'vision']
|
||||
@@ -139,7 +141,11 @@ class ModelState(ModelStateBase):
|
||||
policy_input_shapes = first_policy_metadata['input_shapes']
|
||||
self._vision_input_names = [k for k in vision_input_shapes if 'img' in k]
|
||||
frame_skip = derive_frame_skip(vision_input_shapes, policy_input_shapes)
|
||||
self.input_queues, self.numpy_inputs = make_split_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, device=self.DEV)
|
||||
self.input_queues, self.numpy_inputs = make_split_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, device=self.QUEUE_DEV)
|
||||
|
||||
self._desire_key = next(key for key in self.numpy_inputs if key.startswith('desire'))
|
||||
self._road_key = next(key for key in self._vision_input_names if 'big' not in key)
|
||||
self._wide_key = next(key for key in self._vision_input_names if 'big' in key)
|
||||
|
||||
from openpilot.sunnypilot.modeld_v2.parse_model_outputs_split import Parser as SplitParser
|
||||
from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser as CombinedParser
|
||||
@@ -161,12 +167,11 @@ class ModelState(ModelStateBase):
|
||||
|
||||
self._run_policy = jits[(cam_w, cam_h)]['run_policy']
|
||||
self._warp_enqueue = jits[(cam_w, cam_h)]['warp_enqueue']
|
||||
road_name = next(k for k in self._vision_input_names if 'big' not in k)
|
||||
yuv_size = self.frame_buf_params[road_name][3]
|
||||
yuv_size = self.frame_buf_params[self._road_key][3]
|
||||
self._warp_enqueue(
|
||||
**self.input_queues,
|
||||
frame=Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.DEV).contiguous().realize(),
|
||||
big_frame=Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.DEV).contiguous().realize())
|
||||
frame=Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.WARP_DEV).contiguous().realize(),
|
||||
big_frame=Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.WARP_DEV).contiguous().realize())
|
||||
|
||||
|
||||
@property
|
||||
@@ -179,7 +184,7 @@ class ModelState(ModelStateBase):
|
||||
|
||||
@property
|
||||
def desire_key(self) -> str:
|
||||
return next(k for k in self.numpy_inputs if k.startswith('desire'))
|
||||
return self._desire_key
|
||||
|
||||
def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray],
|
||||
inputs: dict[str, np.ndarray], prepare_only: bool) -> dict[str, np.ndarray] | None:
|
||||
@@ -190,19 +195,19 @@ class ModelState(ModelStateBase):
|
||||
yuv_size = self.frame_buf_params[key][3]
|
||||
cache_key = (key, ptr)
|
||||
if cache_key not in self._blob_cache:
|
||||
self._blob_cache[cache_key] = Tensor.from_blob(ptr, (yuv_size,), dtype='uint8', device=self.DEV)
|
||||
self._blob_cache[cache_key] = Tensor.from_blob(ptr, (yuv_size,), dtype='uint8', device=self.WARP_DEV)
|
||||
self.full_frames[key] = self._blob_cache[cache_key]
|
||||
|
||||
desire_key = self.desire_key
|
||||
inputs[desire_key][0] = 0
|
||||
self.numpy_inputs[desire_key][:] = np.where(inputs[desire_key] - self.prev_desire > .99, inputs[desire_key], 0)
|
||||
self.prev_desire[:] = inputs[desire_key]
|
||||
for key in ('traffic_convention', 'lateral_control_params'):
|
||||
for key in ('traffic_convention', 'lateral_control_params', 'action_t'):
|
||||
if key in self.numpy_inputs and key in inputs:
|
||||
self.numpy_inputs[key][:] = inputs[key]
|
||||
|
||||
road_key = next(n for n in bufs if 'big' not in n)
|
||||
wide_key = next(n for n in bufs if 'big' in n)
|
||||
road_key = self._road_key
|
||||
wide_key = self._wide_key
|
||||
self.numpy_inputs['tfm'][:, :] = transforms[road_key].reshape(3, 3)
|
||||
self.numpy_inputs['big_tfm'][:, :] = transforms[wide_key].reshape(3, 3)
|
||||
|
||||
@@ -225,8 +230,12 @@ class ModelState(ModelStateBase):
|
||||
policy_output = raw_outputs[i + 1].numpy().flatten()
|
||||
policy_sliced = {k: policy_output[np.newaxis, v] for k, v in policy_slices.items()}
|
||||
parsed = self.parser.parse_policy_outputs(policy_sliced)
|
||||
if 'off' in self._policy_keys[i] and self._has_on_policy:
|
||||
if ('off' in self._policy_keys[i]
|
||||
and self._has_on_policy
|
||||
and any('plan' in self._policy_slices_list[j] for j, k in enumerate(self._policy_keys) if 'on' in k.lower())):
|
||||
|
||||
parsed.pop('plan', None)
|
||||
|
||||
outputs.update(parsed)
|
||||
|
||||
if 'planplus' in outputs and 'plan' in outputs:
|
||||
@@ -241,13 +250,20 @@ class ModelState(ModelStateBase):
|
||||
|
||||
def get_action_from_model(self, model_output: dict[str, np.ndarray], prev_action: log.ModelDataV2.Action,
|
||||
lat_action_t: float, long_action_t: float, v_ego: float) -> log.ModelDataV2.Action:
|
||||
plan = model_output['plan'][0]
|
||||
desired_accel, should_stop = get_accel_from_plan(plan[:, Plan.VELOCITY][:, 0], plan[:, Plan.ACCELERATION][:, 0], self.constants.T_IDXS,
|
||||
action_t=long_action_t)
|
||||
desired_accel = smooth_value(desired_accel, prev_action.desiredAcceleration, self.LONG_SMOOTH_SECONDS)
|
||||
if 'action' not in model_output:
|
||||
plan = model_output['plan'][0]
|
||||
desired_accel, should_stop = get_accel_from_plan(plan[:, Plan.VELOCITY][:, 0], plan[:, Plan.ACCELERATION][:, 0], self.constants.T_IDXS,
|
||||
action_t=long_action_t)
|
||||
desired_accel = smooth_value(desired_accel, prev_action.desiredAcceleration, self.LONG_SMOOTH_SECONDS)
|
||||
|
||||
curvature_plan = (plan + (self.PLANPLUS_CONTROL - 1.0) * model_output['planplus'][0]
|
||||
if 'planplus' in model_output and self.PLANPLUS_CONTROL != 1.0 else plan)
|
||||
desired_curvature = get_curvature_from_output(model_output, curvature_plan, v_ego, lat_action_t, self.mlsim)
|
||||
else:
|
||||
desired_accel = model_output['action'][0, 1]
|
||||
desired_curvature = model_output['action'][0, 0] / (max(1.0, v_ego))**2
|
||||
should_stop = (v_ego < 0.3 and desired_accel < 0.1)
|
||||
|
||||
curvature_plan = plan + (self.PLANPLUS_CONTROL - 1.0) * model_output['planplus'][0] if 'planplus' in model_output and self.PLANPLUS_CONTROL != 1.0 else plan
|
||||
desired_curvature = get_curvature_from_output(model_output, curvature_plan, v_ego, lat_action_t, self.mlsim)
|
||||
if self.generation is not None and self.generation >= 10: # smooth curvature for post FOF models
|
||||
if v_ego > self.MIN_LAT_CONTROL_SPEED:
|
||||
desired_curvature = smooth_value(desired_curvature, prev_action.desiredCurvature, self.LAT_SMOOTH_SECONDS)
|
||||
@@ -400,6 +416,12 @@ def main(demo=False):
|
||||
|
||||
bufs = {name: buf_extra if 'big' in name else buf_main for name in model.vision_input_names}
|
||||
transforms = {name: model_transform_extra if 'big' in name else model_transform_main for name in model.vision_input_names}
|
||||
|
||||
frame_delay = DT_MDL # compensate for time passed since the frame was captured: current_time - timestamp_eof is 50ms on average
|
||||
action_delay = DT_MDL / 2 # middle of the interval between model output (current state) and next frame (expected state)
|
||||
lat_action_t = lat_delay + frame_delay + action_delay
|
||||
long_action_t = long_delay + frame_delay + action_delay
|
||||
|
||||
inputs:dict[str, np.ndarray] = {
|
||||
model.desire_key: vec_desire,
|
||||
'traffic_convention': traffic_convention,
|
||||
@@ -408,6 +430,9 @@ def main(demo=False):
|
||||
if 'lateral_control_params' in model.numpy_inputs:
|
||||
inputs['lateral_control_params'] = np.array([v_ego, lat_delay], dtype=np.float32)
|
||||
|
||||
if 'action_t' in model.numpy_inputs:
|
||||
inputs['action_t'] = np.array([lat_action_t, long_action_t], dtype=np.float32)
|
||||
|
||||
mt1 = time.perf_counter()
|
||||
model_output = model.run(bufs, transforms, inputs, prepare_only)
|
||||
mt2 = time.perf_counter()
|
||||
@@ -419,7 +444,7 @@ def main(demo=False):
|
||||
posenet_send = messaging.new_message('cameraOdometry')
|
||||
mdv2sp_send = messaging.new_message('modelDataV2SP')
|
||||
|
||||
action = model.get_action_from_model(model_output, prev_action, lat_delay + DT_MDL, long_delay + DT_MDL, v_ego)
|
||||
action = model.get_action_from_model(model_output, prev_action, lat_action_t, long_action_t, v_ego)
|
||||
prev_action = action
|
||||
fill_model_msg(drivingdata_send, modelv2_send, model_output, action,
|
||||
publish_state, meta_main.frame_id, meta_extra.frame_id, frame_id,
|
||||
|
||||
@@ -137,6 +137,8 @@ class Parser:
|
||||
out_shape=(SplitModelConstants.NUM_ROAD_EDGES,SplitModelConstants.IDX_N,SplitModelConstants.LANE_LINES_WIDTH))
|
||||
if 'sim_pose' in outs:
|
||||
self.parse_mdn('sim_pose', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.POSE_WIDTH,))
|
||||
if 'action' in outs:
|
||||
self.parse_mdn('action', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.ACTION_WIDTH,))
|
||||
|
||||
def parse_vision_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
|
||||
self.parse_mdn('pose', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.POSE_WIDTH,))
|
||||
|
||||
@@ -1,103 +0,0 @@
|
||||
import os
|
||||
os.environ['DEV'] = 'CPU'
|
||||
import pytest
|
||||
import numpy as np
|
||||
from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
||||
from openpilot.sunnypilot.modeld_v2.warp import CAMERA_CONFIGS
|
||||
from openpilot.sunnypilot.modeld_v2.warp import Warp, MODEL_W, MODEL_H
|
||||
|
||||
VISION_NAME_PAIRS = [ # needed to account for supercombos input_imgs
|
||||
('img', 'big_img'),
|
||||
('input_imgs', 'big_input_imgs'),
|
||||
]
|
||||
|
||||
|
||||
class MockVisionBuf:
|
||||
def __init__(self, w, h):
|
||||
self.width = w
|
||||
self.height = h
|
||||
_, _, _, yuv_size = get_nv12_info(w, h)
|
||||
self.data = np.zeros(yuv_size, dtype=np.uint8)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("buffer_length", [2, 5])
|
||||
def test_warp_initialization(buffer_length):
|
||||
warp = Warp(buffer_length)
|
||||
assert warp.buffer_length == buffer_length
|
||||
assert warp.img_buffer_shape == (buffer_length * 6, MODEL_H // 2, MODEL_W // 2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("buffer_length", [2, 5])
|
||||
@pytest.mark.parametrize("cam_w, cam_h", CAMERA_CONFIGS)
|
||||
@pytest.mark.parametrize("road, wide", VISION_NAME_PAIRS)
|
||||
def test_warp_process(buffer_length, cam_w, cam_h, road, wide):
|
||||
warp = Warp(buffer_length)
|
||||
mock_buf = MockVisionBuf(cam_w, cam_h)
|
||||
transform = np.eye(3, dtype=np.float32).flatten()
|
||||
bufs = {road: mock_buf, wide: mock_buf}
|
||||
transforms = {road: transform, wide: transform}
|
||||
|
||||
out = warp.process(bufs, transforms)
|
||||
assert isinstance(out, dict)
|
||||
assert road in out and wide in out
|
||||
assert out[road].shape == (1, 12, MODEL_H // 2, MODEL_W // 2)
|
||||
assert out[wide].shape == (1, 12, MODEL_H // 2, MODEL_W // 2)
|
||||
|
||||
key = (cam_w, cam_h)
|
||||
assert key in warp.jit_cache
|
||||
|
||||
out2 = warp.process(bufs, transforms)
|
||||
assert out2[road].shape == out[road].shape
|
||||
|
||||
|
||||
@pytest.mark.parametrize("road, wide", VISION_NAME_PAIRS)
|
||||
def test_warp_buffer_shift(road, wide):
|
||||
warp = Warp(2)
|
||||
cam_w, cam_h = CAMERA_CONFIGS[1]
|
||||
transform = np.eye(3, dtype=np.float32).flatten()
|
||||
|
||||
buf1 = MockVisionBuf(cam_w, cam_h)
|
||||
buf1.data[0] = 255
|
||||
bufs1 = {road: buf1, wide: buf1}
|
||||
transforms = {road: transform, wide: transform}
|
||||
out1 = warp.process(bufs1, transforms)
|
||||
road1 = out1[road].numpy().copy()
|
||||
|
||||
buf2 = MockVisionBuf(cam_w, cam_h)
|
||||
buf2.data[0] = 128
|
||||
bufs2 = {road: buf2, wide: buf2}
|
||||
out2 = warp.process(bufs2, transforms)
|
||||
assert not np.array_equal(road1, out2[road].numpy())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("buffer_length", [2, 5])
|
||||
@pytest.mark.parametrize("road, wide", VISION_NAME_PAIRS)
|
||||
def test_warp_buffer_accumulation(buffer_length, road, wide):
|
||||
warp = Warp(buffer_length)
|
||||
cam_w, cam_h = CAMERA_CONFIGS[0]
|
||||
transform = np.eye(3, dtype=np.float32).flatten()
|
||||
transforms = {road: transform, wide: transform}
|
||||
outputs = []
|
||||
|
||||
for i in range(buffer_length + 1):
|
||||
buf = MockVisionBuf(cam_w, cam_h)
|
||||
buf.data[:] = i * 10
|
||||
out = warp.process({road: buf, wide: buf}, transforms)
|
||||
outputs.append(out[road].numpy().copy())
|
||||
|
||||
assert warp.full_buffers['img'].shape == (buffer_length * 6, MODEL_H // 2, MODEL_W // 2)
|
||||
for i in range(1, len(outputs)):
|
||||
assert not np.array_equal(outputs[i - 1], outputs[i])
|
||||
|
||||
|
||||
def test_warp_different_cameras_same_instance():
|
||||
warp = Warp(2)
|
||||
transform = np.eye(3, dtype=np.float32).flatten()
|
||||
|
||||
buf1 = MockVisionBuf(*CAMERA_CONFIGS[0])
|
||||
warp.process({'img': buf1, 'big_img': buf1}, {'img': transform, 'big_img': transform})
|
||||
assert len(warp.jit_cache) == 1
|
||||
|
||||
buf2 = MockVisionBuf(*CAMERA_CONFIGS[1])
|
||||
warp.process({'img': buf2, 'big_img': buf2}, {'img': transform, 'big_img': transform})
|
||||
assert len(warp.jit_cache) == 2
|
||||
@@ -1,171 +0,0 @@
|
||||
import pickle
|
||||
import time
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.engine.jit import TinyJit
|
||||
from tinygrad.device import Device
|
||||
|
||||
from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
||||
from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE
|
||||
from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
|
||||
from openpilot.selfdrive.modeld.compile_modeld import NV12Frame, make_frame_prepare as _make_frame_prepare
|
||||
|
||||
CAMERA_CONFIGS = [
|
||||
(_ar_ox_fisheye.width, _ar_ox_fisheye.height),
|
||||
(_os_fisheye.width, _os_fisheye.height),
|
||||
]
|
||||
|
||||
|
||||
def make_frame_prepare(cam_w, cam_h, model_w, model_h):
|
||||
nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
|
||||
return _make_frame_prepare(nv12, model_w, model_h)
|
||||
|
||||
|
||||
def warp_pkl_path(w, h):
|
||||
from openpilot.selfdrive.modeld.helpers import MODELS_DIR
|
||||
return MODELS_DIR / f'warp_{w}x{h}_tinygrad.pkl'
|
||||
|
||||
|
||||
def make_update_img_input(frame_prepare, model_w, model_h):
|
||||
def update_img_input_tinygrad(tensor, frame, M_inv):
|
||||
M_inv = M_inv.to(Device.DEFAULT)
|
||||
new_img = frame_prepare(frame, M_inv)
|
||||
tensor.assign(tensor[6:].cat(new_img, dim=0).contiguous())
|
||||
return Tensor.cat(tensor[:6], tensor[-6:], dim=0).contiguous().reshape(1, 12, model_h//2, model_w//2)
|
||||
return update_img_input_tinygrad
|
||||
|
||||
|
||||
def make_update_both_imgs(frame_prepare, model_w, model_h):
|
||||
update_img = make_update_img_input(frame_prepare, model_w, model_h)
|
||||
def update_both_imgs_tinygrad(calib_img_buffer, new_img, M_inv,
|
||||
calib_big_img_buffer, new_big_img, M_inv_big):
|
||||
calib_img_pair = update_img(calib_img_buffer, new_img, M_inv)
|
||||
calib_big_img_pair = update_img(calib_big_img_buffer, new_big_img, M_inv_big)
|
||||
return calib_img_pair, calib_big_img_pair
|
||||
return update_both_imgs_tinygrad
|
||||
|
||||
MODELS_DIR = Path(__file__).parent / 'models'
|
||||
MODEL_W, MODEL_H = MEDMODEL_INPUT_SIZE
|
||||
UPSTREAM_BUFFER_LENGTH = 5
|
||||
|
||||
|
||||
def v2_warp_pkl_path(cam_w, cam_h, buffer_length):
|
||||
return MODELS_DIR / f'warp_{cam_w}x{cam_h}_b{buffer_length}_tinygrad.pkl'
|
||||
|
||||
|
||||
def compile_v2_warp(cam_w, cam_h, buffer_length):
|
||||
_, _, _, yuv_size = get_nv12_info(cam_w, cam_h)
|
||||
img_buffer_shape = (buffer_length * 6, MODEL_H // 2, MODEL_W // 2)
|
||||
|
||||
print(f"Compiling v2 warp for {cam_w}x{cam_h} buffer_length={buffer_length}...")
|
||||
|
||||
frame_prepare = make_frame_prepare(cam_w, cam_h, MODEL_W, MODEL_H)
|
||||
update_both_imgs = make_update_both_imgs(frame_prepare, MODEL_W, MODEL_H)
|
||||
update_img_jit = TinyJit(update_both_imgs, prune=True)
|
||||
|
||||
full_buffer = Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize()
|
||||
big_full_buffer = Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize()
|
||||
new_frame_np = np.random.default_rng(0).integers(0, 256, yuv_size, dtype=np.uint8)
|
||||
new_big_frame_np = np.random.default_rng(1).integers(0, 256, yuv_size, dtype=np.uint8)
|
||||
for i in range(10):
|
||||
img_inputs = [full_buffer,
|
||||
Tensor.from_blob(new_frame_np.ctypes.data, (yuv_size,), dtype='uint8').realize(),
|
||||
Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')]
|
||||
big_img_inputs = [big_full_buffer,
|
||||
Tensor.from_blob(new_big_frame_np.ctypes.data, (yuv_size,), dtype='uint8').realize(),
|
||||
Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')]
|
||||
inputs = img_inputs + big_img_inputs
|
||||
Device.default.synchronize()
|
||||
|
||||
st = time.perf_counter()
|
||||
_ = update_img_jit(*inputs)
|
||||
mt = time.perf_counter()
|
||||
Device.default.synchronize()
|
||||
et = time.perf_counter()
|
||||
print(f" [{i+1}/10] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms")
|
||||
|
||||
pkl_path = v2_warp_pkl_path(cam_w, cam_h, buffer_length)
|
||||
with open(pkl_path, "wb") as f:
|
||||
pickle.dump(update_img_jit, f)
|
||||
print(f" Saved to {pkl_path}")
|
||||
|
||||
jit = pickle.load(open(pkl_path, "rb"))
|
||||
verify_frame = np.random.default_rng(0).integers(0, 256, yuv_size, dtype=np.uint8)
|
||||
verify_big_frame = np.random.default_rng(1).integers(0, 256, yuv_size, dtype=np.uint8)
|
||||
fresh_inputs = [
|
||||
Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize(),
|
||||
Tensor.from_blob(verify_frame.ctypes.data, (yuv_size,), dtype='uint8').realize(),
|
||||
Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY'),
|
||||
Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize(),
|
||||
Tensor.from_blob(verify_big_frame.ctypes.data, (yuv_size,), dtype='uint8').realize(),
|
||||
Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY'),
|
||||
]
|
||||
jit(*fresh_inputs)
|
||||
|
||||
|
||||
class Warp:
|
||||
def __init__(self, buffer_length=2):
|
||||
self.buffer_length = buffer_length
|
||||
self.img_buffer_shape = (buffer_length * 6, MODEL_H // 2, MODEL_W // 2)
|
||||
|
||||
self.jit_cache = {}
|
||||
self.full_buffers = {k: Tensor.zeros(self.img_buffer_shape, dtype='uint8').contiguous().realize() for k in ['img', 'big_img']}
|
||||
self._blob_cache: dict[int, Tensor] = {}
|
||||
self._nv12_cache: dict[tuple[int, int], int] = {}
|
||||
self.transforms_np = {k: np.zeros((3, 3), dtype=np.float32) for k in ['img', 'big_img']}
|
||||
self.transforms = {k: Tensor(v, device='NPY').realize() for k, v in self.transforms_np.items()}
|
||||
|
||||
def process(self, bufs, transforms):
|
||||
if not bufs:
|
||||
return {}
|
||||
road = next(n for n in bufs if 'big' not in n)
|
||||
wide = next(n for n in bufs if 'big' in n)
|
||||
cam_w, cam_h = bufs[road].width, bufs[road].height
|
||||
key = (cam_w, cam_h)
|
||||
|
||||
if key not in self.jit_cache:
|
||||
v2_pkl = v2_warp_pkl_path(cam_w, cam_h, self.buffer_length)
|
||||
if v2_pkl.exists():
|
||||
with open(v2_pkl, 'rb') as f:
|
||||
self.jit_cache[key] = pickle.load(f)
|
||||
elif self.buffer_length == UPSTREAM_BUFFER_LENGTH:
|
||||
upstream_pkl = warp_pkl_path(cam_w, cam_h)
|
||||
if upstream_pkl.exists():
|
||||
with open(upstream_pkl, 'rb') as f:
|
||||
self.jit_cache[key] = pickle.load(f)
|
||||
if key not in self.jit_cache:
|
||||
frame_prepare = make_frame_prepare(cam_w, cam_h, MODEL_W, MODEL_H)
|
||||
update_both_imgs = make_update_both_imgs(frame_prepare, MODEL_W, MODEL_H)
|
||||
self.jit_cache[key] = TinyJit(update_both_imgs, prune=True)
|
||||
|
||||
if key not in self._nv12_cache:
|
||||
self._nv12_cache[key] = get_nv12_info(cam_w, cam_h)[3]
|
||||
yuv_size = self._nv12_cache[key]
|
||||
|
||||
road_ptr = bufs[road].data.ctypes.data
|
||||
wide_ptr = bufs[wide].data.ctypes.data
|
||||
if road_ptr not in self._blob_cache:
|
||||
self._blob_cache[road_ptr] = Tensor.from_blob(road_ptr, (yuv_size,), dtype='uint8')
|
||||
if wide_ptr not in self._blob_cache:
|
||||
self._blob_cache[wide_ptr] = Tensor.from_blob(wide_ptr, (yuv_size,), dtype='uint8')
|
||||
road_blob = self._blob_cache[road_ptr]
|
||||
wide_blob = self._blob_cache[wide_ptr] if wide_ptr != road_ptr else Tensor.from_blob(wide_ptr, (yuv_size,), dtype='uint8')
|
||||
np.copyto(self.transforms_np['img'], transforms[road].reshape(3, 3))
|
||||
np.copyto(self.transforms_np['big_img'], transforms[wide].reshape(3, 3))
|
||||
|
||||
Device.default.synchronize()
|
||||
res = self.jit_cache[key](
|
||||
self.full_buffers['img'], road_blob, self.transforms['img'],
|
||||
self.full_buffers['big_img'], wide_blob, self.transforms['big_img'],
|
||||
)
|
||||
out_road = res[0].realize()
|
||||
out_wide = res[1].realize()
|
||||
|
||||
return {road: out_road, wide: out_wide}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
for cam_w, cam_h in CAMERA_CONFIGS:
|
||||
for bl in [2, 5]:
|
||||
compile_v2_warp(cam_w, cam_h, bl)
|
||||
@@ -6,11 +6,12 @@ See the LICENSE.md file in the root directory for more details.
|
||||
"""
|
||||
|
||||
import time
|
||||
|
||||
import os
|
||||
import requests
|
||||
from requests.exceptions import (SSLError, RequestException, HTTPError)
|
||||
from openpilot.common.params import Params
|
||||
from openpilot.common.swaglog import cloudlog
|
||||
from openpilot.system.hardware.hw import Paths
|
||||
from openpilot.sunnypilot.models.helpers import is_bundle_version_compatible
|
||||
|
||||
from openpilot.cereal import custom
|
||||
@@ -26,11 +27,35 @@ class ModelParser:
|
||||
download_uri.sha256 = download_uri_data.get("sha256")
|
||||
return download_uri
|
||||
|
||||
@staticmethod
|
||||
def _parse_chunk(chunk_data) -> custom.ModelManagerSP.Chunk:
|
||||
chunk = custom.ModelManagerSP.Chunk()
|
||||
chunk.fileName = chunk_data.get("file_name")
|
||||
chunk.sha256 = chunk_data.get("sha256")
|
||||
return chunk
|
||||
|
||||
@staticmethod
|
||||
def _parse_artifact(artifact_data) -> custom.ModelManagerSP.Artifact:
|
||||
artifact = custom.ModelManagerSP.Artifact()
|
||||
artifact.fileName = artifact_data.get("file_name")
|
||||
artifact.downloadUri = ModelParser._parse_download_uri(artifact_data.get("download_uri", {}))
|
||||
|
||||
if "chunks" in artifact_data:
|
||||
artifact.chunks = [ModelParser._parse_chunk(chunk_data) for chunk_data in artifact_data["chunks"]]
|
||||
|
||||
try:
|
||||
model_dir = Paths.model_root()
|
||||
os.makedirs(model_dir, exist_ok=True)
|
||||
manifest_path = os.path.join(model_dir, f"{artifact.fileName}.chunkmanifest")
|
||||
num_chunks = str(len(artifact.chunks))
|
||||
|
||||
if not os.path.exists(manifest_path) or open(manifest_path).read().strip() != num_chunks:
|
||||
with open(manifest_path, "w") as f:
|
||||
f.write(num_chunks)
|
||||
cloudlog.info(f"Wrote chunk manifest for {artifact.fileName}: {num_chunks} chunks")
|
||||
except Exception as e:
|
||||
cloudlog.warning(f"Failed to write chunk manifest for {artifact.fileName}: {e}")
|
||||
|
||||
return artifact
|
||||
|
||||
@staticmethod
|
||||
@@ -116,7 +141,7 @@ class ModelCache:
|
||||
|
||||
class ModelFetcher:
|
||||
"""Handles fetching and caching of model data from remote source"""
|
||||
MODEL_URL = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_v17.json"
|
||||
MODEL_URL = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_v18.json"
|
||||
|
||||
def __init__(self, params: Params):
|
||||
self.params = params
|
||||
@@ -184,4 +209,7 @@ if __name__ == "__main__":
|
||||
# Print artifact details
|
||||
print(f"Artifact: {model.artifact.fileName}, Download URI: {model.artifact.downloadUri.uri}")
|
||||
# Print metadata details
|
||||
print(f"Metadata: {model.metadata.fileName}, Download URI: {model.metadata.downloadUri.uri}")
|
||||
if model.artifact.chunks:
|
||||
print(f"Contains {len(model.artifact.chunks)} chunks.")
|
||||
if hasattr(model, 'metadata') and model.metadata and model.metadata.fileName:
|
||||
print(f"Metadata: {model.metadata.fileName}, Download URI: {model.metadata.downloadUri.uri}")
|
||||
|
||||
@@ -89,20 +89,16 @@ class ModelManagerSP:
|
||||
del self._download_start_times[model.fileName]
|
||||
|
||||
async def _download_chunked(self, base_url: str, base_path: str, artifact) -> None:
|
||||
from openpilot.common.file_chunker import get_manifest_path, get_chunk_name
|
||||
manifest_url = get_manifest_path(base_url)
|
||||
from openpilot.common.file_chunker import get_chunk_name, get_manifest_path
|
||||
|
||||
num_chunks = len(artifact.chunks)
|
||||
if num_chunks == 0:
|
||||
raise ValueError("No chunks defined in artifact")
|
||||
|
||||
manifest_path = get_manifest_path(base_path)
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get(manifest_url) as resp:
|
||||
if resp.status == 404:
|
||||
raise FileNotFoundError
|
||||
resp.raise_for_status()
|
||||
num_chunks = int((await resp.read()).strip())
|
||||
|
||||
self._download_start_times[artifact.fileName] = time.monotonic()
|
||||
|
||||
for i in range(num_chunks):
|
||||
for i, _ in enumerate(artifact.chunks):
|
||||
chunk_url = get_chunk_name(base_url, i, num_chunks)
|
||||
chunk_path = get_chunk_name(base_path, i, num_chunks)
|
||||
chunk_downloaded = 0
|
||||
@@ -117,7 +113,7 @@ class ModelManagerSP:
|
||||
if self.params.get("ModelManager_DownloadIndex") is None:
|
||||
raise Exception("Download cancelled")
|
||||
intra = chunk_downloaded / max(chunk_size, 1)
|
||||
progress = min(99, (i + intra) / num_chunks * 100)
|
||||
progress = min(99.0, ((i + intra) / num_chunks) * 100)
|
||||
artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.downloading
|
||||
artifact.downloadProgress.progress = progress
|
||||
artifact.downloadProgress.eta = self._calculate_eta(artifact.fileName, progress)
|
||||
@@ -148,9 +144,9 @@ class ModelManagerSP:
|
||||
self._report_status()
|
||||
return
|
||||
|
||||
try:
|
||||
if len(artifact.chunks) > 0:
|
||||
await self._download_chunked(url, full_path, artifact)
|
||||
except (FileNotFoundError, aiohttp.ClientResponseError):
|
||||
else:
|
||||
await self._download_file(url, full_path, artifact)
|
||||
|
||||
if not await verify_file(full_path, expected_hash):
|
||||
@@ -170,18 +166,15 @@ class ModelManagerSP:
|
||||
artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.failed
|
||||
artifact.downloadProgress.eta = 0
|
||||
self._sync_artifact_progress(artifact)
|
||||
self.selected_bundle.status = custom.ModelManagerSP.DownloadStatus.failed
|
||||
if self.selected_bundle:
|
||||
self.selected_bundle.status = custom.ModelManagerSP.DownloadStatus.failed
|
||||
self._report_status()
|
||||
self._download_start_times.pop(artifact.fileName, None)
|
||||
raise
|
||||
|
||||
async def _process_model(self, model, destination_path: str) -> None:
|
||||
"""Processes a single model download including verification"""
|
||||
model_artifact = model.artifact
|
||||
metadata_artifact = model.metadata
|
||||
|
||||
await self._process_artifact(metadata_artifact, destination_path)
|
||||
await self._process_artifact(model_artifact, destination_path)
|
||||
await self._process_artifact(model.artifact, destination_path)
|
||||
|
||||
def _report_status(self) -> None:
|
||||
"""Reports current status through messaging system"""
|
||||
|
||||
@@ -1,28 +0,0 @@
|
||||
from openpilot.sunnypilot.models.helpers import get_active_bundle
|
||||
from openpilot.sunnypilot.models.runners.model_runner import ModelRunner
|
||||
from openpilot.sunnypilot.models.runners.tinygrad.tinygrad_runner import TinygradRunner, TinygradSplitRunner
|
||||
from openpilot.sunnypilot.models.runners.constants import ModelType
|
||||
|
||||
|
||||
def get_model_runner() -> ModelRunner:
|
||||
"""
|
||||
Factory function to create and return the appropriate ModelRunner instance.
|
||||
|
||||
Selects TinygradRunner, choosing TinygradSplitRunner if separate vision/policy
|
||||
models are detected in the active bundle.
|
||||
|
||||
:return: An instance of a ModelRunner subclass (ONNXRunner, TinygradRunner, or TinygradSplitRunner).
|
||||
"""
|
||||
bundle = get_active_bundle()
|
||||
if bundle and bundle.models:
|
||||
model_types = {m.type.raw for m in bundle.models}
|
||||
# Check if the bundle uses separate vision and policy models (legacy or new split format)
|
||||
split_types = {ModelType.vision, ModelType.policy, ModelType.offPolicy, ModelType.onPolicy}
|
||||
if model_types & split_types:
|
||||
return TinygradSplitRunner()
|
||||
# Otherwise, assume a single model (likely supercombo)
|
||||
if bundle.models:
|
||||
return TinygradRunner(bundle.models[0].type.raw)
|
||||
|
||||
# Default fallback to TinygradRunner with the supercombo type if bundle info is missing/incomplete
|
||||
return TinygradRunner(ModelType.supercombo)
|
||||
@@ -1,174 +0,0 @@
|
||||
from abc import abstractmethod, ABC
|
||||
|
||||
import numpy as np
|
||||
from openpilot.sunnypilot.models.helpers import get_active_bundle
|
||||
from openpilot.sunnypilot.models.runners.constants import NumpyDict, ShapeDict, Model, SliceDict, SEND_RAW_PRED
|
||||
from openpilot.common.hardware.hw import Paths
|
||||
import pickle
|
||||
|
||||
CUSTOM_MODEL_PATH = Paths.model_root()
|
||||
|
||||
|
||||
class ModelData:
|
||||
"""
|
||||
Stores metadata and configuration for a specific machine learning model.
|
||||
|
||||
This class loads model metadata (like input shapes and output slices)
|
||||
from a pickle file associated with a model instance.
|
||||
|
||||
:param model: The machine learning model object containing metadata.
|
||||
"""
|
||||
def __init__(self, model: Model):
|
||||
self.model = model
|
||||
self.metadata = model.metadata
|
||||
self.input_shapes: ShapeDict = {}
|
||||
self.output_slices: SliceDict = {}
|
||||
if self.metadata:
|
||||
self._load_metadata()
|
||||
|
||||
def _load_metadata(self) -> None:
|
||||
"""Loads input shapes and output slices from the model's metadata pickle file."""
|
||||
metadata_path = f"{CUSTOM_MODEL_PATH}/{self.metadata.fileName}"
|
||||
with open(metadata_path, 'rb') as f:
|
||||
model_metadata = pickle.load(f)
|
||||
self.input_shapes = model_metadata.get('input_shapes', {})
|
||||
self.output_slices = model_metadata.get('output_slices', {})
|
||||
|
||||
|
||||
class ModularRunner(ABC):
|
||||
"""
|
||||
Represents a modular runner for handling and slicing model outputs.
|
||||
|
||||
This abstract base class is designed to provide an interface for modular
|
||||
parsing and processing of model outputs. Classes inheriting from it must
|
||||
implement the specified abstract methods, defining how model outputs
|
||||
should be handled and stored. The primary goal is to enable structured
|
||||
parsing of outputs through a dictionary-based method mapping.
|
||||
|
||||
:ivar parser_method_dict: Mapping dictionary containing parser methods
|
||||
for handling specific types of outputs.
|
||||
:type parser_method_dict: dict
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def parser_method_dict(self) -> dict:
|
||||
pass
|
||||
|
||||
@parser_method_dict.setter
|
||||
@abstractmethod
|
||||
def parser_method_dict(self, value: dict) -> None:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def _slice_outputs(self, model_outputs: np.ndarray) -> NumpyDict:
|
||||
pass
|
||||
|
||||
|
||||
class ModelRunner(ModularRunner):
|
||||
"""
|
||||
Abstract base class for managing and executing machine learning models.
|
||||
|
||||
Provides a common interface for loading models, preparing inputs, running
|
||||
inference, and slicing/parsing outputs based on model metadata. Derived
|
||||
classes implement the specifics of input preparation and model execution
|
||||
for different frameworks (e.g., Tinygrad, ONNX).
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initializes the model runner, loading the active model bundle."""
|
||||
self.is_20hz: bool | None = None
|
||||
self.is_20hz_3d: bool | None = None
|
||||
self.models: dict[int, ModelData] = {}
|
||||
self._model_data: ModelData | None = None # Active model data for current operation
|
||||
self._parser_method_dict: dict = {}
|
||||
self.inputs: dict = {}
|
||||
self._parser = None
|
||||
self._load_models()
|
||||
self._constants = None
|
||||
|
||||
@property
|
||||
def constants(self):
|
||||
return self._constants
|
||||
|
||||
@property
|
||||
def parser_method_dict(self) -> dict:
|
||||
"""Returns the dictionary mapping model types to their respective parsing methods."""
|
||||
return self._parser_method_dict
|
||||
|
||||
@parser_method_dict.setter
|
||||
def parser_method_dict(self, value: dict) -> None:
|
||||
"""Sets the dictionary mapping model types to their respective parsing methods."""
|
||||
self._parser_method_dict = value
|
||||
|
||||
def _load_models(self) -> None:
|
||||
"""Loads the active model bundle configuration and sets up ModelData."""
|
||||
bundle = get_active_bundle()
|
||||
if not bundle:
|
||||
raise ValueError("No active model bundle found, why are we being executed?")
|
||||
|
||||
self.models = {model.type.raw: ModelData(model) for model in bundle.models}
|
||||
self.is_20hz = bundle.is20hz
|
||||
self.is_20hz_3d = False
|
||||
|
||||
@property
|
||||
def input_shapes(self) -> ShapeDict:
|
||||
"""Returns the input shapes for the currently active model."""
|
||||
if self._model_data:
|
||||
return self._model_data.input_shapes
|
||||
raise ValueError("Model data is not available. Ensure the model is loaded correctly.")
|
||||
|
||||
@property
|
||||
def output_slices(self) -> SliceDict:
|
||||
"""Returns the output slices for the currently active model."""
|
||||
if self._model_data:
|
||||
return self._model_data.output_slices
|
||||
raise ValueError("Model data is not available. Ensure the model is loaded correctly.")
|
||||
|
||||
@property
|
||||
def vision_input_names(self) -> list[str]:
|
||||
"""Returns the list of vision input names from the input shapes."""
|
||||
if self._model_data:
|
||||
return list(self._model_data.input_shapes.keys())
|
||||
raise ValueError("Model data is not available. Ensure the model is loaded correctly.")
|
||||
|
||||
@abstractmethod
|
||||
def prepare_inputs(self, numpy_inputs: NumpyDict) -> dict:
|
||||
"""
|
||||
Abstract method to prepare inputs for model inference.
|
||||
|
||||
:param numpy_inputs: Dictionary of numpy arrays for non-image inputs.
|
||||
:return: Dictionary of prepared inputs ready for the model.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def _run_model(self) -> NumpyDict:
|
||||
"""
|
||||
Abstract method to execute model inference with prepared inputs.
|
||||
|
||||
:return: Dictionary containing the model's raw output arrays.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def _slice_outputs(self, model_outputs: np.ndarray) -> NumpyDict:
|
||||
"""
|
||||
Slices the raw model output array based on the output_slices metadata.
|
||||
|
||||
:param model_outputs: The raw numpy array output from the model.
|
||||
:return: A dictionary where keys are output names and values are sliced numpy arrays.
|
||||
"""
|
||||
if not self._model_data:
|
||||
raise ValueError("Model data is not available. Ensure the model is loaded correctly.")
|
||||
sliced_outputs = {k: model_outputs[np.newaxis, v] for k, v in self._model_data.output_slices.items()}
|
||||
if SEND_RAW_PRED:
|
||||
sliced_outputs['raw_pred'] = model_outputs.copy() # Optionally include the full raw output
|
||||
return sliced_outputs
|
||||
|
||||
def run_model(self) -> NumpyDict:
|
||||
"""
|
||||
Executes the model inference pipeline: runs the model and parses outputs.
|
||||
|
||||
:return: Dictionary containing the final parsed model outputs.
|
||||
"""
|
||||
return self._run_model() # Parsing is handled within specific runner implementations
|
||||
@@ -1,91 +0,0 @@
|
||||
import os
|
||||
from abc import ABC
|
||||
|
||||
import numpy as np
|
||||
from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser as CombinedParser
|
||||
from openpilot.sunnypilot.modeld_v2.parse_model_outputs_split import Parser as SplitParser
|
||||
from openpilot.sunnypilot.models.runners.constants import ModelType, NumpyDict
|
||||
from openpilot.sunnypilot.models.runners.model_runner import ModularRunner
|
||||
from openpilot.common.hardware.hw import Paths
|
||||
|
||||
|
||||
SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
|
||||
CUSTOM_MODEL_PATH = Paths.model_root()
|
||||
|
||||
|
||||
class OffPolicyTinygrad(ModularRunner, ABC):
|
||||
"""
|
||||
A TinygradRunner specialized for off-policy models.
|
||||
|
||||
Uses a SplitParser to handle outputs specific to the off-policy part of a split model setup.
|
||||
"""
|
||||
def __init__(self):
|
||||
self._off_policy_parser = SplitParser()
|
||||
self.parser_method_dict[ModelType.offPolicy] = self._parse_off_policy_outputs
|
||||
|
||||
def _parse_off_policy_outputs(self, model_outputs: np.ndarray) -> NumpyDict:
|
||||
"""Parses off-policy model outputs using SplitParser."""
|
||||
result: NumpyDict = self._off_policy_parser.parse_policy_outputs(self._slice_outputs(model_outputs))
|
||||
return result
|
||||
|
||||
|
||||
class OnPolicyTinygrad(ModularRunner, ABC):
|
||||
"""
|
||||
A TinygradRunner specialized for on-policy models.
|
||||
|
||||
Uses a SplitParser to handle outputs specific to the on-policy part of a split model setup.
|
||||
"""
|
||||
def __init__(self):
|
||||
self._on_policy_parser = SplitParser()
|
||||
self.parser_method_dict[ModelType.onPolicy] = self._parse_on_policy_outputs
|
||||
|
||||
def _parse_on_policy_outputs(self, model_outputs: np.ndarray) -> NumpyDict:
|
||||
"""Parses on-policy model outputs using SplitParser."""
|
||||
result: NumpyDict = self._on_policy_parser.parse_policy_outputs(self._slice_outputs(model_outputs))
|
||||
return result
|
||||
|
||||
|
||||
class PolicyTinygrad(ModularRunner, ABC):
|
||||
"""
|
||||
A TinygradRunner specialized for policy-only models.
|
||||
|
||||
Uses a SplitParser to handle outputs specific to the policy part of a split model setup.
|
||||
"""
|
||||
def __init__(self):
|
||||
self._policy_parser = SplitParser()
|
||||
self.parser_method_dict[ModelType.policy] = self._parse_policy_outputs
|
||||
|
||||
def _parse_policy_outputs(self, model_outputs: np.ndarray) -> NumpyDict:
|
||||
"""Parses policy model outputs using SplitParser."""
|
||||
result: NumpyDict = self._policy_parser.parse_policy_outputs(self._slice_outputs(model_outputs))
|
||||
return result
|
||||
|
||||
class VisionTinygrad(ModularRunner, ABC):
|
||||
"""
|
||||
A TinygradRunner specialized for vision-only models.
|
||||
|
||||
Uses a SplitParser to handle outputs specific to the vision part of a split model setup.
|
||||
"""
|
||||
def __init__(self):
|
||||
self._vision_parser = SplitParser()
|
||||
self.parser_method_dict[ModelType.vision] = self._parse_vision_outputs
|
||||
|
||||
def _parse_vision_outputs(self, model_outputs: np.ndarray) -> NumpyDict:
|
||||
"""Parses vision model outputs using SplitParser."""
|
||||
result: NumpyDict = self._vision_parser.parse_vision_outputs(self._slice_outputs(model_outputs))
|
||||
return result
|
||||
|
||||
class SupercomboTinygrad(ModularRunner, ABC):
|
||||
"""
|
||||
A TinygradRunner specialized for vision-only models.
|
||||
|
||||
Uses a SplitParser to handle outputs specific to the vision part of a split model setup.
|
||||
"""
|
||||
def __init__(self):
|
||||
self._supercombo_parser = CombinedParser()
|
||||
self.parser_method_dict[ModelType.supercombo] = self._parse_supercombo_outputs
|
||||
|
||||
def _parse_supercombo_outputs(self, model_outputs: np.ndarray) -> NumpyDict:
|
||||
"""Parses vision model outputs using SplitParser."""
|
||||
result: NumpyDict = self._supercombo_parser.parse_outputs(self._slice_outputs(model_outputs))
|
||||
return result
|
||||
@@ -1,179 +0,0 @@
|
||||
import pickle
|
||||
|
||||
import numpy as np
|
||||
from openpilot.sunnypilot.models.runners.constants import NumpyDict, ModelType, ShapeDict, CUSTOM_MODEL_PATH, SliceDict
|
||||
from openpilot.sunnypilot.models.runners.model_runner import ModelRunner
|
||||
from openpilot.sunnypilot.models.runners.tinygrad.model_types import PolicyTinygrad, VisionTinygrad, SupercomboTinygrad, OffPolicyTinygrad, OnPolicyTinygrad
|
||||
from openpilot.sunnypilot.models.split_model_constants import SplitModelConstants
|
||||
from openpilot.sunnypilot.modeld_v2.constants import ModelConstants
|
||||
|
||||
from tinygrad.tensor import Tensor
|
||||
|
||||
|
||||
class TinygradRunner(ModelRunner, SupercomboTinygrad, PolicyTinygrad, VisionTinygrad, OffPolicyTinygrad, OnPolicyTinygrad):
|
||||
"""
|
||||
A ModelRunner implementation for executing Tinygrad models.
|
||||
|
||||
Handles loading Tinygrad model artifacts (.pkl), preparing inputs as Tinygrad
|
||||
Tensors (potentially using QCOM extensions on TICI), running inference,
|
||||
and parsing the outputs.
|
||||
|
||||
:param model_type: The type of model (e.g., supercombo) to load and run.
|
||||
"""
|
||||
def __init__(self, model_type: int = ModelType.supercombo):
|
||||
ModelRunner.__init__(self)
|
||||
SupercomboTinygrad.__init__(self)
|
||||
PolicyTinygrad.__init__(self)
|
||||
VisionTinygrad.__init__(self)
|
||||
OffPolicyTinygrad.__init__(self)
|
||||
OnPolicyTinygrad.__init__(self)
|
||||
self._constants = ModelConstants
|
||||
self._model_data = self.models.get(model_type)
|
||||
if not self._model_data or not self._model_data.model:
|
||||
raise ValueError(f"Model data for type {model_type} not available.")
|
||||
|
||||
artifact_filename = self._model_data.model.artifact.fileName
|
||||
assert artifact_filename.endswith('_tinygrad.pkl'), \
|
||||
f"Invalid model file {artifact_filename} for TinygradRunner"
|
||||
|
||||
model_pkl_path = f"{CUSTOM_MODEL_PATH}/{artifact_filename}"
|
||||
with open(model_pkl_path, "rb") as f:
|
||||
try:
|
||||
# Load the compiled Tinygrad model runner function
|
||||
self.model_run = pickle.load(f)
|
||||
except FileNotFoundError as e:
|
||||
# Provide a helpful error message if the model was built for a different platform
|
||||
assert "/dev/kgsl-3d0" not in str(e), "Model was built on C3 or C3X, but is being loaded on PC"
|
||||
raise
|
||||
|
||||
# Map input names to their required dtype and device from the loaded model
|
||||
self.input_to_dtype = {}
|
||||
self.input_to_device = {}
|
||||
for idx, name in enumerate(self.model_run.captured.expected_names):
|
||||
info = self.model_run.captured.expected_input_info[idx]
|
||||
self.input_to_dtype[name] = info[2] # dtype
|
||||
self.input_to_device[name] = info[3] # device
|
||||
self._policy_cached = False
|
||||
|
||||
@property
|
||||
def vision_input_names(self) -> list[str]:
|
||||
"""Returns the list of vision input names from the input shapes."""
|
||||
return [name for name in self.input_shapes.keys() if 'img' in name]
|
||||
|
||||
|
||||
def prepare_policy_inputs(self, numpy_inputs: NumpyDict):
|
||||
if not self._policy_cached:
|
||||
for key, value in numpy_inputs.items():
|
||||
self.inputs[key] = Tensor(value, device='NPY').realize()
|
||||
self._policy_cached = True
|
||||
|
||||
def prepare_inputs(self, numpy_inputs: NumpyDict) -> dict:
|
||||
"""Prepares all vision and policy inputs for the model."""
|
||||
self.prepare_policy_inputs(numpy_inputs)
|
||||
for key in self.vision_input_names:
|
||||
if key in self.inputs:
|
||||
self.inputs[key] = self.inputs[key].cast(self.input_to_dtype[key])
|
||||
return self.inputs
|
||||
|
||||
def _run_model(self) -> NumpyDict:
|
||||
"""Runs the Tinygrad model inference and parses the outputs."""
|
||||
outputs = self.model_run(**self.inputs).contiguous().realize().uop.base.buffer.numpy().flatten()
|
||||
return self._parse_outputs(outputs)
|
||||
|
||||
def _parse_outputs(self, model_outputs: np.ndarray) -> NumpyDict:
|
||||
"""Parses the raw model outputs using the standard Parser."""
|
||||
if self._model_data is None:
|
||||
raise ValueError("Model data is not available. Ensure the model is loaded correctly.")
|
||||
|
||||
result: NumpyDict = self.parser_method_dict[self._model_data.model.type.raw](model_outputs)
|
||||
return result
|
||||
|
||||
|
||||
class TinygradSplitRunner(ModelRunner):
|
||||
"""
|
||||
A ModelRunner that coordinates separate TinygradVisionRunner and TinygradPolicyRunner instances.
|
||||
|
||||
Manages the execution of split vision and policy models, combining their inputs and outputs.
|
||||
"""
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.is_20hz_3d = True
|
||||
self.vision_runner = TinygradRunner(ModelType.vision)
|
||||
self.policy_runner = TinygradRunner(ModelType.policy) if self.models.get(ModelType.policy) else None
|
||||
self.off_policy_runner = TinygradRunner(ModelType.offPolicy) if self.models.get(ModelType.offPolicy) else None
|
||||
self.on_policy_runner = TinygradRunner(ModelType.onPolicy) if self.models.get(ModelType.onPolicy) else None
|
||||
self._constants = SplitModelConstants
|
||||
|
||||
def _run_model(self) -> NumpyDict:
|
||||
"""Runs both vision and policy models and merges their parsed outputs."""
|
||||
vision_output = self.vision_runner.run_model()
|
||||
outputs = {**vision_output}
|
||||
|
||||
if self.policy_runner:
|
||||
policy_output = self.policy_runner.run_model()
|
||||
outputs.update(policy_output)
|
||||
|
||||
if self.off_policy_runner:
|
||||
off_policy_output = self.off_policy_runner.run_model()
|
||||
if self.on_policy_runner:
|
||||
off_policy_output.pop('plan', None)
|
||||
outputs.update(off_policy_output)
|
||||
|
||||
if self.on_policy_runner:
|
||||
on_policy_output = self.on_policy_runner.run_model()
|
||||
outputs.update(on_policy_output)
|
||||
|
||||
if 'planplus' in outputs and 'plan' in outputs:
|
||||
outputs['plan'] = outputs['plan'] + outputs['planplus']
|
||||
|
||||
return outputs
|
||||
|
||||
@property
|
||||
def vision_input_names(self) -> list[str]:
|
||||
"""Returns the list of vision input names from the vision runner."""
|
||||
return list(self.vision_runner.vision_input_names)
|
||||
|
||||
@property
|
||||
def input_shapes(self) -> ShapeDict:
|
||||
"""Returns the combined input shapes from both vision and policy models."""
|
||||
shapes = {**self.vision_runner.input_shapes}
|
||||
if self.policy_runner:
|
||||
shapes.update(self.policy_runner.input_shapes)
|
||||
if self.off_policy_runner:
|
||||
shapes.update(self.off_policy_runner.input_shapes)
|
||||
if self.on_policy_runner:
|
||||
shapes.update(self.on_policy_runner.input_shapes)
|
||||
return shapes
|
||||
|
||||
@property
|
||||
def output_slices(self) -> SliceDict:
|
||||
"""Returns the combined output slices from both vision and policy models."""
|
||||
slices = {**self.vision_runner.output_slices}
|
||||
if self.policy_runner:
|
||||
slices.update(self.policy_runner.output_slices)
|
||||
if self.off_policy_runner:
|
||||
slices.update(self.off_policy_runner.output_slices)
|
||||
if self.on_policy_runner:
|
||||
slices.update(self.on_policy_runner.output_slices)
|
||||
return slices
|
||||
|
||||
def prepare_inputs(self, numpy_inputs: NumpyDict) -> dict:
|
||||
"""Prepares inputs for both vision and policy models."""
|
||||
if self.policy_runner:
|
||||
self.policy_runner.prepare_policy_inputs(numpy_inputs)
|
||||
|
||||
for key in self.vision_input_names:
|
||||
if key in self.inputs:
|
||||
self.vision_runner.inputs[key] = self.inputs[key].cast(self.vision_runner.input_to_dtype[key])
|
||||
|
||||
inputs = {**self.vision_runner.inputs}
|
||||
if self.policy_runner:
|
||||
inputs.update(self.policy_runner.inputs)
|
||||
|
||||
if self.off_policy_runner:
|
||||
self.off_policy_runner.prepare_policy_inputs(numpy_inputs)
|
||||
inputs.update(self.off_policy_runner.inputs)
|
||||
if self.on_policy_runner:
|
||||
self.on_policy_runner.prepare_policy_inputs(numpy_inputs)
|
||||
inputs.update(self.on_policy_runner.inputs)
|
||||
return inputs
|
||||
@@ -43,6 +43,7 @@ class SplitModelConstants:
|
||||
LANE_LINES_WIDTH = 2
|
||||
ROAD_EDGES_WIDTH = 2
|
||||
PLAN_WIDTH = 15
|
||||
ACTION_WIDTH = 2
|
||||
DESIRE_PRED_WIDTH = 8
|
||||
LAT_PLANNER_SOLUTION_WIDTH = 4
|
||||
DESIRED_CURV_WIDTH = 1
|
||||
|
||||
@@ -151,8 +151,8 @@ class SmartCruiseControlMap:
|
||||
a = 0.5 * TARGET_JERK
|
||||
b = self.a_ego
|
||||
c = self.v_ego - tv
|
||||
t_a = -1 * ((b**2 - 4 * a * c) ** 0.5 + b) / 2 * a
|
||||
t_b = ((b**2 - 4 * a * c) ** 0.5 - b) / 2 * a
|
||||
t_a = -1 * ((b**2 - 4 * a * c) ** 0.5 + b) / (2 * a)
|
||||
t_b = ((b**2 - 4 * a * c) ** 0.5 - b) / (2 * a)
|
||||
if not isinstance(t_a, complex) and t_a > 0:
|
||||
t = t_a
|
||||
else:
|
||||
|
||||
+18
-1
@@ -4,13 +4,17 @@ 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 json
|
||||
import math
|
||||
import platform
|
||||
|
||||
import pytest
|
||||
|
||||
from openpilot.cereal import custom
|
||||
from openpilot.common.params import Params
|
||||
from openpilot.common.realtime import DT_MDL
|
||||
from openpilot.selfdrive.car.cruise import V_CRUISE_UNSET
|
||||
from openpilot.sunnypilot.selfdrive.controls.lib.smart_cruise_control.map_controller import SmartCruiseControlMap
|
||||
from openpilot.sunnypilot.selfdrive.controls.lib.smart_cruise_control.map_controller import R, SmartCruiseControlMap
|
||||
|
||||
MapState = VisionState = custom.LongitudinalPlanSP.SmartCruiseControl.MapState
|
||||
|
||||
@@ -55,4 +59,17 @@ class TestSmartCruiseControlMap:
|
||||
self.scc_m.update(True, False, 0., 0., 0.)
|
||||
assert self.scc_m.state == VisionState.enabled
|
||||
|
||||
def test_moderate_curve(self):
|
||||
# Regression: `... / 2 * a` parsed as `(.../2)*a` instead of `.../(2*a)`,
|
||||
# making max_d ~11x too small so the moderate-curve branch never tripped.
|
||||
# v_ego=25, a_ego=0, tv=24: fixed max_d≈45m vs buggy ≈4m at a 40m waypoint.
|
||||
waypoint_lon_deg = (40.0 / R) * (180.0 / math.pi)
|
||||
self.mem_params.put("LastGPSPosition", json.dumps({"latitude": 0.0, "longitude": 0.0}), block=True)
|
||||
self.mem_params.put("MapTargetVelocities",
|
||||
json.dumps([{"latitude": 0.0, "longitude": waypoint_lon_deg, "velocity": 24.0}]), block=True)
|
||||
|
||||
self.scc_m.update(True, False, 25.0, 0.0, 30.0)
|
||||
|
||||
assert self.scc_m.v_target == pytest.approx(24.0)
|
||||
|
||||
# TODO-SP: mock data from modelV2 to test other states
|
||||
|
||||
@@ -95,9 +95,6 @@ def is_stock_model(started, params, CP: car.CarParams) -> bool:
|
||||
"""Check if the active model runner is stock."""
|
||||
return bool(get_active_model_runner(params, not started) == custom.ModelManagerSP.Runner.stock)
|
||||
|
||||
def not_wgpu(started: bool, params: Params, CP: car.CarParams) -> bool:
|
||||
return not params.get_bool("WgpuEnabled")
|
||||
|
||||
def mapd_ready(started: bool, params: Params, CP: car.CarParams) -> bool:
|
||||
return bool(os.path.exists(Paths.mapd_root()))
|
||||
|
||||
@@ -131,7 +128,7 @@ procs = [
|
||||
PythonProcess("micd", "openpilot.system.micd", iscar),
|
||||
PythonProcess("timed", "openpilot.system.timed", always_run, enabled=not PC),
|
||||
|
||||
PythonProcess("modeld", "openpilot.selfdrive.modeld.modeld", and_(and_(only_onroad, is_stock_model), not_wgpu)),
|
||||
PythonProcess("modeld", "openpilot.selfdrive.modeld.modeld", and_(only_onroad, is_stock_model)),
|
||||
PythonProcess("dmonitoringmodeld", "openpilot.selfdrive.modeld.dmonitoringmodeld", driverview, enabled=(WEBCAM or not PC)),
|
||||
|
||||
PythonProcess("sensord", "openpilot.system.sensord.sensord", only_onroad, enabled=not PC),
|
||||
@@ -180,7 +177,7 @@ procs = [
|
||||
procs += [
|
||||
# Models
|
||||
PythonProcess("models_manager", "openpilot.sunnypilot.models.manager", only_offroad),
|
||||
NativeProcess("modeld_tinygrad", "openpilot/sunnypilot/modeld_v2", ["./modeld"], and_(and_(only_onroad, is_tinygrad_model), not_wgpu)),
|
||||
NativeProcess("modeld_tinygrad", "openpilot/sunnypilot/modeld_v2", ["./modeld"], and_(only_onroad, is_tinygrad_model)),
|
||||
|
||||
# Backup
|
||||
PythonProcess("backup_manager", "openpilot.sunnypilot.sunnylink.backups.manager", and_(only_offroad, sunnylink_ready_shim)),
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
#!/usr/bin/env python3
|
||||
import os
|
||||
import argparse
|
||||
import multiprocessing
|
||||
import time
|
||||
@@ -9,7 +10,6 @@ from collections import deque
|
||||
import openpilot.cereal.messaging as messaging
|
||||
from msgq.visionipc import VisionIpcServer, VisionStreamType
|
||||
from openpilot.tools.camerastream.ffmpeg_decoder import Decoder, FFmpegError
|
||||
from openpilot.tools.wgpu.zmq import ZmqSubMaster, ZmqSubSocket
|
||||
|
||||
V4L2_BUF_FLAG_KEYFRAME = 8
|
||||
|
||||
@@ -30,7 +30,10 @@ def decoder(addr, vipc_server, vst, W, H, debug=False):
|
||||
|
||||
codec = Decoder("hevc")
|
||||
|
||||
sock = ZmqSubSocket(sock_name, addr)
|
||||
os.environ["ZMQ"] = "1"
|
||||
messaging.reset_context()
|
||||
sock = messaging.sub_sock(sock_name, None, addr=addr, conflate=False)
|
||||
cnt = 0
|
||||
last_idx = -1
|
||||
seen_iframe = False
|
||||
|
||||
@@ -43,9 +46,8 @@ def decoder(addr, vipc_server, vst, W, H, debug=False):
|
||||
time_q.clear()
|
||||
|
||||
while 1:
|
||||
msgs = sock.drain(wait_for_one=True)
|
||||
for raw in msgs:
|
||||
evt = messaging.log_from_bytes(raw)
|
||||
msgs = messaging.drain_sock(sock, wait_for_one=True)
|
||||
for evt in msgs:
|
||||
evta = getattr(evt, evt.which())
|
||||
if last_idx != -1 and evta.idx.encodeId != (last_idx + 1):
|
||||
if debug:
|
||||
@@ -92,9 +94,8 @@ def decoder(addr, vipc_server, vst, W, H, debug=False):
|
||||
continue
|
||||
|
||||
frame_start_time = time_q.popleft()
|
||||
# Preserve the device camera metadata so remote model outputs line up with
|
||||
# the rest of the device's cereal timeline.
|
||||
vipc_server.send(vst, img_yuv.data, evta.idx.frameId, evta.idx.timestampSof, evta.idx.timestampEof)
|
||||
vipc_server.send(vst, img_yuv.data, cnt, int(frame_start_time*1e9), int(time.monotonic()*1e9))
|
||||
cnt += 1
|
||||
|
||||
pc_latency = (time.monotonic()-frame_start_time)*1000
|
||||
if debug:
|
||||
@@ -104,30 +105,25 @@ def decoder(addr, vipc_server, vst, W, H, debug=False):
|
||||
|
||||
class CompressedVipc:
|
||||
def __init__(self, addr, vision_streams, server_name, debug=False):
|
||||
print("waiting for remote camera stream metadata", flush=True)
|
||||
sm = ZmqSubMaster([ENCODE_SOCKETS[s] for s in vision_streams], addr)
|
||||
print("getting frame sizes")
|
||||
os.environ["ZMQ"] = "1"
|
||||
messaging.reset_context()
|
||||
sm = messaging.SubMaster([ENCODE_SOCKETS[s] for s in vision_streams], addr=addr)
|
||||
while min(sm.recv_frame.values()) == 0:
|
||||
sm.update(100)
|
||||
|
||||
stream_dimensions = {
|
||||
vst: (sm[ENCODE_SOCKETS[vst]].width, sm[ENCODE_SOCKETS[vst]].height)
|
||||
for vst in vision_streams
|
||||
}
|
||||
# The metadata subscribers are setup-only. Leaving them connected creates a
|
||||
# second unread camera subscription whose TCP queues grow for the entire run.
|
||||
sm.close()
|
||||
os.environ.pop("ZMQ")
|
||||
messaging.reset_context()
|
||||
|
||||
self.vipc_server = VisionIpcServer(server_name)
|
||||
for vst in vision_streams:
|
||||
width, height = stream_dimensions[vst]
|
||||
self.vipc_server.create_buffers(vst, 4, width, height)
|
||||
ed = sm[ENCODE_SOCKETS[vst]]
|
||||
self.vipc_server.create_buffers(vst, 4, ed.width, ed.height)
|
||||
self.vipc_server.start_listener()
|
||||
|
||||
self.procs = []
|
||||
process_context = multiprocessing.get_context("fork")
|
||||
for vst in vision_streams:
|
||||
width, height = stream_dimensions[vst]
|
||||
p = process_context.Process(target=decoder, args=(addr, self.vipc_server, vst, width, height, debug))
|
||||
ed = sm[ENCODE_SOCKETS[vst]]
|
||||
p = multiprocessing.Process(target=decoder, args=(addr, self.vipc_server, vst, ed.width, ed.height, debug))
|
||||
p.start()
|
||||
self.procs.append(p)
|
||||
|
||||
|
||||
@@ -46,23 +46,10 @@ def _bind(fn, restype, *argtypes):
|
||||
return fn
|
||||
|
||||
|
||||
def _library_path(name: str, major: int) -> str:
|
||||
candidates = (
|
||||
f"lib{name}.so.{major}",
|
||||
f"lib{name}.{major}.dylib",
|
||||
f"lib{name}.dylib",
|
||||
)
|
||||
for candidate in candidates:
|
||||
path = os.path.join(ffmpeg.LIB_DIR, candidate)
|
||||
if os.path.isfile(path):
|
||||
return path
|
||||
raise FileNotFoundError(f"FFmpeg library not found in {ffmpeg.LIB_DIR}: {', '.join(candidates)}")
|
||||
|
||||
|
||||
def _load_libraries():
|
||||
avutil = ctypes.CDLL(_library_path("avutil", 59), mode=ctypes.RTLD_GLOBAL)
|
||||
avcodec = ctypes.CDLL(_library_path("avcodec", 61), mode=ctypes.RTLD_GLOBAL)
|
||||
swscale = ctypes.CDLL(_library_path("swscale", 8), mode=ctypes.RTLD_GLOBAL)
|
||||
avutil = ctypes.CDLL(os.path.join(ffmpeg.LIB_DIR, "libavutil.so.59"), mode=ctypes.RTLD_GLOBAL)
|
||||
avcodec = ctypes.CDLL(os.path.join(ffmpeg.LIB_DIR, "libavcodec.so.61"), mode=ctypes.RTLD_GLOBAL)
|
||||
swscale = ctypes.CDLL(os.path.join(ffmpeg.LIB_DIR, "libswscale.so.8"), mode=ctypes.RTLD_GLOBAL)
|
||||
|
||||
c_int, c_char_p, c_void_p, c_size_t = ctypes.c_int, ctypes.c_char_p, ctypes.c_void_p, ctypes.c_size_t
|
||||
c_uint8_p = ctypes.POINTER(ctypes.c_uint8)
|
||||
|
||||
@@ -1,74 +0,0 @@
|
||||
# Wireless modeld proof of concept
|
||||
|
||||
This runs driving `modeld` on a laptop and returns its cereal outputs to a comma
|
||||
device over the existing Wi-Fi network. It reuses the existing HEVC camera
|
||||
stream, VisionIPC decoder, and cereal ZMQ bridge.
|
||||
|
||||
This is for controlled bench testing only. Wi-Fi has no deterministic latency
|
||||
or availability guarantee. The device-side helper switches only after receiving
|
||||
a fresh remote model and after manager has stopped the local model publisher.
|
||||
It restores local `modeld` if the remote model is missing for one second.
|
||||
|
||||
## Build
|
||||
|
||||
Use the same commit on the laptop and comma device. Build the cereal bridge on
|
||||
the device:
|
||||
|
||||
```sh
|
||||
scons -u openpilot/cereal/messaging/bridge
|
||||
```
|
||||
|
||||
Build and test the normal model on the laptop first:
|
||||
|
||||
```sh
|
||||
PATH="$PWD/.venv/bin:$PATH" scons -u
|
||||
```
|
||||
|
||||
To compile the big external-GPU model for the laptop's local tinygrad backend:
|
||||
|
||||
```sh
|
||||
PATH="$PWD/.venv/bin:$PATH" WGPU=1 scons -u \
|
||||
openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl.chunkmanifest
|
||||
```
|
||||
|
||||
On macOS, the build selects tinygrad's Metal backend when it is available.
|
||||
On an 8-GPU-core M5 MacBook Air, the small model's compiled policy pass measured
|
||||
about 6–11 ms, while the big model measured about 79–81 ms. The latter already
|
||||
misses the 50 ms model cadence before network and codec latency, so start with
|
||||
the small model on that class of laptop.
|
||||
|
||||
## Run
|
||||
|
||||
Find the laptop's LAN IP address that the comma device can reach. The helper can
|
||||
start while onroad: it forwards camera/state while local `modeld` remains active,
|
||||
then performs an exclusive publisher handoff after the laptop produces a fresh
|
||||
valid model:
|
||||
|
||||
```sh
|
||||
cd /data/openpilot
|
||||
python3 -m openpilot.tools.wgpu.device LAPTOP_IP
|
||||
```
|
||||
|
||||
Keep that terminal open. On the laptop, run:
|
||||
|
||||
```sh
|
||||
cd /path/to/openpilot
|
||||
python3 -m openpilot.tools.wgpu.host COMMA_IP
|
||||
```
|
||||
|
||||
Add `--big-model` after `COMMA_IP` to use the locally compiled big model.
|
||||
|
||||
The device helper sends cached `carParams` over a dedicated startup channel, so
|
||||
the host does not attach to the camera streams and accumulate stale frames while
|
||||
waiting for the periodic state bridge. Stop either side with Ctrl+C. A host
|
||||
disconnect automatically stops remote publication and restores local `modeld`
|
||||
after a one-second timeout; the device helper then stays running and waits for
|
||||
the next host session. Stop the device helper itself with Ctrl+C before changing
|
||||
branches or rebooting.
|
||||
|
||||
While WGPU is active, model lag does not create an engagement-blocking alert.
|
||||
The fixed-column diagnostics panel at the top-right of the mici onroad UI shows
|
||||
the active source (`LOCAL` or `WGPU`), remote model size (`SMALL` or `BIG`),
|
||||
total model-frame age, execution time, remaining I/O and queue time, and dropped
|
||||
frames. The home GPU icon is gray while the device helper is ready and waiting,
|
||||
and green while a valid remote model is active.
|
||||
@@ -1,152 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
import argparse
|
||||
import signal
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import openpilot.cereal.messaging as messaging
|
||||
from openpilot.common.params import Params
|
||||
from openpilot.tools.wgpu.zmq import WGPU_CAR_PARAMS, ZmqPubMaster, ZmqSubSocket
|
||||
|
||||
|
||||
MODEL_OUTPUTS = "modelV2,drivingModelData,cameraOdometry,modelDataV2SP"
|
||||
MODEL_PROCESSES = {"modeld", "modeld_tinygrad"}
|
||||
REMOTE_MODEL_TIMEOUT = 1.0
|
||||
WGPU_STATUS = "wgpuStatus"
|
||||
ROOT = Path(__file__).resolve().parents[3]
|
||||
BRIDGE = ROOT / "openpilot/cereal/messaging/bridge"
|
||||
|
||||
|
||||
def stop_process(proc: subprocess.Popen) -> None:
|
||||
proc.terminate()
|
||||
try:
|
||||
proc.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
proc.wait()
|
||||
|
||||
|
||||
def handle_sigterm(*_) -> None:
|
||||
raise KeyboardInterrupt
|
||||
|
||||
|
||||
def receive_model(sock: ZmqSubSocket) -> tuple[bool, float]:
|
||||
raw = sock.receive(non_blocking=True)
|
||||
if raw is None:
|
||||
return False, float("inf")
|
||||
event = messaging.log_from_bytes(raw)
|
||||
if event.which() != "modelV2" or not event.valid:
|
||||
return True, float("inf")
|
||||
model_age = (time.monotonic_ns() - event.modelV2.timestampEof) / 1e9
|
||||
return True, model_age
|
||||
|
||||
|
||||
def receive_model_name(sock: ZmqSubSocket) -> str | None:
|
||||
raw = sock.receive(non_blocking=True)
|
||||
if raw is None:
|
||||
return None
|
||||
model_name = raw.decode(errors="replace").upper()
|
||||
return model_name if model_name in ("SMALL", "BIG") else None
|
||||
|
||||
|
||||
def wait_for_local_modeld_stop(timeout: float = 10.) -> None:
|
||||
sm = messaging.SubMaster(["managerState"])
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
sm.update(100)
|
||||
if sm.seen["managerState"]:
|
||||
running = {p.name for p in sm["managerState"].processes if p.running}
|
||||
if not running.intersection(MODEL_PROCESSES):
|
||||
return
|
||||
raise RuntimeError("timed out waiting for local modeld publisher to stop")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Route modeld traffic between this device and a wireless host.")
|
||||
parser.add_argument("host", help="Laptop IP address reachable from this device")
|
||||
args = parser.parse_args()
|
||||
|
||||
if not BRIDGE.is_file():
|
||||
raise FileNotFoundError(f"build the cereal bridge first: {BRIDGE}")
|
||||
|
||||
params = Params()
|
||||
forward: subprocess.Popen | None = None
|
||||
reverse: subprocess.Popen | None = None
|
||||
wgpu_enabled = False
|
||||
try:
|
||||
forward = subprocess.Popen([str(BRIDGE)])
|
||||
params.put_bool("WgpuReady", True, block=True)
|
||||
car_params_pub = ZmqPubMaster([WGPU_CAR_PARAMS])
|
||||
remote_model = ZmqSubSocket("modelV2", args.host, conflate=True)
|
||||
remote_status = ZmqSubSocket(WGPU_STATUS, args.host, conflate=True)
|
||||
car_params = params.get("CarParams") or params.get("CarParamsPersistent")
|
||||
|
||||
while True:
|
||||
# Keep local modeld publishing while a host connects and warms up.
|
||||
print(f"forwarding camera/state to {args.host}; waiting for a fresh remote model")
|
||||
model_name = None
|
||||
while True:
|
||||
if car_params is None:
|
||||
car_params = params.get("CarParams") or params.get("CarParamsPersistent")
|
||||
if car_params is not None:
|
||||
# This dedicated conflated startup channel lets remote modeld obtain CP
|
||||
# before it connects to VisionIPC and accumulates stale frame metadata.
|
||||
car_params_pub.send_raw(WGPU_CAR_PARAMS, car_params)
|
||||
received, model_age = receive_model(remote_model)
|
||||
model_name = receive_model_name(remote_status) or model_name
|
||||
if received and 0 <= model_age < REMOTE_MODEL_TIMEOUT and model_name is not None:
|
||||
break
|
||||
if forward.poll() is not None:
|
||||
raise RuntimeError(f"forward bridge exited with status {forward.returncode}")
|
||||
time.sleep(0.05)
|
||||
|
||||
# Stop the local publisher before attaching the reverse bridge. msgq permits
|
||||
# only one publisher for each model service.
|
||||
params.put("WgpuModelName", model_name, block=True)
|
||||
params.put_bool("WgpuEnabled", True, block=True)
|
||||
wgpu_enabled = True
|
||||
wait_for_local_modeld_stop()
|
||||
reverse = subprocess.Popen([str(BRIDGE), args.host, MODEL_OUTPUTS])
|
||||
print("wgpu active; Ctrl+C or loss of the remote model restores local modeld")
|
||||
|
||||
last_remote_model = time.monotonic()
|
||||
while True:
|
||||
received, _ = receive_model(remote_model)
|
||||
if received:
|
||||
last_remote_model = time.monotonic()
|
||||
if forward.poll() is not None:
|
||||
raise RuntimeError(f"forward bridge exited with status {forward.returncode}")
|
||||
if reverse.poll() is not None:
|
||||
print(f"reverse bridge exited with status {reverse.returncode}; restoring local modeld")
|
||||
break
|
||||
if time.monotonic() - last_remote_model > REMOTE_MODEL_TIMEOUT:
|
||||
print("remote model timed out; restoring local modeld")
|
||||
break
|
||||
time.sleep(0.05)
|
||||
|
||||
stop_process(reverse)
|
||||
reverse = None
|
||||
params.put_bool("WgpuEnabled", False, block=True)
|
||||
wgpu_enabled = False
|
||||
params.remove("WgpuModelName")
|
||||
print("wgpu disabled; local modeld restored; ready for the next host run")
|
||||
finally:
|
||||
# Stop remote publication before allowing the local publisher to restart.
|
||||
if reverse is not None and reverse.poll() is None:
|
||||
stop_process(reverse)
|
||||
if wgpu_enabled:
|
||||
params.put_bool("WgpuEnabled", False, block=True)
|
||||
params.put_bool("WgpuReady", False, block=True)
|
||||
params.remove("WgpuModelName")
|
||||
if forward is not None and forward.poll() is None:
|
||||
stop_process(forward)
|
||||
print("wgpu disabled; local modeld restored")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
signal.signal(signal.SIGTERM, handle_sigterm)
|
||||
try:
|
||||
main()
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
@@ -1,63 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
import argparse
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from openpilot.tools.wgpu.zmq import ZmqPubMaster
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[3]
|
||||
CAMERASTREAM = ROOT / "openpilot/tools/camerastream/compressed_vipc.py"
|
||||
WGPU_STATUS = "wgpuStatus"
|
||||
|
||||
|
||||
def stop_process(proc: subprocess.Popen) -> None:
|
||||
try:
|
||||
os.killpg(proc.pid, signal.SIGTERM)
|
||||
except ProcessLookupError:
|
||||
proc.wait()
|
||||
return
|
||||
try:
|
||||
proc.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
os.killpg(proc.pid, signal.SIGKILL)
|
||||
proc.wait()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Run modeld on this host for a remote comma device.")
|
||||
parser.add_argument("device", help="comma device hostname or IP address")
|
||||
parser.add_argument("--big-model", action="store_true", help="use the locally compiled big driving model")
|
||||
args = parser.parse_args()
|
||||
|
||||
camera = subprocess.Popen([sys.executable, str(CAMERASTREAM), args.device, "--cams", "0,2"], start_new_session=True)
|
||||
model_args = [sys.executable, "-m", "openpilot.selfdrive.modeld.modeld", "--remote", args.device]
|
||||
if args.big_model:
|
||||
model_args.append("--big-model")
|
||||
model = subprocess.Popen(model_args, cwd=ROOT, start_new_session=True)
|
||||
status = ZmqPubMaster([WGPU_STATUS])
|
||||
model_name = b"BIG" if args.big_model else b"SMALL"
|
||||
|
||||
procs = {"camera bridge": camera, "modeld": model}
|
||||
try:
|
||||
while all(proc.poll() is None for proc in procs.values()):
|
||||
status.send_raw(WGPU_STATUS, model_name)
|
||||
time.sleep(0.25)
|
||||
failed_name, failed = next((name, proc) for name, proc in procs.items() if proc.poll() is not None)
|
||||
raise RuntimeError(f"wgpu {failed_name} exited with status {failed.returncode}")
|
||||
finally:
|
||||
for proc in procs.values():
|
||||
if proc.poll() is None:
|
||||
stop_process(proc)
|
||||
print("wgpu host stopped")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
main()
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
@@ -1,105 +0,0 @@
|
||||
import time
|
||||
|
||||
import zmq
|
||||
|
||||
import openpilot.cereal.messaging as messaging
|
||||
|
||||
|
||||
WGPU_CAR_PARAMS = "wgpuCarParams"
|
||||
|
||||
|
||||
def service_port(endpoint: str) -> int:
|
||||
# Keep this in sync with cereal/messaging/bridge_zmq.cc.
|
||||
value = 0xcbf29ce484222325
|
||||
for char in endpoint.encode():
|
||||
value ^= char
|
||||
value = (value * 0x100000001b3) & 0xffffffffffffffff
|
||||
return 8023 + (value % (65535 - 8023))
|
||||
|
||||
|
||||
class ZmqSubSocket:
|
||||
def __init__(self, endpoint: str, address: str, conflate: bool = False):
|
||||
self.context = zmq.Context()
|
||||
self.socket = self.context.socket(zmq.SUB)
|
||||
self.socket.setsockopt(zmq.SUBSCRIBE, b"")
|
||||
self.socket.setsockopt(zmq.RECONNECT_IVL_MAX, 500)
|
||||
if conflate:
|
||||
self.socket.setsockopt(zmq.CONFLATE, 1)
|
||||
self.socket.connect(f"tcp://{address}:{service_port(endpoint)}")
|
||||
|
||||
def receive(self, non_blocking: bool = False) -> bytes | None:
|
||||
try:
|
||||
return self.socket.recv(flags=zmq.NOBLOCK if non_blocking else 0)
|
||||
except zmq.Again:
|
||||
return None
|
||||
|
||||
def drain(self, wait_for_one: bool = False) -> list[bytes]:
|
||||
messages = []
|
||||
if wait_for_one:
|
||||
message = self.receive()
|
||||
if message is not None:
|
||||
messages.append(message)
|
||||
while (message := self.receive(non_blocking=True)) is not None:
|
||||
messages.append(message)
|
||||
return messages
|
||||
|
||||
def close(self) -> None:
|
||||
self.socket.close(linger=0)
|
||||
self.context.term()
|
||||
|
||||
|
||||
class ZmqSubMaster:
|
||||
def __init__(self, services: list[str], address: str):
|
||||
self.services = services
|
||||
self.sockets = {service: ZmqSubSocket(service, address, conflate=True) for service in services}
|
||||
self.poller = zmq.Poller()
|
||||
self.socket_to_service = {}
|
||||
for service, sub in self.sockets.items():
|
||||
self.poller.register(sub.socket, zmq.POLLIN)
|
||||
self.socket_to_service[sub.socket] = service
|
||||
|
||||
self.data = {service: getattr(messaging.new_message(service).as_reader(), service) for service in services}
|
||||
self.seen = dict.fromkeys(services, False)
|
||||
self.updated = dict.fromkeys(services, False)
|
||||
self.recv_frame = dict.fromkeys(services, 0)
|
||||
self.frame = -1
|
||||
|
||||
def __getitem__(self, service: str):
|
||||
return self.data[service]
|
||||
|
||||
def update(self, timeout: int = 100) -> None:
|
||||
self.frame += 1
|
||||
self.updated = dict.fromkeys(self.services, False)
|
||||
for socket, _ in self.poller.poll(timeout):
|
||||
service = self.socket_to_service[socket]
|
||||
raw = self.sockets[service].receive(non_blocking=True)
|
||||
if raw is None:
|
||||
continue
|
||||
event = messaging.log_from_bytes(raw)
|
||||
self.data[service] = getattr(event, service)
|
||||
self.seen[service] = True
|
||||
self.updated[service] = True
|
||||
self.recv_frame[service] = self.frame
|
||||
|
||||
def close(self) -> None:
|
||||
for sub in self.sockets.values():
|
||||
self.poller.unregister(sub.socket)
|
||||
sub.close()
|
||||
|
||||
|
||||
class ZmqPubMaster:
|
||||
def __init__(self, services: list[str]):
|
||||
context = zmq.Context.instance()
|
||||
self.sockets = {}
|
||||
for service in services:
|
||||
socket = context.socket(zmq.PUB)
|
||||
socket.bind(f"tcp://*:{service_port(service)}")
|
||||
self.sockets[service] = socket
|
||||
# Give already-running bridge subscribers time to finish their handshake.
|
||||
time.sleep(0.1)
|
||||
|
||||
def send(self, service: str, message) -> None:
|
||||
self.sockets[service].send(message.to_bytes(), flags=zmq.NOBLOCK)
|
||||
|
||||
def send_raw(self, service: str, data: bytes) -> None:
|
||||
self.sockets[service].send(data, flags=zmq.NOBLOCK)
|
||||
+38
-109
@@ -5,8 +5,6 @@ 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 pickle
|
||||
import sys
|
||||
import hashlib
|
||||
import json
|
||||
@@ -14,46 +12,8 @@ import re
|
||||
from pathlib import Path
|
||||
from datetime import datetime, UTC
|
||||
|
||||
REQUIRED_OUTPUT_KEYS = frozenset({
|
||||
"plan",
|
||||
"lane_lines",
|
||||
"road_edges",
|
||||
"lead",
|
||||
"desire_state",
|
||||
"desire_pred",
|
||||
"meta",
|
||||
"lead_prob",
|
||||
"lane_lines_prob",
|
||||
"pose",
|
||||
"wide_from_device_euler",
|
||||
"road_transform",
|
||||
"hidden_state",
|
||||
})
|
||||
OPTIONAL_OUTPUT_KEYS = frozenset({
|
||||
"planplus",
|
||||
"sim_pose",
|
||||
"desired_curvature",
|
||||
})
|
||||
|
||||
|
||||
def validate_model_outputs(metadata_paths: list[Path]) -> None:
|
||||
combined_keys: set[str] = set()
|
||||
for path in metadata_paths:
|
||||
if path.stat().st_size == 0:
|
||||
print(f"skipping empty metadata: {path}")
|
||||
continue
|
||||
with open(path, "rb") as f:
|
||||
metadata = pickle.load(f)
|
||||
combined_keys.update(metadata.get("output_slices", {}).keys())
|
||||
missing = REQUIRED_OUTPUT_KEYS - combined_keys
|
||||
if missing:
|
||||
raise ValueError(f"Combined model metadata is missing required output keys: {sorted(missing)}")
|
||||
detected_optional = sorted(OPTIONAL_OUTPUT_KEYS & combined_keys)
|
||||
if detected_optional:
|
||||
print(f"Optional output keys detected: {detected_optional}")
|
||||
|
||||
|
||||
def create_short_name(full_name):
|
||||
def create_short_name(full_name: str) -> str:
|
||||
# Remove parentheses and extract alphanumeric words
|
||||
clean_name = re.sub(r'\([^)]*\)', '', full_name)
|
||||
words = [re.sub(r'[^a-zA-Z0-9]', '', word) for word in clean_name.split() if re.sub(r'[^a-zA-Z0-9]', '', word)]
|
||||
@@ -121,49 +81,41 @@ def _rename_pkl_with_chunks(old_pkl: Path, new_pkl: Path) -> Path:
|
||||
return old_pkl.rename(new_pkl)
|
||||
|
||||
|
||||
def generate_metadata(model_path: Path, output_dir: Path, short_name: str, driving_pkl: Path):
|
||||
base = model_path.stem
|
||||
metadata_file = output_dir / f"{base}_metadata.pkl"
|
||||
|
||||
if short_name:
|
||||
renamed_meta = output_dir / f"{base}_{short_name.lower()}_metadata.pkl"
|
||||
if metadata_file.exists() and not renamed_meta.exists():
|
||||
metadata_file = metadata_file.rename(renamed_meta)
|
||||
elif renamed_meta.exists():
|
||||
metadata_file = renamed_meta
|
||||
|
||||
if not metadata_file.exists():
|
||||
print(f"Warning: Missing metadata for {base} ({metadata_file}), skipping", file=sys.stderr)
|
||||
return
|
||||
|
||||
def generate_chunked_model(driving_pkl: Path) -> dict:
|
||||
tinygrad_hash = hashlib.sha256(_read_pkl_bytes(driving_pkl)).hexdigest()
|
||||
|
||||
with open(metadata_file, 'rb') as f:
|
||||
metadata_hash = hashlib.sha256(f.read()).hexdigest()
|
||||
chunks_config = []
|
||||
manifest_file = Path(f"{driving_pkl}.chunkmanifest")
|
||||
if manifest_file.exists():
|
||||
num_chunks = int(manifest_file.read_text().strip())
|
||||
for i in range(num_chunks):
|
||||
chunk_path = Path(f"{driving_pkl}.chunk{i + 1:02d}of{num_chunks:02d}")
|
||||
if chunk_path.exists():
|
||||
chunk_hash = hashlib.sha256(chunk_path.read_bytes()).hexdigest()
|
||||
chunks_config.append({
|
||||
"file_name": chunk_path.name,
|
||||
"sha256": chunk_hash
|
||||
})
|
||||
|
||||
model_type = "offPolicy" if "off_policy" in base else "onPolicy" if "on_policy" in base else base.split("_")[-1]
|
||||
|
||||
return {
|
||||
"type": model_type,
|
||||
"artifact": {
|
||||
"file_name": driving_pkl.name,
|
||||
"download_uri": {
|
||||
"url": "https://gitlab.com/sunnypilot/public/docs.sunnypilot.ai/-/raw/main/",
|
||||
"sha256": tinygrad_hash
|
||||
}
|
||||
},
|
||||
"metadata": {
|
||||
"file_name": metadata_file.name,
|
||||
"download_uri": {
|
||||
"url": "https://gitlab.com/sunnypilot/public/docs.sunnypilot.ai/-/raw/main/",
|
||||
"sha256": metadata_hash
|
||||
}
|
||||
artifact_data = {
|
||||
"file_name": driving_pkl.name,
|
||||
"download_uri": {
|
||||
"url": "https://gitlab.com/sunnypilot/public/docs.sunnypilot.ai/-/raw/main/",
|
||||
"sha256": tinygrad_hash
|
||||
}
|
||||
}
|
||||
|
||||
if chunks_config:
|
||||
artifact_data["chunks"] = chunks_config
|
||||
|
||||
def create_metadata_json(models: list, output_dir: Path, custom_name=None, short_name=None, is_20hz=False, upstream_branch="unknown"):
|
||||
metadata_json = {
|
||||
return {
|
||||
"type": "chunked",
|
||||
"artifact": artifact_data,
|
||||
}
|
||||
|
||||
|
||||
def create_metadata_json(models: list, output_dir: Path, custom_name=None, short_name=None, is_20hz=False, upstream_branch="unknown") -> None:
|
||||
bundle_json = {
|
||||
"short_name": short_name,
|
||||
"display_name": custom_name or upstream_branch,
|
||||
"is_20hz": is_20hz,
|
||||
@@ -179,40 +131,26 @@ def create_metadata_json(models: list, output_dir: Path, custom_name=None, short
|
||||
}
|
||||
|
||||
# Write metadata to output_dir
|
||||
metadata_json = {
|
||||
"bundles": [bundle_json]
|
||||
}
|
||||
|
||||
with open(output_dir / "metadata.json", "w") as f:
|
||||
json.dump(metadata_json, f, indent=2)
|
||||
|
||||
print(f"Generated metadata.json with {len(models)} models.")
|
||||
print("Generated metadata.json")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
import glob
|
||||
|
||||
parser = argparse.ArgumentParser(description="Generate metadata for model files")
|
||||
parser.add_argument("--model-dir", default="./models", help="Directory containing ONNX model files")
|
||||
parser = argparse.ArgumentParser(description="Generate metadata JSON for the compiled JIT model")
|
||||
parser.add_argument("--model-dir", default="./models", help="Directory containing the model files")
|
||||
parser.add_argument("--output-dir", default="./output", help="Output directory for metadata")
|
||||
parser.add_argument("--custom-name", help="Custom display name for the model")
|
||||
parser.add_argument("--is-20hz", action="store_true", help="Whether this is a 20Hz model")
|
||||
parser.add_argument("--validate-only", action="store_true")
|
||||
parser.add_argument("--upstream-branch", default="unknown", help="Upstream branch name")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.validate_only:
|
||||
metadata_paths = glob.glob(os.path.join(args.model_dir, "*_metadata.pkl"))
|
||||
if not metadata_paths:
|
||||
print(f"No metadata files found in {args.model_dir}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
validate_model_outputs([Path(p) for p in metadata_paths])
|
||||
print(f"Validated {len(metadata_paths)} metadata files successfully.")
|
||||
sys.exit(0)
|
||||
|
||||
# Find all ONNX files in the given directory
|
||||
model_paths = glob.glob(os.path.join(args.model_dir, "*.onnx"))
|
||||
if not model_paths:
|
||||
print(f"No ONNX files found in {args.model_dir}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
_output_dir = Path(args.output_dir)
|
||||
_output_dir.mkdir(exist_ok=True, parents=True)
|
||||
_short_name = create_short_name(args.custom_name) if args.custom_name else None
|
||||
@@ -229,14 +167,5 @@ if __name__ == "__main__":
|
||||
else:
|
||||
_driving_pkl = new_pkl
|
||||
|
||||
_models = []
|
||||
|
||||
for _model_path in model_paths:
|
||||
_model_metadata = generate_metadata(Path(_model_path), _output_dir, _short_name, _driving_pkl)
|
||||
if _model_metadata:
|
||||
_models.append(_model_metadata)
|
||||
|
||||
if _models:
|
||||
create_metadata_json(_models, _output_dir, args.custom_name, _short_name, args.is_20hz, args.upstream_branch)
|
||||
else:
|
||||
print("No models processed.", file=sys.stderr)
|
||||
_model_metadata = generate_chunked_model(_driving_pkl)
|
||||
create_metadata_json([_model_metadata], _output_dir, args.custom_name, _short_name, args.is_20hz, args.upstream_branch)
|
||||
|
||||
+1
-1
Submodule tinygrad_repo updated: ac1632ab96...2fecac4e4a
Reference in New Issue
Block a user