chalk-remote-call-python 1.5.0__tar.gz → 1.6.0__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.5.0 → chalk_remote_call_python-1.6.0}/PKG-INFO +3 -2
  2. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk-remote-call-rs/Cargo.lock +10 -12
  3. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk-remote-call-rs/chalk-remote-call-server/Cargo.toml +1 -1
  4. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk-remote-call-rs/chalk-remote-call-server/src/lib.rs +1 -1
  5. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk-remote-call-rs/chalk-remote-call-server/src/python_bridge.rs +18 -10
  6. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk-remote-call-rs/chalk-remote-call-server/src/service.rs +6 -1
  7. chalk_remote_call_python-1.6.0/chalk_remote_call/_version.py +1 -0
  8. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/servicer.py +71 -67
  9. chalk_remote_call_python-1.6.0/chalk_remote_call/tracing.py +314 -0
  10. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call_python.egg-info/PKG-INFO +3 -2
  11. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call_python.egg-info/SOURCES.txt +1 -0
  12. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/pyproject.toml +2 -1
  13. chalk_remote_call_python-1.5.0/chalk_remote_call/_version.py +0 -1
  14. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/MANIFEST.in +0 -0
  15. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/README.md +0 -0
  16. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk-remote-call-rs/Cargo.toml +0 -0
  17. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk-remote-call-rs/chalk-remote-call-proto/Cargo.toml +0 -0
  18. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.auth.v1.rs +0 -0
  19. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.common.v1.rs +0 -0
  20. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.runtime.v1.rs +0 -0
  21. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.runtime.v1.tonic.rs +0 -0
  22. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.utils.v1.rs +0 -0
  23. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/descriptor.bin +0 -0
  24. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk-remote-call-rs/chalk-remote-call-proto/src/lib.rs +0 -0
  25. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk-remote-call-rs/chalk-remote-call-server/src/coalesce.rs +0 -0
  26. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk-remote-call-rs/chalk-remote-call-server/src/server.rs +0 -0
  27. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk-remote-call-rs/rust-toolchain.toml +0 -0
  28. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/__init__.py +0 -0
  29. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/__main__.py +0 -0
  30. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/__init__.py +0 -0
  31. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/__init__.py +0 -0
  32. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/auth/__init__.py +0 -0
  33. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/auth/v1/__init__.py +0 -0
  34. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/auth/v1/permissions_pb2.py +0 -0
  35. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/auth/v1/permissions_pb2_grpc.py +0 -0
  36. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/common/__init__.py +0 -0
  37. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/common/v1/__init__.py +0 -0
  38. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/common/v1/chalk_error_pb2.py +0 -0
  39. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/common/v1/chalk_error_pb2_grpc.py +0 -0
  40. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/runtime/__init__.py +0 -0
  41. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/runtime/v1/__init__.py +0 -0
  42. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/runtime/v1/remote_python_call_pb2.py +0 -0
  43. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/runtime/v1/remote_python_call_pb2_grpc.py +0 -0
  44. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/utils/__init__.py +0 -0
  45. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/utils/v1/__init__.py +0 -0
  46. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/utils/v1/encoding_pb2.py +0 -0
  47. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/utils/v1/encoding_pb2_grpc.py +0 -0
  48. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/utils/v1/field_change_pb2.py +0 -0
  49. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/utils/v1/field_change_pb2_grpc.py +0 -0
  50. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/utils/v1/sensitive_pb2.py +0 -0
  51. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_gen/chalk/utils/v1/sensitive_pb2_grpc.py +0 -0
  52. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/_native.pyi +0 -0
  53. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/arrow_utils.py +0 -0
  54. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/cli.py +0 -0
  55. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/handler_loader.py +0 -0
  56. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/input_transform.py +0 -0
  57. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call/server.py +0 -0
  58. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call_python.egg-info/dependency_links.txt +0 -0
  59. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call_python.egg-info/entry_points.txt +0 -0
  60. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call_python.egg-info/requires.txt +0 -0
  61. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/chalk_remote_call_python.egg-info/top_level.txt +0 -0
  62. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/setup.cfg +0 -0
  63. {chalk_remote_call_python-1.5.0 → chalk_remote_call_python-1.6.0}/setup.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: chalk-remote-call-python
