import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import numpy as np

class WineDataset(Dataset):
 def __init__(self, data, targets):
        self.data = data
        self.targets = targets

 def __len__(self):
 return len(self.data)

 def __getitem__(self, idx):
 return torch.tensor(self.data[idx], dtype=torch.float), torch.tensor(self.targets[idx], dtype=torch.long)

class SimpleNN(torch.nn.Module):
 def __init__(self):
 super(SimpleNN, self).__init__()
        self.fc1 = torch.nn.Linear(13, 64)
        self.fc2 = torch.nn.Linear(64, 3)

 def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
 return x

class Trainer():
 def __init__(self, model, train_data, optimizer, gpu_id, save_every):
        self.model = model
        self.train_data = train_data
        self.optimizer = optimizer
        self.gpu_id = gpu_id
        self.save_every = save_every
        self.losses = []

 def _run_batch(self, source, targets):
        self.optimizer.zero_grad()
        output = self.model(source)
        loss = F.cross_entropy(output, targets)
        loss.backward()
        self.optimizer.step()
 return loss.item()

 def _run_epoch(self, epoch):
        total_loss = 0.0
        num_batches = len(self.train_data)
 for source, targets in self.train_data:
            source = source.to(self.gpu_id)
            targets = targets.to(self.gpu_id)
            loss = self._run_batch(source, targets)
            total_loss += loss

        avg_loss = total_loss / num_batches
        self.losses.append(avg_loss)
 print(f"Epoch {epoch}, Loss: {avg_loss:.4f}")

 def _save_checkpoint(self, epoch):
        checkpoint = self.model.state_dict()
        PATH = f"model_{epoch}.pt"
        torch.save(checkpoint, PATH)
 print(f"Epoch {epoch} | Model saved to {PATH}")

 def train(self, max_epochs):
        self.model.train()
 for epoch in range(max_epochs):
            self._run_epoch(epoch)
 if epoch % self.save_every == 0:
                self._save_checkpoint(epoch)

def load_train_objs():
    wine_data = load_wine()
    X = wine_data.data
    y = wine_data.target

 # Normalize and split
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    scaler = StandardScaler().fit(X_train)
    X_train = scaler.transform(X_train)
    X_test = scaler.transform(X_test)

    train_set = WineDataset(X_train, y_train)
    test_set = WineDataset(X_test, y_test)

 print("Sample from dataset:")
    sample_data, sample_target = train_set[0]
 print(f"Data: {sample_data}")
 print(f"Target: {sample_target}")

    model = SimpleNN()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

 return train_set, model, optimizer

def prepare_dataloader(dataset, batch_size):
 return DataLoader(dataset, batch_size=batch_size, pin_memory=True, shuffle=True)

def main(device, total_epochs, save_every, batch_size):
    dataset, model, optimizer = load_train_objs()
    train_data = prepare_dataloader(dataset, batch_size)
    trainer = Trainer(model, train_data, optimizer, device, save_every)
    trainer.train(total_epochs)

