def shunt_torch_communication()

in sample_workloads/megatron-gke/docker/monitor_collectives.py [0:0]

• 8 lines of code
• 2 McCabe index (conditional complexity)

def shunt_torch_communication():
  if _TRACE_MODE == 'none':
    print('Tracing torch.distributed collectives disabled.', flush=True)
    return

  _shunt_torch_communication_objects()
  _shunt_torch_communication_calls()

  print('NVTX and print tracing of torch.distributed collectives enabled.',
        flush=True)