135 lines
3.4 KiB
Python
135 lines
3.4 KiB
Python
import time
|
|
import json
|
|
|
|
from eacgm.bpf import BccBPF
|
|
from eacgm.sampler import eBPFSampler
|
|
from eacgm.collector import to_perfetto
|
|
|
|
text = """
|
|
#include <uapi/linux/ptrace.h>
|
|
|
|
int ncclAllReduceEntry(struct pt_regs *ctx){
|
|
u64 ts = bpf_ktime_get_ns();
|
|
u64 size_count = PT_REGS_PARM3(ctx);
|
|
u64 data_type = PT_REGS_PARM4(ctx);
|
|
u64 reduce_op = PT_REGS_PARM5(ctx);
|
|
bpf_trace_printk("%ld@start@ncclAllReduce@%ld\\n", ts, size_count * 8);
|
|
return 0;
|
|
};
|
|
|
|
int ncclAllReduceExit(struct pt_regs *ctx){
|
|
u64 ts = bpf_ktime_get_ns();
|
|
bpf_trace_printk("%ld@end@ncclAllReduce\\n", ts);
|
|
return 0;
|
|
};
|
|
|
|
int ncclReduceEntry(struct pt_regs *ctx){
|
|
u64 ts = bpf_ktime_get_ns();
|
|
u64 size_count = PT_REGS_PARM3(ctx);
|
|
u64 data_type = PT_REGS_PARM4(ctx);
|
|
u64 reduce_op = PT_REGS_PARM5(ctx);
|
|
bpf_trace_printk("%ld@start@ncclReduce@%ld\\n", ts, size_count * 8);
|
|
return 0;
|
|
};
|
|
|
|
int ncclReduceExit(struct pt_regs *ctx){
|
|
u64 ts = bpf_ktime_get_ns();
|
|
bpf_trace_printk("%ld@end@ncclReduce\\n", ts);
|
|
return 0;
|
|
};
|
|
|
|
int ncclBroadcastEntry(struct pt_regs *ctx){
|
|
u64 ts = bpf_ktime_get_ns();
|
|
u64 size_count = PT_REGS_PARM3(ctx);
|
|
u64 data_type = PT_REGS_PARM4(ctx);
|
|
u64 root_id = PT_REGS_PARM5(ctx);
|
|
bpf_trace_printk("%ld@start@ncclBroadcast@%ld\\n", ts, size_count * 8);
|
|
return 0;
|
|
};
|
|
|
|
int ncclBroadcastExit(struct pt_regs *ctx){
|
|
u64 ts = bpf_ktime_get_ns();
|
|
bpf_trace_printk("%ld@end@ncclBroadcast\\n", ts);
|
|
return 0;
|
|
};
|
|
|
|
int ncclAllGatherEntry(struct pt_regs *ctx){
|
|
u64 ts = bpf_ktime_get_ns();
|
|
u64 size_count = PT_REGS_PARM3(ctx);
|
|
u64 data_type = PT_REGS_PARM4(ctx);
|
|
bpf_trace_printk("%ld@start@ncclAllGather@%ld\\n", ts, size_count * 8);
|
|
return 0;
|
|
};
|
|
|
|
int ncclAllGatherExit(struct pt_regs *ctx){
|
|
u64 ts = bpf_ktime_get_ns();
|
|
bpf_trace_printk("%ld@end@ncclAllGather\\n", ts);
|
|
return 0;
|
|
};
|
|
|
|
int ncclSendEntry(struct pt_regs *ctx){
|
|
u64 ts = bpf_ktime_get_ns();
|
|
u64 size_count = PT_REGS_PARM2(ctx);
|
|
u64 data_type = PT_REGS_PARM3(ctx);
|
|
bpf_trace_printk("%ld@start@ncclSend@%ld\\n", ts, size_count * 8);
|
|
return 0;
|
|
};
|
|
|
|
int ncclSendExit(struct pt_regs *ctx){
|
|
u64 ts = bpf_ktime_get_ns();
|
|
bpf_trace_printk("%ld@end@ncclSend\\n", ts);
|
|
return 0;
|
|
};
|
|
|
|
int ncclRecvEntry(struct pt_regs *ctx){
|
|
u64 ts = bpf_ktime_get_ns();
|
|
u64 size_count = PT_REGS_PARM2(ctx);
|
|
u64 data_type = PT_REGS_PARM3(ctx);
|
|
bpf_trace_printk("%ld@start@ncclRecv@%ld\\n", ts, size_count * 8);
|
|
return 0;
|
|
};
|
|
|
|
int ncclRecvExit(struct pt_regs *ctx){
|
|
u64 ts = bpf_ktime_get_ns();
|
|
bpf_trace_printk("%ld@end@ncclRecv\\n", ts);
|
|
return 0;
|
|
};
|
|
"""
|
|
|
|
bpf = BccBPF("NCCLeBPF", text, ["-w"])
|
|
|
|
attach_config = [
|
|
{
|
|
"name": "NCCLSampler",
|
|
"exe_path": [
|
|
"/home/txx/data/miniconda3/envs/eACGM/lib/python3.12/site-packages/nvidia/nccl/lib/libnccl.so.2",
|
|
],
|
|
"exe_sym": [
|
|
"ncclAllReduce",
|
|
"ncclReduce",
|
|
"ncclBroadcast",
|
|
"ncclAllGather",
|
|
"ncclSend",
|
|
"ncclRecv",
|
|
]
|
|
},
|
|
]
|
|
|
|
sampler = eBPFSampler(bpf)
|
|
|
|
sampler.run(attach_config)
|
|
|
|
states = []
|
|
while True:
|
|
try:
|
|
samples = sampler.sample(time_stamp=1)
|
|
states += samples
|
|
# for sample in samples:
|
|
# print(sample)
|
|
# print("---")
|
|
except KeyboardInterrupt:
|
|
break
|
|
|
|
sampler.close()
|
|
collector = to_perfetto(states)
|
|
json.dump(collector, open("nccl.json", "w", encoding="utf-8"), indent=4) |