make subprocess child shows different process name

src/viztracer/main.py

@@ -263,6 +263,7 @@ def parse(self, argv: List[str]) -> VizProcedureResult:
             "sanitize_function_name": options.sanitize_function_name,
             "dump_raw": True,
             "minimize_memory": options.minimize_memory,
+            "process_name": None,
         }
 
         return True, None
@@ -318,6 +319,12 @@ def run_code(self, code: Any, global_dict: Dict[str, Any]) -> VizProcedureResult
         options = self.options
         self.parent_pid = os.getpid()
 
+        if options.subprocess_child:
+            if options.cmd_string is not None:
+                self.init_kwargs["process_name"] = "python -c"
+            else:
+                self.init_kwargs["process_name"] = sys.argv[0]
+
         tracer = VizTracer(**self.init_kwargs)
         self.tracer = tracer
 

src/viztracer/modules/snaptrace.c

@@ -650,20 +650,27 @@ snaptrace_load(TracerObject* self, PyObject* args)
         PyObject* dict = PyDict_New();
         PyObject* args = PyDict_New();
         PyObject* process_name_string = PyUnicode_FromString("process_name");
-        PyObject* current_process_method = PyObject_GetAttrString(multiprocessing_module, "current_process");
-        if (!current_process_method) {
-            perror("Failed to access multiprocessing.current_process()");
-            exit(-1);
-        }
-        PyObject* current_process = PyObject_CallObject(current_process_method, NULL);
-        if (!current_process_method) {
-            perror("Failed to access multiprocessing.current_process()");
-            exit(-1);
-        }
-        PyObject* process_name = PyObject_GetAttrString(current_process, "name");
+        PyObject* process_name = NULL;
 
-        Py_DECREF(current_process_method);
-        Py_DECREF(current_process);
+        if (self->process_name) {
+            process_name = self->process_name;
+            Py_INCREF(process_name);
+        } else {
+            PyObject* current_process_method = PyObject_GetAttrString(multiprocessing_module, "current_process");
+            if (!current_process_method) {
+                perror("Failed to access multiprocessing.current_process()");
+                exit(-1);
+            }
+            PyObject* current_process = PyObject_CallObject(current_process_method, NULL);
+            if (!current_process_method) {
+                perror("Failed to access multiprocessing.current_process()");
+                exit(-1);
+            }
+            process_name = PyObject_GetAttrString(current_process, "name");
+            Py_DECREF(current_process_method);
+            Py_DECREF(current_process);
+        }
+
         PyDict_SetItemString(dict, "ph", ph_M);
         PyDict_SetItemString(dict, "pid", pid);
         PyDict_SetItemString(dict, "tid", pid);
