chalk-remote-call-python 1.4.0__tar.gz → 1.5.0__tar.gz

{chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.0}/PKG-INFO

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

chalk_remote_call_python-1.5.0/chalk_remote_call/_version.py

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

{chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.0}/chalk_remote_call/cli.py

@@ -80,6 +80,17 @@ def main() -> None:
             "(env: CHALK_REMOTE_CALL_MAX_BUFFER_DURATION_MS)"
         ),
     )
+    parser.add_argument(
+        "--batching-mode",
+        choices=["per_caller", "combined"],
+        default=os.environ.get("CHALK_REMOTE_CALL_BATCHING_MODE") or None,
+        help=(
+            "Coalescing variant when --max-batching-size is set. "
+            "'per_caller' (default) passes events as a list of dicts; "
+            "'combined' passes a single concatenated RecordBatch with "
+            "offsets. (env: CHALK_REMOTE_CALL_BATCHING_MODE)"
+        ),
+    )
     parser.add_argument(
         "--log-level",
         default="INFO",
@@ -135,4 +146,5 @@ def main() -> None:
         arg_names=arg_names,
         max_batching_size=args.max_batching_size,
         max_buffer_duration_ms=args.max_buffer_duration_ms,
+        batching_mode=args.batching_mode,
     )

{chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.0}/chalk_remote_call/server.py

@@ -1,14 +1,38 @@
 from __future__ import annotations
 
 import logging
+import os
 from collections.abc import Callable
-from typing import Any
+from typing import Any, Literal
 
 from chalk_remote_call._native import start_server
-from chalk_remote_call.servicer import process_batches, process_batches_coalesced
+from chalk_remote_call.servicer import (
+    process_batches,
+    process_batches_coalesced,
+    process_batches_coalesced_combined,
+)
 
 logger = logging.getLogger(__name__)
 
