Skip to content

Commit 9297faf

Browse files
authored
Avoid disk I/O when preparing USD stage for OVRTX renderer (#5631)
# Description Avoid disk I/O when preparing USD stage for OVRTX renderer ## Type of change - Bug fix (non-breaking change which fixes an issue) ## Checklist - [x] I have read and understood the [contribution guidelines](https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html) - [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with `./isaaclab.sh --format` - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] I have updated the changelog and the corresponding version in the extension's `config/extension.toml` file - [x] I have added my name to the `CONTRIBUTORS.md` or my name already exists there
1 parent bed2bf9 commit 9297faf

4 files changed

Lines changed: 80 additions & 73 deletions

File tree

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
Changed
2+
^^^^^^^
3+
4+
* Changed :attr:`~isaaclab_ov.renderers.OVRTXRendererCfg.temp_usd_dir` defaults to ``None``. Set it to a writable
5+
directory when you want the combined stage written to disk for debugging.
6+
7+
Removed
8+
^^^^^^^
9+
10+
* Removed :attr:`~isaaclab_ov.renderers.OVRTXRendererCfg.temp_usd_suffix`. When a temp file is written, the renderer
11+
uses ``ovrtx_renderer_stage.usda`` filename under the configured temp directory.
12+
13+
Fixed
14+
^^^^^
15+
16+
* Avoided OVRTX staging disk I/O by exporting the prepared USD to memory and loading it with ``open_usd_from_string``
17+
instead of always writing intermediate scene and combined USD files.

source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py

Lines changed: 38 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
2222
import logging
2323
import math
2424
import os
25+
import tempfile
26+
from pathlib import Path
2527
from typing import TYPE_CHECKING, Any
2628

2729
logger = logging.getLogger(__name__)
@@ -58,9 +60,9 @@
5860
sync_newton_transforms_kernel,
5961
)
6062
from .ovrtx_usd import (
63+
build_render_product_as_string,
6164
create_scene_partition_attributes,
62-
export_stage_for_ovrtx,
63-
inject_cameras_into_usd,
65+
export_stage_to_string,
6466
)
6567

6668
if TYPE_CHECKING:
@@ -164,7 +166,7 @@ def __init__(self, cfg: OVRTXRendererCfg):
164166
self._object_binding = None
165167
self._object_newton_indices: wp.array | None = None
166168
self._initialized_scene = False
167-
self._exported_usd_path: str | None = None
169+
self._exported_usd_string: str | None = None
168170
self._camera_rel_path: str | None = None
169171
self._output_semantic_color_buffer: wp.array | None = None
170172

@@ -192,21 +194,18 @@ def __init__(self, cfg: OVRTXRendererCfg):
192194
logger.info("OVRTX renderer created successfully")
193195

194196
def prepare_stage(self, stage: Any, num_envs: int) -> None:
195-
"""Export the USD stage for OVRTX before create_render_data.
197+
"""Prepare the USD stage for OVRTX before :meth:`create_render_data`.
196198
197-
Adds cloning attributes and exports the stage to a temporary file.
198-
The exported path is used by create_render_data when loading into OVRTX.
199+
Adds cloning attributes and exports the stage to a string held on the renderer until
200+
:meth:`create_render_data` is called.
199201
"""
200202
if stage is None:
201203
return
202204

203205
logger.info("Preparing stage for export (%d envs, cloning=%s)...", num_envs, self._use_ovrtx_cloning)
204206
create_scene_partition_attributes(stage, num_envs, self._use_ovrtx_cloning, not _IS_OVRTX_0_3_0_OR_NEWER)
205207

206-
export_path = "/tmp/stage_before_ovrtx.usda"
207-
export_stage_for_ovrtx(stage, export_path, num_envs, self._use_ovrtx_cloning)
208-
self._exported_usd_path = export_path
209-
logger.info("Exported to %s", export_path)
208+
self._exported_usd_string = export_stage_to_string(stage, num_envs, self._use_ovrtx_cloning)
210209

211210
def _initialize_from_spec(self, spec: CameraRenderSpec):
212211
"""Initialize the OVRTX renderer with internal environment cloning.
@@ -225,14 +224,10 @@ def _initialize_from_spec(self, spec: CameraRenderSpec):
225224
raise RuntimeError(f"Expected camera prim under '{env_0_prefix}', got '{first_cam_path}'")
226225
self._camera_rel_path = spec.camera_path_relative_to_env_0
227226

228-
usd_scene_path = self._exported_usd_path
229-
230-
if usd_scene_path is not None:
227+
if self._exported_usd_string is not None:
231228
logger.info("Injecting camera definitions...")
232229

