chalk-remote-call-python 1.7.0__tar.gz → 1.7.2__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 (65) hide show
  1. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/PKG-INFO +1 -1
  2. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.auth.v1.rs +9 -0
  3. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.runtime.v1.rs +58 -0
  4. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.runtime.v1.tonic.rs +173 -0
  5. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/descriptor.bin +0 -0
  6. chalk_remote_call_python-1.7.2/chalk-remote-call-rs/chalk-remote-call-server/src/async_service.rs +296 -0
  7. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/chalk-remote-call-server/src/lib.rs +1 -0
  8. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/chalk-remote-call-server/src/python_bridge.rs +81 -8
  9. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/chalk-remote-call-server/src/server.rs +13 -0
  10. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/chalk-remote-call-server/src/service.rs +1 -1
  11. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/auth/v1/permissions_pb2.py +9 -3
  12. chalk_remote_call_python-1.7.2/chalk_remote_call/_gen/chalk/runtime/v1/remote_python_call_pb2.py +80 -0
  13. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/runtime/v1/remote_python_call_pb2_grpc.py +68 -0
  14. chalk_remote_call_python-1.7.2/chalk_remote_call/_version.py +1 -0
  15. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call_python.egg-info/PKG-INFO +1 -1
  16. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call_python.egg-info/SOURCES.txt +1 -0
  17. chalk_remote_call_python-1.7.0/chalk_remote_call/_gen/chalk/runtime/v1/remote_python_call_pb2.py +0 -64
  18. chalk_remote_call_python-1.7.0/chalk_remote_call/_version.py +0 -1
  19. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/MANIFEST.in +0 -0
  20. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/README.md +0 -0
  21. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/Cargo.lock +0 -0
  22. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/Cargo.toml +0 -0
  23. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/chalk-remote-call-proto/Cargo.toml +0 -0
  24. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.common.v1.rs +0 -0
  25. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.utils.v1.rs +0 -0
  26. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/chalk-remote-call-proto/src/lib.rs +0 -0
  27. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/chalk-remote-call-server/Cargo.toml +0 -0
  28. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/chalk-remote-call-server/src/coalesce.rs +0 -0
  29. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk-remote-call-rs/rust-toolchain.toml +0 -0
  30. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/__init__.py +0 -0
  31. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/__main__.py +0 -0
  32. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/__init__.py +0 -0
  33. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/__init__.py +0 -0
  34. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/auth/__init__.py +0 -0
  35. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/auth/v1/__init__.py +0 -0
  36. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/auth/v1/permissions_pb2_grpc.py +0 -0
  37. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/common/__init__.py +0 -0
  38. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/common/v1/__init__.py +0 -0
  39. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/common/v1/chalk_error_pb2.py +0 -0
  40. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/common/v1/chalk_error_pb2_grpc.py +0 -0
  41. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/runtime/__init__.py +0 -0
  42. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/runtime/v1/__init__.py +0 -0
  43. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/utils/__init__.py +0 -0
  44. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/utils/v1/__init__.py +0 -0
  45. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/utils/v1/encoding_pb2.py +0 -0
  46. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/utils/v1/encoding_pb2_grpc.py +0 -0
  47. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/utils/v1/field_change_pb2.py +0 -0
  48. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/utils/v1/field_change_pb2_grpc.py +0 -0
  49. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/utils/v1/sensitive_pb2.py +0 -0
  50. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_gen/chalk/utils/v1/sensitive_pb2_grpc.py +0 -0
  51. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/_native.pyi +0 -0
  52. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/arrow_utils.py +0 -0
  53. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/cli.py +0 -0
  54. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/handler_loader.py +0 -0
  55. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/input_transform.py +0 -0
  56. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/server.py +0 -0
  57. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/servicer.py +0 -0
  58. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call/tracing.py +0 -0
  59. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call_python.egg-info/dependency_links.txt +0 -0
  60. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call_python.egg-info/entry_points.txt +0 -0
  61. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call_python.egg-info/requires.txt +0 -0
  62. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/chalk_remote_call_python.egg-info/top_level.txt +0 -0
  63. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/pyproject.toml +0 -0
  64. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/setup.cfg +0 -0
  65. {chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.2}/setup.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: chalk-remote-call-python
