From 23cc7f70bac4d468366be92e4c5bf8db69ab9036 Mon Sep 17 00:00:00 2001
From: koubaa <koubaa@github.com>
Date: Wed, 11 Dec 2024 10:59:58 -0600
Subject: [PATCH] Fix race condition with python logging

---
 python/src/main.cpp                   | 24 ++++++++++++++++++++++++
 src/include/kompute/Core.hpp          |  7 ++++++-
 src/include/kompute/logger/Logger.hpp | 20 +++++++++++---------
 3 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/python/src/main.cpp b/python/src/main.cpp
index 3bddce88..5a3a4441 100644
--- a/python/src/main.cpp
+++ b/python/src/main.cpp
@@ -14,6 +14,30 @@ namespace py = pybind11;
 // used in Core.hpp
 py::object kp_trace, kp_debug, kp_info, kp_warning, kp_error;
 
+static void kp_log(int level, const std::string& msg) {
+  py::gil_scoped_acquire gil;
+    switch (level) {
+      case 0:
+        kp_trace(msg); break;
+      case 1:
+        kp_debug(msg); break;
+      case 2:
+        kp_info(msg); break;
+      case 3:
+        kp_warning(msg); break;
+      case 4:
+        kp_error(msg); break;
+      default:
+        break;
+    }
+}
+
+void py_log_trace(const std::string& msg) {kp_log(0, msg);}
+void py_log_debug(const std::string& msg) {kp_log(1, msg);}
+void py_log_info(const std::string& msg) {kp_log(2, msg);}
+void py_log_warning(const std::string& msg) {kp_log(3, msg);}
+void py_log_error(const std::string& msg) {kp_log(4, msg);}
+
 std::unique_ptr<kp::OpAlgoDispatch>
 opAlgoDispatchPyInit(std::shared_ptr<kp::Algorithm>& algorithm,
                      const py::array& push_consts)
diff --git a/src/include/kompute/Core.hpp b/src/include/kompute/Core.hpp
index 2384e47b..d347a827 100644
--- a/src/include/kompute/Core.hpp
+++ b/src/include/kompute/Core.hpp
@@ -26,5 +26,10 @@ typedef std::vector<float> Constants;
 #include <pybind11/pybind11.h>
 namespace py = pybind11;
 // from python/src/main.cpp
-extern py::object kp_trace, kp_debug, kp_info, kp_warning, kp_error;
+
+extern void py_log_trace(const std::string& msg);
+extern void py_log_debug(const std::string& msg);
+extern void py_log_info(const std::string& msg);
+extern void py_log_warning(const std::string& msg);
+extern void py_log_error(const std::string& msg);
 #endif
diff --git a/src/include/kompute/logger/Logger.hpp b/src/include/kompute/logger/Logger.hpp
index 6e06918a..bfdad41a 100644
--- a/src/include/kompute/logger/Logger.hpp
+++ b/src/include/kompute/logger/Logger.hpp
@@ -25,10 +25,12 @@ static const char* KOMPUTE_LOG_TAG = "KomputeLog";
 #else
 #if KOMPUTE_BUILD_PYTHON
 #include <fmt/core.h>
-#include <pybind11/pybind11.h>
-namespace py = pybind11;
 // from python/src/main.cpp
-extern py::object kp_trace, kp_debug, kp_info, kp_warning, kp_error;
+extern void py_log_trace(const std::string& msg);
+extern void py_log_debug(const std::string& msg);
+extern void py_log_info(const std::string& msg);
+extern void py_log_warning(const std::string& msg);
+extern void py_log_error(const std::string& msg);
 #else
 #include <fmt/core.h>
 #endif // KOMPUTE_BUILD_PYTHON
@@ -57,7 +59,7 @@ setupLogger();
       ANDROID_LOG_VERBOSE, KOMPUTE_LOG_TAG, fmt::format(__VA_ARGS__).c_str()))
 #else
 #if KOMPUTE_BUILD_PYTHON
-#define KP_LOG_TRACE(...) kp_trace(fmt::format(__VA_ARGS__))
+#define KP_LOG_TRACE(...) py_log_trace(fmt::format(__VA_ARGS__))
 #else
 #define KP_LOG_TRACE(...)                                                      \
     fmt::print("[{} {}] [trace] [{}:{}] {}\n",                                 \
@@ -81,7 +83,7 @@ setupLogger();
       ANDROID_LOG_DEBUG, KOMPUTE_LOG_TAG, fmt::format(__VA_ARGS__).c_str()))
 #else
 #if KOMPUTE_BUILD_PYTHON
-#define KP_LOG_DEBUG(...) kp_debug(fmt::format(__VA_ARGS__))
+#define KP_LOG_DEBUG(...) py_log_debug(fmt::format(__VA_ARGS__))
 #else
 #ifdef __FILE_NAME__ // gcc 12 provides only file name without path
 #define KP_LOG_DEBUG(...)                                                      \
@@ -98,7 +100,7 @@ setupLogger();
                __TIME__,                                                       \
                __FILE__,                                                       \
                __LINE__,                                                       \
-               fmt::format(__VA_ARGS__))
+               fmt::format(__VA_ARGS__)); std::cout << std::flush
 #endif // __FILE__NAME__
 #endif // KOMPUTE_BUILD_PYTHON
 #endif // VK_USE_PLATFORM_ANDROID_KHR
@@ -115,7 +117,7 @@ setupLogger();
       ANDROID_LOG_INFO, KOMPUTE_LOG_TAG, fmt::format(__VA_ARGS__).c_str()))
 #else
 #if KOMPUTE_BUILD_PYTHON
-#define KP_LOG_INFO(...) kp_info(fmt::format(__VA_ARGS__))
+#define KP_LOG_INFO(...) py_log_info(fmt::format(__VA_ARGS__))
 #else
 #define KP_LOG_INFO(...)                                                       \
     fmt::print("[{} {}] [info] [{}:{}] {}\n",                                  \
@@ -139,7 +141,7 @@ setupLogger();
       ANDROID_LOG_WARN, KOMPUTE_LOG_TAG, fmt::format(__VA_ARGS__).c_str()))
 #else
 #if KOMPUTE_BUILD_PYTHON
-#define KP_LOG_WARN(...) kp_warning(fmt::format(__VA_ARGS__))
+#define KP_LOG_WARN(...) py_log_warning(fmt::format(__VA_ARGS__))
 #else
 #define KP_LOG_WARN(...)                                                       \
     fmt::print("[{} {}] [warn] [{}:{}] {}\n",                                  \
@@ -163,7 +165,7 @@ setupLogger();
       ANDROID_LOG_ERROR, KOMPUTE_LOG_TAG, fmt::format(__VA_ARGS__).c_str()))
 #else
 #if KOMPUTE_BUILD_PYTHON
-#define KP_LOG_ERROR(...) kp_error(fmt::format(__VA_ARGS__))
+#define KP_LOG_ERROR(...) py_log_error(fmt::format(__VA_ARGS__))
 #else
 #define KP_LOG_ERROR(...)                                                      \
     fmt::print("[{} {}] [error] [{}:{}] {}\n",                                 \