main(device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"), total_epochs=100, save_every=50, batch_size=32)

import torch.multiprocessing as mp
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group, destroy_process_group
import os

def ddp_setup(rank, world_size):
 """
    Set up the distributed environment.
 
    Args:
        rank: The rank of the current process. Unique identifier for each process in the distributed training.
        world_size: Total number of processes participating in the distributed training.
    """
 
 # Address of the main node. Since we are doing single-node training, it's set to localhost.
    os.environ["MASTER_ADDR"] = "localhost"
 
 # Port on which the master node is expected to listen for communications from workers.
    os.environ["MASTER_PORT"] = "12355"
 
 # Initialize the process group. 
 # 'backend' specifies the communication backend to be used, "nccl" is optimized for GPU training.
    init_process_group(backend="nccl", rank=rank, world_size=world_size)
 
 # Set the current CUDA device to the specified device (identified by rank).
 # This ensures that each process uses a different GPU in a multi-GPU setup.
    torch.cuda.set_device(rank)

class Trainer():
 def __init__(self, model, train_data, optimizer, gpu_id, save_every):
        self.model = model.to(gpu_id)
        self.train_data = train_data
        self.optimizer = optimizer
        self.gpu_id = gpu_id
        self.save_every = save_every
        self.losses = []
 
 # This changes
        self.model = DDP(self.model, device_ids=[gpu_id])

def prepare_dataloader(dataset: Dataset, batch_size: int):
 return DataLoader(
        dataset,
        batch_size=batch_size,
        pin_memory=True,
        shuffle=False,
        sampler=DistributedSampler(dataset)
 )

def main(rank: int, world_size: int, save_every: int, total_epochs: int, batch_size: int):
 """
    Main training function for distributed data parallel (DDP) setup.
 
    Args:
        rank (int): The rank of the current process (0 <= rank < world_size). Each process is assigned a unique rank.
        world_size (int): Total number of processes involved in the distributed training.
        save_every (int): Frequency of model checkpoint saving, in terms of epochs.
        total_epochs (int): Total number of epochs for training.
        batch_size (int): Number of samples processed in one iteration (forward and backward pass).
    """
 
 # Set up the distributed environment, including setting the master address, port, and backend.
    ddp_setup(rank, world_size)
 
 # Load the necessary training objects - dataset, model, and optimizer.
    dataset, model, optimizer = load_train_objs()
 
 # Prepare the data loader for distributed training. It partitions the dataset across the processes and handles shuffling.
    train_data = prepare_dataloader(dataset, batch_size)
 
 # Initialize the trainer instance with the loaded model, data, and other configurations.
    trainer = Trainer(model, train_data, optimizer, rank, save_every)
 
 # Train the model for the specified number of epochs.
    trainer.train(total_epochs)
 
 # Cleanup the distributed environment after training is complete.
    destroy_process_group()

if __name__ == "__main__":
 import argparse
    parser = argparse.ArgumentParser(description='simple distributed training job')
    parser.add_argument('total_epochs', type=int, help='Total epochs to train the model')
    parser.add_argument('save_every', type=int, help='How often to save a snapshot')
    parser.add_argument('--batch_size', default=32, type=int, help='Input batch size on each device (default: 32)')
    args = parser.parse_args()
 
    world_size = torch.cuda.device_count()
    mp.spawn(main, args=(world_size, args.save_every, args.total_epochs, args.batch_size), nprocs=world_size)

def ddp_setup(local_rank, world_size_per_node, node_rank):
    os.environ["MASTER_ADDR"] = "MASTER_NODE_IP" # <-- Replace with your master node IP
    os.environ["MASTER_PORT"] = "12355" 
    global_rank = node_rank * world_size_per_node + local_rank
    init_process_group(backend="nccl", rank=global_rank, world_size=world_size_per_node*torch.cuda.device_count())
    torch.cuda.set_device(local_rank)

def main(local_rank: int, world_size_per_node: int, save_every: int, total_epochs: int, batch_size: int, node_rank: int):
    ddp_setup(local_rank, world_size_per_node, node_rank)
 # ... (rest of the main function)

if __name__ == "__main__":
 import argparse
    parser = argparse.ArgumentParser(description='simple distributed training job')
    parser.add_argument('total_epochs', type=int, help='Total epochs to train the model')
    parser.add_argument('save_every', type=int, help='How often to save a snapshot')
    parser.add_argument('--batch_size', default=32, type=int, help='Input batch size on each device (default: 32)')
    parser.add_argument('--node_rank', default=0, type=int, help='The rank of the node in multi-node training')
    args = parser.parse_args()

    world_size_per_node = torch.cuda.device_count()
    mp.spawn(main, args=(world_size_per_node, args.save_every, args.total_epochs, args.batch_size, args.node_rank), nprocs=world_size_per_node)

#!/bin/bash
#SBATCH --job-name=DDPTraining       # Name of the job
#SBATCH --nodes=$1                   # Number of nodes specified by the user
#SBATCH --ntasks-per-node=1          # Ensure only one task runs per node
#SBATCH --cpus-per-task=1            # Number of CPU cores per task
#SBATCH --gres=gpu:1                 # Number of GPUs per node
#SBATCH --time=01:00:00              # Time limit hrs:min:sec (1 hour in this example)
#SBATCH --mem=4GB                    # Memory limit per GPU
#SBATCH --output=training_%j.log     # Output and error log name (%j expands to jobId)
#SBATCH --partition=gpu              # Specify the partition or queue

srun python3 your_python_script.py --total_epochs 10 --save_every 2 --batch_size 32 --node_rank $SLURM_NODEID

sbatch train_net.sh 2 # for using 2 nodes