From 4cfd7748556cee51d008e7ce9dac4502fec632b3 Mon Sep 17 00:00:00 2001 From: Armand du Parc Locmaria Date: Thu, 14 May 2026 23:09:50 -0700 Subject: [PATCH] modeld/dmonitoringmodeld: explicitly set input devices (#38044) * modeld/dmonitoringmodeld: explicitly set input devices * lint * ignore metadata json file --- .gitignore | 2 +- selfdrive/modeld/SConscript | 26 ++++++++++++++++---------- selfdrive/modeld/compile_modeld.py | 12 ++++++------ selfdrive/modeld/dmonitoringmodeld.py | 7 +++---- selfdrive/modeld/helpers.py | 10 ++++------ selfdrive/modeld/modeld.py | 18 +++++++----------- 6 files changed, 37 insertions(+), 38 deletions(-) diff --git a/.gitignore b/.gitignore index 689d1d6c2..8f4442e52 100644 --- a/.gitignore +++ b/.gitignore @@ -44,7 +44,7 @@ bin/ config.json compile_commands.json compare_runtime*.html -selfdrive/modeld/models/tg_compiled_flags.json +selfdrive/modeld/models/tg_input_devices.json # build artifacts docs_site/ diff --git a/selfdrive/modeld/SConscript b/selfdrive/modeld/SConscript index 73dcb4a6f..18e806a6a 100644 --- a/selfdrive/modeld/SConscript +++ b/selfdrive/modeld/SConscript @@ -8,6 +8,7 @@ from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT from openpilot.selfdrive.modeld.constants import ModelConstants from tinygrad import Device from openpilot.system.hardware import HARDWARE, PC +from openpilot.selfdrive.modeld.helpers import TG_INPUT_DEVICES_PATH Import('env', 'arch', 'release') @@ -46,15 +47,20 @@ else: tg_backend = 'CPU' if arch == 'Darwin' else 'CPU:LLVM' tg_flags = f'DEV={tg_backend}' -def write_tg_compiled_flags(target, source, env): +tg_devices = { # which device to put jit inputs to at runtime + 'selfdrive.modeld.modeld': tg_backend, + 'selfdrive.modeld.dmonitoringmodeld': tg_backend, +} + +def write_tg_devices(target, source, env): with open(str(target[0]), "w") as f: - json.dump({"DEV": tg_backend}, f) + json.dump(tg_devices, f) f.write("\n") -compiled_flags_node = lenv.Command( - File("models/tg_compiled_flags.json").abspath, - tinygrad_files + [Value(tg_backend)], - write_tg_compiled_flags, +tg_devices_node = lenv.Command( + str(TG_INPUT_DEVICES_PATH), + [Value(json.dumps(tg_devices, sort_keys=True))], + write_tg_devices, ) # tinygrad calls brew which needs a $HOME in the env @@ -74,7 +80,7 @@ cmd = (f'{tg_flags} {mac_brew_string} python3 {modeld_dir}/compile_modeld.py ' f'--vision-onnx {File("models/driving_vision.onnx").abspath} ' f'--policy-onnx {File("models/driving_policy.onnx").abspath} ' f'--output {pkl_path} --frame-skip {frame_skip}') -node = lenv.Command(pkl_path, tinygrad_files + compile_modeld_script + driving_onnx_deps + [Value(camera_res_args), chunker_file, compiled_flags_node], cmd) +node = lenv.Command(pkl_path, tinygrad_files + compile_modeld_script + driving_onnx_deps + [Value(camera_res_args), chunker_file, tg_devices_node], cmd) onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps) chunk_targets = get_chunk_paths(pkl_path, estimate_pickle_max_size(onnx_sizes_sum)*2) # TODO make weight dedupe work on QCOM def do_chunk(target, source, env, pkl=pkl_path, chunks=chunk_targets): @@ -85,7 +91,7 @@ lenv.Command(chunk_targets, node, do_chunk) fn = File(f"models/dmonitoring_model").abspath script_files = [File(Dir("#selfdrive/modeld").File("get_model_metadata.py").abspath)] cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx' -lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files + script_files + [compiled_flags_node], cmd) +lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files + script_files + [tg_devices_node], cmd) dm_w, dm_h = DM_INPUT_SIZE compile_dm_warp_script = [File(f"{modeld_dir}/compile_dm_warp.py")] @@ -94,7 +100,7 @@ for cam_w, cam_h in CAMERA_CONFIGS: cmd = (f'{tg_flags} {mac_brew_string} python3 {modeld_dir}/compile_dm_warp.py ' f'--camera-resolution {cam_w}x{cam_h} --warp-to {dm_w}x{dm_h} ' f'--output {dm_pkl_path}') - lenv.Command(dm_pkl_path, tinygrad_files + compile_dm_warp_script + compile_modeld_script + [compiled_flags_node], cmd) + lenv.Command(dm_pkl_path, tinygrad_files + compile_dm_warp_script + compile_modeld_script + [tg_devices_node], cmd) driving_metadata_deps = [File(f"models/{m}_metadata.pkl").abspath for m in ['driving_vision', 'driving_policy']] def tg_compile(flags, model_name): @@ -105,7 +111,7 @@ def tg_compile(flags, model_name): chunk_targets = get_chunk_paths(pkl, estimate_pickle_max_size(os.path.getsize(onnx_path))) compile_node = lenv.Command( pkl, - [onnx_path] + tinygrad_files + [chunker_file, compiled_flags_node], + [onnx_path] + tinygrad_files + [chunker_file, tg_devices_node], f'{pythonpath_string} {flags} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {pkl}', ) def do_chunk(target, source, env): diff --git a/selfdrive/modeld/compile_modeld.py b/selfdrive/modeld/compile_modeld.py index f5d92edce..031f13be3 100755 --- a/selfdrive/modeld/compile_modeld.py +++ b/selfdrive/modeld/compile_modeld.py @@ -87,7 +87,7 @@ def make_frame_prepare(nv12: NV12Frame, model_w, model_h): return frame_prepare_tinygrad -def make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip): +def make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, device): img = vision_input_shapes['img'] # (1, 12, 128, 256) n_frames = img[1] // 6 img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img[2], img[3]) @@ -103,10 +103,10 @@ def make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip): 'big_tfm': np.zeros((3, 3), dtype=np.float32), } input_queues = { - 'img_q': Tensor.zeros(img_buf_shape, dtype='uint8').contiguous().realize(), - 'big_img_q': Tensor.zeros(img_buf_shape, dtype='uint8').contiguous().realize(), - 'feat_q': Tensor.zeros(frame_skip * (fb[1] - 1) + 1, fb[0], fb[2]).contiguous().realize(), - 'desire_q': Tensor.zeros(frame_skip * dp[1], dp[0], dp[2]).contiguous().realize(), + 'img_q': Tensor.zeros(img_buf_shape, dtype='uint8', device=device).contiguous().realize(), + 'big_img_q': Tensor.zeros(img_buf_shape, dtype='uint8', device=device).contiguous().realize(), + 'feat_q': Tensor.zeros(frame_skip * (fb[1] - 1) + 1, fb[0], fb[2], device=device).contiguous().realize(), + 'desire_q': Tensor.zeros(frame_skip * dp[1], dp[0], dp[2], device=device).contiguous().realize(), **{k: Tensor(v, device='NPY').realize() for k, v in npy.items()}, } return input_queues, npy @@ -172,7 +172,7 @@ def compile_modeld(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip, SEED = 42 def random_inputs_run_fn(fn, seed, test_val=None, test_buffers=None, expect_match=True): - input_queues, npy = make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip) + input_queues, npy = make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, Device.DEFAULT) np.random.seed(seed) Tensor.manual_seed(seed) diff --git a/selfdrive/modeld/dmonitoringmodeld.py b/selfdrive/modeld/dmonitoringmodeld.py index 618d39ca6..7b0a26f4a 100755 --- a/selfdrive/modeld/dmonitoringmodeld.py +++ b/selfdrive/modeld/dmonitoringmodeld.py @@ -1,8 +1,6 @@ #!/usr/bin/env python3 import os -from openpilot.selfdrive.modeld.helpers import MODELS_DIR, set_tinygrad_backend_from_compiled_flags -set_tinygrad_backend_from_compiled_flags() - +from openpilot.selfdrive.modeld.helpers import MODELS_DIR, get_tg_input_devices from tinygrad.tensor import Tensor import time import pickle @@ -29,6 +27,7 @@ class ModelState: output: np.ndarray def __init__(self, cam_w: int, cam_h: int): + self.DEV = get_tg_input_devices(PROCESS_NAME) with open(METADATA_PATH, 'rb') as f: model_metadata = pickle.load(f) self.input_shapes = model_metadata['input_shapes'] @@ -55,7 +54,7 @@ class ModelState: ptr = buf.data.ctypes.data # There is a ringbuffer of imgs, just cache tensors pointing to all of them if ptr not in self._blob_cache: - self._blob_cache[ptr] = Tensor.from_blob(ptr, (self.frame_buf_params[3],), dtype='uint8') + self._blob_cache[ptr] = Tensor.from_blob(ptr, (self.frame_buf_params[3],), dtype='uint8', device=self.DEV) self.warp_inputs_np['transform'][:] = transform[:] self.tensor_inputs['input_img'] = self.image_warp(self._blob_cache[ptr], self.warp_inputs['transform']) diff --git a/selfdrive/modeld/helpers.py b/selfdrive/modeld/helpers.py index 9c1a0d468..f4b396842 100644 --- a/selfdrive/modeld/helpers.py +++ b/selfdrive/modeld/helpers.py @@ -1,12 +1,10 @@ import json -import os from pathlib import Path MODELS_DIR = Path(__file__).resolve().parent / 'models' -COMPILED_FLAGS_PATH = MODELS_DIR / 'tg_compiled_flags.json' +TG_INPUT_DEVICES_PATH = MODELS_DIR / 'tg_input_devices.json' -def set_tinygrad_backend_from_compiled_flags() -> None: - if os.path.isfile(COMPILED_FLAGS_PATH): - with open(COMPILED_FLAGS_PATH) as f: - os.environ['DEV'] = str(json.load(f)['DEV']) +def get_tg_input_devices(process_name: str) -> dict[str, str]: + with open(TG_INPUT_DEVICES_PATH) as f: + return json.load(f)[process_name] diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py index 527cc5d7a..6aa94be49 100755 --- a/selfdrive/modeld/modeld.py +++ b/selfdrive/modeld/modeld.py @@ -1,12 +1,6 @@ #!/usr/bin/env python3 import os -from openpilot.selfdrive.modeld.helpers import MODELS_DIR, set_tinygrad_backend_from_compiled_flags -set_tinygrad_backend_from_compiled_flags() - -USBGPU = "USBGPU" in os.environ -if USBGPU: - os.environ['DEV'] = 'AMD' - os.environ['AMD_IFACE'] = 'USB' +from openpilot.selfdrive.modeld.helpers import MODELS_DIR, get_tg_input_devices from tinygrad.tensor import Tensor import time import pickle @@ -78,6 +72,7 @@ class ModelState: prev_desire: np.ndarray # for tracking the rising edge of the pulse def __init__(self, cam_w: int, cam_h: int): + self.DEV = get_tg_input_devices(PROCESS_NAME) jits = pickle.loads(read_file_chunked(MODELS_DIR / 'driving_tinygrad.pkl')) vision_metadata = jits['metadata']['vision'] self.vision_input_shapes = vision_metadata['input_shapes'] @@ -91,7 +86,7 @@ class ModelState: self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32) self.frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ - self.input_queues, self.npy = make_input_queues(self.vision_input_shapes, self.policy_input_shapes, self.frame_skip) + self.input_queues, self.npy = make_input_queues(self.vision_input_shapes, self.policy_input_shapes, self.frame_skip, device=self.DEV) self.full_frames : dict[str, Tensor] = {} self._blob_cache : dict[int, Tensor] = {} self.parser = Parser() @@ -100,8 +95,8 @@ class ModelState: self.warp_enqueue = jits[(cam_w,cam_h)]['warp_enqueue'] self.warp_enqueue( **self.input_queues, - frame=Tensor.zeros(self.frame_buf_params['img'][3], dtype='uint8').contiguous().realize(), - big_frame=Tensor.zeros(self.frame_buf_params['big_img'][3], dtype='uint8').contiguous().realize()) + frame=Tensor.zeros(self.frame_buf_params['img'][3], dtype='uint8', device=self.DEV).contiguous().realize(), + big_frame=Tensor.zeros(self.frame_buf_params['big_img'][3], dtype='uint8', device=self.DEV).contiguous().realize()) def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]: parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in output_slices.items()} @@ -115,7 +110,7 @@ class ModelState: # 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') + self._blob_cache[cache_key] = Tensor.from_blob(ptr, (yuv_size,), dtype='uint8', device=self.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 @@ -148,6 +143,7 @@ class ModelState: def main(demo=False): cloudlog.warning("modeld init") + USBGPU = False if not USBGPU: # USB GPU currently saturates a core so can't do this yet, # also need to move the aux USB interrupts for good timings