Support pystack as profiler agent to collect native stacktraces and core dumps

tvalentyn · tvalentyn · commit 0be3fc83084d · 2026-07-24T04:34:46.000-07:00
diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py
@@ -1743,6 +1743,10 @@ def validate(self, validator):
           _LOGGER.info(
               'Setting --profile_location to %s since profiling is enabled.',
               self.profile_location)
+
+      if self.profiler_agent == 'pystack_coredump':
+        debug_options = self.view_as(DebugOptions)
+        debug_options.add_experiment('core_pattern=/tmp/core.%e.%p')
     return errors
 
 
diff --git a/sdks/python/apache_beam/options/pipeline_options_test.py b/sdks/python/apache_beam/options/pipeline_options_test.py
@@ -686,6 +686,14 @@ def test_profiling_agent_is_exclusive_with_legacy_profiling_options(self):
     self.assertTrue(
         any('--profiler_agent is mutually exclusive' in err for err in errors))
 
+  def test_profiling_agent_pystack_coredump_adds_core_pattern(self):
+    options = PipelineOptions(['--profiler_agent=pystack_coredump'])
+    validator = PipelineOptionsValidator(options, None)
+    self.assertEqual(validator.validate(), [])
+    debug_options = options.view_as(DebugOptions)
+    self.assertEqual(
+        debug_options.lookup_experiment('core_pattern'), '/tmp/core.%e.%p')
+
   def test_profile_location_defaulting_and_opt_out(self):
     options = PipelineOptions(
         ['--profiler_agent=memray', '--temp_location=gs://bucket/temp'])
diff --git a/sdks/python/container/base_image_requirements_manual.txt b/sdks/python/container/base_image_requirements_manual.txt
@@ -42,6 +42,7 @@ guppy3
 memray==1.19.3
 mmh3  # Optimizes execution of some Beam codepaths. TODO: Make it Beam's dependency.
 nltk  # Commonly used for natural language processing.
+pystack
 google-crc32c
 scipy
 scikit-learn
diff --git a/sdks/python/container/boot.go b/sdks/python/container/boot.go
@@ -365,6 +365,7 @@ func launchSDKProcess() error {
 
 				logger.Printf(ctx, "Executing Python (%v): %v %v", envStr, currentProg, strings.Join(currentArgs, " "))
 				cmd := StartCommandEnv(currentEnv, os.Stdin, bufLogger, bufLogger, currentProg, currentArgs...)
+				logger.Printf(ctx, "Started worker %s with PID %d", workerId, cmd.Process.Pid)
 				activePids = append(activePids, cmd.Process.Pid)
 				workerMu.Unlock()
 
diff --git a/sdks/python/container/ml/py310/base_image_requirements.txt b/sdks/python/container/ml/py310/base_image_requirements.txt
@@ -181,6 +181,7 @@ PyMySQL==1.2.0
 pyOpenSSL==26.2.0
 pyparsing==3.3.2
 pyproject_hooks==1.2.0
+pystack==1.7.0
 pytest==9.1.1
 pytest-timeout==2.4.0
 pytest-xdist==3.8.0
diff --git a/sdks/python/container/ml/py310/gpu_image_requirements.txt b/sdks/python/container/ml/py310/gpu_image_requirements.txt
@@ -256,6 +256,7 @@ PyMySQL==1.2.0
 pyOpenSSL==26.2.0
 pyparsing==3.3.2
 pyproject_hooks==1.2.0
+pystack==1.7.0
 pytest==9.1.1
 pytest-timeout==2.4.0
 pytest-xdist==3.8.0
diff --git a/sdks/python/container/ml/py311/base_image_requirements.txt b/sdks/python/container/ml/py311/base_image_requirements.txt
@@ -180,6 +180,7 @@ PyMySQL==1.2.0
 pyOpenSSL==26.2.0
 pyparsing==3.3.2
 pyproject_hooks==1.2.0
+pystack==1.7.0
 pytest==9.1.1
 pytest-timeout==2.4.0
 pytest-xdist==3.8.0
diff --git a/sdks/python/container/ml/py311/gpu_image_requirements.txt b/sdks/python/container/ml/py311/gpu_image_requirements.txt
@@ -255,6 +255,7 @@ PyMySQL==1.2.0
 pyOpenSSL==26.2.0
 pyparsing==3.3.2
 pyproject_hooks==1.2.0
+pystack==1.7.0
 pytest==9.1.1
 pytest-timeout==2.4.0
 pytest-xdist==3.8.0
diff --git a/sdks/python/container/ml/py312/base_image_requirements.txt b/sdks/python/container/ml/py312/base_image_requirements.txt
@@ -178,6 +178,7 @@ PyMySQL==1.2.0
 pyOpenSSL==26.2.0
 pyparsing==3.3.2
 pyproject_hooks==1.2.0
+pystack==1.7.0
 pytest==9.1.1
 pytest-timeout==2.4.0
 pytest-xdist==3.8.0
diff --git a/sdks/python/container/ml/py312/gpu_image_requirements.txt b/sdks/python/container/ml/py312/gpu_image_requirements.txt
@@ -253,6 +253,7 @@ PyMySQL==1.2.0
 pyOpenSSL==26.2.0
 pyparsing==3.3.2
 pyproject_hooks==1.2.0
+pystack==1.7.0
 pytest==9.1.1
 pytest-timeout==2.4.0
 pytest-xdist==3.8.0
diff --git a/sdks/python/container/ml/py313/base_image_requirements.txt b/sdks/python/container/ml/py313/base_image_requirements.txt
@@ -177,6 +177,7 @@ PyMySQL==1.2.0
 pyOpenSSL==26.2.0
 pyparsing==3.3.2
 pyproject_hooks==1.2.0
+pystack==1.7.0
 pytest==9.1.1
 pytest-timeout==2.4.0
 pytest-xdist==3.8.0
diff --git a/sdks/python/container/profiler.go b/sdks/python/container/profiler.go
@@ -148,6 +148,14 @@ func startProfilerBackgroundTasks(ctx context.Context, logger *tools.Logger) {
 	if pcfg.Agent == "memray" {
 		go postProcessProfilesLoop(ctx, logger, pcfg.TempLocation, pcfg.PostprocessIntervalSec)
 	}
+
+	if pcfg.Agent == "pystack_coredump" {
+		go monitorCoredumpsLoop(ctx, logger, pcfg)
+	}
+
+	if pcfg.Agent == "pystack" {
+		go monitorActiveProcessesLoop(ctx, logger, pcfg)
+	}
 }
 
 // maybeWithProfiler builds the execution arguments and environment variables if profiling is enabled and active.
@@ -192,6 +200,9 @@ func maybeWithProfiler(
 		}
 		env["HEAPPROFILE"] = tcmallocHeapPath
 		args = currentArgs
+	} else if pcfg.Agent == "pystack_coredump" || pcfg.Agent == "pystack" {
+		// No wrapping is needed for pystack / pystack_coredump.
+		args = currentArgs
 	} else {
 		prog = pcfg.Agent
 		args = append(append([]string{}, pcfg.ExtraArgs...), currentProg)
@@ -223,6 +234,14 @@ func stopProfiling(ctx context.Context) error {
 	return err
 }
 
+// isProfilerDisengaged checks if the stop sentinel file exists.
+func isProfilerDisengaged(pcfg *ProfilerConfig) bool {
+	if _, err := os.Stat(pcfg.StopSentinelPath); err == nil {
+		return true
+	}
+	return false
+}
+
 // syncProfilesToGCS uploads newly created local memory profiles to the designated GCS target path using gcloud storage.
 func syncProfilesToGCS(ctx context.Context, logger *tools.Logger, localDir, gcsDest string) {
 	entries, err := os.ReadDir(localDir)
@@ -334,3 +353,137 @@ func needsProcessing(binInfo os.FileInfo, path string) bool {
 	// Don't regenerate when there were no updates to the profile.
 	return binInfo.ModTime().After(info.ModTime())
 }
+
+
+func monitorCoredumpsLoop(ctx context.Context, logger *tools.Logger, pcfg *ProfilerConfig) {
+	// We require the core file pattern to be configured as "/tmp/core.%e.%p" via pipeline
+	// options experiments (which is automatically set by the Python SDK). This ensures
+	// that core dumps are written in /tmp/ with a prefix matching "core.".
+	coreDir := "/tmp"
+	interval := 5 * time.Second
+	if pcfg.PostprocessIntervalSec > 0 {
+		interval = time.Duration(pcfg.PostprocessIntervalSec) * time.Second
+	}
+	logger.Printf(ctx, "Monitoring directory %s for core dumps matching prefix core. every %v", coreDir, interval)
+
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			processNewCoredumps(ctx, logger, pcfg, coreDir)
+			if isProfilerDisengaged(pcfg) {
+				return
+			}
+		}
+	}
+}
+
+func processNewCoredumps(ctx context.Context, logger *tools.Logger, pcfg *ProfilerConfig, coreDir string) {
+	files, err := os.ReadDir(coreDir)
+	if err != nil {
+		return
+	}
+
+	prefix := "core."
+
+	for _, file := range files {
+		if file.IsDir() {
+			continue
+		}
+		name := file.Name()
+		if !strings.HasPrefix(name, prefix) {
+			continue
+		}
+
+		corePath := filepath.Join(coreDir, name)
+		info, err := os.Stat(corePath)
+		if err != nil {
+			continue
+		}
+
+		fileAge := time.Since(info.ModTime())
+
+		// Skip if the file was modified very recently (might still be writing)
+		if fileAge < 2*time.Second {
+			continue
+		}
+
+		logger.Printf(ctx, "Found core dump file: %s (%d bytes)", name, info.Size())
+
+		// Find python executable. Since the worker might be running in a venv,
+		// we look for "python" in the PATH.
+		pythonProg := "python"
+		if path, err := exec.LookPath("python"); err == nil {
+			pythonProg = path
+		}
+
+		args := []string{"core"}
+		args = append(args, pcfg.ExtraArgs...)
+		args = append(args, corePath, pythonProg)
+
+		logger.Printf(ctx, "Running pystack %s", strings.Join(args, " "))
+		cmd := exec.CommandContext(ctx, "pystack", args...)
+		output, err := cmd.CombinedOutput()
+		if err != nil {
+			logger.Warnf(ctx, "pystack failed on %s: %v. Output:\n%s", name, err, string(output))
+			// Give up and delete if the file has been around for over 60 seconds
+			if fileAge > 60*time.Second {
+				logger.Printf(ctx, "Giving up on %s after 60s and deleting.", corePath)
+				if err := os.Remove(corePath); err != nil {
+					logger.Warnf(ctx, "Failed to delete core dump %s: %v", corePath, err)
+				}
+			}
+		} else {
+			logger.Errorf(ctx, "Pystack coredump analysis for %s:\n%s", name, string(output))
+			if err := os.Remove(corePath); err != nil {
+				logger.Warnf(ctx, "Failed to delete core dump %s: %v", corePath, err)
+			}
+		}
+	}
+}
+
+func monitorActiveProcessesLoop(ctx context.Context, logger *tools.Logger, pcfg *ProfilerConfig) {
+	interval := 60 * time.Second
+	if pcfg.PostprocessIntervalSec > 0 {
+		interval = time.Duration(pcfg.PostprocessIntervalSec) * time.Second
+	}
+
+	logger.Printf(ctx, "Starting pystack remote monitoring loop every %v", interval)
+
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			if isProfilerDisengaged(pcfg) {
+				return
+			}
+			pids := getActivePids()
+			if len(pids) == 0 {
+				continue
+			}
+
+			for _, pid := range pids {
+				args := []string{"remote"}
+				args = append(args, pcfg.ExtraArgs...)
+				args = append(args, fmt.Sprintf("%d", pid))
+
+				logger.Printf(ctx, "Running pystack remote on PID %d: pystack %s", pid, strings.Join(args, " "))
+				cmd := exec.CommandContext(ctx, "pystack", args...)
+				output, err := cmd.CombinedOutput()
+				if err != nil {
+					logger.Warnf(ctx, "pystack remote failed on PID %d: %v. Output:\n%s", pid, err, string(output))
+				} else {
+					logger.Printf(ctx, "Pystack thread dump for PID %d:\n%s", pid, string(output))
+				}
+			}
+		}
+	}
+}
diff --git a/sdks/python/container/py310/base_image_requirements.txt b/sdks/python/container/py310/base_image_requirements.txt
@@ -163,6 +163,7 @@ PyMySQL==1.2.0
 pyOpenSSL==26.2.0
 pyparsing==3.3.2
 pyproject_hooks==1.2.0
+pystack==1.7.0
 pytest==9.1.1
 pytest-timeout==2.4.0
 pytest-xdist==3.8.0
diff --git a/sdks/python/container/py311/base_image_requirements.txt b/sdks/python/container/py311/base_image_requirements.txt
@@ -162,6 +162,7 @@ PyMySQL==1.2.0
 pyOpenSSL==26.2.0
 pyparsing==3.3.2
 pyproject_hooks==1.2.0
+pystack==1.7.0
 pytest==9.1.1
 pytest-timeout==2.4.0
 pytest-xdist==3.8.0
diff --git a/sdks/python/container/py312/base_image_requirements.txt b/sdks/python/container/py312/base_image_requirements.txt
@@ -160,6 +160,7 @@ PyMySQL==1.2.0
 pyOpenSSL==26.2.0
 pyparsing==3.3.2
 pyproject_hooks==1.2.0
+pystack==1.7.0
 pytest==9.1.1
 pytest-timeout==2.4.0
 pytest-xdist==3.8.0
diff --git a/sdks/python/container/py313/base_image_requirements.txt b/sdks/python/container/py313/base_image_requirements.txt
@@ -159,6 +159,7 @@ PyMySQL==1.2.0
 pyOpenSSL==26.2.0
 pyparsing==3.3.2
 pyproject_hooks==1.2.0
+pystack==1.7.0
 pytest==9.1.1
 pytest-timeout==2.4.0
 pytest-xdist==3.8.0
diff --git a/sdks/python/container/py314/base_image_requirements.txt b/sdks/python/container/py314/base_image_requirements.txt
@@ -158,6 +158,7 @@ PyMySQL==1.2.0
 pyOpenSSL==26.2.0
 pyparsing==3.3.2
 pyproject_hooks==1.2.0
+pystack==1.7.0
 pytest==9.1.1
 pytest-timeout==2.4.0
 pytest-xdist==3.8.0