233-
combined_usd_path, render_product_path = inject_cameras_into_usd(
234-
usd_scene_path,
235-
self.cfg,
230+
render_product_string, render_product_path = build_render_product_as_string(
236231
width=width,
237232
height=height,
238233
num_envs=num_envs,
@@ -242,23 +237,44 @@ def _initialize_from_spec(self, spec: CameraRenderSpec):
242237
)
243238
self._render_product_paths.append(render_product_path)
244239

240+
combined_usd_string = self._exported_usd_string + "\n\n" + render_product_string
241+
self._exported_usd_string = None # Free memory
242+
243+
if self.cfg.temp_usd_dir is not None:
244+
temp_usd_dir = Path(self.cfg.temp_usd_dir)
245+
elif not _IS_OVRTX_0_3_0_OR_NEWER:
246+
# OVRTX 0.2.0 is not able to load USD from a string, so we need to write to a temporary file.
247+
temp_usd_dir = Path(tempfile.gettempdir()) / "ovrtx"
248+
else:
249+
temp_usd_dir = None
250+
251+
if temp_usd_dir is not None:
252+
temp_usd_dir.mkdir(parents=True, exist_ok=True)
253+
temp_usd_path = temp_usd_dir / "ovrtx_renderer_stage.usda"
254+
with open(temp_usd_path, "w", encoding="utf-8") as f:
255+
f.write(combined_usd_string)
256+
logger.info("Wrote combined USD stage to %s", temp_usd_path)
257+
else:
258+
temp_usd_path = None
259+
245260
logger.info("Loading USD into OvRTX...")
246261
try:
247262
if _IS_OVRTX_0_3_0_OR_NEWER:
248-
self._renderer.open_usd(combined_usd_path)
249-
logger.info("USD loaded as root layer (path: %s)", combined_usd_path)
263+
self._renderer.open_usd_from_string(combined_usd_string)
264+
logger.info("OVRTX loaded USD from string successfully")
250265
else:
251-
handle = self._renderer.add_usd(combined_usd_path, path_prefix=None)
266+
assert temp_usd_path is not None # OVRTX < 0.3.0 always materializes combined USD on disk.
267+
handle = self._renderer.add_usd(str(temp_usd_path), path_prefix=None)
252268
self._usd_handles.append(handle)
253-
logger.info("USD loaded (path: %s, handle: %s)", combined_usd_path, handle)
269+
logger.info("OVRTX loaded USD from file successfully (path: %s, handle: %s)", temp_usd_path, handle)
254270
except Exception as e:
255271
logger.exception("Error loading USD: %s", e)
256272
raise
257273

258274
if self._use_ovrtx_cloning and num_envs > 1:
259275
logger.info("Using OVRTX internal cloning")
260276
self._clone_environments_in_ovrtx(num_envs)
261-
self._update_scene_partitions_after_clone(combined_usd_path, num_envs)
277+
self._update_scene_partitions_after_clone(num_envs)
262278

263279
self._initialized_scene = True
264280

@@ -299,7 +315,7 @@ def _clone_environments_in_ovrtx(self, num_envs: int):
299315
logger.error("Failed to clone environments: %s", e)
300316
raise RuntimeError(f"OvRTX environment cloning failed: {e}")
301317

302-
def _update_scene_partitions_after_clone(self, usd_file_path: str, num_envs: int):
318+
def _update_scene_partitions_after_clone(self, num_envs: int):
303319
"""Update scene partition attributes on cloned environments and cameras in OvRTX."""
304320
logger.info("Writing scene partitions for %d environments...", num_envs)
305321
partition_tokens = [f"env_{i}" for i in range(num_envs)]

source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer_cfg.py

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,6 @@
55

66
"""Configuration for OVRTX Renderer."""
77

8-
import tempfile
9-
from pathlib import Path
10-
118
from isaaclab.renderers.renderer_cfg import RendererCfg
129
from isaaclab.utils.configclass import configclass
1310

@@ -26,14 +23,11 @@ class OVRTXRendererCfg(RendererCfg):
2623
renderer_type: str = "ovrtx"
2724
"""Type identifier for OVRTX renderer."""
2825

29-
temp_usd_dir: str = str(Path(tempfile.gettempdir()) / "ovrtx")
26+
temp_usd_dir: str | None = None
3027
"""Directory for temporary combined USD files (scene + injected cameras).
3128
Used by the OVRTX renderer when building the render scope; must be writable.
3229
"""
3330

34-
temp_usd_suffix: str = ".usda"
35-
"""File suffix for temporary combined USD files (e.g. '.usda' or '.usdc')."""
36-
3731
use_ovrtx_cloning: bool = True
3832
"""When True, export only env_0 and use OVRTX ``clone_usd``. When False, export full multi-environment stage.
3933

source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_usd.py

Lines changed: 24 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -7,15 +7,9 @@
77

88
import logging
99
import math
10-
import tempfile
11-
from pathlib import Path
12-
from typing import TYPE_CHECKING
1310

1411
from pxr import Sdf, Usd, UsdGeom
1512

16-
if TYPE_CHECKING:
17-
from .ovrtx_renderer_cfg import OVRTXRendererCfg
18-
1913
logger = logging.getLogger(__name__)
2014

2115

@@ -105,34 +99,29 @@ def _tiled_resolution(num_envs: int, width: int, height: int) -> tuple[int, int]
10599
return num_cols * width, num_rows * height
106100

107101

108-
def inject_cameras_into_usd(
109-
usd_scene_path: str,
110-
cfg: "OVRTXRendererCfg",
102+
def build_render_product_as_string(
111103
width: int,
112104
height: int,
113105
num_envs: int,
114106
data_types: list[str],
115107
minimal_mode: int | None = None,
116108
camera_rel_path: str = "Camera",
117109
) -> tuple[str, str]:
118-
"""Inject camera and render product definitions into an existing USD file.
110+
"""Build the render product USD snippet as a string.
119111
120-
Reads the USD file, appends a Render scope (cameras + RenderProduct + Vars),
121-
writes to a temp file in cfg.temp_usd_dir, and returns (path_to_combined_usd, render_product_path).
112+
This string is meant to be appended to an exported stage (ASCII) before loading into OVRTX.
122113
123114
Args:
124-
usd_scene_path: Path to the base USD scene.
125-
cfg: OVRTX renderer config (simple_shading_mode, temp_usd_dir, temp_usd_suffix).
126-
width: Tile width from sensor config.
127-
height: Tile height from sensor config.
115+
width: Tile width from sensor config [px].
116+
height: Tile height from sensor config [px].
128117
num_envs: Number of environments from scene.
129118
data_types: Data types from sensor config.
130119
minimal_mode: RTX minimal mode. None if not requested. Valid values are 1, 2, 3.
131120
camera_rel_path: Camera prim path relative to the env root (e.g. ``"Camera"`` or ``"Robot/head_cam"``).
132-
"""
133-
with open(usd_scene_path) as f:
134-
original_usd = f.read()
135121
122+
Returns:
123+
Tuple of (render product USD snippet as a string, absolute render product prim path).
124+
"""
136125
data_types = data_types if data_types else ["rgb"]
137126
tiled_width, tiled_height = _tiled_resolution(num_envs, width, height)
138127

@@ -152,14 +141,7 @@ def inject_cameras_into_usd(
152141
tiled_height,
153142
minimal_mode,
154143
)
155-
combined_usd = original_usd.rstrip() + "\n\n" + camera_content
156-
157-
Path(cfg.temp_usd_dir).mkdir(parents=True, exist_ok=True)
158-
with tempfile.NamedTemporaryFile(mode="w", suffix=cfg.temp_usd_suffix, delete=False, dir=cfg.temp_usd_dir) as f:
159-
f.write(combined_usd)
160-
temp_path = f.name
161-
logger.info("Created combined USD: %s", temp_path)
162-
return temp_path, render_product_path
144+
return camera_content, render_product_path
163145

164146

165147
def create_scene_partition_attributes(
@@ -207,40 +189,38 @@ def create_scene_partition_attributes(
207189
logger.debug("Set scene partition '%s' on prim '%s'", scene_partition, prim.GetPath())
208190

209191

210-
def export_stage_for_ovrtx(stage, export_path: str, num_envs: int, use_ovrtx_cloning: bool = True) -> str:
211-
"""Export the stage to a USD file; when num_envs > 1, only env_0 is exported for OVRTX cloning.
192+
def export_stage_to_string(stage, num_envs: int, use_ovrtx_cloning: bool = True) -> str:
193+
"""Export the stage to a string; when num_envs > 1, only env_0 is exported for OVRTX cloning.
212194
213195
When num_envs > 1, deactivates env_1..env_{num_envs-1} before export and reactivates
214-
them after, so the file contains only env_0. The stage is modified in place.
196+
them after, so the exported content contains only env_0. The stage is modified in place.
215197
216198
Args:
217199
stage: USD stage to export.
218-
export_path: Path for the exported file.
219200
num_envs: Number of environments.
201+
use_ovrtx_cloning: Whether OVRTX cloning is enabled.
220202
221203
Returns:
222-
export_path (same as input).
204+
The exported stage as a string.
223205
"""
224-
deactivated = []
206+
deactivated_prims = []
225207
if use_ovrtx_cloning and num_envs > 1:
226-
logger.info("Deactivating %d cloned environments...", num_envs - 1)
208+
logger.info("Deactivating %d environment roots...", num_envs - 1)
227209
for env_idx in range(1, num_envs):
228210
env_path = f"/World/envs/env_{env_idx}"
229211
prim = stage.GetPrimAtPath(env_path)
230212
if prim.IsValid() and prim.IsActive():
231213
prim.SetActive(False)
232-
deactivated.append(prim)
233-
if env_idx <= 3 or env_idx == num_envs - 1:
234-
logger.info("Deactivated: %s", env_path)
235-
if num_envs > 5:
236-
logger.info("... (deactivated %d environments total)", len(deactivated))
214+
deactivated_prims.append(prim)
215+
logger.debug("Deactivated environment root: %s", env_path)
216+
217+
logger.info("Deactivated %d environment roots in total", len(deactivated_prims))
237218

238219
try:
239-
stage.Export(export_path)
240-
return export_path
220+
return stage.ExportToString()
241221
finally:
242-
if deactivated:
243-
logger.info("Reactivating %d environments...", len(deactivated))
244-
for prim in deactivated:
222+
if deactivated_prims:
223+
logger.info("Reactivating %d environment roots...", len(deactivated_prims))
224+
for prim in deactivated_prims:
245225
if prim.IsValid():
246226
prim.SetActive(True)

0 commit comments

Comments
 (0)