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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (63) hide show
  1. {chalk_remote_call_python-1.4.0/chalk_remote_call_python.egg-info → chalk_remote_call_python-1.5.1}/PKG-INFO +3 -2
  2. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk-remote-call-rs/Cargo.lock +10 -12
  3. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk-remote-call-rs/chalk-remote-call-server/Cargo.toml +1 -1
  4. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk-remote-call-rs/chalk-remote-call-server/src/lib.rs +1 -1
  5. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk-remote-call-rs/chalk-remote-call-server/src/python_bridge.rs +7 -6
  6. chalk_remote_call_python-1.5.1/chalk_remote_call/_version.py +1 -0
  7. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/cli.py +12 -0
  8. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/server.py +44 -5
  9. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/servicer.py +154 -39
  10. chalk_remote_call_python-1.5.1/chalk_remote_call/tracing.py +150 -0
  11. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1/chalk_remote_call_python.egg-info}/PKG-INFO +3 -2
  12. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call_python.egg-info/SOURCES.txt +1 -0
  13. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/pyproject.toml +2 -1
  14. chalk_remote_call_python-1.4.0/chalk_remote_call/_version.py +0 -1
  15. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/MANIFEST.in +0 -0
  16. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/README.md +0 -0
  17. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk-remote-call-rs/Cargo.toml +0 -0
  18. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk-remote-call-rs/chalk-remote-call-proto/Cargo.toml +0 -0
  19. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.auth.v1.rs +0 -0
  20. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.common.v1.rs +0 -0
  21. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.runtime.v1.rs +0 -0
  22. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.runtime.v1.tonic.rs +0 -0
  23. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.utils.v1.rs +0 -0
  24. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/descriptor.bin +0 -0
  25. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk-remote-call-rs/chalk-remote-call-proto/src/lib.rs +0 -0
  26. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk-remote-call-rs/chalk-remote-call-server/src/coalesce.rs +0 -0
  27. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk-remote-call-rs/chalk-remote-call-server/src/server.rs +0 -0
  28. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk-remote-call-rs/chalk-remote-call-server/src/service.rs +0 -0
  29. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk-remote-call-rs/rust-toolchain.toml +0 -0
  30. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/__init__.py +0 -0
  31. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/__main__.py +0 -0
  32. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/__init__.py +0 -0
  33. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/__init__.py +0 -0
  34. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/auth/__init__.py +0 -0
  35. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/auth/v1/__init__.py +0 -0
  36. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/auth/v1/permissions_pb2.py +0 -0
  37. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/auth/v1/permissions_pb2_grpc.py +0 -0
  38. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/common/__init__.py +0 -0
  39. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/common/v1/__init__.py +0 -0
  40. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/common/v1/chalk_error_pb2.py +0 -0
  41. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/common/v1/chalk_error_pb2_grpc.py +0 -0
  42. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/runtime/__init__.py +0 -0
  43. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/runtime/v1/__init__.py +0 -0
  44. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/runtime/v1/remote_python_call_pb2.py +0 -0
  45. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/runtime/v1/remote_python_call_pb2_grpc.py +0 -0
  46. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/utils/__init__.py +0 -0
  47. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/utils/v1/__init__.py +0 -0
  48. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/utils/v1/encoding_pb2.py +0 -0
  49. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/utils/v1/encoding_pb2_grpc.py +0 -0
  50. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/utils/v1/field_change_pb2.py +0 -0
  51. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/utils/v1/field_change_pb2_grpc.py +0 -0
  52. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/utils/v1/sensitive_pb2.py +0 -0
  53. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_gen/chalk/utils/v1/sensitive_pb2_grpc.py +0 -0
  54. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/_native.pyi +0 -0
  55. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/arrow_utils.py +0 -0
  56. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/handler_loader.py +0 -0
  57. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call/input_transform.py +0 -0
  58. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call_python.egg-info/dependency_links.txt +0 -0
  59. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call_python.egg-info/entry_points.txt +0 -0
  60. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call_python.egg-info/requires.txt +0 -0
  61. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/chalk_remote_call_python.egg-info/top_level.txt +0 -0
  62. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/setup.cfg +0 -0
  63. {chalk_remote_call_python-1.4.0 → chalk_remote_call_python-1.5.1}/setup.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: chalk-remote-call-python
3
- Version: 1.4.0
3
+ Version: 1.5.1
4
4
  Summary: Chalk remote call Python runtime interface client
5
5
  Author: Chalk AI, Inc.
6
6
  Project-URL: Homepage, https://chalk.ai
@@ -11,6 +11,7 @@ Classifier: Programming Language :: Python :: 3.10
11
11
  Classifier: Programming Language :: Python :: 3.11
