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

{chalk_remote_call_python-1.9.0/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.9.0
+Version: 1.9.1
 Summary: Chalk remote call Python runtime interface client
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai

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

@@ -511,7 +511,7 @@ dependencies = [
 [[package]]
 name = "chalk-queue"
 version = "0.0.0"
-source = "git+https://github.com/chalk-ai/chalk-private.git?rev=783e9ce6b2c#783e9ce6b2c17563ab3302e5d11f8a664ed00232"
+source = "git+https://github.com/chalk-ai/chalk-private.git?rev=9c61f6a938d8f885f138217ac64149e29d7a691e#9c61f6a938d8f885f138217ac64149e29d7a691e"
 dependencies = [
  "base64 0.22.1",
  "redis",

{chalk_remote_call_python-1.9.0 → 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 = "783e9ce6b2c" }
+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.9.0 → chalk_remote_call_python-1.9.1}/chalk-remote-call-rs/chalk-remote-call-server/src/self_consumer.rs

@@ -66,6 +66,7 @@ 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
@@ -88,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.
@@ -109,6 +118,7 @@ pub struct SelfConsumerConfig {
     pub poll_backoff_max: Duration,
     pub flush_timeout: Duration,
     pub queue_protocol: QueueProtocol,
+    pub stream_read_count: usize,
 }
 
 impl SelfConsumerConfig {
@@ -180,6 +190,12 @@ impl SelfConsumerConfig {
                 })
             })
             .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,
@@ -191,6 +207,7 @@ impl SelfConsumerConfig {
             poll_backoff_max,
             flush_timeout,
             queue_protocol,
+            stream_read_count,
         })
     }
 }
@@ -225,6 +242,10 @@ pub async fn run(
         Uuid::new_v4()
     );
 
+    // 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(),
@@ -238,10 +259,28 @@ pub async fn run(
         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)
@@ -345,6 +384,165 @@ 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>,
@@ -882,6 +1080,7 @@ mod tests {
         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));
@@ -898,6 +1097,7 @@ mod tests {
             Duration::from_millis(DEFAULT_FLUSH_TIMEOUT_MS)
         );
         assert_eq!(config.queue_protocol, QueueProtocol::ListV1);
+        assert_eq!(config.stream_read_count, 1);
     }
 
     #[test]
@@ -940,10 +1140,29 @@ mod tests {
         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]

chalk_remote_call_python-1.9.1/chalk_remote_call/_version.py

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

{chalk_remote_call_python-1.9.0 → 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.9.0 → 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.9.0
+Version: 1.9.1
 Summary: Chalk remote call Python runtime interface client
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai

chalk_remote_call_python-1.9.0/chalk_remote_call/_version.py

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