re-format output

This commit is contained in:
Ubuntu
2026-03-17 19:59:35 +00:00
parent c84c2ede20
commit 8dc63faedd

View File

@@ -181,7 +181,7 @@ def main(
n_graph_iters: int = 10,
):
mscclpp_group = CommGroup(MPI.COMM_WORLD)
print(f"use GPU {mscclpp_group.my_rank % mscclpp_group.nranks_per_node} on rank {mscclpp_group.my_rank}")
#print(f"use GPU {mscclpp_group.my_rank % mscclpp_group.nranks_per_node} on rank {mscclpp_group.my_rank}")
cp.cuda.Device(mscclpp_group.my_rank % mscclpp_group.nranks_per_node).use()
executor = Executor(mscclpp_group.communicator)
npkit_dump_dir = env().npkit_dump_dir
@@ -230,12 +230,16 @@ def main(
if npkit_dump_dir is not None:
npkit.dump(npkit_dump_dir)
npkit.shutdown()
print(
f"Rank: {mscclpp_group.my_rank} Execution time: {execution_time} us, "
f"data size: {result_buf.nbytes} bytes data type: {dtype().dtype.name} "
f" busBW: {result_buf.nbytes / execution_time / 1e3:.2f} GB/s "
f"packet type: {packet_type}"
)
# Only rank 0 reports output
if mscclpp_group.my_rank == 0:
msg_size = result_buf.nbytes
bw = result_buf.nbytes / execution_time / 1e3 # GB/s
latency = execution_time # us
# Print header once
print(f"{'Message Size (B)':>18} {'BW (GB/s)':>12} {'Latency (us)':>14} {'Packet Type':>12}")
print(f"{msg_size:18d} {bw:12.2f} {latency:14.2f} {str(packet_type):>12}")
executor = None
mscclpp_group = None