code: def run_training_proc(
      local_proc_rank: int,
      num_nodes: int,
      node_rank: int,
      num_training_procs_per_node: int,
      dataset_name: str,
      in_channels: int,
      out_channels: int,
      dataset: glt.distributed.DistDataset,
      train_idx: Tensor,
      test_idx: Tensor,
      epochs: int,
      batch_size: int,
      master_addr: str,
      training_pg_master_port: int,
      train_loader_master_port: int,
      test_loader_master_port: int,
  ):
      # Initialize graphlearn_torch distributed worker group context:
      glt.distributed.init_worker_group(
          world_size=num_nodes * num_training_procs_per_node,
          rank=node_rank * num_training_procs_per_node
          + local_proc_rank,
          group_name="distributed-sage-supervised-trainer",
      )
  
      current_ctx = glt.distributed.get_context()
      current_device = torch.device(
          local_proc_rank % torch.cuda.device_count()
      )
  
      # Initialize training process group of PyTorch:
      torch.distributed.init_process_group(
          backend="nccl",  # or choose 'gloo' if 'nccl' is not supported.
          rank=current_ctx.rank,
          world_size=current_ctx.world_size,
          init_method=f"tcp://{master_addr}:{training_pg_master_port}",
      )
  
      # Create distributed neighbor loader for training.
      # We replace PyG's NeighborLoader with GLT's DistNeighborLoader.
      # GLT parameters for sampling are quite similar to PyG.
      # We only need to configure additional network and device parameters:
      train_idx = train_idx.split(
          train_idx.size(0) // num_training_procs_per_node
      )[local_proc_rank]
      train_loader = glt.distributed.DistNeighborLoader(
          data=dataset,
          num_neighbors=[15, 10, 5],
          input_nodes=train_idx,
          batch_size=batch_size,
          shuffle=True,
          collect_features=True,
          to_device=current_device,
          worker_options=glt.distributed.MpDistSamplingWorkerOptions(
              num_workers=1,
              worker_devices=[current_device],
              worker_concurrency=4,
              master_addr=master_addr,
              master_port=train_loader_master_port,
              channel_size="1GB",
              pin_memory=True,
          ),
      )
  
      # Create distributed neighbor loader for testing.
      test_idx = test_idx.split(
          test_idx.size(0) // num_training_procs_per_node
      )[local_proc_rank]
      test_loader = glt.distributed.DistNeighborLoader(
          data=dataset,
          num_neighbors=[15, 10, 5],
          input_nodes=test_idx,
          batch_size=batch_size,
          shuffle=False,
          collect_features=True,
          to_device=current_device,
          worker_options=glt.distributed.MpDistSamplingWorkerOptions(
              num_workers=2,
              worker_devices=[
                  torch.device("cuda", i % torch.cuda.device_count())
                  for i in range(2)
              ],
              worker_concurrency=4,
              master_addr=master_addr,
              master_port=test_loader_master_port,
              channel_size="2GB",
              pin_memory=True,
          ),
      )
  
      # Define the model and optimizer.
      torch.cuda.set_device(current_device)
      model = GraphSAGE(
          in_channels=in_channels,
          hidden_channels=256,
          num_layers=3,
          out_channels=out_channels,
      ).to(current_device)
      model = DistributedDataParallel(
          model, device_ids=[current_device.index]
      )
  
      optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
  
      # Train and test:
      f = open("dist_sage_sup.txt", "a+")
      for epoch in range(0, epochs):
          model.train()
          start = time.time()
          for batch in train_loader:
              optimizer.zero_grad()
              out = model(batch.x, batch.edge_index)[: batch.batch_size]
              loss = F.cross_entropy(
                  out, batch.y[: batch.batch_size].long()
              )
              loss.backward()
              optimizer.step()
          f.write(
              f"-- [Trainer {current_ctx.rank}] Epoch: {epoch:03d}, "
              f"Loss: {loss:.4f}, Epoch Time: {time.time() - start}\n"
          )
  
          torch.cuda.synchronize()
          torch.distributed.barrier()
  
          if epoch == 0 or epoch > (epochs // 2):
              test_acc = test(model, test_loader, dataset_name)
              f.write(
                  f"-- [Trainer {current_ctx.rank}] "
                  f"Test Acc: {test_acc:.4f}\n"
              )
              torch.cuda.synchronize()
              torch.distributed.barrier()
  
file_path: pytorch_geometric/examples/distributed/graphlearn_for_pytorch/dist_train_sage_supervised.py
imports: import argparse
  import os.path
  import time
  import graphlearn_torch
  import torch
  import torch.distributed
  import torch.nn.functional
  from ogb.nodeproppred import Evaluator
  from torch import Tensor
  from torch.nn.parallel import DistributedDataParallel
  from torch_geometric.nn import GraphSAGE
label_desc: Use image node_id 2 for calling a global function with example usage: run_training_proc(local_proc_rank, num_nodes, node_rank, num_training_procs_per_node, dataset_name, in_channels, out_channels, dataset, train_idx, test_idx, epochs, batch_size, master_addr, training_pg_master_port, train_loader_master_port, test_loader_master_port) without return types