def _emit_call_description()

in sample_workloads/lit-gpt-demo/utilities/monitor_collectives.py [0:0]


def _emit_call_description(
    name, message_size, group=None, peer_rank=None, root_rank=None, correlation_id=None):
  call_description = _TorchDistributedCallDescriptor(
      name, message_size, group, peer_rank, root_rank, correlation_id).to_json()

  nvtx.mark(call_description)
  if _should_rank_print(group, peer_rank, root_rank):
    print(call_description)