Fix benchmark timing validation

Author: AntoineRichard

scripts/benchmarks/rl_games/benchmark_rl_games_play.py

@@ -86,6 +86,8 @@ def _parse_args(argv: list[str]):
         parser.error("--num_frames must be greater than zero")
     if args_cli.warmup_steps < 0:
         parser.error("--warmup_steps must be non-negative")
+    if args_cli.warmup_steps >= args_cli.num_frames:
+        parser.error("--warmup_steps must be less than --num_frames")
     sys.argv = [sys.argv[0]] + remaining_args
 
     return args_cli, remaining_args

scripts/benchmarks/rsl_rl/benchmark_rsl_rl_play.py

@@ -88,6 +88,8 @@ def _parse_args(argv: list[str]) -> tuple[argparse.Namespace, list[str]]:
         parser.error("--num_frames must be greater than zero")
     if args.warmup_steps < 0:
         parser.error("--warmup_steps must be non-negative")
+    if args.warmup_steps >= args.num_frames:
+        parser.error("--warmup_steps must be less than --num_frames")
     sys.argv = [sys.argv[0]] + remaining
     return args, remaining
 

scripts/benchmarks/sb3/benchmark_sb3_play.py

@@ -94,6 +94,8 @@ def _parse_args(argv: list[str]):
         parser.error("--num_frames must be greater than zero")
     if args_cli.warmup_steps < 0:
         parser.error("--warmup_steps must be non-negative")
+    if args_cli.warmup_steps >= args_cli.num_frames:
+        parser.error("--warmup_steps must be less than --num_frames")
     sys.argv = [sys.argv[0]] + remaining_args
 
     return args_cli, remaining_args

scripts/benchmarks/skrl/benchmark_skrl_play.py

@@ -107,6 +107,8 @@ def _parse_args(argv: list[str]):
         parser.error("--num_frames must be greater than zero")
     if args_cli.warmup_steps < 0:
         parser.error("--warmup_steps must be non-negative")
+    if args_cli.warmup_steps >= args_cli.num_frames:
+        parser.error("--warmup_steps must be less than --num_frames")
     sys.argv = [sys.argv[0]] + remaining_args
 
     return args_cli, remaining_args

scripts/benchmarks/test/test_benchmark_smoke.py

@@ -8,6 +8,7 @@
 import importlib.util
 import json
 import subprocess
+import sys
 from pathlib import Path
 
 import pytest
@@ -48,26 +49,46 @@ def _load_adapter(library: str, workflow: str):
 
 @pytest.mark.parametrize("library", ["rsl_rl", "rl_games", "skrl", "sb3"])
 @pytest.mark.parametrize(("workflow", "argument"), [("play", "--num_frames"), ("train", "--max_iterations")])
-def test_adapters_reject_non_positive_workloads(library: str, workflow: str, argument: str):
+def test_adapters_reject_non_positive_workloads(library: str, workflow: str, argument: str, monkeypatch, capsys):
     """Benchmark adapters reject workloads that cannot produce timing samples."""
     module = _load_adapter(library, workflow)
+    argv = ["--task", _TASK, argument, "0", "--headless"]
+    monkeypatch.setattr(sys, "argv", ["benchmark", *argv])
 
     with pytest.raises(SystemExit) as exc_info:
-        module._parse_args(["--task", _TASK, argument, "0", "--headless"])
+        module._parse_args(argv)
 
     assert exc_info.value.code == 2
+    assert f"{argument} must be greater than zero" in capsys.readouterr().err
 
 
 @pytest.mark.parametrize("library", ["rsl_rl", "rl_games", "skrl", "sb3"])
 @pytest.mark.parametrize("workflow", ["play", "train"])
-def test_adapters_reject_negative_warmup_steps(library: str, workflow: str):
+def test_adapters_reject_negative_warmup_steps(library: str, workflow: str, monkeypatch, capsys):
     """Benchmark adapters reject a negative --warmup_steps (the cold-start exclusion is opt-in, N >= 0)."""
     module = _load_adapter(library, workflow)
+    argv = ["--task", _TASK, "--warmup_steps", "-1", "--headless"]
+    monkeypatch.setattr(sys, "argv", ["benchmark", *argv])
 
     with pytest.raises(SystemExit) as exc_info:
-        module._parse_args(["--task", _TASK, "--warmup_steps", "-1", "--headless"])
+        module._parse_args(argv)
 
     assert exc_info.value.code == 2