12
12
  Classifier: Programming Language :: Python :: 3.12
13
13
  Classifier: Programming Language :: Python :: 3.13
14
+ Classifier: Programming Language :: Python :: 3.14
14
15
  Classifier: Programming Language :: Python
15
16
  Classifier: Typing :: Typed
16
17
  Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
@@ -18,7 +19,7 @@ Classifier: Topic :: Scientific/Engineering :: Information Analysis
18
19
  Classifier: Topic :: Software Development :: Code Generators
19
20
  Classifier: Topic :: Software Development :: Libraries :: Application Frameworks
20
21
  Classifier: Topic :: Software Development :: Libraries :: Python Modules
21
- Requires-Python: <3.14,>=3.10
22
+ Requires-Python: <3.15,>=3.10
22
23
  Description-Content-Type: text/markdown
23
24
  Requires-Dist: pyarrow>=14.0.0
24
25
  Provides-Extra: dev
@@ -1066,11 +1066,10 @@ dependencies = [
1066
1066
 
1067
1067
  [[package]]
1068
1068
  name = "pyo3"
1069
- version = "0.24.2"
1069
+ version = "0.27.1"
1070
1070
  source = "registry+https://github.com/rust-lang/crates.io-index"
1071
- checksum = "e5203598f366b11a02b13aa20cab591229ff0a89fd121a308a5df751d5fc9219"
1071
+ checksum = "37a6df7eab65fc7bee654a421404947e10a0f7085b6951bf2ea395f4659fb0cf"
1072
1072
  dependencies = [
1073
- "cfg-if",
1074
1073
  "indoc",
1075
1074
  "libc",
1076
1075
  "memoffset",
@@ -1084,19 +1083,18 @@ dependencies = [
1084
1083
 
1085
1084
  [[package]]
1086
1085
  name = "pyo3-build-config"
1087
- version = "0.24.2"
1086
+ version = "0.27.1"
1088
1087
  source = "registry+https://github.com/rust-lang/crates.io-index"
1089
- checksum = "99636d423fa2ca130fa5acde3059308006d46f98caac629418e53f7ebb1e9999"
1088
+ checksum = "f77d387774f6f6eec64a004eac0ed525aab7fa1966d94b42f743797b3e395afb"
1090
1089
  dependencies = [
1091
- "once_cell",
1092
1090
  "target-lexicon",
1093
1091
  ]
1094
1092
 
1095
1093
  [[package]]
1096
1094
  name = "pyo3-ffi"
1097
- version = "0.24.2"
1095
+ version = "0.27.1"
1098
1096
  source = "registry+https://github.com/rust-lang/crates.io-index"
1099
- checksum = "78f9cf92ba9c409279bc3305b5409d90db2d2c22392d443a87df3a1adad59e33"
1097
+ checksum = "2dd13844a4242793e02df3e2ec093f540d948299a6a77ea9ce7afd8623f542be"
1100
1098
  dependencies = [
1101
1099
  "libc",
1102
1100
  "pyo3-build-config",
@@ -1104,9 +1102,9 @@ dependencies = [
1104
1102
 
1105
1103
  [[package]]
1106
1104
  name = "pyo3-macros"
1107
- version = "0.24.2"
1105
+ version = "0.27.1"
1108
1106
  source = "registry+https://github.com/rust-lang/crates.io-index"
1109
- checksum = "0b999cb1a6ce21f9a6b147dcf1be9ffedf02e0043aec74dc390f3007047cecd9"
1107
+ checksum = "eaf8f9f1108270b90d3676b8679586385430e5c0bb78bb5f043f95499c821a71"
1110
1108
  dependencies = [
1111
1109
  "proc-macro2",
1112
1110
  "pyo3-macros-backend",
@@ -1116,9 +1114,9 @@ dependencies = [
1116
1114
 
1117
1115
  [[package]]
1118
1116
  name = "pyo3-macros-backend"
1119
- version = "0.24.2"
1117
+ version = "0.27.1"
1120
1118
  source = "registry+https://github.com/rust-lang/crates.io-index"
1121
- checksum = "822ece1c7e1012745607d5cf0bcb2874769f0f7cb34c4cde03b9358eb9ef911a"
1119
+ checksum = "70a3b2274450ba5288bc9b8c1b69ff569d1d61189d4bff38f8d22e03d17f932b"
1122
1120
  dependencies = [
1123
1121
  "heck",
1124
1122
  "proc-macro2",
@@ -9,7 +9,7 @@ crate-type = ["cdylib"]
9
9
 
10
10
  [dependencies]
11
11
  chalk-remote-call-proto = { path = "../chalk-remote-call-proto" }
12
- pyo3 = { version = "0.24", features = ["extension-module"] }
12
+ pyo3 = { version = "0.27", features = ["extension-module"] }
13
13
  tonic = "0.12"
14
14
  tonic-health = "0.12"
15
15
  tonic-reflection = "0.12"
@@ -54,7 +54,7 @@ fn start_server(
54
54
  };
55
55
 
56
56
  // Release the GIL while the Rust server runs
57
- py.allow_threads(|| {
57
+ py.detach(|| {
58
58
  let rt = tokio::runtime::Builder::new_multi_thread()
59
59
  .worker_threads(workers)
60
60
  .enable_all()
@@ -44,11 +44,11 @@ impl PythonHandler {
44
44
  peer: String,
45
45
  ) -> Result<Vec<Vec<u8>>, PythonError> {
46
46
  let (handler, process_fn) =
47
- Python::with_gil(|py| (self.handler.clone_ref(py), self.process_fn.clone_ref(py)));
47
+ Python::attach(|py| (self.handler.clone_ref(py), self.process_fn.clone_ref(py)));
48
48
  let arg_names = self.arg_names.clone();
49
49
 
50
50
  tokio::task::spawn_blocking(move || {
51
- Python::with_gil(|py| -> Result<Vec<Vec<u8>>, PythonError> {
51
+ Python::attach(|py| -> Result<Vec<Vec<u8>>, PythonError> {
52
52
  let py_bytes = PyBytes::new(py, &ipc_bytes).into_any().unbind();
53
53
  let py_ctx = build_context(py, &metadata, &peer)?.into_any().unbind();
54
54
  let py_arg_names = build_arg_names(py, arg_names.as_deref())?;
@@ -67,11 +67,11 @@ impl PythonHandler {
67
67
  callers: Vec<CallerInput>,
68
68
  ) -> Result<Vec<Vec<u8>>, PythonError> {
69
69
  let (handler, process_fn) =
70
- Python::with_gil(|py| (self.handler.clone_ref(py), self.process_fn.clone_ref(py)));
70
+ Python::attach(|py| (self.handler.clone_ref(py), self.process_fn.clone_ref(py)));
71
71
  let arg_names = self.arg_names.clone();
72
72
 
73
73
  tokio::task::spawn_blocking(move || {
74
- Python::with_gil(|py| -> Result<Vec<Vec<u8>>, PythonError> {
74
+ Python::attach(|py| -> Result<Vec<Vec<u8>>, PythonError> {
75
75
  let py_inputs = PyList::empty(py);
76
76
  let py_contexts = PyList::empty(py);
77
77
  for c in &callers {
@@ -152,13 +152,14 @@ fn invoke_process_fn(
152
152
  .map_err(|e| into_python_error(py, e))?;
153
153
 
154
154
  let result_list: &Bound<'_, PyList> = result
155
- .downcast_bound::<PyList>(py)
155
+ .bind(py)
156
+ .cast::<PyList>()
156
157
  .map_err(|e| PythonError(format!("bridge must return a list, got: {e}")))?;
157
158
 
158
159
  let mut out = Vec::with_capacity(result_list.len());
159
160
  for item in result_list.iter() {
160
161
  let bytes_ref: &Bound<'_, PyBytes> = item
161
- .downcast::<PyBytes>()
162
+ .cast::<PyBytes>()
162
163
  .map_err(|e| PythonError(format!("bridge must return list[bytes], got item: {e}")))?;
163
164
  out.push(bytes_ref.as_bytes().to_vec());
164
165
  }
@@ -0,0 +1 @@
1
+ __version__ = "1.5.1"
@@ -80,6 +80,17 @@ def main() -> None:
80
80
  "(env: CHALK_REMOTE_CALL_MAX_BUFFER_DURATION_MS)"
81
81
  ),
82
82
  )
83
+ parser.add_argument(
84
+ "--batching-mode",
85
+ choices=["per_caller", "combined"],
86
+ default=os.environ.get("CHALK_REMOTE_CALL_BATCHING_MODE") or None,
87
+ help=(
88
+ "Coalescing variant when --max-batching-size is set. "
89
+ "'per_caller' (default) passes events as a list of dicts; "
90
+ "'combined' passes a single concatenated RecordBatch with "
91
+ "offsets. (env: CHALK_REMOTE_CALL_BATCHING_MODE)"
92
+ ),
93
+ )
83
94
  parser.add_argument(
84
95
  "--log-level",
85
96
  default="INFO",
@@ -135,4 +146,5 @@ def main() -> None:
135
146
  arg_names=arg_names,
136
147
  max_batching_size=args.max_batching_size,
137
148
  max_buffer_duration_ms=args.max_buffer_duration_ms,
149
+ batching_mode=args.batching_mode,
138
150
  )
@@ -1,14 +1,38 @@
1
1
  from __future__ import annotations
2
2
 
3
3
  import logging
4
+ import os
4
5
  from collections.abc import Callable
5
- from typing import Any
6
+ from typing import Any, Literal
6
7
 
7
8
  from chalk_remote_call._native import start_server
8
- from chalk_remote_call.servicer import process_batches, process_batches_coalesced
9
+ from chalk_remote_call.servicer import (
10
+ process_batches,
11
+ process_batches_coalesced,
12
+ process_batches_coalesced_combined,
13
+ )
9
14
 
10
15
  logger = logging.getLogger(__name__)
11
16
 
17
+ _BATCHING_MODE_ENV = "CHALK_REMOTE_CALL_BATCHING_MODE"
18
+ _VALID_BATCHING_MODES = ("per_caller", "combined")
19
+
20
+
21
+ def _resolve_batching_mode(kwarg: str | None) -> str:
22
+ """Resolve batching mode from kwarg or env var.
23
+
24
+ Precedence: explicit kwarg > env var > default "per_caller". The env var
25
+ is the canonical opt-in for combined mode; the kwarg exists so embedders
26
+ can override programmatically.
27
+ """
28
+ if kwarg is not None:
29
+ mode = kwarg
30
+ else:
31
+ mode = (os.environ.get(_BATCHING_MODE_ENV, "") or "per_caller").strip() or "per_caller"
32
+ if mode not in _VALID_BATCHING_MODES:
33
+ raise ValueError(f"batching_mode must be one of {_VALID_BATCHING_MODES}, got {mode!r}")
34
+ return mode
35
+
12
36
 
13
37
  def serve(
14
38
  handler: Callable[..., Any],
@@ -20,15 +44,21 @@ def serve(
20
44
  arg_names: list[str] | None = None,
21
45
  max_batching_size: int | None = None,
22
46
  max_buffer_duration_ms: int | None = None,
47
+ batching_mode: Literal["per_caller", "combined"] | None = None,
23
48
  ) -> None:
24
49
  """Start the gRPC server implementing RemoteCallService.
25
50
 
26
51
  Args:
27
52
  handler: The user's handler function.
28
53
  Single-request mode (default): (event, context) -> result.
29
- Coalesced mode (when max_batching_size is set):
54
+ Coalesced per-caller mode (when max_batching_size is set and
55
+ batching_mode == "per_caller"):
30
56
  (events: list[dict], contexts: list[dict]) -> list[result],
31
57
  one result per input.
58
+ Coalesced combined mode (when max_batching_size is set and
59
+ batching_mode == "combined"):
60
+ (combined: pa.RecordBatch, offsets: list[int], contexts: list[dict])
61
+ -> pa.RecordBatch with combined.num_rows rows.
32
62
  host: The host to bind to.
33
63
  port: The port to bind to.
34
64
  workers: Number of worker threads for the tokio runtime.
@@ -40,21 +70,30 @@ def serve(
40
70
  max_buffer_duration_ms: Max time in milliseconds to buffer incoming
41
71
  requests before flushing even if max_batching_size hasn't been
42
72
  reached. Defaults to 1000ms when batching is enabled.
73
+ batching_mode: "per_caller" (default) or "combined". Only effective
74
+ when max_batching_size > 0. Overrides the
75
+ CHALK_REMOTE_CALL_BATCHING_MODE env var when set; otherwise the
76
+ env var (or "per_caller") wins.
43
77
  """
44
78
  if on_startup is not None:
45
79
  logger.info("Running startup hook...")
46
80
  on_startup()
47
81
 
82
+ mode = _resolve_batching_mode(batching_mode)
83
+
48
84
  if max_batching_size is not None and max_batching_size > 0:
49
- process_fn = process_batches_coalesced
85
+ process_fn = process_batches_coalesced_combined if mode == "combined" else process_batches_coalesced
50
86
  batch_size = max_batching_size
51
87
  duration_ms = max_buffer_duration_ms if max_buffer_duration_ms is not None else 1000
52
88
  logger.info(
53
- "Batching enabled: max_batching_size=%d max_buffer_duration_ms=%d",
89
+ "Batching enabled: mode=%s max_batching_size=%d max_buffer_duration_ms=%d",
90
+ mode,
54
91
  batch_size,
55
92
  duration_ms,
56
93
  )
57
94
  else:
95
+ if mode == "combined":
96
+ raise ValueError("batching_mode='combined' requires max_batching_size > 0")
58
97
  process_fn = process_batches
59
98
  batch_size = 0
60
99
  duration_ms = 0
@@ -10,6 +10,7 @@ import pyarrow as pa
10
10
 
11
11
  from chalk_remote_call.arrow_utils import decode_ipc_stream, encode_record_batch
12
12
  from chalk_remote_call.input_transform import transform
13
+ from chalk_remote_call.tracing import remote_function_invocation_span
13
14
 
14
15
 
15
16
  def _is_structured(obj: Any) -> bool:
@@ -178,30 +179,31 @@ def process_batches(
178
179
  except Exception as e:
179
180
  raise ValueError(f"Input transformation failed: {e}") from e
180
181
 
181
- try:
182
- result = handler(event, context_metadata)
183
- # If the handler is async, await the coroutine.
184
- if inspect.iscoroutine(result):
185
- result = _run_async(result)
186
- except Exception as e:
187
- raise RuntimeError(f"Exception raised during handler execution: {e}") from e
188
-
189
- try:
190
- if inspect.isgenerator(result):
191
- # Sync generator: encode each yielded value as a separate chunk.
192
- for item in result:
193
- result_batch = _coerce_to_record_batch(item)
194
- results.append(encode_record_batch(result_batch))
195
- elif inspect.isasyncgen(result):
196
- # Async generator: collect items and encode each as a separate chunk.
197
- for item in _collect_async_gen(result):
198
- result_batch = _coerce_to_record_batch(item)
182
+ with remote_function_invocation_span(context_metadata):
183
+ try:
184
+ result = handler(event, context_metadata)
185
+ # If the handler is async, await the coroutine.
186
+ if inspect.iscoroutine(result):
187
+ result = _run_async(result)
188
+ except Exception as e:
189
+ raise RuntimeError(f"Exception raised during handler execution: {e}") from e
190
+
191
+ try:
192
+ if inspect.isgenerator(result):
193
+ # Sync generator: encode each yielded value as a separate chunk.
194
+ for item in result:
195
+ result_batch = _coerce_to_record_batch(item)
196
+ results.append(encode_record_batch(result_batch))
197
+ elif inspect.isasyncgen(result):
198
+ # Async generator: collect items and encode each as a separate chunk.
199
+ for item in _collect_async_gen(result):
200
+ result_batch = _coerce_to_record_batch(item)
201
+ results.append(encode_record_batch(result_batch))
202
+ else:
203
+ result_batch = _coerce_to_record_batch(result)
199
204
  results.append(encode_record_batch(result_batch))
200
- else:
201
- result_batch = _coerce_to_record_batch(result)
202
- results.append(encode_record_batch(result_batch))
203
- except Exception as e:
204
- raise RuntimeError(f"Failed to encode response: {e}") from e
205
+ except Exception as e:
206
+ raise RuntimeError(f"Failed to encode response: {e}") from e
205
207
 
206
208
  return results
207
209
 
@@ -273,25 +275,138 @@ def process_batches_coalesced(
273
275
 
274
276
  # Dispatch. Contexts are passed as a parallel list — handler can correlate
275
277
  # per-caller metadata with each event by index.
276
- try:
277
- result = handler(events, context_metadatas)
278
- if inspect.iscoroutine(result):
279
- result = _run_async(result)
280
- except Exception as e:
281
- raise RuntimeError(f"Exception raised during handler execution: {e}") from e
278
+ with remote_function_invocation_span(context_metadatas, coalesced_count=len(context_metadatas)):
279
+ try:
280
+ result = handler(events, context_metadatas)
281
+ if inspect.iscoroutine(result):
282
+ result = _run_async(result)
283
+ except Exception as e:
284
+ raise RuntimeError(f"Exception raised during handler execution: {e}") from e
285
+
286
+ if not isinstance(result, list):
287
+ raise TypeError(f"Coalesced handler must return a list, got {type(result).__name__}")
288
+ if len(result) != len(events):
289
+ raise ValueError(f"Coalesced handler must return one result per input ({len(events)}); got {len(result)}")
290
+
291
+ # Coerce and encode each per-caller result.
292
+ output: list[bytes] = []
293
+ for i, item in enumerate(result):
294
+ try:
295
+ rb = _coerce_to_record_batch(item)
296
+ output.append(encode_record_batch(rb))
297
+ except Exception as e:
298
+ raise RuntimeError(f"Failed to encode response for caller {i}: {e}") from e
299
+
300
+ return output
282
301
 
283
- if not isinstance(result, list):
284
- raise TypeError(f"Coalesced handler must return a list, got {type(result).__name__}")
285
- if len(result) != len(events):
286
- raise ValueError(f"Coalesced handler must return one result per input ({len(events)}); got {len(result)}")
287
302
 
288
- # Coerce and encode each per-caller result.
289
- output: list[bytes] = []
290
- for i, item in enumerate(result):
303
+ def process_batches_coalesced_combined(
304
+ items_ipc_bytes: list[bytes],
305
+ handler: Callable[..., Any],
306
+ arg_names: list[str] | None,
307
+ context_metadatas: list[dict[str, Any]],
308
+ ) -> list[bytes]:
309
+ """Coalesced bridge, combined-RecordBatch variant.
310
+
311
+ Vertically concatenates every caller's batches into one RecordBatch,
312
+ builds a CSR-style offsets array marking caller boundaries, and invokes
313
+ the handler once. The handler returns one RecordBatch with row count
314
+ equal to the combined input; the framework slices it back into per-caller
315
+ responses (zero-copy on Arrow buffers).
316
+
317
+ Handler signature:
318
+ (combined: pa.RecordBatch, offsets: list[int], contexts: list[dict])
319
+ -> pa.RecordBatch
320
+
321
+ `offsets` has length `len(contexts) + 1`; caller `i` owns rows
322
+ `[offsets[i] : offsets[i+1])`.
323
+
324
+ Raises:
325
+ ValueError: empty input, length mismatch, schema mismatch across
326
+ callers, arg_names length mismatch, or handler output row count
327
+ doesn't match combined input.
328
+ TypeError: handler is a generator / async generator (not supported
329
+ in coalesced modes — each caller expects exactly one response
330
+ chunk).
331
+ """
332
+ if not items_ipc_bytes:
333
+ raise ValueError("process_batches_coalesced_combined called with no items")
334
+ if len(items_ipc_bytes) != len(context_metadatas):
335
+ raise ValueError(
336
+ f"items_ipc_bytes ({len(items_ipc_bytes)}) and context_metadatas "
337
+ f"({len(context_metadatas)}) lengths must match"
338
+ )
339
+ if inspect.isgeneratorfunction(handler) or inspect.isasyncgenfunction(handler):
340
+ raise TypeError("Generator/async-generator handlers are not supported when batching is enabled")
341
+
342
+ per_caller: list[pa.RecordBatch] = []
343
+ reference_schema: pa.Schema | None = None
344
+ for i, ipc_bytes in enumerate(items_ipc_bytes):
345
+ batches = decode_ipc_stream(ipc_bytes)
346
+ if not batches:
347
+ raise ValueError(f"Caller {i} produced no batches")
348
+ rb = _concat_batches(batches)
349
+ if reference_schema is None:
350
+ reference_schema = rb.schema
351
+ elif not rb.schema.equals(reference_schema):
352
+ raise ValueError(
353
+ f"Schema mismatch in coalesced batch: caller 0 has {reference_schema}, caller {i} has {rb.schema}"
354
+ )
355
+ per_caller.append(rb)
356
+
357
+ assert reference_schema is not None # guaranteed by the empty-input check above
358
+
359
+ offsets: list[int] = [0]
360
+ for rb in per_caller:
361
+ offsets.append(offsets[-1] + rb.num_rows)
362
+
363
+ table = pa.Table.from_batches(per_caller)
364
+ combined_batches = table.combine_chunks().to_batches()
365
+ if combined_batches:
366
+ combined = combined_batches[0]
367
+ else:
368
+ combined = pa.RecordBatch.from_pydict(
369
+ {name: [] for name in reference_schema.names},
370
+ schema=reference_schema,
371
+ )
372
+
373
+ if arg_names is not None:
374
+ if len(arg_names) != combined.num_columns:
375
+ raise ValueError(
376
+ f"CHALK_INPUT_ARGS specifies {len(arg_names)} names but "
377
+ f"combined batch has {combined.num_columns} columns"
378
+ )
379
+ combined = combined.rename_columns(arg_names)
380
+
381
+ with remote_function_invocation_span(context_metadatas, coalesced_count=len(context_metadatas)):
291
382
  try:
292
- rb = _coerce_to_record_batch(item)
293
- output.append(encode_record_batch(rb))
383
+ result = handler(combined, offsets, context_metadatas)
384
+ if inspect.iscoroutine(result):
385
+ result = _run_async(result)
294
386
  except Exception as e:
295
- raise RuntimeError(f"Failed to encode response for caller {i}: {e}") from e
387
+ raise RuntimeError(f"Exception raised during handler execution: {e}") from e
388
+
389
+ if isinstance(result, pa.RecordBatch):
390
+ result_rb = result
391
+ else:
392
+ try:
393
+ result_rb = _coerce_to_record_batch(result)
394
+ except Exception as e:
395
+ raise RuntimeError(f"Failed to coerce handler output: {e}") from e
396
+
397
+ if result_rb.num_rows != combined.num_rows:
398
+ raise ValueError(
399
+ f"Combined handler must return a RecordBatch with {combined.num_rows} rows "
400
+ f"(matching input); got {result_rb.num_rows}"
401
+ )
296
402
 
403
+ output: list[bytes] = []
404
+ for i in range(len(per_caller)):
405
+ start = offsets[i]
406
+ length = offsets[i + 1] - start
407
+ sliced = result_rb.slice(start, length)
408
+ try:
409
+ output.append(encode_record_batch(sliced))
410
+ except Exception as e:
411
+ raise RuntimeError(f"Failed to encode response for caller {i}: {e}") from e
297
412
  return output
@@ -0,0 +1,150 @@
1
+ from __future__ import annotations
2
+
3
+ import contextlib
4
+ import importlib
5
+ import os
6
+ from collections.abc import Callable, Iterator, Mapping, Sequence
7
+ from typing import Any, cast
8
+
9
+ _TRACER_NAME = "chalk_remote_call"
10
+ _TRACE_CONTEXT_METADATA_KEYS = ("traceparent", "tracestate", "baggage")
11
+ _REMOTE_FUNCTION_NAME_METADATA = "x-chalk-function-name"
12
+ _REMOTE_FUNCTION_TRACING_ENV_VAR = "CHALK_ENABLE_REMOTE_FUNCTION_TRACING"
13
+ _missing_sdk_tracer = object()
14
+ _missing_otel_modules = object()
15
+ _sdk_get_tracer: Callable[[], Any] | object | None = None
16
+ _otel_modules: tuple[Any, Any, Any] | object | None = None
17
+
18
+
19
+ def _raw_metadata(context_metadata: Any) -> Mapping[str, Any] | None:
20
+ if not isinstance(context_metadata, Mapping):
21
+ return None
22
+
23
+ raw = context_metadata.get("metadata", context_metadata)
24
+ if isinstance(raw, Mapping):
25
+ return raw
26
+ return None
27
+
28
+
29
+ def _first_context_with_traceparent(context_metadata: Any) -> Mapping[str, Any] | None:
30
+ if isinstance(context_metadata, Sequence) and not isinstance(context_metadata, str | bytes):
31
+ for item in context_metadata:
32
+ metadata = _raw_metadata(item)
33
+ if metadata is not None and metadata.get("traceparent"):
34
+ return metadata
35
+ return None
36
+
37
+ metadata = _raw_metadata(context_metadata)
38
+ if metadata is not None and metadata.get("traceparent"):
39
+ return metadata
40
+ return None
41
+
42
+
43
+ def _trace_carrier(metadata: Mapping[str, Any]) -> dict[str, str]:
44
+ return {
45
+ str(key).lower(): str(value)
46
+ for key, value in metadata.items()
47
+ if str(key).lower() in _TRACE_CONTEXT_METADATA_KEYS and value
48
+ }
49
+
50
+
51
+ def _function_name(metadata: Mapping[str, Any]) -> str:
52
+ return str(metadata.get(_REMOTE_FUNCTION_NAME_METADATA) or "") or "unknown"
53
+
54
+
55
+ def _get_tracer(trace_api: Any) -> Any:
56
+ global _sdk_get_tracer
57
+ if _sdk_get_tracer is None:
58
+ try:
59
+ sdk_tracing = importlib.import_module("chalkcompute._tracing")
60
+ except ImportError:
61
+ _sdk_get_tracer = _missing_sdk_tracer
62
+ else:
63
+ _sdk_get_tracer = sdk_tracing.get_tracer
64
+
65
+ sdk_get_tracer = _sdk_get_tracer
66
+ if callable(sdk_get_tracer):
67
+ try:
68
+ return sdk_get_tracer()
69
+ except ImportError:
70
+ return trace_api.get_tracer(_TRACER_NAME)
71
+ return trace_api.get_tracer(_TRACER_NAME)
72
+
73
+
74
+ def _get_otel_modules() -> tuple[Any, Any, Any] | None:
75
+ global _otel_modules
76
+ if _otel_modules is None:
77
+ try:
78
+ _otel_modules = (
79
+ importlib.import_module("opentelemetry.context"),
80
+ importlib.import_module("opentelemetry.propagate"),
81
+ importlib.import_module("opentelemetry.trace"),
82
+ )
83
+ except ImportError:
84
+ _otel_modules = _missing_otel_modules
85
+
86
+ if _otel_modules is _missing_otel_modules:
87
+ return None
88
+ return cast(tuple[Any, Any, Any], _otel_modules)
89
+
90
+
91
+ @contextlib.contextmanager
92
+ def remote_function_invocation_span(
93
+ context_metadata: Any,
94
+ coalesced_count: int | None = None,
95
+ ) -> Iterator[None]:
96
+ if os.environ.get(_REMOTE_FUNCTION_TRACING_ENV_VAR, "") != "1":
97
+ yield
98
+ return
99
+
100
+ metadata = _first_context_with_traceparent(context_metadata)
101
+ if metadata is None:
102
+ yield
103
+ return
104
+
105
+ carrier = _trace_carrier(metadata)
106
+ if not carrier:
107
+ yield
108
+ return
109
+
110
+ otel_modules = _get_otel_modules()
111
+ if otel_modules is None:
112
+ yield
113
+ return
114
+ otel_context, propagate, trace = otel_modules
115
+
116
+ parent_context = propagate.extract(carrier)
117
+ parent_span_context = trace.get_current_span(parent_context).get_span_context()
118
+ if not parent_span_context.is_valid:
119
+ yield
120
+ return
121
+
122
+ SpanKind = trace.SpanKind
123
+ Status = trace.Status
124
+ StatusCode = trace.StatusCode
125
+ function_name = _function_name(metadata)
126
+ tracer = _get_tracer(trace)
127
+ with tracer.start_as_current_span(
128
+ "chalkcompute.remote_function.invoke",
129
+ context=parent_context,
130
+ kind=SpanKind.SERVER,
131
+ ) as span:
132
+ span.set_attribute("chalk.remote_function.name", function_name)
133
+ if coalesced_count is not None:
134
+ span.set_attribute(
135
+ "chalk.remote_function.coalesced_count",
136
+ coalesced_count,
137
+ )
138
+ if not span.get_span_context().is_valid:
139
+ token = otel_context.attach(parent_context)
140
+ try:
141
+ yield
142
+ finally:
143
+ otel_context.detach(token)
144
+ return
145
+ try:
146
+ yield
147
+ except Exception as exc:
148
+ span.record_exception(exc)
149
+ span.set_status(Status(StatusCode.ERROR, str(exc)))
150
+ raise
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: chalk-remote-call-python
3
- Version: 1.4.0
3
+ Version: 1.5.1
4
4
  Summary: Chalk remote call Python runtime interface client
5
5
  Author: Chalk AI, Inc.
6
6
  Project-URL: Homepage, https://chalk.ai
@@ -11,6 +11,7 @@ Classifier: Programming Language :: Python :: 3.10
11
11
  Classifier: Programming Language :: Python :: 3.11
12
12
  Classifier: Programming Language :: Python :: 3.12
13
13
  Classifier: Programming Language :: Python :: 3.13
14
+ Classifier: Programming Language :: Python :: 3.14
14
15
  Classifier: Programming Language :: Python
15
16
  Classifier: Typing :: Typed
16
17
  Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
@@ -18,7 +19,7 @@ Classifier: Topic :: Scientific/Engineering :: Information Analysis
18
19
  Classifier: Topic :: Software Development :: Code Generators
19
20
  Classifier: Topic :: Software Development :: Libraries :: Application Frameworks
20
21
  Classifier: Topic :: Software Development :: Libraries :: Python Modules
21
- Requires-Python: <3.14,>=3.10
22
+ Requires-Python: <3.15,>=3.10
22
23
  Description-Content-Type: text/markdown
23
24
  Requires-Dist: pyarrow>=14.0.0
24
25
  Provides-Extra: dev
@@ -29,6 +29,7 @@ chalk_remote_call/handler_loader.py
29
29
  chalk_remote_call/input_transform.py
30
30
  chalk_remote_call/server.py
31
31
  chalk_remote_call/servicer.py
32
+ chalk_remote_call/tracing.py
32
33
  chalk_remote_call/_gen/__init__.py
33
34
  chalk_remote_call/_gen/chalk/__init__.py
34
35
  chalk_remote_call/_gen/chalk/auth/__init__.py
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "chalk-remote-call-python"
7
- requires-python = ">=3.10,<3.14"
7
+ requires-python = ">=3.10,<3.15"
8
8
  authors = [
9
9
  { name = "Chalk AI, Inc." }
10
10
  ]
@@ -18,6 +18,7 @@ classifiers = [
18
18
  "Programming Language :: Python :: 3.11",
19
19
  "Programming Language :: Python :: 3.12",
20
20
  "Programming Language :: Python :: 3.13",
21
+ "Programming Language :: Python :: 3.14",
21
22
  "Programming Language :: Python",
22
23
  "Typing :: Typed",
23
24
  "Topic :: Scientific/Engineering :: Artificial Intelligence",
@@ -1 +0,0 @@
1
- __version__ = "1.4.0"