@@ -898,22 +905,30 @@ snaptrace_dump(TracerObject* self, PyObject* args)
     // == Load the metadata first ==
     //    Process Name
     {
-        PyObject* current_process_method = PyObject_GetAttrString(multiprocessing_module, "current_process");
-        if (!current_process_method) {
-            perror("Failed to access multiprocessing.current_process()");
-            exit(-1);
-        }
-        PyObject* current_process = PyObject_CallObject(current_process_method, NULL);
-        if (!current_process_method) {
-            perror("Failed to access multiprocessing.current_process()");
-            exit(-1);
+        PyObject* process_name = NULL;
+        if (self->process_name) {
+            process_name = self->process_name;
+            Py_INCREF(process_name);
+        } else {
+            PyObject* current_process_method = PyObject_GetAttrString(multiprocessing_module, "current_process");
+            if (!current_process_method) {
+                perror("Failed to access multiprocessing.current_process()");
+                exit(-1);
+            }
+            PyObject* current_process = PyObject_CallObject(current_process_method, NULL);
+            if (!current_process_method) {
+                perror("Failed to access multiprocessing.current_process()");
+                exit(-1);
+            }
+            process_name = PyObject_GetAttrString(current_process, "name");
+            Py_DECREF(current_process_method);
+            Py_DECREF(current_process);
         }
-        PyObject* process_name = PyObject_GetAttrString(current_process, "name");
 
-        Py_DECREF(current_process_method);
-        Py_DECREF(current_process);
-        fprintf(fptr, "{\"ph\":\"M\",\"pid\":%lu,\"tid\":%lu,\"name\":\"process_name\",\"args\":{\"name\":\"%s\"}},",
-                pid, pid, PyUnicode_AsUTF8(process_name));
+        fprintf(fptr, "{\"ph\":\"M\",\"pid\":%lu,\"tid\":%lu,\"name\":\"process_name\",\"args\":{\"name\":\"",
+                pid, pid);
+        fprint_escape(fptr, PyUnicode_AsUTF8(process_name));
+        fprintf(fptr, "\"}},");
         Py_DECREF(process_name);
     }
 
@@ -1140,11 +1155,12 @@ snaptrace_config(TracerObject* self, PyObject* args, PyObject* kw)
     static char* kwlist[] = {"verbose", "lib_file_path", "max_stack_depth", 
             "include_files", "exclude_files", "ignore_c_function", "ignore_frozen",
             "log_func_retval", "log_func_args", "log_async", "trace_self",
-            "min_duration",
+            "min_duration", "process_name",
             NULL};
     int kw_verbose = -1;
     int kw_max_stack_depth = 0;
     char* kw_lib_file_path = NULL;
+    PyObject* kw_process_name = NULL;
     PyObject* kw_include_files = NULL;
     PyObject* kw_exclude_files = NULL;
     int kw_ignore_c_function = -1;
@@ -1154,7 +1170,7 @@ snaptrace_config(TracerObject* self, PyObject* args, PyObject* kw)
     int kw_log_async = -1;
     int kw_trace_self = -1;
     double kw_min_duration = 0;
-    if (!PyArg_ParseTupleAndKeywords(args, kw, "|isiOOppppppd", kwlist,
+    if (!PyArg_ParseTupleAndKeywords(args, kw, "|isiOOppppppdO", kwlist,
             &kw_verbose,
             &kw_lib_file_path,
             &kw_max_stack_depth,
@@ -1166,7 +1182,8 @@ snaptrace_config(TracerObject* self, PyObject* args, PyObject* kw)
             &kw_log_func_args,
             &kw_log_async,
             &kw_trace_self,
-            &kw_min_duration)) {
+            &kw_min_duration,
+            &kw_process_name)) {
         return NULL;
     }
 
@@ -1190,6 +1207,15 @@ snaptrace_config(TracerObject* self, PyObject* args, PyObject* kw)
         strcpy(self->lib_file_path, kw_lib_file_path);
     }
 
+    if (kw_process_name && !Py_IsNone(kw_process_name)) {
+        if (!PyUnicode_CheckExact(kw_process_name)) {
+            PyErr_SetString(PyExc_TypeError, "process_name must be a string");
+            return NULL;
+        }
+        Py_INCREF(kw_process_name);
+        Py_XSETREF(self->process_name, kw_process_name); 
+    }
+
     if (kw_ignore_c_function == 1) {
         SET_FLAG(self->check_flags, SNAPTRACE_IGNORE_C_FUNCTION);
     } else if (kw_ignore_c_function == 0) {

src/viztracer/modules/snaptrace.h

@@ -69,6 +69,7 @@ typedef struct {
     int verbose;
     char* lib_file_path;
     int max_stack_depth;
+    PyObject* process_name;
     PyObject* include_files;
     PyObject* exclude_files;
     double min_duration;

src/viztracer/tracer.py

@@ -28,7 +28,8 @@ def __init__(
             log_gc: bool = False,
             log_async: bool = False,
             trace_self: bool = False,
-            min_duration: float = 0) -> None:
+            min_duration: float = 0,
+            process_name: Optional[str] = None) -> None:
         self.initialized = False
         self.enable = False
         self.parsed = False
@@ -53,6 +54,7 @@ def __init__(
         self.total_entries = 0
         self.gc_start_args: Dict[str, int] = {}
         self.initialized = True
+        self.process_name = process_name
 
     @property
     def max_stack_depth(self) -> int:
@@ -237,6 +239,7 @@ def config(self) -> None:
             "log_async": self.log_async,
             "trace_self": self.trace_self,
             "min_duration": self.min_duration,
+            "process_name": self.process_name,
         }
 
         self._tracer.config(**cfg)

src/viztracer/viztracer.py

@@ -42,6 +42,7 @@ def __init__(self,
                  minimize_memory: bool = False,
                  dump_raw: bool = False,
                  sanitize_function_name: bool = False,
+                 process_name: Optional[str] = None,
                  output_file: str = "result.json",
                  plugins: Sequence[Union[VizPluginBase, str]] = []) -> None:
         super().__init__(
@@ -58,6 +59,7 @@ def __init__(self,
             log_async=log_async,
             trace_self=trace_self,
             min_duration=min_duration,
+            process_name=process_name,
         )
         self._tracer: Any
         self.verbose = verbose

tests/test_multiprocess.py

@@ -7,6 +7,7 @@
 import sys
 import tempfile
 import unittest
+import json
 
 from .cmdline_tmpl import CmdlineTmpl
 
@@ -44,6 +45,12 @@ def fib(n):
 print(subprocess.call(["python", "-m", "timeit", "-n", "100", "'1+1'"]))
 """
 
+file_subprocess_code_string = """
+import subprocess
+p = subprocess.Popen(['python', '-c', 'import time;time.sleep(0.5)'])
+p.wait()
+"""
+
 file_fork = """
 import os
 import time
@@ -226,6 +233,13 @@ def setUp(self):
     def tearDown(self):
         os.remove("child.py")
 
+    def assertSubprocessName(self, name, data):
+        for entry in data["traceEvents"]:
+            if entry["name"] == "process_name" and entry["args"]["name"] == name:
+                break
+        else:
+            self.fail("no matching subprocess name")
+
     def test_basic(self):
         def check_func(data):
             pids = set()
@@ -240,12 +254,32 @@ def test_child_process(self):
             self.template(["viztracer", "-o", os.path.join(tmpdir, "result.json"), "--subprocess_child", "child.py"],
                           expected_output_file=None)
             self.assertEqual(len(os.listdir(tmpdir)), 1)
+            with open(os.path.join(tmpdir, os.listdir(tmpdir)[0])) as f:
+                self.assertSubprocessName("child.py", json.load(f))
 
     def test_module(self):
-        self.template(["viztracer", "-o", "result.json", "cmdline_test.py"],
-                      expected_output_file="result.json",
-                      expected_stdout=".*100 loops.*",
-                      script=file_subprocess_module)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            output_file = os.path.join(tmpdir, "result.json")
+            self.template(["viztracer", "-o", output_file, "cmdline_test.py"],
+                          expected_output_file=output_file,
+                          expected_stdout=".*100 loops.*",
+                          script=file_subprocess_module,
+                          check_func=lambda data: self.assertSubprocessName("timeit", data))
+
+    def test_code_string(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            output_path = os.path.join(tmpdir, "result.json")
+            self.template(["viztracer", "-o", output_path, "cmdline_test.py"],
+                          expected_output_file=output_path,
+                          script=file_subprocess_code_string,
+                          check_func=lambda data: self.assertSubprocessName("python -c", data))
+
+            # this is for coverage
+            self.template(['viztracer', '-o', os.path.join(tmpdir, "result.json"), '--subprocess_child',
+                           '-c', 'import time;time.sleep(0.5)'], expected_output_file=None)
+            self.assertEqual(len(os.listdir(tmpdir)), 1)
+            with open(os.path.join(tmpdir, os.listdir(tmpdir)[0])) as f:
+                self.assertSubprocessName("python -c", json.load(f))
 
     def test_nested(self):
         def check_func(data):