3
- Version: 1.7.0
3
+ Version: 1.7.2
4
4
  Summary: Chalk remote call Python runtime interface client
5
5
  Author: Chalk AI, Inc.
6
6
  Project-URL: Homepage, https://chalk.ai
@@ -39,6 +39,9 @@ pub enum Permission {
39
39
  InfrastructureWrite = 31,
40
40
  EnvironmentCreate = 32,
41
41
  InternalDataplaneStatusUpdate = 33,
42
+ InternalWorkingTokenExchange = 34,
43
+ BillingWrite = 35,
44
+ QueryOfflineRead = 36,
42
45
  }
43
46
  impl Permission {
44
47
  /// String value of the enum field names used in the ProtoBuf definition.
@@ -81,6 +84,9 @@ impl Permission {
81
84
  Permission::InfrastructureWrite => "PERMISSION_INFRASTRUCTURE_WRITE",
82
85
  Permission::EnvironmentCreate => "PERMISSION_ENVIRONMENT_CREATE",
83
86
  Permission::InternalDataplaneStatusUpdate => "PERMISSION_INTERNAL_DATAPLANE_STATUS_UPDATE",
87
+ Permission::InternalWorkingTokenExchange => "PERMISSION_INTERNAL_WORKING_TOKEN_EXCHANGE",
88
+ Permission::BillingWrite => "PERMISSION_BILLING_WRITE",
89
+ Permission::QueryOfflineRead => "PERMISSION_QUERY_OFFLINE_READ",
84
90
  }
85
91
  }
86
92
  /// Creates an enum from field names used in the ProtoBuf definition.
@@ -120,6 +126,9 @@ impl Permission {
120
126
  "PERMISSION_INFRASTRUCTURE_WRITE" => Some(Self::InfrastructureWrite),
121
127
  "PERMISSION_ENVIRONMENT_CREATE" => Some(Self::EnvironmentCreate),
122
128
  "PERMISSION_INTERNAL_DATAPLANE_STATUS_UPDATE" => Some(Self::InternalDataplaneStatusUpdate),
129
+ "PERMISSION_INTERNAL_WORKING_TOKEN_EXCHANGE" => Some(Self::InternalWorkingTokenExchange),
130
+ "PERMISSION_BILLING_WRITE" => Some(Self::BillingWrite),
131
+ "PERMISSION_QUERY_OFFLINE_READ" => Some(Self::QueryOfflineRead),
123
132
  _ => None,
124
133
  }
125
134
  }
@@ -86,6 +86,38 @@ pub struct PollRemoteCallResponse {
86
86
  #[prost(message, repeated, tag="4")]
87
87
  pub errors: ::prost::alloc::vec::Vec<super::super::common::v1::ChalkError>,
88
88
  }
89
+ /// Drop pending items from one or all per-function queues for the tenant.
90
+ /// Already-popped (in-flight) calls continue to completion; only items still
91
+ /// queued in Redis are removed.
92
+ #[allow(clippy::derive_partial_eq_without_eq)]
93
+ #[derive(Clone, PartialEq, ::prost::Message)]
94
+ pub struct PurgeQueueRequest {
95
+ #[prost(oneof="purge_queue_request::Request", tags="1, 2")]
96
+ pub request: ::core::option::Option<purge_queue_request::Request>,
97
+ }
98
+ /// Nested message and enum types in `PurgeQueueRequest`.
99
+ pub mod purge_queue_request {
100
+ #[allow(clippy::derive_partial_eq_without_eq)]
101
+ #[derive(Clone, PartialEq, ::prost::Oneof)]
102
+ pub enum Request {
103
+ /// Drop pending items from the named function's queue.
104
+ #[prost(string, tag="1")]
105
+ FunctionName(::prost::alloc::string::String),
106
+ /// If true, drop pending items from every per-function queue for the tenant.
107
+ #[prost(bool, tag="2")]
108
+ All(bool),
109
+ }
110
+ }
111
+ #[allow(clippy::derive_partial_eq_without_eq)]
112
+ #[derive(Clone, PartialEq, ::prost::Message)]
113
+ pub struct PurgeQueueResponse {
114
+ /// Number of pending items dropped per function. For a single-function
115
+ /// purge, contains exactly one entry (or none, if the queue was already
116
+ /// empty). For a tenant-wide purge, contains every queue that had at
117
+ /// least one pending item. The total dropped is `sum(values)`.
118
+ #[prost(map="string, uint64", tag="1")]
119
+ pub items_removed_by_function: ::std::collections::HashMap<::prost::alloc::string::String, u64>,
120
+ }
89
121
  // ── Function Queue Meta Service ───────────────────────────────────
