diff --git a/integrations/omnidreams/README.md b/integrations/omnidreams/README.md index f58d892f3..e2f0b12c3 100644 --- a/integrations/omnidreams/README.md +++ b/integrations/omnidreams/README.md @@ -80,6 +80,9 @@ From the workspace root, run: uv run --package flashdreams-omnidreams torchrun --nproc_per_node 1 -m omnidreams.webrtc.server --pipeline_config_name omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf --scene-uuid 065dcac9-ee67-4434-a835-c6b816c88e48 --port 8089 ``` +For issue 195 renderer/serving profiling, see +[docs/issue195_renderer_profile_repro.md](docs/issue195_renderer_profile_repro.md). + When `--scene_dir` is omitted, the server downloads the selected scene from the configured Hugging Face org, extracts its `clipgt-.usdz` archive, and stages it under `FLASHDREAMS_CACHE_DIR` (or `~/.cache/flashdreams`). If diff --git a/integrations/omnidreams/docs/issue195_renderer_profile_repro.md b/integrations/omnidreams/docs/issue195_renderer_profile_repro.md new file mode 100644 index 000000000..25caf0250 --- /dev/null +++ b/integrations/omnidreams/docs/issue195_renderer_profile_repro.md @@ -0,0 +1,105 @@ +# Issue 195 Renderer Profiling Repro + +This branch includes temporary profiling instrumentation for OmniDreams WebRTC +serving. It records whether end-to-end serving time is spent in the model +pipeline, queue/enqueue logic, or Ludus render conditioning. + +## Base + +The original repro was run from: + +```text +upstream/main +3816e32d99a001eab91996f1e52fe488be1ee9cf +``` + +Use this branch directly after it is published. No patch application is needed. + +## Launch On HSG + +Set `HF_TOKEN` or `HF_TOKEN_FILE` first. Then request a full GB200 node, but run +the WebRTC server with one visible GPU, matching the OmniDreams WebRTC README +path. + +```bash +cd /lustre/fsw/portfolios/nvr/projects/nvr_torontoai_videogen/users/junchenl/flashdreams +git fetch issue-195-render-profile-repro +git switch issue-195-render-profile-repro + +uv sync --package flashdreams-omnidreams --extra interactive-drive + +srun -A nvr_torontoai_videogen -p batch --qos=interactive \ + -N1 --gpus-per-node=4 --ntasks=1 --cpus-per-task=16 --time=02:00:00 \ + --chdir "$PWD" \ + bash -lc ' + export OMNIDREAMS_LUDUS_RENDER_PROFILE=1 + export OMNIDREAMS_LUDUS_RENDER_PROFILE_CUDA_EVENTS=1 + unset OMNIDREAMS_LUDUS_MSAA_SAMPLES + CUDA_VISIBLE_DEVICES=0 uv run --no-sync --package flashdreams-omnidreams \ + torchrun --nproc_per_node 1 -m omnidreams.webrtc.server \ + --pipeline_config_name omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf \ + --scene-uuid 065dcac9-ee67-4434-a835-c6b816c88e48 \ + --port 8095 2>&1 | tee issue195_webrtc_profile.log + ' +``` + +`OMNIDREAMS_LUDUS_MSAA_SAMPLES` is intentionally unset so the renderer uses its +default `4` samples. The profile log should include `renderer_msaa_samples: 4.0`. + +## Forward And Drive The Viewer + +In another shell, replace `` with the Slurm node name: + +```bash +socat -d -d TCP-LISTEN:8097,reuseaddr,fork,bind=0.0.0.0 \ + TCP:.cm.cluster:8095 +``` + +Open: + +```text +http://oci-hsg-cs-001-vscode-01:8097/request_session +``` + +Click Connect Session and drive for 20-30 seconds. The first few chunks may +include warmup or compile noise, so summarize later chunks. + +## Summarize + +```bash +python integrations/omnidreams/scripts/summarize_webrtc_profile.py \ + issue195_webrtc_profile.log \ + --min-chunk 4 +``` + +Use `--session latest` for the most recent browser run, or `--session all` to +include loopback warmup and browser sessions together. + +## Expected Signature + +The model pipeline is fast enough for more than 40 FPS, but the end-to-end +viewer generation loop is around 18-20 FPS because render conditioning dominates. + +Observed browser-session summary from the original repro: + +```text +selected_chunks=31 session=latest min_chunk=4 +gen_ms avg= 439.3 ms / 8 frames +wrapper_render_condition_ms avg= 259.8 ms +renderer_ctx_render_ms avg= 258.6 ms +ctx_render_plugin_cuda_ms_sum avg= 254.9 ms +pipeline_total_ms avg= 175.0 ms +pipeline_total_ms_wo_finalize avg= 126.2 ms +enqueue_ms avg= 26.4 ms +``` + +Interpretation: + +```text +wrapper_render_condition_ms ~= renderer_ctx_render_ms +renderer_ctx_render_ms ~= ctx_render_plugin_cuda_ms_sum +``` + +So the extra serving time is inside `self.ctx.render(...)`, and nearly all of +that time is accumulated by Ludus CUDA render plugin calls across the 8 +conditioning frames. diff --git a/integrations/omnidreams/ludus-renderer/ludus_renderer/_ops/context.py b/integrations/omnidreams/ludus-renderer/ludus_renderer/_ops/context.py index 7e2f78ed1..125d68230 100644 --- a/integrations/omnidreams/ludus-renderer/ludus_renderer/_ops/context.py +++ b/integrations/omnidreams/ludus-renderer/ludus_renderer/_ops/context.py @@ -15,6 +15,7 @@ """Ludus rendering context - ``LudusCudaTimestampedContext``.""" +import time from typing import List, Optional, Tuple import torch @@ -27,6 +28,24 @@ ) +def _record_ms_summary( + profile: dict[str, float], prefix: str, values: list[float] +) -> None: + if not values: + return + profile[f"{prefix}_count"] = float(len(values)) + profile[f"{prefix}_sum"] = float(sum(values)) + profile[f"{prefix}_avg"] = float(sum(values) / len(values)) + profile[f"{prefix}_min"] = float(min(values)) + profile[f"{prefix}_max"] = float(max(values)) + + +def _create_event_profiler(): + from flashdreams.infra.profiler import EventProfiler + + return EventProfiler() + + def _compute_element_aabbs( vertices: torch.Tensor, prefix_sum: torch.Tensor, device: torch.device ) -> torch.Tensor: @@ -77,6 +96,9 @@ def __init__(self, device=None): self._cameras: List[FThetaCamera] = [] self._camera_intrinsics: Optional[torch.Tensor] = None self.needs_vflip = False # CUDA renders top-down (standard image convention) + self.enable_render_profiling = False + self.enable_render_profile_cuda_events = False + self.last_render_profile: dict[str, float] = {} # Per-scene flat buffers self._scenes: List[dict] = [] @@ -424,21 +446,47 @@ def render( """ assert self._camera_intrinsics is not None, "call upload_cameras first" + total_t0 = time.perf_counter() n = scene_ids.shape[0] all_images = [] + profile_enabled = bool(getattr(self, "enable_render_profiling", False)) + cuda_event_enabled = ( + profile_enabled + and bool(getattr(self, "enable_render_profile_cuda_events", False)) + and scene_ids.is_cuda + and torch.cuda.is_available() + ) + profile: dict[str, float] = {} + scalar_item_ms: list[float] = [] + query_prep_ms: list[float] = [] + plugin_host_ms: list[float] = [] + event_profiler = _create_event_profiler() if cuda_event_enabled else None + + if profile_enabled: + profile["ctx_render_batch_size"] = float(n) plugin = _get_plugin() + loop_t0 = time.perf_counter() for i in range(n): + scalar_item_t0 = time.perf_counter() sid = int(scene_ids[i].item()) cid = int(camera_ids[i].item()) ts_val = int(timestamps_us[i].item()) cam_type = int(camera_type_ids[i].item()) + if profile_enabled: + scalar_item_ms.append((time.perf_counter() - scalar_item_t0) * 1e3) + query_prep_t0 = time.perf_counter() sd = self._scenes[sid] cam_intrinsics = self._camera_intrinsics[cid : cid + 1].contiguous() pose = camera_poses[i : i + 1].contiguous() + if profile_enabled: + query_prep_ms.append((time.perf_counter() - query_prep_t0) * 1e3) + if event_profiler is not None: + event_profiler.record(f"before_plugin_{i}") + plugin_t0 = time.perf_counter() img = plugin.ludus_render_fwd_cuda_timestamped( self.cpp_wrapper, sd["timestamps"], @@ -459,6 +507,46 @@ def render( resolution, self._tessellation_threshold, ) + if event_profiler is not None: + event_profiler.record(f"plugin_{i}") + if profile_enabled: + plugin_host_ms.append((time.perf_counter() - plugin_t0) * 1e3) all_images.append(img) - return torch.cat(all_images, dim=0) + if profile_enabled: + profile["ctx_render_loop_host_ms"] = (time.perf_counter() - loop_t0) * 1e3 + _record_ms_summary( + profile, "ctx_render_scalar_item_host_ms", scalar_item_ms + ) + _record_ms_summary(profile, "ctx_render_query_prep_host_ms", query_prep_ms) + _record_ms_summary(profile, "ctx_render_plugin_host_ms", plugin_host_ms) + + if event_profiler is not None: + event_profiler.record("before_cat") + cat_t0 = time.perf_counter() + output = torch.cat(all_images, dim=0) + if event_profiler is not None: + event_profiler.record("cat") + if profile_enabled: + profile["ctx_render_cat_host_ms"] = (time.perf_counter() - cat_t0) * 1e3 + + if event_profiler is not None: + sync_t0 = time.perf_counter() + cuda_elapsed_ms = event_profiler.sync_and_summarize() + profile["ctx_render_cuda_event_sync_host_ms"] = ( + time.perf_counter() - sync_t0 + ) * 1e3 + _record_ms_summary( + profile, + "ctx_render_plugin_cuda_ms", + [cuda_elapsed_ms[f"plugin_{i}"] for i in range(n)], + ) + profile["ctx_render_cat_cuda_ms"] = cuda_elapsed_ms["cat"] + + if profile_enabled: + profile["ctx_render_total_host_ms"] = (time.perf_counter() - total_t0) * 1e3 + self.last_render_profile = profile + else: + self.last_render_profile = {} + + return output diff --git a/integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py b/integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py index 2c7118dee..8bafd2956 100644 --- a/integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py +++ b/integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py @@ -23,6 +23,7 @@ from __future__ import annotations +import time from dataclasses import dataclass import numpy as np @@ -61,6 +62,22 @@ class TextPrompt: ) +def _numeric_stats(stats: dict[str, float] | None) -> dict[str, float]: + if not stats: + return {} + + numeric_stats: dict[str, float] = {} + for key, value in stats.items(): + try: + numeric_value = float(value) + except (TypeError, ValueError): + continue + if not np.isfinite(numeric_value): + continue + numeric_stats[key] = numeric_value + return numeric_stats + + @dataclass class OmnidreamsConditioningState: """State for generation, including renderer and pipeline state.""" @@ -79,6 +96,7 @@ class GenerationOutput: None # Generated video frames [B, V, T, 3, H, W] (None if skip_video_generation) ) finalization_state: dict | None = None # Finalization state from the video model + stats: dict[str, float] | None = None # Host-wall profiling stats for this call class OmnidreamsConditioningWrapper(nn.Module): @@ -174,11 +192,11 @@ def finalize_block_generation( self, pipeline_cache: OmnidreamsPipelineCache, finalization_state: dict | None, - ) -> None: + ) -> dict[str, float] | None: if finalization_state is None: - return + return None block_idx = int(finalization_state["autoregressive_index"]) - self.pipeline.finalize( + return self.pipeline.finalize( autoregressive_index=block_idx, cache=pipeline_cache, ) @@ -397,18 +415,22 @@ def start_generation( and ``rgb_frames`` ``[B, V, T, 3, H, W]`` (or None). """ + total_t0 = time.perf_counter() assert len(text_prompts) == 1, ( "Only one text prompt (batch size == 1) is supported for now" ) + validate_t0 = time.perf_counter() self._validate_camera_inputs( camera_names=camera_names, camera_poses_per_view=camera_poses_per_view, frame_timestamps_us=frame_timestamps_us, expected_length=self.initial_frame_chunk_size, ) + stats = {"wrapper_validate_ms": (time.perf_counter() - validate_t0) * 1e3} # condition_frames: [B, V, T, 3, H, W] + render_t0 = time.perf_counter() condition_frames = self._render_condition_frames( renderer, camera_names, @@ -416,15 +438,22 @@ def start_generation( frame_timestamps_us, dynamic_actor_pool, ) + stats["wrapper_render_condition_ms"] = (time.perf_counter() - render_t0) * 1e3 + stats.update(_numeric_stats(getattr(renderer, "last_render_stats", None))) if skip_video_generation: state = OmnidreamsConditioningState( renderer=renderer, ) + stats["wrapper_total_ms"] = (time.perf_counter() - total_t0) * 1e3 return GenerationOutput( - state=state, condition_frames=condition_frames, rgb_frames=None + state=state, + condition_frames=condition_frames, + rgb_frames=None, + stats=stats, ) + prepare_t0 = time.perf_counter() initial_rgb_frames, condition_frames = self._normalize_start_inputs( initial_rgb_frames, condition_frames ) @@ -433,16 +462,32 @@ def start_generation( first_frame = self._to_model_range(initial_rgb_frames).unsqueeze(2) condition = self._to_model_range(condition_frames) + stats["wrapper_prepare_condition_ms"] = ( + time.perf_counter() - prepare_t0 + ) * 1e3 + initialize_cache_t0 = time.perf_counter() pipeline_cache = self.pipeline.initialize_cache( text=text, image=first_frame, view_names=camera_names ) + stats["wrapper_initialize_cache_ms"] = ( + time.perf_counter() - initialize_cache_t0 + ) * 1e3 + + pipeline_generate_t0 = time.perf_counter() rgb_frames = self.pipeline.generate( autoregressive_index=0, hdmap=condition, cache=pipeline_cache, ) + stats["wrapper_pipeline_generate_ms"] = ( + time.perf_counter() - pipeline_generate_t0 + ) * 1e3 + + to_uint8_t0 = time.perf_counter() rgb_frames = self._to_uint8(rgb_frames).contiguous() + stats["wrapper_to_uint8_ms"] = (time.perf_counter() - to_uint8_t0) * 1e3 + stats["wrapper_total_ms"] = (time.perf_counter() - total_t0) * 1e3 state = OmnidreamsConditioningState( renderer=renderer, @@ -453,6 +498,7 @@ def start_generation( condition_frames=condition_frames, rgb_frames=rgb_frames, finalization_state={"autoregressive_index": 0}, + stats=stats, ) def continue_generation( @@ -480,6 +526,8 @@ def continue_generation( GenerationOutput with ``condition_frames`` ``[B, V, T, 3, H, W]`` and ``rgb_frames`` ``[B, V, T, 3, H, W]`` (or None). """ + total_t0 = time.perf_counter() + validate_t0 = time.perf_counter() self._validate_camera_inputs( camera_names=camera_names, camera_poses_per_view=camera_poses_per_view, @@ -487,11 +535,13 @@ def continue_generation( expected_length=self.frame_chunk_size, ) renderer = state.renderer + stats = {"wrapper_validate_ms": (time.perf_counter() - validate_t0) * 1e3} profiler = get_profiler() session_id, chunk_idx = get_profiling_context() # condition_frames: [B, V, T, 3, H, W] + render_t0 = time.perf_counter() with profiler.measure( "render_condition_frames", session_id=session_id, chunk_idx=chunk_idx ): @@ -502,10 +552,16 @@ def continue_generation( frame_timestamps_us, dynamic_actor_pool, ) + stats["wrapper_render_condition_ms"] = (time.perf_counter() - render_t0) * 1e3 + stats.update(_numeric_stats(getattr(renderer, "last_render_stats", None))) if skip_video_generation: + stats["wrapper_total_ms"] = (time.perf_counter() - total_t0) * 1e3 return GenerationOutput( - state=state, condition_frames=condition_frames, rgb_frames=None + state=state, + condition_frames=condition_frames, + rgb_frames=None, + stats=stats, ) if state.pipeline_cache is None: @@ -514,11 +570,16 @@ def continue_generation( "(session was started with skip_video_generation=True)" ) + prepare_t0 = time.perf_counter() model_cond = self._normalize_condition_input(condition_frames) condition = self._to_model_range(model_cond) prev_block_idx = state.pipeline_cache.autoregressive_index block_idx = 0 if prev_block_idx is None else prev_block_idx + 1 + stats["wrapper_prepare_condition_ms"] = ( + time.perf_counter() - prepare_t0 + ) * 1e3 + pipeline_generate_t0 = time.perf_counter() with profiler.measure( "pipeline.continue_generation", session_id=session_id, @@ -530,7 +591,16 @@ def continue_generation( hdmap=condition, cache=state.pipeline_cache, ) + stats["wrapper_pipeline_generate_ms"] = ( + time.perf_counter() - pipeline_generate_t0 + ) * 1e3 + + to_uint8_t0 = time.perf_counter() rgb_frames = self._to_uint8(rgb_frames).contiguous() + stats["wrapper_to_uint8_ms"] = ( + time.perf_counter() - to_uint8_t0 + ) * 1e3 + stats["wrapper_total_ms"] = (time.perf_counter() - total_t0) * 1e3 new_state = OmnidreamsConditioningState( renderer=renderer, @@ -541,6 +611,7 @@ def continue_generation( condition_frames=condition_frames, rgb_frames=rgb_frames, finalization_state={"autoregressive_index": block_idx}, + stats=stats, ) def cleanup(self, state: OmnidreamsConditioningState) -> None: diff --git a/integrations/omnidreams/omnidreams/conditioning/renderer.py b/integrations/omnidreams/omnidreams/conditioning/renderer.py index 48dcf7576..3d959ee50 100644 --- a/integrations/omnidreams/omnidreams/conditioning/renderer.py +++ b/integrations/omnidreams/omnidreams/conditioning/renderer.py @@ -21,6 +21,7 @@ from __future__ import annotations +import os import time from pathlib import Path from typing import Literal @@ -48,7 +49,27 @@ from omnidreams.conditioning.world_scenario.pinhole import PinholeCamera +def _env_flag(name: str, default: bool = False) -> bool: + value = os.environ.get(name) + if value is None: + return default + return value.strip().lower() in {"1", "true", "yes", "on"} + + +def _env_int(name: str, default: int) -> int: + value = os.environ.get(name) + if value is None: + return default + return int(value.strip()) + + class LudusRenderer: + @staticmethod + def _to_int(value: torch.Tensor | int | float) -> int: + if isinstance(value, torch.Tensor): + return int(value.detach().cpu().item()) + return int(value) + @staticmethod def to_ludus_camera( camera: PinholeCamera | FThetaCamera, @@ -121,8 +142,15 @@ def __init__( ) self.ctx = LudusCudaTimestampedContext(device=self.device) + self.ctx.enable_render_profiling = _env_flag( + "OMNIDREAMS_LUDUS_RENDER_PROFILE" + ) + self.ctx.enable_render_profile_cuda_events = _env_flag( + "OMNIDREAMS_LUDUS_RENDER_PROFILE_CUDA_EVENTS" + ) self.ctx.set_depth_scaling(True) - self.ctx.set_msaa_samples(4) + self.msaa_samples = _env_int("OMNIDREAMS_LUDUS_MSAA_SAMPLES", 4) + self.ctx.set_msaa_samples(self.msaa_samples) self.ctx.set_max_tessellation_levels(cube=0) # Create and upload cameras @@ -143,6 +171,7 @@ def __init__( scene = self.scene_data.metadata["ludus_scene"] self._base_timestamped_scene = scene.timestamped_scene self.scene_id = self.ctx.upload_scene(scene.timestamped_scene) + self.last_render_stats: dict[str, float] = {} def _scene_id_for_dynamic_actor_pool( self, dynamic_actor_pool: CubePool | None @@ -183,8 +212,19 @@ def render_all_frames_and_cameras( n_frames = len(frame_timestamps_us) assert n_frames > 0, "Number of frames must be greater than 0" + total_t0 = time.perf_counter() + stats: dict[str, float] = { + "renderer_num_frames": float(n_frames), + "renderer_num_cameras": float(n_cameras), + "renderer_msaa_samples": float(self.msaa_samples), + } + # Create batch tensors + scene_id_t0 = time.perf_counter() scene_id = self._scene_id_for_dynamic_actor_pool(dynamic_actor_pool) + stats["renderer_scene_id_ms"] = (time.perf_counter() - scene_id_t0) * 1e3 + + batch_setup_t0 = time.perf_counter() scene_id_batch = torch.full( (n_frames * n_cameras,), scene_id, @@ -200,12 +240,16 @@ def render_all_frames_and_cameras( timestamps_batch = torch.tensor( frame_timestamps_us, dtype=torch.int64, device=self.device ).repeat(n_cameras) + stats["renderer_batch_setup_ms"] = ( + time.perf_counter() - batch_setup_t0 + ) * 1e3 H, W = None, None camera_id_batch = [] camera_poses_batch = [] + camera_pose_t0 = time.perf_counter() for camera_name in camera_names: # Get camera ID, model and check resolution c = self.all_camera_map[camera_name] @@ -229,7 +273,11 @@ def render_all_frames_and_cameras( "Camera poses must have the same length as frame timestamps" ) camera_poses_batch.append(self.to_ludus_camera_pose(poses)) + stats["renderer_camera_pose_ms"] = ( + time.perf_counter() - camera_pose_t0 + ) * 1e3 + batch_finalize_t0 = time.perf_counter() camera_id_batch = ( torch.tensor(camera_id_batch, dtype=torch.int32, device=self.device) .unsqueeze(1) @@ -237,9 +285,15 @@ def render_all_frames_and_cameras( .flatten() ) camera_poses_batch = torch.stack(camera_poses_batch, dim=0).reshape(-1, 4, 4) + stats["renderer_batch_finalize_ms"] = ( + time.perf_counter() - batch_finalize_t0 + ) * 1e3 assert H is not None and W is not None, "No cameras provided" + stats["renderer_height"] = float(self._to_int(H)) + stats["renderer_width"] = float(self._to_int(W)) + ctx_render_t0 = time.perf_counter() images = self.ctx.render( scene_id_batch, camera_id_batch, @@ -248,13 +302,22 @@ def render_all_frames_and_cameras( camera_poses_batch, resolution=(H, W), # ty:ignore[invalid-argument-type] ) + stats["renderer_ctx_render_ms"] = (time.perf_counter() - ctx_render_t0) * 1e3 + stats.update(getattr(self.ctx, "last_render_profile", {})) + output_layout_t0 = time.perf_counter() rgb = images[:, :, :, :3] if self.ctx.needs_vflip: rgb = rgb.flip(1) # rgb is [N, H, W, 3] where N = n_cameras * n_frames. # Rearrange to [n_cameras, n_frames, 3, H, W]. - return rgb.permute(0, 3, 1, 2).contiguous().view(n_cameras, n_frames, 3, H, W) # ty:ignore[invalid-argument-type] + output = rgb.permute(0, 3, 1, 2).contiguous().view(n_cameras, n_frames, 3, H, W) # ty:ignore[invalid-argument-type] + stats["renderer_output_layout_ms"] = ( + time.perf_counter() - output_layout_t0 + ) * 1e3 + stats["renderer_total_ms"] = (time.perf_counter() - total_t0) * 1e3 + self.last_render_stats = stats + return output def cleanup(self) -> None: """Cleanup the renderer.""" diff --git a/integrations/omnidreams/omnidreams/webrtc/session.py b/integrations/omnidreams/omnidreams/webrtc/session.py index 1b8495f6d..37840d63f 100644 --- a/integrations/omnidreams/omnidreams/webrtc/session.py +++ b/integrations/omnidreams/omnidreams/webrtc/session.py @@ -74,6 +74,33 @@ WEBRTC_SCENE_IMAGE_SUFFIXES = SCENE_IMAGE_SUFFIXES +def _numeric_profile_values( + stats: dict[str, float] | None, + *, + prefix: str = "", +) -> dict[str, float]: + if not stats: + return {} + + numeric_stats: dict[str, float] = {} + for key, value in stats.items(): + try: + numeric_value = float(value) + except (TypeError, ValueError): + continue + if not np.isfinite(numeric_value): + continue + numeric_stats[f"{prefix}{key}"] = numeric_value + return numeric_stats + + +def _rounded_profile_values(stats: dict[str, float] | None) -> dict[str, float]: + return { + key: round(value, 1) + for key, value in _numeric_profile_values(stats).items() + } + + def _choose_existing_asset( directory: Path, *, @@ -738,6 +765,8 @@ def _generate_one_chunk_sync( f"Chunk={self.autoregressive_index} received empty segments." ) + total_t0 = time.perf_counter() + pose_t0 = total_t0 ego_poses = self.pose_integrator.integrate_chunk( segments=segments, frame_times=frame_times ) @@ -750,6 +779,8 @@ def _generate_one_chunk_sync( camera_names = [self.config.camera_name] camera_poses_per_view = {self.config.camera_name: camera_poses} serve_hdmaps = self.config.debug_serve_hdmaps + pose_ms = (time.perf_counter() - pose_t0) * 1e3 + wrapper_t0 = time.perf_counter() if self._state is None: output = self._wrapper.start_generation( text_prompts=self._text_prompts, @@ -771,24 +802,48 @@ def _generate_one_chunk_sync( ) self._state = output.state + wrapper_ms = (time.perf_counter() - wrapper_t0) * 1e3 + finalize_ms = 0.0 + pipeline_stats: dict[str, float] | None = None if self._state.pipeline_cache is not None: - self._wrapper.finalize_block_generation( + finalize_t0 = time.perf_counter() + pipeline_stats = self._wrapper.finalize_block_generation( self._state.pipeline_cache, output.finalization_state, ) + finalize_ms = (time.perf_counter() - finalize_t0) * 1e3 + select_output_t0 = time.perf_counter() if serve_hdmaps: video_chunk = output.condition_frames elif output.rgb_frames is None: raise OmnidreamsRuntimeError("Omnidreams WebRTC received no RGB frames.") else: video_chunk = output.rgb_frames + num_output_frames = int(video_chunk.shape[2]) + select_output_ms = (time.perf_counter() - select_output_t0) * 1e3 + + detach_cpu_t0 = time.perf_counter() + video_chunk_cpu = video_chunk.detach().cpu() + detach_cpu_ms = (time.perf_counter() - detach_cpu_t0) * 1e3 + total_ms = (time.perf_counter() - total_t0) * 1e3 + + stats = { + "pose_ms": pose_ms, + "wrapper_ms": wrapper_ms, + "finalize_ms": finalize_ms, + "select_output_ms": select_output_ms, + "detach_cpu_ms": detach_cpu_ms, + "total_ms": total_ms, + } + stats.update(_numeric_profile_values(getattr(output, "stats", None))) + stats.update(_numeric_profile_values(pipeline_stats, prefix="pipeline_")) result = OmnidreamsStepResult( chunk_index=self.autoregressive_index, - num_frames=int(video_chunk.shape[2]), - video_chunk=video_chunk.detach().cpu(), - stats=None, + num_frames=num_output_frames, + video_chunk=video_chunk_cpu, + stats=stats, ) self.autoregressive_index += 1 return result @@ -1245,10 +1300,11 @@ async def _generation_worker( if consumed_action_arrivals else None ) + profile = _rounded_profile_values(result.stats) logger.info( "Chunk done chunk={} num_frames={} segments={} " "enqueued={} gen_ms={:.1f} enqueue_ms={:.1f} play_ms={:.1f} " - "queue_depth={} lag_ms={:.1f}", + "queue_depth={} lag_ms={:.1f} profile={}", result.chunk_index, result.num_frames, len(segments), @@ -1258,6 +1314,7 @@ async def _generation_worker( play_ms, video_track.qsize(), lag_ms, + profile, ) channel = managed_session.control_channel @@ -1282,6 +1339,8 @@ async def _generation_worker( "queue_depth": video_track.qsize(), "lag_ms": round(lag_ms, 1), } + if profile: + payload["profile"] = profile if control_latency_ms is not None: payload["latency_ms"] = round(control_latency_ms, 1) payload["control_latency_ms"] = round(control_latency_ms, 1) diff --git a/integrations/omnidreams/scripts/summarize_webrtc_profile.py b/integrations/omnidreams/scripts/summarize_webrtc_profile.py new file mode 100644 index 000000000..e81e5424b --- /dev/null +++ b/integrations/omnidreams/scripts/summarize_webrtc_profile.py @@ -0,0 +1,149 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import argparse +import ast +import re +import statistics as st +from pathlib import Path + + +CHUNK_RE = re.compile( + r"Chunk done chunk=(?P\d+).*?" + r"gen_ms=(?P[0-9.]+).*?" + r"enqueue_ms=(?P[0-9.]+).*?" + r"play_ms=(?P[0-9.]+).*?" + r"queue_depth=(?P\d+).*?" + r"lag_ms=(?P[0-9.]+).*?" + r"profile=(?P\{.*\})" +) + +DEFAULT_KEYS = [ + "gen_ms", + "wrapper_render_condition_ms", + "renderer_ctx_render_ms", + "ctx_render_plugin_cuda_ms_sum", + "ctx_render_plugin_cuda_ms_avg", + "pipeline_total_ms", + "pipeline_total_ms_wo_finalize", + "enqueue_ms", + "total_ms", + "play_ms", + "lag_ms", +] + + +def percentile(values: list[float], q: float) -> float: + if not values: + raise ValueError("percentile requires at least one value") + ordered = sorted(values) + index = int(q * (len(ordered) - 1)) + return ordered[index] + + +def load_rows(path: Path) -> list[dict[str, float]]: + rows: list[dict[str, float]] = [] + session = -1 + prev_chunk: int | None = None + for line in path.read_text(errors="replace").splitlines(): + match = CHUNK_RE.search(line) + if match is None: + continue + chunk = int(match.group("chunk")) + if prev_chunk is None or chunk <= prev_chunk: + session += 1 + prev_chunk = chunk + profile = ast.literal_eval(match.group("profile")) + row = { + "session": float(session), + "chunk": float(chunk), + "gen_ms": float(match.group("gen_ms")), + "enqueue_ms": float(match.group("enqueue_ms")), + "play_ms": float(match.group("play_ms")), + "queue_depth": float(match.group("queue_depth")), + "lag_ms": float(match.group("lag_ms")), + } + row.update({key: float(value) for key, value in profile.items()}) + rows.append(row) + return rows + + +def print_stats(rows: list[dict[str, float]], keys: list[str]) -> None: + for key in keys: + values = [row[key] for row in rows if key in row] + if not values: + continue + print( + f"{key:34s} avg={st.mean(values):7.1f} " + f"p50={st.median(values):7.1f} p90={percentile(values, 0.9):7.1f} " + f"min={min(values):7.1f} max={max(values):7.1f}" + ) + + +def print_top_profile_keys(rows: list[dict[str, float]], limit: int) -> None: + skip = {"session", "chunk", "queue_depth", "renderer_height", "renderer_width"} + means = [] + keys = sorted(set().union(*(row.keys() for row in rows))) + for key in keys: + if key in skip or key.endswith("_count"): + continue + if not key.endswith("_ms") and not key.endswith("_ms_sum"): + continue + values = [row[key] for row in rows if key in row] + if values: + means.append((st.mean(values), key)) + for mean, key in sorted(means, reverse=True)[:limit]: + print(f"{key:34s} avg={mean:7.1f}") + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Summarize OmniDreams WebRTC chunk profile logs." + ) + parser.add_argument("log", type=Path) + parser.add_argument("--min-chunk", type=int, default=4) + parser.add_argument( + "--session", + default="latest", + help="'latest', 'all', or a numeric session index inferred from chunk resets.", + ) + parser.add_argument("--top", type=int, default=12) + args = parser.parse_args() + + rows = load_rows(args.log) + print(f"parsed_chunks={len(rows)}") + if not rows: + return + + sessions = sorted({int(row["session"]) for row in rows}) + print(f"sessions={sessions}") + if args.session == "latest": + wanted_sessions = {sessions[-1]} + elif args.session == "all": + wanted_sessions = set(sessions) + else: + wanted_sessions = {int(args.session)} + + selected = [ + row + for row in rows + if int(row["session"]) in wanted_sessions and row["chunk"] >= args.min_chunk + ] + print( + f"selected_chunks={len(selected)} " + f"session={args.session} min_chunk={args.min_chunk}" + ) + if not selected: + return + + print("\nprimary stats") + print_stats(selected, DEFAULT_KEYS) + print("\ntop profile timings") + print_top_profile_keys(selected, args.top) + + +if __name__ == "__main__": + main()