code: def main(args):
      import logging
  
      logging.basicConfig(
          level=logging.INFO,
      )
      logger = logging.getLogger("bench_cugraph_training")
      logger.setLevel(logging.INFO)
  
      local_rank = int(os.environ["LOCAL_RANK"])
      global_rank = int(os.environ["RANK"])
  
      init_pytorch_worker(
          local_rank,
          use_rmm_torch_allocator=(args.framework == "cuGraph"),
      )
      enable_spilling()
      print(f"worker initialized")
      dist.barrier()
  
      world_size = (
          int(os.environ["SLURM_JOB_NUM_NODES"]) * args.gpus_per_node
      )
  
      dataset = OGBNPapers100MDataset(
          replication_factor=args.replication_factor,
          dataset_dir=args.dataset_dir,
          train_split=args.train_split,
          val_split=args.val_split,
          load_edge_index=(args.framework == "PyG"),
      )
  
      if global_rank == 0:
          dataset.download()
      dist.barrier()
  
      fanout = [int(f) for f in args.fanout.split("_")]
  
      if args.framework == "PyG":
          from trainers.pyg import PyGNativeTrainer
  
          trainer = PyGNativeTrainer(
              model=args.model,
              dataset=dataset,
              device=local_rank,
              rank=global_rank,
              world_size=world_size,
              num_epochs=args.num_epochs,
              shuffle=True,
              replace=False,
              num_neighbors=fanout,
              batch_size=args.batch_size,
          )
      elif args.framework == "cuGraphPyG":
          sample_dir = os.path.join(
              args.sample_dir,
              f"ogbn_papers100M[{args.replication_factor}]_b{args.batch_size}_f{fanout}",
          )
          from trainers.pyg import PyGCuGraphTrainer
  
          trainer = PyGCuGraphTrainer(
              model=args.model,
              dataset=dataset,
              sample_dir=sample_dir,
              device=local_rank,
              rank=global_rank,
              world_size=world_size,
              num_epochs=args.num_epochs,
              shuffle=True,
              replace=False,
              num_neighbors=fanout,
              batch_size=args.batch_size,
          )
      else:
          raise ValueError("unsupported framework")
  
      logger.info(f"Trainer ready on rank {global_rank}")
      stats = trainer.train()
      logger.info(stats)
  
      with open(f"{args.output_file}[{global_rank}]", "w") as f:
          json.dump(stats, f)
  
file_path: cugraph/benchmarks/cugraph/standalone/bulk_sampling/bench_cugraph_training.py
imports: import os
  import argparse
  import json
  import warnings
  import torch
  import numpy
  import pandas
  import torch.distributed
  from datasets import OGBNPapers100MDataset
  from cugraph.testing.mg_utils import enable_spilling
label_desc: Use image node_id 3 for calling a global function with example usage: main(args) without return types