From f3dcbe2a756abd6bc3a1d26161bc9929afb344a8 Mon Sep 17 00:00:00 2001 From: pv-nvidia <197907000+pv-nvidia@users.noreply.github.com> Date: Wed, 19 Aug 2026 06:35:26 +0000 Subject: [PATCH 1/4] Order OVRTX mapped-attribute commits after the producing Warp stream Add map_attribute_for_warp_writes, a reusable context manager that maps an OVRTX attribute binding for CUDA writes and unmaps it with the producing Warp stream as the CUDA sync, so OVRTX's commit of the mapped data waits for the fill on the GPU. The binding's own context manager cannot carry that sync: its __exit__ takes no arguments and unmaps with an empty cuda_sync, which the OVRTX API treats as no synchronization. Route the renderer's object and camera transform writes, the only GPU attribute mappings in the codebase, through the helper. Previously their ordering held only through CUDA legacy default-stream serialization against Warp's blocking-type streams, an implementation detail the OVRTX contract does not promise. Render-var CUDA maps already order their reads via _map_render_var_to_dlpack, and render-var CPU maps are synchronized by OVRTX before map() returns, so neither needs changes. --- .../api/lab_ov/isaaclab_ov.renderers.rst | 3 + .../pv-ovrtx-mapping-stream-sync.minor.rst | 16 +++++ .../isaaclab_ov/renderers/__init__.pyi | 2 + .../isaaclab_ov/renderers/ovrtx_mapping.py | 61 +++++++++++++++++++ .../isaaclab_ov/renderers/ovrtx_renderer.py | 17 ++---- .../test/test_ovrtx_renderer_contract.py | 61 ++++++++++++++++++- 6 files changed, 147 insertions(+), 13 deletions(-) create mode 100644 source/isaaclab_ov/changelog.d/pv-ovrtx-mapping-stream-sync.minor.rst create mode 100644 source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py diff --git a/docs/source/api/lab_ov/isaaclab_ov.renderers.rst b/docs/source/api/lab_ov/isaaclab_ov.renderers.rst index 3ad45b7b4e93..a774c1ad1448 100644 --- a/docs/source/api/lab_ov/isaaclab_ov.renderers.rst +++ b/docs/source/api/lab_ov/isaaclab_ov.renderers.rst @@ -15,9 +15,12 @@ The following classes are part of the public :mod:`isaaclab_ov.renderers` API. OVRTXRenderer OVRTXRendererCfg + map_attribute_for_warp_writes .. autoclass:: OVRTXRenderer :show-inheritance: .. autoclass:: OVRTXRendererCfg :show-inheritance: + +.. autofunction:: map_attribute_for_warp_writes diff --git a/source/isaaclab_ov/changelog.d/pv-ovrtx-mapping-stream-sync.minor.rst b/source/isaaclab_ov/changelog.d/pv-ovrtx-mapping-stream-sync.minor.rst new file mode 100644 index 000000000000..d425463225fc --- /dev/null +++ b/source/isaaclab_ov/changelog.d/pv-ovrtx-mapping-stream-sync.minor.rst @@ -0,0 +1,16 @@ +Added +^^^^^ + +* Added :func:`~isaaclab_ov.renderers.map_attribute_for_warp_writes`, a context manager that maps + an OVRTX attribute binding for CUDA writes and unmaps it with the producing Warp stream as the + CUDA sync. Use it instead of ``with binding.map(...)`` for GPU writes: the binding's own context + manager unmaps without a CUDA sync, so OVRTX's commit is not ordered against the fill. + +Fixed +^^^^^ + +* Fixed the OVRTX renderer's GPU transform writes (object and camera ``omni:xform`` mappings) + committing without CUDA synchronization against the Warp kernels that fill the mapped buffers. + The commit is now ordered on the producing Warp stream, as the OVRTX API contract requires; + previously the ordering held only through CUDA legacy default-stream serialization, an + implementation detail the contract does not promise. diff --git a/source/isaaclab_ov/isaaclab_ov/renderers/__init__.pyi b/source/isaaclab_ov/isaaclab_ov/renderers/__init__.pyi index 70526feeb84a..bd16abc54924 100644 --- a/source/isaaclab_ov/isaaclab_ov/renderers/__init__.pyi +++ b/source/isaaclab_ov/isaaclab_ov/renderers/__init__.pyi @@ -7,8 +7,10 @@ __all__ = [ "OVRTXRenderer", "OVRTXRendererCfg", "Renderer", + "map_attribute_for_warp_writes", ] +from .ovrtx_mapping import map_attribute_for_warp_writes from .ovrtx_renderer import OVRTXRenderer from .ovrtx_renderer import OVRTXRenderer as Renderer from .ovrtx_renderer_cfg import OVRTXRendererCfg diff --git a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py new file mode 100644 index 000000000000..ec8e07fe4a81 --- /dev/null +++ b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py @@ -0,0 +1,61 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Stream-safe helpers for OVRTX attribute-binding mappings. + +Filling a CUDA-mapped OVRTX attribute buffer from Warp is only correct when the commit at unmap +time is ordered against the Warp stream that produced the data: OVRTX's API contract requires the +mapped data to be ready when the unmap's CUDA sync signals, and an unmap without a sync performs +no synchronization at all. The ``with binding.map(...)`` form cannot carry that sync -- its +``__exit__`` takes no arguments -- so GPU writes through a mapping should use +:func:`map_attribute_for_warp_writes` instead of the binding's own context manager. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager +from typing import Any + +import warp as wp +from ovrtx import Device + + +def cuda_device_id(device: str) -> int: + """CUDA device index parsed from a Warp device string, e.g. ``"cuda:1"`` -> ``1``. + + Args: + device: Warp CUDA device string (``"cuda"`` or ``"cuda:"``). + + Returns: + The device index, ``0`` when the string carries none. + """ + parts = device.split(":") + return int(parts[1]) if len(parts) > 1 else 0 + + +@contextmanager +def map_attribute_for_warp_writes(binding: Any, device: str, dtype: Any) -> Iterator[wp.array]: + """Map ``binding`` for CUDA writes and yield its buffer as a Warp array; commit after the fill. + + The caller fills the yielded array with Warp work enqueued on ``device``'s current stream (the + default for ``wp.launch``/``wp.copy``). On exit -- error or not -- the mapping is unmapped with + that stream as the CUDA sync, so OVRTX's commit of the mapped data waits for the fill on the + GPU instead of racing it. OVRTX has no discard path (unmap always commits), so a failed fill + still publishes whatever landed in the buffer. + + Args: + binding: OVRTX attribute binding (from ``bind_attribute``) whose buffer is written. + device: Warp CUDA device string the fill work runs on (e.g. ``"cuda:0"``). + dtype: Warp dtype the mapped tensor is viewed as (e.g. ``wp.mat44d``). + + Yields: + The mapped buffer as a zero-copy Warp array, valid only inside the ``with`` block. + """ + attr_mapping = binding.map(device=Device.CUDA, device_id=cuda_device_id(device)) + try: + yield wp.from_dlpack(attr_mapping.tensor, dtype=dtype) + finally: + attr_mapping.unmap(stream=wp.get_stream(device).cuda_stream) diff --git a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py index 5d7d380bbe34..fefcc840f0d6 100644 --- a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py +++ b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py @@ -87,6 +87,7 @@ decode_stable_id_map, decode_stable_id_semantic_id_map, ) +from .ovrtx_mapping import map_attribute_for_warp_writes from .ovrtx_renderer_cfg import OVRTXRendererCfg from .ovrtx_renderer_kernels import ( create_camera_transforms_kernel, @@ -341,12 +342,6 @@ def supported_output_types(self) -> dict[RenderBufferKind, RenderBufferSpec]: RenderBufferKind.MOTION_VECTORS: RenderBufferSpec(2, wp.float32), } - @property - def _device_id(self) -> int: - """CUDA device index extracted from ``self._device`` for OVRTX ``binding.map()`` calls.""" - parts = self._device.split(":") - return int(parts[1]) if len(parts) > 1 else 0 - def __init__(self, cfg: OVRTXRendererCfg): self.cfg = cfg self._device = "cuda:0" # default; overridden by create_render_data(spec) @@ -852,8 +847,7 @@ def _update_transforms_legacy(self) -> None: if body_q is None: return - with self._object_xform_binding.map(device=Device.CUDA, device_id=self._device_id) as attr_mapping: - ovrtx_transforms = wp.from_dlpack(attr_mapping.tensor, dtype=wp.mat44d) + with map_attribute_for_warp_writes(self._object_xform_binding, self._device, wp.mat44d) as ovrtx_transforms: wp.launch( kernel=sync_newton_transforms_kernel, dim=len(self._object_newton_indices), @@ -955,9 +949,8 @@ def _update_camera_legacy( device=self._device, ) if self._camera_xform_binding is not None: - with self._camera_xform_binding.map(device=Device.CUDA, device_id=self._device_id) as attr_mapping: - wp_transforms_view = wp.from_dlpack(attr_mapping.tensor, dtype=wp.mat44d) - wp.copy(wp_transforms_view, camera_transforms) + with map_attribute_for_warp_writes(self._camera_xform_binding, self._device, wp.mat44d) as transforms_view: + wp.copy(transforms_view, camera_transforms) def read_output( self, @@ -1260,7 +1253,7 @@ def _prepare_ppisp_hdr_source( # FIXME: OVRTX render var mapping can select a different CUDA device # than the camera/output buffers on MGPU systems. Keep this PPISP-only # bridge until render var mapping can be constrained like transform - # bindings, which use ``device_id=self._device_id``. + # bindings, whose maps pin ``device_id`` (see ``ovrtx_mapping``). return wp.clone(tiled_data, device=output_device) def _process_render_frame(self, render_data: OVRTXRenderData, frame, output_buffers: dict) -> None: diff --git a/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py b/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py index 1cb9de1386f0..3f61c4f76366 100644 --- a/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py +++ b/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py @@ -30,7 +30,10 @@ ] if not _MISSING_MODULES: - from isaaclab_ov.renderers import OVRTXRendererCfg # noqa: E402 + from isaaclab_ov.renderers import ( + OVRTXRendererCfg, # noqa: E402 + ovrtx_mapping, # noqa: E402 + ) from isaaclab_ov.renderers import ovrtx_renderer as ovrtx_renderer_module # noqa: E402 from isaaclab_ov.renderers.ovrtx_renderer import ( # noqa: E402 _DISABLE_LINUX_CUDA_CPU_SYNC_ENV, @@ -476,6 +479,62 @@ def test_ovrtx_map_render_var_orders_the_read_against_render_completion(monkeypa assert render_var.ordering == [expected] +class _RecordingMappedBinding: + """Stand-in for an OVRTX attribute binding that records how its mapping is committed.""" + + def __init__(self): + self.map_calls: list[dict] = [] + self.unmap_calls: list[dict] = [] + + def map(self, *, device, device_id): + self.map_calls.append({"device": device, "device_id": device_id}) + binding = self + + class _Mapping: + tensor = object() + + def unmap(self, *, event=None, stream=None): + binding.unmap_calls.append({"event": event, "stream": stream}) + + return _Mapping() + + +def test_map_attribute_for_warp_writes_commits_on_the_producer_stream(monkeypatch): + """The unmap names the Warp stream that produced the data, so the commit cannot race the fill. + + An unmap without a CUDA sync performs no synchronization at all, so the assertion is on the + unmap's ``stream`` argument, not merely on the unmap happening. + """ + sentinel = object() + binding = _RecordingMappedBinding() + monkeypatch.setattr(ovrtx_mapping.wp, "get_stream", lambda device: types.SimpleNamespace(cuda_stream=99)) + monkeypatch.setattr(ovrtx_mapping.wp, "from_dlpack", lambda tensor, dtype: sentinel) + + with ovrtx_mapping.map_attribute_for_warp_writes(binding, "cuda:1", wp.mat44d) as array: + assert array is sentinel + + assert binding.map_calls == [{"device": ovrtx_mapping.Device.CUDA, "device_id": 1}] + assert binding.unmap_calls == [{"event": None, "stream": 99}] + + +def test_map_attribute_for_warp_writes_unmaps_when_the_fill_raises(monkeypatch): + """A failed fill must still release the mapping exactly once, with the same stream ordering. + + Skipping the unmap would leak the mapping to OVRTX's ``__del__`` safety net, which commits + fire-and-forget without any CUDA sync. + """ + binding = _RecordingMappedBinding() + monkeypatch.setattr(ovrtx_mapping.wp, "get_stream", lambda device: types.SimpleNamespace(cuda_stream=7)) + monkeypatch.setattr(ovrtx_mapping.wp, "from_dlpack", lambda tensor, dtype: object()) + + with pytest.raises(ValueError, match="fill failed"): + with ovrtx_mapping.map_attribute_for_warp_writes(binding, "cuda", wp.mat44d): + raise ValueError("fill failed") + + assert binding.map_calls == [{"device": ovrtx_mapping.Device.CUDA, "device_id": 0}] + assert binding.unmap_calls == [{"event": None, "stream": 7}] + + def test_ovrtx_cleanup_releases_only_the_given_render_data(): """``cleanup`` releases the render data's own buffers and leaves the renderer usable. From 7e85e0b97f224ef4b4ac638cd4530de6afd6dca4 Mon Sep 17 00:00:00 2001 From: pv-nvidia <197907000+pv-nvidia@users.noreply.github.com> Date: Wed, 19 Aug 2026 07:29:33 +0000 Subject: [PATCH 2/4] Resolve the mapping device through Warp and defer the ovrtx import wp.get_device resolves the device ident once, so the mapped CUDA device and the sync stream come from the same device: a bare "cuda" previously parsed to index 0 while the stream lookup used Warp's current CUDA device, silently voiding the ordering guarantee on non-zero devices. The string parser is gone; the resolved device's ordinal and current stream are used for map and unmap. The ovrtx import moves inside the helper so importing the module through the package's lazy exports cannot initialize ovrtx without the guarded environment ovrtx_renderer establishes (OVRTX_SKIP_USD_CHECK, actionable install error). Any real binding is created through that path, so ovrtx is already imported when the helper runs. --- .../isaaclab_ov/renderers/ovrtx_mapping.py | 30 ++++++++--------- .../test/test_ovrtx_renderer_contract.py | 33 ++++++++++++++++--- 2 files changed, 41 insertions(+), 22 deletions(-) diff --git a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py index ec8e07fe4a81..0f58deb3b563 100644 --- a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py +++ b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py @@ -20,20 +20,6 @@ from typing import Any import warp as wp -from ovrtx import Device - - -def cuda_device_id(device: str) -> int: - """CUDA device index parsed from a Warp device string, e.g. ``"cuda:1"`` -> ``1``. - - Args: - device: Warp CUDA device string (``"cuda"`` or ``"cuda:"``). - - Returns: - The device index, ``0`` when the string carries none. - """ - parts = device.split(":") - return int(parts[1]) if len(parts) > 1 else 0 @contextmanager @@ -46,16 +32,26 @@ def map_attribute_for_warp_writes(binding: Any, device: str, dtype: Any) -> Iter GPU instead of racing it. OVRTX has no discard path (unmap always commits), so a failed fill still publishes whatever landed in the buffer. + ``device`` is resolved once through Warp, so the mapped CUDA device and the sync stream cannot + diverge -- a bare ``"cuda"`` maps and syncs Warp's current CUDA device. + Args: binding: OVRTX attribute binding (from ``bind_attribute``) whose buffer is written. - device: Warp CUDA device string the fill work runs on (e.g. ``"cuda:0"``). + device: Warp CUDA device the fill work runs on (e.g. ``"cuda:0"``). dtype: Warp dtype the mapped tensor is viewed as (e.g. ``wp.mat44d``). Yields: The mapped buffer as a zero-copy Warp array, valid only inside the ``with`` block. """ - attr_mapping = binding.map(device=Device.CUDA, device_id=cuda_device_id(device)) + # Deferred so importing this module (e.g. through the package's lazy exports) cannot initialize + # ovrtx without the guarded environment ``ovrtx_renderer`` establishes (OVRTX_SKIP_USD_CHECK, + # actionable install error). Any real ``binding`` was created through that path, so ovrtx is + # already imported by the time this runs. + from ovrtx import Device # noqa: PLC0415 + + warp_device = wp.get_device(device) + attr_mapping = binding.map(device=Device.CUDA, device_id=warp_device.ordinal) try: yield wp.from_dlpack(attr_mapping.tensor, dtype=dtype) finally: - attr_mapping.unmap(stream=wp.get_stream(device).cuda_stream) + attr_mapping.unmap(stream=warp_device.stream.cuda_stream) diff --git a/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py b/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py index 3f61c4f76366..d3abe9e20a3c 100644 --- a/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py +++ b/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py @@ -499,6 +499,12 @@ def unmap(self, *, event=None, stream=None): return _Mapping() +def _patch_warp_device(monkeypatch, *, ordinal: int, cuda_stream: int) -> None: + """Resolve any device ident to one fake Warp device with the given ordinal and current stream.""" + warp_device = types.SimpleNamespace(ordinal=ordinal, stream=types.SimpleNamespace(cuda_stream=cuda_stream)) + monkeypatch.setattr(ovrtx_mapping.wp, "get_device", lambda device: warp_device) + + def test_map_attribute_for_warp_writes_commits_on_the_producer_stream(monkeypatch): """The unmap names the Warp stream that produced the data, so the commit cannot race the fill. @@ -507,16 +513,33 @@ def test_map_attribute_for_warp_writes_commits_on_the_producer_stream(monkeypatc """ sentinel = object() binding = _RecordingMappedBinding() - monkeypatch.setattr(ovrtx_mapping.wp, "get_stream", lambda device: types.SimpleNamespace(cuda_stream=99)) + _patch_warp_device(monkeypatch, ordinal=1, cuda_stream=99) monkeypatch.setattr(ovrtx_mapping.wp, "from_dlpack", lambda tensor, dtype: sentinel) with ovrtx_mapping.map_attribute_for_warp_writes(binding, "cuda:1", wp.mat44d) as array: assert array is sentinel - assert binding.map_calls == [{"device": ovrtx_mapping.Device.CUDA, "device_id": 1}] + assert binding.map_calls == [{"device": ovrtx_renderer_module.Device.CUDA, "device_id": 1}] assert binding.unmap_calls == [{"event": None, "stream": 99}] +def test_map_attribute_for_warp_writes_resolves_bare_cuda_through_warp(monkeypatch): + """A bare ``"cuda"`` maps the device Warp resolves it to, not index 0. + + Parsing the string instead would map GPU 0 while syncing another GPU's stream whenever Warp's + current CUDA device is non-zero, silently voiding the ordering guarantee. + """ + binding = _RecordingMappedBinding() + _patch_warp_device(monkeypatch, ordinal=3, cuda_stream=42) + monkeypatch.setattr(ovrtx_mapping.wp, "from_dlpack", lambda tensor, dtype: object()) + + with ovrtx_mapping.map_attribute_for_warp_writes(binding, "cuda", wp.mat44d): + pass + + assert binding.map_calls == [{"device": ovrtx_renderer_module.Device.CUDA, "device_id": 3}] + assert binding.unmap_calls == [{"event": None, "stream": 42}] + + def test_map_attribute_for_warp_writes_unmaps_when_the_fill_raises(monkeypatch): """A failed fill must still release the mapping exactly once, with the same stream ordering. @@ -524,14 +547,14 @@ def test_map_attribute_for_warp_writes_unmaps_when_the_fill_raises(monkeypatch): fire-and-forget without any CUDA sync. """ binding = _RecordingMappedBinding() - monkeypatch.setattr(ovrtx_mapping.wp, "get_stream", lambda device: types.SimpleNamespace(cuda_stream=7)) + _patch_warp_device(monkeypatch, ordinal=0, cuda_stream=7) monkeypatch.setattr(ovrtx_mapping.wp, "from_dlpack", lambda tensor, dtype: object()) with pytest.raises(ValueError, match="fill failed"): - with ovrtx_mapping.map_attribute_for_warp_writes(binding, "cuda", wp.mat44d): + with ovrtx_mapping.map_attribute_for_warp_writes(binding, "cuda:0", wp.mat44d): raise ValueError("fill failed") - assert binding.map_calls == [{"device": ovrtx_mapping.Device.CUDA, "device_id": 0}] + assert binding.map_calls == [{"device": ovrtx_renderer_module.Device.CUDA, "device_id": 0}] assert binding.unmap_calls == [{"event": None, "stream": 7}] From 72ff1d2a4a74809cb37cc36ebd8d0117115b1664 Mon Sep 17 00:00:00 2001 From: pv-nvidia <197907000+pv-nvidia@users.noreply.github.com> Date: Wed, 19 Aug 2026 07:39:59 +0000 Subject: [PATCH 3/4] Extract get_warp_device_id and give it its own unit test The device-index resolution is one testable contract: it delegates to Warp so a bare "cuda" means Warp's current CUDA device, never a parse of the string. Pin that in a focused unit test and drop the map-level test that duplicated it; the map tests keep asserting only what the context manager owns (device forwarding, producer-stream unmap, and release on a failed fill). --- .../isaaclab_ov/renderers/ovrtx_mapping.py | 24 +++++++++++--- .../test/test_ovrtx_renderer_contract.py | 32 ++++++++----------- 2 files changed, 32 insertions(+), 24 deletions(-) diff --git a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py index 0f58deb3b563..0870213d964c 100644 --- a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py +++ b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py @@ -22,6 +22,21 @@ import warp as wp +def get_warp_device_id(device: str) -> int: + """CUDA device index of the Warp device ``device`` resolves to. + + Resolving through Warp keeps the mapped device consistent with where Warp enqueues the fill + work: a bare ``"cuda"`` means Warp's current CUDA device, not index ``0``. + + Args: + device: Warp CUDA device identifier (e.g. ``"cuda"`` or ``"cuda:1"``). + + Returns: + The CUDA device index. + """ + return wp.get_device(device).ordinal + + @contextmanager def map_attribute_for_warp_writes(binding: Any, device: str, dtype: Any) -> Iterator[wp.array]: """Map ``binding`` for CUDA writes and yield its buffer as a Warp array; commit after the fill. @@ -32,8 +47,8 @@ def map_attribute_for_warp_writes(binding: Any, device: str, dtype: Any) -> Iter GPU instead of racing it. OVRTX has no discard path (unmap always commits), so a failed fill still publishes whatever landed in the buffer. - ``device`` is resolved once through Warp, so the mapped CUDA device and the sync stream cannot - diverge -- a bare ``"cuda"`` maps and syncs Warp's current CUDA device. + ``device`` is resolved through Warp (see :func:`get_warp_device_id`), so the mapped CUDA device + matches where the fill work runs -- a bare ``"cuda"`` maps and syncs Warp's current CUDA device. Args: binding: OVRTX attribute binding (from ``bind_attribute``) whose buffer is written. @@ -49,9 +64,8 @@ def map_attribute_for_warp_writes(binding: Any, device: str, dtype: Any) -> Iter # already imported by the time this runs. from ovrtx import Device # noqa: PLC0415 - warp_device = wp.get_device(device) - attr_mapping = binding.map(device=Device.CUDA, device_id=warp_device.ordinal) + attr_mapping = binding.map(device=Device.CUDA, device_id=get_warp_device_id(device)) try: yield wp.from_dlpack(attr_mapping.tensor, dtype=dtype) finally: - attr_mapping.unmap(stream=warp_device.stream.cuda_stream) + attr_mapping.unmap(stream=wp.get_stream(device).cuda_stream) diff --git a/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py b/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py index d3abe9e20a3c..0bbbbd07b23f 100644 --- a/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py +++ b/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py @@ -501,8 +501,19 @@ def unmap(self, *, event=None, stream=None): def _patch_warp_device(monkeypatch, *, ordinal: int, cuda_stream: int) -> None: """Resolve any device ident to one fake Warp device with the given ordinal and current stream.""" - warp_device = types.SimpleNamespace(ordinal=ordinal, stream=types.SimpleNamespace(cuda_stream=cuda_stream)) - monkeypatch.setattr(ovrtx_mapping.wp, "get_device", lambda device: warp_device) + monkeypatch.setattr(ovrtx_mapping.wp, "get_device", lambda device: types.SimpleNamespace(ordinal=ordinal)) + monkeypatch.setattr(ovrtx_mapping.wp, "get_stream", lambda device: types.SimpleNamespace(cuda_stream=cuda_stream)) + + +def test_get_warp_device_id_resolves_through_warp(monkeypatch): + """The device index comes from Warp's resolution of the ident, never from parsing the string. + + Parsing would send a bare ``"cuda"`` to index 0 while Warp enqueues the fill on its current + CUDA device, so the mapping and the sync stream could target different GPUs. + """ + monkeypatch.setattr(ovrtx_mapping.wp, "get_device", lambda device: types.SimpleNamespace(ordinal=3)) + + assert ovrtx_mapping.get_warp_device_id("cuda") == 3 def test_map_attribute_for_warp_writes_commits_on_the_producer_stream(monkeypatch): @@ -523,23 +534,6 @@ def test_map_attribute_for_warp_writes_commits_on_the_producer_stream(monkeypatc assert binding.unmap_calls == [{"event": None, "stream": 99}] -def test_map_attribute_for_warp_writes_resolves_bare_cuda_through_warp(monkeypatch): - """A bare ``"cuda"`` maps the device Warp resolves it to, not index 0. - - Parsing the string instead would map GPU 0 while syncing another GPU's stream whenever Warp's - current CUDA device is non-zero, silently voiding the ordering guarantee. - """ - binding = _RecordingMappedBinding() - _patch_warp_device(monkeypatch, ordinal=3, cuda_stream=42) - monkeypatch.setattr(ovrtx_mapping.wp, "from_dlpack", lambda tensor, dtype: object()) - - with ovrtx_mapping.map_attribute_for_warp_writes(binding, "cuda", wp.mat44d): - pass - - assert binding.map_calls == [{"device": ovrtx_renderer_module.Device.CUDA, "device_id": 3}] - assert binding.unmap_calls == [{"event": None, "stream": 42}] - - def test_map_attribute_for_warp_writes_unmaps_when_the_fill_raises(monkeypatch): """A failed fill must still release the mapping exactly once, with the same stream ordering. From 309c3454a644680916198ccba15d8507144dbdda Mon Sep 17 00:00:00 2001 From: pv-nvidia <197907000+pv-nvidia@users.noreply.github.com> Date: Wed, 19 Aug 2026 07:51:51 +0000 Subject: [PATCH 4/4] Keep the pre-existing bare-cuda parse; defer device resolution to a follow-up Resolving the mapping device through Warp changed behavior for bare "cuda" strings (current CUDA device instead of index 0). That change is worth making, but not inside the stream-ordering fix: restore the string parse so this PR alters no device placement, and mark the bare-cuda / current-device split with a TODO. The follow-up caches the resolved Warp device on the renderer instead of re-deriving it from strings. --- .../isaaclab_ov/renderers/ovrtx_mapping.py | 22 +++++++++---------- .../test/test_ovrtx_renderer_contract.py | 16 ++++++-------- 2 files changed, 18 insertions(+), 20 deletions(-) diff --git a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py index 0870213d964c..a0d272d6abd4 100644 --- a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py +++ b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_mapping.py @@ -22,19 +22,22 @@ import warp as wp -def get_warp_device_id(device: str) -> int: - """CUDA device index of the Warp device ``device`` resolves to. +def _cuda_device_id(device: str) -> int: + """CUDA device index parsed from a Warp device string, e.g. ``"cuda:1"`` -> ``1``. - Resolving through Warp keeps the mapped device consistent with where Warp enqueues the fill - work: a bare ``"cuda"`` means Warp's current CUDA device, not index ``0``. + TODO: A bare ``"cuda"`` parses to ``0`` while Warp enqueues fill work on its *current* CUDA + device, so the mapping and the fill can target different GPUs on multi-GPU processes. The + split predates this helper and is kept here to avoid a behavior change; a follow-up caches + the resolved Warp device on the renderer instead of re-deriving it from strings. Args: - device: Warp CUDA device identifier (e.g. ``"cuda"`` or ``"cuda:1"``). + device: Warp CUDA device string (``"cuda"`` or ``"cuda:"``). Returns: - The CUDA device index. + The parsed CUDA device index, ``0`` when the string carries none. """ - return wp.get_device(device).ordinal + parts = device.split(":") + return int(parts[1]) if len(parts) > 1 else 0 @contextmanager @@ -47,9 +50,6 @@ def map_attribute_for_warp_writes(binding: Any, device: str, dtype: Any) -> Iter GPU instead of racing it. OVRTX has no discard path (unmap always commits), so a failed fill still publishes whatever landed in the buffer. - ``device`` is resolved through Warp (see :func:`get_warp_device_id`), so the mapped CUDA device - matches where the fill work runs -- a bare ``"cuda"`` maps and syncs Warp's current CUDA device. - Args: binding: OVRTX attribute binding (from ``bind_attribute``) whose buffer is written. device: Warp CUDA device the fill work runs on (e.g. ``"cuda:0"``). @@ -64,7 +64,7 @@ def map_attribute_for_warp_writes(binding: Any, device: str, dtype: Any) -> Iter # already imported by the time this runs. from ovrtx import Device # noqa: PLC0415 - attr_mapping = binding.map(device=Device.CUDA, device_id=get_warp_device_id(device)) + attr_mapping = binding.map(device=Device.CUDA, device_id=_cuda_device_id(device)) try: yield wp.from_dlpack(attr_mapping.tensor, dtype=dtype) finally: diff --git a/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py b/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py index 0bbbbd07b23f..c24815628c8a 100644 --- a/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py +++ b/source/isaaclab_ov/test/test_ovrtx_renderer_contract.py @@ -500,20 +500,18 @@ def unmap(self, *, event=None, stream=None): def _patch_warp_device(monkeypatch, *, ordinal: int, cuda_stream: int) -> None: - """Resolve any device ident to one fake Warp device with the given ordinal and current stream.""" - monkeypatch.setattr(ovrtx_mapping.wp, "get_device", lambda device: types.SimpleNamespace(ordinal=ordinal)) + """Fake the current Warp stream; ``ordinal`` documents the device the test pretends to run on.""" monkeypatch.setattr(ovrtx_mapping.wp, "get_stream", lambda device: types.SimpleNamespace(cuda_stream=cuda_stream)) -def test_get_warp_device_id_resolves_through_warp(monkeypatch): - """The device index comes from Warp's resolution of the ident, never from parsing the string. +@pytest.mark.parametrize(("device", "expected"), [("cuda:1", 1), ("cuda", 0)]) +def test_cuda_device_id_parses_the_device_string(device, expected): + """The mapping device index is parsed from the string; a bare ``"cuda"`` parses to 0. - Parsing would send a bare ``"cuda"`` to index 0 while Warp enqueues the fill on its current - CUDA device, so the mapping and the sync stream could target different GPUs. + The bare-``"cuda"`` case intentionally preserves pre-existing behavior even though Warp + resolves it to its current CUDA device -- see the TODO on ``_cuda_device_id``. """ - monkeypatch.setattr(ovrtx_mapping.wp, "get_device", lambda device: types.SimpleNamespace(ordinal=3)) - - assert ovrtx_mapping.get_warp_device_id("cuda") == 3 + assert ovrtx_mapping._cuda_device_id(device) == expected def test_map_attribute_for_warp_writes_commits_on_the_producer_stream(monkeypatch):