multi_host_test.py

#!/usr/bin/env python3
"""
Multi-host TPU pod test script.
This script should be run on ALL hosts simultaneously.

Usage with gcloud:
  gcloud compute tpus tpu-vm ssh YOUR_TPU \
    --zone YOUR_ZONE \
    --worker=all \
    --command="cd ~/tunix && python3 multi_host_test.py"
"""

import os
import sys
import jax
import jax.numpy as jnp

def print_separator(char="=", length=60):
    print(char * length)

def main():
    print_separator()
    print(f"JAX Multi-Host TPU Test")
    print_separator()
    
    try:
        # Basic JAX info
        print(f"\n[Process {jax.process_index()}/{jax.process_count()}]")
        print(f"JAX version: {jax.__version__}")
        print(f"Hostname: {os.environ.get('HOSTNAME', 'unknown')}")
        
        # Device counts
        print(f"\nDevice Information:")
        print(f"  Local devices (this host): {jax.local_device_count()}")
        print(f"  Global devices (all hosts): {jax.device_count()}")
        print(f"  Process index: {jax.process_index()}")
        print(f"  Process count: {jax.process_count()}")
        
        # Show local devices
        print(f"\nLocal devices on this host:")
        for i, device in enumerate(jax.local_devices()):
            print(f"  {i}: {device}")
        
        # Only process 0 shows global info to avoid spam
        if jax.process_index() == 0:
            print(f"\nAll devices across all hosts:")
            for i, device in enumerate(jax.devices()):
                print(f"  {i}: {device}")
        
        # Synchronization test
        print(f"\n[Process {jax.process_index()}] Synchronizing...")
        jax.experimental.multihost_utils.sync_global_devices("test_sync")
        
        if jax.process_index() == 0:
            print(f"\n✓ All {jax.process_count()} hosts synchronized!")
        
        # Simple computation test
        print(f"\n[Process {jax.process_index()}] Running computation...")
        
        # Create array on local devices
        local_size = 100
        x = jnp.ones((local_size, local_size))
        y = x @ x  # Matrix multiplication
        result = y.sum()
        
        print(f"[Process {jax.process_index()}] Local computation result: {result}")
        
        # Final sync
        jax.experimental.multihost_utils.sync_global_devices("test_complete")
        
        if jax.process_index() == 0:
            print(f"\n" + "=" * 60)
            print(f"✓ SUCCESS: Multi-host TPU pod is working!")
            print(f"=" * 60)
            print(f"\nSummary:")
            print(f"  Total hosts: {jax.process_count()}")
            print(f"  Total TPU chips: {jax.device_count()}")
            print(f"  Chips per host: {jax.local_device_count()}")
            print(f"\nYour TPU pod is ready for training!")
        
        return 0
        
    except Exception as e:
        print(f"\n✗ ERROR on process {jax.process_index()}: {e}")
        import traceback
        traceback.print_exc()
        return 1

if __name__ == "__main__":
    sys.exit(main())