[Fix] Automatically work around OVPhysX multi-GPU device selection

source/isaaclab/changelog.d/jichuanh-mgpu-virtual-local-rank.minor.rst

@@ -0,0 +1,8 @@
+Added
+^^^^^
+
+* Added automatic single-GPU-per-worker execution to the multi-GPU launchers when the selected presets
+  include OVPhysX, which otherwise hangs runs on more than one GPU whenever an RTX renderer shares the
+  process. Every worker sees its own GPU as ``cuda:0`` and reports ``LOCAL_RANK=0`` while this is active,
+  so use the global rank to name per-rank files. This is applied automatically and needs no flag; it will
+  be removed once the pinned OVPhysX version contains the fix.

source/isaaclab/isaaclab/cli/multigpu.py

@@ -50,6 +50,12 @@
 SKRL_JAX_TORCHRUN_ONLY_ARGS = tuple(name for name in TORCHRUN_ARGS if name not in SKRL_JAX_ARGS)
 """``torchrun`` options that have no skrl JAX equivalent."""
 
+PRESET_SELECTOR_PREFIXES = ("presets=", "physics=")
+"""Hydra-style tokens that can name a physics backend, forwarded to the worker rather than parsed here."""
+
+VIRTUAL_LOCAL_RANK_PRESET = "ovphysx"
+"""Preset that needs one visible GPU per worker. See :func:`_use_virtual_local_rank`."""
+
 # torchelastic gives its own workers 30 s before it SIGKILLs them, so the graceful window stays
 # above that rather than preempting a shutdown that is already making progress.
 _POLL_INTERVAL_S = 0.2
@@ -206,6 +212,9 @@ def build_launch_command(args_cli: argparse.Namespace, worker_args: list[str], c
     if not args_cli.log_all_ranks and args_cli.local_ranks_filter is None:
         command.extend(("--local_ranks_filter", "0"))
 
+    if _use_virtual_local_rank(worker_args):
+        command.append("--virtual_local_rank")
+
     return command + _worker_argv(args_cli, worker_args, cfg)
 
 
@@ -229,7 +238,8 @@ def run_multigpu_cli(argv: list[str] | None, cfg: MultiGpuLauncherCfg) -> int:
         print(shlex.join(command))
         return 0
 
-    print(f"[INFO] Launching distributed workers with: {shlex.join(command)}")
+    # Flushed so the launcher preamble precedes the workers' inherited-fd output in a redirected log.
+    print(f"[INFO] Launching distributed workers with: {shlex.join(command)}", flush=True)
     return run_launch_command(command)
 
 
@@ -280,6 +290,28 @@ def _is_skrl_jax_launcher(args_cli: argparse.Namespace, worker_args: list[str],
     return _forwarded_arg_value(worker_args, "--ml_framework") == "jax"
 
 
+def _use_virtual_local_rank(worker_args: list[str]) -> bool:
+    """Return whether the presets select OVPhysX, which needs one visible GPU per worker.
+
+    Works around nvbug 6573426 in ovphysx <= 0.5.10; delete with the pin. Matches the physics backend
+    only, because ``renderer=rtx`` resolves to OVRTX inside the worker, and matches the raw name, so a
+    future ``ovphysx`` alias would bypass it.
+    """
+    selected: set[str] = set()
+    for arg in worker_args:
+        for prefix in PRESET_SELECTOR_PREFIXES:
+            if arg.startswith(prefix):
+                selected.update(name.strip() for name in arg[len(prefix) :].split(","))
+    if VIRTUAL_LOCAL_RANK_PRESET not in selected:
+        return False
+    print(
+        "[INFO] Presets select OVPhysX; giving every worker a single GPU as cuda:0 to avoid the "
+        "multi-GPU device-selection hang (nvbug 6573426).",
+        flush=True,
+    )
+    return True
+
+
 def _visible_cuda_device_count() -> int | None:
     """Return the number of visible CUDA devices on this node, or ``None`` if undetermined."""
     visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")

source/isaaclab/test/benchmark/test_multigpu_launcher.py

@@ -76,6 +76,47 @@ def test_multi_node_rendezvous_options_reach_torchrun():
     assert command[command.index("--rdzv_endpoint") + 1] == "host0:29400"
 
 
+def test_virtual_local_rank_reaches_torchrun_without_a_value():
+    """The OVPhysX workaround reaches torchrun as a bare flag and never reaches the workers."""
+    enabled = _command("training", ["--num_gpus", "2", "--task", "X", "presets=ovphysx,ovrtx"])
+
+    assert "--virtual_local_rank" in enabled[: enabled.index(multigpu.WORKER_SCRIPT)]
+    assert "True" not in enabled  # torchrun rejects the flag if it is given a value
+    assert "--virtual_local_rank" not in _worker_argv(enabled)
+
+
+def test_virtual_local_rank_is_not_a_launcher_option():
+    """The workaround is inferred from the presets, so the launcher honors no opt-in flag."""
+    command = _command("training", ["--num_gpus", "2", "--task", "X", "--virtual_local_rank"])
+
+    assert "--virtual_local_rank" not in command[: command.index(multigpu.WORKER_SCRIPT)]
+
+
+@pytest.mark.parametrize(
+    "preset_argv, expected",
+    [
+        (["presets=ovphysx,ovrtx"], True),
+        (["physics=ovphysx", "renderer=ovrtx"], True),
+        (["presets=ovrtx", "physics=ovphysx"], True),
+        # ``rtx`` only becomes OVRTX once the worker resolves it, so the launcher cannot match on it.
+        (["physics=ovphysx", "renderer=rtx"], True),
+        # Engaged but inert: no RTX renderer shares the process here.
+        (["presets=ovphysx"], True),
+        (["presets=ovphysx,newton_renderer"], True),
+        (["presets=newton_mjwarp,ovrtx"], False),
+        (["presets=newton_mjwarp"], False),
+        ([], False),
+    ],
+)
+def test_virtual_local_rank_follows_the_ovphysx_preset(preset_argv: list[str], expected: bool):
+    """Selecting OVPhysX enables the workaround regardless of which renderer the worker resolves."""
+    command = _command("training", ["--num_gpus", "2", "--task", "X", *preset_argv])
+
+    assert ("--virtual_local_rank" in command) is expected
+    # The presets are read, not claimed, so the worker still receives them verbatim.
+    assert all(token in _worker_argv(command) for token in preset_argv)
+
+
 def test_dry_run_prints_a_shell_parsable_command(capsys: pytest.CaptureFixture[str]):
     """``--dry_run`` reports the exact command instead of launching workers."""
     status = multigpu.run_multigpu_benchmark_cli("startup", ["--dry_run", "--num_gpus", "2", "--task", "X"])