90
122
 
91
123
  /// Information about a single function call.
@@ -130,6 +162,32 @@ pub struct GetRecentCallsResponse {
130
162
  #[prost(message, repeated, tag="1")]
131
163
  pub calls: ::prost::alloc::vec::Vec<FunctionCallInfo>,
132
164
  }
165
+ #[allow(clippy::derive_partial_eq_without_eq)]
166
+ #[derive(Clone, PartialEq, ::prost::Message)]
167
+ pub struct GetCallResultsRequest {
168
+ /// Call IDs to fetch from the function queue.
169
+ #[prost(string, repeated, tag="1")]
170
+ pub call_ids: ::prost::alloc::vec::Vec<::prost::alloc::string::String>,
171
+ }
172
+ #[allow(clippy::derive_partial_eq_without_eq)]
173
+ #[derive(Clone, PartialEq, ::prost::Message)]
174
+ pub struct CallResult {
175
+ /// The requested call ID.
176
+ #[prost(string, tag="1")]
177
+ pub call_id: ::prost::alloc::string::String,
178
+ /// Present when the function queue returned a result for this call.
179
+ #[prost(message, optional, tag="2")]
180
+ pub response: ::core::option::Option<PollRemoteCallResponse>,
181
+ /// Populated when the queue lookup failed for this specific call.
182
+ #[prost(string, tag="3")]
183
+ pub error_message: ::prost::alloc::string::String,
184
+ }
185
+ #[allow(clippy::derive_partial_eq_without_eq)]
186
+ #[derive(Clone, PartialEq, ::prost::Message)]
187
+ pub struct GetCallResultsResponse {
188
+ #[prost(message, repeated, tag="1")]
189
+ pub results: ::prost::alloc::vec::Vec<CallResult>,
190
+ }
133
191
  /// Request the number of calls submitted in the past hour.
134
192
  #[allow(clippy::derive_partial_eq_without_eq)]
135
193
  #[derive(Clone, PartialEq, ::prost::Message)]
@@ -443,6 +443,36 @@ pub mod async_remote_call_service_client {
443
443
  );
444
444
  self.inner.unary(req, path, codec).await
445
445
  }
446
+ pub async fn purge_queue(
447
+ &mut self,
448
+ request: impl tonic::IntoRequest<super::PurgeQueueRequest>,
449
+ ) -> std::result::Result<
450
+ tonic::Response<super::PurgeQueueResponse>,
451
+ tonic::Status,
452
+ > {
453
+ self.inner
454
+ .ready()
455
+ .await
456
+ .map_err(|e| {
457
+ tonic::Status::new(
458
+ tonic::Code::Unknown,
459
+ format!("Service was not ready: {}", e.into()),
460
+ )
461
+ })?;
462
+ let codec = tonic::codec::ProstCodec::default();
463
+ let path = http::uri::PathAndQuery::from_static(
464
+ "/chalk.runtime.v1.AsyncRemoteCallService/PurgeQueue",
465
+ );
466
+ let mut req = request.into_request();
467
+ req.extensions_mut()
468
+ .insert(
469
+ GrpcMethod::new(
470
+ "chalk.runtime.v1.AsyncRemoteCallService",
471
+ "PurgeQueue",
472
+ ),
473
+ );
474
+ self.inner.unary(req, path, codec).await
475
+ }
446
476
  }
447
477
  }
448
478
  /// Generated server implementations.
@@ -466,6 +496,13 @@ pub mod async_remote_call_service_server {
466
496
  tonic::Response<super::PollRemoteCallResponse>,
467
497
  tonic::Status,
468
498
  >;
499
+ async fn purge_queue(
500
+ &self,
501
+ request: tonic::Request<super::PurgeQueueRequest>,
502
+ ) -> std::result::Result<
503
+ tonic::Response<super::PurgeQueueResponse>,
504
+ tonic::Status,
505
+ >;
469
506
  }
