run_parameter_server_training.py

from __future__ import annotations

import argparse
import time
from pathlib import Path
from typing import Dict

import torch
from torch.utils.data import DataLoader, Subset

from src.config import SystemConfig, TrainingConfig
from src.coordinator import Coordinator
from src.dataset import load_dataset
from src.parameter_server import ParameterServer
from src.trainer import SimpleCNN


def _subset_dataset(dataset, max_samples: int | None):
    if max_samples is None or max_samples >= len(dataset):
        return dataset
    return Subset(dataset, list(range(max_samples)))


def _evaluate_model(model: torch.nn.Module, dataset, batch_size: int) -> Dict[str, float]:
    model.eval()
    criterion = torch.nn.CrossEntropyLoss()
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0)
    total_loss = 0.0
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in loader:
            out = model(data)
            loss = criterion(out, target)
            total_loss += float(loss.item())
            pred = out.argmax(dim=1)
            correct += int((pred == target).sum().item())
            total += int(target.size(0))
    return {
        'accuracy': 100.0 * correct / max(total, 1),
        'loss': total_loss / max(len(loader), 1),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description='Run parameter-server training.')
    parser.add_argument('--workers', type=int, default=2, help='Number of worker processes.')
    parser.add_argument('--dataset', type=str, default='mnist', choices=['mnist', 'fashion_mnist', 'cifar10'])
    parser.add_argument('--data-dir', type=str, default='./data')
    parser.add_argument('--batch-size', type=int, default=64)
    parser.add_argument('--learning-rate', type=float, default=0.001)
    parser.add_argument('--num-epochs', type=int, default=2)
    parser.add_argument('--aggregation-mode', type=str, default='sync', choices=['sync', 'async'])
    parser.add_argument('--compression-enabled', action='store_true')
    parser.add_argument('--compression-type', type=str, default='quantization', choices=['quantization', 'topk'])
    parser.add_argument('--compression-ratio', type=float, default=0.1)
    parser.add_argument('--heartbeat-interval', type=float, default=0.5)
    parser.add_argument('--heartbeat-timeout', type=float, default=2.0)
    parser.add_argument('--timeout-s', type=float, default=3600.0)
    parser.add_argument('--max-train-samples', type=int, default=None)
    parser.add_argument('--max-test-samples', type=int, default=None)
    parser.add_argument('--checkpoint-dir', type=str, default='./checkpoints/ps_cli')
    parser.add_argument('--log-dir', type=str, default='./logs/ps_cli')
    args = parser.parse_args()

    if args.workers <= 0:
        raise ValueError('--workers must be positive')

    Path(args.checkpoint_dir).mkdir(parents=True, exist_ok=True)
    Path(args.log_dir).mkdir(parents=True, exist_ok=True)

    train_cfg = TrainingConfig(
        batch_size=args.batch_size,
        learning_rate=args.learning_rate,
        num_epochs=args.num_epochs,
        dataset=args.dataset,
        model_architecture='simple_cnn',
        checkpoint_interval=1,
    )
    system_cfg = SystemConfig(
        num_workers=args.workers,
        architecture='parameter_server',
        aggregation_mode=args.aggregation_mode,
        compression_enabled=bool(args.compression_enabled),
        compression_type=args.compression_type,
        compression_ratio=args.compression_ratio,
        heartbeat_interval=args.heartbeat_interval,
        heartbeat_timeout=args.heartbeat_timeout,
        checkpoint_dir=args.checkpoint_dir,
        log_dir=args.log_dir,
    )

    train_ds = _subset_dataset(load_dataset(args.dataset, train=True, data_dir=args.data_dir), args.max_train_samples)
    test_ds = _subset_dataset(load_dataset(args.dataset, train=False, data_dir=args.data_dir), args.max_test_samples)

    print("Starting parameter-server training on cpu", flush=True)
    print(f"Dataset: {args.dataset}", flush=True)
    print(f"Epochs: {args.num_epochs}", flush=True)
    print(f"Batch size: {args.batch_size}", flush=True)
    print(f"Learning rate: {args.learning_rate}", flush=True)
    print(f"Workers: {args.workers}", flush=True)
    print(f"Aggregation mode: {args.aggregation_mode}", flush=True)
    print(f"Compression enabled: {args.compression_enabled}", flush=True)
    if args.compression_enabled:
        print(f"Compression type: {args.compression_type}", flush=True)
        print(f"Compression ratio: {args.compression_ratio}", flush=True)

    ps = ParameterServer(
        model=SimpleCNN(args.dataset),
        num_workers=args.workers,
        aggregation_mode=args.aggregation_mode,
        learning_rate=args.learning_rate,
    )
    ps.start_server(host='127.0.0.1', port=0, max_workers=32)
    assert ps.bound_port is not None

    coord = Coordinator(
        training_config=train_cfg,
        system_config=system_cfg,
        parameter_server=ps,
        parameter_server_address=f'127.0.0.1:{ps.bound_port}',
        data_dir=args.data_dir,
        worker_log_root=str(Path(args.log_dir) / 'workers'),
    )
    coord.start_server(host='127.0.0.1', port=0, max_workers=32)
    coord.start_heartbeat_monitor()

    start = time.time()
    try:
        coord.start_workers(
            model_factory=lambda: SimpleCNN(args.dataset),
            dataset=train_ds,
            start_training=True,
        )
        ok = coord.wait_for_workers(timeout_s=args.timeout_s)
        if not ok:
            raise RuntimeError(f'Parameter-server training timed out after {args.timeout_s} seconds')
        if coord.worker_errors:
            raise RuntimeError(f'Worker errors: {coord.worker_errors}')

        metrics = _evaluate_model(ps.model, test_ds, batch_size=args.batch_size)
        total_time = time.time() - start
        total_samples = len(train_ds) * args.num_epochs
        throughput = float(total_samples) / max(total_time, 1e-9)

        ckpt_path = Path(args.checkpoint_dir) / 'parameter_server_final.pt'
        ps.save_checkpoint(str(ckpt_path))

        print(f"\nTraining completed in {total_time:.2f} seconds")
        print(f"Final test accuracy: {metrics['accuracy']:.2f}%")
        print(f"Throughput: {throughput:.2f} samples/second")

        print("\n" + "=" * 60)
        print("TRAINING COMPLETE")
        print("=" * 60)
        print(f"Final Accuracy: {metrics['accuracy']:.2f}%")
        print(f"Final Loss: {metrics['loss']:.4f}")
        print(f"Total Time: {total_time:.2f} seconds")
        print(f"Throughput: {throughput:.2f} samples/second")
        print("=" * 60)
        print(f"Checkpoint: {ckpt_path.resolve()}")
        print(f"Logs: {Path(args.log_dir).resolve()}")
        print("=" * 60)
    finally:
        coord.shutdown()
        ps.stop_server(0)


if __name__ == '__main__':
    main()