3
- Version: 1.5.0
3
+ Version: 1.6.0
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()
@@ -15,11 +15,12 @@ pub struct PythonHandler {
15
15
  arg_names: Option<Vec<String>>,
16
16
  }
17
17
 
18
- /// One caller's input: request bytes + gRPC metadata + peer address.
18
+ /// One caller's input: request bytes + function name + gRPC metadata + peer address.
19
19
  /// Used identically in both the single-request and coalesced paths; the
20
20
  /// coalesced path just has N of them.
21
21
  pub struct CallerInput {
22
22
  pub ipc_bytes: Vec<u8>,
23
+ pub function_name: String,
23
24
  pub metadata: HashMap<String, String>,
24
25
  pub peer: String,
25
26
  }
@@ -40,17 +41,20 @@ impl PythonHandler {
40
41
  pub async fn call(
41
42
  &self,
42
43
  ipc_bytes: Vec<u8>,
44
+ function_name: String,
43
45
  metadata: HashMap<String, String>,
44
46
  peer: String,
45
47
  ) -> Result<Vec<Vec<u8>>, PythonError> {
46
48
  let (handler, process_fn) =
47
- Python::with_gil(|py| (self.handler.clone_ref(py), self.process_fn.clone_ref(py)));
49
+ Python::attach(|py| (self.handler.clone_ref(py), self.process_fn.clone_ref(py)));
48
50
  let arg_names = self.arg_names.clone();
49
51
 
50
52
  tokio::task::spawn_blocking(move || {
51
- Python::with_gil(|py| -> Result<Vec<Vec<u8>>, PythonError> {
53
+ Python::attach(|py| -> Result<Vec<Vec<u8>>, PythonError> {
52
54
  let py_bytes = PyBytes::new(py, &ipc_bytes).into_any().unbind();
53
- let py_ctx = build_context(py, &metadata, &peer)?.into_any().unbind();
55
+ let py_ctx = build_context(py, &function_name, &metadata, &peer)?
56
+ .into_any()
57
+ .unbind();
54
58
  let py_arg_names = build_arg_names(py, arg_names.as_deref())?;
55
59
  invoke_process_fn(py, &process_fn, &handler, py_bytes, py_arg_names, py_ctx)
56
60
  })
@@ -67,11 +71,11 @@ impl PythonHandler {
67
71
  callers: Vec<CallerInput>,
68
72
  ) -> Result<Vec<Vec<u8>>, PythonError> {
69
73
  let (handler, process_fn) =
70
- Python::with_gil(|py| (self.handler.clone_ref(py), self.process_fn.clone_ref(py)));
74
+ Python::attach(|py| (self.handler.clone_ref(py), self.process_fn.clone_ref(py)));
71
75
  let arg_names = self.arg_names.clone();
72
76
 
73
77
  tokio::task::spawn_blocking(move || {
74
- Python::with_gil(|py| -> Result<Vec<Vec<u8>>, PythonError> {
78
+ Python::attach(|py| -> Result<Vec<Vec<u8>>, PythonError> {
75
79
  let py_inputs = PyList::empty(py);
76
80
  let py_contexts = PyList::empty(py);
77
81
  for c in &callers {
@@ -79,7 +83,7 @@ impl PythonHandler {
79
83
  .append(PyBytes::new(py, &c.ipc_bytes))
80
84
  .map_err(|e| into_python_error(py, e))?;
81
85
  py_contexts
82
- .append(build_context(py, &c.metadata, &c.peer)?)
86
+ .append(build_context(py, &c.function_name, &c.metadata, &c.peer)?)
83
87
  .map_err(|e| into_python_error(py, e))?;
84
88
  }
85
89
  let py_arg_names = build_arg_names(py, arg_names.as_deref())?;
@@ -100,9 +104,10 @@ impl PythonHandler {
100
104
 
101
105
  // ---- Shared helpers -------------------------------------------------------
102
106
 
103
- /// Build a `{"peer": str, "metadata": dict}` context dict for one caller.
107
+ /// Build a `{"function_name": str, "peer": str, "metadata": dict}` context dict for one caller.
104
108
  fn build_context<'py>(
105
109
  py: Python<'py>,
110
+ function_name: &str,
106
111
  metadata: &HashMap<String, String>,
107
112
  peer: &str,
108
113
  ) -> Result<Bound<'py, PyDict>, PythonError> {
@@ -113,6 +118,8 @@ fn build_context<'py>(
113
118
  .map_err(|e| into_python_error(py, e))?;
114
119
  }
115
120
  let ctx = PyDict::new(py);
121
+ ctx.set_item("function_name", function_name)
122
+ .map_err(|e| into_python_error(py, e))?;
116
123
  ctx.set_item("peer", peer)
117
124
  .map_err(|e| into_python_error(py, e))?;
118
125
  ctx.set_item("metadata", meta_dict)
@@ -152,13 +159,14 @@ fn invoke_process_fn(
152
159
  .map_err(|e| into_python_error(py, e))?;
153
160
 
154
161
  let result_list: &Bound<'_, PyList> = result
155
- .downcast_bound::<PyList>(py)
162
+ .bind(py)
163
+ .cast::<PyList>()
156
164
  .map_err(|e| PythonError(format!("bridge must return a list, got: {e}")))?;
157
165
 
158
166
  let mut out = Vec::with_capacity(result_list.len());
159
167
  for item in result_list.iter() {
160
168
  let bytes_ref: &Bound<'_, PyBytes> = item
161
- .downcast::<PyBytes>()
169
+ .cast::<PyBytes>()
162
170
  .map_err(|e| PythonError(format!("bridge must return list[bytes], got item: {e}")))?;
163
171
  out.push(bytes_ref.as_bytes().to_vec());
164
172
  }
@@ -40,7 +40,11 @@ impl RemoteCallService for RemoteCallServiceImpl {
40
40
 
41
41
  // Accumulate all feather_stream bytes from the request stream
42
42
  let mut all_bytes: Vec<u8> = Vec::new();
43
+ let mut function_name = String::new();
43
44
  while let Some(req) = stream.message().await? {
45
+ if function_name.is_empty() && !req.name.is_empty() {
46
+ function_name = req.name.clone();
47
+ }
44
48
  all_bytes.extend_from_slice(&req.feather_stream);
45
49
  if all_bytes.len() > MAX_REQUEST_SIZE {
46
50
  return Err(Status::resource_exhausted(format!(
@@ -74,6 +78,7 @@ impl RemoteCallService for RemoteCallServiceImpl {
74
78
  .submit(BufferedCall {
75
79
  input: CallerInput {
76
80
  ipc_bytes: all_bytes,
81
+ function_name,
77
82
  metadata,
78
83
  peer,
79
84
  },
@@ -86,7 +91,7 @@ impl RemoteCallService for RemoteCallServiceImpl {
86
91
  // Single-request path.
87
92
  let handler = self.python_handler.clone();
88
93
  tokio::spawn(async move {
89
- match handler.call(all_bytes, metadata, peer).await {
94
+ match handler.call(all_bytes, function_name, metadata, peer).await {
90
95
  Ok(response_chunks) => {
91
96
  for chunk in response_chunks {
92
97
  let msg = CallFunctionResponse {
@@ -0,0 +1 @@
1
+ __version__ = "1.6.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,26 +275,27 @@ 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
282
-
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
-
288
- # Coerce and encode each per-caller result.
289
- output: list[bytes] = []
290
- for i, item in enumerate(result):
278
+ with remote_function_invocation_span(context_metadatas, coalesced_count=len(context_metadatas)):
291
279
  try:
292
- rb = _coerce_to_record_batch(item)
293
- output.append(encode_record_batch(rb))
280
+ result = handler(events, context_metadatas)
281
+ if inspect.iscoroutine(result):
282
+ result = _run_async(result)
294
283
  except Exception as e:
295
- raise RuntimeError(f"Failed to encode response for caller {i}: {e}") from 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
296
299
 
297
300
  return output
298
301
 
@@ -375,34 +378,35 @@ def process_batches_coalesced_combined(
375
378
  )
376
379
  combined = combined.rename_columns(arg_names)
377
380
 
378
- try:
379
- result = handler(combined, offsets, context_metadatas)
380
- if inspect.iscoroutine(result):
381
- result = _run_async(result)
382
- except Exception as e:
383
- raise RuntimeError(f"Exception raised during handler execution: {e}") from e
384
-
385
- if isinstance(result, pa.RecordBatch):
386
- result_rb = result
387
- else:
381
+ with remote_function_invocation_span(context_metadatas, coalesced_count=len(context_metadatas)):
388
382
  try:
389
- result_rb = _coerce_to_record_batch(result)
383
+ result = handler(combined, offsets, context_metadatas)
384
+ if inspect.iscoroutine(result):
385
+ result = _run_async(result)
390
386
  except Exception as e:
391
- raise RuntimeError(f"Failed to coerce handler output: {e}") from e
387
+ raise RuntimeError(f"Exception raised during handler execution: {e}") from e
392
388
 
393
- if result_rb.num_rows != combined.num_rows:
394
- raise ValueError(
395
- f"Combined handler must return a RecordBatch with {combined.num_rows} rows "
396
- f"(matching input); got {result_rb.num_rows}"
397
- )
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
398
396
 
399
- output: list[bytes] = []
400
- for i in range(len(per_caller)):
401
- start = offsets[i]
402
- length = offsets[i + 1] - start
403
- sliced = result_rb.slice(start, length)
404
- try:
405
- output.append(encode_record_batch(sliced))
406
- except Exception as e:
407
- raise RuntimeError(f"Failed to encode response for caller {i}: {e}") from e
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
+ )
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
408
412
  return output
@@ -0,0 +1,314 @@
1
+ from __future__ import annotations
2
+
3
+ import contextlib
4
+ import importlib
5
+ import os
6
+ from collections.abc import 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_TRACE_POLICY_ENV_VAR = "CHALK_REMOTE_FUNCTION_TRACE_POLICY"
12
+ _REMOTE_FUNCTION_TRACE_SAMPLE_RATE_ENV_VAR = "CHALK_REMOTE_FUNCTION_TRACE_SAMPLE_RATE"
13
+ _TRACE_POLICY_PARENT_BASED_ALWAYS_OFF = "parentbased_always_off"
14
+ _TRACE_POLICY_PARENT_BASED_TRACE_ID_RATIO = "parentbased_traceidratio"
15
+ _TRACE_POLICY_ALWAYS_OFF = "always_off"
16
+ _TRACE_POLICIES = frozenset(
17
+ {
18
+ _TRACE_POLICY_PARENT_BASED_ALWAYS_OFF,
19
+ _TRACE_POLICY_PARENT_BASED_TRACE_ID_RATIO,
20
+ _TRACE_POLICY_ALWAYS_OFF,
21
+ }
22
+ )
23
+ _missing_otel_modules = object()
24
+ _missing_always_off_tracer = object()
25
+ _otel_modules: tuple[Any, Any, Any] | object | None = None
26
+ _always_off_tracer: Any | object | None = None
27
+ _runtime_tracing_configured = False
28
+ _runtime_tracer_provider: Any | None = None
29
+
30
+
31
+ def _raw_metadata(context_metadata: Any) -> Mapping[str, Any] | None:
32
+ if not isinstance(context_metadata, Mapping):
33
+ return None
34
+
35
+ raw = context_metadata.get("metadata", context_metadata)
36
+ if isinstance(raw, Mapping):
37
+ return raw
38
+ return None
39
+
40
+
41
+ def _collect_trace_metadata(
42
+ context_metadata: Any,
43
+ ) -> tuple[Mapping[str, Any] | None, Mapping[str, Any] | None]:
44
+ if isinstance(context_metadata, Mapping):
45
+ first = _raw_metadata(context_metadata)
46
+ if first is not None and first.get("traceparent"):
47
+ return first, first
48
+ return first, None
49
+
50
+ first: Mapping[str, Any] | None = None
51
+ if isinstance(context_metadata, Sequence) and not isinstance(context_metadata, str | bytes):
52
+ for item in context_metadata:
53
+ metadata = _raw_metadata(item)
54
+ if metadata is None:
55
+ continue
56
+
57
+ if first is None:
58
+ first = metadata
59
+ if metadata.get("traceparent"):
60
+ return first, metadata
61
+
62
+ return first, None
63
+
64
+
65
+ def _trace_policy() -> str | None:
66
+ raw_policy = os.environ.get(_REMOTE_FUNCTION_TRACE_POLICY_ENV_VAR)
67
+ if raw_policy is None:
68
+ return None
69
+
70
+ policy = raw_policy.strip().lower()
71
+ if policy in _TRACE_POLICIES:
72
+ return policy
73
+ return None
74
+
75
+
76
+ def _trace_sample_rate() -> float | None:
77
+ raw_rate = os.environ.get(_REMOTE_FUNCTION_TRACE_SAMPLE_RATE_ENV_VAR)
78
+ if raw_rate is None:
79
+ return None
80
+
81
+ try:
82
+ rate = float(raw_rate)
83
+ except ValueError:
84
+ return None
85
+
86
+ if rate < 0.0 or rate > 1.0:
87
+ return None
88
+ return rate
89
+
90
+
91
+ def _otel_exporter_configured() -> bool:
92
+ return bool(os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") or os.environ.get("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"))
93
+
94
+
95
+ def _configure_runtime_tracing(sample_rate: float) -> None:
96
+ global _runtime_tracing_configured, _runtime_tracer_provider
97
+ if _runtime_tracing_configured:
98
+ return
99
+
100
+ try:
101
+ resources = importlib.import_module("opentelemetry.sdk.resources")
102
+ sdk_trace = importlib.import_module("opentelemetry.sdk.trace")
103
+ sdk_sampling = importlib.import_module("opentelemetry.sdk.trace.sampling")
104
+ trace_api = importlib.import_module("opentelemetry.trace")
105
+ except Exception:
106
+ _runtime_tracing_configured = True
107
+ return
108
+
109
+ try:
110
+ provider = sdk_trace.TracerProvider(
111
+ resource=resources.Resource.create({"service.name": os.environ.get("CHALK_SERVICE") or _TRACER_NAME}),
112
+ sampler=sdk_sampling.ParentBased(sdk_sampling.TraceIdRatioBased(sample_rate)),
113
+ )
114
+
115
+ if _otel_exporter_configured():
116
+ try:
117
+ otlp_exporter = importlib.import_module("opentelemetry.exporter.otlp.proto.grpc.trace_exporter")
118
+ sdk_export = importlib.import_module("opentelemetry.sdk.trace.export")
119
+ provider.add_span_processor(sdk_export.BatchSpanProcessor(otlp_exporter.OTLPSpanExporter()))
120
+ except Exception:
121
+ pass
122
+
123
+ trace_api.set_tracer_provider(provider)
124
+ _runtime_tracer_provider = provider
125
+ except Exception:
126
+ pass
127
+ _runtime_tracing_configured = True
128
+
129
+
130
+ def _span_metadata(
131
+ first_metadata: Mapping[str, Any] | None,
132
+ traceparent_metadata: Mapping[str, Any] | None,
133
+ policy: str,
134
+ ) -> tuple[Mapping[str, Any] | None, Any | None, Any]:
135
+ if traceparent_metadata is not None:
136
+ carrier = _trace_carrier(traceparent_metadata)
137
+ if not carrier:
138
+ return None, None, None
139
+
140
+ otel_modules = _get_otel_modules()
141
+ if otel_modules is None:
142
+ return None, None, None
143
+ _, propagate, trace = otel_modules
144
+
145
+ parent_context = propagate.extract(carrier)
146
+ parent_span_context = trace.get_current_span(parent_context).get_span_context()
147
+ if not parent_span_context.is_valid:
148
+ return None, None, None
149
+ return traceparent_metadata, parent_context, otel_modules
150
+
151
+ if policy != _TRACE_POLICY_PARENT_BASED_TRACE_ID_RATIO:
152
+ return None, None, None
153
+
154
+ sample_rate = _trace_sample_rate()
155
+ if sample_rate is None:
156
+ return None, None, None
157
+
158
+ _configure_runtime_tracing(sample_rate)
159
+ otel_modules = _get_otel_modules()
160
+ if otel_modules is None:
161
+ return None, None, None
162
+ return first_metadata, None, otel_modules
163
+
164
+
165
+ def _always_off_span_metadata(
166
+ first_metadata: Mapping[str, Any] | None,
167
+ traceparent_metadata: Mapping[str, Any] | None,
168
+ ) -> tuple[Mapping[str, Any] | None, Any | None, Any | None, Any | None]:
169
+ otel_modules = _get_otel_modules()
170
+ if otel_modules is None:
171
+ return None, None, None, None
172
+ _, propagate, trace = otel_modules
173
+
174
+ parent_context = None
175
+ if traceparent_metadata is not None:
176
+ carrier = _trace_carrier(traceparent_metadata)
177
+ if carrier:
178
+ extracted_context = propagate.extract(carrier)
179
+ parent_span_context = trace.get_current_span(extracted_context).get_span_context()
180
+ if parent_span_context.is_valid:
181
+ parent_context = extracted_context
182
+
183
+ tracer = _get_always_off_tracer()
184
+ if tracer is None:
185
+ return None, None, None, None
186
+ return traceparent_metadata or first_metadata, parent_context, otel_modules, tracer
187
+
188
+
189
+ def _trace_carrier(metadata: Mapping[str, Any]) -> dict[str, str]:
190
+ return {
191
+ str(key).lower(): str(value)
192
+ for key, value in metadata.items()
193
+ if str(key).lower() in _TRACE_CONTEXT_METADATA_KEYS and value
194
+ }
195
+
196
+
197
+ def _function_name(context_metadata: Any) -> str:
198
+ if isinstance(context_metadata, Mapping):
199
+ return str(context_metadata.get("function_name") or "") or "unknown"
200
+
201
+ if isinstance(context_metadata, Sequence) and not isinstance(context_metadata, str | bytes):
202
+ for item in context_metadata:
203
+ if isinstance(item, Mapping):
204
+ function_name = str(item.get("function_name") or "")
205
+ if function_name:
206
+ return function_name
207
+
208
+ return "unknown"
209
+
210
+
211
+ def _get_tracer(trace_api: Any) -> Any:
212
+ return trace_api.get_tracer(_TRACER_NAME)
213
+
214
+
215
+ def _get_otel_modules() -> tuple[Any, Any, Any] | None:
216
+ global _otel_modules
217
+ if _otel_modules is None:
218
+ try:
219
+ _otel_modules = (
220
+ importlib.import_module("opentelemetry.context"),
221
+ importlib.import_module("opentelemetry.propagate"),
222
+ importlib.import_module("opentelemetry.trace"),
223
+ )
224
+ except ImportError:
225
+ _otel_modules = _missing_otel_modules
226
+
227
+ if _otel_modules is _missing_otel_modules:
228
+ return None
229
+ return cast(tuple[Any, Any, Any], _otel_modules)
230
+
231
+
232
+ def _get_always_off_tracer() -> Any | None:
233
+ global _always_off_tracer
234
+ if _always_off_tracer is None:
235
+ try:
236
+ sdk_trace = importlib.import_module("opentelemetry.sdk.trace")
237
+ sdk_sampling = importlib.import_module("opentelemetry.sdk.trace.sampling")
238
+ except ImportError:
239
+ _always_off_tracer = _missing_always_off_tracer
240
+ else:
241
+ # Keep always-off local to this invocation instead of replacing the
242
+ # process-wide tracer provider.
243
+ provider = sdk_trace.TracerProvider(sampler=sdk_sampling.ALWAYS_OFF)
244
+ _always_off_tracer = provider.get_tracer(_TRACER_NAME)
245
+
246
+ if _always_off_tracer is _missing_always_off_tracer:
247
+ return None
248
+ return _always_off_tracer
249
+
250
+
251
+ @contextlib.contextmanager
252
+ def remote_function_invocation_span(
253
+ context_metadata: Any,
254
+ coalesced_count: int | None = None,
255
+ ) -> Iterator[None]:
256
+ policy = _trace_policy()
257
+ if policy is None:
258
+ yield
259
+ return
260
+
261
+ first_metadata, traceparent_metadata = _collect_trace_metadata(context_metadata)
262
+ if policy == _TRACE_POLICY_PARENT_BASED_ALWAYS_OFF and traceparent_metadata is None:
263
+ yield
264
+ return
265
+
266
+ if policy == _TRACE_POLICY_ALWAYS_OFF:
267
+ _metadata, parent_context, otel_modules, tracer = _always_off_span_metadata(
268
+ first_metadata,
269
+ traceparent_metadata,
270
+ )
271
+ else:
272
+ _metadata, parent_context, otel_modules = _span_metadata(
273
+ first_metadata,
274
+ traceparent_metadata,
275
+ policy,
276
+ )
277
+ tracer = _get_tracer(otel_modules[2]) if otel_modules is not None else None
278
+
279
+ if otel_modules is None:
280
+ yield
281
+ return
282
+ if tracer is None:
283
+ yield
284
+ return
285
+ otel_context, _, trace = otel_modules
286
+
287
+ SpanKind = trace.SpanKind
288
+ Status = trace.Status
289
+ StatusCode = trace.StatusCode
290
+ function_name = _function_name(context_metadata)
291
+ span_kwargs: dict[str, Any] = {"kind": SpanKind.SERVER}
292
+ if parent_context is not None:
293
+ span_kwargs["context"] = parent_context
294
+
295
+ with tracer.start_as_current_span("chalkcompute.remote_function.invoke", **span_kwargs) as span:
296
+ span.set_attribute("chalk.remote_function.name", function_name)
297
+ if coalesced_count is not None:
298
+ span.set_attribute(
299
+ "chalk.remote_function.coalesced_count",
300
+ coalesced_count,
301
+ )
302
+ if parent_context is not None and not span.get_span_context().is_valid:
303
+ token = otel_context.attach(parent_context)
304
+ try:
305
+ yield
306
+ finally:
307
+ otel_context.detach(token)
308
+ return
309
+ try:
310
+ yield
311
+ except Exception as exc:
312
+ span.record_exception(exc)
313
+ span.set_status(Status(StatusCode.ERROR, str(exc)))
314
+ raise
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: chalk-remote-call-python
3
- Version: 1.5.0
3
+ Version: 1.6.0
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.5.0"