update

Author: msaroufim

.github/workflows/amd-health.yml

@@ -6,7 +6,6 @@ on:
     - cron: '0 2 * * *'
   workflow_dispatch:
   push:
-    branches: [main]
 
 jobs:
   health-check:
@@ -28,64 +27,5 @@ jobs:
 
     - name: Distributed Health Check
       run: |
-        # Check how many GPUs are available
         python -c "import torch; print(f'Available GPUs: {torch.cuda.device_count()}')"
-        
-        # Test process group initialization in a loop to debug hanging issues
-        python -c "
-        import torch
-        import torch.distributed as dist
-        import os
-        import time
-        import signal
-        
-        def timeout_handler(signum, frame):
-            print('✗ Process group initialization timed out after 30 seconds')
-            exit(1)
-        
-        # Set timeout for process group initialization
-        signal.signal(signal.SIGALRM, timeout_handler)
-        
-        num_gpus = torch.cuda.device_count()
-        print(f'Testing process group initialization on {num_gpus} GPUs')
-        
-        for attempt in range(3):  # Try 3 times
-            try:
-                print(f'Attempt {attempt + 1}: Initializing process group...')
-                
-                # Set environment variables
-                os.environ['MASTER_ADDR'] = '127.0.0.1'
-                os.environ['MASTER_PORT'] = str(12345 + attempt)
-                os.environ['WORLD_SIZE'] = '1'
-                os.environ['RANK'] = '0'
-                
-                # Set 30 second timeout
-                signal.alarm(30)
-                
-                # Test single-process initialization first
-                dist.init_process_group('nccl', rank=0, world_size=1)
-                
-                # Cancel timeout
-                signal.alarm(0)
-                
-                print(f'✓ Attempt {attempt + 1}: Process group initialized successfully')
-                
-                # Test basic tensor operations
-                device = torch.device('cuda:0')
-                tensor = torch.ones(10, device=device)
-                print(f'✓ Tensor operations work: {tensor.sum().item()}')
-                
-                dist.destroy_process_group()
-                print(f'✓ Attempt {attempt + 1}: Process group destroyed successfully')
-                break
-                
-            except Exception as e:
-                signal.alarm(0)  # Cancel timeout
-                print(f'✗ Attempt {attempt + 1} failed: {type(e).__name__}: {e}')
-                if attempt == 2:  # Last attempt
-                    print('✗ All initialization attempts failed')
-                    exit(1)
-                time.sleep(2)  # Wait before retry
-        
-        print('✓ Distributed health check passed')
-        "
+        python scripts/test_distributed.py

scripts/test_distributed.py

@@ -0,0 +1,75 @@
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import os
+import signal
+import sys
+
+def timeout_handler(signum, frame):
+    print('✗ TIMEOUT: Process hung')
+    sys.exit(1)
+
+def test_worker(rank, world_size, master_port):
+    try:
+        os.environ['MASTER_ADDR'] = '127.0.0.1'
+        os.environ['MASTER_PORT'] = str(master_port)
+        os.environ['RANK'] = str(rank)
+        os.environ['WORLD_SIZE'] = str(world_size)
+        
+        signal.signal(signal.SIGALRM, timeout_handler)
+        signal.alarm(30)
+        
+        print(f'Rank {rank}: Init NCCL...')
+        dist.init_process_group('nccl', rank=rank, world_size=world_size)
+        signal.alarm(0)
+        
+        device = torch.device(f'cuda:{rank}')
+        tensor = torch.ones(100, device=device) * rank
+        
+        signal.alarm(15)
+        dist.all_reduce(tensor)
+        signal.alarm(0)
+        
+        print(f'✓ Rank {rank}: sum = {tensor[0].item()}')
+        dist.destroy_process_group()
+        
+    except Exception as e:
+        signal.alarm(0)
+        print(f'✗ Rank {rank}: {e}')
+        sys.exit(1)
+
+def main():
+    num_gpus = torch.cuda.device_count()
+    print(f'Testing {num_gpus} GPUs - 4 rounds')
+    
+    for round_num in range(4):
+        print(f'=== ROUND {round_num + 1} ===')
+        master_port = 29500 + round_num
+        
+        mp.set_start_method('spawn', force=True)
+        processes = []
+        
+        for rank in range(num_gpus):
+            p = mp.Process(target=test_worker, args=(rank, num_gpus, master_port))
+            p.start()
+            processes.append(p)
+        
+        for i, p in enumerate(processes):
+            p.join(timeout=60)
+            if p.exitcode != 0:
+                print(f'✗ ROUND {round_num + 1} FAILED')
+                for rp in processes:
+                    if rp.is_alive():
+                        rp.terminate()
+                sys.exit(1)
+            elif p.is_alive():
+                print(f'✗ ROUND {round_num + 1} HUNG')
+                p.terminate()
+                sys.exit(1)
+        
+        print(f'✓ ROUND {round_num + 1} PASSED')
+    
+    print('✓ ALL ROUNDS PASSED')
+
+if __name__ == '__main__':
+    main()