chalk-remote-call-python 1.10.7__tar.gz → 1.10.9__tar.gz

{chalk_remote_call_python-1.10.7/chalk_remote_call_python.egg-info → chalk_remote_call_python-1.10.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalk-remote-call-python
-Version: 1.10.7
+Version: 1.10.9
 Summary: Chalk remote call Python runtime interface client
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai

{chalk_remote_call_python-1.10.7 → chalk_remote_call_python-1.10.9}/chalk-remote-call-rs/Cargo.lock

@@ -548,6 +548,7 @@ dependencies = [
  "futures-util",
  "prost 0.13.5",
  "pyo3",
+ "pyo3-build-config",
  "redis",
  "serde_json",
  "tokio",

{chalk_remote_call_python-1.10.7 → chalk_remote_call_python-1.10.9}/chalk-remote-call-rs/chalk-remote-call-server/Cargo.toml

@@ -11,6 +11,9 @@ crate-type = ["cdylib"]
 default = ["extension-module"]
 extension-module = ["pyo3/extension-module"]
 
+[build-dependencies]
+pyo3-build-config = { version = "0.29", features = ["resolve-config"] }
+
 [dependencies]
 chalk-remote-call-proto = { path = "../chalk-remote-call-proto" }
 redis = { version = "1.0.2", features = ["aio", "tokio-comp", "streams", "connection-manager"] }

chalk_remote_call_python-1.10.9/chalk-remote-call-rs/chalk-remote-call-server/build.rs

@@ -0,0 +1,15 @@
+fn main() {
+    println!("cargo:rerun-if-env-changed=CARGO_FEATURE_EXTENSION_MODULE");
+    if std::env::var_os("CARGO_FEATURE_EXTENSION_MODULE").is_some() {
+        return;
+    }
+
+    let target_os = std::env::var("CARGO_CFG_TARGET_OS").unwrap_or_default();
+    if matches!(target_os.as_str(), "linux" | "macos") {
+        // Embedding test binaries need the same Python library that PyO3 links,
+        // even when a managed interpreter installs it outside system paths.
+        if let Some(lib_dir) = pyo3_build_config::get().lib_dir() {
+            println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}");
+        }
+    }
+}

{chalk_remote_call_python-1.10.7 → chalk_remote_call_python-1.10.9}/chalk-remote-call-rs/chalk-remote-call-server/src/python_bridge.rs

@@ -1,3 +1,4 @@
+use std::cell::OnceCell;
 use std::collections::HashMap;
 use std::sync::{Arc, Mutex};
 
@@ -8,6 +9,46 @@ use pyo3::types::{PyBytes, PyDict, PyList};
 use tokio::sync::{mpsc, Notify};
 use tonic::Status;
 
+thread_local! {
+    // This value belongs only to its OS thread, including its destructor.
+    static THREAD_ATTACHMENT: OnceCell<ThreadAttachment> = const { OnceCell::new() };
+}
+
+struct ThreadAttachment {
+    state: pyo3::ffi::PyGILState_STATE,
+}
+
+impl ThreadAttachment {
+    fn new(_py: Python<'_>) -> Self {
+        // SAFETY: the Python token guarantees an attached thread. This nested
+        // Ensure returns LOCKED and pins its state after the outer guard detaches.
+        // The outer PyO3 guard still releases the GIL between invocations.
+        let state = unsafe { pyo3::ffi::PyGILState_Ensure() };
+        Self { state }
+    }
+}
+
+impl Drop for ThreadAttachment {
+    fn drop(&mut self) {
+        // Tokio joins its blocking threads while start_server's caller is detached.
+        // Reattach before releasing our matching LOCKED token; PyO3 then releases
+        // its outer token, deleting the state without retaining the GIL.
+        let _ = Python::try_attach(|_| unsafe {
+            // SAFETY: creation and destruction run on this same OS thread under
+            // the GIL, and this releases exactly the nested Ensure above.
+            pyo3::ffi::PyGILState_Release(self.state);
+        });
+    }
+}
+
+fn retain_thread_state(py: Python<'_>) {
+    // pybind11 can cache this state in native TLS. Its lifetime must match the
+    // blocking worker rather than an individual Python invocation.
+    THREAD_ATTACHMENT.with(|attachment| {
+        attachment.get_or_init(|| ThreadAttachment::new(py));
+    });
+}
+
 /// Holds references to the Python handler and bridge function.
 pub struct PythonHandler {
     /// The user's handler callable.
@@ -45,7 +86,7 @@ impl PythonHandler {
     /// into a shared buffer). The only difference between the two paths is which
     /// emitter object `process_batches` is handed.
     async fn call_with_emitter<F>(
-        &self,
+        self: &Arc<Self>,
         ipc_bytes: Vec<u8>,
         function_name: String,
         metadata: HashMap<String, String>,
@@ -55,12 +96,14 @@ impl PythonHandler {
     where
         F: FnOnce(Python<'_>) -> PyResult<Py<PyAny>> + Send + 'static,
     {
-        let (handler, process_fn) =
-            Python::attach(|py| (self.handler.clone_ref(py), self.process_fn.clone_ref(py)));
-        let arg_names = self.arg_names.clone();
-
+        // GIL acquisition can block, so only clone the Rust owner on the runtime worker.
+        let owner = Arc::clone(self);
         tokio::task::spawn_blocking(move || {
             Python::attach(|py| -> Result<(), PythonError> {
+                retain_thread_state(py);
+                let handler = owner.handler.clone_ref(py);
+                let process_fn = &owner.process_fn;
+                let arg_names = &owner.arg_names;
                 let py_bytes = PyBytes::new(py, &ipc_bytes).into_any().unbind();
                 let py_ctx = build_context(py, &function_name, &metadata, &peer)?
                     .into_any()
@@ -81,7 +124,7 @@ impl PythonHandler {
     /// Non-batched streaming dispatch. Calls `process_batches(..., emit)` so
     /// generator results are sent to the gRPC stream as each yield is encoded.
     pub async fn call_streaming(
-        &self,
+        self: &Arc<Self>,
         ipc_bytes: Vec<u8>,
         function_name: String,
         metadata: HashMap<String, String>,
@@ -100,7 +143,7 @@ impl PythonHandler {
     /// value, so completion never depends on emitter teardown timing (which is
     /// what caused a Python 3.10 hang in the channel-drain version).
     pub async fn call_into_buffer(
-        &self,
+        self: &Arc<Self>,
         ipc_bytes: Vec<u8>,
         function_name: String,
         metadata: HashMap<String, String>,
@@ -117,15 +160,17 @@ impl PythonHandler {
     /// `process_batches_coalesced`. Returns one IPC-encoded response per
     /// caller, in input order.
     pub async fn call_coalesced(
-        &self,
+        self: &Arc<Self>,
         callers: Vec<CallerInput>,
     ) -> Result<Vec<Vec<u8>>, PythonError> {
-        let (handler, process_fn) =
-            Python::attach(|py| (self.handler.clone_ref(py), self.process_fn.clone_ref(py)));
-        let arg_names = self.arg_names.clone();
-
+        // GIL acquisition can block, so only clone the Rust owner on the runtime worker.
+        let owner = Arc::clone(self);
         tokio::task::spawn_blocking(move || {
             Python::attach(|py| -> Result<Vec<Vec<u8>>, PythonError> {
+                retain_thread_state(py);
+                let handler = &owner.handler;
+                let process_fn = &owner.process_fn;
+                let arg_names = &owner.arg_names;
                 let py_inputs = PyList::empty(py);
                 let py_contexts = PyList::empty(py);
                 for c in &callers {
@@ -139,8 +184,8 @@ impl PythonHandler {
                 let py_arg_names = build_arg_names(py, arg_names.as_deref())?;
                 invoke_process_fn(
                     py,
-                    &process_fn,
-                    &handler,
+                    process_fn,
+                    handler,
                     py_inputs.into_any().unbind(),
                     py_arg_names,
                     py_contexts.into_any().unbind(),
@@ -389,3 +434,187 @@ fn format_python_traceback(py: Python<'_>, err: &PyErr) -> PyResult<String> {
     let lines: Vec<String> = formatted.extract()?;
     Ok(lines.concat().trim_end().to_string())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::future::{poll_fn, Future};
+    use std::task::Poll;
+    use std::time::Duration;
+
+    #[derive(Clone, Copy)]
+    enum Dispatch {
+        Streaming,
+        Buffered,
+        Coalesced,
+    }
+
+    // Polling a dispatch must yield even while another thread owns the GIL:
+    // otherwise Python contention blocks the worker that also serves health RPCs.
+    fn dispatch_yields_while_gil_is_held(dispatch: Dispatch) {
+        Python::initialize();
+        let (yielded, worker) = Python::attach(|py| {
+            let handler = Arc::new(PythonHandler::new(
+                py.None(),
+                py.eval(pyo3::ffi::c_str!("lambda *args: []"), None, None)
+                    .unwrap()
+                    .unbind(),
+                None,
+            ));
+            let (polled_tx, polled_rx) = std::sync::mpsc::channel();
+            let worker = std::thread::spawn(move || {
+                let runtime = tokio::runtime::Builder::new_current_thread()
+                    .build()
+                    .unwrap();
+                runtime.block_on(async move {
+                    let call = async {
+                        match dispatch {
+                            Dispatch::Streaming => {
+                                let (tx, _rx) = mpsc::channel(1);
+                                handler
+                                    .call_streaming(
+                                        Vec::new(),
+                                        String::new(),
+                                        HashMap::new(),
+                                        String::new(),
+                                        tx,
+                                    )
+                                    .await
+                            }
+                            Dispatch::Buffered => {
+                                handler
+                                    .call_into_buffer(
+                                        Vec::new(),
+                                        String::new(),
+                                        HashMap::new(),
+                                        String::new(),
+                                        ChunkBuffer::new().into_shared(),
+                                    )
+                                    .await
+                            }
+                            Dispatch::Coalesced => {
+                                handler.call_coalesced(Vec::new()).await.map(|_| ())
+                            }
+                        }
+                    };
+                    tokio::pin!(call);
+                    poll_fn(|cx| {
+                        assert!(call.as_mut().poll(cx).is_pending());
+                        let _ = polled_tx.send(());
+                        Poll::Ready(())
+                    })
+                    .await;
+                    call.await.unwrap();
+                });
+            });
+            // No Python bytecode runs here, so CPython cannot periodically release
+            // the GIL and accidentally let a blocking attach pass the test.
+            let yielded = polled_rx.recv_timeout(Duration::from_secs(5)).is_ok();
+            (yielded, worker)
+        });
+        // Release the GIL before joining so even a regressed implementation exits.
+        worker.join().unwrap();
+        assert!(
+            yielded,
+            "dispatch blocked its Tokio worker acquiring the GIL"
+        );
+    }
+
+    #[test]
+    fn streaming_dispatch_yields_under_gil_contention() {
+        dispatch_yields_while_gil_is_held(Dispatch::Streaming);
+    }
+
+    #[test]
+    fn buffered_dispatch_yields_under_gil_contention() {
+        dispatch_yields_while_gil_is_held(Dispatch::Buffered);
+    }
+
+    #[test]
+    fn coalesced_dispatch_yields_under_gil_contention() {
+        dispatch_yields_while_gil_is_held(Dispatch::Coalesced);
+    }
+    fn contains_thread_state(_py: Python<'_>, id: i64) -> bool {
+        // SAFETY: interpreter state iteration and ID access occur under the GIL.
+        unsafe {
+            let interpreter =
+                pyo3::ffi::PyThreadState_GetInterpreter(pyo3::ffi::PyThreadState_Get());
+            let mut state = pyo3::ffi::PyInterpreterState_ThreadHead(interpreter);
+            while !state.is_null() {
+                if pyo3::ffi::PyThreadState_GetID(state) == id {
+                    return true;
+                }
+                state = pyo3::ffi::PyThreadState_Next(state);
+            }
+        }
+        false
+    }
+
+    async fn blocking_thread_state_id() -> i64 {
+        tokio::task::spawn_blocking(|| {
+            Python::attach(|py| {
+                retain_thread_state(py);
+                // SAFETY: this call runs under the live attachment above.
+                unsafe { pyo3::ffi::PyThreadState_GetID(pyo3::ffi::PyThreadState_Get()) }
+            })
+        })
+        .await
+        .unwrap()
+    }
+
+    #[test]
+    fn blocking_worker_reuses_state_and_releases_it_at_shutdown() {
+        Python::initialize();
+        let runtime = tokio::runtime::Builder::new_current_thread()
+            .max_blocking_threads(1)
+            .build()
+            .unwrap();
+        let ids = runtime.block_on(async {
+            let mut ids = Vec::new();
+            for _ in 0..40 {
+                ids.push(blocking_thread_state_id().await);
+            }
+            ids
+        });
+        assert!(ids.iter().all(|id| *id == ids[0]));
+        assert!(Python::attach(|py| contains_thread_state(py, ids[0])));
+        let failed = runtime.block_on(async {
+            tokio::task::spawn_blocking(|| {
+                Python::attach(|py| {
+                    retain_thread_state(py);
+                    panic!("test blocking job panic");
+                });
+            })
+            .await
+            .unwrap_err()
+        });
+        assert!(failed.is_panic());
+        assert_eq!(ids[0], runtime.block_on(blocking_thread_state_id()));
+        // Joining must happen with this test thread detached so TLS can clean up.
+        drop(runtime);
+        assert!(!Python::attach(|py| contains_thread_state(py, ids[0])));
+    }
+
+    #[test]
+    fn idle_blocking_worker_releases_state_before_replacement() {
+        Python::initialize();
+        let runtime = tokio::runtime::Builder::new_current_thread()
+            .max_blocking_threads(1)
+            .thread_keep_alive(Duration::from_millis(20))
+            .build()
+            .unwrap();
+        let original = runtime.block_on(blocking_thread_state_id());
+        let deadline = std::time::Instant::now() + Duration::from_secs(5);
+        while Python::attach(|py| contains_thread_state(py, original)) {
+            assert!(
+                std::time::Instant::now() < deadline,
+                "idle worker state was not released"
+            );
+            std::thread::sleep(Duration::from_millis(5));
+        }
+        let replacement = runtime.block_on(blocking_thread_state_id());
+        assert_ne!(original, replacement);
+        drop(runtime);
+        assert!(!Python::attach(|py| contains_thread_state(py, replacement)));
+    }
+}

{chalk_remote_call_python-1.10.7 → chalk_remote_call_python-1.10.9}/chalk-remote-call-rs/chalk-remote-call-server/src/self_consumer.rs

@@ -1351,7 +1351,7 @@ fn ensure_traceparent(headers: &mut HashMap<String, String>) -> String {
 /// `end`, or `error`) to the call's result stream, updating meta status along
 /// the way.
 async fn execute_and_write(
-    handler: &PythonHandler,
+    handler: &Arc<PythonHandler>,
     output: &mut CallOutput,
     work: QueuedCall,
     flush_timeout: Duration,

chalk_remote_call_python-1.10.9/chalk_remote_call/_version.py

@@ -0,0 +1 @@
+__version__ = "1.10.9"

{chalk_remote_call_python-1.10.7 → chalk_remote_call_python-1.10.9/chalk_remote_call_python.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalk-remote-call-python
-Version: 1.10.7
+Version: 1.10.9
 Summary: Chalk remote call Python runtime interface client
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai

{chalk_remote_call_python-1.10.7 → chalk_remote_call_python-1.10.9}/chalk_remote_call_python.egg-info/SOURCES.txt

@@ -14,6 +14,7 @@ chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.runtime.v1.tonic.rs
 chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.utils.v1.rs
 chalk-remote-call-rs/chalk-remote-call-proto/src/gen/descriptor.bin
 chalk-remote-call-rs/chalk-remote-call-server/Cargo.toml
+chalk-remote-call-rs/chalk-remote-call-server/build.rs
 chalk-remote-call-rs/chalk-remote-call-server/src/async_service.rs
 chalk-remote-call-rs/chalk-remote-call-server/src/coalesce.rs
 chalk-remote-call-rs/chalk-remote-call-server/src/fnq_client.rs

chalk_remote_call_python-1.10.7/chalk_remote_call/_version.py

@@ -1 +0,0 @@
-__version__ = "1.10.7"