diff --git a/THIRD-PARTY-NOTICES b/THIRD-PARTY-NOTICES index bc8f3ff47..10ff5fb72 100644 --- a/THIRD-PARTY-NOTICES +++ b/THIRD-PARTY-NOTICES @@ -56,6 +56,7 @@ Optional integration: integrations/omnidreams mediapy Apache-2.0 https://github.com/google/mediapy nvidia-cudnn-frontend MIT https://github.com/NVIDIA/cudnn-frontend +nvtx Apache-2.0 WITH LLVM-exception https://github.com/NVIDIA/NVTX opencv-python-headless Apache-2.0 https://github.com/opencv/opencv-python grpcio, grpcio-tools Apache-2.0 https://github.com/grpc/grpc shapely BSD-3-Clause https://github.com/shapely/shapely diff --git a/integrations/omnidreams/ludus-renderer/ludus_renderer/_cpp/bindings/torch_bindings_cuda.cpp b/integrations/omnidreams/ludus-renderer/ludus_renderer/_cpp/bindings/torch_bindings_cuda.cpp index 53c84e9d0..91c3b961d 100644 --- a/integrations/omnidreams/ludus-renderer/ludus_renderer/_cpp/bindings/torch_bindings_cuda.cpp +++ b/integrations/omnidreams/ludus-renderer/ludus_renderer/_cpp/bindings/torch_bindings_cuda.cpp @@ -62,9 +62,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { pybind11::int_((int)CR::CudaRaster::RenderModeFlag_EnableDepthPeeling); // CUDA rendering ops - m.def("ludus_render_fwd_cuda", &ludus_render_fwd_cuda, "ludus f-theta CUDA rendering"); - m.def("ludus_render_fwd_cuda_ts", &ludus_render_fwd_cuda_ts, "ludus f-theta CUDA rendering with timestamped cube pools"); - m.def("ludus_render_fwd_cuda_timestamped", &ludus_render_fwd_cuda_timestamped, "ludus f-theta CUDA timestamped rendering from flat buffers"); + // Rendering can wait behind world-model work on the same GPU. Keep those + // waits from blocking Python's presentation thread. + m.def("ludus_render_fwd_cuda", &ludus_render_fwd_cuda, "ludus f-theta CUDA rendering", + pybind11::call_guard()); + m.def("ludus_render_fwd_cuda_ts", &ludus_render_fwd_cuda_ts, "ludus f-theta CUDA rendering with timestamped cube pools", + pybind11::call_guard()); + m.def("ludus_render_fwd_cuda_timestamped", &ludus_render_fwd_cuda_timestamped, "ludus f-theta CUDA timestamped rendering from flat buffers", + pybind11::call_guard()); } //------------------------------------------------------------------------ diff --git a/integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py b/integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py index 2c7118dee..4f4ba9c9e 100644 --- a/integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py +++ b/integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py @@ -455,6 +455,7 @@ def start_generation( finalization_state={"autoregressive_index": 0}, ) + @nvtx.annotate() def continue_generation( self, state: OmnidreamsConditioningState, diff --git a/integrations/omnidreams/omnidreams/interactive_drive/README.md b/integrations/omnidreams/omnidreams/interactive_drive/README.md index ff06081b8..923e90bc0 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/README.md +++ b/integrations/omnidreams/omnidreams/interactive_drive/README.md @@ -624,3 +624,44 @@ tests run with no checkpoint downloads. The hook does not auto-fix. Use `./scripts/check.sh --fix` to clean up lint and format issues explicitly. + +### To Profile + +```bash +sudo sh -c 'echo 0 > /proc/sys/kernel/perf_event_paranoid' + +# unsafe, adds unpriv access to GPU perf counters +cat > tmp.sh <<'EOF' +#!/usr/bin/env bash +set -euo pipefail +# Unsafe: allows unprivileged access to NVIDIA GPU perf/profiling counters. +# Run only on a trusted profiling machine. +for cap in profiler-device profiler-context trace-device; do + minor="$(awk -v c="$cap" '$1 == c { print $2 }' /proc/driver/nvidia-caps/sys-minors)" + if [[ -z "${minor}" ]]; then + echo "Could not find sys-minor for ${cap}" >&2 + exit 1 + fi + sudo nvidia-modprobe -f "/proc/driver/nvidia/capabilities/${cap}" + sudo chmod a+r "/dev/nvidia-caps/nvidia-cap${minor}" + echo "DeviceFileModify: 0" | sudo tee "/proc/driver/nvidia/capabilities/${cap}" >/dev/null + echo "Enabled ${cap} via /dev/nvidia-caps/nvidia-cap${minor}" +done +EOF +chmod +x tmp.sh +./tmp.sh +rm ./tmp.sh + +nsys profile --output=/tmp/flashdreams-nsys \ + --force-overwrite=true \ + --trace=cuda,nvtx,vulkan,python-gil \ + --sample=process-tree --backtrace=dwarf --samples-per-backtrace=1 \ + --python-sampling=true \ + --python-sampling-frequency=1000 \ + --python-backtrace=cuda \ + --cudabacktrace=all \ + --gpu-metrics-devices=all \ + --gpu-metrics-frequency=10000 \ + --stop-on-exit=true \ + uv run --package flashdreams-omnidreams interactive-drive --manifest example_world_model_perf.yaml --auto-start --stop-after-chunks 48 +``` \ No newline at end of file diff --git a/integrations/omnidreams/omnidreams/interactive_drive/app.py b/integrations/omnidreams/omnidreams/interactive_drive/app.py index 8ac0b0ff9..1db847fec 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/app.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/app.py @@ -16,6 +16,7 @@ KeyboardState, ) from omnidreams.interactive_drive.presenter import SlangPyPresenter +import nvtx from omnidreams.interactive_drive.runtime.loop import ( LoopConfig, PresenterBackend, @@ -167,6 +168,7 @@ def first_chunk_produced(self) -> bool: """``True`` once the model has produced its first generated chunk.""" return self._pipeline.first_chunk_produced.is_set() + @nvtx.annotate() def load_scene( self, scene_path: object, variant: str, prompt_override: str | None ) -> bool: @@ -303,6 +305,7 @@ def _preload_worker(self, pending: list[tuple[object, str, str | None]]) -> None f"{Path(str(scene_path)).name} variant={variant!r}", ) + @nvtx.annotate() def _resolve_scene_assets( self, scene_path: object, variant: str, prompt_override: str | None ) -> tuple[SceneBundle, MapBounds | None, GroundSnapper | None]: @@ -417,6 +420,7 @@ def _loading_base_frame(self) -> np.ndarray: self._loading_base_rgb = np.zeros((height, width, 3), dtype=np.uint8) return self._loading_base_rgb + @nvtx.annotate() def run_scene(self) -> None: """Drive the current scene until the presenter closes or switches. @@ -502,6 +506,7 @@ def run_scene(self) -> None: # rather than after the rebuild completes. self._present_loading_once(loading_status) + @nvtx.annotate() def _present_loading_once(self, loading_status: Callable[[], str]) -> None: """Render a single loading-overlay frame immediately (used on reset).""" if self._scene is None: @@ -535,6 +540,7 @@ def _resetting_status_message(self) -> str: """Phase text shown while a reset / respawn re-primes the rollout.""" return "Resetting..." + @nvtx.annotate() def run(self) -> None: """Single-scene convenience: load the configured scene, run, tear down. diff --git a/integrations/omnidreams/omnidreams/interactive_drive/backends/raster.py b/integrations/omnidreams/omnidreams/interactive_drive/backends/raster.py index e5d19c2b1..352729e52 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/backends/raster.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/backends/raster.py @@ -5,6 +5,7 @@ from omnidreams.interactive_drive.backends.base import RenderBackend from omnidreams.interactive_drive.config import BevConfig, ChunkConfig, RasterConfig +import nvtx from omnidreams.interactive_drive.rasterizer import LudusConditionRasterizer from omnidreams.interactive_drive.types import FrameChunk, SceneBundle, TrajectoryChunk @@ -24,16 +25,20 @@ def warmup_model(self) -> None: # No model to load; per-scene work happens in load_scene. return + @nvtx.annotate() def load_scene(self, scene: SceneBundle) -> None: self._scene = scene self._rasterizer.load_scene(scene) + @nvtx.annotate() def render_first_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk: return self._render_chunk(trajectory) + @nvtx.annotate() def render_next_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk: return self._render_chunk(trajectory) + @nvtx.annotate() def _render_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk: raster_chunk = self._rasterizer.render_chunk( rig_poses_world=trajectory.rig_poses_world, diff --git a/integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py b/integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py index d4fae0b58..d046a2018 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py @@ -17,6 +17,7 @@ RasterConfig, WorldModelProfileConfig, ) +import nvtx from omnidreams.interactive_drive.rasterizer import LudusConditionRasterizer from omnidreams.interactive_drive.types import ( FrameChunk, @@ -66,6 +67,7 @@ def optimizes_on_first_chunk(self) -> bool: # which can take minutes on the first launch. return True + @nvtx.annotate() def warmup_model(self) -> None: if self._manifest.resolution_wh != self._raster.resolution_wh: raise ValueError( @@ -91,6 +93,7 @@ def warmup_model(self) -> None: f"[world-model] model warmup session_ms={(time.perf_counter() - start) * 1000.0:.1f}", ) + @nvtx.annotate() def load_scene(self, scene: SceneBundle) -> None: self._scene = scene self._next_chunk_count = 0 @@ -115,6 +118,7 @@ def load_scene(self, scene: SceneBundle) -> None: f"total_ms={(prepare_end - load_start) * 1000.0:.1f}", ) + @nvtx.annotate() def render_first_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk: scene = self._require_scene() chunk_start = time.perf_counter() @@ -128,23 +132,24 @@ def render_first_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk: display_frames = raster_chunk.frames else: raster_end = time.perf_counter() - condition_frames = [ - frame.copy() for frame in self._debug_first_chunk_condition_frames - ] - display_frames = tuple( - PresentedFrame( - timestamp_us=int(timestamp_us), - rgb_host_uint8=frame.copy(), - depth_host_f32=None, - rgb_native=None, - depth_native=None, - ) - for timestamp_us, frame in zip( - trajectory.timestamps_us, - self._debug_first_chunk_condition_frames, - strict=True, + with nvtx.annotate("backend.world_model.debug_condition_frames", color="yellow"): + condition_frames = [ + frame.copy() for frame in self._debug_first_chunk_condition_frames + ] + display_frames = tuple( + PresentedFrame( + timestamp_us=int(timestamp_us), + rgb_host_uint8=frame.copy(), + depth_host_f32=None, + rgb_native=None, + depth_native=None, + ) + for timestamp_us, frame in zip( + trajectory.timestamps_us, + self._debug_first_chunk_condition_frames, + strict=True, + ) ) - ) logger.info( "[world-model] first_chunk using official hdmap override " f"dir={self._manifest.debug_condition_frame_dir}", @@ -182,6 +187,7 @@ def render_first_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk: ), ) + @nvtx.annotate() def render_next_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk: self._require_scene() chunk_start = time.perf_counter() @@ -225,10 +231,12 @@ def render_next_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk: ), ) + @nvtx.annotate() def reset(self) -> None: self._session.reset() self._next_chunk_count = 0 + @nvtx.annotate() def reset_scene_conditioning(self) -> None: self._session.reset(clear_precomputed_embeddings=True) self._next_chunk_count = 0 @@ -265,6 +273,7 @@ def _load_debug_condition_frames( frames.append(np.array(rgb, dtype=np.uint8)) return tuple(frames) + @nvtx.annotate() def _merge_frames( self, raster_frames: Sequence[PresentedFrame], diff --git a/integrations/omnidreams/omnidreams/interactive_drive/cuda_host_prefetch.py b/integrations/omnidreams/omnidreams/interactive_drive/cuda_host_prefetch.py index f994db374..d2e749d9f 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/cuda_host_prefetch.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/cuda_host_prefetch.py @@ -7,6 +7,7 @@ from typing import Any import numpy as np +import nvtx _STREAMS_LOCK = threading.Lock() _HOST_COPY_STREAMS: dict[int, Any] = {} @@ -22,6 +23,7 @@ def __init__(self, tensor: Any, *, source_event: Any | None = None) -> None: self._done_event: Any | None = None self._started = False + @nvtx.annotate() def start(self) -> bool: if self._started: return self._host_tensor is not None @@ -48,10 +50,11 @@ def start(self) -> bool: if self._source_event is not None: copy_stream.wait_event(self._source_event) with torch.cuda.stream(copy_stream): - host_tensor.copy_(tensor, non_blocking=True) - tensor.record_stream(copy_stream) - done_event = torch.cuda.Event() - done_event.record(copy_stream) + with nvtx.annotate("cuda_host_prefetch.copy_to_host", color="yellow"): + host_tensor.copy_(tensor, non_blocking=True) + tensor.record_stream(copy_stream) + done_event = torch.cuda.Event() + done_event.record(copy_stream) except Exception: self._host_tensor = None self._done_event = None @@ -61,6 +64,7 @@ def start(self) -> bool: self._done_event = done_event return True + @nvtx.annotate() def to_numpy(self) -> np.ndarray: host_tensor = self._host_tensor if host_tensor is None: diff --git a/integrations/omnidreams/omnidreams/interactive_drive/input/keyboard.py b/integrations/omnidreams/omnidreams/interactive_drive/input/keyboard.py index 18c610875..d2e825bdd 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/input/keyboard.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/input/keyboard.py @@ -5,6 +5,7 @@ import time from omnidreams.interactive_drive.input.backend import InputBackend, SampledInput +import nvtx from omnidreams.interactive_drive.types import ( ControlSnapshot, DriverCommand, @@ -149,6 +150,7 @@ class KeyboardInputBackend(InputBackend): def __init__(self, keyboard: KeyboardState) -> None: self._keyboard = keyboard + @nvtx.annotate() def sample(self) -> SampledInput: sample_time = time.perf_counter() return SampledInput(command=self._keyboard.command(), sample_time=sample_time) diff --git a/integrations/omnidreams/omnidreams/interactive_drive/presenter.py b/integrations/omnidreams/omnidreams/interactive_drive/presenter.py index 623dd756f..14c6cf1a7 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/presenter.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/presenter.py @@ -13,6 +13,7 @@ ) from omnidreams.interactive_drive.input.keyboard import KeyboardState from omnidreams.interactive_drive.loading_overlay import render_loading_overlay +import nvtx from omnidreams.interactive_drive.types import PresentedFrame @@ -72,9 +73,11 @@ def close(self) -> None: self._cuda_rgb_interop = None self._window.close() + @nvtx.annotate() def process_events(self) -> None: self._window.process_events() + @nvtx.annotate() def prepare_frame(self, frame: PresentedFrame, view_mode: str) -> None: if ( view_mode == "model_rgb" @@ -86,8 +89,13 @@ def prepare_frame(self, frame: PresentedFrame, view_mode: str) -> None: if view_mode != "model_rgb": _prefetch_to_numpy(frame.rgb_host_uint8) + @nvtx.annotate() def present_frame(self, frame: PresentedFrame, view_mode: str) -> None: - if view_mode == "model_rgb" and frame.model_rgb_host_uint8 is not None: + with nvtx.annotate("presenter.select_view_rgb", color="yellow"): + use_model_rgb = ( + view_mode == "model_rgb" and frame.model_rgb_host_uint8 is not None + ) + if use_model_rgb: if self._present_cuda_rgb( frame.model_rgb_host_uint8, status_message=frame.status_message, @@ -96,9 +104,8 @@ def present_frame(self, frame: PresentedFrame, view_mode: str) -> None: rgb = _with_status_overlay(frame.model_rgb_host_uint8, frame.status_message) self._present_array(rgb) return - self._present_array( - _with_status_overlay(frame.rgb_host_uint8, frame.status_message) - ) + rgb = _with_status_overlay(frame.rgb_host_uint8, frame.status_message) + self._present_array(rgb) def _create_device(self): existing_device_handles = self._cuda_existing_device_handles() @@ -181,6 +188,7 @@ def _create_cuda_rgb_interop(self): logger.info("[presenter] cuda_interop=enabled") return interop + @nvtx.annotate() def _present_cuda_rgb( self, rgb_frame: object, *, status_message: str | None ) -> bool: @@ -207,13 +215,14 @@ def _present_cuda_rgb( submitted = self._submit_ready_cuda_rgb() return True + @nvtx.annotate() def _submit_ready_cuda_rgb(self) -> bool: if self._cuda_rgb_interop is None: return False interop_frame = self._cuda_rgb_interop.ready_rgba_buffer() if interop_frame is None: return False - rgba_buffer, cuda_stream = interop_frame + rgba_buffer, _cuda_stream = interop_frame if not self._surface.config: return False surface_texture = self._surface.acquire_next_image() @@ -221,28 +230,32 @@ def _submit_ready_cuda_rgb(self) -> bool: time.sleep(0.001) return False - command_encoder = self._device.create_command_encoder() - command_encoder.copy_buffer_to_texture( - self._display_texture, - 0, - 0, - [0, 0, 0], - rgba_buffer.buffer, - 0, - rgba_buffer.size_bytes, - rgba_buffer.row_pitch, - [self._raster.width, self._raster.height, 1], - ) - command_encoder.blit(surface_texture, self._display_texture) - submit_id = self._device.submit_command_buffer( - command_encoder.finish(), - cuda_stream=cuda_stream, - ) + with nvtx.annotate("presenter.submit_ready_cuda_rgb.encode_commands", color="green"): + command_encoder = self._device.create_command_encoder() + command_encoder.copy_buffer_to_texture( + self._display_texture, + 0, + 0, + [0, 0, 0], + rgba_buffer.buffer, + 0, + rgba_buffer.size_bytes, + rgba_buffer.row_pitch, + [self._raster.width, self._raster.height, 1], + ) + command_encoder.blit(surface_texture, self._display_texture) + with nvtx.annotate("presenter.submit_ready_cuda_rgb.submit", color="green"): + # ``ready_rgba_buffer`` only returns after its CUDA completion + # event succeeds. Passing the producer stream here would make + # Slang wait on newer work appended after that ready buffer. + submit_id = self._device.submit_command_buffer(command_encoder.finish()) self._cuda_rgb_interop.mark_submitted(rgba_buffer, submit_id) + with nvtx.annotate("presenter.submit_ready_cuda_rgb.surface_present", color="green"): + self._surface.present() del surface_texture - self._surface.present() return True + @nvtx.annotate() def _present_array(self, rgb_host_uint8: np.ndarray) -> None: if not self._surface.config: return @@ -252,13 +265,16 @@ def _present_array(self, rgb_host_uint8: np.ndarray) -> None: return upload = self._pack_surface_pixels(rgb_host_uint8) - self._display_texture.copy_from_numpy(upload) - - command_encoder = self._device.create_command_encoder() - command_encoder.blit(surface_texture, self._display_texture) - self._device.submit_command_buffer(command_encoder.finish()) + with nvtx.annotate("presenter.present_array.upload_texture", color="green"): + self._display_texture.copy_from_numpy(upload) + + with nvtx.annotate("presenter.present_array.submit", color="green"): + command_encoder = self._device.create_command_encoder() + command_encoder.blit(surface_texture, self._display_texture) + self._device.submit_command_buffer(command_encoder.finish()) + with nvtx.annotate("presenter.present_array.surface_present", color="green"): + self._surface.present() del surface_texture - self._surface.present() def _choose_surface_format(self): linear_pairs = { @@ -285,6 +301,7 @@ def _choose_surface_format(self): f"Presenter requires a linear swapchain, but the surface only supports: {supported}" ) + @nvtx.annotate() def _pack_surface_pixels(self, rgb_host_uint8: np.ndarray) -> np.ndarray: upload = np.zeros((self._raster.height, self._raster.width, 4), dtype=np.uint8) upload[..., :3] = rgb_host_uint8 @@ -432,6 +449,7 @@ def as_cuda_rgb_source(self, rgb_frame: object) -> "_CudaRGBFrame | None": ready=_cuda_event_ready(source_event), ) + @nvtx.annotate() def enqueue_rgb_to_shared_rgba(self, rgb_frame: "_CudaRGBFrame") -> bool: shared_buffer = self._acquire_buffer() if shared_buffer is None: @@ -454,6 +472,7 @@ def enqueue_rgb_to_shared_rgba(self, rgb_frame: "_CudaRGBFrame") -> bool: shared_buffer.copy_done_event = copy_done_event return True + @nvtx.annotate() def enqueue_camera_to_shared_rgba( self, rgb_frame: "_CudaRGBFrame", @@ -712,6 +731,7 @@ def _prefetch_to_numpy(frame: object) -> None: prefetch() +@nvtx.annotate() def _with_status_overlay(rgb_host_uint8: object, message: str | None) -> np.ndarray: rgb_host_uint8 = _as_rgb_host_uint8(rgb_host_uint8) if message is None: diff --git a/integrations/omnidreams/omnidreams/interactive_drive/rasterizer.py b/integrations/omnidreams/omnidreams/interactive_drive/rasterizer.py index d5363910b..0cd68a293 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/rasterizer.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/rasterizer.py @@ -5,8 +5,8 @@ When :class:`BevConfig` is enabled it also renders a top-down BEV via a synthetic ``FThetaCamera`` above the rig (pinhole projection + a fixed -straight-down sensor-to-rig matrix); the BEV rides alongside the main RGB on -each :class:`PresentedFrame`. +straight-down sensor-to-rig matrix). The public facade pipelines BEV one chunk +behind RGB so presentation never waits for a current-chunk BEV. """ import concurrent.futures @@ -32,6 +32,7 @@ from omnidreams.interactive_drive.config import BevConfig, RasterConfig from omnidreams.interactive_drive.cuda_env import DISABLE_CUDA_INTEROP_ENV, env_truthy from omnidreams.interactive_drive.cuda_host_prefetch import CudaHostPrefetch +import nvtx from omnidreams.interactive_drive.types import PresentedFrame, RasterChunk, SceneBundle from torch import Tensor @@ -68,6 +69,7 @@ def __init__( self._host: np.ndarray | None = None self._prefetch: CudaHostPrefetch | None = None + @nvtx.annotate() def prefetch_to_numpy(self) -> None: if ( self._host is not None @@ -80,6 +82,7 @@ def prefetch_to_numpy(self) -> None: if prefetch.start(): self._prefetch = prefetch + @nvtx.annotate() def to_numpy(self) -> np.ndarray: if self._host is None: if self._prefetch is not None: @@ -99,6 +102,7 @@ def to_numpy(self) -> np.ndarray: self._frames_hwc_uint8 = None return self._host + @nvtx.annotate() def to_cuda_tensor(self) -> Tensor: if self._frames_hwc_uint8 is None: raise RuntimeError( @@ -193,6 +197,7 @@ def _to_ludus_camera_pose(self, camera_poses: Tensor) -> Tensor: """Convert sensor-to-world camera poses to Ludus' world-to-sensor format.""" return torch.linalg.inv(camera_poses) + @nvtx.annotate() def load_scene(self, scene: SceneBundle) -> None: """Load a scene from the USDZ bundle. @@ -201,13 +206,14 @@ def load_scene(self, scene: SceneBundle) -> None: """ self.ctx.clear_scenes() - clipgt_scene = load_ludus_scene( - scene.scene_path, - device=self._device, - target_resolution=(self._raster.width, self._raster.height), - include_ego_trajectory=False, - include_ego_obstacle=False, - ) + with nvtx.annotate("rasterizer.load_ludus_scene", color="purple"): + clipgt_scene = load_ludus_scene( + scene.scene_path, + device=self._device, + target_resolution=(self._raster.width, self._raster.height), + include_ego_trajectory=False, + include_ego_obstacle=False, + ) scene_adapter = SceneAdapter(clipgt_scene) self._scene_data = _LoadedSceneData( @@ -238,11 +244,14 @@ def load_scene(self, scene: SceneBundle) -> None: self._bev_camera_id = None self._bev_sensor_to_rig = None - self.ctx.upload_cameras(self._all_cameras) + with nvtx.annotate("rasterizer.upload_cameras", color="purple"): + self.ctx.upload_cameras(self._all_cameras) # Single scene upload shared by the main camera and the BEV minimap. - self._scene_id = self.ctx.upload_scene(clipgt_scene.timestamped_scene) + with nvtx.annotate("rasterizer.upload_scene", color="purple"): + self._scene_id = self.ctx.upload_scene(clipgt_scene.timestamped_scene) + @nvtx.annotate() def render_chunk( self, rig_poses_world: npt.NDArray[np.float32], @@ -282,24 +291,90 @@ def render_chunk( np.ascontiguousarray(timestamps_us, dtype=np.int64) ).to(device=self._device) - rgb_frames = self._render_one_camera( - rig_poses=rig_poses_torch, + with nvtx.annotate("render_one_camera RGB", color="orange"): + rgb_frames = self._render_one_camera( + rig_poses=rig_poses_torch, + timestamps_batch=timestamps_batch, + scene_id=self._scene_id, + camera_id=self._all_camera_map[camera_name], + sensor_to_rig=self._sensor_to_rig[camera_name], + camera_type=CAMERA_TYPE_REGULAR, + resolution=(self._raster.height, self._raster.width), + ) + + bev_frames = self.render_bev_frames( + rig_poses_torch=rig_poses_torch, timestamps_batch=timestamps_batch, - scene_id=self._scene_id, - camera_id=self._all_camera_map[camera_name], - sensor_to_rig=self._sensor_to_rig[camera_name], - camera_type=CAMERA_TYPE_REGULAR, - resolution=(self._raster.height, self._raster.width), + ) + return self.build_chunk( + timestamps_us=timestamps_us, + rgb_frames=rgb_frames, + bev_frames=bev_frames, ) - bev_frames: _RenderedCameraFrames | None = None + @nvtx.annotate() + def render_rgb_frames( + self, + rig_poses_world: npt.NDArray[np.float32], + timestamps_us: npt.NDArray[np.int64], + ) -> tuple[npt.NDArray[np.int64], Tensor, Tensor, _RenderedCameraFrames]: + """Render only the main camera frames needed for model conditioning.""" if ( - self._bev is not None + self._scene_data is None + or self._scene_id is None + or self._selected_camera_name is None + ): + raise RuntimeError("load_scene() must be called before render_chunk().") + + camera_name = self._selected_camera_name + if camera_name not in self._all_camera_map: + available = sorted(self._all_camera_map.keys()) + raise RuntimeError( + f"Camera {camera_name!r} not found. Available: {available}" + ) + + rig_poses_torch = torch.from_numpy( + np.ascontiguousarray(rig_poses_world, dtype=np.float32) + ).to(device=self._device) + timestamps_batch = torch.from_numpy( + np.ascontiguousarray(timestamps_us, dtype=np.int64) + ).to(device=self._device) + + with nvtx.annotate("render_one_camera RGB", color="orange"): + rgb_frames = self._render_one_camera( + rig_poses=rig_poses_torch, + timestamps_batch=timestamps_batch, + scene_id=self._scene_id, + camera_id=self._all_camera_map[camera_name], + sensor_to_rig=self._sensor_to_rig[camera_name], + camera_type=CAMERA_TYPE_REGULAR, + resolution=(self._raster.height, self._raster.width), + ) + return ( + np.asarray(timestamps_us, dtype=np.int64), + rig_poses_torch, + timestamps_batch, + rgb_frames, + ) + + @nvtx.annotate() + def render_bev_frames( + self, + *, + rig_poses_torch: Tensor, + timestamps_batch: Tensor, + ) -> _RenderedCameraFrames | None: + """Render BEV frames for an already-prepared pose batch.""" + if not ( + self._scene_id is not None + and self._bev is not None and self._bev.enabled and self._bev_camera_id is not None and self._bev_sensor_to_rig is not None ): - bev_frames = self._render_one_camera( + return None + with nvtx.annotate("render_one_camera BEV", color="orange"): + return self._render_one_camera( rig_poses=rig_poses_torch, timestamps_batch=timestamps_batch, scene_id=self._scene_id, @@ -309,47 +384,74 @@ def render_chunk( resolution=(self._bev.height, self._bev.width), ) + @nvtx.annotate() + def build_chunk( + self, + *, + timestamps_us: npt.NDArray[np.int64], + rgb_frames: _RenderedCameraFrames, + bev_frames: _RenderedCameraFrames | None, + ) -> RasterChunk: + """Wrap rendered camera tensors in lazy frame objects.""" + if ( + bev_frames is not None + and int(bev_frames.frames_hwc_uint8.shape[0]) == 0 + ): + bev_frames = None + bev_frame_indices = _resampled_frame_indices( + source_count=( + int(bev_frames.frames_hwc_uint8.shape[0]) + if bev_frames is not None + else 0 + ), + target_count=len(timestamps_us), + ) if self._use_cuda_frames: + with nvtx.annotate("rasterizer.build_lazy_presented_frames", color="yellow"): + frames = [ + PresentedFrame( + timestamp_us=int(timestamps_us[idx]), + rgb_host_uint8=_LazyRasterFrame( + rgb_frames.frames_hwc_uint8, + idx, + source_event=rgb_frames.ready_event, + ), + depth_host_f32=None, + bev_host_uint8=( + _LazyRasterFrame( + bev_frames.frames_hwc_uint8, + bev_frame_indices[idx], + source_event=bev_frames.ready_event, + ) + if bev_frames is not None + else None + ), + ) + for idx in range(len(timestamps_us)) + ] + return RasterChunk(frames=tuple(frames)) + + rgb_host_frames = _rendered_frames_to_numpy(rgb_frames) + bev_host_frames = ( + _rendered_frames_to_numpy(bev_frames) if bev_frames is not None else None + ) + with nvtx.annotate("rasterizer.build_host_presented_frames", color="yellow"): frames = [ PresentedFrame( timestamp_us=int(timestamps_us[idx]), - rgb_host_uint8=_LazyRasterFrame( - rgb_frames.frames_hwc_uint8, - idx, - source_event=rgb_frames.ready_event, - ), + rgb_host_uint8=rgb_host_frames[idx], depth_host_f32=None, bev_host_uint8=( - _LazyRasterFrame( - bev_frames.frames_hwc_uint8, - idx, - source_event=bev_frames.ready_event, - ) - if bev_frames is not None + bev_host_frames[bev_frame_indices[idx]] + if bev_host_frames is not None else None ), ) for idx in range(len(timestamps_us)) ] - return RasterChunk(frames=tuple(frames)) - - rgb_host_frames = _rendered_frames_to_numpy(rgb_frames) - bev_host_frames = ( - _rendered_frames_to_numpy(bev_frames) if bev_frames is not None else None - ) - frames = [ - PresentedFrame( - timestamp_us=int(timestamps_us[idx]), - rgb_host_uint8=rgb_host_frames[idx], - depth_host_f32=None, - bev_host_uint8=( - bev_host_frames[idx] if bev_host_frames is not None else None - ), - ) - for idx in range(len(timestamps_us)) - ] return RasterChunk(frames=tuple(frames)) + @nvtx.annotate() def _render_one_camera( self, *, @@ -366,41 +468,44 @@ def _render_one_camera( Frames stay CUDA-backed so the world model consumes HDMap conditioning without a GPU->CPU->GPU round trip (presenters materialize NumPy lazily). """ - n_frames = timestamps_batch.shape[0] - camera_poses_world = torch.einsum( - "nij,jk->nik", rig_poses, sensor_to_rig.to(self._device) - ) - camera_poses_ludus = self._to_ludus_camera_pose(camera_poses_world) - scene_id_batch = torch.full( - (n_frames,), scene_id, dtype=torch.int32, device=self._device - ) - camera_id_batch = torch.full( - (n_frames,), camera_id, dtype=torch.int32, device=self._device - ) - camera_type_id_batch = torch.full( - (n_frames,), camera_type, dtype=torch.int32, device=self._device - ) + with nvtx.annotate("rasterizer.one_camera.prepare_inputs", color="orange"): + n_frames = timestamps_batch.shape[0] + camera_poses_world = torch.einsum( + "nij,jk->nik", rig_poses, sensor_to_rig.to(self._device) + ) + camera_poses_ludus = self._to_ludus_camera_pose(camera_poses_world) + scene_id_batch = torch.full( + (n_frames,), scene_id, dtype=torch.int32, device=self._device + ) + camera_id_batch = torch.full( + (n_frames,), camera_id, dtype=torch.int32, device=self._device + ) + camera_type_id_batch = torch.full( + (n_frames,), camera_type, dtype=torch.int32, device=self._device + ) height, width = resolution - images = self.ctx.render( - scene_id_batch, - camera_id_batch, - timestamps_batch, - camera_type_id_batch, - camera_poses_ludus, - resolution=(height, width), - ) + with nvtx.annotate("rasterizer.render", color="orange"): + images = self.ctx.render( + scene_id_batch, + camera_id_batch, + timestamps_batch, + camera_type_id_batch, + camera_poses_ludus, + resolution=(height, width), + ) - rgb = images[:, :, :, :3] - if self.ctx.needs_vflip: - rgb = rgb.flip(1) - if rgb.dtype != torch.uint8: - rgb = (rgb.clamp(0.0, 1.0) * 255.0 + 0.5).to(torch.uint8) - rgb = rgb.detach().contiguous() - ready_event = None - if rgb.is_cuda: - ready_event = torch.cuda.Event() - ready_event.record(torch.cuda.current_stream(rgb.device)) + with nvtx.annotate("rasterizer.one_camera.postprocess_rgb", color="yellow"): + rgb = images[:, :, :, :3] + if self.ctx.needs_vflip: + rgb = rgb.flip(1) + if rgb.dtype != torch.uint8: + rgb = (rgb.clamp(0.0, 1.0) * 255.0 + 0.5).to(torch.uint8) + rgb = rgb.detach().contiguous() + ready_event = None + if rgb.is_cuda: + ready_event = torch.cuda.Event() + ready_event.record(torch.cuda.current_stream(rgb.device)) return _RenderedCameraFrames(frames_hwc_uint8=rgb, ready_event=ready_event) def cleanup(self) -> None: @@ -433,18 +538,58 @@ def __init__(self, raster: RasterConfig, bev: BevConfig | None = None) -> None: self._impl: _LudusConditionRasterizerImpl | None = self._exec.submit( _LudusConditionRasterizerImpl, raster, bev ).result() + self._bev_enabled = bool(bev is not None and bev.enabled) + self._pending_bev: ( + concurrent.futures.Future[_RenderedCameraFrames | None] | None + ) = None + self._latest_bev: _RenderedCameraFrames | None = None + @nvtx.annotate() def load_scene(self, scene: SceneBundle) -> None: exec_, impl = self._require_alive() + self._clear_pending_bev() + self._latest_bev = None return exec_.submit(impl.load_scene, scene).result() + @nvtx.annotate() def render_chunk( self, rig_poses_world: npt.NDArray[np.float32], timestamps_us: npt.NDArray[np.int64], ) -> "RasterChunk": exec_, impl = self._require_alive() - return exec_.submit(impl.render_chunk, rig_poses_world, timestamps_us).result() + if not self._bev_enabled: + return exec_.submit( + impl.render_chunk, rig_poses_world, timestamps_us + ).result() + + with nvtx.annotate("rasterizer.poll_lagged_bev", color="yellow"): + lagged_bev = self._poll_ready_bev() + + ( + chunk_timestamps_us, + rig_poses_torch, + timestamps_batch, + rgb_frames, + ) = exec_.submit(impl.render_rgb_frames, rig_poses_world, timestamps_us).result() + + # RGB rendering gives the previous BEV another chance to finish without + # ever making it part of the critical path. If it is still in flight, + # reuse the latest completed BEV and skip this refresh so work cannot + # queue up behind presentation. + with nvtx.annotate("rasterizer.poll_lagged_bev_after_rgb", color="yellow"): + lagged_bev = self._poll_ready_bev() + if self._pending_bev is None: + self._pending_bev = exec_.submit( + impl.render_bev_frames, + rig_poses_torch=rig_poses_torch, + timestamps_batch=timestamps_batch, + ) + return impl.build_chunk( + timestamps_us=chunk_timestamps_us, + rgb_frames=rgb_frames, + bev_frames=lagged_bev, + ) def _require_alive( self, @@ -458,6 +603,7 @@ def cleanup(self) -> None: exec_ = getattr(self, "_exec", None) if exec_ is None: return + self._clear_pending_bev() impl = self._impl self._impl = None if impl is not None: @@ -470,7 +616,24 @@ def __del__(self) -> None: with contextlib.suppress(Exception): self.cleanup() + def _clear_pending_bev(self) -> None: + pending = getattr(self, "_pending_bev", None) + if pending is None: + return + pending.cancel() + with contextlib.suppress(Exception): + pending.result(timeout=0) + self._pending_bev = None + + def _poll_ready_bev(self) -> _RenderedCameraFrames | None: + pending = self._pending_bev + if pending is not None and pending.done(): + self._latest_bev = pending.result() + self._pending_bev = None + return self._latest_bev + +@nvtx.annotate() def _rendered_frames_to_numpy(rendered: _RenderedCameraFrames) -> list[np.ndarray]: synchronize = getattr(rendered.ready_event, "synchronize", None) if callable(synchronize): @@ -480,6 +643,16 @@ def _rendered_frames_to_numpy(rendered: _RenderedCameraFrames) -> list[np.ndarra return [frames[idx] for idx in range(frames.shape[0])] +def _resampled_frame_indices(*, source_count: int, target_count: int) -> list[int]: + """Map a lagged chunk across the current chunk without assuming equal sizes.""" + if source_count <= 0 or target_count <= 0: + return [] + if source_count == 1 or target_count == 1: + return [0] * target_count + scale = (source_count - 1) / (target_count - 1) + return [round(index * scale) for index in range(target_count)] + + def _build_bev_camera(bev: BevConfig, device: torch.device) -> FThetaCamera: """Construct a synthetic pinhole-as-FTheta camera for BEV rendering. diff --git a/integrations/omnidreams/omnidreams/interactive_drive/runtime/loop.py b/integrations/omnidreams/omnidreams/interactive_drive/runtime/loop.py index 581b07709..f0d396171 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/runtime/loop.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/runtime/loop.py @@ -11,6 +11,7 @@ from loguru import logger from omnidreams.interactive_drive.input.backend import InputBackend +import nvtx from omnidreams.interactive_drive.runtime.runtime_controls import RuntimeControls from omnidreams.interactive_drive.runtime.timing import ( ChunkHistory, @@ -38,6 +39,8 @@ _PROFILE_E2E_SUM_ADJ_MS: float = 0.0 _PROFILE_E2E_COUNT: int = 0 _PROFILE_E2E_WINDOW_START: float | None = None +_PROFILE_LAST_PRESENT_TIME: float | None = None +_PROFILE_PRESENT_INTERVALS_MS: list[float] = [] def _profile_input_to_present_enabled() -> bool: @@ -60,11 +63,25 @@ def reset_input_to_present_profile_window() -> None: global _PROFILE_E2E_SUM_ADJ_MS global _PROFILE_E2E_COUNT global _PROFILE_E2E_WINDOW_START + global _PROFILE_LAST_PRESENT_TIME _PROFILE_E2E_SUM_RAW_MS = 0.0 _PROFILE_E2E_SUM_ADJ_MS = 0.0 _PROFILE_E2E_COUNT = 0 _PROFILE_E2E_WINDOW_START = None + _PROFILE_LAST_PRESENT_TIME = None + _PROFILE_PRESENT_INTERVALS_MS.clear() + + +def _profile_percentile(values: list[float], quantile: float) -> float: + if not values: + return 0.0 + ordered = sorted(values) + position = (len(ordered) - 1) * quantile + lower = int(position) + upper = min(lower + 1, len(ordered) - 1) + fraction = position - lower + return ordered[lower] + (ordered[upper] - ordered[lower]) * fraction def _chunk_frame_interval_s(chunk_times: ChunkTimes) -> float: @@ -87,7 +104,13 @@ def _record_input_to_present_for_profile( global _PROFILE_E2E_SUM_ADJ_MS global _PROFILE_E2E_COUNT global _PROFILE_E2E_WINDOW_START + global _PROFILE_LAST_PRESENT_TIME + if _PROFILE_LAST_PRESENT_TIME is not None: + _PROFILE_PRESENT_INTERVALS_MS.append( + (present_time - _PROFILE_LAST_PRESENT_TIME) * 1000.0 + ) + _PROFILE_LAST_PRESENT_TIME = present_time raw_ms = (present_time - input_sample_time) * 1000.0 scheduled_ms = frame_index * (frame_interval_s * 1000.0) adj_ms = raw_ms - scheduled_ms @@ -108,9 +131,23 @@ def _record_input_to_present_for_profile( wall_present_fps = float(count) / window_s if window_s > 1e-9 else 0.0 avg_raw_ms = _PROFILE_E2E_SUM_RAW_MS / float(count) avg_adj_ms = _PROFILE_E2E_SUM_ADJ_MS / float(count) + present_interval_p50_ms = _profile_percentile( + _PROFILE_PRESENT_INTERVALS_MS, 0.50 + ) + present_interval_p95_ms = _profile_percentile( + _PROFILE_PRESENT_INTERVALS_MS, 0.95 + ) + present_interval_p99_ms = _profile_percentile( + _PROFILE_PRESENT_INTERVALS_MS, 0.99 + ) + present_interval_max_ms = max(_PROFILE_PRESENT_INTERVALS_MS, default=0.0) logger.info( "[profile] e2e " f"wall_present_fps={wall_present_fps:.1f} " + f"present_interval_p50_ms={present_interval_p50_ms:.2f} " + f"present_interval_p95_ms={present_interval_p95_ms:.2f} " + f"present_interval_p99_ms={present_interval_p99_ms:.2f} " + f"present_interval_max_ms={present_interval_max_ms:.2f} " f"avg_adj_control_to_present_ms={avg_adj_ms:.2f} " f"avg_raw_control_to_present_ms={avg_raw_ms:.2f} " f"samples={count}", @@ -119,6 +156,7 @@ def _record_input_to_present_for_profile( _PROFILE_E2E_SUM_ADJ_MS = 0.0 _PROFILE_E2E_COUNT = 0 _PROFILE_E2E_WINDOW_START = present_time + _PROFILE_PRESENT_INTERVALS_MS.clear() class PresenterBackend(Protocol): @@ -194,6 +232,7 @@ def should_request_chunk(state: MainLoopState) -> bool: return state.chunks_outstanding < 1 +@nvtx.annotate() def make_chunk_request( state: MainLoopState, simulation: SimulationBackend, @@ -254,6 +293,7 @@ def make_chunk_request( ) +@nvtx.annotate() def present_queued_frame( queued_frame: QueuedFrame, presenter: PresenterBackend, @@ -272,6 +312,7 @@ def present_queued_frame( re-presents that intersperse the warmup window with an OOB warning. """ frame_times = queued_frame.chunk_times.frames[queued_frame.frame_index] + nvtx.mark("loop.present_queued_frame") frame_times.sample_display_pose_time = time.perf_counter() display_frame = _frame_with_overlay(queued_frame.frame, oob_message) present_call_begin_time = time.perf_counter() @@ -431,6 +472,7 @@ def _prepare_queued_frame( prepare_frame(queued_frame.frame, view_mode=view_mode) +@nvtx.annotate() def _drain_pipeline_frames( *, pipeline: ChunkPipeline, @@ -452,6 +494,7 @@ def _drain_pipeline_frames( ready_frames.append(queued_frame) +@nvtx.annotate() def run_main_loop( presenter: PresenterBackend, runtime_controls: RuntimeControls, @@ -538,9 +581,13 @@ def run_main_loop( now = time.perf_counter() if now < state.next_present_time: wait_begin = now - time.sleep( - min(config.poll_timeout_s, max(0.0, state.next_present_time - now)) - ) + with nvtx.annotate("loop.present_wait", color="gray"): + time.sleep( + min( + config.poll_timeout_s, + max(0.0, state.next_present_time - now), + ) + ) wait_end = time.perf_counter() last_present_wait_event = _trace_main_range( active_trace, diff --git a/integrations/omnidreams/omnidreams/interactive_drive/slangpy_hud_presenter.py b/integrations/omnidreams/omnidreams/interactive_drive/slangpy_hud_presenter.py index 89f720ad0..eb9ff2c35 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/slangpy_hud_presenter.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/slangpy_hud_presenter.py @@ -12,6 +12,7 @@ from __future__ import annotations +import concurrent.futures import contextlib import math as _math import time @@ -28,6 +29,7 @@ _CudaRGBInterop, _env_truthy, ) +import nvtx from omnidreams.interactive_drive.types import DriverCommand, PresentedFrame from PIL import Image, ImageDraw, ImageFont @@ -75,6 +77,8 @@ # :class:`SlangPyHudPresenter`. DRIVE_KEY_RELEASE_DEBOUNCE_S = 0.08 +_BevPanelKey = tuple[int, int, int, int] + def _allocate_canvas(width: int, height: int) -> tuple[np.ndarray, Image.Image]: """Allocate the chrome buffer and a PIL Image view sharing its memory. @@ -92,6 +96,30 @@ def _allocate_canvas(width: int, height: int) -> tuple[np.ndarray, Image.Image]: return buf, img +@nvtx.annotate() +def _build_bev_panel_image( + key: _BevPanelKey, + bev_source: object, + target_size: tuple[int, int], +) -> tuple[_BevPanelKey, Image.Image]: + """Materialize, resize, and recolor BEV away from presentation.""" + from omnidreams.interactive_drive.demo import _apply_googlemaps_filter + + bev_rgb = _as_rgb_host_uint8(bev_source) + bev = Image.fromarray(bev_rgb, mode="RGB") + target_w, target_h = target_size + scale = max(target_w / bev.width, target_h / bev.height) + scaled_w = max(1, int(bev.width * scale)) + scaled_h = max(1, int(bev.height * scale)) + scaled = bev.resize((scaled_w, scaled_h), Image.Resampling.BILINEAR) + crop_left = (scaled_w - target_w) // 2 + crop_top = (scaled_h - target_h) // 2 + cropped = scaled.crop( + (crop_left, crop_top, crop_left + target_w, crop_top + target_h) + ) + return key, _apply_googlemaps_filter(cropped) + + class _LRUCache(OrderedDict): """Tiny ordered-dict-backed LRU. @@ -331,11 +359,20 @@ def __init__( self._pedal_cache: _LRUCache = _LRUCache(maxsize=16) self._scene_thumb_cache: dict[Any, Image.Image | None] = {} self._variant_thumb_cache: dict[tuple[Any, str], Image.Image | None] = {} - self._bev_panel_cache_key: tuple[int, int, int] | None = None + self._bev_panel_cache_key: _BevPanelKey | None = None self._bev_panel_cache: Image.Image | None = None + self._bev_panel_epoch = 0 + self._bev_panel_exec = concurrent.futures.ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="interactive-drive-bev-panel", + ) + self._bev_panel_future: concurrent.futures.Future[ + tuple[_BevPanelKey, Image.Image] + ] | None = None self._latest_camera_pil: Image.Image | None = None - self._latest_bev_pil: Image.Image | None = None + self._latest_bev_source: object | None = None + self._bev_source_generation = 0 # Numpy view of the latest world-model frame (RGBA8 with alpha # padded to 255) used by the GPU camera path. Lazily filled on # demand from ``_latest_camera_pil`` so we don't pay for the @@ -443,6 +480,7 @@ def __init__( def should_close(self) -> bool: return self._should_close_flag or self._window.should_close() + @nvtx.annotate() def process_events(self) -> None: self._window.process_events() # A wheel/controller's bound exit button posts its request onto the @@ -453,6 +491,7 @@ def process_events(self) -> None: if self._keyboard.consume_exit_scene_request(): self.exit_scene() + @nvtx.annotate() def prepare_frame(self, frame: PresentedFrame, view_mode: str) -> None: rgb = self._select_view_rgb(frame, view_mode) if self._cuda_hud_interop is None or not _has_cuda_tensor(rgb): @@ -460,6 +499,7 @@ def prepare_frame(self, frame: PresentedFrame, view_mode: str) -> None: if frame.bev_host_uint8 is not None: _prefetch_to_numpy(frame.bev_host_uint8) + @nvtx.annotate() def present_frame(self, frame: PresentedFrame, view_mode: str) -> None: # Apply any pending resize before touching the display texture # this frame. Done here (not inside on_resize) so Vulkan @@ -490,6 +530,7 @@ def present_frame(self, frame: PresentedFrame, view_mode: str) -> None: self._render_canvas(frame.status_message) self._present_canvas(use_gpu_camera=frame.status_message is None) + @nvtx.annotate() def present_world_model_loading(self, *, process_events: bool = True) -> None: """Paint the HUD's world-model loading state during blocking setup work.""" if process_events: @@ -498,6 +539,7 @@ def present_world_model_loading(self, *, process_events: bool = True) -> None: self._render_canvas("Loading World Model") self._present_canvas(use_gpu_camera=False) + @nvtx.annotate() def _present_cuda_hud_frame(self, frame: PresentedFrame, rgb: object) -> bool: if self._cuda_hud_interop is None: return False @@ -514,7 +556,8 @@ def _present_cuda_hud_frame(self, frame: PresentedFrame, rgb: object) -> bool: self._update_bev_pil(frame.bev_host_uint8) self._has_camera_frame = True self._render_canvas(frame.status_message, camera_transparent=True) - overlay = np.array(self._canvas, dtype=np.uint8) + with nvtx.annotate("hud_presenter.cuda_hud.overlay_to_numpy", color="yellow"): + overlay = np.array(self._canvas, dtype=np.uint8) camera_area, _panel_rect = self._layout_regions() submitted = self._submit_ready_cuda_hud() @@ -532,6 +575,10 @@ def _present_cuda_hud_frame(self, frame: PresentedFrame, rgb: object) -> bool: def close(self) -> None: self._should_close_flag = True + bev_panel_exec = getattr(self, "_bev_panel_exec", None) + if bev_panel_exec is not None: + bev_panel_exec.shutdown(wait=True, cancel_futures=True) + self._bev_panel_exec = None if self._cuda_hud_interop is not None: with contextlib.suppress(Exception): self._cuda_hud_interop.close() @@ -558,6 +605,7 @@ def _select_view_rgb(frame: PresentedFrame, view_mode: str) -> object: return frame.model_rgb_host_uint8 return frame.rgb_host_uint8 + @nvtx.annotate() def _update_camera_pil(self, rgb: object) -> None: rgb = _as_rgb_host_uint8(rgb) # ``Image.fromarray`` over a contiguous numpy buffer is zero-copy @@ -585,19 +633,12 @@ def _update_camera_pil(self, rgb: object) -> None: self._camera_resize_cache = None self._has_camera_frame = True + @nvtx.annotate() def _update_bev_pil(self, bev_rgb: object) -> None: - bev_rgb = _as_rgb_host_uint8(bev_rgb) - # Wrap the raw BEV; the GoogleMaps recolour runs in - # :meth:`_get_bev_panel_image` *after* the panel-sized resize so the - # float32 pipeline processes ~0.22 MP instead of 1 MP. - if not bev_rgb.flags["C_CONTIGUOUS"]: - bev_rgb = np.ascontiguousarray(bev_rgb) - try: - self._latest_bev_pil = Image.fromarray(bev_rgb, mode="RGB") - except (ValueError, OSError): - return - self._bev_panel_cache_key = None - self._bev_panel_cache = None + # Keep CUDA event synchronization and host materialization off the + # presentation thread. The panel worker consumes this lazy source. + self._latest_bev_source = bev_rgb + self._bev_source_generation += 1 # -- Vulkan / surface plumbing --------------------------------- @@ -779,6 +820,7 @@ def _on_resize(self, width: int, height: int) -> None: # race with whatever frame is in flight. self._pending_resize = self._normalise_present_size(width, height) + @nvtx.annotate() def _submit_ready_cuda_hud(self) -> bool: interop = self._cuda_hud_interop if interop is None: @@ -786,7 +828,7 @@ def _submit_ready_cuda_hud(self) -> bool: interop_frame = interop.ready_rgba_buffer() if interop_frame is None: return False - rgba_buffer, cuda_stream = interop_frame + rgba_buffer, _cuda_stream = interop_frame self._sync_window_size() if self._cuda_hud_interop is not interop: return False @@ -819,13 +861,13 @@ def _submit_ready_cuda_hud(self) -> bool: [width, height, 1], ) encoder.blit(surface_texture, self._display_texture) - submit_id = self._device.submit_command_buffer( - encoder.finish(), - cuda_stream=cuda_stream, - ) + # The interop buffer's CUDA event was queried successfully before + # this point. Forwarding its producer stream would also wait on + # newer frame copies queued behind the completed buffer. + submit_id = self._device.submit_command_buffer(encoder.finish()) interop.mark_submitted(rgba_buffer, submit_id) - del surface_texture self._surface.present() + del surface_texture except RuntimeError as exc: logger.warning( f"[presenter] swapchain present failed ({exc}); reconfiguring", @@ -834,6 +876,7 @@ def _submit_ready_cuda_hud(self) -> bool: return False return True + @nvtx.annotate() def _present_canvas(self, use_gpu_camera: bool = False) -> None: # Sync to the window's CURRENT size before every present. # SDL3 doesn't always fire on_resize for compositor-side rezies @@ -871,8 +914,8 @@ def _present_canvas(self, use_gpu_camera: bool = False) -> None: self._composite_camera_gpu(encoder) encoder.blit(surface_texture, self._display_texture) self._device.submit_command_buffer(encoder.finish()) - del surface_texture self._surface.present() + del surface_texture except RuntimeError as exc: logger.warning( f"[presenter] swapchain present failed ({exc}); reconfiguring", @@ -881,6 +924,7 @@ def _present_canvas(self, use_gpu_camera: bool = False) -> None: # -- GPU camera composite -------------------------------------- + @nvtx.annotate() def _composite_camera_gpu(self, encoder: Any) -> None: """Stamp the camera frame into the display texture on the GPU. @@ -940,6 +984,7 @@ def _compute_camera_fit(self) -> tuple[int, int, int, int] | None: offset_y = (cam_h - fit_h) // 2 return (fit_w, fit_h, offset_x, offset_y) + @nvtx.annotate() def _ensure_camera_texture_uploaded(self) -> bool: """Upload the latest world-model frame to the GPU camera texture. @@ -986,6 +1031,7 @@ def _ensure_camera_texture_uploaded(self) -> bool: self._camera_texture.copy_from_numpy(self._latest_camera_rgba) return True + @nvtx.annotate() def _ensure_camera_fit_texture(self, fit_w: int, fit_h: int) -> None: """Lazily (re)allocate the fit-sized GPU camera texture.""" if self._camera_fit_texture is not None and self._camera_fit_size == ( @@ -1047,6 +1093,7 @@ def _recreate_cuda_hud_interop_after_resize(self, width: int, height: int) -> No # -- Render ------------------------------------------------------ + @nvtx.annotate() def _layout_regions( self, ) -> tuple[tuple[int, int, int, int], tuple[int, int, int, int]]: @@ -1058,6 +1105,7 @@ def _layout_regions( panel_rect = (camera_area[2], 0, screen_w, screen_h) return camera_area, panel_rect + @nvtx.annotate() def _render_canvas( self, status_message: str | None, @@ -1748,6 +1796,7 @@ def _build() -> Image.Image: # -- BEV minimap ------------------------------------------------- + @nvtx.annotate() def _draw_bev( self, canvas: Image.Image, @@ -1765,8 +1814,9 @@ def _draw_bev( inner = (bev_rect[0] + 4, bev_rect[1] + 4, bev_rect[2] - 4, bev_rect[3] - 4) inner_w = inner[2] - inner[0] inner_h = inner[3] - inner[1] + self._bev_panel_target_size = (inner_w, inner_h) - if self._latest_bev_pil is None: + if self._latest_bev_source is None: text = "WAITING FOR BEV..." tbox = _measure_text(self._font_tiny, text) tw = tbox[2] - tbox[0] @@ -1790,38 +1840,46 @@ def _draw_bev( marker_size = max(10, min(inner_w, inner_h) // 14) self._draw_bev_marker(draw, marker_cx, marker_cy, marker_size) + @nvtx.annotate() def _get_bev_panel_image(self, target_size: tuple[int, int]) -> Image.Image | None: - if self._latest_bev_pil is None: + if self._latest_bev_source is None: return None target_w, target_h = target_size if target_w <= 0 or target_h <= 0: return None - key = (id(self._latest_bev_pil), target_w, target_h) - if key == self._bev_panel_cache_key and self._bev_panel_cache is not None: - return self._bev_panel_cache - from omnidreams.interactive_drive.demo import _apply_googlemaps_filter - - bev = self._latest_bev_pil - # Cover-fit + crop, then GoogleMaps filter. Two ordering choices keep - # this cheap (~10 ms vs ~57 ms / tick at 1024 -> ~470 panel): - # 1) Resize before filtering so the per-pixel float32 filter runs on - # ~0.22 MP not 1 MP; it commutes with bilinear resampling to - # within ~2 channel units (visually identical). - # 2) BILINEAR not LANCZOS -- far cheaper, and the tint blend masks - # the sharpness difference at minimap scale. - scale = max(target_w / bev.width, target_h / bev.height) - scaled_w = max(1, int(bev.width * scale)) - scaled_h = max(1, int(bev.height * scale)) - scaled = bev.resize((scaled_w, scaled_h), Image.Resampling.BILINEAR) - crop_left = (scaled_w - target_w) // 2 - crop_top = (scaled_h - target_h) // 2 - cropped = scaled.crop( - (crop_left, crop_top, crop_left + target_w, crop_top + target_h) + key = ( + self._bev_panel_epoch, + self._bev_source_generation, + target_w, + target_h, ) - filtered = _apply_googlemaps_filter(cropped) - self._bev_panel_cache = filtered - self._bev_panel_cache_key = key - return filtered + future = self._bev_panel_future + if future is not None and future.done(): + try: + completed_key, completed_image = future.result() + except Exception as exc: + logger.warning(f"[presenter] BEV panel processing failed: {exc}") + else: + self._bev_panel_cache_key = completed_key + self._bev_panel_cache = completed_image + self._bev_panel_future = None + + if key != self._bev_panel_cache_key and self._bev_panel_future is None: + self._bev_panel_future = self._bev_panel_exec.submit( + _build_bev_panel_image, + key, + self._latest_bev_source, + target_size, + ) + + cache_key = self._bev_panel_cache_key + if ( + cache_key is not None + and cache_key[0] == self._bev_panel_epoch + and cache_key[2:] == target_size + ): + return self._bev_panel_cache + return None @staticmethod def _draw_bev_marker( @@ -2415,7 +2473,14 @@ def _reset_scene_view_state(self) -> None: self._camera_resize_cache_key = None self._camera_resize_cache = None self._latest_camera_pil = None - self._latest_bev_pil = None + self._latest_bev_source = None + self._prepared_bev_source_key = None + self._bev_source_generation = 0 + self._bev_panel_epoch = getattr(self, "_bev_panel_epoch", 0) + 1 + bev_panel_future = getattr(self, "_bev_panel_future", None) + if bev_panel_future is not None: + bev_panel_future.cancel() + self.bev_panel_future = None self._bev_panel_cache_key = None self._bev_panel_cache = None # Panel chrome shows the scene label, so its cache key changes diff --git a/integrations/omnidreams/omnidreams/interactive_drive/streaming_presenter.py b/integrations/omnidreams/omnidreams/interactive_drive/streaming_presenter.py index 8de0cd46e..6f45fb87c 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/streaming_presenter.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/streaming_presenter.py @@ -10,6 +10,8 @@ from __future__ import annotations +import concurrent.futures +import contextlib import io import json import shutil @@ -27,6 +29,7 @@ from omnidreams.interactive_drive.config import RasterConfig from omnidreams.interactive_drive.input.keyboard import KeyboardState from omnidreams.interactive_drive.loading_overlay import render_loading_overlay +import nvtx from omnidreams.interactive_drive.types import DriverCommand, PresentedFrame from PIL import Image @@ -553,6 +556,11 @@ def __init__( # to either waiter are always safe. self._latest_bev_jpeg: bytes | None = None self._bev_frame_count = 0 + self._bev_publish_exec = concurrent.futures.ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="interactive_drive-bev-publish", + ) + self._bev_publish_future: concurrent.futures.Future[None] | None = None # Scene options surfaced to the browser dropdown via /scenes. # Each entry is a dict with ``label``, ``path``, ``variants``; # the demo wrapper builds these from its scene-discovery layer @@ -715,27 +723,43 @@ def bind_keyboard(self, keyboard: KeyboardState) -> None: _KeyboardDriveSink(keyboard) ) + @nvtx.annotate() def process_events(self) -> None: # Per-tick integrator update so auto-crawl smoothing advances at sim # cadence regardless of how often the browser posts /control events. self._keyboard_drive.update() + @nvtx.annotate() + def prepare_frame(self, frame: PresentedFrame, view_mode: str) -> None: + if view_mode == "model_rgb" and frame.model_rgb_host_uint8 is not None: + _prefetch_to_numpy(frame.model_rgb_host_uint8) + else: + _prefetch_to_numpy(frame.rgb_host_uint8) + if frame.bev_host_uint8 is not None: + _prefetch_to_numpy(frame.bev_host_uint8) + + @nvtx.annotate() def present_frame(self, frame: PresentedFrame, view_mode: str) -> None: # Mirror SlangPyPresenter.present_frame's view-mode branching so # the user's `1`/`2` toggles behave identically. - if view_mode == "model_rgb" and frame.model_rgb_host_uint8 is not None: - self._publish( - _with_status_overlay(frame.model_rgb_host_uint8, frame.status_message) - ) - else: - self._publish( - _with_status_overlay(frame.rgb_host_uint8, frame.status_message) - ) + with nvtx.annotate("streaming_presenter.select_view_rgb", color="yellow"): + if view_mode == "model_rgb" and frame.model_rgb_host_uint8 is not None: + rgb = _with_status_overlay( + frame.model_rgb_host_uint8, frame.status_message + ) + else: + rgb = _with_status_overlay(frame.rgb_host_uint8, frame.status_message) + self._publish(rgb) if frame.bev_host_uint8 is not None: - self._publish_bev(frame.bev_host_uint8) + self._submit_bev_publish(frame.bev_host_uint8) def close(self) -> None: self._stop_event.set() + future = self._bev_publish_future + if future is not None: + with contextlib.suppress(Exception): + future.result(timeout=1.0) + self._bev_publish_exec.shutdown(wait=True, cancel_futures=True) # Wake any /stream handlers blocked in ``_frame_cond.wait`` so # they observe ``should_close`` and exit their per-connection loop. with self._frame_cond: @@ -747,17 +771,34 @@ def close(self) -> None: # -- Internals -------------------------------------------------- + @nvtx.annotate() def _publish(self, rgb_host_uint8: object) -> None: buf = io.BytesIO() - Image.fromarray(_as_rgb_host_uint8(rgb_host_uint8)).save( - buf, format="JPEG", quality=self._jpeg_quality + with nvtx.annotate("streaming_presenter.publish_jpeg.materialize_rgb", color="yellow"): + image = Image.fromarray(_as_rgb_host_uint8(rgb_host_uint8)) + with nvtx.annotate("streaming_presenter.publish_jpeg.encode", color="green"): + image.save(buf, format="JPEG", quality=self._jpeg_quality) + jpeg = buf.getvalue() + with nvtx.annotate("streaming_presenter.publish_jpeg.store", color="green"): + with self._frame_cond: + self._latest_jpeg = jpeg + self._frame_count += 1 + self._frame_cond.notify_all() + + @nvtx.annotate() + def _submit_bev_publish(self, bev_rgb_host_uint8: object) -> None: + future = self._bev_publish_future + if future is not None: + if not future.done(): + return + with contextlib.suppress(Exception): + future.result() + self._bev_publish_future = self._bev_publish_exec.submit( + self._publish_bev, + bev_rgb_host_uint8, ) - jpeg = buf.getvalue() - with self._frame_cond: - self._latest_jpeg = jpeg - self._frame_count += 1 - self._frame_cond.notify_all() + @nvtx.annotate() def _publish_bev(self, bev_rgb_host_uint8: object) -> None: """Encode the BEV minimap (quality 95, not 85) and stash it for ``/bev_stream``. @@ -765,14 +806,16 @@ def _publish_bev(self, bev_rgb_host_uint8: object) -> None: otherwise surface as grey halos around lane / vehicle edges. """ buf = io.BytesIO() - Image.fromarray(_as_rgb_host_uint8(bev_rgb_host_uint8)).save( - buf, format="JPEG", quality=95 - ) - jpeg = buf.getvalue() - with self._frame_cond: - self._latest_bev_jpeg = jpeg - self._bev_frame_count += 1 - self._frame_cond.notify_all() + with nvtx.annotate("streaming_presenter.publish_bev.materialize_rgb", color="yellow"): + image = Image.fromarray(_as_rgb_host_uint8(bev_rgb_host_uint8)) + with nvtx.annotate("streaming_presenter.publish_bev.encode", color="yellow"): + image.save(buf, format="JPEG", quality=95) + jpeg = buf.getvalue() + with nvtx.annotate("streaming_presenter.publish_bev.store", color="yellow"): + with self._frame_cond: + self._latest_bev_jpeg = jpeg + self._bev_frame_count += 1 + self._frame_cond.notify_all() def _wait_for_new_frame(self, last_seen_count: int) -> tuple[bytes, int] | None: """Block until a frame newer than ``last_seen_count`` is ready or @@ -1079,6 +1122,13 @@ def _as_rgb_host_uint8(frame: object) -> np.ndarray: return np.ascontiguousarray(np.asarray(frame, dtype=np.uint8)[..., :3]) +def _prefetch_to_numpy(frame: object) -> None: + prefetch = getattr(frame, "prefetch_to_numpy", None) + if callable(prefetch): + prefetch() + + +@nvtx.annotate() def _with_status_overlay(rgb_host_uint8: object, message: str | None) -> np.ndarray: rgb_host_uint8 = _as_rgb_host_uint8(rgb_host_uint8) if message is None: diff --git a/integrations/omnidreams/omnidreams/interactive_drive/video_model/chunk_pipeline.py b/integrations/omnidreams/omnidreams/interactive_drive/video_model/chunk_pipeline.py index ec45e0812..30ed81a2e 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/video_model/chunk_pipeline.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/video_model/chunk_pipeline.py @@ -9,6 +9,7 @@ from typing import Protocol from loguru import logger +import nvtx from omnidreams.interactive_drive.runtime.timing import ( ChunkTimes, TraceContext, @@ -150,6 +151,7 @@ def frame_queue(self) -> "queue.Queue[QueuedFrame]": self._raise_worker_error_if_any() return self._frame_queue + @nvtx.annotate() def request_scene(self, scene: SceneBundle) -> None: """Bind ``scene`` on the worker thread. Non-blocking. @@ -173,6 +175,7 @@ def request_scene(self, scene: SceneBundle) -> None: f"frames={cleared} generation={submit_generation}", ) + @nvtx.annotate() def load_scene_command(backend: VideoModelBackend) -> bool: if submit_generation != self.current_generation: logger.info( @@ -187,6 +190,7 @@ def load_scene_command(backend: VideoModelBackend) -> bool: self._command_queue.put(load_scene_command) + @nvtx.annotate() def request_pose_chunk(self, request: ChunkRequest) -> None: self._raise_worker_error_if_any() @@ -195,6 +199,7 @@ def request_pose_chunk(self, request: ChunkRequest) -> None: trace_dependency_event = request.trace_dependency_event submit_generation = self.current_generation + @nvtx.annotate() def render_command(backend: VideoModelBackend) -> bool: trace_context = ( self._trace_context if chunk_times.chunk_index >= 1 else None @@ -271,22 +276,25 @@ def render_command(backend: VideoModelBackend) -> bool: # the first frame while first_chunk_produced() still reads False. if frame_chunk.frames: self._first_chunk_produced.set() - for frame_index, frame in enumerate(frame_chunk.frames): - frame_times = chunk_times.frames[frame_index] - frame_times.image_ready_time = time.perf_counter() - self._frame_queue.put( - QueuedFrame( - frame=frame, - chunk_times=chunk_times, - frame_index=frame_index, - generation=submit_generation, - worker_ready_event_id=worker_ready_event_id, - ) - ) + with nvtx.annotate("pipeline.enqueue_frames", color="yellow"): + for frame_index, frame in enumerate(frame_chunk.frames): + frame_times = chunk_times.frames[frame_index] + frame_times.image_ready_time = time.perf_counter() + with nvtx.annotate("pipeline.enqueue_frame", color="yellow"): + self._frame_queue.put( + QueuedFrame( + frame=frame, + chunk_times=chunk_times, + frame_index=frame_index, + generation=submit_generation, + worker_ready_event_id=worker_ready_event_id, + ) + ) return True self._command_queue.put(render_command) + @nvtx.annotate() def reset(self) -> None: """Signal the worker to start a new rollout. Non-blocking. @@ -303,6 +311,7 @@ def reset(self) -> None: f"frames={cleared} generation={generation}", ) + @nvtx.annotate() def reset_command(backend: VideoModelBackend) -> bool: backend.reset() return True @@ -314,6 +323,7 @@ def shutdown(self) -> None: self._thread.join() self._raise_worker_error_if_any() + @nvtx.annotate() def _worker(self) -> None: try: warmup_start = time.perf_counter() @@ -333,7 +343,8 @@ def _worker(self) -> None: ) self._model_ready.set() while True: - command = self._command_queue.get() + with nvtx.annotate("pipeline.worker.get_command", color="gray"): + command = self._command_queue.get() if not command(self._backend): return except BaseException as exc: diff --git a/integrations/omnidreams/omnidreams/interactive_drive/video_model/local.py b/integrations/omnidreams/omnidreams/interactive_drive/video_model/local.py index f7795be62..93101c91d 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/video_model/local.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/video_model/local.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from omnidreams.interactive_drive.backends.base import RenderBackend +import nvtx from omnidreams.interactive_drive.types import FrameChunk, SceneBundle, TrajectoryChunk @@ -23,9 +24,11 @@ def __init__(self, backend: RenderBackend) -> None: def can_prewarm(self) -> bool: return self._backend.can_prewarm + @nvtx.annotate() def warmup_model(self) -> None: self._backend.warmup_model() + @nvtx.annotate() def load_scene(self, scene: SceneBundle) -> None: # Scene/variant switches must not carry over rollout cache, text/image # embeddings, or first-chunk state from the previous selection. @@ -36,12 +39,14 @@ def load_scene(self, scene: SceneBundle) -> None: # from the new scene's initial frame and prompt. self._is_first_chunk = True + @nvtx.annotate() def render_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk: if self._is_first_chunk: self._is_first_chunk = False return self._backend.render_first_chunk(trajectory) return self._backend.render_next_chunk(trajectory) + @nvtx.annotate() def reset(self) -> None: self._backend.reset() self._is_first_chunk = True diff --git a/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py b/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py index 22bdb012d..52dfe3fd2 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py @@ -14,6 +14,7 @@ from loguru import logger from omnidreams.interactive_drive.config import WorldModelProfileConfig from omnidreams.interactive_drive.cuda_host_prefetch import CudaHostPrefetch +import nvtx from omnidreams.interactive_drive.world_model.manifest import WorldModelManifest from omnidreams.interactive_drive.world_model.synthetic_fixture import ( build_synthetic_world_model_assets, @@ -103,6 +104,7 @@ def _pipeline_config_log_line( ) +@nvtx.annotate() def _build_pipeline_config( manifest: WorldModelManifest, profile: WorldModelProfileConfig ) -> Any: @@ -284,6 +286,7 @@ def _transformer_overrides(manifest: WorldModelManifest) -> dict[str, object]: } +@nvtx.annotate() def _setup_pipeline_from_config(config: Any, manifest: WorldModelManifest) -> Any: pipeline = config.setup().to(device=torch.device(manifest.device)) if manifest.seed_for_every_rollout is None: @@ -292,6 +295,7 @@ def _setup_pipeline_from_config(config: Any, manifest: WorldModelManifest) -> An return pipeline +@nvtx.annotate() def _precompute_embeddings_from_config( config: Any, manifest: WorldModelManifest, @@ -379,6 +383,7 @@ def _default_pipeline_factory( return _setup_pipeline_from_config(config, manifest) +@nvtx.annotate() def _initial_rgb_tensor(frame: object, *, device: torch.device) -> torch.Tensor: tensor = torch.from_numpy(_rgb_hwc_uint8(frame)) tensor = tensor.permute(2, 0, 1).unsqueeze(0).unsqueeze(0).unsqueeze(2) @@ -492,6 +497,7 @@ def can_prewarm(self) -> bool: or not self._offload_text_encoder ) + @nvtx.annotate() def warmup_model(self) -> None: """Build the scene-independent diffusion pipeline (weights + compile). @@ -514,6 +520,7 @@ def warmup_model(self) -> None: f"[flashdreams-session] model warmup runtime_ms={elapsed_ms:.1f}", ) + @nvtx.annotate() def prepare_for_scene( self, *, initial_rgb: object | None = None, prompt: str | None = None ) -> None: @@ -564,6 +571,7 @@ def _validate_chunk_sizes(self) -> None: f"{steady_chunk_frames} vs {self.manifest.num_frames_per_block}" ) + @nvtx.annotate() def _release_pipeline(self) -> None: if self._pipeline is None: return @@ -574,6 +582,7 @@ def _release_pipeline(self) -> None: torch.cuda.synchronize(device) torch.cuda.empty_cache() + @nvtx.annotate() def start( self, initial_rgb: object, @@ -590,11 +599,13 @@ def start( start = time.perf_counter() with torch.no_grad(): self._cache = self._initialize_cache(initial_rgb, prompt) - video = self.pipeline.generate( - autoregressive_index=0, - cache=self._cache, - hdmap=self._condition_tensor(condition_frames), - ) + hdmap = self._condition_tensor(condition_frames) + with nvtx.annotate("flashdreams_session.pipeline_generate", color="red"): + video = self.pipeline.generate( + autoregressive_index=0, + cache=self._cache, + hdmap=hdmap, + ) model_frames = self._video_tensor_to_frames(video) _synchronize_cuda_frame_event(model_frames) self._pending_finalization_index = 0 @@ -603,6 +614,7 @@ def start( logger.info(f"[flashdreams-session] start total_ms={elapsed_ms:.1f}") return model_frames + @nvtx.annotate() def continue_generation(self, condition_frames: list[object]) -> list[object]: if self._cache is None: raise RuntimeError("start() must be called before continue_generation()") @@ -616,13 +628,21 @@ def continue_generation(self, condition_frames: list[object]) -> list[object]: start = time.perf_counter() with torch.no_grad(): if self._pending_finalization_index is not None: - self.pipeline.finalize(self._pending_finalization_index, self._cache) + with nvtx.annotate("flashdreams_session.pipeline_finalize", color="red"): + self.pipeline.finalize( + self._pending_finalization_index, self._cache + ) self._pending_finalization_index = None - video = self.pipeline.generate( - autoregressive_index=self._next_block_index, - cache=self._cache, - hdmap=self._condition_tensor(condition_frames), - ) + hdmap = self._condition_tensor(condition_frames) + with nvtx.annotate( + "flashdreams_session.continuing_generation.pipeline_generate", + color="red", + ): + video = self.pipeline.generate( + autoregressive_index=self._next_block_index, + cache=self._cache, + hdmap=hdmap, + ) model_frames = self._video_tensor_to_frames(video) _synchronize_cuda_frame_event(model_frames) block_index = self._next_block_index @@ -635,6 +655,7 @@ def continue_generation(self, condition_frames: list[object]) -> list[object]: ) return model_frames + @nvtx.annotate() def reset(self, *, clear_precomputed_embeddings: bool = False) -> None: self._cache = None self._pending_finalization_index = None @@ -646,13 +667,16 @@ def reset(self, *, clear_precomputed_embeddings: bool = False) -> None: "will rerun text/image encoders for the next scene", ) + @nvtx.annotate() def close(self) -> None: if self._cache is not None and self._pending_finalization_index is not None: - self.pipeline.finalize(self._pending_finalization_index, self._cache) + with nvtx.annotate("flashdreams_session.pipeline_finalize", color="red"): + self.pipeline.finalize(self._pending_finalization_index, self._cache) self._pending_finalization_index = None self._cache = None self._pipeline = None + @nvtx.annotate() def _initialize_cache(self, initial_rgb: object, prompt: str) -> Any: if self.manifest.synthetic_model: return self._initialize_synthetic_cache() @@ -697,6 +721,7 @@ def _initialize_synthetic_cache(self) -> Any: view_names=_VIEW_NAMES, ) + @nvtx.annotate() def _ensure_precomputed_embeddings( self, initial_rgb: object, prompt: str ) -> dict[str, torch.Tensor | None]: @@ -709,10 +734,11 @@ def _ensure_precomputed_embeddings( "offload_text_encoder requires flashdreams precompute_embeddings()." ) - embeddings = precompute_embeddings( - text=[[prompt]], - image=self._initial_rgb_tensor(initial_rgb), - ) + with nvtx.annotate("flashdreams_session.precompute_embeddings_call", color="purple"): + embeddings = precompute_embeddings( + text=[[prompt]], + image=self._initial_rgb_tensor(initial_rgb), + ) self._precomputed_embeddings = { "text_embeddings": embeddings["text_embeddings"].cpu(), "image_embeddings": embeddings["image_embeddings"].cpu(), @@ -730,9 +756,11 @@ def _ensure_precomputed_embeddings( logger.info("[flashdreams-session] release_oneshot_encoders done") return self._precomputed_embeddings + @nvtx.annotate() def _initial_rgb_tensor(self, initial_rgb: object) -> torch.Tensor: return _initial_rgb_tensor(initial_rgb, device=self.pipeline.device) + @nvtx.annotate() def _condition_tensor(self, condition_frames: Sequence[object]) -> torch.Tensor: cuda_video = _condition_cuda_video(condition_frames) if cuda_video is not None: @@ -747,6 +775,7 @@ def _to_model_range(self, tensor: torch.Tensor) -> torch.Tensor: return _to_model_range(tensor, device=self.pipeline.device) @staticmethod + @nvtx.annotate() def _video_tensor_to_frames(video: torch.Tensor) -> list[object]: if video.ndim != 6: raise ValueError( @@ -783,6 +812,7 @@ def __init__( self._host: np.ndarray | None = None self._prefetch: CudaHostPrefetch | None = None + @nvtx.annotate() def prefetch_to_numpy(self) -> None: if ( self._host is not None @@ -795,6 +825,7 @@ def prefetch_to_numpy(self) -> None: if prefetch.start(): self._prefetch = prefetch + @nvtx.annotate() def to_numpy(self) -> np.ndarray: if self._host is None: if self._prefetch is not None: @@ -811,6 +842,7 @@ def to_numpy(self) -> np.ndarray: self._frames_hwc_uint8 = None return self._host + @nvtx.annotate() def to_cuda_tensor(self) -> torch.Tensor: if self._frames_hwc_uint8 is None: raise RuntimeError("Lazy RGB frame was already materialized on the host.") @@ -876,6 +908,7 @@ def _condition_cuda_video(condition_frames: Sequence[object]) -> torch.Tensor | return torch.stack(tensors, dim=0) +@nvtx.annotate() def _synchronize_cuda_frame_event(frames: Sequence[object]) -> None: for frame in frames: to_cuda_event = getattr(frame, "to_cuda_event", None) diff --git a/integrations/omnidreams/omnidreams/webrtc/session.py b/integrations/omnidreams/omnidreams/webrtc/session.py index b307c7d21..6ee70a9a9 100644 --- a/integrations/omnidreams/omnidreams/webrtc/session.py +++ b/integrations/omnidreams/omnidreams/webrtc/session.py @@ -805,6 +805,7 @@ def _close_sync(self) -> None: torch.cuda.synchronize(device=self._device) torch.cuda.empty_cache() + @nvtx.annotate() def _generate_one_chunk_sync( self, *, diff --git a/integrations/omnidreams/pyproject.toml b/integrations/omnidreams/pyproject.toml index 5f796eb79..af5209110 100644 --- a/integrations/omnidreams/pyproject.toml +++ b/integrations/omnidreams/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ "grpcio>=1.50", "nvidia-cudnn-frontend==1.22.1", "nvidia-cudnn-cu13==9.23.2.1; sys_platform == 'win32'", + "nvtx>=0.2.15", "opencv-python-headless>=4.5", "shapely>=2.0", # Runtime deps for the ``omnidreams.interactive_drive`` desktop demo diff --git a/integrations/omnidreams/tests/interactive_drive/test_latency_loop.py b/integrations/omnidreams/tests/interactive_drive/test_latency_loop.py index f7140b20a..15db613e6 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_latency_loop.py +++ b/integrations/omnidreams/tests/interactive_drive/test_latency_loop.py @@ -360,6 +360,10 @@ def test_input_to_present_profile_prints_window_summary( output = "\n".join(messages) assert "[profile] e2e" in output assert "wall_present_fps=" in output + assert "present_interval_p50_ms=300.00" in output + assert "present_interval_p95_ms=300.00" in output + assert "present_interval_p99_ms=300.00" in output + assert "present_interval_max_ms=300.00" in output assert "avg_adj_control_to_present_ms=" in output assert "avg_raw_control_to_present_ms=" in output assert "samples=2" in output diff --git a/integrations/omnidreams/tests/interactive_drive/test_presenter.py b/integrations/omnidreams/tests/interactive_drive/test_presenter.py index 36f80b678..f9cc0b201 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_presenter.py +++ b/integrations/omnidreams/tests/interactive_drive/test_presenter.py @@ -3,6 +3,7 @@ from __future__ import annotations +import concurrent.futures import sys from types import SimpleNamespace @@ -16,6 +17,7 @@ ) from omnidreams.interactive_drive.slangpy_hud_presenter import SlangPyHudPresenter from omnidreams.interactive_drive.types import PresentedFrame +from PIL import Image class _LazyFrame: @@ -353,6 +355,80 @@ def test_hud_prepare_frame_keeps_cuda_model_rgb_lazy() -> None: assert bev.prefetch_calls == 1 +def test_hud_prepare_frame_prefetches_one_bev_per_raster_batch() -> None: + presenter = _hud_presenter_without_window() + first = _LazyFrame() + second = _LazyFrame() + batch_key = object() + first.source_group_key = lambda: batch_key # type: ignore[attr-defined] + second.source_group_key = lambda: batch_key # type: ignore[attr-defined] + presenter._cuda_hud_interop = object() + + for bev in (first, second): + presenter.prepare_frame( + PresentedFrame( + timestamp_us=0, + rgb_host_uint8=np.zeros((4, 4, 3), dtype=np.uint8), + depth_host_f32=None, + bev_host_uint8=bev, + ), + view_mode="rgb", + ) + + assert first.prefetch_calls == 1 + assert second.prefetch_calls == 0 + + +def test_hud_bev_panel_reuses_completed_image_while_refresh_is_in_flight() -> None: + presenter = _hud_presenter_without_window() + pending: concurrent.futures.Future[object] = concurrent.futures.Future() + cached = Image.new("RGB", (4, 3), (1, 2, 3)) + presenter._latest_bev_source = np.full((8, 8, 3), 5, dtype=np.uint8) + presenter._bev_source_generation = 4 + presenter._bev_panel_epoch = 2 + presenter._bev_panel_future = pending + presenter._bev_panel_cache_key = (2, 123, 4, 3) + presenter._bev_panel_cache = cached + + assert presenter._get_bev_panel_image((4, 3)) is cached + assert not pending.done() + + +def test_hud_bev_panel_build_runs_outside_draw_path() -> None: + presenter = _hud_presenter_without_window() + presenter._latest_bev_source = np.zeros((8, 8, 3), dtype=np.uint8) + presenter._bev_source_generation = 1 + presenter._bev_panel_epoch = 0 + presenter._bev_panel_future = None + presenter._bev_panel_cache_key = None + presenter._bev_panel_cache = None + presenter._bev_panel_exec = concurrent.futures.ThreadPoolExecutor(max_workers=1) + try: + assert presenter._get_bev_panel_image((4, 3)) is None + assert presenter._bev_panel_future is not None + presenter._bev_panel_future.result(timeout=2.0) + + panel = presenter._get_bev_panel_image((4, 3)) + + assert panel is not None + assert panel.size == (4, 3) + finally: + presenter._bev_panel_exec.shutdown(wait=True, cancel_futures=True) + + +def test_hud_bev_update_keeps_lazy_source_unmaterialized() -> None: + presenter = _hud_presenter_without_window() + lazy = _LazyFrame() + presenter._latest_bev_source = None + presenter._bev_source_generation = 0 + + presenter._update_bev_pil(lazy) + + assert presenter._latest_bev_source is lazy + assert presenter._bev_source_generation == 1 + assert lazy.numpy_calls == 0 + + def test_hud_model_rgb_uses_cuda_path_without_materializing_host_frame() -> None: presenter = _hud_presenter_without_window() lazy = _LazyFrame() @@ -485,6 +561,69 @@ def sync_window_size() -> None: assert mark_calls == 0 +def test_hud_cuda_submit_does_not_forward_copy_stream() -> None: + submitted_kwargs: list[dict[str, object]] = [] + marked: list[tuple[object, int]] = [] + surface_events: list[str] = [] + buffer = SimpleNamespace( + buffer=object(), + size_bytes=16, + row_pitch=8, + ) + + class _Interop: + def ready_rgba_buffer(self) -> tuple[object, object]: + return buffer, object() + + def mark_submitted(self, submitted_buffer: object, submit_id: int) -> None: + marked.append((submitted_buffer, submit_id)) + + class _Encoder: + def copy_buffer_to_texture(self, *args: object) -> None: + del args + + def blit(self, *args: object) -> None: + del args + + def finish(self) -> object: + return object() + + class _Device: + def create_command_encoder(self) -> _Encoder: + return _Encoder() + + def submit_command_buffer(self, command: object, **kwargs: object) -> int: + del command + submitted_kwargs.append(kwargs) + return 7 + + class _Surface: + config = object() + + def acquire_next_image(self) -> object: + class _SurfaceTexture: + def __del__(self) -> None: + surface_events.append("release") + + return _SurfaceTexture() + + def present(self) -> None: + surface_events.append("present") + + presenter = _hud_presenter_without_window() + presenter._cuda_hud_interop = _Interop() + presenter._sync_window_size = lambda: None + presenter._configured_size = (2, 2) + presenter._surface = _Surface() + presenter._device = _Device() + presenter._display_texture = object() + + assert presenter._submit_ready_cuda_hud() + assert submitted_kwargs == [{}] + assert marked == [(buffer, 7)] + assert surface_events == ["present", "release"] + + def test_hud_model_rgb_falls_back_to_host_when_cuda_path_declines() -> None: presenter = _hud_presenter_without_window() lazy = _LazyFrame() @@ -579,7 +718,8 @@ def _hud_presenter_for_exit(selected_variant: str) -> SlangPyHudPresenter: presenter._camera_resize_cache_key = object() presenter._camera_resize_cache = object() presenter._latest_camera_pil = object() - presenter._latest_bev_pil = object() + presenter._latest_bev_source = object() + presenter._bev_source_generation = 1 presenter._bev_panel_cache_key = object() presenter._bev_panel_cache = object() presenter._panel_chrome_cache_key = object() diff --git a/integrations/omnidreams/tests/interactive_drive/test_rasterizer.py b/integrations/omnidreams/tests/interactive_drive/test_rasterizer.py index f6e8af38d..11452e668 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_rasterizer.py +++ b/integrations/omnidreams/tests/interactive_drive/test_rasterizer.py @@ -3,16 +3,21 @@ from __future__ import annotations +import concurrent.futures from types import SimpleNamespace import numpy as np +import pytest import torch from omnidreams.interactive_drive.rasterizer import ( _LoadedSceneData, _LudusConditionRasterizerImpl, _RenderedCameraFrames, + LudusConditionRasterizer, ) +pytestmark = pytest.mark.ci_cpu + class _Event: def __init__(self) -> None: @@ -72,3 +77,60 @@ def test_raster_chunk_can_disable_cuda_backed_frames() -> None: assert isinstance(first, np.ndarray) assert not callable(getattr(first, "to_cuda_tensor", None)) assert np.array_equal(first, np.arange(18, dtype=np.uint8).reshape(2, 3, 3)) + + +def test_lagged_bev_poll_does_not_wait_for_in_flight_render() -> None: + rasterizer = LudusConditionRasterizer.__new__(LudusConditionRasterizer) + pending: concurrent.futures.Future[_RenderedCameraFrames | None] = ( + concurrent.futures.Future() + ) + latest = _RenderedCameraFrames( + frames_hwc_uint8=torch.zeros((1, 1, 1, 3), dtype=torch.uint8), + ready_event=None, + ) + rasterizer._pending_bev = pending + rasterizer._latest_bev = latest + + assert rasterizer._poll_ready_bev() is latest + assert rasterizer._pending_bev is pending + + +def test_lagged_bev_poll_promotes_completed_render() -> None: + rasterizer = LudusConditionRasterizer.__new__(LudusConditionRasterizer) + pending: concurrent.futures.Future[_RenderedCameraFrames | None] = ( + concurrent.futures.Future() + ) + rendered = _RenderedCameraFrames( + frames_hwc_uint8=torch.ones((1, 1, 1, 3), dtype=torch.uint8), + ready_event=None, + ) + pending.set_result(rendered) + rasterizer._pending_bev = pending + rasterizer._latest_bev = None + + assert rasterizer._poll_ready_bev() is rendered + assert rasterizer._pending_bev is None + + +def test_build_chunk_resamples_lagged_bev_with_different_frame_count() -> None: + impl = _impl_for_render_chunk(use_cuda_frames=True) + rgb_frames = _RenderedCameraFrames( + frames_hwc_uint8=torch.zeros((7, 1, 1, 3), dtype=torch.uint8), + ready_event=None, + ) + bev_frames = _RenderedCameraFrames( + frames_hwc_uint8=torch.arange(5, dtype=torch.uint8).reshape(5, 1, 1, 1), + ready_event=None, + ) + + chunk = impl.build_chunk( + timestamps_us=np.arange(7, dtype=np.int64), + rgb_frames=rgb_frames, + bev_frames=bev_frames, + ) + + bev_values = [ + int(frame.bev_host_uint8.to_cuda_tensor()[0, 0, 0]) + for frame in chunk.frames + ] + assert bev_values == [0, 1, 1, 2, 3, 3, 4] diff --git a/uv.lock b/uv.lock index 2b0d03215..d1dea933e 100644 --- a/uv.lock +++ b/uv.lock @@ -1509,6 +1509,7 @@ dependencies = [ { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, { name = "nvidia-cudnn-cu13", version = "9.23.2.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'win32' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, { name = "nvidia-cudnn-frontend" }, + { name = "nvtx" }, { name = "opencv-python-headless" }, { name = "pillow" }, { name = "pyarrow" }, @@ -1550,6 +1551,7 @@ requires-dist = [ { name = "numpy", specifier = ">=1.26" }, { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'win32'", specifier = "==9.23.2.1" }, { name = "nvidia-cudnn-frontend", specifier = "==1.22.1" }, + { name = "nvtx", specifier = ">=0.2.15" }, { name = "opencv-python-headless", specifier = ">=4.5" }, { name = "pillow", specifier = ">=10.0" }, { name = "pyarrow", specifier = ">=16.0" }, @@ -3200,26 +3202,10 @@ name = "numpy" version = "2.4.6" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.14' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version == '3.13.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version == '3.12.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform == 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.14' and sys_platform != 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform != 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.14' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.13.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.12.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", + "python_full_version >= '3.12' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version >= '3.12' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform != 'win32'", ] sdist = { url = "https://files.pythonhosted.org/packages/d0/ad/fed0499ce6a338d2a03ebae59cd15093910c8875328855781952abf6c2fe/numpy-2.4.6.tar.gz", hash = "sha256:f3a3570c4a2a16746ac2c31a7c7c7b0c186b95ce902e33db6f28094ed7387dda", size = 20735807, upload-time = "2026-05-18T23:37:14.07Z" } wheels = [ @@ -3731,6 +3717,35 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9f/99/4c9c0c329bf9fc125008c3b54c7c94c0023518d06fc025ae36431375e1fe/nvidia_nvtx_cu12-12.8.90-py3-none-win_amd64.whl", hash = "sha256:619c8304aedc69f02ea82dd244541a83c3d9d40993381b3b590f1adaed3db41e", size = 56492, upload-time = "2025-03-07T01:52:24.69Z" }, ] +[[package]] +name = "nvtx" +version = "0.2.15" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/92/dd/692765e87de30bae1522cdffaa0f2b52949658a92a0fa6d96b1a01eae9d2/nvtx-0.2.15.tar.gz", hash = "sha256:2287d3be05b85661deb386f878d1f536c2e532774aa9ec7a50c434942ed81ae5", size = 121230, upload-time = "2026-03-18T10:01:25.547Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9d/ef/ea1e9d92afd07fdf2a2390e508f1d214e5ba890561d7849d6ca708534b9d/nvtx-0.2.15-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a4f50832fd90a1b480a9deef6e4cd48015b61869095b54dd1a7afe87b4138c6a", size = 768543, upload-time = "2026-03-18T10:07:21.819Z" }, + { url = "https://files.pythonhosted.org/packages/32/8e/b42c05cf3cc43c51f21fdda6f7c4fe28a595c6d2bdb0cfbf0477dc5805f2/nvtx-0.2.15-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5f3362f0db4252514719326c9d5662b0f93d254659ba97b9c8dbe556286e0e3e", size = 771975, upload-time = "2026-03-18T10:12:23.772Z" }, + { url = "https://files.pythonhosted.org/packages/60/77/fc000055b5bb1651cdd772f0fe1fd9a16c7773b28dfc5624eea331d1415d/nvtx-0.2.15-cp310-cp310-win_amd64.whl", hash = "sha256:d71f934e580d4572f382712b6da464ab69e4c212981506f781f927d5c6d935d6", size = 134503, upload-time = "2026-03-18T10:04:05.773Z" }, + { url = "https://files.pythonhosted.org/packages/80/65/435d10b2041ee082c07d5aed129afd504012c8908796d695f10e66bcc716/nvtx-0.2.15-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:157b80ea9b4db6c8f47f8dbe2fa2e81e7a7f1445bb87f8268f43dec9210b78a1", size = 806443, upload-time = "2026-03-18T10:05:49.308Z" }, + { url = "https://files.pythonhosted.org/packages/47/bc/be94576ba33af75bcc68a857daade64cb86481764d4fb0f36308b1f6fc85/nvtx-0.2.15-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:02bca69ee55e0be41eabf908de9dbcdd18e702c7f49f9aa63fd396ce684ff5d5", size = 808183, upload-time = "2026-03-18T10:11:16.262Z" }, + { url = "https://files.pythonhosted.org/packages/f6/7a/42109f1cfb1ff9913201cb2b804956a4f003db4c018c2522a3c8066b3a1c/nvtx-0.2.15-cp311-cp311-win_amd64.whl", hash = "sha256:dbe41f78f5a811bd4cdad0a237e5b41a4937d8c2c6c9abdd161091671a598bc0", size = 134631, upload-time = "2026-03-18T10:02:11.247Z" }, + { url = "https://files.pythonhosted.org/packages/c2/07/698355285a03a366ef63ea9762fc1feef3f9f25483e1655408f72d827090/nvtx-0.2.15-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2cc530cd0f1a2c14a3a7e683833db509888ac5ed4ead94e5c9e2c7317c6937a7", size = 807159, upload-time = "2026-03-18T10:09:49.232Z" }, + { url = "https://files.pythonhosted.org/packages/c0/d1/08f22448d83481408d663065764ba583df091a7de629ed38fc97e522f1af/nvtx-0.2.15-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3ca8030a6d197952318013dd1c12c22da1d4b9feb76ba72e0fcd449961183c2c", size = 806187, upload-time = "2026-03-18T10:13:32.972Z" }, + { url = "https://files.pythonhosted.org/packages/54/23/c97c39e3b7ba256aa343cb828ca0d1c8421f705ca84795658ecd14ca95ed/nvtx-0.2.15-cp312-cp312-win_amd64.whl", hash = "sha256:70a1e768964e0520b68ccabc4df391cc227537c45936a7eba6507bc65e617e00", size = 129178, upload-time = "2026-03-18T10:02:55.299Z" }, + { url = "https://files.pythonhosted.org/packages/05/c9/8341224b8284f7deb6a634119939de5885adc421e64b6743693b30da2186/nvtx-0.2.15-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d28660d9c46f8ba750d781572b6aa5a1e6221abba224ab32d7fb32c2d0fd67df", size = 780787, upload-time = "2026-03-18T10:10:40.634Z" }, + { url = "https://files.pythonhosted.org/packages/b1/c0/4a5bb7897918de7c7e0191d9342df8ae4cb797ff07276e0f20d13e497ce7/nvtx-0.2.15-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10749686633f880ad53dcdbb2179fad41b45dcf5b7631d4a1070a577577bd386", size = 782575, upload-time = "2026-03-18T10:13:57.3Z" }, + { url = "https://files.pythonhosted.org/packages/38/b9/6b381ac7c5a3ded331aebbf25f8959d19b51d320fb2514c76c6b6edddaaa/nvtx-0.2.15-cp313-cp313-win_amd64.whl", hash = "sha256:a6650b029263d12f8427a4dee8bd59cb9c91bccb60543bfcb20bc2b00fdcd672", size = 128764, upload-time = "2026-03-18T10:02:33.343Z" }, + { url = "https://files.pythonhosted.org/packages/75/69/a9acb6d95d2e0e381b2956544768528dd8d7a9e827af8c2014169d838284/nvtx-0.2.15-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:25813ead4fff4d3a6e04f69a72507b096a6bdbecefa369f1100b0e584767bca8", size = 833375, upload-time = "2026-03-18T10:06:31.955Z" }, + { url = "https://files.pythonhosted.org/packages/38/56/c7e8645061cc2fc23f3a54f33e1e340df59216f07dcfb97d46b8ae7dd26c/nvtx-0.2.15-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3741edac4678b92f03d22a3f0a2dfd469f422f85e63db71b038e02525b2404ad", size = 788639, upload-time = "2026-03-18T10:12:01.69Z" }, + { url = "https://files.pythonhosted.org/packages/96/03/fadd82acdbca6d1c49ac517081a0c3714346f52f4c7e1d4449d77605b4aa/nvtx-0.2.15-cp313-cp313t-win_amd64.whl", hash = "sha256:8be06c3c8c267eba56a0396366b9593092e0b75ea8d3702b303d48c0a1662f0e", size = 142609, upload-time = "2026-03-18T10:01:48.832Z" }, + { url = "https://files.pythonhosted.org/packages/e0/5b/ca0ba6fa769d08174b7a5b4775c279e2e26611cdd5e7833aa699187871c7/nvtx-0.2.15-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b5171b8283dd3ea9ae688a86d16901b4c2c142c4eb0a4bdbf6c222f5f67f9524", size = 781769, upload-time = "2026-03-18T10:08:59.357Z" }, + { url = "https://files.pythonhosted.org/packages/f7/e1/e02fafc01c18f1868a2d2c030953f49e38d65f2d95884789a6c46ff308f1/nvtx-0.2.15-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3c6d0f27d4f8a2f479eb64a6b842c13aee32120348a1715d995b9bb9f75b35cf", size = 774614, upload-time = "2026-03-18T10:12:46.979Z" }, + { url = "https://files.pythonhosted.org/packages/20/77/a2b64335bab7c75fe1c054cc4ebe2d3b3234cbdb04d2e1d6ca73551c54f5/nvtx-0.2.15-cp314-cp314-win_amd64.whl", hash = "sha256:9934fad0b441cfa6e896a848b092498ba23e2ff205c2b9a7b60520ff8367ffef", size = 130932, upload-time = "2026-03-18T10:03:43.507Z" }, + { url = "https://files.pythonhosted.org/packages/db/24/528619230976c18364eda2340906ea67b3bf7588b7ce59e054723614abae/nvtx-0.2.15-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:aca61135c76b8107ae3c994325613afa661e1336a991c59cc9c6176829b3b32c", size = 834439, upload-time = "2026-03-18T10:05:01.181Z" }, + { url = "https://files.pythonhosted.org/packages/ef/7b/c1b96f13ef89bdf2a8c2f326a97bed89699271990d7c8624fda3fedc6e61/nvtx-0.2.15-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:58653bf6fd8453947b9e5153da2ad7aeb0ceafa030de7f133efb3eada5da7ca7", size = 790247, upload-time = "2026-03-18T10:11:39.124Z" }, + { url = "https://files.pythonhosted.org/packages/14/5d/e000de781d92b732d52c572517db0e9e3a0085795f8bdc18201713c52d1f/nvtx-0.2.15-cp314-cp314t-win_amd64.whl", hash = "sha256:9d1d10db4fb4a3b0ffd6ed37bf25f0a966a3b4d34b3c9abb1f6572732959a6e5", size = 149109, upload-time = "2026-03-18T10:03:21.615Z" }, +] + [[package]] name = "onnx" version = "1.22.0"