diff --git a/systrace/build.sh b/systrace/build.sh index 2ad64d0c99e72a980b610d0542b0e2dce7161baa..b5e3727e770782ac2ce176d843d15a8901e78ce7 100644 --- a/systrace/build.sh +++ b/systrace/build.sh @@ -14,7 +14,7 @@ cleanup() { } setup_config() { - [ -f "$PY_FUNC_LIST" ] && sudo cp -f "$PY_FUNC_LIST" "$CONFIG_DIR/" + [ -f "$PY_FUNC_LIST" ] && cp -f "$PY_FUNC_LIST" "$CONFIG_DIR/" } compile_proto() { diff --git a/systrace/config/PyFuncList b/systrace/config/PyFuncList index 49931b64920db081dec0432739205ca96baaf4d0..4a33382dfaa3856b70dc20d518f6cfa41401341f 100644 --- a/systrace/config/PyFuncList +++ b/systrace/config/PyFuncList @@ -1,15 +1,21 @@ -# PyTorch Tracing Functions Configuration File -# ------------------------------------------- -# This file lists Python functions that should be traced by sysTrace. -# Each line represents one function to be traced, specified in the format: +# ============================================ +# PyTorch Tracing Functions +# Configuration File +# ============================================ +# +# This file lists Python functions to be traced by sysTrace. +# Specify functions using the format: # # module.path@ClassName@method_name (for class methods) -# or -# module.path@function_name (for standalone functions) +# module.path@function_name (for standalone functions) # -# Examples: +# Example: # torch.utils.data.dataloader@_BaseDataLoaderIter@__next__ -# ------------------------------------------- +# ============================================ + +# ============================================ +# TORCH_NPU FUNCTIONS +# ============================================ torch.utils.data.dataloader@_BaseDataLoaderIter@__next__ torch_npu@npu@synchronize torch_npu.npu@Event@synchronize @@ -20,4 +26,28 @@ torch_npu.npu@Stream@wait_stream torch@autograd@backward torch@autograd@grad megatron.core.pipeline_parallel@schedules@forward_step -megatron.core.pipeline_parallel@schedules@backward_step \ No newline at end of file +megatron.core.pipeline_parallel@schedules@backward_step + +# ============================================ +# MINDSPORE FUNCTIONS +# ============================================ +mindspore.dataset.engine.datasets@Dataset@__iter__ +mindspore.hal@synchronize +mindspore.hal@Event@synchronize +mindspore.hal@Event@wait +mindspore.hal@Stream@synchronize +mindspore.hal@Stream@wait_event +mindspore.hal@Stream@wait_stream +mindspore.ops@GradOperation@__call__ +mindspore.ops@GradOperation@__call__ +mindspore.nn@Cell@__call__ +mindspore.ops@GradOperation@__call__ +mindspore.nn.TrainOneStepCell@__call__ +mindspore.parallel.nn.Pipeline@__call__ +mindspore.nn.PipelineCell@__call__ +mindspore.train.model@Model@_train_dataset_sink_process +mindspore.train.model@Model@_train_process +mindspore.train.model@Model@_train_network +mindspore.train.model@Model@_exec_preprocess +mindformers.core@MFLossMonitor@print_output_info +mindspore.train@Callback@on_train_step_begin \ No newline at end of file diff --git a/systrace/include/common/shared_constants.c b/systrace/include/common/shared_constants.c index f5f0ddff4a98bc27051b9312b7e88a6a56f06abd..7fc06b43c382b0e22fea4e92e98dc9b916e8ecee 100644 --- a/systrace/include/common/shared_constants.c +++ b/systrace/include/common/shared_constants.c @@ -11,14 +11,12 @@ static int shm_fd = -1; int init_shared_memory() { shm_fd = shm_open(SHM_NAME, O_CREAT | O_RDWR, 0666); - if (shm_fd == -1) - { + if (shm_fd == -1) { perror("shm_open failed"); return -1; } - if (ftruncate(shm_fd, sizeof(SharedData)) == -1) - { + if (ftruncate(shm_fd, sizeof(SharedData)) == -1) { perror("ftruncate failed"); close(shm_fd); return -1; @@ -26,29 +24,34 @@ int init_shared_memory() shared_data = mmap(NULL, sizeof(SharedData), PROT_READ | PROT_WRITE, MAP_SHARED, shm_fd, 0); - if (shared_data == MAP_FAILED) - { + if (shared_data == MAP_FAILED) { perror("mmap failed"); close(shm_fd); return -1; } + pthread_mutex_lock(&shared_data->g_trace_mutex); + if (shared_data->initialized) { + pthread_mutex_unlock(&shared_data->g_trace_mutex); + return 0; + } + static pthread_mutexattr_t mutex_attr; - if (pthread_mutexattr_init(&mutex_attr) != 0) - { + if (pthread_mutexattr_init(&mutex_attr) != 0) { perror("pthread_mutexattr_init failed"); + pthread_mutex_unlock(&shared_data->g_trace_mutex); return -1; } - if (pthread_mutexattr_setpshared(&mutex_attr, PTHREAD_PROCESS_SHARED) != 0) - { + if (pthread_mutexattr_setpshared(&mutex_attr, PTHREAD_PROCESS_SHARED) != 0) { perror("pthread_mutexattr_setpshared failed"); + pthread_mutex_unlock(&shared_data->g_trace_mutex); return -1; } - if (pthread_mutex_init(&shared_data->g_trace_mutex, &mutex_attr) != 0) - { + if (pthread_mutex_init(&shared_data->g_trace_mutex, &mutex_attr) != 0) { perror("pthread_mutex_init failed"); + pthread_mutex_unlock(&shared_data->g_trace_mutex); return -1; } @@ -72,9 +75,11 @@ int init_shared_memory() shared_data->need_dump_L2_once = false; shared_data->need_dump_L3_once = false; + shared_data->initialized = true; + + pthread_mutex_unlock(&shared_data->g_trace_mutex); return 0; } - SharedData *get_shared_data() { if (shared_data == NULL) @@ -117,8 +122,6 @@ bool checkAndUpdateTimer(int level) { bool* timer_active = NULL; time_t* start_time = NULL; const char* level_name = ""; - bool *dumped = false; - bool *need_dump_once = NULL; switch(level) { case 1: // L1 @@ -127,8 +130,6 @@ bool checkAndUpdateTimer(int level) { timer_active = &shared_data->g_L1_timer_active; start_time = &shared_data->g_L1_start_time; level_name = "L1"; - dumped = &shared_data->dumped_L1; - need_dump_once = &shared_data->need_dump_L1_once; break; case 2: // L2 dump_flag = &shared_data->g_dump_L2; @@ -136,17 +137,13 @@ bool checkAndUpdateTimer(int level) { timer_active = &shared_data->g_L2_timer_active; start_time = &shared_data->g_L2_start_time; level_name = "L2"; - dumped = &shared_data->dumped_L2; - need_dump_once = &shared_data->need_dump_L2_once; break; - case 3: // L2 + case 3: // L3 dump_flag = &shared_data->g_dump_L3; interval = &shared_data->g_dump_L3_interval; timer_active = &shared_data->g_L3_timer_active; start_time = &shared_data->g_L3_start_time; level_name = "L3"; - dumped = &shared_data->dumped_L3; - need_dump_once = &shared_data->need_dump_L3_once; break; default: pthread_mutex_unlock(&shared_data->g_trace_mutex); @@ -162,14 +159,11 @@ bool checkAndUpdateTimer(int level) { } else if (*timer_active) { time_t now = time(NULL); - double elapsed = difftime(now, *start_time) / 60; + double elapsed = difftime(now, *start_time) / 60.0; if (elapsed >= *interval) { *dump_flag = false; *timer_active = false; - if (!*dumped) { - *need_dump_once = true; - } } else { result = true; } diff --git a/systrace/include/common/shared_constants.h b/systrace/include/common/shared_constants.h index 22ebb564a36cbb397b4cc26a9fcb5bac2e712ca2..d53b7a95868f78a538527979a029e72c7094325c 100644 --- a/systrace/include/common/shared_constants.h +++ b/systrace/include/common/shared_constants.h @@ -37,6 +37,7 @@ extern "C" time_t g_L2_start_time; time_t g_L3_start_time; pthread_mutex_t g_trace_mutex; + bool initialized; } SharedData; int init_shared_memory(); diff --git a/systrace/include/common/util.cc b/systrace/include/common/util.cc index 4c2a6e0eedacffa41cbe69a9bddf8ec19610fd73..8fb61e4bdb607d9c719bd260826da8a31bb1e3ba 100644 --- a/systrace/include/common/util.cc +++ b/systrace/include/common/util.cc @@ -123,11 +123,11 @@ void LoadEnvironmentVariables() auto loadStr = [](const char *name) { return env::EnvVarRegistry::GetEnvVar(name); }; - config.rank = loadInt("RANK"); + config.rank = loadInt("RANK") ? loadInt("RANK") : loadInt("RANK_ID"); config.job_name = loadStr("ENV_ARGO_WORKFLOW_NAME"); - config.local_rank = loadInt("LOCAL_RANK"); + config.local_rank = loadInt("LOCAL_RANK") ? loadInt("LOCAL_RANK") : loadInt("DEVICE_ID"); config.local_world_size = loadInt("LOCAL_WORLD_SIZE"); - config.world_size = loadInt("WORLD_SIZE"); + config.world_size = loadInt("WORLD_SIZE") ? loadInt("WORLD_SIZE") : loadInt("RANK_SIZE"); config.rank_str = "[RANK " + std::to_string(config.rank) + "] "; } diff --git a/systrace/server/monitor_server.cpp b/systrace/server/monitor_server.cpp index 671c674fddbd08067435ef43b4df5a887af9a142..f81f007bef8fb648cf47150ba5ecc3c7675ed07b 100644 --- a/systrace/server/monitor_server.cpp +++ b/systrace/server/monitor_server.cpp @@ -16,6 +16,7 @@ MonitorServer &MonitorServer::getInstance() { instance_ = new MonitorServer(); instance_->start(); + std::atexit(cleanup); }); return *instance_; } @@ -75,8 +76,10 @@ void MonitorServer::start() void MonitorServer::stop() { + should_run_ = false; if (server_fd_ != -1) { + shutdown(server_fd_, SHUT_RDWR); close(server_fd_); server_fd_ = -1; } @@ -91,7 +94,7 @@ void MonitorServer::stop() void MonitorServer::server_thread_func() { - while (true) + while (should_run_) { struct sockaddr_un client_addr; socklen_t client_len = sizeof(client_addr); diff --git a/systrace/server/monitor_server.hpp b/systrace/server/monitor_server.hpp index 5af92bde5328f60a4a6d3cbb9d42b5676814f39f..dffa76a67d96eb06a2ff3d5b9a81f55f4b3addbd 100644 --- a/systrace/server/monitor_server.hpp +++ b/systrace/server/monitor_server.hpp @@ -4,6 +4,7 @@ #include "../include/common/shared_constants.h" #include #include +#include #include class MonitorServer @@ -13,6 +14,14 @@ class MonitorServer void start(); void stop(); + static void cleanup() + { + if (instance_) + { + instance_->stop(); + cleanup_shared_memory(); + } + } MonitorServer(const MonitorServer &) = delete; MonitorServer &operator=(const MonitorServer &) = delete; @@ -31,6 +40,7 @@ class MonitorServer int server_fd_{-1}; std::thread server_thread_; + std::atomic should_run_{true}; static MonitorServer *instance_; inline static std::once_flag init_flag_; diff --git a/systrace/src/ascend/hook.cc b/systrace/src/ascend/hook.cc index c1681285b5a96d85dcc0b3a2191f8611e7a66224..01c55b1b7dbda196a969fe402c81702f63d54168 100644 --- a/systrace/src/ascend/hook.cc +++ b/systrace/src/ascend/hook.cc @@ -1,8 +1,56 @@ -#include "hook.h" -#include "../src/trace/systrace_manager.h" -#include #include +#include +#include +#include #include +#include +#include +#include "../src/trace/systrace_manager.h" +#include "hook.h" + +static std::string get_mindspore_lib_path() { + const char* cmd = "python -c \"import mindspore as ms; import os; print(os.path.join(os.path.dirname(ms.__file__), 'lib/libmindspore_backend.so'))\""; + FILE* pipe = popen(cmd, "r"); + if (!pipe) return ""; + + char buffer[1024]; + std::string result; + if (fgets(buffer, sizeof(buffer), pipe) != nullptr) { + result = buffer; + result.erase(result.find_last_not_of("\n") + 1); + } + pclose(pipe); + return result; +} + +extern "C" void _ZN9mindspore11distributed10InitializeEv() { + std::call_once(init_flag, []() { + std::string so_path = get_mindspore_lib_path(); + if (so_path.empty()) { + fprintf(stderr, "[ERROR] Failed to find libmindspore_backend.so\n"); + return; + } + + void* handle = dlopen(so_path.c_str(), RTLD_LAZY); + if (!handle) { + fprintf(stderr, "[ERROR] dlopen failed: %s\n", dlerror()); + return; + } + + original_Initialize = (void (*)())dlsym(handle, "_ZN9mindspore11distributed10InitializeEv"); + if (!original_Initialize) { + fprintf(stderr, "[ERROR] dlsym failed: %s\n", dlerror()); + return; + } + ::systrace::SysTrace::getInstance(); + }); + + if (!original_Initialize) { + fprintf(stderr, "[ERROR] Original function not loaded\n"); + return; + } + original_Initialize(); +} #ifdef __cplusplus extern "C" diff --git a/systrace/src/ascend/hook.h b/systrace/src/ascend/hook.h index 31f2581beca83b7663654b3a9c66e634a2f1803a..e0e1f34a09edc35c80e27de14630f5bb7c2b5325 100644 --- a/systrace/src/ascend/hook.h +++ b/systrace/src/ascend/hook.h @@ -37,6 +37,9 @@ extern "C" aclInitFn orig_aclInit = nullptr; aclrtMapMemFn orig_aclrtMapMem = nullptr; aclrtLaunchKernelFn orig_aclrtLaunchKernel = nullptr; + + static std::once_flag init_flag; + static void (*original_Initialize)() = nullptr; #ifdef __cplusplus } #endif \ No newline at end of file diff --git a/systrace/src/cann/common_hook.c b/systrace/src/cann/common_hook.c index a6964529b0578d6a1727a957773ec5e66da4cc89..8d8410ecc7add04f66f08c4c15403745368209dd 100644 --- a/systrace/src/cann/common_hook.c +++ b/systrace/src/cann/common_hook.c @@ -35,7 +35,7 @@ unw_word_t get_so_base(unw_word_t addr) } void get_log_filename(char *buf, size_t buf_size, const char *path_suffix) { - const char *rank_str = getenv("RANK"); + const char *rank_str = getenv("RANK") ? getenv("RANK") : getenv("RANK_ID"); int rank = rank_str ? atoi(rank_str) : 0; char path[PATH_MAX] = {0}; diff --git a/systrace/src/cann/io_hook.c b/systrace/src/cann/io_hook.c index 986f1a11fa3d89c867830f18cde551a3e45c81a6..7e22e21ce53f44f978074f405ebc536a99ed4195 100644 --- a/systrace/src/cann/io_hook.c +++ b/systrace/src/cann/io_hook.c @@ -30,7 +30,6 @@ typedef int (*halFCloseFunc_t)(FILE *stream); typedef int (*halFFlushFunc_t)(FILE *stream); typedef int (*halRemoveFunc_t)(const char *filename); typedef int (*halRenameFunc_t)(const char *oldname, const char *newname); -typedef int (*halOpenFunc_t)(const char *pathname, int flags, mode_t mode); typedef int (*halCloseFunc_t)(int fd); typedef int (*halFsyncFunc_t)(int fd); typedef int (*halMkdirFunc_t)(const char *path, mode_t mode); @@ -48,7 +47,6 @@ static halFCloseFunc_t orig_fclose = NULL; static halFFlushFunc_t orig_fflush = NULL; static halRemoveFunc_t orig_remove = NULL; static halRenameFunc_t orig_rename = NULL; -static halOpenFunc_t orig_open = NULL; static halCloseFunc_t orig_close = NULL; static halFsyncFunc_t orig_fsync = NULL; static halMkdirFunc_t orig_mkdir = NULL; @@ -182,7 +180,7 @@ static void add_io_entry(int fd, uint64_t start_us, uint64_t duration, IOType op entry->file_name.data = (uint8_t *)filename; entry->file_name.len = strlen(filename); - const char *rank_str = getenv("RANK"); + const char *rank_str = getenv("RANK") ? getenv("RANK") : getenv("RANK_ID"); entry->rank = rank_str ? atoi(rank_str) : 0; td->io->n_io_entries++; @@ -206,7 +204,6 @@ int init_io_trace() { orig_fflush = (halFFlushFunc_t)dlsym(lib, "fflush"); orig_remove = (halRemoveFunc_t)dlsym(lib, "remove"); orig_rename = (halRenameFunc_t)dlsym(lib, "rename"); - orig_open = (halOpenFunc_t)dlsym(lib, "open"); orig_close = (halCloseFunc_t)dlsym(lib, "close"); orig_fsync = (halFsyncFunc_t)dlsym(lib, "fsync"); orig_mkdir = (halMkdirFunc_t)dlsym(lib, "mkdir"); @@ -217,7 +214,7 @@ int init_io_trace() { if (!orig_fread || !orig_fwrite || !orig_read || !orig_write || !orig_fopen || !orig_fclose || !orig_fflush || !orig_remove || - !orig_rename || !orig_open || !orig_close || !orig_fsync || + !orig_rename || !orig_close || !orig_fsync || !orig_mkdir || !orig_rmdir || !orig_unlink || !orig_opendir || !orig_closedir) { fprintf(stderr, "dlsym failed: %s\n", dlerror()); @@ -388,31 +385,6 @@ int rename(const char *oldname, const char *newname) { return ret; } -int open(const char *pathname, int flags, ...) { - if (!orig_open) { - init_io_trace(); - } - - mode_t mode = 0; - if (flags & O_CREAT) { - va_list ap; - va_start(ap, flags); - mode = va_arg(ap, mode_t); - va_end(ap); - } - - uint64_t start_us = get_current_us(); - int ret = orig_open(pathname, flags, mode); - uint64_t end_us = get_current_us(); - - if (ret >= 0) { - add_io_entry(ret, start_us, end_us - start_us, IOTYPE__IO_OPEN); - } - - write_protobuf_to_file(); - return ret; -} - int close(int fd) { if (!orig_close) { init_io_trace(); diff --git a/systrace/src/cann/mem_hook.c b/systrace/src/cann/mem_hook.c index 6569fe5f81e6d8a74badd00acecfd01524320266..f9cda8fa55aeb51a6c5612905c01f65777f6cc95 100644 --- a/systrace/src/cann/mem_hook.c +++ b/systrace/src/cann/mem_hook.c @@ -138,7 +138,7 @@ static ThreadData *get_thread_data() td = calloc(1, sizeof(ThreadData)); td->proc_mem = calloc(1, sizeof(ProcMem)); proc_mem__init(td->proc_mem); - const char *rank_str = getenv("RANK"); + const char *rank_str = getenv("RANK") ? getenv("RANK") : getenv("RANK_ID"); int rank = rank_str ? atoi(rank_str) : 0; td->proc_mem->pid = rank; td->last_log_time = time(NULL); diff --git a/systrace/src/mspti/json_file_writer.h b/systrace/src/mspti/json_file_writer.h index 2cc265817e60f4c15717954b4ad7c890fb957bfb..cfa8eddca7a16c36496099c2adc5eb4730453334 100644 --- a/systrace/src/mspti/json_file_writer.h +++ b/systrace/src/mspti/json_file_writer.h @@ -43,7 +43,7 @@ public: std::string filenameWithRank = saveFilename; this->markerActivityBuffer = std::make_unique>(); - const char* localRankCStr = std::getenv("RANK"); + const char* localRankCStr = std::getenv("RANK") ? std::getenv("RANK") : std::getenv("RANK_ID"); if (localRankCStr == nullptr) { localRankCStr = "-1"; } @@ -135,7 +135,7 @@ public: } void hcclActivityFormatToCSV() { - if (!checkAndUpdateTimer(1) && !need_dump_L1_once()) { + if (!checkAndUpdateTimer(1)) { return; } std::lock_guard lock(this->buffermtx); @@ -156,13 +156,6 @@ public: } } this->markerActivityBuffer->clear(); - SharedData *shared_data = get_shared_data(); - if (!shared_data) - { - return; - } - shared_data->dumped_L1 = true; - shared_data->need_dump_L1_once = false; } else { std::cout << "File is not open" << std::endl; } diff --git a/systrace/src/os/os_probe.c b/systrace/src/os/os_probe.c index 3784f0cd66f0c0e81f9db263d397fd00e5248437..3aac1ee0363e00fdd51e181467d970531a55610b 100644 --- a/systrace/src/os/os_probe.c +++ b/systrace/src/os/os_probe.c @@ -137,8 +137,8 @@ static __u64 get_unix_time_from_uptime(__u64 uptime) } void initialize_osprobe() { - const char *rank_str = getenv("RANK"); - const char *local_rank_str = getenv("LOCAL_RANK"); + const char *rank_str = getenv("RANK") ? getenv("RANK") : getenv("RANK_ID"); + const char *local_rank_str = getenv("LOCAL_RANK") ? getenv("LOCAL_RANK") : getenv("DEVICE_ID"); rank = rank_str ? atoi(rank_str) : 0; local_rank = local_rank_str? atoi(local_rank_str) : 0; get_sys_boot_time(); diff --git a/systrace/src/trace/systrace_manager.cc b/systrace/src/trace/systrace_manager.cc index c678545e494ebf3bb0aaa1b484013241ad814959..cb9fb0515d8d5d5b9deeee3ef46a5cfde02cd9bd 100644 --- a/systrace/src/trace/systrace_manager.cc +++ b/systrace/src/trace/systrace_manager.cc @@ -202,6 +202,7 @@ SysTrace &SysTrace::getInstance() { instance_ = new SysTrace(); instance_->initializeSystem(); + std::atexit(cleanup); }); return *instance_; } @@ -249,6 +250,7 @@ void SysTrace::stopOsProbePoller() os_probe_.join(); } } +#endif void SysTrace::stopEventPoller() { @@ -258,7 +260,6 @@ void SysTrace::stopEventPoller() event_poller_.join(); } } -#endif void SysTrace::eventPollerMain() { diff --git a/systrace/src/trace/systrace_manager.h b/systrace/src/trace/systrace_manager.h index 03c5c4233542ef18036edc02d5355ece384077e9..64fea22ac7832579bfc57bb6079876c038678cfd 100644 --- a/systrace/src/trace/systrace_manager.h +++ b/systrace/src/trace/systrace_manager.h @@ -67,6 +67,12 @@ class SysTrace void startEventPoller(); void stopEventPoller(); void eventPollerMain(); + static void cleanup() { + #ifdef HAS_BTF_SUPPORT + stopOsProbePoller(); + #endif + instance_->stopEventPoller(); + } inline static SysTrace *instance_ = nullptr; inline static std::once_flag init_flag_;