From 3759a81e8093b8532bf3e9f1f021768dd5d9981b Mon Sep 17 00:00:00 2001 From: Junchen Liu Date: Fri, 29 May 2026 07:49:45 -0700 Subject: [PATCH 1/2] add profiling script & repro md --- integrations/omnidreams/README.md | 3 + .../docs/issue195_renderer_profile_repro.md | 105 ++++++++++++ .../ludus_renderer/_ops/context.py | 90 ++++++++++- .../tests/test_context_render_profile.py | 67 ++++++++ .../conditioning/conditioning_wrapper.py | 81 +++++++++- .../omnidreams/conditioning/renderer.py | 67 +++++++- .../omnidreams/omnidreams/webrtc/session.py | 69 +++++++- .../scripts/summarize_webrtc_profile.py | 149 ++++++++++++++++++ .../tests/test_conditioning_wrapper.py | 26 +++ .../omnidreams/tests/test_ludus_renderer.py | 116 ++++++++++++++ .../omnidreams/tests/test_webrtc_runtime.py | 109 ++++++++++++- 11 files changed, 868 insertions(+), 14 deletions(-) create mode 100644 integrations/omnidreams/docs/issue195_renderer_profile_repro.md create mode 100644 integrations/omnidreams/ludus-renderer/tests/test_context_render_profile.py create mode 100644 integrations/omnidreams/scripts/summarize_webrtc_profile.py 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/ludus-renderer/tests/test_context_render_profile.py b/integrations/omnidreams/ludus-renderer/tests/test_context_render_profile.py new file mode 100644 index 000000000..3127cb475 --- /dev/null +++ b/integrations/omnidreams/ludus-renderer/tests/test_context_render_profile.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest +import torch + + +@pytest.mark.ci_cpu +def test_context_render_records_internal_profile(monkeypatch) -> None: + from ludus_renderer._ops import context + + class _FakePlugin: + def __init__(self) -> None: + self.calls = 0 + + def ludus_render_fwd_cuda_timestamped(self, *args, **kwargs) -> torch.Tensor: + del args, kwargs + self.calls += 1 + return torch.full((1, 2, 3, 4), self.calls, dtype=torch.uint8) + + fake_plugin = _FakePlugin() + monkeypatch.setattr(context, "_get_plugin", lambda: fake_plugin) + + ctx = context.LudusCudaTimestampedContext.__new__( + context.LudusCudaTimestampedContext + ) + ctx._camera_intrinsics = torch.zeros((1, 4), dtype=torch.float32) + ctx._max_extrapolation_us = 500_000 + ctx._tessellation_threshold = 1.0 + ctx.cpp_wrapper = object() + ctx.enable_render_profiling = True + ctx.enable_render_profile_cuda_events = False + empty = torch.empty(0) + ctx._scenes = [ + { + "timestamps": empty, + "int32": empty, + "vertices": empty, + "triangles": empty, + "floats": empty, + "polyline_pools": empty, + "polygon_pools": empty, + "cube_pools": empty, + "max_varrays_per_ts_polyline": 0, + "max_varrays_per_ts_polygon": 0, + } + ] + + image = ctx.render( + scene_ids=torch.zeros(2, dtype=torch.int32), + camera_ids=torch.zeros(2, dtype=torch.int32), + timestamps_us=torch.tensor([1_000_000, 1_033_333], dtype=torch.int64), + camera_type_ids=torch.zeros(2, dtype=torch.int32), + camera_poses=torch.eye(4).repeat(2, 1, 1), + resolution=(2, 3), + ) + + assert image.shape == (2, 2, 3, 4) + assert fake_plugin.calls == 2 + assert ctx.last_render_profile["ctx_render_batch_size"] == 2.0 + assert ctx.last_render_profile["ctx_render_scalar_item_host_ms_count"] == 2.0 + assert ctx.last_render_profile["ctx_render_query_prep_host_ms_count"] == 2.0 + assert ctx.last_render_profile["ctx_render_plugin_host_ms_count"] == 2.0 + assert "ctx_render_cat_host_ms" in ctx.last_render_profile + assert "ctx_render_total_host_ms" in ctx.last_render_profile 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() diff --git a/integrations/omnidreams/tests/test_conditioning_wrapper.py b/integrations/omnidreams/tests/test_conditioning_wrapper.py index caffbe8fb..fa01070ee 100644 --- a/integrations/omnidreams/tests/test_conditioning_wrapper.py +++ b/integrations/omnidreams/tests/test_conditioning_wrapper.py @@ -80,6 +80,7 @@ def __init__(self, *, height: int, width: int) -> None: self.height = height self.width = width self.cleaned_up = False + self.last_render_stats: dict[str, float] = {} def render_all_frames_and_cameras( self, @@ -92,6 +93,10 @@ def render_all_frames_and_cameras( del camera_poses_per_view, dynamic_actor_pool n_views = len(camera_names) n_frames = len(frame_timestamps_us) + self.last_render_stats = { + "renderer_ctx_render_ms": 12.0, + "renderer_output_layout_ms": 3.0, + } # [V, T, 3, H, W] return torch.full( (n_views, n_frames, 3, self.height, self.width), @@ -153,6 +158,17 @@ def test_start_continue_and_finalize_flow() -> None: assert start_output.finalization_state == {"autoregressive_index": 0} assert start_output.state.pipeline_cache is not None assert start_output.state.pipeline_cache.autoregressive_index == 0 + assert start_output.stats is not None + assert { + "wrapper_render_condition_ms", + "renderer_ctx_render_ms", + "renderer_output_layout_ms", + "wrapper_prepare_condition_ms", + "wrapper_initialize_cache_ms", + "wrapper_pipeline_generate_ms", + "wrapper_to_uint8_ms", + "wrapper_total_ms", + } <= start_output.stats.keys() continue_output = wrapper.continue_generation( state=start_output.state, @@ -167,6 +183,16 @@ def test_start_continue_and_finalize_flow() -> None: assert continue_output.finalization_state == {"autoregressive_index": 1} assert continue_output.state.pipeline_cache is not None assert continue_output.state.pipeline_cache.autoregressive_index == 1 + assert continue_output.stats is not None + assert { + "wrapper_render_condition_ms", + "renderer_ctx_render_ms", + "renderer_output_layout_ms", + "wrapper_prepare_condition_ms", + "wrapper_pipeline_generate_ms", + "wrapper_to_uint8_ms", + "wrapper_total_ms", + } <= continue_output.stats.keys() assert continue_output.state.pipeline_cache is not None wrapper.finalize_block_generation( diff --git a/integrations/omnidreams/tests/test_ludus_renderer.py b/integrations/omnidreams/tests/test_ludus_renderer.py index 9828fdd06..555c4b5be 100644 --- a/integrations/omnidreams/tests/test_ludus_renderer.py +++ b/integrations/omnidreams/tests/test_ludus_renderer.py @@ -29,6 +29,7 @@ import tempfile import zipfile from pathlib import Path +from types import SimpleNamespace import pytest import torch @@ -137,6 +138,121 @@ def test_ludus_cuda_context_renders_frame(clipgt_scene_dir: Path) -> None: # --------------------------------------------------------------------------- +@pytest.mark.ci_cpu +@pytest.mark.parametrize(("env_value", "expected"), [("0", 0), ("4", 4)]) +def test_ludus_renderer_msaa_samples_can_be_configured( + monkeypatch: pytest.MonkeyPatch, env_value: str, expected: int +) -> None: + from omnidreams.conditioning import renderer as renderer_module + + class _FakeCtx: + instances = [] + + def __init__(self, device: torch.device) -> None: + self.device = device + self.msaa_samples = [] + _FakeCtx.instances.append(self) + + def set_depth_scaling(self, enabled: bool) -> None: + self.depth_scaling = enabled + + def set_msaa_samples(self, samples: int) -> None: + self.msaa_samples.append(samples) + + def set_max_tessellation_levels(self, **kwargs: int) -> None: + self.max_tessellation_levels = kwargs + + def upload_cameras(self, cameras: list[object]) -> None: + self.cameras = cameras + + def upload_scene(self, scene: object) -> int: + self.scene = scene + return 7 + + monkeypatch.setenv("OMNIDREAMS_LUDUS_MSAA_SAMPLES", env_value) + monkeypatch.setattr(renderer_module, "LudusCudaTimestampedContext", _FakeCtx) + monkeypatch.setattr( + renderer_module.LudusRenderer, + "to_ludus_camera", + staticmethod(lambda camera: camera), + ) + + scene_data = SimpleNamespace( + metadata={"ludus_scene": SimpleNamespace(timestamped_scene=object())} + ) + + renderer_module.LudusRenderer( + scene_data=scene_data, + camera_models={"front": SimpleNamespace()}, + device=torch.device("cpu"), + ) + + assert _FakeCtx.instances[-1].msaa_samples == [expected] + + +@pytest.mark.ci_cpu +def test_ludus_renderer_records_render_substage_stats_with_cpu_fake_ctx() -> None: + from omnidreams.conditioning.renderer import LudusRenderer + + class _FakeCtx: + needs_vflip = False + last_render_profile = { + "ctx_render_scalar_item_host_ms_sum": 2.0, + "ctx_render_plugin_host_ms_sum": 5.0, + "ctx_render_cat_host_ms": 1.0, + } + + def render( + self, + scene_id_batch: torch.Tensor, + camera_id_batch: torch.Tensor, + timestamps_batch: torch.Tensor, + camera_type_id_batch: torch.Tensor, + camera_poses_batch: torch.Tensor, + *, + resolution: tuple[int, int], + ) -> torch.Tensor: + del camera_id_batch, timestamps_batch, camera_type_id_batch + assert tuple(scene_id_batch.shape) == (2,) + assert tuple(camera_poses_batch.shape) == (2, 4, 4) + height, width = resolution + return torch.zeros((2, height, width, 4), dtype=torch.uint8) + + renderer = LudusRenderer.__new__(LudusRenderer) + renderer.device = torch.device("cpu") + renderer.scene_id = 7 + renderer.msaa_samples = 0 + renderer.ctx = _FakeCtx() + renderer.all_camera_map = {"front": 0} + renderer.all_cameras = [ + SimpleNamespace(image_size=torch.tensor([4, 3], dtype=torch.int64)) + ] + + output = renderer.render_all_frames_and_cameras( + camera_names=["front"], + camera_poses_per_camera={"front": torch.eye(4).repeat(2, 1, 1)}, + frame_timestamps_us=[1_000_000, 1_033_333], + ) + + assert output.shape == (1, 2, 3, 3, 4) + assert { + "renderer_scene_id_ms", + "renderer_batch_setup_ms", + "renderer_camera_pose_ms", + "renderer_ctx_render_ms", + "renderer_output_layout_ms", + "renderer_total_ms", + "renderer_num_frames", + "renderer_num_cameras", + "renderer_msaa_samples", + "renderer_height", + "renderer_width", + "ctx_render_scalar_item_host_ms_sum", + "ctx_render_plugin_host_ms_sum", + "ctx_render_cat_host_ms", + } <= renderer.last_render_stats.keys() + + @pytest.mark.manual @pytest.mark.parametrize( "n_frames", [1, 2, 3], ids=["single-frame", "two-frame", "multi-frame"] diff --git a/integrations/omnidreams/tests/test_webrtc_runtime.py b/integrations/omnidreams/tests/test_webrtc_runtime.py index f36c1ec34..1df0d77d3 100644 --- a/integrations/omnidreams/tests/test_webrtc_runtime.py +++ b/integrations/omnidreams/tests/test_webrtc_runtime.py @@ -5,6 +5,7 @@ import argparse import asyncio +import json import sys import zipfile from dataclasses import dataclass @@ -97,9 +98,16 @@ def continue_generation(self, **kwargs: Any) -> _FakeOutput: def finalize_block_generation( self, pipeline_cache: object, finalization_state: dict[str, int] - ) -> None: + ) -> dict[str, float]: del pipeline_cache self.finalized.append(finalization_state) + return { + "encode_ms": 1.0, + "diffuse_ms": 2.0, + "decode_ms": 3.0, + "finalize_ms": 4.0, + "total_ms": 10.0, + } def _build_fake_runtime() -> tuple[OmnidreamsInferenceRuntime, _FakeWrapper]: @@ -142,6 +150,18 @@ def test_generate_chunk_dispatches_start_then_continue() -> None: assert wrapper.calls[1][1] == (3, 4, 4) assert len(wrapper.finalized) == 2 assert wrapper.skip_video_generation_flags == [False, False] + assert result1.stats is not None + assert { + "pose_ms", + "wrapper_ms", + "finalize_ms", + "select_output_ms", + "detach_cpu_ms", + "total_ms", + "pipeline_total_ms", + } <= result1.stats.keys() + assert result1.stats["pipeline_total_ms"] == 10.0 + assert result1.stats["total_ms"] >= 0.0 def test_generate_chunk_can_stream_debug_hdmaps_without_rgb_frames() -> None: @@ -686,6 +706,93 @@ async def test_disconnect_message_closes_active_session( assert peer_connection.closed +@pytest.mark.asyncio +async def test_generation_worker_includes_profile_in_chunk_done_payload() -> None: + class _Runtime: + def __init__(self) -> None: + self.generate_calls = 0 + + def peek_next_chunk_num_frames(self) -> int: + return 1 + + async def generate_chunk( + self, + *, + segments: list[tuple[float, float, frozenset[str]]], + frame_times: list[float], + ) -> OmnidreamsStepResult: + del segments, frame_times + self.generate_calls += 1 + return OmnidreamsStepResult( + chunk_index=0, + num_frames=1, + video_chunk=torch.zeros((1, 1, 1, 3, 2, 2), dtype=torch.uint8), + stats={ + "wrapper_ms": 12.34, + "pipeline_total_ms": 10.0, + "pipeline_mem_alloc_gib": 25.445, + }, + ) + + class _Resampler: + dt = 0.0 + next_chunk_start_v = 0.0 + + def sample_chunk( + self, num_frames: int + ) -> tuple[list[tuple[float, float, frozenset[str]]], list[float]]: + assert num_frames == 1 + return [(0.0, 0.0, frozenset({"w"}))], [0.0] + + class _VideoTrack: + fps = 30 + + def __init__(self) -> None: + self.managed_session: session._ManagedOmnidreamsSession | None = None + + async def enqueue_chunk(self, video_chunk: torch.Tensor) -> int: + assert video_chunk.shape == (1, 1, 1, 3, 2, 2) + assert self.managed_session is not None + self.managed_session.closed = True + return 1 + + def qsize(self) -> int: + return 0 + + class _Channel: + def __init__(self) -> None: + self.messages: list[str] = [] + + def send(self, message: str) -> None: + self.messages.append(message) + + manager = OmnidreamsWebRTCSessionManager( + runtime_config=OmnidreamsRuntimeConfig(device="cpu", warmup_chunks=0) + ) + video_track = _VideoTrack() + channel = _Channel() + first_action_received = asyncio.Event() + first_action_received.set() + managed_session = session._ManagedOmnidreamsSession( + runtime=_Runtime(), # ty:ignore[invalid-argument-type] + video_track=video_track, # ty:ignore[invalid-argument-type] + peer_connection=_FakeCloseable(), + resampler=_Resampler(), # ty:ignore[invalid-argument-type] + control_channel=channel, + first_action_received=first_action_received, + ) + video_track.managed_session = managed_session + + await manager._generation_worker(managed_session=managed_session) + + assert len(channel.messages) == 1 + payload = json.loads(channel.messages[0]) + assert payload["type"] == "chunk_done" + assert payload["profile"]["wrapper_ms"] == 12.3 + assert payload["profile"]["pipeline_total_ms"] == 10.0 + assert payload["profile"]["pipeline_mem_alloc_gib"] == 25.4 + + @pytest.mark.asyncio async def test_generation_worker_closes_session_after_generation_failure() -> None: class _FailingRuntime: From 068867f8ed02e40ba1ea3f1598f11ef7ca9556a0 Mon Sep 17 00:00:00 2001 From: Junchen Liu Date: Fri, 29 May 2026 08:08:04 -0700 Subject: [PATCH 2/2] drop repro test changes --- .../tests/test_context_render_profile.py | 67 ---------- .../tests/test_conditioning_wrapper.py | 26 ---- .../omnidreams/tests/test_ludus_renderer.py | 116 ------------------ .../omnidreams/tests/test_webrtc_runtime.py | 109 +--------------- 4 files changed, 1 insertion(+), 317 deletions(-) delete mode 100644 integrations/omnidreams/ludus-renderer/tests/test_context_render_profile.py diff --git a/integrations/omnidreams/ludus-renderer/tests/test_context_render_profile.py b/integrations/omnidreams/ludus-renderer/tests/test_context_render_profile.py deleted file mode 100644 index 3127cb475..000000000 --- a/integrations/omnidreams/ludus-renderer/tests/test_context_render_profile.py +++ /dev/null @@ -1,67 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import pytest -import torch - - -@pytest.mark.ci_cpu -def test_context_render_records_internal_profile(monkeypatch) -> None: - from ludus_renderer._ops import context - - class _FakePlugin: - def __init__(self) -> None: - self.calls = 0 - - def ludus_render_fwd_cuda_timestamped(self, *args, **kwargs) -> torch.Tensor: - del args, kwargs - self.calls += 1 - return torch.full((1, 2, 3, 4), self.calls, dtype=torch.uint8) - - fake_plugin = _FakePlugin() - monkeypatch.setattr(context, "_get_plugin", lambda: fake_plugin) - - ctx = context.LudusCudaTimestampedContext.__new__( - context.LudusCudaTimestampedContext - ) - ctx._camera_intrinsics = torch.zeros((1, 4), dtype=torch.float32) - ctx._max_extrapolation_us = 500_000 - ctx._tessellation_threshold = 1.0 - ctx.cpp_wrapper = object() - ctx.enable_render_profiling = True - ctx.enable_render_profile_cuda_events = False - empty = torch.empty(0) - ctx._scenes = [ - { - "timestamps": empty, - "int32": empty, - "vertices": empty, - "triangles": empty, - "floats": empty, - "polyline_pools": empty, - "polygon_pools": empty, - "cube_pools": empty, - "max_varrays_per_ts_polyline": 0, - "max_varrays_per_ts_polygon": 0, - } - ] - - image = ctx.render( - scene_ids=torch.zeros(2, dtype=torch.int32), - camera_ids=torch.zeros(2, dtype=torch.int32), - timestamps_us=torch.tensor([1_000_000, 1_033_333], dtype=torch.int64), - camera_type_ids=torch.zeros(2, dtype=torch.int32), - camera_poses=torch.eye(4).repeat(2, 1, 1), - resolution=(2, 3), - ) - - assert image.shape == (2, 2, 3, 4) - assert fake_plugin.calls == 2 - assert ctx.last_render_profile["ctx_render_batch_size"] == 2.0 - assert ctx.last_render_profile["ctx_render_scalar_item_host_ms_count"] == 2.0 - assert ctx.last_render_profile["ctx_render_query_prep_host_ms_count"] == 2.0 - assert ctx.last_render_profile["ctx_render_plugin_host_ms_count"] == 2.0 - assert "ctx_render_cat_host_ms" in ctx.last_render_profile - assert "ctx_render_total_host_ms" in ctx.last_render_profile diff --git a/integrations/omnidreams/tests/test_conditioning_wrapper.py b/integrations/omnidreams/tests/test_conditioning_wrapper.py index fa01070ee..caffbe8fb 100644 --- a/integrations/omnidreams/tests/test_conditioning_wrapper.py +++ b/integrations/omnidreams/tests/test_conditioning_wrapper.py @@ -80,7 +80,6 @@ def __init__(self, *, height: int, width: int) -> None: self.height = height self.width = width self.cleaned_up = False - self.last_render_stats: dict[str, float] = {} def render_all_frames_and_cameras( self, @@ -93,10 +92,6 @@ def render_all_frames_and_cameras( del camera_poses_per_view, dynamic_actor_pool n_views = len(camera_names) n_frames = len(frame_timestamps_us) - self.last_render_stats = { - "renderer_ctx_render_ms": 12.0, - "renderer_output_layout_ms": 3.0, - } # [V, T, 3, H, W] return torch.full( (n_views, n_frames, 3, self.height, self.width), @@ -158,17 +153,6 @@ def test_start_continue_and_finalize_flow() -> None: assert start_output.finalization_state == {"autoregressive_index": 0} assert start_output.state.pipeline_cache is not None assert start_output.state.pipeline_cache.autoregressive_index == 0 - assert start_output.stats is not None - assert { - "wrapper_render_condition_ms", - "renderer_ctx_render_ms", - "renderer_output_layout_ms", - "wrapper_prepare_condition_ms", - "wrapper_initialize_cache_ms", - "wrapper_pipeline_generate_ms", - "wrapper_to_uint8_ms", - "wrapper_total_ms", - } <= start_output.stats.keys() continue_output = wrapper.continue_generation( state=start_output.state, @@ -183,16 +167,6 @@ def test_start_continue_and_finalize_flow() -> None: assert continue_output.finalization_state == {"autoregressive_index": 1} assert continue_output.state.pipeline_cache is not None assert continue_output.state.pipeline_cache.autoregressive_index == 1 - assert continue_output.stats is not None - assert { - "wrapper_render_condition_ms", - "renderer_ctx_render_ms", - "renderer_output_layout_ms", - "wrapper_prepare_condition_ms", - "wrapper_pipeline_generate_ms", - "wrapper_to_uint8_ms", - "wrapper_total_ms", - } <= continue_output.stats.keys() assert continue_output.state.pipeline_cache is not None wrapper.finalize_block_generation( diff --git a/integrations/omnidreams/tests/test_ludus_renderer.py b/integrations/omnidreams/tests/test_ludus_renderer.py index 555c4b5be..9828fdd06 100644 --- a/integrations/omnidreams/tests/test_ludus_renderer.py +++ b/integrations/omnidreams/tests/test_ludus_renderer.py @@ -29,7 +29,6 @@ import tempfile import zipfile from pathlib import Path -from types import SimpleNamespace import pytest import torch @@ -138,121 +137,6 @@ def test_ludus_cuda_context_renders_frame(clipgt_scene_dir: Path) -> None: # --------------------------------------------------------------------------- -@pytest.mark.ci_cpu -@pytest.mark.parametrize(("env_value", "expected"), [("0", 0), ("4", 4)]) -def test_ludus_renderer_msaa_samples_can_be_configured( - monkeypatch: pytest.MonkeyPatch, env_value: str, expected: int -) -> None: - from omnidreams.conditioning import renderer as renderer_module - - class _FakeCtx: - instances = [] - - def __init__(self, device: torch.device) -> None: - self.device = device - self.msaa_samples = [] - _FakeCtx.instances.append(self) - - def set_depth_scaling(self, enabled: bool) -> None: - self.depth_scaling = enabled - - def set_msaa_samples(self, samples: int) -> None: - self.msaa_samples.append(samples) - - def set_max_tessellation_levels(self, **kwargs: int) -> None: - self.max_tessellation_levels = kwargs - - def upload_cameras(self, cameras: list[object]) -> None: - self.cameras = cameras - - def upload_scene(self, scene: object) -> int: - self.scene = scene - return 7 - - monkeypatch.setenv("OMNIDREAMS_LUDUS_MSAA_SAMPLES", env_value) - monkeypatch.setattr(renderer_module, "LudusCudaTimestampedContext", _FakeCtx) - monkeypatch.setattr( - renderer_module.LudusRenderer, - "to_ludus_camera", - staticmethod(lambda camera: camera), - ) - - scene_data = SimpleNamespace( - metadata={"ludus_scene": SimpleNamespace(timestamped_scene=object())} - ) - - renderer_module.LudusRenderer( - scene_data=scene_data, - camera_models={"front": SimpleNamespace()}, - device=torch.device("cpu"), - ) - - assert _FakeCtx.instances[-1].msaa_samples == [expected] - - -@pytest.mark.ci_cpu -def test_ludus_renderer_records_render_substage_stats_with_cpu_fake_ctx() -> None: - from omnidreams.conditioning.renderer import LudusRenderer - - class _FakeCtx: - needs_vflip = False - last_render_profile = { - "ctx_render_scalar_item_host_ms_sum": 2.0, - "ctx_render_plugin_host_ms_sum": 5.0, - "ctx_render_cat_host_ms": 1.0, - } - - def render( - self, - scene_id_batch: torch.Tensor, - camera_id_batch: torch.Tensor, - timestamps_batch: torch.Tensor, - camera_type_id_batch: torch.Tensor, - camera_poses_batch: torch.Tensor, - *, - resolution: tuple[int, int], - ) -> torch.Tensor: - del camera_id_batch, timestamps_batch, camera_type_id_batch - assert tuple(scene_id_batch.shape) == (2,) - assert tuple(camera_poses_batch.shape) == (2, 4, 4) - height, width = resolution - return torch.zeros((2, height, width, 4), dtype=torch.uint8) - - renderer = LudusRenderer.__new__(LudusRenderer) - renderer.device = torch.device("cpu") - renderer.scene_id = 7 - renderer.msaa_samples = 0 - renderer.ctx = _FakeCtx() - renderer.all_camera_map = {"front": 0} - renderer.all_cameras = [ - SimpleNamespace(image_size=torch.tensor([4, 3], dtype=torch.int64)) - ] - - output = renderer.render_all_frames_and_cameras( - camera_names=["front"], - camera_poses_per_camera={"front": torch.eye(4).repeat(2, 1, 1)}, - frame_timestamps_us=[1_000_000, 1_033_333], - ) - - assert output.shape == (1, 2, 3, 3, 4) - assert { - "renderer_scene_id_ms", - "renderer_batch_setup_ms", - "renderer_camera_pose_ms", - "renderer_ctx_render_ms", - "renderer_output_layout_ms", - "renderer_total_ms", - "renderer_num_frames", - "renderer_num_cameras", - "renderer_msaa_samples", - "renderer_height", - "renderer_width", - "ctx_render_scalar_item_host_ms_sum", - "ctx_render_plugin_host_ms_sum", - "ctx_render_cat_host_ms", - } <= renderer.last_render_stats.keys() - - @pytest.mark.manual @pytest.mark.parametrize( "n_frames", [1, 2, 3], ids=["single-frame", "two-frame", "multi-frame"] diff --git a/integrations/omnidreams/tests/test_webrtc_runtime.py b/integrations/omnidreams/tests/test_webrtc_runtime.py index 1df0d77d3..f36c1ec34 100644 --- a/integrations/omnidreams/tests/test_webrtc_runtime.py +++ b/integrations/omnidreams/tests/test_webrtc_runtime.py @@ -5,7 +5,6 @@ import argparse import asyncio -import json import sys import zipfile from dataclasses import dataclass @@ -98,16 +97,9 @@ def continue_generation(self, **kwargs: Any) -> _FakeOutput: def finalize_block_generation( self, pipeline_cache: object, finalization_state: dict[str, int] - ) -> dict[str, float]: + ) -> None: del pipeline_cache self.finalized.append(finalization_state) - return { - "encode_ms": 1.0, - "diffuse_ms": 2.0, - "decode_ms": 3.0, - "finalize_ms": 4.0, - "total_ms": 10.0, - } def _build_fake_runtime() -> tuple[OmnidreamsInferenceRuntime, _FakeWrapper]: @@ -150,18 +142,6 @@ def test_generate_chunk_dispatches_start_then_continue() -> None: assert wrapper.calls[1][1] == (3, 4, 4) assert len(wrapper.finalized) == 2 assert wrapper.skip_video_generation_flags == [False, False] - assert result1.stats is not None - assert { - "pose_ms", - "wrapper_ms", - "finalize_ms", - "select_output_ms", - "detach_cpu_ms", - "total_ms", - "pipeline_total_ms", - } <= result1.stats.keys() - assert result1.stats["pipeline_total_ms"] == 10.0 - assert result1.stats["total_ms"] >= 0.0 def test_generate_chunk_can_stream_debug_hdmaps_without_rgb_frames() -> None: @@ -706,93 +686,6 @@ async def test_disconnect_message_closes_active_session( assert peer_connection.closed -@pytest.mark.asyncio -async def test_generation_worker_includes_profile_in_chunk_done_payload() -> None: - class _Runtime: - def __init__(self) -> None: - self.generate_calls = 0 - - def peek_next_chunk_num_frames(self) -> int: - return 1 - - async def generate_chunk( - self, - *, - segments: list[tuple[float, float, frozenset[str]]], - frame_times: list[float], - ) -> OmnidreamsStepResult: - del segments, frame_times - self.generate_calls += 1 - return OmnidreamsStepResult( - chunk_index=0, - num_frames=1, - video_chunk=torch.zeros((1, 1, 1, 3, 2, 2), dtype=torch.uint8), - stats={ - "wrapper_ms": 12.34, - "pipeline_total_ms": 10.0, - "pipeline_mem_alloc_gib": 25.445, - }, - ) - - class _Resampler: - dt = 0.0 - next_chunk_start_v = 0.0 - - def sample_chunk( - self, num_frames: int - ) -> tuple[list[tuple[float, float, frozenset[str]]], list[float]]: - assert num_frames == 1 - return [(0.0, 0.0, frozenset({"w"}))], [0.0] - - class _VideoTrack: - fps = 30 - - def __init__(self) -> None: - self.managed_session: session._ManagedOmnidreamsSession | None = None - - async def enqueue_chunk(self, video_chunk: torch.Tensor) -> int: - assert video_chunk.shape == (1, 1, 1, 3, 2, 2) - assert self.managed_session is not None - self.managed_session.closed = True - return 1 - - def qsize(self) -> int: - return 0 - - class _Channel: - def __init__(self) -> None: - self.messages: list[str] = [] - - def send(self, message: str) -> None: - self.messages.append(message) - - manager = OmnidreamsWebRTCSessionManager( - runtime_config=OmnidreamsRuntimeConfig(device="cpu", warmup_chunks=0) - ) - video_track = _VideoTrack() - channel = _Channel() - first_action_received = asyncio.Event() - first_action_received.set() - managed_session = session._ManagedOmnidreamsSession( - runtime=_Runtime(), # ty:ignore[invalid-argument-type] - video_track=video_track, # ty:ignore[invalid-argument-type] - peer_connection=_FakeCloseable(), - resampler=_Resampler(), # ty:ignore[invalid-argument-type] - control_channel=channel, - first_action_received=first_action_received, - ) - video_track.managed_session = managed_session - - await manager._generation_worker(managed_session=managed_session) - - assert len(channel.messages) == 1 - payload = json.loads(channel.messages[0]) - assert payload["type"] == "chunk_done" - assert payload["profile"]["wrapper_ms"] == 12.3 - assert payload["profile"]["pipeline_total_ms"] == 10.0 - assert payload["profile"]["pipeline_mem_alloc_gib"] == 25.4 - - @pytest.mark.asyncio async def test_generation_worker_closes_session_after_generation_failure() -> None: class _FailingRuntime: