[feat] Support core dump analysis with pystack.

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
 
 

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'])

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

sdks/python/container/boot.go

@@ -60,6 +60,9 @@ var (
 	provisionEndpoint = flag.String("provision_endpoint", "", "Provision endpoint (required).")
 	controlEndpoint   = flag.String("control_endpoint", "", "Control endpoint (required).")
 	semiPersistDir    = flag.String("semi_persist_dir", "/tmp", "Local semi-persistent directory (optional).")
+
+	workerMu     sync.Mutex
+	shuttingDown bool
 )
 
 const (
@@ -307,19 +310,12 @@ func launchSDKProcess() error {
 
 	workerIds := append([]string{*id}, info.GetSiblingWorkerIds()...)
 
-	// Keep track of child PIDs for clean shutdown without zombies
-	childPids := struct {
-		v        []int
-		canceled bool
-		mu       sync.Mutex
-	}{v: make([]int, 0, len(workerIds))}
-
 	// Forward trapped signals to child process groups in order to terminate them gracefully and avoid zombies
 	go func() {
 		logger.Printf(ctx, "Received signal: %v", <-signalChannel)
-		childPids.mu.Lock()
-		childPids.canceled = true
-		for _, pid := range childPids.v {
+		workerMu.Lock()
+		shuttingDown = true
+		for _, pid := range activePids {
 			go func(pid int) {
 				// This goroutine will be canceled if the main process exits before the 5 seconds
 				// have elapsed, i.e., as soon as all subprocesses have returned from Wait().
@@ -330,7 +326,7 @@ func launchSDKProcess() error {
 			}(pid)
 			syscall.Kill(-pid, syscall.SIGTERM)
 		}
-		childPids.mu.Unlock()
+		workerMu.Unlock()
 	}()
 
 	var wg sync.WaitGroup
@@ -342,9 +338,9 @@ func launchSDKProcess() error {
 			bufLogger := tools.NewBufferedLogger(logger)
 			errorCount := 0
 			for {
-				childPids.mu.Lock()
-				if childPids.canceled {
-					childPids.mu.Unlock()
+				workerMu.Lock()
+				if shuttingDown {
+					workerMu.Unlock()
 					return
 				}
 
@@ -369,8 +365,9 @@ 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...)
-				childPids.v = append(childPids.v, cmd.Process.Pid)
-				childPids.mu.Unlock()
+				logger.Printf(ctx, "Started worker %s with PID %d", workerId, cmd.Process.Pid)
+				activePids = append(activePids, cmd.Process.Pid)
+				workerMu.Unlock()
 
 				var timer *time.Timer
 				var profilingTimedOut atomic.Bool
@@ -379,8 +376,8 @@ func launchSDKProcess() error {
 				if profilingActive && pcfg.StopAfterSec > 0 {
 					duration := time.Duration(pcfg.StopAfterSec) * time.Second
 					timer = time.AfterFunc(duration, func() {
-						childPids.mu.Lock()
-						defer childPids.mu.Unlock()
+						workerMu.Lock()
+						defer workerMu.Unlock()
 						if cmd.Process != nil {
 							logger.Printf(ctx, "Profiling timeout of %d seconds reached. Sending SIGINT to worker %s",
 								pcfg.StopAfterSec, workerId)
@@ -391,6 +388,7 @@ func launchSDKProcess() error {
 				}
 
 				err := cmd.Wait()
+				unregisterPid(cmd.Process.Pid)
 				if timer != nil {
 					timer.Stop()
 				}
@@ -603,4 +601,22 @@ func logSubmissionEnvDependencies(ctx context.Context, bufLogger *tools.Buffered
 	return nil
 }
 
+var (
+	activePids []int
+)
 
+func unregisterPid(pid int) {
+	workerMu.Lock()
+	defer workerMu.Unlock()
+	activePids = slices.DeleteFunc(activePids, func(p int) bool {
+		return p == pid
+	})
+}
+
+func getActivePids() []int {
+	workerMu.Lock()
+	defer workerMu.Unlock()
+	pids := make([]int, len(activePids))
+	copy(pids, activePids)
+	return pids
+}

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

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

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

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

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

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

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

sdks/python/container/profiler.go

@@ -146,8 +146,13 @@ func startProfilerBackgroundTasks(ctx context.Context, logger *tools.Logger) {
 	}
 
 	if pcfg.Agent == "memray" {
-		go postProcessProfilesLoop(ctx, logger, pcfg.TempLocation, pcfg.PostprocessIntervalSec)
+		go postProcessProfilesLoop(ctx, logger, pcfg)
 	}
+
+	if pcfg.Agent == "pystack_coredump" {
+		go monitorCoredumpsLoop(ctx, logger, pcfg)
+	}
+
 }
 
 // maybeWithProfiler builds the execution arguments and environment variables if profiling is enabled and active.
@@ -192,6 +197,9 @@ func maybeWithProfiler(
 		}
 		env["HEAPPROFILE"] = tcmallocHeapPath
 		args = currentArgs
+	} else if pcfg.Agent == "pystack_coredump" {
+		// No wrapping is needed for pystack_coredump.
+		args = currentArgs
 	} else {
 		prog = pcfg.Agent
 		args = append(append([]string{}, pcfg.ExtraArgs...), currentProg)
@@ -223,6 +231,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)
@@ -241,18 +257,22 @@ func syncProfilesToGCS(ctx context.Context, logger *tools.Logger, localDir, gcsD
 }
 
 // postProcessProfilesLoop runs a background loop that periodically triggers profile post-processing if enabled.
-func postProcessProfilesLoop(ctx context.Context, logger *tools.Logger, profilesDir string, intervalSec int) {
-	if intervalSec <= 0 {
+func postProcessProfilesLoop(ctx context.Context, logger *tools.Logger, pcfg *ProfilerConfig) {
+	if pcfg.PostprocessIntervalSec <= 0 {
 		return
 	}
 
 	for {
-		runPostProcessingSweep(ctx, logger, profilesDir, intervalSec)
+		runPostProcessingSweep(ctx, logger, pcfg.TempLocation, pcfg.PostprocessIntervalSec)
+
+		if isProfilerDisengaged(pcfg) {
+			return
+		}
 
 		select {
 		case <-ctx.Done():
 			return
-		case <-time.After(time.Duration(intervalSec) * time.Second):
+		case <-time.After(time.Duration(pcfg.PostprocessIntervalSec) * time.Second):
 			// Block until the sleep completes before starting the next sweep
 		}
 	}
@@ -334,3 +354,95 @@ 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)
+			}
+		}
+	}
+}