formatting

Author: brb-nv

tests/unittest/_torch/distributed/test_cp_broadcast.py

@@ -5,14 +5,13 @@
 when CP (context parallelism) is enabled (e.g., in Helix parallelism).
 
 For MPIDist tests, run with mpirun:
-    mpirun -n 2 python -m pytest tests/unittest/_torch/distributed/test_cp_broadcast.py -v
+mpirun -n 2 python -m pytest tests/unittest/_torch/distributed/test_cp_broadcast.py -v
 
 For TorchDist tests, see test_ops.py which uses Ray for distributed testing.
 """
 
 import numpy as np
 import pytest
-import torch
 
 from tensorrt_llm._torch.distributed import MPIDist
 from tensorrt_llm.mapping import Mapping
@@ -22,6 +21,7 @@ def get_mpi_info():
     """Get MPI rank and world size, returns (0, 1) if MPI is not available."""
     try:
         from mpi4py import MPI
+
         comm = MPI.COMM_WORLD
         return comm.Get_rank(), comm.Get_size()
     except ImportError:
@@ -68,9 +68,9 @@ def test_broadcast_numpy_array(self):
 
         # Store original data from root for verification
         from mpi4py import MPI
+
         expected = np.zeros(shape, dtype=np.float32)
-        MPI.COMM_WORLD.Bcast(data if self.mapping.cp_rank == root else expected,
-                             root=root)
+        MPI.COMM_WORLD.Bcast(data if self.mapping.cp_rank == root else expected, root=root)
         if self.mapping.cp_rank == root:
             expected = data.copy()
 
@@ -89,7 +89,7 @@ def test_broadcast_python_dict(self):
                 "model_name": "llama",
                 "batch_size": 32,
                 "tokens": [1, 2, 3, 4, 5],
-                "config": {"hidden_size": 4096, "num_layers": 32}
+                "config": {"hidden_size": 4096, "num_layers": 32},
             }
         else:
             obj = None
@@ -168,8 +168,6 @@ def test_broadcast_string(self):
 def test_mpi_cp_broadcast_integration():
     """
     Integration test for MPIDist cp_broadcast.
-
-    Run with: mpirun -n 2 python -m pytest tests/unittest/_torch/distributed/test_cp_broadcast.py::test_mpi_cp_broadcast_integration -v
     """
     rank, world_size = get_mpi_info()
     if world_size < 2:
@@ -210,4 +208,3 @@ def test_mpi_cp_broadcast_integration():
 if __name__ == "__main__":
     # Allow running directly with mpirun
     pytest.main([__file__, "-v"])
-

tests/unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py

@@ -339,8 +339,7 @@ def run_object_broadcast(self, root_obj, root: int = 0):
 
 
 @pytest.mark.gpu2
-@pytest.mark.parametrize("hidden_size", [128, 512],
-                         ids=lambda x: f"hidden:{x}")
+@pytest.mark.parametrize("hidden_size", [128, 512], ids=lambda x: f"hidden:{x}")
 @pytest.mark.parametrize("seq_len", [16, 32], ids=lambda x: f"seqlen:{x}")
 def test_cp_broadcast_tensor(setup_ray_cluster, seq_len, hidden_size):
     """Test TorchDist.cp_broadcast with tensor data."""
@@ -382,10 +381,16 @@ def test_cp_broadcast_tensor(setup_ray_cluster, seq_len, hidden_size):
 
 @pytest.mark.gpu2
 @pytest.mark.parametrize("test_object", [
-    {"key1": "value1", "key2": [1, 2, 3]},
-    ["item1", "item2", {"nested": True}],
+    {
+        "key1": "value1",
+        "key2": [1, 2, 3]
+    },
+    ["item1", "item2", {
+        "nested": True
+    }],
     "simple_string",
-], ids=["dict", "list", "string"])
+],
+                         ids=["dict", "list", "string"])
 def test_cp_broadcast_object(setup_ray_cluster, test_object):
     """Test TorchDist.cp_broadcast with non-tensor objects."""
     world_size = 2