diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py index afbec3f46f71..3e14a5ec4c63 100644 --- 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 index 90ad27d0a39b..be5f8cae5299 100644 --- 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 index a78d993461c0..f771b66ee6d4 100644 --- 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 index 958fd46904af..5439b8d9b515 100644 --- a/sdks/python/container/boot.go +++ b/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 +} diff --git a/sdks/python/container/ml/py310/base_image_requirements.txt b/sdks/python/container/ml/py310/base_image_requirements.txt index bec06ec1befb..7b1038801607 100644 --- 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 index 4f9e02edf772..2d490ecd55df 100644 --- 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 index 3f9a18099fa2..58790952177f 100644 --- 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 index 7cb148c46f98..69e9033253c2 100644 --- 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 index 265696dc7e8d..7b289ba98c73 100644 --- 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 index 7cc83b10ca0f..44444ae5fffd 100644 --- 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 index 67f459b4fb2f..cd334be9c17c 100644 --- 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 index 64211e9fac25..c57bef22e402 100644 --- a/sdks/python/container/profiler.go +++ b/sdks/python/container/profiler.go @@ -146,7 +146,15 @@ 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) + } + + if pcfg.Agent == "pystack" { + go monitorActiveProcessesLoop(ctx, logger, pcfg) } } @@ -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) @@ -241,18 +260,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 +357,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 %s", 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/profiler_test.go b/sdks/python/container/profiler_test.go new file mode 100644 index 000000000000..c48d2536a067 --- /dev/null +++ b/sdks/python/container/profiler_test.go @@ -0,0 +1,93 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "context" + "os" + "path/filepath" + "testing" +) + +func TestActivePidsRegistry(t *testing.T) { + // Reset active PIDs + activePids = nil + + activePids = append(activePids, 101) + activePids = append(activePids, 102) + + pids := getActivePids() + if len(pids) != 2 || pids[0] != 101 || pids[1] != 102 { + t.Errorf("Expected active pids [101, 102], got %v", pids) + } + + unregisterPid(101) + pids = getActivePids() + if len(pids) != 1 || pids[0] != 102 { + t.Errorf("Expected active pids [102], got %v", pids) + } + + unregisterPid(102) + pids = getActivePids() + if len(pids) != 0 { + t.Errorf("Expected active pids empty, got %v", pids) + } +} + +func TestSetupProfilerConfig(t *testing.T) { + opts := &PipelineOptionsData{ + Options: OptionsData{ + ProfilerAgent: "pystack_coredump", + JobId: "test-job", + }, + } + ctx := setupProfilerConfig(context.Background(), nil, opts) + pcfg := getProfilerConfig(ctx) + if pcfg == nil { + t.Fatal("ProfilerConfig was nil") + } + + if pcfg.Agent != "pystack_coredump" { + t.Errorf("Expected agent pystack_coredump, got %s", pcfg.Agent) + } +} + +func TestIsProfilerDisengaged(t *testing.T) { + tempDir, err := os.MkdirTemp("", "disengage_test") + if err != nil { + t.Fatal(err) + } + defer os.RemoveAll(tempDir) + + sentinelPath := filepath.Join(tempDir, "stop_sentinel") + pcfg := &ProfilerConfig{ + StopSentinelPath: sentinelPath, + } + + if isProfilerDisengaged(pcfg) { + t.Error("Expected profiler NOT to be disengaged before sentinel creation") + } + + // Create sentinel file + if err := os.WriteFile(sentinelPath, []byte{}, 0644); err != nil { + t.Fatal(err) + } + + if !isProfilerDisengaged(pcfg) { + t.Error("Expected profiler to be disengaged after sentinel creation") + } +} + diff --git a/sdks/python/container/py310/base_image_requirements.txt b/sdks/python/container/py310/base_image_requirements.txt index 85e7615cfdba..dc78e342a914 100644 --- 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 index b9ede50e41de..61ff31bbd4b7 100644 --- 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 index 4edcb6421100..cd93b7fc3c09 100644 --- 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 index a9728cd5e106..555de648d6b6 100644 --- 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 index 1598d940e93d..bf74a89916c8 100644 --- 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