[libshortfin] Move the logging setup

This patch moves logging functions from
nod-ai#181

libshortfin/python/lib_ext.cc

@@ -339,6 +339,13 @@ NB_MODULE(lib, m) {
       });
 
   py::class_<iree::vm_opaque_ref>(m, "_OpaqueVmRef");
+
+  // Logging entrypoints.
+  m.def("log_debug", [](std::string_view sv) { logging::debug("{}", sv); });
+  m.def("log_info", [](std::string_view sv) { logging::info("{}", sv); });
+  m.def("log_warn", [](std::string_view sv) { logging::warn("{}", sv); });
+  m.def("log_error", [](std::string_view sv) { logging::error("{}", sv); });
+
   auto local_m = m.def_submodule("local");
   BindLocal(local_m);
   BindHostSystem(local_m);

libshortfin/python/shortfin/interop/support/logging_setup.py

@@ -0,0 +1,52 @@
+# Copyright 2024 Advanced Micro Devices, Inc.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+import logging
+import sys
+
+from _shortfin import lib as _sfl
+
+_LOG_FUNCTIONS = {
+    logging.DEBUG: _sfl.log_debug,
+    logging.INFO: _sfl.log_info,
+    logging.WARNING: _sfl.log_warn,
+    logging.ERROR: _sfl.log_error,
+    logging.CRITICAL: _sfl.log_error,
+}
+
+logger = logging.getLogger("shortfin")
+logger.propagate = False
+
+
+class NativeHandler(logging.Handler):
+    def emit(self, record):
+        formatted = self.format(record)
+        f = _LOG_FUNCTIONS.get(record.levelno)
+        if f is not None:
+            f(formatted)
+
+
+class NativeFormatter(logging.Formatter):
+    def __init__(self):
+        super().__init__("[%(filename)s:%(lineno)d] %(message)s")
+
+
+native_handler = NativeHandler()
+native_handler.setFormatter(NativeFormatter())
+
+# TODO: Source from env vars.
+logger.setLevel(logging.DEBUG)
+logger.addHandler(native_handler)
+
+
+def configure_main_logger(module_suffix: str = "__main__") -> logging.Logger:
+    """Configures logging from a main entrypoint.
+    Returns a logger that can be used for the main module itself.
+    """
+    logging.root.addHandler(native_handler)
+    logging.root.setLevel(logging.DEBUG)  # TODO: source from env vars
+    main_module = sys.modules["__main__"]
+    return logging.getLogger(f"{main_module.__package__}.{module_suffix}")

libshortfin/src/shortfin/support/logging.cc

@@ -5,3 +5,14 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 #include "shortfin/support/logging.h"
+
+#include "spdlog/cfg/env.h"
+
+namespace shortfin::logging {
+
+void InitializeFromEnv() {
+  // TODO: Also support our own env vars.
+  spdlog::cfg::load_env_levels();
+}
+
+}  // namespace shortfin::logging