code: def run(
      rank: int,
      world_size: int,
      args: argparse.ArgumentParser,
      num_classes: int,
      data: Union[Data, HeteroData],
      custom_optimizer: Callable[[Any, Any], Tuple[Any, Any]] = None,
  ):
      if not device_conditions[args.device]():
          raise RuntimeError(f"{args.device.upper()} is not available")
  
      device = torch.device(f"{args.device}:{rank}")
  
      if rank == 0:
          print("BENCHMARK STARTS")
          print(f"Running on {args.device.upper()}")
          print(f"Dataset: {args.dataset}")
  
      hetero = True if args.dataset == "ogbn-mag" else False
      mask, val_mask, test_mask = get_split_masks(data, args.dataset)
      mask = create_mask_per_rank(mask, rank, world_size, hetero)
      degree = None
  
      inputs_channels = (
          data["paper"].num_features
          if args.dataset == "ogbn-mag"
          else data.num_features
      )
  
      if args.model not in supported_sets[args.dataset]:
          err_msg = (
              f"Configuration of {args.dataset} + {args.model}"
              "not supported"
          )
          raise RuntimeError(err_msg)
      if rank == 0:
          print(f"Training bench for {args.model}:")
  
      num_nodes = int(mask[-1].sum()) if hetero else int(mask.sum())
      num_neighbors = args.num_neighbors
  
      if type(num_neighbors) is list:
          if len(num_neighbors) == 1:
              num_neighbors = num_neighbors * args.num_layers
      elif type(num_neighbors) is int:
          num_neighbors = [num_neighbors] * args.num_layers
  
      if len(num_neighbors) != args.num_layers:
          err_msg = (
              f"num_neighbors={num_neighbors} lenght != num of"
              "layers={args.num_layers}"
          )
  
      kwargs = {
          "num_neighbors": num_neighbors,
          "batch_size": args.batch_size,
          "num_workers": args.num_workers,
      }
      subgraph_loader = NeighborLoader(
          data,
          input_nodes=mask,
          sampler=None,
          **kwargs,
      )
      if rank == 0 and args.evaluate:
          val_loader = NeighborLoader(
              data,
              input_nodes=val_mask,
              sampler=None,
              **kwargs,
          )
          test_loader = NeighborLoader(
              data,
              input_nodes=test_mask,
              sampler=None,
              **kwargs,
          )
  
      if rank == 0:
          print("----------------------------------------------")
          print(
              f"Batch size={args.batch_size}, "
              f"Layers amount={args.num_layers}, "
              f"Num_neighbors={num_neighbors}, "
              f"Hidden features size={args.num_hidden_channels}",
              flush=True,
          )
  
      params = {
          "inputs_channels": inputs_channels,
          "hidden_channels": args.num_hidden_channels,
          "output_channels": num_classes,
          "num_heads": args.num_heads,
          "num_layers": args.num_layers,
      }
  
      if args.model == "pna" and degree is None:
          degree = PNAConv.get_degree_histogram(subgraph_loader)
          print(
              f"Rank: {rank}, calculated degree for {args.dataset}.",
              flush=True,
          )
          params["degree"] = degree
      dist.barrier()
  
      torch.manual_seed(12345)
      model = get_model(
          args.model,
          params,
          metadata=data.metadata() if hetero else None,
      )
      model = model.to(device)
  
      if hetero:
          model.eval()
          x_keys = data.metadata()[0]
          edge_index_keys = data.metadata()[1]
          fake_x_dict = {
              k: torch.rand((32, inputs_channels), device=device)
              for k in x_keys
          }
          fake_edge_index_dict = {
              k: torch.randint(0, 32, (2, 8), device=device)
              for k in edge_index_keys
          }
          model.forward(fake_x_dict, fake_edge_index_dict)
  
      model = DDP(
          model, device_ids=[device], find_unused_parameters=hetero
      )
      model.train()
  
      optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
  
      if custom_optimizer:
          model, optimizer = custom_optimizer(model, optimizer)
  
      train = train_hetero if hetero else train_homo
  
      maybe_synchronize(args.device)
      dist.barrier()
      if rank == 0:
          beg = perf_counter()
  
      for epoch in range(args.num_epochs):
          loss = train(
              model,
              subgraph_loader,
              optimizer,
              device,
          )
  
          dist.barrier()
  
          if rank == 0:
              print(f"Epoch: {epoch:02d}, Loss: {loss:.4f}", flush=True)
  
          if rank == 0 and args.evaluate:
              # In evaluate, throughput and
              # latency are not accurate.
              val_acc = test(
                  model, val_loader, device, hetero, progress_bar=False
              )
              print(f"Val Accuracy: {val_acc:.4f}")
  
          dist.barrier()
  
      maybe_synchronize(args.device)
      dist.barrier()
      if rank == 0:
          end = perf_counter()
          duration = end - beg
  
      if rank == 0 and args.evaluate:
          test_acc = test(
              model, test_loader, device, hetero, progress_bar=False
          )
          print(f"Test Accuracy: {test_acc:.4f}")
  
      dist.barrier()
  
      if rank == 0:
          num_nodes_total = num_nodes * world_size
          duration_per_epoch = duration / args.num_epochs
          throughput = num_nodes_total / duration_per_epoch
          latency = duration_per_epoch / num_nodes_total * 1000
          print(f"Time: {duration_per_epoch:.4f}s")
          print(f"Throughput: {throughput:.3f} samples/s")
          print(f"Latency: {latency:.3f} ms", flush=True)
  
      dist.destroy_process_group()
  
file_path: pytorch_geometric/benchmark/multi_gpu/training/common.py
imports: import argparse
  import ast
  from time import perf_counter
  from typing import Any, Callable, Tuple, Union
  import torch
  import torch.distributed
  import torch.nn.functional
  from torch.nn.parallel import DistributedDataParallel
  from benchmark.utils import get_model, get_split_masks, test
  from torch_geometric.data import Data, HeteroData
  from torch_geometric.loader import NeighborLoader
  from torch_geometric.nn import PNAConv
label_desc: Use image node_id 5 for calling a global function with example usage: run(rank, world_size, args, num_classes, data, custom_optimizer) without return types