Repository: gaogaotiantian/viztracer (language: python)

Issue:
Viztracer result becomes wrong after some pytorch function call with `exclude_files`
## Environment

torch 2.5.1, python 3.11.10, ubuntu 22.04

## Problem

I don't need trace results inside `torch.export` module, so I use `exclude_files` to ignore it. The problem is if I trace with `exclude_files`, after calling `torch.distributed.pipelining.pipeline`, trace result of following calls only keep stack depth <= 4 result.

with `exclude_files`:
![image](https://github.com/user-attachments/assets/ca92be91-6b92-46c8-835c-8d80bf0ee59b)

without `exclude_files`:
![image](https://github.com/user-attachments/assets/b19fd8d4-9605-4cb7-ad27-0e8fcbe50f3d)

## Reproduce

1. install pytorch >= 2.5

2. get `exclude_files` by `import torch.export; print(torch.export.__file__)`

3. launch viztracer with following script and command:

```Python
import torch
import torch.export 
from torch.distributed.pipelining import pipeline
from torch.distributed.pipelining import SplitPoint
in_dim = 16
C = 64
HIDDEN = 256
layer_dims = [C] * 8
out_dim = 10
STAGE = 8


class MyNetworkBlock(torch.nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        mid_dim = HIDDEN
        self.lin1 = torch.nn.Conv2d(in_dim, mid_dim, 3, padding=1)
        self.lin2 = torch.nn.Conv2d(mid_dim, mid_dim, 3, padding=1)
        self.lin3 = torch.nn.Conv2d(mid_dim, out_dim, 3, padding=1)

    def forward(self, x):
        x = self.lin1(x)
        x = torch.relu(x)
        x = self.lin2(x)
        x = torch.relu(x)
        x = self.lin3(x)
        x = torch.relu(x)
        return x

class MyNetwork(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.num_layers = len(layer_dims)

        prev_dim = in_dim
        for i, dim in enumerate(layer_dims):
            super().add_module(f"layer{i}", MyNetworkBlock(prev_dim, dim))
            prev_dim = dim
        self.output_proj = torch.nn.Conv2d(layer_dims[-1], out_dim, 3, padding=1)

    def forward(self, x):
        for i in range(self.num_layers):
            layer = getattr(self, f"layer{i}")
            x = layer(x)

        return self.output_proj(x)

print(torch.export.__file__)
example_input = torch.randn(16, in_dim, 240, 320)

split_spec = {}
for j in range(STAGE - 1):
    split_spec[f"layer{j}"] = SplitPoint.END


pipe = pipeline(MyNetwork(), mb_args=(example_input,), split_spec=split_spec)

# after pytorch pipeline, the trace result (stack depth) of viztracer becomes wrong,
# we can only see results with stack depth <= 4.
class A:
    def a(self):
        print("a")
        self.b()
    def b(self):
        print("b")
        self.c()
    
    def c(self):
        print("c")
        self.d()
    def d(self):
        print("d")
        self.e()
    def e(self):
        print("e")
        # self.f()

a = A()
a.a()
```

```Bash
viztracer --exclude_files /path/to/torch/export --max_stack_depth 15 /path/to/script.py

```


Discussion / hints:
Which version of viztracer are you using? I can't repro this on 3.12, the stack looks fine. To be honest, I don't recommend the `--exclude_files` command for heavy filtering. The command itself has a relatively large overhead. It's helpful for understanding the code, but might be a bit slow for profiling.
I tried on python 3.11/3.12 on Mac OS, this problem exists on python 3.11, OK on 3.12.
I convert result from viztracer to tree view and tree filtering to analysis complex project such as PyTorch, the tree filtering is done in frontend and become slow if number of entries larger than 100000, so I need `exclude_files ` to reduce size. Another option is remove them after tracing, maybe slower than `exclude_files` since your code handle this in C.
Okay it only repros on 3.11, interesting. I can try to repro it and see if it's a VizTracer issue. For the latest version of Perfetto, you can just select the part you need and do a flamegraph, not sure if it's helpful. Removing after tracing is not as fast, but it does not use your program time so your result would be more accurate. If you can narrow down your interest area, `log_sparse` would be very helpful to reduce the entries.

Fix patch:
diff --git a/src/viztracer/modules/snaptrace.c b/src/viztracer/modules/snaptrace.c
index 7a333572..4658bd22 100644
--- a/src/viztracer/modules/snaptrace.c
+++ b/src/viztracer/modules/snaptrace.c
@@ -338,25 +338,17 @@ prepare_before_trace(TracerObject* self, int is_call, struct ThreadInfo** info_o
     }
 
     if (info->ignore_stack_depth > 0) {
-        if (is_call) {
-            info->ignore_stack_depth += 1;
-            return 0;
-        } else {
-            info->ignore_stack_depth -= 1;
-            return 0;
-        }
+        return 0;
     }
 
     if (CHECK_FLAG(self->check_flags, SNAPTRACE_MAX_STACK_DEPTH)) {
         if (is_call) {
             if (info->curr_stack_depth >= self->max_stack_depth) {
-                info->curr_stack_depth += 1;
                 return 0;
             }
         } else {
             if (info->curr_stack_depth > 0) {
                 if (info->curr_stack_depth > self->max_stack_depth) {
-                    info->curr_stack_depth -= 1;
                     return 0;
                 }
             }
@@ -374,7 +366,7 @@ tracer_pycall_callback(TracerObject* self, PyCodeObject* code)
     if (prepare_before_trace(self, 1, &info) <= 0) {
         // For now we think -1 and 0 should both return because we should not
         // have the -1 case.
-        return 0;
+        goto cleanup_ignore;
     }
 
     PyObject* co_filename = code->co_filename;
@@ -382,8 +374,7 @@ tracer_pycall_callback(TracerObject* self, PyCodeObject* code)
     if (!CHECK_FLAG(self->check_flags, SNAPTRACE_TRACE_SELF)) {
         if (self->lib_file_path && co_filename && PyUnicode_Check(co_filename) &&
                 startswith(PyUnicode_AsUTF8(co_filename), self->lib_file_path)) {
-            info->ignore_stack_depth += 1;
-            goto cleanup;
+            goto cleanup_ignore;
         }
     }
 
@@ -408,18 +399,16 @@ tracer_pycall_callback(TracerObject* self, PyCodeObject* code)
                 }
             }
             if (record == 0) {
-                info->ignore_stack_depth += 1;
-                goto cleanup;
+                goto cleanup_ignore;
             }
         } else {
-            goto cleanup;
+            goto cleanup_ignore;
         }
     }
 
     if (CHECK_FLAG(self->check_flags, SNAPTRACE_IGNORE_FROZEN)) {
         if (startswith(PyUnicode_AsUTF8(co_filename), "<frozen")) {
-            info->ignore_stack_depth += 1;
-            goto cleanup;
+            goto cleanup_ignore;
         }
     }
 
