From cfc35c855bd61050ea7fd37fee646c243772ad31 Mon Sep 17 00:00:00 2001
From: Gregory Comer <gjcomer@meta.com>
Date: Thu, 12 Jun 2025 14:10:19 -0700
Subject: [PATCH] Use PAL table override for pybinding log routing

Summary: This PR updates the pybinding log routing to use the new PAL table overrides to reliably route logs. Previously, the PAL is not overridden in optimized bento kernels, preventing users from seeing logs.

Differential Revision: D74122422
---
 extension/pybindings/pybindings.cpp | 41 ++++++++++++++++++-----------
 1 file changed, 26 insertions(+), 15 deletions(-)

diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp
index 9cd5adae679..db0871657f6 100644
--- a/extension/pybindings/pybindings.cpp
+++ b/extension/pybindings/pybindings.cpp
@@ -68,21 +68,6 @@
     }                                                             \
   })
 
-// Our logs work by writing to stderr. By default this is done through fprintf
-// (as defined in posix.cpp) which then does not show up in python environments.
-// Here we override the pal to use std::cerr which can be properly redirected by
-// scoped_estream_redirect.
-void et_pal_emit_log_message(
-    et_timestamp_t timestamp,
-    et_pal_log_level_t level,
-    const char* filename,
-    ET_UNUSED const char* function,
-    size_t line,
-    const char* message,
-    ET_UNUSED size_t length) {
-  std::cerr << "[" << filename << ":" << line << "] " << message << std::endl;
-}
-
 namespace py = pybind11;
 using executorch::BUNDLED_PROGRAM_NAMESPACE::verify_method_outputs;
 using ::executorch::ET_RUNTIME_NAMESPACE::BackendInterface;
@@ -1168,6 +1153,32 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) {
       .def("__repr__", &PyMethodMeta::repr, call_guard);
 }
 
+namespace {
+
+// Our logs work by writing to stderr. By default this is done through fprintf
+// (as defined in posix.cpp) which then does not show up in python environments.
+// Here we override the pal to use std::cerr which can be properly redirected by
+// scoped_estream_redirect.
+void emit_log_message(
+    et_timestamp_t timestamp,
+    et_pal_log_level_t level,
+    const char* filename,
+    ET_UNUSED const char* function,
+    size_t line,
+    const char* message,
+    ET_UNUSED size_t length) {
+  std::cerr << "[" << filename << ":" << line << "] " << message << std::endl;
+}
+
+runtime::PalImpl build_pal() {
+  return runtime::PalImpl::create(emit_log_message, __FILE__);
+}
+
+// Update PAL to redirect logs.
+ET_UNUSED bool registration_result = runtime::register_pal(build_pal());
+
+} // namespace
+
 } // namespace pybindings
 } // namespace extension
 } // namespace executorch