chalk-remote-call-python 1.7.1__tar.gz → 1.7.2__tar.gz

{chalk_remote_call_python-1.7.1 → chalk_remote_call_python-1.7.2}/PKG-INFO

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

{chalk_remote_call_python-1.7.1 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/chalk-remote-call-server/src/async_service.rs

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant};
 
 use chalk_remote_call_proto::chalk::common::v1::ChalkError;
 use chalk_remote_call_proto::chalk::runtime::v1::async_remote_call_service_server::AsyncRemoteCallService;
@@ -14,6 +15,14 @@ use tracing::{debug, error, info, instrument};
 use crate::python_bridge::{ChunkBuffer, PythonHandler};
 use crate::service::extract_metadata;
 
+/// How often the orphan sweeper runs.
+const SWEEP_INTERVAL: Duration = Duration::from_secs(30);
+/// How long a *terminal* entry is retained before the sweeper reclaims it. Only
+/// matters for orphans — calls the consumer drained are evicted immediately on
+/// their terminal poll. Comfortably larger than the consumer's poll cadence so
+/// the sweeper never races a live consumer that is about to read a result.
+const RETENTION: Duration = Duration::from_mins(30);
+
 /// In-memory execution record for one enqueued call.
 ///
 /// Lives only on the sg pod that accepted the `EnqueueRemoteCall`, and is lost
@@ -31,6 +40,9 @@ struct CallState {
     chunks: ChunkBuffer,
     /// Populated only when `status == Failed`; surfaced as a `ChalkError`.
     error: Option<String>,
+    /// When the call reached a terminal status. Used only by the orphan sweeper
+    /// to decide when an un-drained terminal entry is safe to reclaim.
+    completed_at: Option<Instant>,
 }
 
 impl CallState {
@@ -39,12 +51,20 @@ impl CallState {
             status: RemoteCallStatus::Pending,
             chunks: Arc::new(Mutex::new(Vec::new())),
             error: None,
+            completed_at: None,
         }
     }
 }
 
 type Registry = Arc<Mutex<HashMap<String, CallState>>>;
 
+fn is_terminal(status: RemoteCallStatus) -> bool {
+    matches!(
+        status,
+        RemoteCallStatus::Completed | RemoteCallStatus::Failed
+    )
+}
+
 /// Async (enqueue + poll) executor for the scaling-group hop.
 ///
 /// `EnqueueRemoteCall` runs the handler in a detached background task that
@@ -53,6 +73,11 @@ type Registry = Arc<Mutex<HashMap<String, CallState>>>;
 /// function execution from any single connection's lifetime, so the path is
 /// immune to gateway request/stream timeouts — every RPC here is sub-second
 /// regardless of how long the function runs.
+///
+/// Entries are reclaimed two ways: the terminal poll that delivers a result
+/// evicts its own entry (the consumer drains a call exactly once, then persists
+/// to Redis and stops polling it), and a background sweeper reclaims orphans —
+/// calls that finished but whose consumer died before the terminal poll.
 pub struct AsyncRemoteCallServiceImpl {
     python_handler: Arc<PythonHandler>,
     registry: Registry,
@@ -66,9 +91,13 @@ pub struct AsyncRemoteCallServiceImpl {
 
 impl AsyncRemoteCallServiceImpl {
     pub fn new(python_handler: Arc<PythonHandler>, batching_enabled: bool) -> Self {
+        let registry: Registry = Arc::new(Mutex::new(HashMap::new()));
+        // Runs inside the server's tokio runtime (the service is constructed
+        // within `run_server`); reclaims orphaned terminal entries.
+        spawn_sweeper(registry.clone());
         Self {
             python_handler,
-            registry: Arc::new(Mutex::new(HashMap::new())),
+            registry,
             next_id: AtomicU64::new(1),
             batching_enabled,
         }
@@ -138,6 +167,7 @@ impl AsyncRemoteCallService for AsyncRemoteCallServiceImpl {
 
             let mut reg = registry.lock().expect("registry mutex poisoned");
             if let Some(s) = reg.get_mut(&id) {
+                s.completed_at = Some(Instant::now());
                 match result {
                     Ok(()) => {
                         s.status = RemoteCallStatus::Completed;
@@ -170,16 +200,23 @@ impl AsyncRemoteCallService for AsyncRemoteCallServiceImpl {
                 .map_err(|_| Status::invalid_argument(format!("invalid cursor: {}", req.cursor)))?
         };
 
-        // Read status/error and grab the chunk-buffer handle under the map lock,
-        // then release it before locking the buffer. `status` is set to a
-        // terminal value only after the handler returns (all chunks already
-        // appended), so seeing `Completed` guarantees the buffer is complete.
+        // Read status/error and grab the chunk-buffer handle under the map lock.
+        // `status` reaches a terminal value only after the handler returns (all
+        // chunks already appended), so seeing `Completed`/`Failed` guarantees
+        // the buffer is complete — and this poll delivers it in full. Since the
+        // consumer drains a call to terminal exactly once (then persists to
+        // Redis and never re-polls this call_id), we evict the entry here. The
+        // cloned `chunks` Arc keeps the buffer alive for the read below.
         let (status, error, chunks) = {
-            let reg = self.registry.lock().expect("registry mutex poisoned");
+            let mut reg = self.registry.lock().expect("registry mutex poisoned");
             let state = reg
                 .get(&req.call_id)
                 .ok_or_else(|| Status::not_found(format!("unknown call_id: {}", req.call_id)))?;
-            (state.status, state.error.clone(), state.chunks.clone())
+            let snapshot = (state.status, state.error.clone(), state.chunks.clone());
+            if is_terminal(snapshot.0) {
+                reg.remove(&req.call_id);
+            }
+            snapshot
         };
 
         let buf = chunks.lock().expect("chunk buffer poisoned");
@@ -236,3 +273,24 @@ fn set_status(registry: &Registry, call_id: &str, status: RemoteCallStatus) {
         }
     }
 }
+
+/// Background task: periodically reclaim *orphaned* terminal entries — calls
+/// that finished but whose consumer died before issuing the terminal poll that
+/// would normally evict them. Entries the consumer drains are removed on their
+/// terminal poll and never reach the sweeper; non-terminal entries (still
+/// running) are never reclaimed here.
+fn spawn_sweeper(registry: Registry) {
+    tokio::spawn(async move {
+        let mut ticker = tokio::time::interval(SWEEP_INTERVAL);
+        ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
+        loop {
+            ticker.tick().await;
+            let Ok(mut reg) = registry.lock() else {
+                continue;
+            };
+            reg.retain(|_, s| {
+                !(is_terminal(s.status) && s.completed_at.is_some_and(|t| t.elapsed() > RETENTION))
+            });
+        }
+    });
+}

chalk_remote_call_python-1.7.2/chalk_remote_call/_version.py

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

{chalk_remote_call_python-1.7.1 → chalk_remote_call_python-1.7.2}/chalk_remote_call_python.egg-info/PKG-INFO

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

chalk_remote_call_python-1.7.1/chalk_remote_call/_version.py

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