From f7bc3ce0de34bdcaddf3bbe1eb18c3596d1c0bff Mon Sep 17 00:00:00 2001 From: Wenyuan Chi Date: Tue, 19 May 2026 13:29:42 -0700 Subject: [PATCH] Fix NCU profiling under Buck PAR via KERNEL_PROFILER_PYTHON (#136) Summary: NCU profiling subprocesses failed inside a Buck PAR with: ImportError: .../platform010/lib/python3.12/lib-dynload/ _posixsubprocess.cpython-312-x86_64-linux-gnu.so: undefined symbol: _Py_NoneStruct ModuleNotFoundError: No module named 'torch' Root cause: inside a PAR, sys.executable points at the static-linked native-main binary. Re-exec'ing it from a subprocess skips the env setup that _bootstrap.sh performs (LD_LIBRARY_PATH for CUDA, LD_PRELOAD for the allocator, PYTHONPATH/FB_PAR_* for the embedded import system), so the spawned Python falls back to system platform010 stdlib whose lib-dynload .so files are ABI-incompatible with the static libpython. Fix is split along the OSS/fb boundary: * OSS ncu_profiler.profile_triton_kernel honors a generic env-var override KERNEL_PROFILER_PYTHON, falling back to sys.executable. No Meta/buck/PAR knowledge in OSS code. * fb utils/fb/internal_env.setup_internal_environment detects the PAR via "#native-main#" in sys.executable.name and sets the env var to /_bootstrap.sh, which rebuilds the full env from \$0 before exec'ing the native main. Respects an explicit user override. Differential Revision: D105739069 --- kernel_perf_agent/kernel_opt/profiler/ncu_profiler.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/kernel_perf_agent/kernel_opt/profiler/ncu_profiler.py b/kernel_perf_agent/kernel_opt/profiler/ncu_profiler.py index 69a3a69f..ef9ae9f6 100644 --- a/kernel_perf_agent/kernel_opt/profiler/ncu_profiler.py +++ b/kernel_perf_agent/kernel_opt/profiler/ncu_profiler.py @@ -123,7 +123,12 @@ def profile_triton_kernel( # Resolve paths if python_executable is None: - python_executable = sys.executable + # Generic override hook: a wrapping environment (e.g. a packaged + # launcher that re-establishes its own runtime env vars) can set + # KERNEL_PROFILER_PYTHON to a Python entrypoint more appropriate for + # spawning torch-importable subprocesses than `sys.executable`. + # Falls back to `sys.executable` when unset. + python_executable = os.environ.get("KERNEL_PROFILER_PYTHON") or sys.executable if ncu_bin is None: ncu_bin = shutil.which("ncu") or "/usr/local/cuda/bin/ncu"