470
507
  #[derive(Debug)]
471
508
  pub struct AsyncRemoteCallServiceServer<T: AsyncRemoteCallService> {
@@ -642,6 +679,52 @@ pub mod async_remote_call_service_server {
642
679
  };
643
680
  Box::pin(fut)
644
681
  }
682
+ "/chalk.runtime.v1.AsyncRemoteCallService/PurgeQueue" => {
683
+ #[allow(non_camel_case_types)]
684
+ struct PurgeQueueSvc<T: AsyncRemoteCallService>(pub Arc<T>);
685
+ impl<
686
+ T: AsyncRemoteCallService,
687
+ > tonic::server::UnaryService<super::PurgeQueueRequest>
688
+ for PurgeQueueSvc<T> {
689
+ type Response = super::PurgeQueueResponse;
690
+ type Future = BoxFuture<
691
+ tonic::Response<Self::Response>,
692
+ tonic::Status,
693
+ >;
694
+ fn call(
695
+ &mut self,
696
+ request: tonic::Request<super::PurgeQueueRequest>,
697
+ ) -> Self::Future {
698
+ let inner = Arc::clone(&self.0);
699
+ let fut = async move {
700
+ <T as AsyncRemoteCallService>::purge_queue(&inner, request)
701
+ .await
702
+ };
703
+ Box::pin(fut)
704
+ }
705
+ }
706
+ let accept_compression_encodings = self.accept_compression_encodings;
707
+ let send_compression_encodings = self.send_compression_encodings;
708
+ let max_decoding_message_size = self.max_decoding_message_size;
709
+ let max_encoding_message_size = self.max_encoding_message_size;
710
+ let inner = self.inner.clone();
711
+ let fut = async move {
712
+ let method = PurgeQueueSvc(inner);
713
+ let codec = tonic::codec::ProstCodec::default();
714
+ let mut grpc = tonic::server::Grpc::new(codec)
715
+ .apply_compression_config(
716
+ accept_compression_encodings,
717
+ send_compression_encodings,
718
+ )
719
+ .apply_max_message_size_config(
720
+ max_decoding_message_size,
721
+ max_encoding_message_size,
722
+ );
723
+ let res = grpc.unary(method, req).await;
724
+ Ok(res)
725
+ };
726
+ Box::pin(fut)
727
+ }
645
728
  _ => {
646
729
  Box::pin(async move {
647
730
  Ok(
@@ -797,6 +880,38 @@ pub mod function_queue_meta_service_client {
797
880
  );
798
881
  self.inner.unary(req, path, codec).await
799
882
  }
883
+ /** Return current status and accumulated result chunks for specific calls.
884
+ */
885
+ pub async fn get_call_results(
886
+ &mut self,
887
+ request: impl tonic::IntoRequest<super::GetCallResultsRequest>,
888
+ ) -> std::result::Result<
889
+ tonic::Response<super::GetCallResultsResponse>,
890
+ tonic::Status,
891
+ > {
892
+ self.inner
893
+ .ready()
894
+ .await
895
+ .map_err(|e| {
896
+ tonic::Status::new(
897
+ tonic::Code::Unknown,
898
+ format!("Service was not ready: {}", e.into()),
899
+ )
900
+ })?;
901
+ let codec = tonic::codec::ProstCodec::default();
902
+ let path = http::uri::PathAndQuery::from_static(
903
+ "/chalk.runtime.v1.FunctionQueueMetaService/GetCallResults",
904
+ );
905
+ let mut req = request.into_request();
906
+ req.extensions_mut()
907
+ .insert(
908
+ GrpcMethod::new(
909
+ "chalk.runtime.v1.FunctionQueueMetaService",
910
+ "GetCallResults",
911
+ ),
912
+ );
913
+ self.inner.unary(req, path, codec).await
914
+ }
800
915
  /** Return the number of calls submitted to a function in the past hour.
801
916
  */
802
917
  pub async fn get_call_count(
@@ -847,6 +962,15 @@ pub mod function_queue_meta_service_server {
847
962
  tonic::Response<super::GetRecentCallsResponse>,
848
963
  tonic::Status,
849
964
  >;
965
+ /** Return current status and accumulated result chunks for specific calls.
966
+ */
967
+ async fn get_call_results(
968
+ &self,
969
+ request: tonic::Request<super::GetCallResultsRequest>,
970
+ ) -> std::result::Result<
971
+ tonic::Response<super::GetCallResultsResponse>,
972
+ tonic::Status,
973
+ >;
850
974
  /** Return the number of calls submitted to a function in the past hour.
851
975
  */
852
976
  async fn get_call_count(
@@ -984,6 +1108,55 @@ pub mod function_queue_meta_service_server {
984
1108
  };
985
1109
  Box::pin(fut)
986
1110
  }
1111
+ "/chalk.runtime.v1.FunctionQueueMetaService/GetCallResults" => {
1112
+ #[allow(non_camel_case_types)]
1113
+ struct GetCallResultsSvc<T: FunctionQueueMetaService>(pub Arc<T>);
1114
+ impl<
1115
+ T: FunctionQueueMetaService,
1116
+ > tonic::server::UnaryService<super::GetCallResultsRequest>
1117
+ for GetCallResultsSvc<T> {
1118
+ type Response = super::GetCallResultsResponse;
1119
+ type Future = BoxFuture<
1120
+ tonic::Response<Self::Response>,
1121
+ tonic::Status,
1122
+ >;
1123
+ fn call(
1124
+ &mut self,
1125
+ request: tonic::Request<super::GetCallResultsRequest>,
1126
+ ) -> Self::Future {
1127
+ let inner = Arc::clone(&self.0);
1128
+ let fut = async move {
1129
+ <T as FunctionQueueMetaService>::get_call_results(
1130
+ &inner,
1131
+ request,
1132
+ )
1133
+ .await
1134
+ };
1135
+ Box::pin(fut)
1136
+ }
1137
+ }
1138
+ let accept_compression_encodings = self.accept_compression_encodings;
1139
+ let send_compression_encodings = self.send_compression_encodings;
1140
+ let max_decoding_message_size = self.max_decoding_message_size;
1141
+ let max_encoding_message_size = self.max_encoding_message_size;
1142
+ let inner = self.inner.clone();
1143
+ let fut = async move {
1144
+ let method = GetCallResultsSvc(inner);
1145
+ let codec = tonic::codec::ProstCodec::default();
1146
+ let mut grpc = tonic::server::Grpc::new(codec)
1147
+ .apply_compression_config(
1148
+ accept_compression_encodings,
1149
+ send_compression_encodings,
1150
+ )
1151
+ .apply_max_message_size_config(
1152
+ max_decoding_message_size,
1153
+ max_encoding_message_size,
1154
+ );
1155
+ let res = grpc.unary(method, req).await;
1156
+ Ok(res)
1157
+ };
1158
+ Box::pin(fut)
1159
+ }
987
1160
  "/chalk.runtime.v1.FunctionQueueMetaService/GetCallCount" => {
988
1161
  #[allow(non_camel_case_types)]
989
1162
  struct GetCallCountSvc<T: FunctionQueueMetaService>(pub Arc<T>);
@@ -0,0 +1,296 @@
1
+ use std::collections::HashMap;
2
+ use std::sync::atomic::{AtomicU64, Ordering};
3
+ use std::sync::{Arc, Mutex};
4
+ use std::time::{Duration, Instant};
5
+
6
+ use chalk_remote_call_proto::chalk::common::v1::ChalkError;
7
+ use chalk_remote_call_proto::chalk::runtime::v1::async_remote_call_service_server::AsyncRemoteCallService;
8
+ use chalk_remote_call_proto::chalk::runtime::v1::{
9
+ remote_call_args, EnqueueRemoteCallRequest, EnqueueRemoteCallResponse, PollRemoteCallRequest,
10
+ PollRemoteCallResponse, PurgeQueueRequest, PurgeQueueResponse, RemoteCallStatus,
11
+ };
12
+ use tonic::{Request, Response, Status};
13
+ use tracing::{debug, error, info, instrument};
14
+
15
+ use crate::python_bridge::{ChunkBuffer, PythonHandler};
16
+ use crate::service::extract_metadata;
17
+
18
+ /// How often the orphan sweeper runs.
19
+ const SWEEP_INTERVAL: Duration = Duration::from_secs(30);
20
+ /// How long a *terminal* entry is retained before the sweeper reclaims it. Only
21
+ /// matters for orphans — calls the consumer drained are evicted immediately on
22
+ /// their terminal poll. Comfortably larger than the consumer's poll cadence so
23
+ /// the sweeper never races a live consumer that is about to read a result.
24
+ const RETENTION: Duration = Duration::from_mins(30);
25
+
26
+ /// In-memory execution record for one enqueued call.
27
+ ///
28
+ /// Lives only on the sg pod that accepted the `EnqueueRemoteCall`, and is lost
29
+ /// on pod restart. That is the accepted limitation of the base async path:
30
+ /// arbitrary user handlers cannot be checkpointed, so an sg-pod death loses the
31
+ /// computation regardless. A reclaiming consumer will re-enqueue (and the
32
+ /// function re-runs). Durable reattach across *consumer* restarts is a later
33
+ /// phase and lives outside this struct.
34
+ struct CallState {
35
+ status: RemoteCallStatus,
36
+ /// Result chunks, written directly by the handler's `BufferEmitter` as it
37
+ /// runs and read by `PollRemoteCall` by cursor. Held behind its own lock
38
+ /// (not the registry map lock) so the hot emit path only contends per-call,
39
+ /// and `status`/`error` can be updated without blocking emits.
40
+ chunks: ChunkBuffer,
41
+ /// Populated only when `status == Failed`; surfaced as a `ChalkError`.
42
+ error: Option<String>,
43
+ /// When the call reached a terminal status. Used only by the orphan sweeper
44
+ /// to decide when an un-drained terminal entry is safe to reclaim.
45
+ completed_at: Option<Instant>,
46
+ }
47
+
48
+ impl CallState {
49
+ fn new() -> Self {
50
+ Self {
51
+ status: RemoteCallStatus::Pending,
52
+ chunks: Arc::new(Mutex::new(Vec::new())),
53
+ error: None,
54
+ completed_at: None,
55
+ }
56
+ }
57
+ }
58
+
59
+ type Registry = Arc<Mutex<HashMap<String, CallState>>>;
60
+
61
+ fn is_terminal(status: RemoteCallStatus) -> bool {
62
+ matches!(
63
+ status,
64
+ RemoteCallStatus::Completed | RemoteCallStatus::Failed
65
+ )
66
+ }
67
+
68
+ /// Async (enqueue + poll) executor for the scaling-group hop.
69
+ ///
70
+ /// `EnqueueRemoteCall` runs the handler in a detached background task that
71
+ /// writes result chunks straight into the call's buffer, and returns
72
+ /// immediately; `PollRemoteCall` reads those chunks by cursor. This decouples
73
+ /// function execution from any single connection's lifetime, so the path is
74
+ /// immune to gateway request/stream timeouts — every RPC here is sub-second
75
+ /// regardless of how long the function runs.
76
+ ///
77
+ /// Entries are reclaimed two ways: the terminal poll that delivers a result
78
+ /// evicts its own entry (the consumer drains a call exactly once, then persists
79
+ /// to Redis and stops polling it), and a background sweeper reclaims orphans —
80
+ /// calls that finished but whose consumer died before the terminal poll.
81
+ pub struct AsyncRemoteCallServiceImpl {
82
+ python_handler: Arc<PythonHandler>,
83
+ registry: Registry,
84
+ next_id: AtomicU64,
85
+ /// True when the server was started with coalescing enabled. In that mode
86
+ /// `python_handler`'s bridge expects the batched calling convention, which
87
+ /// the per-call dispatch does not satisfy, so async enqueue is refused
88
+ /// rather than silently calling the handler with the wrong shape.
89
+ batching_enabled: bool,
90
+ }
91
+
92
+ impl AsyncRemoteCallServiceImpl {
93
+ pub fn new(python_handler: Arc<PythonHandler>, batching_enabled: bool) -> Self {
94
+ let registry: Registry = Arc::new(Mutex::new(HashMap::new()));
95
+ // Runs inside the server's tokio runtime (the service is constructed
96
+ // within `run_server`); reclaims orphaned terminal entries.
97
+ spawn_sweeper(registry.clone());
98
+ Self {
99
+ python_handler,
100
+ registry,
101
+ next_id: AtomicU64::new(1),
102
+ batching_enabled,
103
+ }
104
+ }
105
+ }
106
+
107
+ #[tonic::async_trait]
108
+ impl AsyncRemoteCallService for AsyncRemoteCallServiceImpl {
109
+ #[instrument(skip_all)]
110
+ async fn enqueue_remote_call(
111
+ &self,
112
+ request: Request<EnqueueRemoteCallRequest>,
113
+ ) -> Result<Response<EnqueueRemoteCallResponse>, Status> {
114
+ if self.batching_enabled {
115
+ return Err(Status::unimplemented(
116
+ "async enqueue/poll is not supported when request coalescing is enabled",
117
+ ));
118
+ }
119
+
120
+ let metadata = extract_metadata(request.metadata());
121
+ let peer = request
122
+ .remote_addr()
123
+ .map(|a| a.to_string())
124
+ .unwrap_or_default();
125
+ let req = request.into_inner();
126
+ let function_name = req.name;
127
+
128
+ // Only inline feather bytes are supported today; the object-store
129
+ // variant is reserved for large payloads and not wired up here.
130
+ let ipc_bytes = match req.args.and_then(|a| a.args) {
131
+ Some(remote_call_args::Args::FeatherBytes(b)) => b.to_vec(),
132
+ Some(remote_call_args::Args::StorageObjectId(_)) => {
133
+ return Err(Status::unimplemented(
134
+ "storage_object_id args are not supported by the scaling-group executor",
135
+ ));
136
+ }
137
+ None => Vec::new(),
138
+ };
139
+
140
+ // Pod-scoped id: the registry only knows about calls accepted by this
141
+ // pod, so a monotonic per-process counter is sufficient for uniqueness.
142
+ let call_id = format!("sg-{}", self.next_id.fetch_add(1, Ordering::Relaxed));
143
+
144
+ let state = CallState::new();
145
+ // The handler writes chunks into this buffer directly; the registry
146
+ // keeps the same Arc so polls can read them back.
147
+ let buffer = state.chunks.clone();
148
+ self.registry
149
+ .lock()
150
+ .expect("registry mutex poisoned")
151
+ .insert(call_id.clone(), state);
152
+
153
+ info!(call_id = %call_id, function_name = %function_name, "enqueued async call");
154
+
155
+ let handler = self.python_handler.clone();
156
+ let registry = self.registry.clone();
157
+ let id = call_id.clone();
158
+ tokio::spawn(async move {
159
+ set_status(&registry, &id, RemoteCallStatus::Running);
160
+
161
+ // The handler appends each chunk straight into `buffer`; completion
162
+ // is signalled purely by this returning, so there is no channel to
163
+ // drain and nothing that waits on Python object teardown.
164
+ let result = handler
165
+ .call_into_buffer(ipc_bytes, function_name, metadata, peer, buffer)
166
+ .await;
167
+
168
+ let mut reg = registry.lock().expect("registry mutex poisoned");
169
+ if let Some(s) = reg.get_mut(&id) {
170
+ s.completed_at = Some(Instant::now());
171
+ match result {
172
+ Ok(()) => {
173
+ s.status = RemoteCallStatus::Completed;
174
+ let n = s.chunks.lock().map(|b| b.len()).unwrap_or(0);
175
+ info!(call_id = %id, chunks = n, "async call completed");
176
+ }
177
+ Err(e) => {
178
+ s.status = RemoteCallStatus::Failed;
179
+ s.error = Some(e.to_string());
180
+ error!(call_id = %id, error = %e, "async call failed");
181
+ }
182
+ }
183
+ }
184
+ });
185
+
186
+ Ok(Response::new(EnqueueRemoteCallResponse { call_id }))
187
+ }
188
+
189
+ #[instrument(skip_all)]
190
+ async fn poll_remote_call(
191
+ &self,
192
+ request: Request<PollRemoteCallRequest>,
193
+ ) -> Result<Response<PollRemoteCallResponse>, Status> {
194
+ let req = request.into_inner();
195
+ let cursor: usize = if req.cursor.is_empty() {
196
+ 0
197
+ } else {
198
+ req.cursor
199
+ .parse()
200
+ .map_err(|_| Status::invalid_argument(format!("invalid cursor: {}", req.cursor)))?
201
+ };
202
+
203
+ // Read status/error and grab the chunk-buffer handle under the map lock.
204
+ // `status` reaches a terminal value only after the handler returns (all
205
+ // chunks already appended), so seeing `Completed`/`Failed` guarantees
206
+ // the buffer is complete — and this poll delivers it in full. Since the
207
+ // consumer drains a call to terminal exactly once (then persists to
208
+ // Redis and never re-polls this call_id), we evict the entry here. The
209
+ // cloned `chunks` Arc keeps the buffer alive for the read below.
210
+ let (status, error, chunks) = {
211
+ let mut reg = self.registry.lock().expect("registry mutex poisoned");
212
+ let state = reg
213
+ .get(&req.call_id)
214
+ .ok_or_else(|| Status::not_found(format!("unknown call_id: {}", req.call_id)))?;
215
+ let snapshot = (state.status, state.error.clone(), state.chunks.clone());
216
+ if is_terminal(snapshot.0) {
217
+ reg.remove(&req.call_id);
218
+ }
219
+ snapshot
220
+ };
221
+
222
+ let buf = chunks.lock().expect("chunk buffer poisoned");
223
+ // Clamp so a stale/over-large cursor returns an empty suffix rather than
224
+ // panicking on the slice.
225
+ let total = buf.len();
226
+ let from = cursor.min(total);
227
+ let results = buf[from..].to_vec();
228
+ drop(buf);
229
+
230
+ let errors = if status == RemoteCallStatus::Failed {
231
+ vec![ChalkError {
232
+ message: error.unwrap_or_else(|| "remote call failed".to_string()),
233
+ ..Default::default()
234
+ }]
235
+ } else {
236
+ Vec::new()
237
+ };
238
+
239
+ debug!(
240
+ call_id = %req.call_id,
241
+ status = ?status,
242
+ from,
243
+ returned = results.len(),
244
+ "poll"
245
+ );
246
+
247
+ Ok(Response::new(PollRemoteCallResponse {
248
+ status: status as i32,
249
+ results,
250
+ cursor: total.to_string(),
251
+ errors,
252
+ }))
253
+ }
254
+
255
+ async fn purge_queue(
256
+ &self,
257
+ _request: Request<PurgeQueueRequest>,
258
+ ) -> Result<Response<PurgeQueueResponse>, Status> {
259
+ // PurgeQueue operates on the tenant's Redis-backed function queues,
260
+ // which live in the function-queue server — not on an individual
261
+ // scaling-group executor pod.
262
+ Err(Status::unimplemented(
263
+ "PurgeQueue is a function-queue-server operation; not supported on the scaling-group executor",
264
+ ))
265
+ }
266
+ }
267
+
268
+ /// Set the status of an in-flight call, if it still exists in the registry.
269
+ fn set_status(registry: &Registry, call_id: &str, status: RemoteCallStatus) {
270
+ if let Ok(mut reg) = registry.lock() {
271
+ if let Some(s) = reg.get_mut(call_id) {
272
+ s.status = status;
273
+ }
274
+ }
275
+ }
276
+
277
+ /// Background task: periodically reclaim *orphaned* terminal entries — calls
278
+ /// that finished but whose consumer died before issuing the terminal poll that
279
+ /// would normally evict them. Entries the consumer drains are removed on their
280
+ /// terminal poll and never reach the sweeper; non-terminal entries (still
281
+ /// running) are never reclaimed here.
282
+ fn spawn_sweeper(registry: Registry) {
283
+ tokio::spawn(async move {
284
+ let mut ticker = tokio::time::interval(SWEEP_INTERVAL);
285
+ ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
286
+ loop {
287
+ ticker.tick().await;
288
+ let Ok(mut reg) = registry.lock() else {
289
+ continue;
290
+ };
291
+ reg.retain(|_, s| {
292
+ !(is_terminal(s.status) && s.completed_at.is_some_and(|t| t.elapsed() > RETENTION))
293
+ });
294
+ }
295
+ });
296
+ }
@@ -1,3 +1,4 @@
1
+ mod async_service;
1
2
  mod coalesce;
2
3
  mod python_bridge;
3
4
  mod server;