code: def run(proc_id, nprocs, devices, g, data, args):
      # Find corresponding device for current process.
      device = devices[proc_id]
      torch.cuda.set_device(device)
      #########################################################################
      # (HIGHLIGHT) Build a data-parallel distributed GraphSAGE model.
      #
      # DDP in PyTorch provides data parallelism across the devices specified
      # by the `process_group`. Gradients are synchronized across each model
      # replica.
      #
      # To prepare a training sub-process, there are four steps involved:
      # 1. Initialize the process group
      # 2. Unpack data for the sub-process.
      # 3. Instantiate a GraphSAGE model on the corresponding device.
      # 4. Parallelize the model with `DistributedDataParallel`.
      #
      # For the detailed usage of `DistributedDataParallel`, please refer to
      # PyTorch documentation.
      #########################################################################
      dist.init_process_group(
          backend="nccl",  # Use NCCL backend for distributed GPU training
          init_method="tcp://127.0.0.1:12345",
          world_size=nprocs,
          rank=proc_id,
      )
      num_classes, train_idx, val_idx, test_idx = data
      if args.mode != "benchmark":
          train_idx = train_idx.to(device)
          val_idx = val_idx.to(device)
          g = g.to(device if args.mode == "puregpu" else "cpu")
      in_size = g.ndata["feat"].shape[1]
      model = SAGE(in_size, 256, num_classes).to(device)
      model = DistributedDataParallel(
          model, device_ids=[device], output_device=device
      )
  
      # Training.
      use_uva = args.mode == "mixed"
  
      if proc_id == 0:
          print("Training...")
      train(
          proc_id,
          nprocs,
          device,
          args,
          g,
          num_classes,
          train_idx,
          val_idx,
          model,
          use_uva,
      )
  
      # Testing.
      if proc_id == 0:
          print("Testing...")
      layerwise_infer(
          proc_id, device, g, num_classes, test_idx, model, use_uva
      )
  
      # Cleanup the process group.
      dist.destroy_process_group()
  
imports: import argparse
  import os
  import time
  import dgl
  import dgl.nn
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  import torch.nn.functional
  import torchmetrics.functional
  import tqdm
  from dgl.data import AsNodePredDataset
  from dgl.dataloading import DataLoader, MultiLayerFullNeighborSampler, NeighborSampler
  from dgl.multiprocessing import shared_tensor
  from ogb.nodeproppred import DglNodePropPredDataset
  from torch.nn.parallel import DistributedDataParallel
label_desc: Use image node_id 4 for calling a global function with example usage: run(proc_id, nprocs, devices, g, data, args) without return types