Traffic routing strategy

import os
import torch
from torch.utils.data import DataLoader, Dataset
from dataclasses import dataclass
from typing import Optional
import psutil

class SimpleDataset(Dataset):
    def __init__(self, size: int = 100):
        self.size = size

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        return {
            'input_ids': torch.randint(0, 32000, (512,)),
            'labels': torch.randint(0, 2, (512,)),
            'sequence_length': 512,
        }

@dataclass
class DeviceLoad:
    device_id: int
    memory_used_gb: float
    memory_total_gb: float
    utilization_percent: float

class TrafficRouter:
    def __init__(self, num_devices: int = 2):
        self.num_devices = num_devices
        self.device_loads = [DeviceLoad(
            device_id=i,
            memory_used_gb=0.0,
            memory_total_gb=float(torch.cuda.get_device_properties(i).total_memory / 1e9),
            utilization_percent=0.0
        ) for i in range(num_devices)]
        self.routing_history = []

    def get_device_loads(self) -> list[DeviceLoad]:
        """Poll current GPU memory usage."""
        loads = []
        for i in range(self.num_devices):
            if torch.cuda.is_available():
                reserved = torch.cuda.memory_reserved(i) / 1e9
                allocated = torch.cuda.memory_allocated(i) / 1e9
                total = torch.cuda.get_device_properties(i).total_memory / 1e9
                utilization = (allocated / total) * 100 if total > 0 else 0
                loads.append(DeviceLoad(
                    device_id=i,
                    memory_used_gb=allocated,
                    memory_total_gb=total,
                    utilization_percent=utilization
                ))
            else:
                loads.append(DeviceLoad(
                    device_id=i,
                    memory_used_gb=0.0,
                    memory_total_gb=16.0,
                    utilization_percent=0.0
                ))
        return loads

    def route_batch(self, batch: dict, strategy: str = 'least_loaded') -> int:
        """Route batch to device based on strategy.
        
        Args:
            batch: Training batch dict
            strategy: 'least_loaded' (memory), 'round_robin', or 'affinity' (by sequence length)
        
        Returns:
            Device ID (0 or 1 typically)
        """
        loads = self.get_device_loads()
        
        if strategy == 'least_loaded':
            selected_device = min(loads, key=lambda x: x.utilization_percent).device_id
        elif strategy == 'round_robin':
            selected_device = len(self.routing_history) % self.num_devices
        elif strategy == 'affinity':
            seq_len = batch.get('sequence_length', 512)
            if seq_len > 1024:
                selected_device = 0
            else:
                selected_device = 1
        else:
            selected_device = 0
        
        self.routing_history.append({
            'device_id': selected_device,
            'strategy': strategy,
            'loads_snapshot': loads
        })
        return selected_device

    def report_routing_stats(self) -> dict:
        """Summarize routing decisions made."""
        if not self.routing_history:
            return {'total_routed': 0, 'device_distribution': {}}
        
        distribution = {}
        for entry in self.routing_history:
            device_id = entry['device_id']
            distribution[device_id] = distribution.get(device_id, 0) + 1
        
        return {
            'total_routed': len(self.routing_history),
            'device_distribution': distribution,
            'balance_ratio': max(distribution.values()) / min(distribution.values()) if distribution else 1.0
        }

def main():
    # Simulate fine-tuning with traffic routing
    num_gpus = 2
    router = TrafficRouter(num_devices=num_gpus)
    dataset = SimpleDataset(size=50)
    dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
    
    print(f"Routing {len(dataset)} samples across {num_gpus} GPUs\n")
    
    for batch_idx, batch in enumerate(dataloader):
        selected_device = router.route_batch(batch, strategy='least_loaded')
        loads = router.get_device_loads()
        gpu_util = [f"GPU{l.device_id}:{l.utilization_percent:.1f}%" for l in loads]
        print(f"Batch {batch_idx:2d} → Device {selected_device} | {' | '.join(gpu_util)}")
        
        if batch_idx >= 9:
            break
    
    stats = router.report_routing_stats()
    print(f"\nRouting Summary:")
    print(f"  Total batches routed: {stats['total_routed']}")
    print(f"  Distribution: {stats['device_distribution']}")
    print(f"  Load balance ratio: {stats['balance_ratio']:.2f}")

if __name__ == '__main__':
    main()