chalk-remote-call-python 1.8.7__tar.gz → 1.9.1__tar.gz

{chalk_remote_call_python-1.8.7/chalk_remote_call_python.egg-info → chalk_remote_call_python-1.9.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalk-remote-call-python
-Version: 1.8.7
+Version: 1.9.1
 Summary: Chalk remote call Python runtime interface client
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai
@@ -113,6 +113,7 @@ python -m chalk_remote_call --handler my_handler.handler
 | Variable | Description |
 |----------|-------------|
 | `CHALK_INPUT_ARGS` | Comma-separated list of column names (e.g. `x,y,z`). Renames incoming RecordBatch columns by index. If unset, the original column names are used. Reserved columns (below) are held out of this mapping and keep their own name. |
+| `CHALK_FNQ_QUEUE_PROTOCOL` | Self-consumer work transport: `list_v1` (default/backward compatible) or `stream_v1`. Deployment tooling sets this from immutable scaling-group revision metadata. |
 
 #### Reserved columns
 

{chalk_remote_call_python-1.8.7 → chalk_remote_call_python-1.9.1}/README.md

@@ -75,6 +75,7 @@ python -m chalk_remote_call --handler my_handler.handler
 | Variable | Description |
 |----------|-------------|
 | `CHALK_INPUT_ARGS` | Comma-separated list of column names (e.g. `x,y,z`). Renames incoming RecordBatch columns by index. If unset, the original column names are used. Reserved columns (below) are held out of this mapping and keep their own name. |
+| `CHALK_FNQ_QUEUE_PROTOCOL` | Self-consumer work transport: `list_v1` (default/backward compatible) or `stream_v1`. Deployment tooling sets this from immutable scaling-group revision metadata. |
 
 #### Reserved columns
 

{chalk_remote_call_python-1.8.7 → chalk_remote_call_python-1.9.1}/chalk-remote-call-rs/Cargo.lock

@@ -484,7 +484,7 @@ dependencies = [
 [[package]]
 name = "chalk-core"
 version = "0.0.0"
-source = "git+https://github.com/chalk-ai/chalk-private.git?rev=ae8e792f746f77253f682ee0c56b246cfdd88f51#ae8e792f746f77253f682ee0c56b246cfdd88f51"
+source = "git+https://github.com/chalk-ai/chalk-private.git?rev=783e9ce6b2c#783e9ce6b2c17563ab3302e5d11f8a664ed00232"
 dependencies = [
  "anyhow",
  "envconfig",
@@ -499,7 +499,7 @@ dependencies = [
 [[package]]
 name = "chalk-protos"
 version = "0.0.0"
-source = "git+https://github.com/chalk-ai/chalk-private.git?rev=ae8e792f746f77253f682ee0c56b246cfdd88f51#ae8e792f746f77253f682ee0c56b246cfdd88f51"
+source = "git+https://github.com/chalk-ai/chalk-private.git?rev=783e9ce6b2c#783e9ce6b2c17563ab3302e5d11f8a664ed00232"
 dependencies = [
  "prost 0.11.9",
  "prost-build",
@@ -511,7 +511,7 @@ dependencies = [
 [[package]]
 name = "chalk-queue"
 version = "0.0.0"
-source = "git+https://github.com/chalk-ai/chalk-private.git?rev=ae8e792f746f77253f682ee0c56b246cfdd88f51#ae8e792f746f77253f682ee0c56b246cfdd88f51"
+source = "git+https://github.com/chalk-ai/chalk-private.git?rev=9c61f6a938d8f885f138217ac64149e29d7a691e#9c61f6a938d8f885f138217ac64149e29d7a691e"
 dependencies = [
  "base64 0.22.1",
  "redis",
@@ -560,7 +560,7 @@ dependencies = [
 [[package]]
 name = "chalk_metric_definitions"
 version = "0.0.0"
-source = "git+https://github.com/chalk-ai/chalk-private.git?rev=ae8e792f746f77253f682ee0c56b246cfdd88f51#ae8e792f746f77253f682ee0c56b246cfdd88f51"
+source = "git+https://github.com/chalk-ai/chalk-private.git?rev=783e9ce6b2c#783e9ce6b2c17563ab3302e5d11f8a664ed00232"
 dependencies = [
  "chalk-protos",
  "metrics-generator",
@@ -570,7 +570,7 @@ dependencies = [
 [[package]]
 name = "chalk_metrics"
 version = "0.0.0"
-source = "git+https://github.com/chalk-ai/chalk-private.git?rev=ae8e792f746f77253f682ee0c56b246cfdd88f51#ae8e792f746f77253f682ee0c56b246cfdd88f51"
+source = "git+https://github.com/chalk-ai/chalk-private.git?rev=783e9ce6b2c#783e9ce6b2c17563ab3302e5d11f8a664ed00232"
 dependencies = [
  "async-trait",
  "chalk-core",
@@ -604,7 +604,7 @@ dependencies = [
 [[package]]
 name = "chalk_statsd"
 version = "0.0.0"
-source = "git+https://github.com/chalk-ai/chalk-private.git?rev=ae8e792f746f77253f682ee0c56b246cfdd88f51#ae8e792f746f77253f682ee0c56b246cfdd88f51"
+source = "git+https://github.com/chalk-ai/chalk-private.git?rev=783e9ce6b2c#783e9ce6b2c17563ab3302e5d11f8a664ed00232"
 dependencies = [
  "chalk-core",
  "chrono",
@@ -1867,7 +1867,7 @@ dependencies = [
 [[package]]
 name = "metrics-generator"
 version = "0.1.0"
-source = "git+https://github.com/chalk-ai/chalk-private.git?rev=ae8e792f746f77253f682ee0c56b246cfdd88f51#ae8e792f746f77253f682ee0c56b246cfdd88f51"
+source = "git+https://github.com/chalk-ai/chalk-private.git?rev=783e9ce6b2c#783e9ce6b2c17563ab3302e5d11f8a664ed00232"
 dependencies = [
  "codegen",
  "serde",

{chalk_remote_call_python-1.8.7 → chalk_remote_call_python-1.9.1}/chalk-remote-call-rs/chalk-remote-call-server/Cargo.toml

@@ -36,6 +36,6 @@ futures-util = "0.3"
 # `[patch.crates-io]` here it builds against vanilla crates.io redis (it only
 # uses standard APIs). Same rev across all three so cargo reuses one checkout.
 [target.'cfg(target_os = "linux")'.dependencies]
-chalk-queue = { git = "https://github.com/chalk-ai/chalk-private.git", rev = "ae8e792f746f77253f682ee0c56b246cfdd88f51" }
-chalk_metrics = { git = "https://github.com/chalk-ai/chalk-private.git", rev = "ae8e792f746f77253f682ee0c56b246cfdd88f51" }
-chalk_statsd = { git = "https://github.com/chalk-ai/chalk-private.git", rev = "ae8e792f746f77253f682ee0c56b246cfdd88f51" }
+chalk-queue = { git = "https://github.com/chalk-ai/chalk-private.git", rev = "9c61f6a938d8f885f138217ac64149e29d7a691e" }
+chalk_metrics = { git = "https://github.com/chalk-ai/chalk-private.git", rev = "783e9ce6b2c" }
+chalk_statsd = { git = "https://github.com/chalk-ai/chalk-private.git", rev = "783e9ce6b2c" }

{chalk_remote_call_python-1.8.7 → chalk_remote_call_python-1.9.1}/chalk-remote-call-rs/chalk-remote-call-server/src/self_consumer.rs

@@ -3,7 +3,10 @@
 //! This runs *alongside* the gRPC server, which is untouched: the synchronous
 //! / direct `CallFunction` path keeps serving inbound calls.
 //!
-//! Concurrency is capped by a Redis-backed distributed semaphore (per function,
+//! Work is read from the revision's advertised Redis List or Stream transport.
+//! Stream deliveries are acknowledged only after result writes complete and
+//! reclaimed after a crashed owner stops heartbeating. Concurrency is capped
+//! by a Redis-backed distributed semaphore (per function,
 //! across every pod in the scaling group). The dispatch loop dispatches as fast
 //! as it can acquire a slot and pop work; in-flight worker tasks are therefore
 //! bounded by `max_concurrent`, not by traffic. A single per-pod renewer task
@@ -30,14 +33,17 @@
 
 use std::collections::HashMap;
 use std::panic::AssertUnwindSafe;
+use std::str::FromStr;
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
 
 use chalk_queue::semaphore::renew_many;
 use chalk_queue::stream_writer;
 use chalk_queue::types::QueuedCall;
-use chalk_queue::{FunctionQueue, SemaphoreLease};
-use dashmap::DashSet;
+use chalk_queue::{
+    BoundFunctionQueue, DequeuedCall, FunctionQueue, QueueProtocol, QueueReceipt, SemaphoreLease,
+};
+use dashmap::DashMap;
 use futures_util::FutureExt;
 use redis::aio::ConnectionManager;
 use redis::Client;
@@ -59,6 +65,8 @@ const ENV_SEM_TTL_SECS: &str = "CHALK_FNQ_SEM_TTL_SECS";
 const ENV_POLL_BACKOFF_INITIAL_MS: &str = "CHALK_FNQ_POLL_BACKOFF_INITIAL_MS";
 const ENV_POLL_BACKOFF_MAX_SECS: &str = "CHALK_FNQ_POLL_BACKOFF_MAX_SECS";
 const ENV_FLUSH_TIMEOUT_MS: &str = "CHALK_FNQ_FLUSH_TIMEOUT_MS";
+const ENV_QUEUE_PROTOCOL: &str = "CHALK_FNQ_QUEUE_PROTOCOL";
+const ENV_STREAM_READ_COUNT: &str = "CHALK_FNQ_STREAM_READ_COUNT";
 
 /// Lease lifetime (default). Intentionally short relative to how long a call may run: the
 /// per-pod renewer keeps a live call's lease alive (see module docs), so the
@@ -81,6 +89,14 @@ const DEFAULT_POLL_BACKOFF_MAX_SECS: u64 = 2;
 /// chunk append races with the wait setup.
 /// Configurable via `CHALK_FNQ_FLUSH_TIMEOUT_MS`.
 const DEFAULT_FLUSH_TIMEOUT_MS: u64 = 100;
+/// A bounded prefetch window for Redis Streams. This is deliberately much
+/// smaller than a logical handler batch: several reads can feed the shared
+/// coalescer, while one read remains cheap to reserve, decode and recover.
+const DEFAULT_STREAM_READ_COUNT: usize = 64;
+/// Recovery does not belong on the normal dequeue path. A dead consumer is
+/// still recovered promptly, while healthy traffic avoids an XAUTOCLAIM RPC
+/// for every read window.
+const STREAM_RECLAIM_INTERVAL: Duration = Duration::from_secs(10);
 
 /// `peer` recorded in the handler context for queue-consumed calls; there is no
 /// inbound network peer on this path.
@@ -101,6 +117,8 @@ pub struct SelfConsumerConfig {
     pub poll_backoff_initial: Duration,
     pub poll_backoff_max: Duration,
     pub flush_timeout: Duration,
+    pub queue_protocol: QueueProtocol,
+    pub stream_read_count: usize,
 }
 
 impl SelfConsumerConfig {
@@ -162,6 +180,22 @@ impl SelfConsumerConfig {
             .filter(|&n| n > 0)
             .unwrap_or(DEFAULT_FLUSH_TIMEOUT_MS);
         let flush_timeout = Duration::from_millis(flush_timeout_ms);
+        let queue_protocol = std::env::var(ENV_QUEUE_PROTOCOL)
+            .ok()
+            .filter(|value| !value.is_empty())
+            .map(|value| {
+                QueueProtocol::from_str(&value).unwrap_or_else(|error| {
+                    warn!(%error, "invalid queue protocol; using legacy List transport");
+                    QueueProtocol::ListV1
+                })
+            })
+            .unwrap_or_default();
+        let stream_read_count = std::env::var(ENV_STREAM_READ_COUNT)
+            .ok()
+            .and_then(|s| s.parse::<usize>().ok())
+            .filter(|&n| n > 0)
+            .unwrap_or(DEFAULT_STREAM_READ_COUNT)
+            .min(max_concurrent as usize);
 
         Some(Self {
             redis_url,
@@ -172,6 +206,8 @@ impl SelfConsumerConfig {
             poll_backoff_initial,
             poll_backoff_max,
             flush_timeout,
+            queue_protocol,
+            stream_read_count,
         })
     }
 }
@@ -179,7 +215,7 @@ impl SelfConsumerConfig {
 /// Set of lease tokens this pod currently holds. The dispatch loop inserts a
 /// token when it admits a call; the worker removes it on completion. The
 /// renewer heartbeats whatever is in the set on each tick.
-type LiveTokens = Arc<DashSet<String>>;
+type LiveCalls = Arc<DashMap<String, QueueReceipt>>;
 
 /// Connect to Redis and run the consume loop until the process exits. Returns
 /// `Err` only on a fatal setup failure (bad URL / unreachable Redis at startup);
@@ -194,25 +230,64 @@ pub async fn run(
     let client = Client::open(config.redis_url.clone())?;
     let conn = ConnectionManager::new(client).await?;
     let queue = FunctionQueue::new(conn.clone(), config.tenant.clone());
+    let input_queue = queue.bind(&config.function_name, config.queue_protocol);
+    let legacy_queue = (config.queue_protocol == QueueProtocol::StreamV1)
+        .then(|| queue.bind(&config.function_name, QueueProtocol::ListV1));
     let sem = queue.semaphore(&config.function_name, config.max_concurrent, config.sem_ttl);
     let sem_key = queue.sem_key_for(&config.function_name);
-    let live_tokens: LiveTokens = Arc::new(DashSet::new());
+    let live_calls: LiveCalls = Arc::new(DashMap::new());
+    let consumer_name = format!(
+        "{}:{}",
+        std::env::var("HOSTNAME").unwrap_or_else(|_| "fnq-shim".to_string()),
+        Uuid::new_v4()
+    );
 
-    spawn_renewer(conn.clone(), sem_key, live_tokens.clone(), config.sem_ttl);
+    // XGROUP CREATE is a startup operation. Keeping it out of every XREAD and
+    // XAUTOCLAIM avoids two avoidable Redis round trips per dispatch window.
+    input_queue.initialize_consumer_group().await?;
+
+    spawn_renewer(
+        input_queue.clone(),
+        conn.clone(),
+        sem_key,
+        live_calls.clone(),
+        config.sem_ttl,
+    );
 
     info!(
         function = %config.function_name,
         tenant = %config.tenant,
         max_concurrent = config.max_concurrent,
+        queue_protocol = config.queue_protocol.as_str(),
+        stream_read_count = config.stream_read_count,
+        consumer = %consumer_name,
         "fnq self-consumer started"
     );
 
+    if config.queue_protocol == QueueProtocol::StreamV1 {
+        return run_stream_dispatch_loop(
+            &config,
+            queue,
+            input_queue,
+            legacy_queue,
+            sem,
+            conn,
+            consumer_name,
+            live_calls,
+            handler,
+            metrics_pipeline,
+            coalescing_queue,
+        )
+        .await;
+    }
+
     let mut backoff = config.poll_backoff_initial;
     loop {
         // Idle-gate: if there is no work, don't take (then immediately release)
         // a semaphore slot — that churn would race genuine acquirers on other
-        // pods. LLEN is as cheap as the RPOP we'd otherwise do.
-        match queue.current_depth(&config.function_name).await {
+        // pods. During a Stream rollout, also observe the legacy List so a new
+        // shim can drain work queued before its revision metadata changed.
+        match pending_depth(&input_queue, legacy_queue.as_ref()).await {
             Ok(0) => {
                 tokio::time::sleep(backoff).await;
                 backoff = (backoff * 2).min(config.poll_backoff_max);
@@ -243,8 +318,15 @@ pub async fn run(
             }
         };
 
-        let work = match queue.pop(&config.function_name).await {
-            Ok(Some(work)) => work,
+        let delivery = match dequeue_next(
+            &input_queue,
+            legacy_queue.as_ref(),
+            &consumer_name,
+            config.sem_ttl,
+        )
+        .await
+        {
+            Ok(Some(delivery)) => delivery,
             // Raced another consumer (or our own idle-gate) to an empty queue:
             // give the slot back and retry.
             Ok(None) => {
@@ -253,7 +335,7 @@ pub async fn run(
                 continue;
             }
             Err(e) => {
-                warn!(error = %e, "pop failed; releasing slot");
+                warn!(error = %e, "dequeue failed; releasing slot");
                 let _ = lease.release().await;
                 backoff = (backoff * 2).min(config.poll_backoff_max);
                 continue;
@@ -264,7 +346,7 @@ pub async fn run(
 
         // Register the token so the pod-level renewer heartbeats this lease
         // while the call runs.
-        live_tokens.insert(lease.token().to_string());
+        live_calls.insert(lease.token().to_string(), delivery.receipt.clone());
 
         match &coalescing_queue {
             // Batching on: route through the shared coalescing queue so this
@@ -273,11 +355,12 @@ pub async fn run(
             Some(cq) => {
                 tokio::spawn(run_one_coalesced_call(
                     queue.clone(),
+                    input_queue.clone(),
                     cq.clone(),
                     conn.clone(),
-                    work,
+                    delivery,
                     lease,
-                    live_tokens.clone(),
+                    live_calls.clone(),
                     metrics_pipeline.clone(),
                 ));
             }
@@ -286,10 +369,11 @@ pub async fn run(
                 tokio::spawn(run_one_call(
                     handler.clone(),
                     queue.clone(),
+                    input_queue.clone(),
                     conn.clone(),
-                    work,
+                    delivery,
                     lease,
-                    live_tokens.clone(),
+                    live_calls.clone(),
                     config.flush_timeout,
                     metrics_pipeline.clone(),
                 ));
@@ -300,6 +384,202 @@ pub async fn run(
     }
 }
 
+#[allow(clippy::too_many_arguments)]
+async fn run_stream_dispatch_loop(
+    config: &SelfConsumerConfig,
+    queue: FunctionQueue,
+    input_queue: BoundFunctionQueue,
+    legacy_queue: Option<BoundFunctionQueue>,
+    sem: chalk_queue::DistributedSemaphore,
+    conn: ConnectionManager,
+    consumer_name: String,
+    live_calls: LiveCalls,
+    handler: Arc<PythonHandler>,
+    metrics_pipeline: Option<Arc<PublishingMetricsPipeline>>,
+    coalescing_queue: Option<Arc<CoalescingQueue>>,
+) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    let mut backoff = config.poll_backoff_initial;
+    // Reclaim once immediately to recover work from a prior crashed pod, then
+    // periodically. Healthy messages are kept alive by `spawn_renewer`.
+    let mut last_reclaim = Instant::now() - STREAM_RECLAIM_INTERVAL;
+
+    loop {
+        // Capacity must be reserved before XREADGROUP. Reading first would put
+        // a potentially huge number of messages into this consumer's PEL while
+        // they wait for semaphore slots.
+        let mut leases = match sem.try_acquire_many(config.stream_read_count).await {
+            Ok(leases) if !leases.is_empty() => leases,
+            Ok(_) => {
+                tokio::time::sleep(backoff).await;
+                backoff = (backoff * 2).min(config.poll_backoff_max);
+                continue;
+            }
+            Err(error) => {
+                warn!(%error, "batch semaphore acquire failed; backing off");
+                tokio::time::sleep(backoff).await;
+                backoff = (backoff * 2).min(config.poll_backoff_max);
+                continue;
+            }
+        };
+
+        let mut deliveries = Vec::with_capacity(leases.len());
+        if last_reclaim.elapsed() >= STREAM_RECLAIM_INTERVAL {
+            last_reclaim = Instant::now();
+            match input_queue
+                .reclaim_stale(&consumer_name, config.sem_ttl, leases.len())
+                .await
+            {
+                Ok(reclaimed) => deliveries.extend(reclaimed),
+                Err(error) => {
+                    warn!(%error, "Stream recovery scan failed; continuing with new work")
+                }
+            }
+        }
+
+        // During the one-way List → Stream migration, drain old List entries
+        // before reading new Stream traffic. This is one bounded operation per
+        // read window, not an LLEN on every Stream message.
+        if deliveries.len() < leases.len() {
+            if let Some(legacy_queue) = &legacy_queue {
+                match legacy_queue
+                    .dequeue_many(&consumer_name, leases.len() - deliveries.len())
+                    .await
+                {
+                    Ok(legacy) if !legacy.is_empty() => {
+                        info!(function = %input_queue.function_name(), count = legacy.len(), "draining legacy List items during Stream migration");
+                        deliveries.extend(legacy);
+                    }
+                    Ok(_) => {}
+                    Err(error) => {
+                        warn!(%error, "legacy List dequeue failed; continuing with Stream work")
+                    }
+                }
+            }
+        }
+        if deliveries.len() < leases.len() {
+            match input_queue
+                .dequeue_many(&consumer_name, leases.len() - deliveries.len())
+                .await
+            {
+                Ok(new_deliveries) => deliveries.extend(new_deliveries),
+                Err(error) => warn!(%error, "batched Stream dequeue failed"),
+            }
+        }
+
+        if deliveries.is_empty() {
+            for lease in leases {
+                let _ = lease.release().await;
+            }
+            tokio::time::sleep(backoff).await;
+            backoff = (backoff * 2).min(config.poll_backoff_max);
+            continue;
+        }
+
+        backoff = config.poll_backoff_initial;
+        // `dequeue_many` can return fewer records than requested. Only retain
+        // one lease per record and immediately return the unneeded capacity.
+        let used = deliveries.len().min(leases.len());
+        let unused = leases.split_off(used);
+        for lease in unused {
+            let _ = lease.release().await;
+        }
+        for (delivery, lease) in deliveries.into_iter().zip(leases) {
+            dispatch_delivery(
+                queue.clone(),
+                input_queue.clone(),
+                coalescing_queue.clone(),
+                handler.clone(),
+                conn.clone(),
+                delivery,
+                lease,
+                live_calls.clone(),
+                config.flush_timeout,
+                metrics_pipeline.clone(),
+            );
+        }
+    }
+}
+
+#[allow(clippy::too_many_arguments)]
+fn dispatch_delivery(
+    queue: FunctionQueue,
+    input_queue: BoundFunctionQueue,
+    coalescing_queue: Option<Arc<CoalescingQueue>>,
+    handler: Arc<PythonHandler>,
+    conn: ConnectionManager,
+    delivery: DequeuedCall,
+    lease: SemaphoreLease,
+    live_calls: LiveCalls,
+    flush_timeout: Duration,
+    metrics_pipeline: Option<Arc<PublishingMetricsPipeline>>,
+) {
+    live_calls.insert(lease.token().to_string(), delivery.receipt.clone());
+    match coalescing_queue {
+        Some(cq) => {
+            tokio::spawn(run_one_coalesced_call(
+                queue,
+                input_queue,
+                cq,
+                conn,
+                delivery,
+                lease,
+                live_calls,
+                metrics_pipeline,
+            ));
+        }
+        None => {
+            tokio::spawn(run_one_call(
+                handler,
+                queue,
+                input_queue,
+                conn,
+                delivery,
+                lease,
+                live_calls,
+                flush_timeout,
+                metrics_pipeline,
+            ));
+        }
+    }
+}
+
+async fn pending_depth(
+    queue: &BoundFunctionQueue,
+    legacy_queue: Option<&BoundFunctionQueue>,
+) -> Result<u64, redis::RedisError> {
+    let selected = queue.current_depth().await?;
+    let legacy = match legacy_queue {
+        Some(queue) => queue.current_depth().await?,
+        None => 0,
+    };
+    Ok(selected + legacy)
+}
+
+async fn dequeue_next(
+    queue: &BoundFunctionQueue,
+    legacy_queue: Option<&BoundFunctionQueue>,
+    consumer_name: &str,
+    reclaim_idle: Duration,
+) -> Result<Option<DequeuedCall>, Box<dyn std::error::Error + Send + Sync>> {
+    // Recovered work wins first. The legacy List is then drained before new
+    // Stream traffic; producers never route new work to both transports.
+    if let Some(reclaimed) = queue
+        .reclaim_stale(consumer_name, reclaim_idle, 1)
+        .await?
+        .into_iter()
+        .next()
+    {
+        return Ok(Some(reclaimed));
+    }
+    if let Some(legacy_queue) = legacy_queue {
+        if let Some(legacy) = legacy_queue.dequeue(consumer_name).await? {
+            info!(function = %queue.function_name(), "draining legacy List item during Stream migration");
+            return Ok(Some(legacy));
+        }
+    }
+    queue.dequeue(consumer_name).await
+}
+
 /// Run a single queued call to completion, write its results to the call's
 /// stream, then deregister + release its lease. Always releases the lease,
 /// even if execution or stream writes fail or panic.
@@ -307,15 +587,20 @@ pub async fn run(
 async fn run_one_call(
     handler: Arc<PythonHandler>,
     queue: FunctionQueue,
+    input_queue: BoundFunctionQueue,
     mut conn: ConnectionManager,
-    work: QueuedCall,
+    delivery: DequeuedCall,
     lease: SemaphoreLease,
-    live_tokens: LiveTokens,
+    live_calls: LiveCalls,
     flush_timeout: Duration,
     metrics_pipeline: Option<Arc<PublishingMetricsPipeline>>,
 ) {
+    let DequeuedCall {
+        call: work,
+        receipt,
+    } = delivery;
     let call_id = work.call_id.clone();
-    match AssertUnwindSafe(execute_and_write(
+    let completed = match AssertUnwindSafe(execute_and_write(
         &handler,
         &queue,
         &mut conn,
@@ -326,11 +611,12 @@ async fn run_one_call(
     .catch_unwind()
     .await
     {
-        Ok(Ok(())) => {}
+        Ok(Ok(())) => true,
         Ok(Err(e)) => {
             // The call's result stream may be incomplete; the call's 24h stream TTL
             // and the meta status are the backstop for a stuck poller.
             error!(call_id = %call_id, error = %e, "failed to execute/write queued call");
+            false
         }
         Err(panic_payload) => {
             let panic_message = panic_payload
@@ -343,12 +629,41 @@ async fn run_one_call(
                 panic = %panic_message,
                 "panic while executing/writing queued call"
             );
+            false
         }
-    }
+    };
 
-    // Deregister before releasing so the renewer stops heartbeating this token
-    // immediately; the subsequent release ZREMs it from the semaphore ZSET.
-    live_tokens.remove(lease.token());
+    finish_delivery(
+        &input_queue,
+        &receipt,
+        completed,
+        &call_id,
+        lease,
+        &live_calls,
+    )
+    .await;
+}
+
+async fn finish_delivery(
+    queue: &BoundFunctionQueue,
+    receipt: &QueueReceipt,
+    completed: bool,
+    call_id: &str,
+    lease: SemaphoreLease,
+    live_calls: &LiveCalls,
+) {
+    // Stop renewing before ack. If completion or ack failed, the Stream item
+    // becomes reclaimable after the idle timeout; List behavior is unchanged.
+    live_calls.remove(lease.token());
+    if completed {
+        match queue.acknowledge(receipt).await {
+            Ok(true) => {}
+            Ok(false) => warn!(%call_id, "delivery ownership was lost before acknowledgement"),
+            Err(error) => {
+                warn!(%call_id, %error, "delivery acknowledgement failed; item will be reclaimed")
+            }
+        }
+    }
     if let Err(e) = lease.release().await {
         warn!(call_id = %call_id, error = %e, "failed to release semaphore slot (will auto-recover after ttl)");
     }
@@ -357,17 +672,23 @@ async fn run_one_call(
 /// Coalesced counterpart to [`run_one_call`]: submit the call into the shared
 /// [`CoalescingQueue`], write its single result blob to the call's stream, then
 /// deregister + release its lease (always, even on failure).
+#[allow(clippy::too_many_arguments)]
 async fn run_one_coalesced_call(
     queue: FunctionQueue,
+    input_queue: BoundFunctionQueue,
     coalescing_queue: Arc<CoalescingQueue>,
     mut conn: ConnectionManager,
-    work: QueuedCall,
+    delivery: DequeuedCall,
     lease: SemaphoreLease,
-    live_tokens: LiveTokens,
+    live_calls: LiveCalls,
     metrics_pipeline: Option<Arc<PublishingMetricsPipeline>>,
 ) {
+    let DequeuedCall {
+        call: work,
+        receipt,
+    } = delivery;
     let call_id = work.call_id.clone();
-    if let Err(e) = coalesced_execute_and_write(
+    let completed = match coalesced_execute_and_write(
         &coalescing_queue,
         &queue,
         &mut conn,
@@ -376,17 +697,24 @@ async fn run_one_coalesced_call(
     )
     .await
     {
-        // The call's result stream may be incomplete; the 24h stream TTL and
-        // the meta status are the backstop for a stuck poller.
-        error!(call_id = %call_id, error = %e, "failed to execute/write coalesced queued call");
-    }
+        Ok(()) => true,
+        Err(e) => {
+            // The call's result stream may be incomplete; the 24h stream TTL and
+            // the meta status are the backstop for a stuck poller.
+            error!(call_id = %call_id, error = %e, "failed to execute/write coalesced queued call");
+            false
+        }
+    };
 
-    // Deregister before releasing so the renewer stops heartbeating this token
-    // immediately; the subsequent release ZREMs it from the semaphore ZSET.
-    live_tokens.remove(lease.token());
-    if let Err(e) = lease.release().await {
-        warn!(call_id = %call_id, error = %e, "failed to release semaphore slot (will auto-recover after ttl)");
-    }
+    finish_delivery(
+        &input_queue,
+        &receipt,
+        completed,
+        &call_id,
+        lease,
+        &live_calls,
+    )
+    .await;
 }
 
 /// Submit one call into the shared coalescing queue and relay its single
@@ -666,14 +994,15 @@ async fn flush_new(
 
 /// Spawn the single per-pod renewer: every `SEM_RENEW_INTERVAL` it heartbeats
 /// all currently-held lease tokens in one `renew_many` round-trip. A reclaimed
-/// token means this pod stalled past `sem_ttl` without renewing (the slot was
-/// then handed to another consumer); the work was already popped so it won't be
-/// re-run from the queue, and a GIL-blocked handler can't be aborted, so we only
-/// surface it.
+/// token means this pod stalled past `sem_ttl` without renewing. Stream
+/// receipts are renewed in grouped round trips too; ownership loss is surfaced
+/// because a GIL-blocked handler cannot be aborted, while the reclaimed call may
+/// already be running on another pod.
 fn spawn_renewer(
+    queue: BoundFunctionQueue,
     mut conn: ConnectionManager,
     sem_key: String,
-    live_tokens: LiveTokens,
+    live_calls: LiveCalls,
     sem_ttl: Duration,
 ) {
     tokio::spawn(async move {
@@ -681,7 +1010,11 @@ fn spawn_renewer(
         ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
         loop {
             ticker.tick().await;
-            let tokens: Vec<String> = live_tokens.iter().map(|r| r.key().clone()).collect();
+            let active: Vec<(String, QueueReceipt)> = live_calls
+                .iter()
+                .map(|entry| (entry.key().clone(), entry.value().clone()))
+                .collect();
+            let tokens: Vec<String> = active.iter().map(|(token, _)| token.clone()).collect();
             if tokens.is_empty() {
                 continue;
             }
@@ -695,6 +1028,16 @@ fn spawn_renewer(
                 Ok(_) => {}
                 Err(e) => warn!(error = %e, "batch lease renew failed"),
             }
+            let receipts: Vec<QueueReceipt> =
+                active.into_iter().map(|(_, receipt)| receipt).collect();
+            match queue.heartbeat_many(&receipts).await {
+                Ok(lost) if !lost.is_empty() => warn!(
+                    count = lost.len(),
+                    "Stream deliveries reclaimed while calls were still running"
+                ),
+                Ok(_) => {}
+                Err(error) => warn!(%error, "batch Stream delivery renewal failed"),
+            }
         }
     });
 }
@@ -736,6 +1079,8 @@ mod tests {
         std::env::remove_var(ENV_POLL_BACKOFF_INITIAL_MS);
         std::env::remove_var(ENV_POLL_BACKOFF_MAX_SECS);
         std::env::remove_var(ENV_FLUSH_TIMEOUT_MS);
+        std::env::remove_var(ENV_QUEUE_PROTOCOL);
+        std::env::remove_var(ENV_STREAM_READ_COUNT);
 
         let config = SelfConsumerConfig::from_env().expect("config should parse");
         assert_eq!(config.sem_ttl, Duration::from_secs(DEFAULT_SEM_TTL_SECS));
@@ -751,6 +1096,8 @@ mod tests {
             config.flush_timeout,
             Duration::from_millis(DEFAULT_FLUSH_TIMEOUT_MS)
         );
+        assert_eq!(config.queue_protocol, QueueProtocol::ListV1);
+        assert_eq!(config.stream_read_count, 1);
     }
 
     #[test]
@@ -785,6 +1132,39 @@ mod tests {
         assert_eq!(config.max_concurrent, 1);
     }
 
+    #[test]
+    fn test_stream_queue_protocol_from_env() {
+        let _env = lock_env();
+
+        std::env::set_var(ENV_REDIS_URL, "redis://localhost");
+        std::env::set_var(ENV_TENANT, "test_tenant");
+        std::env::set_var(ENV_FUNCTION_NAME, "test_fn");
+        std::env::set_var(ENV_QUEUE_PROTOCOL, "stream_v1");
+        std::env::set_var(ENV_MAX_CONCURRENT, "128");
+        std::env::remove_var(ENV_STREAM_READ_COUNT);
+
+        let config = SelfConsumerConfig::from_env().expect("config should parse");
+        assert_eq!(config.queue_protocol, QueueProtocol::StreamV1);
+        assert_eq!(config.stream_read_count, DEFAULT_STREAM_READ_COUNT);
+        std::env::remove_var(ENV_QUEUE_PROTOCOL);
+        std::env::remove_var(ENV_STREAM_READ_COUNT);
+    }
+
+    #[test]
+    fn test_stream_read_count_is_configurable_and_capped_by_concurrency() {
+        let _env = lock_env();
+
+        std::env::set_var(ENV_REDIS_URL, "redis://localhost");
+        std::env::set_var(ENV_TENANT, "test_tenant");
+        std::env::set_var(ENV_FUNCTION_NAME, "test_fn");
+        std::env::set_var(ENV_MAX_CONCURRENT, "32");
+        std::env::set_var(ENV_STREAM_READ_COUNT, "128");
+
+        let config = SelfConsumerConfig::from_env().expect("config should parse");
+        assert_eq!(config.stream_read_count, 32);
+        std::env::remove_var(ENV_STREAM_READ_COUNT);
+    }
+
     #[test]
     fn test_invalid_timeout_values_use_defaults() {
         let _env = lock_env();
@@ -894,7 +1274,7 @@ mod tests {
 
         // Verify notify_handle can be cloned (Arc behavior)
         let _notify_handle_clone = notify_handle.clone();
-        assert_eq!(std::ptr::eq(&notify_handle, &_notify_handle_clone), false);
+        assert!(!std::ptr::eq(&notify_handle, &_notify_handle_clone));
         // But they should point to the same underlying Notify
     }
 
@@ -977,42 +1357,42 @@ mod tests {
     // Integration tests for live token tracking and core dispatch logic
     #[tokio::test]
     async fn test_live_tokens_insert_and_remove() {
-        let live_tokens: LiveTokens = Arc::new(DashSet::new());
+        let live_calls: LiveCalls = Arc::new(DashMap::new());
         let token1 = "lease-token-1";
         let token2 = "lease-token-2";
 
         // Simulate dispatch: register tokens as calls are admitted
-        live_tokens.insert(token1.to_string());
-        live_tokens.insert(token2.to_string());
+        live_calls.insert(token1.to_string(), QueueReceipt::ListV1);
+        live_calls.insert(token2.to_string(), QueueReceipt::ListV1);
 
-        assert!(live_tokens.contains(token1));
-        assert!(live_tokens.contains(token2));
-        assert_eq!(live_tokens.iter().count(), 2, "should track both tokens");
+        assert!(live_calls.contains_key(token1));
+        assert!(live_calls.contains_key(token2));
+        assert_eq!(live_calls.iter().count(), 2, "should track both tokens");
 
         // Simulate cleanup: remove as calls complete
-        live_tokens.remove(token1);
-        assert!(!live_tokens.contains(token1));
-        assert!(live_tokens.contains(token2));
-        assert_eq!(live_tokens.iter().count(), 1);
+        live_calls.remove(token1);
+        assert!(!live_calls.contains_key(token1));
+        assert!(live_calls.contains_key(token2));
+        assert_eq!(live_calls.iter().count(), 1);
 
         // Final cleanup
-        live_tokens.remove(token2);
-        assert_eq!(live_tokens.iter().count(), 0, "all tokens cleaned up");
+        live_calls.remove(token2);
+        assert_eq!(live_calls.iter().count(), 0, "all tokens cleaned up");
     }
 
     #[tokio::test]
     async fn test_live_tokens_concurrent_access() {
-        let live_tokens: LiveTokens = Arc::new(DashSet::new());
+        let live_calls: LiveCalls = Arc::new(DashMap::new());
         let mut handles = vec![];
 
         // Spawn 10 tasks that each insert and remove a token
         for i in 0..10 {
-            let tokens = live_tokens.clone();
+            let calls = live_calls.clone();
             let handle = tokio::spawn(async move {
                 let token = format!("token-{}", i);
-                tokens.insert(token.clone());
+                calls.insert(token.clone(), QueueReceipt::ListV1);
                 tokio::time::sleep(Duration::from_millis(1)).await;
-                tokens.remove(&token);
+                calls.remove(&token);
             });
             handles.push(handle);
         }
@@ -1024,7 +1404,7 @@ mod tests {
 
         // All tokens should be cleaned up
         assert_eq!(
-            live_tokens.iter().count(),
+            live_calls.iter().count(),
             0,
             "concurrent inserts/removes should leave set empty"
         );

chalk_remote_call_python-1.9.1/chalk_remote_call/_version.py

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

{chalk_remote_call_python-1.8.7 → chalk_remote_call_python-1.9.1}/chalk_remote_call/tracing.py

@@ -11,6 +11,9 @@ _TRACER_NAME = "chalk_remote_call"
 _TRACE_CONTEXT_METADATA_KEYS = ("traceparent", "tracestate", "baggage")
 _REMOTE_FUNCTION_TRACE_POLICY_ENV_VAR = "CHALK_REMOTE_FUNCTION_TRACE_POLICY"
 _REMOTE_FUNCTION_TRACE_SAMPLE_RATE_ENV_VAR = "CHALK_REMOTE_FUNCTION_TRACE_SAMPLE_RATE"
+_REMOTE_FUNCTION_TRACE_OVERRIDE_KEY = "x-chalk-remote-function-tracing"
+_TRACE_OVERRIDE_ON = "on"
+_TRACE_OVERRIDE_OFF = "off"
 _TRACE_POLICY_PARENT_BASED_ALWAYS_OFF = "parentbased_always_off"
 _TRACE_POLICY_PARENT_BASED_TRACE_ID_RATIO = "parentbased_traceidratio"
 _TRACE_POLICY_ALWAYS_OFF = "always_off"
@@ -74,7 +77,9 @@ def _trace_policy() -> str | None:
     return None
 
 
-def _trace_sample_rate() -> float | None:
+def _trace_sample_rate(force_trace: bool) -> float | None:
+    if force_trace:
+        return 1.0
     raw_rate = os.environ.get(_REMOTE_FUNCTION_TRACE_SAMPLE_RATE_ENV_VAR)
     if raw_rate is None:
         return None
@@ -136,11 +141,11 @@ def _configure_runtime_tracing(sample_rate: float) -> None:
     _runtime_tracing_configured = True
 
 
-def _sample_rate_for_inherited_parent(policy: str) -> float | None:
+def _sample_rate_for_inherited_parent(policy: str, force_trace: bool) -> float | None:
     if policy == _TRACE_POLICY_PARENT_BASED_ALWAYS_OFF:
         return 0.0
     if policy == _TRACE_POLICY_PARENT_BASED_TRACE_ID_RATIO:
-        return _trace_sample_rate()
+        return _trace_sample_rate(force_trace)
     return None
 
 
@@ -148,6 +153,7 @@ def _span_metadata(
     first_metadata: Mapping[str, Any] | None,
     traceparent_metadata: Mapping[str, Any] | None,
     policy: str,
+    force_trace: bool,
 ) -> tuple[Mapping[str, Any] | None, Any | None, Any]:
     if traceparent_metadata is not None:
         carrier = _trace_carrier(traceparent_metadata)
@@ -164,7 +170,7 @@ def _span_metadata(
         if not parent_span_context.is_valid:
             return None, None, None
 
-        sample_rate = _sample_rate_for_inherited_parent(policy)
+        sample_rate = _sample_rate_for_inherited_parent(policy, force_trace)
         if sample_rate is None:
             return None, None, None
         _configure_runtime_tracing(sample_rate)
@@ -173,7 +179,7 @@ def _span_metadata(
     if policy != _TRACE_POLICY_PARENT_BASED_TRACE_ID_RATIO:
         return None, None, None
 
-    sample_rate = _trace_sample_rate()
+    sample_rate = _trace_sample_rate(force_trace)
     if sample_rate is None:
         return None, None, None
 
@@ -184,6 +190,17 @@ def _span_metadata(
     return first_metadata, None, otel_modules
 
 
+def _force_sampled_traceparent(metadata: Mapping[str, Any]) -> Mapping[str, Any]:
+    parts = str(metadata.get("traceparent") or "").split("-")
+    if len(parts) != 4 or len(parts[3]) != 2:
+        return metadata
+    try:
+        parts[3] = f"{int(parts[3], 16) | 1:02x}"
+    except ValueError:
+        return metadata
+    return {**metadata, "traceparent": "-".join(parts)}
+
+
 def _always_off_span_metadata(
     first_metadata: Mapping[str, Any] | None,
     traceparent_metadata: Mapping[str, Any] | None,
@@ -275,12 +292,21 @@ def remote_function_invocation_span(
     context_metadata: Any,
     coalesced_count: int | None = None,
 ) -> Iterator[None]:
-    policy = _trace_policy()
+    first_metadata, traceparent_metadata = _collect_trace_metadata(context_metadata)
+    override = str((first_metadata or {}).get(_REMOTE_FUNCTION_TRACE_OVERRIDE_KEY) or "").strip().lower()
+    if override == _TRACE_OVERRIDE_ON:
+        policy = _TRACE_POLICY_PARENT_BASED_TRACE_ID_RATIO
+    elif override == _TRACE_OVERRIDE_OFF:
+        policy = _TRACE_POLICY_ALWAYS_OFF
+    else:
+        policy = _trace_policy()
     if policy is None:
         yield
         return
 
-    first_metadata, traceparent_metadata = _collect_trace_metadata(context_metadata)
+    force_trace = override == _TRACE_OVERRIDE_ON
+    if force_trace and traceparent_metadata is not None:
+        traceparent_metadata = _force_sampled_traceparent(traceparent_metadata)
     if policy == _TRACE_POLICY_PARENT_BASED_ALWAYS_OFF and traceparent_metadata is None:
         yield
         return
@@ -295,6 +321,7 @@ def remote_function_invocation_span(
             first_metadata,
             traceparent_metadata,
             policy,
+            force_trace,
         )
         tracer = _get_tracer(otel_modules[2]) if otel_modules is not None else None
 

{chalk_remote_call_python-1.8.7 → chalk_remote_call_python-1.9.1/chalk_remote_call_python.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalk-remote-call-python
-Version: 1.8.7
+Version: 1.9.1
 Summary: Chalk remote call Python runtime interface client
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai
@@ -113,6 +113,7 @@ python -m chalk_remote_call --handler my_handler.handler
 | Variable | Description |
 |----------|-------------|
 | `CHALK_INPUT_ARGS` | Comma-separated list of column names (e.g. `x,y,z`). Renames incoming RecordBatch columns by index. If unset, the original column names are used. Reserved columns (below) are held out of this mapping and keep their own name. |
+| `CHALK_FNQ_QUEUE_PROTOCOL` | Self-consumer work transport: `list_v1` (default/backward compatible) or `stream_v1`. Deployment tooling sets this from immutable scaling-group revision metadata. |
 
 #### Reserved columns
 

chalk_remote_call_python-1.8.7/chalk_remote_call/_version.py

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