fix: make sure pyo3 runtime is initialized before all others (#1063)

Author: georgeh0

python/cocoindex/__init__.py

@@ -2,6 +2,7 @@
 Cocoindex is a framework for building and running indexing pipelines.
 """
 
+from . import _engine  # type: ignore
 from . import functions, sources, targets, cli, utils
 
 from . import targets as storages  # Deprecated: Use targets instead
@@ -42,6 +43,8 @@
     Json,
 )
 
+_engine.init_pyo3_runtime()
+
 __all__ = [
     # Submodules
     "_engine",

src/lib_context.rs

@@ -162,6 +162,13 @@ impl FlowContext {
 static TOKIO_RUNTIME: LazyLock<Runtime> = LazyLock::new(|| Runtime::new().unwrap());
 static AUTH_REGISTRY: LazyLock<Arc<AuthRegistry>> = LazyLock::new(|| Arc::new(AuthRegistry::new()));
 
+pub fn get_runtime() -> &'static Runtime {
+    &TOKIO_RUNTIME
+}
+pub fn get_auth_registry() -> &'static Arc<AuthRegistry> {
+    &AUTH_REGISTRY
+}
+
 type PoolKey = (String, Option<String>);
 type PoolValue = Arc<tokio::sync::OnceCell<PgPool>>;
 
@@ -271,21 +278,10 @@ impl LibContext {
     }
 }
 
-pub fn get_runtime() -> &'static Runtime {
-    &TOKIO_RUNTIME
-}
-
-pub fn get_auth_registry() -> &'static Arc<AuthRegistry> {
-    &AUTH_REGISTRY
-}
-
 static LIB_INIT: OnceLock<()> = OnceLock::new();
 pub async fn create_lib_context(settings: settings::Settings) -> Result<LibContext> {
     LIB_INIT.get_or_init(|| {
         let _ = env_logger::try_init();
-
-        pyo3_async_runtimes::tokio::init_with_runtime(get_runtime()).unwrap();
-
         let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
     });
 

src/py/mod.rs

@@ -115,6 +115,11 @@ fn set_settings_fn(get_settings_fn: Py<PyAny>) -> PyResult<()> {
     Ok(())
 }
 
+#[pyfunction]
+fn init_pyo3_runtime() {
+    pyo3_async_runtimes::tokio::init_with_runtime(get_runtime()).unwrap();
+}
+
 #[pyfunction]
 fn init(py: Python<'_>, settings: Pythonized<Option<Settings>>) -> PyResult<()> {
     py.allow_threads(|| -> anyhow::Result<()> {
@@ -666,6 +671,7 @@ fn seder_roundtrip<'py>(
 #[pymodule]
 #[pyo3(name = "_engine")]
 fn cocoindex_engine(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(init_pyo3_runtime, m)?)?;
     m.add_function(wrap_pyfunction!(init, m)?)?;
     m.add_function(wrap_pyfunction!(set_settings_fn, m)?)?;
     m.add_function(wrap_pyfunction!(start_server, m)?)?;