@@ -456,10 +445,17 @@ tracer_pycall_callback(TracerObject* self, PyCodeObject* code)
         log_func_args(info->stack_top, PyEval_GetFrame(), self->log_func_repr);
     }
 
-cleanup:
-
     info->curr_stack_depth += 1;
 
+    return 0;
+
+cleanup_ignore:
+
+    if (info) {
+        info->ignore_stack_depth += 1;
+        info->curr_stack_depth += 1;
+    }
+
     return 0;
 }
 
@@ -471,15 +467,13 @@ tracer_ccall_callback(TracerObject* self, PyCodeObject* code, PyObject* arg)
     if (prepare_before_trace(self, 1, &info) <= 0) {
         // For now we think -1 and 0 should both return because we should not
         // have the -1 case.
-        return 0;
+        goto cleanup_ignore;
     }
 
     PyCFunctionObject* cfunc = (PyCFunctionObject*) arg;
 
     if (cfunc->m_self == (PyObject*)self) {
-        info->ignore_stack_depth += 1;
-        info->curr_stack_depth += 1;
-        return 0;
+        goto cleanup_ignore;
     }
 
     // If it's a call, we need a new node, and we need to update the stack
@@ -496,6 +490,15 @@ tracer_ccall_callback(TracerObject* self, PyCodeObject* code, PyObject* arg)
 
     info->curr_stack_depth += 1;
 
+    return 0;
+
+cleanup_ignore:
+
+    if (info) {
+        info->ignore_stack_depth += 1;
+        info->curr_stack_depth += 1;
+    }
+
     return 0;
 }
 
@@ -507,7 +510,7 @@ tracer_pyreturn_callback(TracerObject* self, PyCodeObject* code, PyObject* arg)
     if (prepare_before_trace(self, 0, &info) <= 0) {
         // For now we think -1 and 0 should both return because we should not
         // have the -1 case.
-        return 0;
+        goto cleanup_ignore;
     }
 
     struct FunctionNode* stack_top = info->stack_top;
@@ -571,6 +574,20 @@ tracer_pyreturn_callback(TracerObject* self, PyCodeObject* code, PyObject* arg)
         info->curr_stack_depth -= 1;
     } 
 
+    return 0;
+
+cleanup_ignore:
+
+    if (info) {
+        if (info->curr_stack_depth > 0) {
+            info->curr_stack_depth -= 1;
+        }
+
+        if (info->ignore_stack_depth > 0) {
+            info->ignore_stack_depth -= 1;
+        }
+    }
+
     return 0;
 }
 
@@ -582,7 +599,7 @@ tracer_creturn_callback(TracerObject* self, PyCodeObject* code, PyObject* arg)
     if (prepare_before_trace(self, 0, &info) <= 0) {
         // For now we think -1 and 0 should both return because we should not
         // have the -1 case.
-        return 0;
+        goto cleanup_ignore;
     }
 
     struct FunctionNode* stack_top = info->stack_top;
@@ -639,6 +656,20 @@ tracer_creturn_callback(TracerObject* self, PyCodeObject* code, PyObject* arg)
         info->curr_stack_depth -= 1;
     }
 
+    return 0;
+
+cleanup_ignore:
+
+    if (info) {
+        if (info->curr_stack_depth > 0) {
+            info->curr_stack_depth -= 1;
+        }
+
+        if (info->ignore_stack_depth > 0) {
+            info->ignore_stack_depth -= 1;
+        }
+    }
+
     return 0;
 }
 


Test patch:
diff --git a/tests/test_regression.py b/tests/test_regression.py
index 89e90657..08eaf998 100644
--- a/tests/test_regression.py
+++ b/tests/test_regression.py
@@ -308,6 +308,30 @@ def test_issue162_os_popen(self):
                       script=issue162_code_os_popen, expected_stdout=r".*test_issue162.*")
 
 
+class TestIssue508(CmdlineTmpl):
+    def test_issue508(self):
+        script = """
+            import inspect
+            import os
+            import viztracer
+
+            exclude = os.path.dirname(inspect.__file__)
+
+            def call_self(n):
+                if n == 0:
+                    return
+                call_self(n - 1)
+
+            with viztracer.VizTracer(exclude_files=[exclude], max_stack_depth=6):
+                inspect.getsource(call_self)
+                call_self(10)
+        """
+
+        self.template(["python", "cmdline_test.py"], script=script,
+                      expected_output_file="result.json",
+                      expected_entries=6)
+
+
 file_timestamp_disorder = """
 def g():
     pass