From 4f999dab29dc1c4fc42ef1607098d3e32859a225 Mon Sep 17 00:00:00 2001
From: weli-l <1289113577@qq.com>
Date: Thu, 19 Jun 2025 08:58:58 +0800
Subject: [PATCH] support python function configurable

Signed-off-by: weli-l <1289113577@qq.com>
---
 systrace/build.sh                           |  3 +-
 systrace/config/PyFuncList                  | 24 ++++++++++++++++
 systrace/src/trace/python/pytorch_tracing.c |  8 ++++--
 systrace/src/trace/systrace_manager.cc      | 31 +++++++++++----------
 systrace/src/trace/systrace_manager.h       |  1 +
 5 files changed, 50 insertions(+), 17 deletions(-)
 create mode 100644 systrace/config/PyFuncList

diff --git a/systrace/build.sh b/systrace/build.sh
index 3040c97..9c177e1 100644
--- a/systrace/build.sh
+++ b/systrace/build.sh
@@ -2,6 +2,7 @@
 
 sudo dnf remove -y libunwind libunwind-devel 2>/dev/null || true
 mkdir -p build
+cp -f config/PyFuncList /etc/systrace/config/PyFuncList
 
 cd protos
 protoc --c_out=. systrace.proto
@@ -10,4 +11,4 @@ protoc --python_out=. systrace.proto
 cd ..
 cd build
 cmake ..
-make -j $(nproc)
+make -j $(nproc)
\ No newline at end of file
diff --git a/systrace/config/PyFuncList b/systrace/config/PyFuncList
new file mode 100644
index 0000000..9cf2a8a
--- /dev/null
+++ b/systrace/config/PyFuncList
@@ -0,0 +1,24 @@
+# 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:
+#
+#   module.path@ClassName@method_name  (for class methods)
+#   or
+#   module.path@function_name          (for standalone functions)
+#
+# Examples:
+#   torch.utils.data.dataloader@_BaseDataLoaderIter@__next__
+# -------------------------------------------
+GC
+torch.utils.data.dataloader@_BaseDataLoaderIter@__next__
+torch_npu@npu@synchronize
+torch_npu.npu@Event@synchronize
+torch_npu.npu@Event@wait
+torch_npu.npu@Stream@synchronize
+torch_npu.npu@Stream@wait_event
+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
diff --git a/systrace/src/trace/python/pytorch_tracing.c b/systrace/src/trace/python/pytorch_tracing.c
index ce3b651..95e85e6 100644
--- a/systrace/src/trace/python/pytorch_tracing.c
+++ b/systrace/src/trace/python/pytorch_tracing.c
@@ -605,10 +605,14 @@ void systrace_register_tracing(const char **names, int count, char **errors)
     PyGILState_STATE gstate = PyGILState_Ensure();
 
     init_tracing_data_array(count);
-    systrace_register_gc(errors);
 
-    for (int i = 1; i < count; i++)
+    for (int i = 0; i < count; i++)
     {
+        if (strcmp(names[i], "GC") == 0)
+        {
+            systrace_register_gc(errors);
+            continue;
+        }
         register_tracing_function(names[i], i, errors);
     }
 
diff --git a/systrace/src/trace/systrace_manager.cc b/systrace/src/trace/systrace_manager.cc
index 9b1606c..856bbb8 100644
--- a/systrace/src/trace/systrace_manager.cc
+++ b/systrace/src/trace/systrace_manager.cc
@@ -44,26 +44,29 @@ void PyTorchTrace::initialize()
 
 void PyTorchTrace::registerTracingFunctions()
 {
-    pytorch_tracing_functions_ = {
-        "GC",
-        "torch.utils.data.dataloader@_BaseDataLoaderIter@__next__",
-        "torch_npu@npu@synchronize",
-        "torch_npu.npu@Event@synchronize",
-        "torch_npu.npu@Event@wait",
-        "torch_npu.npu@Stream@synchronize",
-        "torch_npu.npu@Stream@wait_event",
-        "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"};
+    std::ifstream funcListFile(PyFuncListPath_);
+    std::string line;
+    if (!funcListFile.is_open())
+    {
+        STLOG(ERROR) << "Failed to open PyFuncList file";
+        return;
+    }
+    while (std::getline(funcListFile, line))
+    {
+        if (!line.empty() && line[0] != '#')
+        {
+            pytorch_tracing_functions_.push_back(line);
+        }
+    }
+
+    funcListFile.close();
 
     auto errors =
         pytorch_tracing_library_->Register(pytorch_tracing_functions_);
     for (size_t i = 0; i < pytorch_tracing_functions_.size(); ++i)
     {
         STLOG(INFO) << "Registered function: " << pytorch_tracing_functions_[i]
-                    << ", status: " << errors[i];
+                    << ", status: " << errors[i] << std::endl;
     }
 }
 
diff --git a/systrace/src/trace/systrace_manager.h b/systrace/src/trace/systrace_manager.h
index 46ee359..85e43d5 100644
--- a/systrace/src/trace/systrace_manager.h
+++ b/systrace/src/trace/systrace_manager.h
@@ -47,6 +47,7 @@ class PyTorchTrace
     std::mutex trace_mutex_;
 
     std::vector<std::string> pytorch_tracing_functions_;
+    std::string PyFuncListPath_ = "/etc/systrace/config/PyFuncList";
     pytorch_tracing::PyTorchTracingLibrary *pytorch_tracing_library_;
 };
 
-- 
Gitee