fix: address review feedback on multi-GPU PR

- Fix resolve_cuda_device() multi-node bug: compare local_rank against
  num_visible (not WORLD_SIZE) so 2-node x 4-GPU setups work correctly.
  Remove unused os import and WORLD_SIZE env var.

- Revert presets.py to upstream: remove newton and physx fields from
  MultiBackendRendererCfg. The global 'newton' preset resolves physics,
  not renderer — renderer resolves via 'newton_renderer' field. Fix
  isaacsim_rtx_renderer to use proper type annotation (independent
  instance, not shared alias).

- Revert _is_kit_camera() PresetCfg escape hatch: PresetCfg renderers
  are always resolved by Hydra before launch_simulation() is called,
  so the check is unnecessary. Restore original upstream logic.

- Rewrite tests: test resolve_cuda_device() by calling the actual
  function with mocked torch.cuda.device_count(). Add multi-node test
  case. Remove vacuous inline-logic test and redundant skipif.

Author: ClawLabby

source/isaaclab/isaaclab/utils/distributed.py

@@ -7,8 +7,6 @@
 
 from __future__ import annotations
 
-import os
-
 import torch
 
 
@@ -27,9 +25,10 @@ def resolve_cuda_device(local_rank: int) -> tuple[str, int]:
         ``("cuda:2", 2)``.
     """
     num_visible = torch.cuda.device_count()
-    world_size = int(os.getenv("WORLD_SIZE", "1"))
-    if num_visible >= world_size:
+    if local_rank < num_visible:
         device_id = local_rank
     else:
+        # CUDA_VISIBLE_DEVICES restricts this process to fewer GPUs than
+        # local_rank implies — fall back to cuda:0.
         device_id = 0
     return f"cuda:{device_id}", device_id

source/isaaclab/test/test_multigpu_newton.py

@@ -19,39 +19,58 @@ def _get_num_gpus() -> int:
     return torch.cuda.device_count()
 
 
-@pytest.mark.skipif(_get_num_gpus() < 2, reason="Requires at least 2 GPUs")
-class TestMultiGPUDeviceAssignment:
-    """Tests for multi-GPU device assignment in distributed training."""
+class TestResolveCudaDevice:
+    """Tests for resolve_cuda_device across different GPU visibility scenarios."""
 
-    def test_resolve_cuda_device_full_visibility(self):
+    def test_full_visibility(self):
         """When all GPUs are visible, each rank gets its own device."""
         from isaaclab.utils.distributed import resolve_cuda_device
 
-        # Simulate 4 visible GPUs, world_size=4
-        with unittest.mock.patch("torch.cuda.device_count", return_value=4), \
-             unittest.mock.patch.dict(os.environ, {"WORLD_SIZE": "4"}):
+        with unittest.mock.patch("torch.cuda.device_count", return_value=4):
             for local_rank in range(4):
                 device, device_id = resolve_cuda_device(local_rank)
                 assert device == f"cuda:{local_rank}"
                 assert device_id == local_rank
 
-    def test_resolve_cuda_device_restricted_visibility(self):
+    def test_restricted_visibility(self):
         """When CUDA_VISIBLE_DEVICES restricts each process to 1 GPU, all ranks use cuda:0."""
         from isaaclab.utils.distributed import resolve_cuda_device
 
-        # Simulate 1 visible GPU (CUDA_VISIBLE_DEVICES restricted), world_size=4
-        with unittest.mock.patch("torch.cuda.device_count", return_value=1), \
-             unittest.mock.patch.dict(os.environ, {"WORLD_SIZE": "4"}):
+        with unittest.mock.patch("torch.cuda.device_count", return_value=1):
             for local_rank in range(4):
                 device, device_id = resolve_cuda_device(local_rank)
                 assert device == "cuda:0"
                 assert device_id == 0
 
+    def test_multi_node(self):
+        """Multi-node: 4 visible GPUs per node, world_size=8, each rank gets its own GPU."""
+        from isaaclab.utils.distributed import resolve_cuda_device
+
+        with unittest.mock.patch("torch.cuda.device_count", return_value=4):
+            for local_rank in range(4):
+                device, device_id = resolve_cuda_device(local_rank)
+                assert device == f"cuda:{local_rank}", (
+                    f"Multi-node: local_rank={local_rank} should map to cuda:{local_rank}"
+                )
+
+    def test_local_rank_exceeds_visible(self):
+        """When local_rank >= num_visible, fall back to cuda:0."""
+        from isaaclab.utils.distributed import resolve_cuda_device
+
+        with unittest.mock.patch("torch.cuda.device_count", return_value=2):
+            device, device_id = resolve_cuda_device(5)
+            assert device == "cuda:0"
+            assert device_id == 0
+
+
+@pytest.mark.skipif(_get_num_gpus() < 2, reason="Requires at least 2 GPUs")
+class TestMultiGPUDeviceAssignment:
+    """Integration tests for multi-GPU training."""
+
     def test_cartpole_newton_multigpu(self):
         """Test that multi-GPU cartpole training with Newton physics runs without error."""
         num_gpus = min(_get_num_gpus(), 4)
 
-        # Run a quick 2-iteration training to verify setup works
         cmd = [
             "torchrun",
             f"--nproc_per_node={num_gpus}",
@@ -65,12 +84,9 @@ def test_cartpole_newton_multigpu(self):
         ]
 
         env = os.environ.copy()
-        # Required for containers with IOMMU where P2P transport fails,
-        # and for environments without InfiniBand.
         env["NCCL_P2P_DISABLE"] = "1"
         env["NCCL_IB_DISABLE"] = "1"
 
-        # Resolve IsaacLab root from test file location
         test_dir = os.path.dirname(os.path.abspath(__file__))
         isaaclab_root = os.path.dirname(os.path.dirname(os.path.dirname(test_dir)))
 
@@ -83,7 +99,6 @@ def test_cartpole_newton_multigpu(self):
             cwd=isaaclab_root,
         )
 
-        # Verify training actually ran (not just a clean exit from early skip)
         has_training_output = "Learning iteration 1" in result.stdout
         no_cuda_errors = "CUDA error" not in result.stderr
         assert result.returncode == 0 and has_training_output and no_cuda_errors, (
@@ -97,42 +112,25 @@ def test_cartpole_newton_multigpu(self):
 class TestMultiGPUCameraRendering:
     """Tests for multi-GPU training with camera observations."""
 
-    def test_preset_renderer_matching(self):
-        """Test that newton preset correctly matches the renderer."""
+    def test_preset_renderer_has_newton_renderer_field(self):
+        """Test that MultiBackendRendererCfg has newton_renderer field."""
         from isaaclab_tasks.utils.presets import MultiBackendRendererCfg
 
         cfg = MultiBackendRendererCfg()
-        assert hasattr(cfg, "newton"), "MultiBackendRendererCfg should have 'newton' field"
-        assert "NewtonWarp" in type(cfg.newton).__name__, (
-            f"newton field should be NewtonWarpRendererCfg, got {type(cfg.newton).__name__}"
+        assert hasattr(cfg, "newton_renderer"), (
+            "MultiBackendRendererCfg should have 'newton_renderer' field"
         )
-
-    def test_physx_preset_is_independent_instance(self):
-        """Test that physx preset is a separate instance from default (not a shared alias)."""
-        from isaaclab_tasks.utils.presets import MultiBackendRendererCfg
-
-        cfg = MultiBackendRendererCfg()
-        assert cfg.physx is not cfg.default, (
-            "physx and default should be independent instances, not shared aliases"
+        assert "NewtonWarp" in type(cfg.newton_renderer).__name__, (
+            f"newton_renderer should be NewtonWarpRendererCfg, got {type(cfg.newton_renderer).__name__}"
         )
 
-    def test_camera_not_kit_with_preset_renderer(self):
-        """Test that PresetCfg renderers are not detected as Kit cameras."""
-        from isaaclab.sensors import TiledCameraCfg
+    def test_alias_fields_are_independent_instances(self):
+        """Test that aliased fields are separate instances (not shared mutable references)."""
         from isaaclab_tasks.utils.presets import MultiBackendRendererCfg
-        from isaaclab_tasks.utils.sim_launcher import _is_kit_camera
-
-        cam = TiledCameraCfg(
-            prim_path="/test",
-            data_types=["rgb"],
-            width=64,
-            height=64,
-            renderer_cfg=MultiBackendRendererCfg(),
-        )
 
-        # PresetCfg renderers resolve to match the physics backend, so not necessarily Kit
-        assert not _is_kit_camera(cam), (
-            "Camera with PresetCfg renderer should not be detected as Kit camera"
+        cfg = MultiBackendRendererCfg()
+        assert cfg.isaacsim_rtx_renderer is not cfg.default, (
+            "isaacsim_rtx_renderer and default should be independent instances"
         )
 
 

source/isaaclab_tasks/isaaclab_tasks/utils/presets.py

@@ -15,9 +15,6 @@
 @configclass
 class MultiBackendRendererCfg(PresetCfg):
     default: IsaacRtxRendererCfg = IsaacRtxRendererCfg()
-    newton: NewtonWarpRendererCfg = NewtonWarpRendererCfg()
     newton_renderer: NewtonWarpRendererCfg = NewtonWarpRendererCfg()
     ovrtx_renderer: OVRTXRendererCfg = OVRTXRendererCfg()
     isaacsim_rtx_renderer: IsaacRtxRendererCfg = IsaacRtxRendererCfg()
-    # PhysX uses RTX rendering by default in Isaac Lab.
-    physx: IsaacRtxRendererCfg = IsaacRtxRendererCfg()

source/isaaclab_tasks/isaaclab_tasks/utils/sim_launcher.py

@@ -122,12 +122,6 @@ def _is_kit_camera(node) -> bool:
         return True
     if isinstance(renderer_cfg, RendererCfg):
         return renderer_cfg.renderer_type in ("default", "isaac_rtx")
-    # PresetCfg renderers (e.g. MultiBackendRendererCfg) are resolved later;
-    # assume they will match the physics backend, so not necessarily Kit.
-    from isaaclab_tasks.utils import PresetCfg
-
-    if isinstance(renderer_cfg, PresetCfg):
-        return False
     # Unknown renderer type — conservatively assume Kit is required.
     return True