chalk-remote-call-python 1.6.2__tar.gz → 1.7.0__tar.gz

{chalk_remote_call_python-1.6.2 → chalk_remote_call_python-1.7.0}/PKG-INFO

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

{chalk_remote_call_python-1.6.2 → chalk_remote_call_python-1.7.0}/chalk-remote-call-rs/chalk-remote-call-server/src/python_bridge.rs

@@ -1,7 +1,10 @@
 use std::collections::HashMap;
 
+use chalk_remote_call_proto::chalk::runtime::v1::CallFunctionResponse;
 use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyDict, PyList};
+use tokio::sync::mpsc;
+use tonic::Status;
 
 /// Holds references to the Python handler and bridge function.
 pub struct PythonHandler {
@@ -34,29 +37,35 @@ impl PythonHandler {
         }
     }
 
-    /// Non-batched: one caller per invocation. Dispatches to `process_batches`.
-    ///
-    /// Returns a list of IPC-encoded response chunks — multiple when the
-    /// handler is a generator, one otherwise.
-    pub async fn call(
+    /// 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,
         ipc_bytes: Vec<u8>,
         function_name: String,
         metadata: HashMap<String, String>,
         peer: String,
-    ) -> Result<Vec<Vec<u8>>, PythonError> {
+        response_tx: mpsc::Sender<Result<CallFunctionResponse, Status>>,
+    ) -> Result<(), 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();
 
         tokio::task::spawn_blocking(move || {
-            Python::attach(|py| -> Result<Vec<Vec<u8>>, PythonError> {
+            Python::attach(|py| -> Result<(), PythonError> {
                 let py_bytes = PyBytes::new(py, &ipc_bytes).into_any().unbind();
                 let py_ctx = build_context(py, &function_name, &metadata, &peer)?
                     .into_any()
                     .unbind();
                 let py_arg_names = build_arg_names(py, arg_names.as_deref())?;
-                invoke_process_fn(py, &process_fn, &handler, py_bytes, py_arg_names, py_ctx)
+                let emitter = Py::new(py, ChunkEmitter { response_tx })
+                    .map_err(|e| into_python_error(py, e))?
+                    .into_any();
+
+                process_fn
+                    .call1(py, (py_bytes, handler, py_arg_names, py_ctx, emitter))
+                    .map_err(|e| into_python_error(py, e))?;
+                Ok(())
             })
         })
         .await
@@ -104,6 +113,28 @@ impl PythonHandler {
 
 // ---- Shared helpers -------------------------------------------------------
 
+#[pyclass]
+struct ChunkEmitter {
+    response_tx: mpsc::Sender<Result<CallFunctionResponse, Status>>,
+}
+
+#[pymethods]
+impl ChunkEmitter {
+    fn __call__(&self, py: Python<'_>, chunk: &Bound<'_, PyBytes>) -> PyResult<()> {
+        // Copy the bytes out of Python memory before releasing the GIL.
+        let msg = CallFunctionResponse {
+            feather_stream: chunk.as_bytes().to_vec().into(),
+        };
+        // Release the GIL while parked in blocking_send so a backpressured
+        // client can't stall every other request's Python execution. Collapse
+        // the (large) SendError to `()` inside the closure so it isn't carried
+        // across the detach boundary.
+        py.detach(|| self.response_tx.blocking_send(Ok(msg)).map_err(|_| ()))
+            .map_err(|()| pyo3::exceptions::PyRuntimeError::new_err("client disconnected"))?;
+        Ok(())
+    }
+}
+
 /// Build a `{"function_name": str, "peer": str, "metadata": dict}` context dict for one caller.
 fn build_context<'py>(
     py: Python<'py>,

{chalk_remote_call_python-1.6.2 → chalk_remote_call_python-1.7.0}/chalk-remote-call-rs/chalk-remote-call-server/src/service.rs

@@ -5,7 +5,7 @@ use chalk_remote_call_proto::chalk::runtime::v1::remote_call_service_server::Rem
 use chalk_remote_call_proto::chalk::runtime::v1::{CallFunctionRequest, CallFunctionResponse};
 use tokio_stream::wrappers::ReceiverStream;
 use tonic::{Request, Response, Status, Streaming};
-use tracing::{error, instrument};
+use tracing::{debug, error, instrument};
 
 use crate::coalesce::{BufferedCall, CoalescingQueue};
 use crate::python_bridge::{CallerInput, PythonHandler};
@@ -91,25 +91,26 @@ impl RemoteCallService for RemoteCallServiceImpl {
         // Single-request path.
         let handler = self.python_handler.clone();
         tokio::spawn(async move {
-            match handler.call(all_bytes, function_name, metadata, peer).await {
-                Ok(response_chunks) => {
-                    for chunk in response_chunks {
-                        let msg = CallFunctionResponse {
-                            feather_stream: chunk.into(),
-                        };
-                        if tx.send(Ok(msg)).await.is_err() {
-                            // Client disconnected
-                            break;
-                        }
-                    }
-                }
+            let error_tx = tx.clone();
+            match handler
+                .call_streaming(all_bytes, function_name, metadata, peer, tx)
+                .await
+            {
+                Ok(()) => {}
                 Err(e) => {
-                    error!("Python handler error: {}", e);
-                    let _ = tx
-                        .send(Err(Status::unknown(format!(
-                            "Exception raised during handler execution: {e}"
-                        ))))
-                        .await;
+                    if error_tx.is_closed() {
+                        // The only reason blocking_send fails is a dropped
+                        // receiver — the client went away mid-stream. Not a
+                        // handler failure; nothing to send.
+                        debug!("client disconnected during streaming response");
+                    } else {
+                        error!("Python handler error: {}", e);
+                        let _ = error_tx
+                            .send(Err(Status::unknown(format!(
+                                "Exception raised during handler execution: {e}"
+                            ))))
+                            .await;
+                    }
                 }
             }
         });

chalk_remote_call_python-1.7.0/chalk_remote_call/_version.py

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

{chalk_remote_call_python-1.6.2 → chalk_remote_call_python-1.7.0}/chalk_remote_call/servicer.py

@@ -158,11 +158,22 @@ def _collect_async_gen(agen) -> list:
     return _run_async(_collect())
 
 
+def _drain_async_gen(agen, emit: Callable[[bytes], None]) -> None:
+    """Drain an async generator and emit each encoded item immediately."""
+
+    async def _drain():
+        async for item in agen:
+            emit(encode_record_batch(_coerce_to_record_batch(item)))
+
+    _run_async(_drain())
+
+
 def process_batches(
     ipc_bytes: bytes,
     handler: Callable[..., Any],
     arg_names: list[str] | None,
     context_metadata: dict[str, Any],
+    emit: Callable[[bytes], None] | None = None,
 ) -> list[bytes]:
     """Bridge function called from Rust via PyO3.
 
@@ -170,9 +181,18 @@ def process_batches(
     coerces the result, and encodes the response.
 
     Returns a list of IPC-encoded response bytes (one per input batch).
+    When ``emit`` is provided, each encoded response is emitted immediately and
+    the returned list is empty.
     """
     batches = decode_ipc_stream(ipc_bytes)
     results = []
+
+    def append_or_emit(chunk: bytes) -> None:
+        if emit is None:
+            results.append(chunk)
+        else:
+            emit(chunk)
+
     for batch in batches:
         try:
             event = transform(batch, arg_names)
@@ -193,15 +213,19 @@ def process_batches(
                     # Sync generator: encode each yielded value as a separate chunk.
                     for item in result:
                         result_batch = _coerce_to_record_batch(item)
-                        results.append(encode_record_batch(result_batch))
+                        append_or_emit(encode_record_batch(result_batch))
                 elif inspect.isasyncgen(result):
-                    # Async generator: collect items and encode each as a separate chunk.
-                    for item in _collect_async_gen(result):
-                        result_batch = _coerce_to_record_batch(item)
-                        results.append(encode_record_batch(result_batch))
+                    # Async generator: stream when possible, otherwise preserve
+                    # the existing list-returning bridge contract.
+                    if emit is not None:
+                        _drain_async_gen(result, emit)
+                    else:
+                        for item in _collect_async_gen(result):
+                            result_batch = _coerce_to_record_batch(item)
+                            results.append(encode_record_batch(result_batch))
                 else:
                     result_batch = _coerce_to_record_batch(result)
-                    results.append(encode_record_batch(result_batch))
+                    append_or_emit(encode_record_batch(result_batch))
             except Exception as e:
                 raise RuntimeError(f"Failed to encode response: {e}") from e
 

{chalk_remote_call_python-1.6.2 → chalk_remote_call_python-1.7.0}/chalk_remote_call_python.egg-info/PKG-INFO

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

chalk_remote_call_python-1.6.2/chalk_remote_call/_version.py

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