+_BATCHING_MODE_ENV = "CHALK_REMOTE_CALL_BATCHING_MODE"
+_VALID_BATCHING_MODES = ("per_caller", "combined")
+
+
+def _resolve_batching_mode(kwarg: str | None) -> str:
+    """Resolve batching mode from kwarg or env var.
+
+    Precedence: explicit kwarg > env var > default "per_caller". The env var
+    is the canonical opt-in for combined mode; the kwarg exists so embedders
+    can override programmatically.
+    """
+    if kwarg is not None:
+        mode = kwarg
+    else:
+        mode = (os.environ.get(_BATCHING_MODE_ENV, "") or "per_caller").strip() or "per_caller"
+    if mode not in _VALID_BATCHING_MODES:
+        raise ValueError(f"batching_mode must be one of {_VALID_BATCHING_MODES}, got {mode!r}")
+    return mode
+
 
 def serve(
     handler: Callable[..., Any],
@@ -20,15 +44,21 @@ def serve(
     arg_names: list[str] | None = None,
     max_batching_size: int | None = None,
     max_buffer_duration_ms: int | None = None,
+    batching_mode: Literal["per_caller", "combined"] | None = None,
 ) -> None:
     """Start the gRPC server implementing RemoteCallService.
 
     Args:
         handler: The user's handler function.
             Single-request mode (default): (event, context) -> result.
-            Coalesced mode (when max_batching_size is set):
+            Coalesced per-caller mode (when max_batching_size is set and
+                batching_mode == "per_caller"):
                 (events: list[dict], contexts: list[dict]) -> list[result],
                 one result per input.
+            Coalesced combined mode (when max_batching_size is set and
+                batching_mode == "combined"):
+                (combined: pa.RecordBatch, offsets: list[int], contexts: list[dict])
+                    -> pa.RecordBatch with combined.num_rows rows.
         host: The host to bind to.
         port: The port to bind to.
         workers: Number of worker threads for the tokio runtime.
@@ -40,21 +70,30 @@ def serve(
         max_buffer_duration_ms: Max time in milliseconds to buffer incoming
             requests before flushing even if max_batching_size hasn't been
             reached. Defaults to 1000ms when batching is enabled.
+        batching_mode: "per_caller" (default) or "combined". Only effective
+            when max_batching_size > 0. Overrides the
+            CHALK_REMOTE_CALL_BATCHING_MODE env var when set; otherwise the
+            env var (or "per_caller") wins.
     """
     if on_startup is not None:
         logger.info("Running startup hook...")
         on_startup()
 
+    mode = _resolve_batching_mode(batching_mode)
+
     if max_batching_size is not None and max_batching_size > 0:
-        process_fn = process_batches_coalesced
+        process_fn = process_batches_coalesced_combined if mode == "combined" else process_batches_coalesced
         batch_size = max_batching_size
         duration_ms = max_buffer_duration_ms if max_buffer_duration_ms is not None else 1000
         logger.info(
-            "Batching enabled: max_batching_size=%d max_buffer_duration_ms=%d",
+            "Batching enabled: mode=%s max_batching_size=%d max_buffer_duration_ms=%d",
+            mode,
             batch_size,
             duration_ms,
         )
     else:
+        if mode == "combined":
+            raise ValueError("batching_mode='combined' requires max_batching_size > 0")
         process_fn = process_batches
         batch_size = 0
         duration_ms = 0

{chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.0}/chalk_remote_call/servicer.py

@@ -295,3 +295,114 @@ def process_batches_coalesced(
             raise RuntimeError(f"Failed to encode response for caller {i}: {e}") from e
 
     return output
+
+
+def process_batches_coalesced_combined(
+    items_ipc_bytes: list[bytes],
+    handler: Callable[..., Any],
+    arg_names: list[str] | None,
+    context_metadatas: list[dict[str, Any]],
+) -> list[bytes]:
+    """Coalesced bridge, combined-RecordBatch variant.
+
+    Vertically concatenates every caller's batches into one RecordBatch,
+    builds a CSR-style offsets array marking caller boundaries, and invokes
+    the handler once. The handler returns one RecordBatch with row count
+    equal to the combined input; the framework slices it back into per-caller
+    responses (zero-copy on Arrow buffers).
+
+    Handler signature:
+        (combined: pa.RecordBatch, offsets: list[int], contexts: list[dict])
+            -> pa.RecordBatch
+
+    `offsets` has length `len(contexts) + 1`; caller `i` owns rows
+    `[offsets[i] : offsets[i+1])`.
+
+    Raises:
+        ValueError: empty input, length mismatch, schema mismatch across
+            callers, arg_names length mismatch, or handler output row count
+            doesn't match combined input.
+        TypeError: handler is a generator / async generator (not supported
+            in coalesced modes — each caller expects exactly one response
+            chunk).
+    """
+    if not items_ipc_bytes:
+        raise ValueError("process_batches_coalesced_combined called with no items")
+    if len(items_ipc_bytes) != len(context_metadatas):
+        raise ValueError(
+            f"items_ipc_bytes ({len(items_ipc_bytes)}) and context_metadatas "
+            f"({len(context_metadatas)}) lengths must match"
+        )
+    if inspect.isgeneratorfunction(handler) or inspect.isasyncgenfunction(handler):
+        raise TypeError("Generator/async-generator handlers are not supported when batching is enabled")
+
+    per_caller: list[pa.RecordBatch] = []
+    reference_schema: pa.Schema | None = None
+    for i, ipc_bytes in enumerate(items_ipc_bytes):
+        batches = decode_ipc_stream(ipc_bytes)
+        if not batches:
+            raise ValueError(f"Caller {i} produced no batches")
+        rb = _concat_batches(batches)
+        if reference_schema is None:
+            reference_schema = rb.schema
+        elif not rb.schema.equals(reference_schema):
+            raise ValueError(
+                f"Schema mismatch in coalesced batch: caller 0 has {reference_schema}, caller {i} has {rb.schema}"
+            )
+        per_caller.append(rb)
+
+    assert reference_schema is not None  # guaranteed by the empty-input check above
+
+    offsets: list[int] = [0]
+    for rb in per_caller:
+        offsets.append(offsets[-1] + rb.num_rows)
+
+    table = pa.Table.from_batches(per_caller)
+    combined_batches = table.combine_chunks().to_batches()
+    if combined_batches:
+        combined = combined_batches[0]
+    else:
+        combined = pa.RecordBatch.from_pydict(
+            {name: [] for name in reference_schema.names},
+            schema=reference_schema,
+        )
+
+    if arg_names is not None:
+        if len(arg_names) != combined.num_columns:
+            raise ValueError(
+                f"CHALK_INPUT_ARGS specifies {len(arg_names)} names but "
+                f"combined batch has {combined.num_columns} columns"
+            )
+        combined = combined.rename_columns(arg_names)
+
+    try:
+        result = handler(combined, offsets, context_metadatas)
+        if inspect.iscoroutine(result):
+            result = _run_async(result)
+    except Exception as e:
+        raise RuntimeError(f"Exception raised during handler execution: {e}") from e
+
+    if isinstance(result, pa.RecordBatch):
+        result_rb = result
+    else:
+        try:
+            result_rb = _coerce_to_record_batch(result)
+        except Exception as e:
+            raise RuntimeError(f"Failed to coerce handler output: {e}") from e
+
+    if result_rb.num_rows != combined.num_rows:
+        raise ValueError(
+            f"Combined handler must return a RecordBatch with {combined.num_rows} rows "
+            f"(matching input); got {result_rb.num_rows}"
+        )
+
+    output: list[bytes] = []
+    for i in range(len(per_caller)):
+        start = offsets[i]
+        length = offsets[i + 1] - start
+        sliced = result_rb.slice(start, length)
+        try:
+            output.append(encode_record_batch(sliced))
+        except Exception as e:
+            raise RuntimeError(f"Failed to encode response for caller {i}: {e}") from e
+    return output

{chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.0}/chalk_remote_call_python.egg-info/PKG-INFO

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

chalk_remote_call_python-1.4.0/chalk_remote_call/_version.py

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