+    assert "--warmup_steps must be non-negative" in capsys.readouterr().err
+
+
+@pytest.mark.parametrize("library", ["rsl_rl", "rl_games", "skrl", "sb3"])
+def test_play_adapters_reject_warmup_that_exhausts_workload(library: str, monkeypatch, capsys):
+    """Play requires at least one measured environment step after warm-up."""
+    module = _load_adapter(library, "play")
+    argv = ["--task", _TASK, "--num_frames", "2", "--warmup_steps", "2", "--headless"]
+    monkeypatch.setattr(sys, "argv", ["benchmark", *argv])
+
+    with pytest.raises(SystemExit) as exc_info:
+        module._parse_args(argv)
+
+    assert exc_info.value.code == 2
+    assert "--warmup_steps must be less than --num_frames" in capsys.readouterr().err
 
 
 @pytest.mark.parametrize(

source/isaaclab/isaaclab/test/benchmark/builders.py

@@ -178,12 +178,19 @@ def build_runtime(
         collection_fps_agg = _with_effective_mean(collection_fps_agg, collection_fps, effective_fps)
         total_fps_agg = _with_effective_mean(total_fps_agg, total_fps, effective_fps)
         iterations_per_s_agg = _with_effective_mean(iterations_per_s_agg, iter_per_s, effective_iterations_per_s)
+    if environment_step_times_s is None and (simulation_step_times_s is not None or simulation_step_calls is not None):
+        raise ValueError("environment_step_times_s is required with simulation timing")
     environment_step_timing = None
     if environment_step_times_s is not None:
         if frames_per_environment_step is None or frames_per_environment_step <= 0:
             raise ValueError("frames_per_environment_step must be greater than zero")
         environment_samples = list(environment_step_times_s)
-        if not environment_samples or any(value <= 0 for value in environment_samples):
+        if not environment_samples:
+            raise ValueError(
+                "environment_step_times_s must contain only positive samples; no samples remained after warm-up, "
+                "so reduce warmup_steps or increase the workload"
+            )
+        if any(value <= 0 for value in environment_samples):
             raise ValueError("environment_step_times_s must contain only positive samples")
         total_environment_time_s = sum(environment_samples)
         environment_fps_samples = [frames_per_environment_step / value for value in environment_samples]

source/isaaclab/isaaclab/test/benchmark/schema.py

@@ -246,6 +246,13 @@ def __post_init__(self) -> None:
         """Validate that the timing fields form one consistent measurement mode."""
         if self.environment_step_calls <= 0:
             raise ValueError("environment_step_calls must be greater than zero")
+        if not (
+            math.isfinite(self.environment_step_time_s.mean)
+            and self.environment_step_time_s.mean > 0.0
+            and math.isfinite(self.environment_step_fps.mean)
+            and self.environment_step_fps.mean > 0.0
+        ):
+            raise ValueError("environment step time and FPS must be greater than zero")
 
         breakdown = (
             self.simulation_step_time_s,

source/isaaclab/isaaclab/test/benchmark/stepping.py

@@ -60,6 +60,35 @@ def sample_random_actions(env) -> torch.Tensor | dict[str, torch.Tensor]:
         return 2.0 * torch.rand(u.num_envs, u.single_action_space.shape[0], device=u.device) - 1.0
 
 
+def _find_cuda_devices(value) -> set[str]:
+    """Collect CUDA device names from nested tensor-like values."""
+    import torch  # noqa: PLC0415
+
+    devices: set[str] = set()
+
+    def collect(item) -> None:
+        if isinstance(item, dict):
+            for nested in item.values():
+                collect(nested)
+            return
+        if isinstance(item, (list, tuple)):
+            for nested in item:
+                collect(nested)
+            return
+        device = item if isinstance(item, (str, torch.device)) else getattr(item, "device", None)
+        if device is None:
+            return
+        try:
+            device = torch.device(device)
+        except (RuntimeError, TypeError):
+            return
+        if device.type == "cuda":
+            devices.add(str(device))
+
+    collect(value)
+    return devices
+
+
 class EnvironmentStepTimingRecorder(AbstractContextManager):
     """Record host-return step time or an optional synchronized step breakdown.
 
@@ -119,6 +148,7 @@ def __enter__(self) -> EnvironmentStepTimingRecorder:
         self._original_env_step = self._env.step
 
         if self._measure_synchronized_step_breakdown:
+            import torch  # noqa: PLC0415
             import warp as wp  # noqa: PLC0415
 
             from isaaclab.utils.timer import Timer  # noqa: PLC0415
@@ -130,34 +160,50 @@ def __enter__(self) -> EnvironmentStepTimingRecorder:
             self._simulation_step_calls = 0
             assert self._simulation_context is not None
             self._original_sim_step = self._simulation_context.step
+            environment_cuda_devices = _find_cuda_devices(getattr(self._env.unwrapped, "device", None))
+            active_cuda_devices = environment_cuda_devices
+
+            def synchronize_torch(devices: set[str]) -> None:
+                for device in sorted(devices):
+                    torch.cuda.synchronize(device)
 
             def timed_simulation_step(*args, **kwargs):
                 if not self._inside_environment_step:
                     return self._original_sim_step(*args, **kwargs)
+                synchronize_torch(active_cuda_devices)
                 wp.synchronize()
                 timer = Timer()
+                timer.start()
                 try:
-                    with timer:
-                        return self._original_sim_step(*args, **kwargs)
+                    return self._original_sim_step(*args, **kwargs)
                 finally:
+                    synchronize_torch(active_cuda_devices)
+                    timer.stop()
                     self._simulation_total_time_s += timer.total_run_time
                     self._simulation_step_calls += 1
 
             self._simulation_context.step = timed_simulation_step
 
             def timed_environment_step(*args, **kwargs):
+                nonlocal active_cuda_devices
                 recording = self._environment_step_index >= self._warmup_steps
                 self._environment_step_index += 1
                 simulation_start_time_s = self._simulation_total_time_s
                 simulation_start_calls = self._simulation_step_calls
+                previous_cuda_devices = active_cuda_devices
+                active_cuda_devices = environment_cuda_devices | _find_cuda_devices(args) | _find_cuda_devices(kwargs)
+                synchronize_torch(active_cuda_devices)
                 wp.synchronize()
                 timer = Timer()
+                timer.start()
                 self._inside_environment_step = True
                 try:
-                    with timer:
-                        return self._original_env_step(*args, **kwargs)
+                    return self._original_env_step(*args, **kwargs)
                 finally:
                     self._inside_environment_step = False
+                    synchronize_torch(active_cuda_devices)
+                    timer.stop()
+                    active_cuda_devices = previous_cuda_devices
                     if recording:
                         self.step_times_s.append(timer.total_run_time)
                         self.simulation_step_times_s.append(self._simulation_total_time_s - simulation_start_time_s)

source/isaaclab/test/benchmark/test_builders.py

@@ -221,6 +221,25 @@ def test_build_runtime_rejects_non_positive_environment_step_times(environment_s
         )
 
 
+@pytest.mark.parametrize(
+    ("simulation_step_times_s", "simulation_step_calls"),
+    [([0.5], None), (None, 1), ([0.5], 1)],
+)
+def test_build_runtime_rejects_simulation_timing_without_environment_timing(
+    simulation_step_times_s, simulation_step_calls
+):
+    with pytest.raises(ValueError, match="environment_step_times_s is required with simulation timing"):
+        builders.build_runtime(
+            startup_time_s=StartupTime(0.1, 0.2, 0.3),
+            iteration_times_s=[1.0],
+            collection_fps=[8.0],
+            total_fps=[8.0],
+            steps_per_iteration=8,
+            simulation_step_times_s=simulation_step_times_s,
+            simulation_step_calls=simulation_step_calls,
+        )
+
+
 def test_build_training_bundle_round_trips(tmp_path):
     run = builders.build_run_identity(
         run_id="x",

source/isaaclab/test/benchmark/test_schema.py

@@ -172,6 +172,23 @@ def test_environment_step_timing_rejects_incomplete_measurement_modes():
         dataclasses.replace(timing, simulation_step_time_s=None)
 
 
+@pytest.mark.parametrize("field", ["environment_step_time_s", "environment_step_fps"])
+def test_environment_step_timing_rejects_non_positive_environment_metrics(field):
+    timing = EnvironmentStepTiming(
+        environment_step_time_s=MeanStd(mean=0.1, std=0.01),
+        environment_step_fps=MeanStd(mean=100.0, std=1.0),
+        simulation_step_time_s=None,
+        outside_simulation_step_time_s=None,
+        outside_simulation_step_fraction=None,
+        environment_step_calls=1,
+        simulation_step_calls=None,
+        measurement_mode="host_return",
+    )
+
+    with pytest.raises(ValueError, match="environment step time and FPS must be greater than zero"):
+        dataclasses.replace(timing, **{field: MeanStd(mean=-1.0, std=0.0)})
+
+
 def test_environment_step_timing_rejects_inconsistent_partition():
     timing = _runtime().environment_step_timing
     assert timing is not None