Fixed the CI issue

Author: git-flyer

.github/workflows/nvidia3.6-build-and-test.yml

@@ -132,7 +132,8 @@ jobs:
           python3 -m pytest -s python/test/tle/unit \
               --ignore=python/test/tle/unit/test_tle_distributed_d2d.py \
               --ignore=python/test/tle/unit/test_tle_get_local_pe.py \
-              --ignore=python/test/tle/unit/test_tle_d2d_barrier.py
+              --ignore=python/test/tle/unit/test_tle_d2d_barrier.py \
+              --ignore=python/test/tle/unit/test_tle_get_node_rank.py \
           ## flagtree hints python tutorials
           python3 python/tutorials/hints/01/01-vector-add.py --only_unit_test
           # python3 python/tutorials/hints/02/02-fused-softmax.py --only_unit_test

python/test/tle/unit/test_tle_get_node_rank.py

@@ -6,12 +6,15 @@
 import triton.experimental.tle.language as tle
 import triton.language as tl
 
+
 LOCAL_WORLD_SIZE = int(os.environ["LOCAL_WORLD_SIZE"])
 WORLD_SIZE = int(os.environ["WORLD_SIZE"])
 if WORLD_SIZE % LOCAL_WORLD_SIZE != 0:
     raise ValueError("WORLD_SIZE must be divisible by LOCAL_WORLD_SIZE")
 
-DEVICE_MESH = tle.device_mesh(tle.MeshConfig(node=WORLD_SIZE // LOCAL_WORLD_SIZE, device=LOCAL_WORLD_SIZE))
+DEVICE_MESH = tle.device_mesh(
+    tle.MeshConfig(node=WORLD_SIZE // LOCAL_WORLD_SIZE, device=LOCAL_WORLD_SIZE)
+)
 
 
 @triton.jit
@@ -49,7 +52,7 @@ def test_tle_get_node_rank():
     torch.cuda.synchronize()
 
     rank = dist.get_rank()
-    expected_node_rank = int(os.environ["GROUP_RANK"])
+    expected_node_rank = rank // LOCAL_WORLD_SIZE
     actual_node_ranks = node_rank_out.cpu().tolist()
     try:
         torch.testing.assert_close(
@@ -73,4 +76,5 @@ def test_tle_get_node_rank():
         tle.cleanup_communicator()
 
 
-test_tle_get_node_rank()
+if __name__ == "__main__":
+    test_tle_get_node_rank()