diff --git a/src/unilab/base/backend/mjwarp/backend.py b/src/unilab/base/backend/mjwarp/backend.py index de84edf78..f57d5d3e1 100644 --- a/src/unilab/base/backend/mjwarp/backend.py +++ b/src/unilab/base/backend/mjwarp/backend.py @@ -9,8 +9,11 @@ from __future__ import annotations +import gc import time -from collections.abc import Sequence +import warnings +from collections.abc import Iterator, Sequence +from contextlib import contextmanager from os import PathLike from typing import Any @@ -33,6 +36,47 @@ from .materialization import materialize_mjwarp_scene from .playback import run_mjwarp_playback, validate_mjwarp_visual_model +_GRAPH_CAPTURE_MIN_DRIVER = (12, 4) + + +@contextmanager +def _suspend_gc() -> Iterator[None]: + """Keep graph finalizers from running inside a new Warp capture.""" + enabled = gc.isenabled() + gc.disable() + try: + yield + finally: + if enabled: + gc.enable() + + +def _cuda_graph_eligibility(warp: Any, device: Any) -> tuple[bool, str | None]: + """Return the cold-path CUDA graph decision and a fallback diagnostic.""" + if not bool(device.is_cuda): + return False, "active Warp device is not CUDA" + + try: + driver_version = warp.get_cuda_driver_version() + except Exception as exc: + return False, f"CUDA driver query failed: {type(exc).__name__}: {exc}" + if driver_version is None: + return False, "CUDA driver version is unavailable" + + try: + mempool_enabled = bool(warp.is_mempool_enabled(device)) + except Exception as exc: + return False, f"CUDA mempool query failed: {type(exc).__name__}: {exc}" + + reasons: list[str] = [] + if tuple(driver_version) < _GRAPH_CAPTURE_MIN_DRIVER: + reasons.append(f"CUDA driver {driver_version[0]}.{driver_version[1]} is older than 12.4") + if not mempool_enabled: + reasons.append("CUDA mempool is disabled") + if reasons: + return False, "; ".join(reasons) + return True, None + class MjwarpBackend(SimBackend): """Independent CUDA backend exposed through the host NumPy profile. @@ -176,6 +220,7 @@ def __init__( self._mujoco_warp.forward(self._device_model, self._device_data) self._synchronize() self._refresh_host_cache() + self._initialize_cuda_graphs(device) # ------------------------------------------------------------------ # # Cold-path model binding # @@ -255,6 +300,96 @@ def _download(self, device_array: Any, host_array: np.ndarray) -> None: def _synchronize(self) -> None: self._warp.synchronize_device() + def _disable_cuda_graphs(self, reason: str) -> None: + """Atomically select the eager path and release any captured graphs.""" + self._cuda_graph_enabled = False + self._step_graph = None + self._forward_graph = None + self._reset_graph = None + self._cuda_graph_disable_reason: str | None = reason + + def _initialize_cuda_graphs(self, device: Any) -> None: + """Capture fixed-address device operations or retain the eager fallback. + + Current uploads mutate existing Warp arrays with ``assign``. Any future + owner-layer operation that replaces a model or data array must call this + method afterward so captured pointers cannot become stale. + """ + self._disable_cuda_graphs("CUDA graph capture has not been initialized") + eligible, reason = _cuda_graph_eligibility(self._warp, device) + if not eligible: + assert reason is not None + self._cuda_graph_disable_reason = reason + warnings.warn( + f"mjwarp CUDA graphs disabled; using eager execution: {reason}", + RuntimeWarning, + stacklevel=2, + ) + return + + try: + # Assign only after all captures succeed. This keeps step/reset on + # one execution mode if any MJWarp operation is not capturable. + with _suspend_gc(), self._warp.ScopedDevice(device): + with self._warp.ScopedCapture() as step_capture: + self._mujoco_warp.step(self._device_model, self._device_data) + with self._warp.ScopedCapture() as forward_capture: + self._mujoco_warp.forward(self._device_model, self._device_data) + with self._warp.ScopedCapture() as reset_capture: + self._mujoco_warp.reset_data( + self._device_model, + self._device_data, + reset=self._reset_mask_device, + ) + step_graph = step_capture.graph + forward_graph = forward_capture.graph + reset_graph = reset_capture.graph + except Exception as exc: + reason = f"capture failed: {type(exc).__name__}: {exc}" + self._disable_cuda_graphs(reason) + warnings.warn( + f"mjwarp CUDA graphs disabled; using eager execution: {reason}", + RuntimeWarning, + stacklevel=2, + ) + return + + self._step_graph = step_graph + self._forward_graph = forward_graph + self._reset_graph = reset_graph + self._cuda_graph_enabled = True + self._cuda_graph_disable_reason = None + + def _execute_device_steps(self, nsteps: int) -> None: + """Advance fixed-shape device state through graph replay or eager calls.""" + if self._cuda_graph_enabled: + assert self._step_graph is not None + for _ in range(nsteps): + self._warp.capture_launch(self._step_graph) + return + for _ in range(nsteps): + self._mujoco_warp.step(self._device_model, self._device_data) + + def _execute_device_reset(self) -> None: + """Clear selected device rows before the host state upload.""" + if self._cuda_graph_enabled: + assert self._reset_graph is not None + self._warp.capture_launch(self._reset_graph) + return + self._mujoco_warp.reset_data( + self._device_model, + self._device_data, + reset=self._reset_mask_device, + ) + + def _execute_device_forward(self) -> None: + """Refresh kinematics after the host state upload.""" + if self._cuda_graph_enabled: + assert self._forward_graph is not None + self._warp.capture_launch(self._forward_graph) + return + self._mujoco_warp.forward(self._device_model, self._device_data) + def _validate_rows(self, env_indices: np.ndarray) -> np.ndarray: rows = np.asarray(env_indices, dtype=np.intp) if rows.ndim != 1: @@ -441,8 +576,7 @@ def _execute_host_step( control_upload_ms = (time.perf_counter() - t0) * 1000.0 t0 = time.perf_counter() - for _ in range(nsteps): - self._mujoco_warp.step(self._device_model, self._device_data) + self._execute_device_steps(nsteps) self._synchronize() physics_ms = (time.perf_counter() - t0) * 1000.0 @@ -473,11 +607,7 @@ def _execute_host_reset( self._reset_mask_host.fill(False) self._reset_mask_host[row_ids] = True self._upload(self._reset_mask_device, self._reset_mask_host) - self._mujoco_warp.reset_data( - self._device_model, - self._device_data, - reset=self._reset_mask_device, - ) + self._execute_device_reset() # Full-cache uploads are intentional for the host compatibility # profile: they preserve complement worlds after reset_data cleared # selected transient state, while keeping all D2H materialization at @@ -487,7 +617,7 @@ def _execute_host_reset( reset_upload_ms = (time.perf_counter() - t0) * 1000.0 t0 = time.perf_counter() - self._mujoco_warp.forward(self._device_model, self._device_data) + self._execute_device_forward() self._synchronize() reset_forward_ms = (time.perf_counter() - t0) * 1000.0 diff --git a/tests/base/test_mjwarp_backend.py b/tests/base/test_mjwarp_backend.py index 2a50e0a3d..5f1ee6b3b 100644 --- a/tests/base/test_mjwarp_backend.py +++ b/tests/base/test_mjwarp_backend.py @@ -48,18 +48,38 @@ def _stand_state(backend: Any, count: int) -> tuple[np.ndarray, np.ndarray]: return qpos.astype(np.float32), qvel -def test_real_cuda_init_reset_step() -> None: +def test_real_cuda_init_reset_step(monkeypatch: pytest.MonkeyPatch) -> None: backend = _backend(2) assert backend.backend_type == "mjwarp" assert backend.num_actuators == 29 assert backend.num_dof_vel == 29 + graph_launches: list[Any] = [] + original_capture_launch = backend._warp.capture_launch + + def capture_launch(graph: Any) -> None: + graph_launches.append(graph) + original_capture_launch(graph) + + monkeypatch.setattr(backend._warp, "capture_launch", capture_launch) + qpos, qvel = _stand_state(backend, 2) backend.set_state(np.asarray([0, 1], dtype=np.int32), qpos, qvel) before = backend.get_base_pos().copy() - result = backend.step(np.tile(qpos[0, -backend.num_actuators :], (2, 1)), nsteps=1) + result = backend.step(np.tile(qpos[0, -backend.num_actuators :], (2, 1)), nsteps=3) assert set(result["timing"]) == {"control_upload_ms", "physics_ms", "host_cache_refresh_ms"} + if backend._cuda_graph_enabled: + assert graph_launches == [ + backend._reset_graph, + backend._forward_graph, + backend._step_graph, + backend._step_graph, + backend._step_graph, + ] + else: + assert graph_launches == [] + assert backend._cuda_graph_disable_reason assert np.isfinite(backend.get_base_pos()).all() assert np.isfinite(backend.get_dof_pos()).all() assert np.isfinite(backend.get_sensor_data("torso_upvector")).all() diff --git a/tests/base/test_mjwarp_cuda_graph.py b/tests/base/test_mjwarp_cuda_graph.py new file mode 100644 index 000000000..463da920b --- /dev/null +++ b/tests/base/test_mjwarp_cuda_graph.py @@ -0,0 +1,202 @@ +"""Cold-path CUDA graph eligibility, capture, replay, and fallback tests.""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from unilab.base.backend.mjwarp.backend import MjwarpBackend, _cuda_graph_eligibility + + +class _FakeDevice: + def __init__(self, *, is_cuda: bool = True) -> None: + self.is_cuda = is_cuda + + +class _FakeContext: + def __enter__(self) -> "_FakeContext": + return self + + def __exit__(self, *_args: Any) -> None: + return None + + +class _FakeCapture(_FakeContext): + def __init__(self, graph: str, *, failure: Exception | None = None) -> None: + self.graph = graph + self._failure = failure + + def __enter__(self) -> "_FakeCapture": + if self._failure is not None: + raise self._failure + return self + + +class _FakeWarp: + def __init__( + self, + *, + driver: tuple[int, int] | None = (13, 0), + mempool: bool = True, + fail_capture_index: int | None = None, + ) -> None: + self.driver = driver + self.mempool = mempool + self.fail_capture_index = fail_capture_index + self.capture_count = 0 + self.launches: list[str] = [] + + def get_cuda_driver_version(self) -> tuple[int, int] | None: + return self.driver + + def is_mempool_enabled(self, _device: Any) -> bool: + return self.mempool + + def ScopedDevice(self, _device: Any) -> _FakeContext: # noqa: N802 + return _FakeContext() + + def ScopedCapture(self) -> _FakeCapture: # noqa: N802 + capture_index = self.capture_count + self.capture_count += 1 + failure = ( + RuntimeError("synthetic capture failure") + if capture_index == self.fail_capture_index + else None + ) + return _FakeCapture(f"graph-{capture_index}", failure=failure) + + def capture_launch(self, graph: str) -> None: + self.launches.append(graph) + + +class _FakeMujocoWarp: + def __init__(self) -> None: + self.step_calls = 0 + self.forward_calls = 0 + self.reset_calls = 0 + + def step(self, _model: Any, _data: Any) -> None: + self.step_calls += 1 + + def forward(self, _model: Any, _data: Any) -> None: + self.forward_calls += 1 + + def reset_data(self, _model: Any, _data: Any, *, reset: Any) -> None: + del reset + self.reset_calls += 1 + + +def _backend(warp: _FakeWarp, mujoco_warp: _FakeMujocoWarp) -> MjwarpBackend: + backend = MjwarpBackend.__new__(MjwarpBackend) + backend._warp = warp + backend._mujoco_warp = mujoco_warp + backend._device_model = object() + backend._device_data = object() + backend._reset_mask_device = object() + return backend + + +@pytest.mark.parametrize( + ("device", "driver", "mempool", "expected_reason"), + [ + (_FakeDevice(is_cuda=False), (13, 0), True, "not CUDA"), + (_FakeDevice(), None, True, "unavailable"), + (_FakeDevice(), (12, 3), True, "older than 12.4"), + (_FakeDevice(), (12, 4), False, "mempool is disabled"), + ], +) +def test_cuda_graph_eligibility_fails_closed( + device: _FakeDevice, + driver: tuple[int, int] | None, + mempool: bool, + expected_reason: str, +) -> None: + eligible, reason = _cuda_graph_eligibility( + _FakeWarp(driver=driver, mempool=mempool), + device, + ) + + assert eligible is False + assert reason is not None and expected_reason in reason + + +def test_cuda_graph_capture_replays_fixed_address_operations() -> None: + warp = _FakeWarp(driver=(12, 4), mempool=True) + mujoco_warp = _FakeMujocoWarp() + backend = _backend(warp, mujoco_warp) + + backend._initialize_cuda_graphs(_FakeDevice()) + + assert backend._cuda_graph_enabled is True + assert backend._cuda_graph_disable_reason is None + assert (backend._step_graph, backend._forward_graph, backend._reset_graph) == ( + "graph-0", + "graph-1", + "graph-2", + ) + capture_calls = ( + mujoco_warp.step_calls, + mujoco_warp.forward_calls, + mujoco_warp.reset_calls, + ) + + backend._execute_device_steps(3) + backend._execute_device_reset() + backend._execute_device_forward() + + assert warp.launches == ["graph-0", "graph-0", "graph-0", "graph-2", "graph-1"] + assert ( + mujoco_warp.step_calls, + mujoco_warp.forward_calls, + mujoco_warp.reset_calls, + ) == capture_calls + + +def test_ineligible_cuda_graph_warns_and_uses_eager_operations() -> None: + warp = _FakeWarp(driver=(12, 3), mempool=True) + mujoco_warp = _FakeMujocoWarp() + backend = _backend(warp, mujoco_warp) + + with pytest.warns(RuntimeWarning, match="older than 12.4"): + backend._initialize_cuda_graphs(_FakeDevice()) + + backend._execute_device_steps(2) + backend._execute_device_reset() + backend._execute_device_forward() + + assert backend._cuda_graph_enabled is False + assert warp.capture_count == 0 + assert warp.launches == [] + assert mujoco_warp.step_calls == 2 + assert mujoco_warp.reset_calls == 1 + assert mujoco_warp.forward_calls == 1 + + +def test_cuda_graph_capture_failure_atomically_falls_back_to_eager() -> None: + warp = _FakeWarp(fail_capture_index=1) + mujoco_warp = _FakeMujocoWarp() + backend = _backend(warp, mujoco_warp) + + with pytest.warns(RuntimeWarning, match="synthetic capture failure"): + backend._initialize_cuda_graphs(_FakeDevice()) + + assert backend._cuda_graph_enabled is False + assert backend._step_graph is None + assert backend._forward_graph is None + assert backend._reset_graph is None + assert "synthetic capture failure" in backend._cuda_graph_disable_reason + captured_calls = ( + mujoco_warp.step_calls, + mujoco_warp.forward_calls, + mujoco_warp.reset_calls, + ) + + backend._execute_device_steps(2) + backend._execute_device_reset() + backend._execute_device_forward() + + assert warp.launches == [] + assert mujoco_warp.step_calls == captured_calls[0] + 2 + assert mujoco_warp.forward_calls == captured_calls[1] + 1 + assert mujoco_warp.reset_calls == captured_calls[2] + 1