diff --git a/Cargo.lock b/Cargo.lock index fa6e9346c..26b3c1b3a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4336,7 +4336,7 @@ dependencies = [ [[package]] name = "pulsing-actor" -version = "0.1.2" +version = "0.1.3" dependencies = [ "anyhow", "async-trait", @@ -4403,7 +4403,7 @@ dependencies = [ [[package]] name = "pulsing-bench-py" -version = "0.1.2" +version = "0.1.3" dependencies = [ "pulsing-bench", "pyo3", @@ -4415,7 +4415,7 @@ dependencies = [ [[package]] name = "pulsing-bindings-core" -version = "0.1.2" +version = "0.1.3" dependencies = [ "bincode", "futures", @@ -4426,7 +4426,7 @@ dependencies = [ [[package]] name = "pulsing-cli" -version = "0.1.2" +version = "0.1.3" dependencies = [ "anyhow", "clap", @@ -4443,7 +4443,7 @@ dependencies = [ [[package]] name = "pulsing-forge" -version = "0.1.2" +version = "0.1.3" dependencies = [ "anyhow", "base64 0.22.1", @@ -4477,7 +4477,7 @@ dependencies = [ [[package]] name = "pulsing-gui" -version = "0.1.2" +version = "0.1.3" dependencies = [ "anyhow", "eframe", @@ -4489,11 +4489,12 @@ dependencies = [ [[package]] name = "pulsing-py" -version = "0.1.2" +version = "0.1.3" dependencies = [ "anyhow", "async-trait", "bincode", + "bytes", "crossbeam-channel", "futures", "pulsing-actor", @@ -4512,7 +4513,7 @@ dependencies = [ [[package]] name = "pulsing-rpymod" -version = "0.1.2" +version = "0.1.3" dependencies = [ "anyhow", "async-trait", @@ -4529,7 +4530,7 @@ dependencies = [ [[package]] name = "pulsing-workspace" -version = "0.1.2" +version = "0.1.3" dependencies = [ "anyhow", "chrono", diff --git a/Cargo.toml b/Cargo.toml index b5335e63b..de3d51da6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,7 +28,7 @@ resolver = "3" unexpected_cfgs = { level = "warn", check-cfg = ['cfg(coverage_nightly)'] } [workspace.package] -version = "0.1.2" +version = "0.1.3" edition = "2021" description = "Pulsing - Distributed Actor Framework" authors = ["Reiase "] diff --git a/README.md b/README.md index 55613fa4b..f0d09107f 100644 --- a/README.md +++ b/README.md @@ -188,6 +188,35 @@ from pulsing.langgraph import with_pulsing distributed_app = with_pulsing(app, seeds=["gpu-server:8001"]) ``` +### 5. Fast TensorDict Transport + +`TensorMessage` carries opaque metadata and ordered, contiguous CPU buffers +without putting tensor payloads in the normal pickle envelope. The caller (for +example PulsingQueue) moves CUDA tensors to CPU, makes them contiguous, and +encodes dtype, shape, byte order, and TensorDict structure in metadata. Pulsing +only transports those bytes and buffers. + +```python +from array import array + +import pulsing as pul + +cpu_buffer = array("f", range(6)) +message = pul.TensorMessage( + metadata=b"...", # dtype, shape, byte order, and TensorDict structure + buffers=[memoryview(cpu_buffer).cast("B")], + version=1, +) +``` + +Clear-text remote connections use a pooled raw TCP path by default. It sends +the header, metadata, and original buffers with vectored I/O and reads every +payload directly into its final receive allocation. TLS and +`PULSING_TENSOR_TRANSPORT=http2` use the packed HTTP/2 compatibility path. + +See the [complete transport design](docs/src/design/tensor-message-transport.md) +and the [runnable TCP example](examples/python/tensor_message_fast_path.py). + ## 📚 Example Guide ``` @@ -203,6 +232,7 @@ examples/ ├── python/ # ⭐⭐ Basic examples │ ├── ping_pong.py # Actor basics │ ├── cluster.py # Cluster communication +│ ├── tensor_message_fast_path.py # Tensor transport │ └── ... └── rust/ # Rust examples ``` diff --git a/README.zh.md b/README.zh.md index 5f069d410..d4e1a8a0c 100644 --- a/README.zh.md +++ b/README.zh.md @@ -188,6 +188,54 @@ from pulsing.langgraph import with_pulsing distributed_app = with_pulsing(app, seeds=["gpu-server:8001"]) ``` +### 5. TensorDict 快速传输 + +`TensorMessage` 用一段不透明 metadata 和多段连续 CPU buffer 传输 TensorDict。PulsingQueue +负责把 CUDA Tensor 搬到 CPU，并在交给 Pulsing 前保证每个 Tensor 连续；Pulsing 不解析 +Tensor 的 dtype、shape 或 stride，这些信息由 PulsingQueue 编码到 metadata。 + +```python +from array import array + +import pulsing as pul + +cpu_buffer = array("f", range(6)) +message = pul.TensorMessage( + metadata=b"...", # dtype / shape / TensorDict 结构 + buffers=[memoryview(cpu_buffer).cast("B")], # 必须连续 + version=1, +) + +# @remote Actor 收到 TensorMessage 时调用该专用入口。 +@pul.remote +class Storage: + async def receive_tensor(self, message: pul.TensorMessage): + return message +``` + +明文远程连接默认使用同一监听端口上的 raw TCP 数据面：长连接池复用连接，发送端用 +`write_vectored` 直接提交 header、metadata 和各原始 buffer，接收端把每个 payload +直接读入最终由 Python Tensor 持有的分配。按应用 payload 的 CPU copy 口径，主路径只有 +`应用 buffer -> TCP kernel buffer` 和 `TCP kernel buffer -> 最终接收 buffer` 两次， +不会先拼成一个大包，也不会在接收后再复制一次。远程调用的 `ask`/`tell` 完成前，发送方 +不得修改或 resize 源 buffer；完成后远端已经持有独立接收分配。若目标 actor 位于同一进程， +双方共享原 storage，接收方仍持有引用期间原地修改会彼此可见；需要快照语义时由调用方先 +`clone()`。返回的接收 buffer 生命周期由 `TensorMessage`/下游 Tensor 引用管理。 + +TLS 连接，或设置 `PULSING_TENSOR_TRANSPORT=http2` 时，使用兼容 HTTP/2 路径。该路径会 +打包和聚合 payload，copy 次数高于 raw TCP。可通过 `pulsing.tensor_transport_stats()` +检查 `active_copy_model`、raw frame/byte 计数和 HTTP/2 fallback 计数。接收边界可用以下 +环境变量限制，超限会在分配大 payload 前拒绝： + +- `PULSING_MAX_TENSOR_WIRE_BYTES`：单帧总大小，默认 64 GiB +- `PULSING_MAX_TENSOR_METADATA_BYTES`：metadata 大小，默认 64 MiB +- `PULSING_MAX_TENSOR_BUFFERS`：buffer 数量，默认 65536 + +同节点共享内存后端已在 `TensorTransport` 抽象中预留，当前尚未实现。 + +完整内容见 [TensorMessage 快速传输设计](docs/src/design/tensor-message-transport.zh.md) +和 [可运行的 TCP 示例](examples/python/tensor_message_fast_path.py)。 + ## 📚 示例导航 ``` @@ -203,6 +251,7 @@ examples/ ├── python/ # ⭐⭐ 基础示例 │ ├── ping_pong.py # Actor 基础 │ ├── cluster.py # 集群通信 +│ ├── tensor_message_fast_path.py # Tensor 传输 │ └── ... └── rust/ # Rust 示例 ``` diff --git a/crates/pulsing-actor/src/actor/mod.rs b/crates/pulsing-actor/src/actor/mod.rs index 70c34b264..80a0dfbec 100644 --- a/crates/pulsing-actor/src/actor/mod.rs +++ b/crates/pulsing-actor/src/actor/mod.rs @@ -138,4 +138,7 @@ pub use mailbox::{Envelope, Mailbox, MailboxSender, ResponseChannel, DEFAULT_MAI pub use reference::{ ActorRef, ActorRefInner, ActorResolver, LazyActorRef, RemoteActorRef, RemoteTransport, }; -pub use traits::{Actor, ActorId, IntoActor, Message, MessageStream, NodeId, StopReason}; +pub use traits::{ + max_tensor_buffers, max_tensor_metadata_bytes, max_tensor_wire_bytes, Actor, ActorId, + IntoActor, Message, MessageStream, NodeId, StopReason, TensorMessage, TENSOR_MESSAGE_TYPE, +}; diff --git a/crates/pulsing-actor/src/actor/traits.rs b/crates/pulsing-actor/src/actor/traits.rs index 6c1594969..f89da9e65 100644 --- a/crates/pulsing-actor/src/actor/traits.rs +++ b/crates/pulsing-actor/src/actor/traits.rs @@ -2,6 +2,7 @@ use crate::error::{PulsingError, Result, RuntimeError}; use async_trait::async_trait; +use bytes::Bytes; use futures::Stream; use serde::{de::DeserializeOwned, Deserialize, Serialize}; use serde_json; @@ -128,6 +129,301 @@ impl Format { /// Message stream type (stream of Single messages). pub type MessageStream = Pin> + Send>>; +/// Stable message type used when a [`TensorMessage`] crosses a remote transport. +pub const TENSOR_MESSAGE_TYPE: &str = "__pulsing_tensor_message__"; + +const TENSOR_WIRE_MAGIC: &[u8; 4] = b"PTM1"; +const TENSOR_WIRE_FIXED_HEADER: usize = 4 + 4 + 8 + 4; +const DEFAULT_MAX_TENSOR_WIRE_BYTES: u64 = 64 * 1024 * 1024 * 1024; +const DEFAULT_MAX_TENSOR_METADATA_BYTES: usize = 64 * 1024 * 1024; +const DEFAULT_MAX_TENSOR_BUFFERS: usize = 65_536; + +fn tensor_env_limit(name: &str, default: usize) -> usize { + std::env::var(name) + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(default) +} + +/// Maximum compatibility-wire body accepted before HTTP body aggregation. +pub fn max_tensor_wire_bytes() -> usize { + let platform_default = usize::try_from(DEFAULT_MAX_TENSOR_WIRE_BYTES).unwrap_or(usize::MAX); + tensor_env_limit("PULSING_MAX_TENSOR_WIRE_BYTES", platform_default) +} + +pub fn max_tensor_metadata_bytes() -> usize { + tensor_env_limit( + "PULSING_MAX_TENSOR_METADATA_BYTES", + DEFAULT_MAX_TENSOR_METADATA_BYTES, + ) +} + +pub fn max_tensor_buffers() -> usize { + tensor_env_limit("PULSING_MAX_TENSOR_BUFFERS", DEFAULT_MAX_TENSOR_BUFFERS) +} + +/// Transport-neutral tensor payload. +/// +/// Pulsing intentionally treats `metadata` as opaque bytes. Tensor schemas, +/// dtypes and shapes belong to the caller (for example PulsingQueue), while +/// Pulsing only owns routing and buffer transport. `Bytes` lets an in-process +/// message retain borrowed Python buffer owners without copying the payload. +#[derive(Clone)] +pub struct TensorMessage { + pub version: u32, + pub metadata: Bytes, + pub buffers: Vec, + origin: TensorBufferOrigin, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum TensorBufferOrigin { + Borrowed, + PackedWire, + OwnedReceive, +} + +impl TensorMessage { + pub fn new(version: u32, metadata: impl Into, buffers: Vec) -> Self { + Self { + version, + metadata: metadata.into(), + buffers, + origin: TensorBufferOrigin::Borrowed, + } + } + + /// Build a message from raw TCP receive allocations. `Vec -> Bytes` moves + /// ownership without copying; Python then exposes these as writable final + /// buffers retained by the resulting Tensor objects. + pub fn from_owned_receive( + version: u32, + metadata: Vec, + buffers: Vec>, + ) -> Result { + if metadata.len() > max_tensor_metadata_bytes() { + return Err(PulsingError::from(RuntimeError::Serialization( + "TensorMessage metadata exceeds configured maximum".into(), + ))); + } + if buffers.len() > max_tensor_buffers() { + return Err(PulsingError::from(RuntimeError::Serialization( + "TensorMessage buffer count exceeds configured maximum".into(), + ))); + } + let total = buffers.iter().try_fold(metadata.len(), |total, buffer| { + total.checked_add(buffer.len()) + }); + // Keep this explicit match for the Rust 1.75 MSRV; Option::is_none_or + // is newer than the version declared by the project. + let exceeds_limit = match total { + Some(total) => total > max_tensor_wire_bytes(), + None => true, + }; + if exceeds_limit { + return Err(PulsingError::from(RuntimeError::Serialization( + "TensorMessage payload exceeds configured maximum".into(), + ))); + } + Ok(Self { + version, + metadata: Bytes::from(metadata), + buffers: buffers.into_iter().map(Bytes::from).collect(), + origin: TensorBufferOrigin::OwnedReceive, + }) + } + + pub fn requires_owned_receive_copy(&self) -> bool { + self.origin == TensorBufferOrigin::PackedWire + } + + pub fn owns_receive_buffers(&self) -> bool { + self.origin == TensorBufferOrigin::OwnedReceive + } + + /// Sum of tensor payload bytes (metadata is deliberately excluded). + pub fn total_bytes(&self) -> usize { + self.buffers.iter().map(Bytes::len).sum() + } + + /// Encode the current HTTP/2 compatibility representation. + /// + /// This performs one payload packing copy. The dedicated tensor data plane + /// can replace this codec later without changing the public Message API. + pub fn encode_wire(&self) -> Result> { + if self.buffers.len() > max_tensor_buffers() { + return Err(PulsingError::from(RuntimeError::Serialization(format!( + "TensorMessage buffer count {} exceeds configured maximum {}", + self.buffers.len(), + max_tensor_buffers() + )))); + } + if self.metadata.len() > max_tensor_metadata_bytes() { + return Err(PulsingError::from(RuntimeError::Serialization(format!( + "TensorMessage metadata size {} exceeds configured maximum {}", + self.metadata.len(), + max_tensor_metadata_bytes() + )))); + } + let buffer_count = u32::try_from(self.buffers.len()).map_err(|_| { + PulsingError::from(RuntimeError::Serialization( + "TensorMessage has too many buffers".into(), + )) + })?; + let metadata_len = u64::try_from(self.metadata.len()).map_err(|_| { + PulsingError::from(RuntimeError::Serialization( + "TensorMessage metadata is too large".into(), + )) + })?; + + let lengths_bytes = self.buffers.len().checked_mul(8).ok_or_else(|| { + PulsingError::from(RuntimeError::Serialization( + "TensorMessage header size overflow".into(), + )) + })?; + let total_len = TENSOR_WIRE_FIXED_HEADER + .checked_add(lengths_bytes) + .and_then(|n| n.checked_add(self.metadata.len())) + .and_then(|n| { + self.buffers + .iter() + .try_fold(n, |acc, buffer| acc.checked_add(buffer.len())) + }) + .ok_or_else(|| { + PulsingError::from(RuntimeError::Serialization( + "TensorMessage payload size overflow".into(), + )) + })?; + + if total_len > max_tensor_wire_bytes() { + return Err(PulsingError::from(RuntimeError::Serialization(format!( + "TensorMessage wire size {total_len} exceeds configured maximum {}", + max_tensor_wire_bytes() + )))); + } + let mut wire = Vec::with_capacity(total_len); + wire.extend_from_slice(TENSOR_WIRE_MAGIC); + wire.extend_from_slice(&self.version.to_le_bytes()); + wire.extend_from_slice(&metadata_len.to_le_bytes()); + wire.extend_from_slice(&buffer_count.to_le_bytes()); + for buffer in &self.buffers { + let len = u64::try_from(buffer.len()).map_err(|_| { + PulsingError::from(RuntimeError::Serialization( + "TensorMessage buffer is too large".into(), + )) + })?; + wire.extend_from_slice(&len.to_le_bytes()); + } + wire.extend_from_slice(&self.metadata); + for buffer in &self.buffers { + wire.extend_from_slice(buffer); + } + Ok(wire) + } + + /// Decode the compatibility wire format without copying metadata or + /// individual tensor buffers. All returned `Bytes` are slices of `wire`. + pub fn decode_wire(wire: Bytes) -> Result { + let malformed = |reason: &str| { + PulsingError::from(RuntimeError::Serialization(format!( + "Malformed TensorMessage: {reason}" + ))) + }; + + if wire.len() > max_tensor_wire_bytes() { + return Err(malformed("wire size exceeds configured maximum")); + } + if wire.len() < TENSOR_WIRE_FIXED_HEADER { + return Err(malformed("header is truncated")); + } + if &wire[..4] != TENSOR_WIRE_MAGIC { + return Err(malformed("invalid magic")); + } + + let version = u32::from_le_bytes(wire[4..8].try_into().expect("fixed-width slice")); + let metadata_len_u64 = + u64::from_le_bytes(wire[8..16].try_into().expect("fixed-width slice")); + let metadata_len = usize::try_from(metadata_len_u64) + .map_err(|_| malformed("metadata length exceeds this platform"))?; + let buffer_count = + u32::from_le_bytes(wire[16..20].try_into().expect("fixed-width slice")) as usize; + if buffer_count > max_tensor_buffers() { + return Err(malformed("buffer count exceeds configured maximum")); + } + if metadata_len > max_tensor_metadata_bytes() { + return Err(malformed("metadata size exceeds configured maximum")); + } + let lengths_bytes = buffer_count + .checked_mul(8) + .ok_or_else(|| malformed("buffer count overflow"))?; + let payload_start = TENSOR_WIRE_FIXED_HEADER + .checked_add(lengths_bytes) + .ok_or_else(|| malformed("header size overflow"))?; + if payload_start > wire.len() { + return Err(malformed("buffer length table is truncated")); + } + + let mut lengths = Vec::with_capacity(buffer_count); + let mut total_payload = metadata_len; + for index in 0..buffer_count { + let offset = TENSOR_WIRE_FIXED_HEADER + index * 8; + let len_u64 = u64::from_le_bytes( + wire[offset..offset + 8] + .try_into() + .expect("validated length table"), + ); + let len = usize::try_from(len_u64) + .map_err(|_| malformed("buffer length exceeds this platform"))?; + total_payload = total_payload + .checked_add(len) + .ok_or_else(|| malformed("payload size overflow"))?; + lengths.push(len); + } + + let expected_len = payload_start + .checked_add(total_payload) + .ok_or_else(|| malformed("message size overflow"))?; + if expected_len != wire.len() { + return Err(malformed(if expected_len > wire.len() { + "payload is truncated" + } else { + "payload has trailing bytes" + })); + } + + let metadata_end = payload_start + metadata_len; + let metadata = wire.slice(payload_start..metadata_end); + let mut offset = metadata_end; + let buffers = lengths + .into_iter() + .map(|len| { + let end = offset + len; + let buffer = wire.slice(offset..end); + offset = end; + buffer + }) + .collect(); + + Ok(Self { + version, + metadata, + buffers, + origin: TensorBufferOrigin::PackedWire, + }) + } +} + +impl fmt::Debug for TensorMessage { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TensorMessage") + .field("version", &self.version) + .field("metadata_len", &self.metadata.len()) + .field("buffer_count", &self.buffers.len()) + .field("total_bytes", &self.total_bytes()) + .finish() + } +} + /// Unified message type for both requests and responses. pub enum Message { Single { @@ -138,6 +434,8 @@ pub enum Message { default_msg_type: String, stream: MessageStream, }, + /// Opaque tensor metadata plus one or more contiguous payload buffers. + Tensor(TensorMessage), } impl Message { @@ -157,6 +455,10 @@ impl Message { }) } + pub fn tensor(version: u32, metadata: impl Into, buffers: Vec) -> Self { + Message::Tensor(TensorMessage::new(version, metadata, buffers)) + } + pub fn unpack(self) -> Result { match self { Message::Single { data, .. } => bincode::deserialize(&data) @@ -164,6 +466,9 @@ impl Message { Message::Stream { .. } => Err(PulsingError::from(RuntimeError::Other( "Cannot unpack stream message".into(), ))), + Message::Tensor(_) => Err(PulsingError::from(RuntimeError::Other( + "Cannot unpack tensor message".into(), + ))), } } @@ -174,6 +479,9 @@ impl Message { Message::Stream { .. } => Err(PulsingError::from(RuntimeError::Other( "Cannot parse stream message".into(), ))), + Message::Tensor(_) => Err(PulsingError::from(RuntimeError::Other( + "Cannot parse tensor message".into(), + ))), } } @@ -204,6 +512,7 @@ impl Message { Message::Stream { default_msg_type, .. } => default_msg_type, + Message::Tensor(_) => TENSOR_MESSAGE_TYPE, } } @@ -214,6 +523,10 @@ impl Message { pub fn is_stream(&self) -> bool { matches!(self, Message::Stream { .. }) } + + pub fn is_tensor(&self) -> bool { + matches!(self, Message::Tensor(_)) + } } impl fmt::Debug for Message { @@ -231,6 +544,7 @@ impl fmt::Debug for Message { .debug_struct("Message::Stream") .field("default_msg_type", default_msg_type) .finish_non_exhaustive(), + Message::Tensor(tensor) => tensor.fmt(f), } } } @@ -415,6 +729,59 @@ mod tests { assert_eq!(request.msg_type(), "Echo"); } + #[test] + fn test_tensor_message_wire_roundtrip() { + let original = TensorMessage::new( + 7, + Bytes::from_static(b"opaque-schema"), + vec![ + Bytes::from_static(b"abc"), + Bytes::new(), + Bytes::from_static(b"defg"), + ], + ); + let wire = Bytes::from(original.encode_wire().unwrap()); + let wire_start = wire.as_ptr() as usize; + let wire_end = wire_start + wire.len(); + + let decoded = TensorMessage::decode_wire(wire).unwrap(); + assert_eq!(decoded.version, 7); + assert_eq!(decoded.metadata, Bytes::from_static(b"opaque-schema")); + assert_eq!(decoded.buffers.len(), 3); + assert_eq!(&decoded.buffers[0][..], b"abc"); + assert!(decoded.buffers[1].is_empty()); + assert_eq!(&decoded.buffers[2][..], b"defg"); + assert_eq!(decoded.total_bytes(), 7); + + // decode_wire returns slices of the received allocation rather than + // copying every tensor into a fresh Vec. + let first_ptr = decoded.buffers[0].as_ptr() as usize; + assert!((wire_start..wire_end).contains(&first_ptr)); + } + + #[test] + fn test_tensor_message_allows_control_only_payload() { + let original = TensorMessage::new(1, Bytes::from_static(b"control"), Vec::new()); + let decoded = + TensorMessage::decode_wire(Bytes::from(original.encode_wire().unwrap())).unwrap(); + assert_eq!(decoded.metadata, Bytes::from_static(b"control")); + assert!(decoded.buffers.is_empty()); + assert_eq!(decoded.total_bytes(), 0); + } + + #[test] + fn test_tensor_message_rejects_truncated_wire() { + let original = TensorMessage::new( + 1, + Bytes::from_static(b"meta"), + vec![Bytes::from_static(b"payload")], + ); + let mut wire = original.encode_wire().unwrap(); + wire.pop(); + let error = TensorMessage::decode_wire(Bytes::from(wire)).unwrap_err(); + assert!(error.to_string().contains("payload is truncated")); + } + #[tokio::test] async fn test_message_server_streaming() { // Simulate a server streaming response with Message stream diff --git a/crates/pulsing-actor/src/lib.rs b/crates/pulsing-actor/src/lib.rs index edb6aaf86..d4f82a38d 100644 --- a/crates/pulsing-actor/src/lib.rs +++ b/crates/pulsing-actor/src/lib.rs @@ -192,7 +192,7 @@ pub mod test_helper; pub use performance_store::{PerformanceSnapshot, PerformanceStore}; pub mod prelude { - pub use crate::actor::{Actor, ActorContext, ActorRef, IntoActor, Message}; + pub use crate::actor::{Actor, ActorContext, ActorRef, IntoActor, Message, TensorMessage}; pub use crate::supervision::{BackoffStrategy, RestartPolicy, SupervisionSpec}; pub use crate::system::{ ActorSystem, ActorSystemCoreExt, ActorSystemOpsExt, ResolveOptions, SpawnOptions, diff --git a/crates/pulsing-actor/src/system/handler.rs b/crates/pulsing-actor/src/system/handler.rs index 56cd2b03a..37f3607ce 100644 --- a/crates/pulsing-actor/src/system/handler.rs +++ b/crates/pulsing-actor/src/system/handler.rs @@ -161,6 +161,10 @@ impl Http2ServerHandler for SystemMessageHandler { self.dispatch_tell(path, msg).await } + async fn handle_tell_full(&self, path: &str, msg: Message) -> Result<()> { + self.dispatch_tell(path, msg).await + } + async fn handle_gossip( &self, payload: Vec, diff --git a/crates/pulsing-actor/src/system_actor/mod.rs b/crates/pulsing-actor/src/system_actor/mod.rs index 931589daa..3c5b1cb2a 100644 --- a/crates/pulsing-actor/src/system_actor/mod.rs +++ b/crates/pulsing-actor/src/system_actor/mod.rs @@ -470,6 +470,9 @@ impl Actor for SystemActor { Message::Stream { .. } => { return self.json_error_response("Stream messages not supported by SystemActor"); } + Message::Tensor(_) => { + return self.json_error_response("Tensor messages not supported by SystemActor"); + } }; let response = match sys_msg { diff --git a/crates/pulsing-actor/src/transport/http2/client.rs b/crates/pulsing-actor/src/transport/http2/client.rs index 2408c384a..74986b89d 100644 --- a/crates/pulsing-actor/src/transport/http2/client.rs +++ b/crates/pulsing-actor/src/transport/http2/client.rs @@ -5,22 +5,30 @@ use super::pool::{ConnectionPool, PoolConfig}; use super::retry::{RetryConfig, RetryExecutor}; use super::stream::{BinaryFrameParser, StreamFrame, StreamHandle}; use super::{headers, MessageMode, RequestType}; -use crate::actor::{Message, MessageStream}; +use crate::actor::{ + max_tensor_wire_bytes, Message, MessageStream, TensorMessage, TENSOR_MESSAGE_TYPE, +}; use crate::error::{PulsingError, Result, RuntimeError}; use crate::tracing::{ active_span_traceparent, active_span_tracestate, OpenTelemetrySpanExt, TRACEPARENT_HEADER, TRACESTATE_HEADER, }; +use crate::transport::tensor::{ + raw_tensor_transport_requested, read_raw_tensor_frame, record_tensor_http2_fallback, + write_raw_tensor_frame, RawTensorKind, TensorCopyModel, +}; use bytes::Bytes; use futures::{Stream, StreamExt, TryStreamExt}; -use http_body_util::{BodyExt, Full, StreamBody}; +use http_body_util::{BodyExt, Full, Limited, StreamBody}; use hyper::body::{Frame, Incoming}; use hyper::{Method, Request}; use opentelemetry::Context; +use std::collections::HashMap; use std::convert::Infallible; use std::net::SocketAddr; use std::sync::Arc; use std::time::Duration; +use tokio::net::TcpStream; use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; use tracing::Instrument; @@ -55,6 +63,7 @@ pub struct Http2Client { retry_config: RetryConfig, cancel: CancellationToken, fault_injector: Option>, + tensor_pool: Arc>>>, } impl Http2Client { @@ -65,6 +74,7 @@ impl Http2Client { retry_config: RetryConfig::default(), cancel: CancellationToken::new(), fault_injector: None, + tensor_pool: Arc::new(Mutex::new(HashMap::new())), } } @@ -82,6 +92,7 @@ impl Http2Client { retry_config, cancel: CancellationToken::new(), fault_injector: None, + tensor_pool: Arc::new(Mutex::new(HashMap::new())), } } @@ -333,6 +344,137 @@ impl Http2Client { .await?; self.parse_response(response).await } + Message::Tensor(tensor) => { + if self.raw_tensor_enabled() { + self.exchange_raw_tensor(addr, path, RawTensorKind::Ask, &tensor) + .await + } else { + // TLS or explicit compatibility mode: pack once into the + // existing HTTP/2 request body. + let wire = tensor.encode_wire()?; + record_tensor_http2_fallback(wire.len()); + self.send_message(addr, path, TENSOR_MESSAGE_TYPE, wire) + .await + } + } + } + } + + pub(crate) fn tensor_copy_model(&self) -> TensorCopyModel { + if self.raw_tensor_enabled() { + TensorCopyModel::DirectTcp + } else { + TensorCopyModel::PackedHttp2Compatibility + } + } + + fn raw_tensor_enabled(&self) -> bool { + if !raw_tensor_transport_requested() { + return false; + } + #[cfg(feature = "tls")] + if self.config.tls.is_some() { + return false; + } + true + } + + async fn take_tensor_connection(&self, addr: SocketAddr) -> Result { + loop { + let pooled = self + .tensor_pool + .lock() + .await + .get_mut(&addr) + .and_then(Vec::pop); + let Some(stream) = pooled else { + break; + }; + let mut probe = [0u8; 1]; + match stream.try_read(&mut probe) { + Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => return Ok(stream), + // EOF or unexpected unread protocol bytes: discard before a + // request starts, then safely establish a fresh connection. + Ok(0) | Ok(_) | Err(_) => continue, + } + } + + let stream = tokio::time::timeout(self.config.connect_timeout, TcpStream::connect(addr)) + .await + .map_err(|_| RuntimeError::connection_failed(addr.to_string(), "Timeout"))? + .map_err(|error| { + RuntimeError::connection_failed(addr.to_string(), error.to_string()) + })?; + stream + .set_nodelay(true) + .map_err(|error| RuntimeError::Io(error.to_string()))?; + Ok(stream) + } + + async fn return_tensor_connection(&self, addr: SocketAddr, stream: TcpStream) { + let mut pool = self.tensor_pool.lock().await; + let connections = pool.entry(addr).or_default(); + if connections.len() < self.config.max_connections_per_host { + connections.push(stream); + } + } + + async fn exchange_raw_tensor( + &self, + addr: SocketAddr, + path: &str, + kind: RawTensorKind, + message: &TensorMessage, + ) -> Result { + let mut stream = self.take_tensor_connection(addr).await?; + let exchange = async { + write_raw_tensor_frame(&mut stream, kind, path, message).await?; + read_raw_tensor_frame(&mut stream).await + }; + let frame = tokio::time::timeout(self.config.stream_timeout, exchange) + .await + .map_err(|_| { + PulsingError::from(RuntimeError::request_timeout( + self.config.stream_timeout.as_millis() as u64, + )) + })??; + + let result = match frame.kind { + RawTensorKind::TensorResponse if kind == RawTensorKind::Ask => { + Ok(Message::Tensor(frame.message)) + } + RawTensorKind::SingleResponse if kind == RawTensorKind::Ask => { + Ok(Message::single(frame.path, frame.message.metadata.to_vec())) + } + RawTensorKind::Ack if kind == RawTensorKind::Tell => { + Ok(Message::single("", Vec::new())) + } + RawTensorKind::Error => { + let error = String::from_utf8_lossy(&frame.message.metadata); + Err(PulsingError::from(RuntimeError::Other(error.into_owned()))) + } + other => Err(PulsingError::from(RuntimeError::Other(format!( + "Unexpected raw tensor response: {other:?}" + )))), + }; + self.return_tensor_connection(addr, stream).await; + result + } + + pub(crate) async fn tell_tensor( + &self, + addr: SocketAddr, + path: &str, + message: TensorMessage, + ) -> Result<()> { + if self.raw_tensor_enabled() { + self.exchange_raw_tensor(addr, path, RawTensorKind::Tell, &message) + .await?; + Ok(()) + } else { + let wire = message.encode_wire()?; + record_tensor_http2_fallback(wire.len()); + self.tell(addr, path, TENSOR_MESSAGE_TYPE, wire).await } } @@ -567,6 +709,12 @@ impl Http2Client { .get(headers::RESPONSE_TYPE) .and_then(|v| v.to_str().ok()) .unwrap_or("single"); + let response_msg_type = response + .headers() + .get(headers::MESSAGE_TYPE) + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .to_string(); if response_type == "stream" { let cancel = CancellationToken::new(); @@ -592,6 +740,18 @@ impl Http2Client { default_msg_type: String::new(), stream: Box::pin(msg_stream), }) + } else if response_type == "tensor" { + let body = tokio::time::timeout( + self.config.request_timeout, + Limited::new(response.into_body(), max_tensor_wire_bytes()).collect(), + ) + .await + .map_err(|_| { + RuntimeError::request_timeout(self.config.request_timeout.as_millis() as u64) + })? + .map_err(|e| RuntimeError::Io(e.to_string()))? + .to_bytes(); + Ok(Message::Tensor(TensorMessage::decode_wire(body)?)) } else { let body = tokio::time::timeout(self.config.request_timeout, response.collect()) .await @@ -601,7 +761,7 @@ impl Http2Client { .map_err(|e| RuntimeError::Io(e.to_string()))? .to_bytes(); - Ok(Message::single("", body.to_vec())) + Ok(Message::single(response_msg_type, body.to_vec())) } } @@ -710,6 +870,7 @@ impl Clone for Http2Client { retry_config: self.retry_config.clone(), cancel: self.cancel.clone(), fault_injector: self.fault_injector.clone(), + tensor_pool: self.tensor_pool.clone(), } } } diff --git a/crates/pulsing-actor/src/transport/http2/mod.rs b/crates/pulsing-actor/src/transport/http2/mod.rs index e5dcbd470..e876681ca 100644 --- a/crates/pulsing-actor/src/transport/http2/mod.rs +++ b/crates/pulsing-actor/src/transport/http2/mod.rs @@ -26,6 +26,7 @@ pub use tls::TlsConfig; use crate::actor::{ActorId, ActorPath, Message, RemoteTransport}; use crate::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig}; +use crate::transport::tensor::{TensorCopyModel, TensorTransport}; use std::net::SocketAddr; use std::sync::Arc; use tokio_util::sync::CancellationToken; @@ -99,20 +100,28 @@ impl Http2Transport { pub async fn tell(&self, addr: SocketAddr, actor_name: &str, msg: Message) -> Result<()> { let path = format!("/actors/{}", actor_name); - let Message::Single { msg_type, data } = msg else { - return Err(RuntimeError::protocol_error("Streaming not supported for tell").into()); - }; - - self.client.tell(addr, &path, &msg_type, data).await + match msg { + Message::Single { msg_type, data } => { + self.client.tell(addr, &path, &msg_type, data).await + } + Message::Tensor(tensor) => self.client.tell_tensor(addr, &path, tensor).await, + Message::Stream { .. } => { + Err(RuntimeError::protocol_error("Streaming not supported for tell").into()) + } + } } pub async fn tell_named(&self, addr: SocketAddr, path: &ActorPath, msg: Message) -> Result<()> { let url_path = format!("/named/{}", path.as_str()); - let Message::Single { msg_type, data } = msg else { - return Err(RuntimeError::protocol_error("Streaming not supported for tell").into()); - }; - - self.client.tell(addr, &url_path, &msg_type, data).await + match msg { + Message::Single { msg_type, data } => { + self.client.tell(addr, &url_path, &msg_type, data).await + } + Message::Tensor(tensor) => self.client.tell_tensor(addr, &url_path, tensor).await, + Message::Stream { .. } => { + Err(RuntimeError::protocol_error("Streaming not supported for tell").into()) + } + } } /// Send a gossip message @@ -205,6 +214,7 @@ impl RequestType { pub enum ResponseType { Single, Stream, + Tensor, } impl ResponseType { @@ -212,6 +222,7 @@ impl ResponseType { match self { ResponseType::Single => "single", ResponseType::Stream => "stream", + ResponseType::Tensor => "tensor", } } @@ -219,6 +230,7 @@ impl ResponseType { match s.to_lowercase().as_str() { "single" => Some(ResponseType::Single), "stream" => Some(ResponseType::Stream), + "tensor" => Some(ResponseType::Tensor), _ => None, } } @@ -397,12 +409,47 @@ impl RemoteTransport for Http2RemoteTransport { /// Send a one-way message (unified interface) async fn send_oneway(&self, actor_id: &ActorId, msg: Message) -> Result<()> { - let Message::Single { msg_type, data } = msg else { - return Err(PulsingError::from(RuntimeError::Other( + match msg { + Message::Single { msg_type, data } => self.send(actor_id, &msg_type, data).await, + Message::Tensor(tensor) => { + let _ = actor_id; + self.client + .tell_tensor(self.remote_addr, &self.path, tensor) + .await + } + Message::Stream { .. } => Err(PulsingError::from(RuntimeError::Other( "Streaming not supported for fire-and-forget (use ask pattern instead)".into(), - ))); - }; - self.send(actor_id, &msg_type, data).await + ))), + } + } +} + +#[async_trait::async_trait] +impl TensorTransport for Http2RemoteTransport { + async fn request_tensor( + &self, + actor_id: &ActorId, + message: crate::actor::TensorMessage, + ) -> Result { + match RemoteTransport::send_message(self, actor_id, Message::Tensor(message)).await? { + Message::Tensor(response) => Ok(response), + other => Err(PulsingError::from(RuntimeError::Other(format!( + "Tensor transport returned unexpected message type: {}", + other.msg_type() + )))), + } + } + + async fn send_tensor( + &self, + actor_id: &ActorId, + message: crate::actor::TensorMessage, + ) -> Result<()> { + RemoteTransport::send_oneway(self, actor_id, Message::Tensor(message)).await + } + + fn copy_model(&self) -> TensorCopyModel { + self.client.tensor_copy_model() } } diff --git a/crates/pulsing-actor/src/transport/http2/server.rs b/crates/pulsing-actor/src/transport/http2/server.rs index a9c3b4270..c9de3883d 100644 --- a/crates/pulsing-actor/src/transport/http2/server.rs +++ b/crates/pulsing-actor/src/transport/http2/server.rs @@ -3,12 +3,16 @@ use super::config::Http2Config; use super::stream::{BinaryFrameParser, StreamFrame}; use super::{headers, MessageMode, RequestType}; -use crate::actor::Message; +use crate::actor::{max_tensor_wire_bytes, Message, TensorMessage, TENSOR_MESSAGE_TYPE}; use crate::error::{PulsingError, Result, RuntimeError}; use crate::tracing::{OpenTelemetrySpanExt, TraceContext, TRACEPARENT_HEADER, TRACESTATE_HEADER}; +use crate::transport::tensor::{ + raw_tensor_transport_requested, read_raw_tensor_frame, record_raw_tensor_connection, + write_raw_tensor_frame, RawTensorKind, RAW_TENSOR_MAGIC, +}; use bytes::Bytes; use futures::StreamExt; -use http_body_util::{BodyExt, Full, StreamBody}; +use http_body_util::{BodyExt, Full, Limited, StreamBody}; use hyper::body::{Frame, Incoming}; use hyper::server::conn::http2; use hyper::service::service_fn; @@ -35,6 +39,9 @@ pub trait Http2ServerHandler: Send + Sync + 'static { Message::Stream { .. } => Err(PulsingError::from(RuntimeError::Other( "Streaming requests not supported by this handler".into(), ))), + Message::Tensor(_) => Err(PulsingError::from(RuntimeError::Other( + "Tensor requests not supported by this handler".into(), + ))), } } @@ -54,6 +61,20 @@ pub trait Http2ServerHandler: Send + Sync + 'static { /// Handle tell (fire-and-forget) message. async fn handle_tell(&self, path: &str, msg_type: &str, payload: Vec) -> Result<()>; + /// Unified tell handler. Implementations which route arbitrary actor + /// messages should override this to preserve TensorMessage buffers. + async fn handle_tell_full(&self, path: &str, msg: Message) -> Result<()> { + match msg { + Message::Single { msg_type, data } => self.handle_tell(path, &msg_type, data).await, + Message::Stream { .. } => Err(PulsingError::from(RuntimeError::Other( + "Streaming requests not supported for tell".into(), + ))), + Message::Tensor(_) => Err(PulsingError::from(RuntimeError::Other( + "Tensor requests not supported for tell by this handler".into(), + ))), + } + } + /// Handle gossip message. async fn handle_gossip( &self, @@ -211,11 +232,154 @@ impl Http2Server { return Self::serve_h2_generic(io, peer_addr, handler, config, cancel).await; } + if raw_tensor_transport_requested() + && Self::is_raw_tensor_connection(&stream, config.connect_timeout).await? + { + return Self::serve_raw_tensor(stream, peer_addr, handler, cancel).await; + } + // Plain TCP mode (h2c) - HTTP/2 only (prior knowledge) let io = TokioIo::new(stream); Self::serve_h2(io, peer_addr, handler, config, cancel).await } + async fn is_raw_tensor_connection( + stream: &tokio::net::TcpStream, + timeout: std::time::Duration, + ) -> anyhow::Result { + let mut prefix = [0u8; 4]; + let peek = async { + loop { + let count = stream.peek(&mut prefix).await?; + if count == 0 { + return Ok::(false); + } + if count >= prefix.len() { + return Ok(&prefix == RAW_TENSOR_MAGIC); + } + tokio::time::sleep(std::time::Duration::from_millis(1)).await; + } + }; + tokio::time::timeout(timeout, peek) + .await + .map_err(|_| anyhow::anyhow!("Connection preface timeout"))? + .map_err(Into::into) + } + + async fn serve_raw_tensor( + mut stream: tokio::net::TcpStream, + peer_addr: SocketAddr, + handler: Arc, + cancel: CancellationToken, + ) -> anyhow::Result<()> { + stream.set_nodelay(true)?; + record_raw_tensor_connection(); + tracing::debug!(peer = %peer_addr, "Raw tensor connection established"); + let max_requests = std::env::var("PULSING_TENSOR_MAX_REQUESTS_PER_CONNECTION") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(usize::MAX) + .max(1); + let mut requests = 0usize; + + loop { + let frame = tokio::select! { + frame = read_raw_tensor_frame(&mut stream) => frame, + _ = cancel.cancelled() => return Ok(()), + } + .map_err(|error| anyhow::anyhow!(error.to_string()))?; + + let version = frame.message.version; + match frame.kind { + RawTensorKind::Ask => { + match handler + .handle_message_full(&frame.path, Message::Tensor(frame.message)) + .await + { + Ok(Message::Tensor(response)) => { + write_raw_tensor_frame( + &mut stream, + RawTensorKind::TensorResponse, + "", + &response, + ) + .await + .map_err(|error| anyhow::anyhow!(error.to_string()))?; + } + Ok(Message::Single { msg_type, data }) => { + let response = + TensorMessage::new(version, Bytes::from(data), Vec::new()); + write_raw_tensor_frame( + &mut stream, + RawTensorKind::SingleResponse, + &msg_type, + &response, + ) + .await + .map_err(|error| anyhow::anyhow!(error.to_string()))?; + } + Ok(Message::Stream { .. }) => { + Self::write_raw_tensor_error( + &mut stream, + version, + "Tensor request returned a streaming response", + ) + .await?; + } + Err(error) => { + Self::write_raw_tensor_error(&mut stream, version, &error.to_string()) + .await?; + } + } + } + RawTensorKind::Tell => { + match handler + .handle_tell_full(&frame.path, Message::Tensor(frame.message)) + .await + { + Ok(()) => { + let ack = TensorMessage::new(version, Bytes::new(), Vec::new()); + write_raw_tensor_frame(&mut stream, RawTensorKind::Ack, "", &ack) + .await + .map_err(|error| anyhow::anyhow!(error.to_string()))?; + } + Err(error) => { + Self::write_raw_tensor_error(&mut stream, version, &error.to_string()) + .await?; + } + } + } + unexpected => { + Self::write_raw_tensor_error( + &mut stream, + version, + &format!("Unexpected request frame: {unexpected:?}"), + ) + .await?; + } + } + requests += 1; + if requests >= max_requests { + return Ok(()); + } + } + } + + async fn write_raw_tensor_error( + stream: &mut tokio::net::TcpStream, + version: u32, + error: &str, + ) -> anyhow::Result<()> { + let message = TensorMessage::new( + version, + Bytes::copy_from_slice(error.as_bytes()), + Vec::new(), + ); + write_raw_tensor_frame(stream, RawTensorKind::Error, "", &message) + .await + .map_err(|error| anyhow::anyhow!(error.to_string())) + } + /// Serve HTTP/2 only (prior knowledge mode) async fn serve_h2( io: TokioIo, @@ -449,33 +613,51 @@ impl Http2Server { match mode { MessageMode::Tell => { - // Tell doesn't support streaming requests - let body_bytes = match req.collect().await { - Ok(collected) => collected.to_bytes().to_vec(), + let is_tensor = msg_type == TENSOR_MESSAGE_TYPE; + let body = match Self::collect_actor_body(req.into_body(), is_tensor).await { + Ok(body) => body, + Err(e) => { + return Ok(error_response( + StatusCode::BAD_REQUEST, + format!("Failed to read body: {e}").into_bytes(), + )); + } + }; + let msg = match Self::decode_single_or_tensor(&msg_type, body) { + Ok(msg) => msg, Err(e) => { return Ok(error_response( StatusCode::BAD_REQUEST, - format!("Failed to read body: {}", e).into_bytes(), + e.to_string().into_bytes(), )); } }; - Self::handle_tell_request(&handler, &path, &msg_type, body_bytes).await + Self::handle_tell_request_full(&handler, &path, msg).await } // Ask and Stream mode - check for streaming request MessageMode::Ask | MessageMode::Stream => { match request_type { RequestType::Single => { - // Single request - read body as bytes - let body_bytes = match req.collect().await { - Ok(collected) => collected.to_bytes().to_vec(), + let is_tensor = msg_type == TENSOR_MESSAGE_TYPE; + let body = + match Self::collect_actor_body(req.into_body(), is_tensor).await { + Ok(body) => body, Err(e) => { return Ok(error_response( StatusCode::BAD_REQUEST, - format!("Failed to read body: {}", e).into_bytes(), + format!("Failed to read body: {e}").into_bytes(), + )); + } + }; + let msg = match Self::decode_single_or_tensor(&msg_type, body) { + Ok(msg) => msg, + Err(e) => { + return Ok(error_response( + StatusCode::BAD_REQUEST, + e.to_string().into_bytes(), )); } }; - let msg = Message::single(&msg_type, body_bytes); Self::handle_message_request_full(&handler, &path, msg).await } RequestType::Stream => { @@ -491,6 +673,32 @@ impl Http2Server { .await } + fn decode_single_or_tensor(msg_type: &str, body: Bytes) -> Result { + if msg_type == TENSOR_MESSAGE_TYPE { + TensorMessage::decode_wire(body).map(Message::Tensor) + } else { + Ok(Message::single(msg_type, body.to_vec())) + } + } + + async fn collect_actor_body( + body: Incoming, + is_tensor: bool, + ) -> std::result::Result { + if is_tensor { + Limited::new(body, max_tensor_wire_bytes()) + .collect() + .await + .map(|collected| collected.to_bytes()) + .map_err(|error| error.to_string()) + } else { + body.collect() + .await + .map(|collected| collected.to_bytes()) + .map_err(|error| error.to_string()) + } + } + /// Parse a streaming request body (binary frames) into Message::Stream fn parse_streaming_request(body: Incoming, default_msg_type: &str) -> Message { let (tx, rx) = mpsc::channel::>(32); @@ -572,12 +780,17 @@ impl Http2Server { msg: Message, ) -> std::result::Result, Infallible> { match handler.handle_message_full(path, msg).await { - Ok(Message::Single { data, .. }) => { - // Single response - return directly with response type header + Ok(Message::Single { msg_type, data }) => { + // Preserve the message type so the client can select the + // matching decoder (for example Python pickle for a small + // application-level tensor ACK). Ok(octet_response_with_headers( StatusCode::OK, full_body(data), - &[(headers::RESPONSE_TYPE, "single")], + &[ + (headers::RESPONSE_TYPE, "single"), + (headers::MESSAGE_TYPE, &msg_type), + ], )) } Ok(Message::Stream { @@ -616,6 +829,17 @@ impl Http2Server { &[(headers::RESPONSE_TYPE, "stream")], )) } + Ok(Message::Tensor(tensor)) => match tensor.encode_wire() { + Ok(wire) => Ok(octet_response_with_headers( + StatusCode::OK, + full_body(wire), + &[(headers::RESPONSE_TYPE, "tensor")], + )), + Err(e) => Ok(error_response( + StatusCode::INTERNAL_SERVER_ERROR, + e.to_string().into_bytes(), + )), + }, Err(e) => Ok(error_response( StatusCode::INTERNAL_SERVER_ERROR, e.to_string().into_bytes(), @@ -623,13 +847,12 @@ impl Http2Server { } } - async fn handle_tell_request( + async fn handle_tell_request_full( handler: &Arc, path: &str, - msg_type: &str, - payload: Vec, + msg: Message, ) -> std::result::Result, Infallible> { - match handler.handle_tell(path, msg_type, payload).await { + match handler.handle_tell_full(path, msg).await { Ok(()) => Ok(empty_response(StatusCode::ACCEPTED)), Err(e) => Ok(error_response( StatusCode::INTERNAL_SERVER_ERROR, diff --git a/crates/pulsing-actor/src/transport/http2/stream.rs b/crates/pulsing-actor/src/transport/http2/stream.rs index 31cf4738a..13b9a479e 100644 --- a/crates/pulsing-actor/src/transport/http2/stream.rs +++ b/crates/pulsing-actor/src/transport/http2/stream.rs @@ -73,6 +73,7 @@ impl StreamFrame { Self::data(frame_msg_type, data) } Message::Stream { .. } => Self::error("Nested streams are not supported"), + Message::Tensor(_) => Self::error("Tensor messages cannot be nested in streams"), } } diff --git a/crates/pulsing-actor/src/transport/mod.rs b/crates/pulsing-actor/src/transport/mod.rs index c38e38bc3..8fe4e29b1 100644 --- a/crates/pulsing-actor/src/transport/mod.rs +++ b/crates/pulsing-actor/src/transport/mod.rs @@ -13,6 +13,7 @@ //! - Circuit breaker for fault tolerance pub mod http2; +pub mod tensor; // HTTP/2 exports pub use http2::{ @@ -21,3 +22,6 @@ pub use http2::{ PoolConfig, PoolStats, RequestType, RetryConfig, RetryableError, StreamFrame, StreamHandle, TransportTarget, FLAG_END, FLAG_ERROR, }; +pub use tensor::{ + raw_tensor_transport_stats, RawTensorTransportStats, TensorCopyModel, TensorTransport, +}; diff --git a/crates/pulsing-actor/src/transport/tensor.rs b/crates/pulsing-actor/src/transport/tensor.rs new file mode 100644 index 000000000..6a9bc4811 --- /dev/null +++ b/crates/pulsing-actor/src/transport/tensor.rs @@ -0,0 +1,502 @@ +//! Tensor data-plane abstraction. +//! +//! Clear-text tensor traffic uses a pooled, long-lived raw TCP connection. +//! `write_vectored` writes the small header, opaque metadata, and the original +//! buffers without first packing them into a combined body; `read_exact` reads +//! each payload directly into its final owned receive allocation. TLS and the +//! explicit `PULSING_TENSOR_TRANSPORT=http2` mode retain a compatibility HTTP/2 +//! backend, which packs the buffers and therefore has additional payload copies. +//! A future same-host shared-memory implementation can implement the same trait +//! without changing PulsingQueue. + +use crate::actor::{ + max_tensor_buffers, max_tensor_metadata_bytes, max_tensor_wire_bytes, ActorId, TensorMessage, +}; +use crate::error::{PulsingError, Result, RuntimeError}; +use std::io::{ErrorKind, IoSlice}; +use std::sync::atomic::{AtomicU64, Ordering}; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; + +pub(crate) const RAW_TENSOR_MAGIC: &[u8; 4] = b"PTR1"; +const RAW_FIXED_HEADER_LEN: usize = 4 + 1 + 3 + 4 + 4 + 8 + 4; +const MAX_RAW_PATH_BYTES: usize = 64 * 1024; +const MAX_VECTORED_SLICES: usize = 64; + +static RAW_FRAMES_SENT: AtomicU64 = AtomicU64::new(0); +static RAW_FRAMES_RECEIVED: AtomicU64 = AtomicU64::new(0); +static RAW_BYTES_SENT: AtomicU64 = AtomicU64::new(0); +static RAW_BYTES_RECEIVED: AtomicU64 = AtomicU64::new(0); +static HTTP2_FALLBACK_FRAMES: AtomicU64 = AtomicU64::new(0); +static HTTP2_FALLBACK_BYTES: AtomicU64 = AtomicU64::new(0); +static LAST_COPY_MODEL: AtomicU64 = AtomicU64::new(0); +static RAW_CONNECTIONS_ACCEPTED: AtomicU64 = AtomicU64::new(0); + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct RawTensorTransportStats { + pub frames_sent: u64, + pub frames_received: u64, + pub bytes_sent: u64, + pub bytes_received: u64, + pub http2_fallback_frames: u64, + pub http2_fallback_bytes: u64, + pub active_copy_model: &'static str, + pub raw_connections_accepted: u64, +} + +pub fn raw_tensor_transport_stats() -> RawTensorTransportStats { + let active_copy_model = match LAST_COPY_MODEL.load(Ordering::Relaxed) { + 1 => "direct_tcp", + 2 => "packed_http2_compatibility", + _ => "unused", + }; + RawTensorTransportStats { + frames_sent: RAW_FRAMES_SENT.load(Ordering::Relaxed), + frames_received: RAW_FRAMES_RECEIVED.load(Ordering::Relaxed), + bytes_sent: RAW_BYTES_SENT.load(Ordering::Relaxed), + bytes_received: RAW_BYTES_RECEIVED.load(Ordering::Relaxed), + http2_fallback_frames: HTTP2_FALLBACK_FRAMES.load(Ordering::Relaxed), + http2_fallback_bytes: HTTP2_FALLBACK_BYTES.load(Ordering::Relaxed), + active_copy_model, + raw_connections_accepted: RAW_CONNECTIONS_ACCEPTED.load(Ordering::Relaxed), + } +} + +pub(crate) fn record_tensor_http2_fallback(bytes: usize) { + HTTP2_FALLBACK_FRAMES.fetch_add(1, Ordering::Relaxed); + HTTP2_FALLBACK_BYTES.fetch_add(bytes as u64, Ordering::Relaxed); + LAST_COPY_MODEL.store(2, Ordering::Relaxed); +} + +pub(crate) fn record_raw_tensor_connection() { + RAW_CONNECTIONS_ACCEPTED.fetch_add(1, Ordering::Relaxed); +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[repr(u8)] +pub(crate) enum RawTensorKind { + Ask = 1, + Tell = 2, + TensorResponse = 3, + Ack = 4, + Error = 5, + SingleResponse = 6, +} + +impl TryFrom for RawTensorKind { + type Error = PulsingError; + + fn try_from(value: u8) -> Result { + match value { + 1 => Ok(Self::Ask), + 2 => Ok(Self::Tell), + 3 => Ok(Self::TensorResponse), + 4 => Ok(Self::Ack), + 5 => Ok(Self::Error), + 6 => Ok(Self::SingleResponse), + _ => Err(protocol_error(format!( + "unknown raw tensor frame kind {value}" + ))), + } + } +} + +#[derive(Debug)] +pub(crate) struct RawTensorFrame { + pub kind: RawTensorKind, + pub path: String, + pub message: TensorMessage, +} + +fn protocol_error(message: impl Into) -> PulsingError { + PulsingError::from(RuntimeError::Other(format!( + "Raw tensor protocol error: {}", + message.into() + ))) +} + +/// Write header, metadata and each tensor buffer directly to the socket. +/// No combined payload allocation is constructed on this path. +pub(crate) async fn write_raw_tensor_frame( + writer: &mut W, + kind: RawTensorKind, + path: &str, + message: &TensorMessage, +) -> Result<()> +where + W: AsyncWrite + Unpin, +{ + let path_bytes = path.as_bytes(); + if path_bytes.len() > MAX_RAW_PATH_BYTES { + return Err(protocol_error("actor path is too large")); + } + if message.metadata.len() > max_tensor_metadata_bytes() { + return Err(protocol_error("metadata exceeds configured maximum")); + } + if message.buffers.len() > max_tensor_buffers() { + return Err(protocol_error("buffer count exceeds configured maximum")); + } + + let lengths_bytes = message + .buffers + .len() + .checked_mul(8) + .ok_or_else(|| protocol_error("length table overflow"))?; + let total = RAW_FIXED_HEADER_LEN + .checked_add(lengths_bytes) + .and_then(|size| size.checked_add(path_bytes.len())) + .and_then(|size| size.checked_add(message.metadata.len())) + .and_then(|size| { + message + .buffers + .iter() + .try_fold(size, |size, buffer| size.checked_add(buffer.len())) + }) + .ok_or_else(|| protocol_error("frame size overflow"))?; + if total > max_tensor_wire_bytes() { + return Err(protocol_error("frame exceeds configured maximum")); + } + + let mut header = Vec::with_capacity(RAW_FIXED_HEADER_LEN + lengths_bytes); + header.extend_from_slice(RAW_TENSOR_MAGIC); + header.push(kind as u8); + header.extend_from_slice(&[0; 3]); + header.extend_from_slice(&message.version.to_le_bytes()); + header.extend_from_slice(&(path_bytes.len() as u32).to_le_bytes()); + header.extend_from_slice(&(message.metadata.len() as u64).to_le_bytes()); + header.extend_from_slice(&(message.buffers.len() as u32).to_le_bytes()); + for buffer in &message.buffers { + header.extend_from_slice(&(buffer.len() as u64).to_le_bytes()); + } + + let mut slices = Vec::with_capacity(3 + message.buffers.len()); + slices.push(IoSlice::new(&header)); + if !path_bytes.is_empty() { + slices.push(IoSlice::new(path_bytes)); + } + if !message.metadata.is_empty() { + slices.push(IoSlice::new(&message.metadata)); + } + slices.extend( + message + .buffers + .iter() + .filter(|buffer| !buffer.is_empty()) + .map(|buffer| IoSlice::new(buffer)), + ); + write_all_vectored(writer, &mut slices).await?; + writer + .flush() + .await + .map_err(|error| protocol_error(error.to_string()))?; + RAW_FRAMES_SENT.fetch_add(1, Ordering::Relaxed); + RAW_BYTES_SENT.fetch_add(total as u64, Ordering::Relaxed); + LAST_COPY_MODEL.store(1, Ordering::Relaxed); + Ok(()) +} + +async fn write_all_vectored(writer: &mut W, slices: &mut [IoSlice<'_>]) -> Result<()> +where + W: AsyncWrite + Unpin, +{ + let mut remaining = slices; + while !remaining.is_empty() { + let batch_len = remaining.len().min(MAX_VECTORED_SLICES); + let written = writer + .write_vectored(&remaining[..batch_len]) + .await + .map_err(|error| protocol_error(error.to_string()))?; + if written == 0 { + return Err(protocol_error( + std::io::Error::new(ErrorKind::WriteZero, "failed to write tensor frame") + .to_string(), + )); + } + IoSlice::advance_slices(&mut remaining, written); + } + Ok(()) +} + +/// Read directly into the final per-buffer allocations. The returned +/// TensorMessage moves those Vec allocations into Bytes without a payload copy. +pub(crate) async fn read_raw_tensor_frame(reader: &mut R) -> Result +where + R: AsyncRead + Unpin, +{ + let mut fixed = [0u8; RAW_FIXED_HEADER_LEN]; + reader + .read_exact(&mut fixed) + .await + .map_err(|error| protocol_error(error.to_string()))?; + if &fixed[..4] != RAW_TENSOR_MAGIC { + return Err(protocol_error("invalid magic")); + } + let kind = RawTensorKind::try_from(fixed[4])?; + let version = u32::from_le_bytes(fixed[8..12].try_into().expect("fixed-width slice")); + let path_len = + u32::from_le_bytes(fixed[12..16].try_into().expect("fixed-width slice")) as usize; + let metadata_len_u64 = u64::from_le_bytes(fixed[16..24].try_into().expect("fixed-width slice")); + let metadata_len = usize::try_from(metadata_len_u64) + .map_err(|_| protocol_error("metadata length exceeds this platform"))?; + let buffer_count = + u32::from_le_bytes(fixed[24..28].try_into().expect("fixed-width slice")) as usize; + + if path_len > MAX_RAW_PATH_BYTES { + return Err(protocol_error("actor path exceeds configured maximum")); + } + if metadata_len > max_tensor_metadata_bytes() { + return Err(protocol_error("metadata exceeds configured maximum")); + } + if buffer_count > max_tensor_buffers() { + return Err(protocol_error("buffer count exceeds configured maximum")); + } + + let lengths_bytes = buffer_count + .checked_mul(8) + .ok_or_else(|| protocol_error("length table overflow"))?; + let mut length_table = vec![0u8; lengths_bytes]; + reader + .read_exact(&mut length_table) + .await + .map_err(|error| protocol_error(error.to_string()))?; + let mut lengths = Vec::with_capacity(buffer_count); + let mut total = RAW_FIXED_HEADER_LEN + .checked_add(lengths_bytes) + .and_then(|size| size.checked_add(path_len)) + .and_then(|size| size.checked_add(metadata_len)) + .ok_or_else(|| protocol_error("frame size overflow"))?; + for index in 0..buffer_count { + let offset = index * 8; + let len_u64 = u64::from_le_bytes( + length_table[offset..offset + 8] + .try_into() + .expect("fixed-width slice"), + ); + let len = usize::try_from(len_u64) + .map_err(|_| protocol_error("buffer length exceeds this platform"))?; + total = total + .checked_add(len) + .ok_or_else(|| protocol_error("frame size overflow"))?; + lengths.push(len); + } + if total > max_tensor_wire_bytes() { + return Err(protocol_error("frame exceeds configured maximum")); + } + + let mut path_bytes = vec![0u8; path_len]; + reader + .read_exact(&mut path_bytes) + .await + .map_err(|error| protocol_error(error.to_string()))?; + let path = String::from_utf8(path_bytes) + .map_err(|error| protocol_error(format!("invalid actor path: {error}")))?; + + let mut metadata = vec![0u8; metadata_len]; + reader + .read_exact(&mut metadata) + .await + .map_err(|error| protocol_error(error.to_string()))?; + let mut buffers = Vec::with_capacity(buffer_count); + for len in lengths { + let mut buffer = vec![0u8; len]; + reader + .read_exact(&mut buffer) + .await + .map_err(|error| protocol_error(error.to_string()))?; + buffers.push(buffer); + } + + RAW_FRAMES_RECEIVED.fetch_add(1, Ordering::Relaxed); + RAW_BYTES_RECEIVED.fetch_add(total as u64, Ordering::Relaxed); + Ok(RawTensorFrame { + kind, + path, + message: TensorMessage::from_owned_receive(version, metadata, buffers)?, + }) +} + +pub(crate) fn raw_tensor_transport_requested() -> bool { + !matches!( + std::env::var("PULSING_TENSOR_TRANSPORT") + .unwrap_or_else(|_| "auto".to_string()) + .to_ascii_lowercase() + .as_str(), + "http2" | "legacy" | "off" + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use bytes::Bytes; + use std::pin::Pin; + use std::task::{Context, Poll}; + + struct PartialVectoredWriter { + data: Vec, + max_write: usize, + vectored_calls: usize, + max_slices_seen: usize, + } + + impl PartialVectoredWriter { + fn new(max_write: usize) -> Self { + Self { + data: Vec::new(), + max_write, + vectored_calls: 0, + max_slices_seen: 0, + } + } + } + + impl AsyncWrite for PartialVectoredWriter { + fn poll_write( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buffer: &[u8], + ) -> Poll> { + let count = buffer.len().min(self.max_write); + self.data.extend_from_slice(&buffer[..count]); + Poll::Ready(Ok(count)) + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buffers: &[IoSlice<'_>], + ) -> Poll> { + self.vectored_calls += 1; + self.max_slices_seen = self.max_slices_seen.max(buffers.len()); + let mut remaining = self.max_write; + let mut written = 0; + for buffer in buffers { + if remaining == 0 { + break; + } + let count = buffer.len().min(remaining); + self.data.extend_from_slice(&buffer[..count]); + written += count; + remaining -= count; + } + Poll::Ready(Ok(written)) + } + + fn is_write_vectored(&self) -> bool { + true + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + + #[tokio::test] + async fn vectored_writer_handles_partial_writes_and_many_buffers() { + let buffers = (0..200) + .map(|index| Bytes::from(vec![index as u8; 3])) + .collect(); + let message = TensorMessage::new(9, Bytes::from_static(b"manifest"), buffers); + let mut writer = PartialVectoredWriter::new(7); + + write_raw_tensor_frame(&mut writer, RawTensorKind::Ask, "/actors/test", &message) + .await + .unwrap(); + + assert!(writer.vectored_calls > 1); + assert_eq!(writer.max_slices_seen, MAX_VECTORED_SLICES); + + let mut wire = writer.data.as_slice(); + let decoded = read_raw_tensor_frame(&mut wire).await.unwrap(); + assert_eq!(decoded.kind, RawTensorKind::Ask); + assert_eq!(decoded.path, "/actors/test"); + assert_eq!(decoded.message.version, 9); + assert_eq!(&decoded.message.metadata[..], b"manifest"); + assert_eq!(decoded.message.buffers.len(), 200); + for (index, buffer) in decoded.message.buffers.iter().enumerate() { + assert_eq!(&buffer[..], &[index as u8; 3]); + } + } + + #[tokio::test] + async fn reader_rejects_excessive_buffer_count_before_allocating_table() { + let mut header = Vec::with_capacity(RAW_FIXED_HEADER_LEN); + header.extend_from_slice(RAW_TENSOR_MAGIC); + header.push(RawTensorKind::Ask as u8); + header.extend_from_slice(&[0; 3]); + header.extend_from_slice(&1u32.to_le_bytes()); + header.extend_from_slice(&0u32.to_le_bytes()); + header.extend_from_slice(&0u64.to_le_bytes()); + header.extend_from_slice(&((max_tensor_buffers() + 1) as u32).to_le_bytes()); + + let mut input = header.as_slice(); + let error = read_raw_tensor_frame(&mut input).await.unwrap_err(); + assert!(error.to_string().contains("buffer count exceeds")); + } + + #[tokio::test] + async fn reader_rejects_excessive_metadata_before_allocating_payload() { + let mut header = Vec::with_capacity(RAW_FIXED_HEADER_LEN); + header.extend_from_slice(RAW_TENSOR_MAGIC); + header.push(RawTensorKind::Ask as u8); + header.extend_from_slice(&[0; 3]); + header.extend_from_slice(&1u32.to_le_bytes()); + header.extend_from_slice(&0u32.to_le_bytes()); + header.extend_from_slice(&((max_tensor_metadata_bytes() + 1) as u64).to_le_bytes()); + header.extend_from_slice(&0u32.to_le_bytes()); + + let mut input = header.as_slice(); + let error = read_raw_tensor_frame(&mut input).await.unwrap_err(); + assert!(error.to_string().contains("metadata exceeds")); + } + + #[tokio::test] + async fn reader_rejects_excessive_total_before_allocating_payload() { + let mut input = Vec::with_capacity(RAW_FIXED_HEADER_LEN + 8); + input.extend_from_slice(RAW_TENSOR_MAGIC); + input.push(RawTensorKind::Ask as u8); + input.extend_from_slice(&[0; 3]); + input.extend_from_slice(&1u32.to_le_bytes()); + input.extend_from_slice(&0u32.to_le_bytes()); + input.extend_from_slice(&0u64.to_le_bytes()); + input.extend_from_slice(&1u32.to_le_bytes()); + // No payload follows. The advertised size must be rejected after the + // tiny length table is read and before any payload Vec is allocated. + input.extend_from_slice(&(max_tensor_wire_bytes() as u64).to_le_bytes()); + + let error = read_raw_tensor_frame(&mut input.as_slice()) + .await + .unwrap_err(); + assert!(error.to_string().contains("frame exceeds")); + } +} + +/// Observable payload-copy contract of a tensor transport backend. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum TensorCopyModel { + /// Compatibility fallback: one sender-side wire packing copy, plus possible + /// HTTP/2 receive-body coalescing before zero-copy Bytes slices are exposed. + PackedHttp2Compatibility, + /// Raw TCP data plane: application buffer to kernel and kernel to final + /// owned receive buffer only. + DirectTcp, + /// Reserved same-host process transport. + SharedMemory, +} + +/// Backend boundary for opaque metadata plus contiguous tensor buffers. +#[async_trait::async_trait] +pub trait TensorTransport: Send + Sync { + async fn request_tensor( + &self, + actor_id: &ActorId, + message: TensorMessage, + ) -> Result; + + async fn send_tensor(&self, actor_id: &ActorId, message: TensorMessage) -> Result<()>; + + fn copy_model(&self) -> TensorCopyModel; +} diff --git a/crates/pulsing-actor/src/watch.rs b/crates/pulsing-actor/src/watch.rs index 2cd97c8ac..cf2259155 100644 --- a/crates/pulsing-actor/src/watch.rs +++ b/crates/pulsing-actor/src/watch.rs @@ -201,6 +201,7 @@ impl ActorLifecycle { Message::Stream { default_msg_type, .. } => (default_msg_type, Vec::new()), + Message::Tensor(_) => unreachable!("termination messages are always packed singles"), }; // Send to all watchers diff --git a/crates/pulsing-actor/tests/common/fixtures.rs b/crates/pulsing-actor/tests/common/fixtures.rs index e04cb1d30..60cdfa87f 100644 --- a/crates/pulsing-actor/tests/common/fixtures.rs +++ b/crates/pulsing-actor/tests/common/fixtures.rs @@ -166,12 +166,20 @@ impl Http2ServerHandler for StreamingHandler { ), )); } + Ok(Message::Tensor(_)) => { + return Err(pulsing_actor::error::PulsingError::from( + pulsing_actor::error::RuntimeError::Other( + "Tensor messages cannot be nested in streams".into(), + ), + )); + } Err(e) => return Err(e), } } let response = format!("{}:collected:{} bytes", path, collected.len()).into_bytes(); Ok(Message::single("stream_echo", response)) } + Message::Tensor(message) => Ok(Message::Tensor(message)), } } diff --git a/crates/pulsing-bench-py/pyproject.toml b/crates/pulsing-bench-py/pyproject.toml index 99d7d98eb..0707135dd 100644 --- a/crates/pulsing-bench-py/pyproject.toml +++ b/crates/pulsing-bench-py/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "pulsing-bench" -version = "0.1.2" +version = "0.1.3" description = "Pulsing Benchmark - LLM Inference Benchmark Tool" readme = "../../README.md" authors = [ diff --git a/crates/pulsing-py/Cargo.toml b/crates/pulsing-py/Cargo.toml index 8531c9cf2..08a573474 100644 --- a/crates/pulsing-py/Cargo.toml +++ b/crates/pulsing-py/Cargo.toml @@ -28,6 +28,7 @@ futures = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } bincode = { workspace = true } +bytes = { workspace = true } thiserror = { workspace = true } tokio = { workspace = true } tracing = { workspace = true } diff --git a/crates/pulsing-py/src/actor.rs b/crates/pulsing-py/src/actor.rs index e8da2eaa7..a662e085e 100644 --- a/crates/pulsing-py/src/actor.rs +++ b/crates/pulsing-py/src/actor.rs @@ -1,7 +1,8 @@ //! Python bindings for the Pulsing Actor System +use bytes::Bytes; use futures::StreamExt; -use pulsing_actor::actor::{ActorId, ActorPath, NodeId}; +use pulsing_actor::actor::{ActorId, ActorPath, NodeId, TensorMessage}; use pulsing_actor::prelude::*; use pulsing_actor::supervision::{BackoffStrategy, RestartPolicy, SupervisionSpec}; use pyo3::exceptions::{PyRuntimeError, PyStopAsyncIteration, PyValueError}; @@ -9,7 +10,10 @@ use pyo3::prelude::*; use pyo3::types::{PyBytes, PyDict}; use serde::{Deserialize, Serialize}; use std::cmp::min; +use std::ffi::{c_char, c_int, c_void}; use std::net::SocketAddr; +use std::ptr; +use std::slice; use std::sync::Arc; use std::sync::Mutex as StdMutex; use tokio::sync::mpsc; @@ -403,10 +407,338 @@ impl PyMessage { payload: None, stream_reader: Some(Arc::new(TokioMutex::new(Some(stream)))), }, + Message::Tensor(tensor) => Self { + msg_type: pulsing_actor::actor::TENSOR_MESSAGE_TYPE.to_string(), + payload: Some(tensor.encode_wire().unwrap_or_default()), + stream_reader: None, + }, + } + } +} + +/// Holds a Python buffer export for as long as `Bytes` references it. +/// +/// Safety contract: callers must not mutate or resize the source buffer until +/// the send completes. PulsingQueue prepares immutable contiguous CPU buffers +/// before constructing TensorMessage. +struct PythonBufferOwner { + view: Box, +} + +// Py_buffer entered Python's Stable ABI in Python 3.11, while Pulsing still +// builds an abi3-py310 extension. CPython 3.10 nevertheless exposes the same +// public C buffer ABI, so declare the two functions locally instead of +// disabling abi3 or silently copying through PyBytes. +#[repr(C)] +struct RawPyBuffer { + buf: *mut c_void, + obj: *mut c_void, + len: isize, + itemsize: isize, + readonly: c_int, + ndim: c_int, + format: *mut c_char, + shape: *mut isize, + strides: *mut isize, + suboffsets: *mut isize, + internal: *mut c_void, +} + +impl RawPyBuffer { + fn new() -> Self { + Self { + buf: ptr::null_mut(), + obj: ptr::null_mut(), + len: 0, + itemsize: 0, + readonly: 0, + ndim: 0, + format: ptr::null_mut(), + shape: ptr::null_mut(), + strides: ptr::null_mut(), + suboffsets: ptr::null_mut(), + internal: ptr::null_mut(), } } } +unsafe extern "C" { + fn PyObject_GetBuffer(obj: *mut c_void, view: *mut RawPyBuffer, flags: c_int) -> c_int; + fn PyBuffer_Release(view: *mut RawPyBuffer); +} + +const PYBUF_SIMPLE: c_int = 0; + +// CPython buffer exports may be moved between the Python executor and Tokio +// send task. The exporter remains pinned by `view.obj`; callers must obey the +// documented no-mutation-while-sending contract. +unsafe impl Send for PythonBufferOwner {} + +impl PythonBufferOwner { + fn acquire(item: &Bound<'_, pyo3::PyAny>) -> PyResult { + let mut view = Box::new(RawPyBuffer::new()); + // SAFETY: `view` points to writable Py_buffer storage and `item` is a + // live Python object. On success CPython gives view.obj an owned ref. + let result = unsafe { PyObject_GetBuffer(item.as_ptr().cast(), &mut *view, PYBUF_SIMPLE) }; + if result != 0 { + return Err(PyErr::fetch(item.py())); + } + if view.len < 0 || (view.len > 0 && view.buf.is_null()) { + // SAFETY: acquisition succeeded, so it must be released exactly once. + unsafe { PyBuffer_Release(&mut *view) }; + return Err(PyValueError::new_err("Invalid Python buffer export")); + } + Ok(Self { view }) + } + + fn len(&self) -> usize { + self.view.len as usize + } +} + +impl AsRef<[u8]> for PythonBufferOwner { + fn as_ref(&self) -> &[u8] { + let len = self.len(); + if len == 0 { + return &[]; + } + // SAFETY: PyBuffer owns an active CPython buffer export, so its pointer + // remains valid until this owner is dropped by Bytes. The constructor + // only accepts C-contiguous byte views. + unsafe { slice::from_raw_parts(self.view.buf.cast::(), len) } + } +} + +impl Drop for PythonBufferOwner { + fn drop(&mut self) { + Python::with_gil(|_| { + // SAFETY: this owner is constructed only after a successful + // PyObject_GetBuffer and owns the one matching release. + unsafe { PyBuffer_Release(&mut *self.view) }; + }); + } +} + +/// Python owner used behind a ctypes array for received Rust `Bytes`. +/// +/// PyO3's custom buffer protocol slots are unavailable for abi3-py310. A +/// ctypes array is therefore used as the stable buffer exporter and retains +/// this object through a private attribute. No PyBytes payload copy occurs. +#[pyclass] +struct PyOwnedTensorBuffer { + _data: TensorBufferOwnerData, +} + +enum TensorBufferOwnerData { + /// Local mailbox path: retain the original Python export through Bytes. + Shared { _bytes: Bytes }, + /// Remote compatibility path: final independent, writable user buffer. + Owned { _bytes: Vec }, +} + +/// Turn an arbitrary contiguous buffer exporter into a one-dimensional byte +/// memoryview without copying. Unlike the legacy ZeroCopyDescriptor helper, +/// this deliberately rejects non-contiguous input instead of calling bytes(). +fn normalize_tensor_buffer( + py: Python<'_>, + item: &Bound<'_, pyo3::PyAny>, +) -> PyResult<(PyObject, usize)> { + let builtins = py.import("builtins")?; + let view = builtins + .getattr("memoryview")? + .call1((item,)) + .map_err(|_| { + PyValueError::new_err( + "TensorMessage.buffers items must implement the Python buffer protocol", + ) + })?; + let contiguous = view.getattr("c_contiguous")?.extract::()?; + if !contiguous { + return Err(PyValueError::new_err( + "TensorMessage.buffers items must be C-contiguous; make tensors contiguous in PulsingQueue", + )); + } + + let format = view.getattr("format")?.extract::()?; + let itemsize = view.getattr("itemsize")?.extract::()?; + let ndim = view.getattr("ndim")?.extract::()?; + let byte_view = if format == "B" && itemsize == 1 && ndim == 1 { + view + } else { + view.call_method1("cast", ("B",)).map_err(|_| { + PyValueError::new_err("TensorMessage.buffers items must support a contiguous byte view") + })? + }; + let buffer = PythonBufferOwner::acquire(&byte_view).map_err(|_| { + PyValueError::new_err("TensorMessage.buffers items must expose a contiguous byte buffer") + })?; + let len = buffer.len(); + drop(buffer); + Ok((byte_view.unbind(), len)) +} + +fn python_buffer_to_bytes(item: &PyObject, py: Python<'_>) -> PyResult { + let owner = PythonBufferOwner::acquire(item.bind(py))?; + Ok(Bytes::from_owner(owner)) +} + +/// Create a Python buffer exporter with explicit ownership semantics. +fn rust_bytes_to_python_buffer( + py: Python<'_>, + data: Bytes, + copy_into_owned_receive_buffer: bool, + writable: bool, +) -> PyResult { + if data.is_empty() { + let empty = if writable { + py.import("builtins")?.getattr("bytearray")?.call0()? + } else { + PyBytes::new(py, &[]).into_any() + }; + let view = py + .import("builtins")? + .getattr("memoryview")? + .call1((empty,))?; + return Ok(view.unbind()); + } + + let len = data.len(); + let (address, owner_data) = if copy_into_owned_receive_buffer { + // This is the compatibility HTTP/2 receive copy into the final + // user-owned allocation. The raw TCP backend will read directly into + // this allocation and remove this intermediate copy. + let owned = data.to_vec(); + ( + owned.as_ptr() as usize, + TensorBufferOwnerData::Owned { _bytes: owned }, + ) + } else { + ( + data.as_ptr() as usize, + TensorBufferOwnerData::Shared { _bytes: data }, + ) + }; + let owner = Py::new(py, PyOwnedTensorBuffer { _data: owner_data })?; + let ctypes = py.import("ctypes")?; + let array_type = ctypes.getattr("c_ubyte")?.call_method1("__mul__", (len,))?; + let array = array_type.call_method1("from_address", (address,))?; + array.setattr("_pulsing_owner", owner)?; + + let view = py + .import("builtins")? + .getattr("memoryview")? + .call1((array,))?; + let byte_view = view.call_method1("cast", ("B",))?; + let byte_view = if writable { + byte_view + } else { + byte_view.call_method0("toreadonly")? + }; + Ok(byte_view.unbind()) +} + +/// Opaque metadata plus contiguous CPU buffers optimized for tensor transport. +#[pyclass(name = "TensorMessage")] +#[derive(Clone)] +pub struct PyTensorMessage { + metadata: Vec, + buffers: Vec, + buffer_lengths: Vec, + version: u32, +} + +impl PyTensorMessage { + fn to_rust_message(&self, py: Python<'_>) -> PyResult { + let buffers = self + .buffers + .iter() + .map(|buffer| python_buffer_to_bytes(buffer, py)) + .collect::>>()?; + Ok(Message::Tensor(TensorMessage::new( + self.version, + Bytes::copy_from_slice(&self.metadata), + buffers, + ))) + } + + fn from_rust_message(py: Python<'_>, message: TensorMessage) -> PyResult { + let copy_into_owned_receive_buffer = message.requires_owned_receive_copy(); + let writable = copy_into_owned_receive_buffer || message.owns_receive_buffers(); + let metadata = message.metadata.to_vec(); + let buffer_lengths = message.buffers.iter().map(Bytes::len).collect(); + let buffers = message + .buffers + .into_iter() + .map(|buffer| { + rust_bytes_to_python_buffer(py, buffer, copy_into_owned_receive_buffer, writable) + }) + .collect::>>()?; + Ok(Self { + metadata, + buffers, + buffer_lengths, + version: message.version, + }) + } +} + +#[pymethods] +impl PyTensorMessage { + #[new] + #[pyo3(signature = (metadata, buffers, version=1))] + fn new( + py: Python<'_>, + metadata: Vec, + buffers: Vec, + version: u32, + ) -> PyResult { + let normalized = buffers + .into_iter() + .map(|buffer| normalize_tensor_buffer(py, buffer.bind(py))) + .collect::>>()?; + let (buffers, buffer_lengths): (Vec<_>, Vec<_>) = normalized.into_iter().unzip(); + Ok(Self { + metadata, + buffers, + buffer_lengths, + version, + }) + } + + #[getter] + fn metadata<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> { + PyBytes::new(py, &self.metadata) + } + + #[getter] + fn buffers(&self, py: Python<'_>) -> Vec { + self.buffers + .iter() + .map(|buffer| buffer.clone_ref(py)) + .collect() + } + + #[getter] + fn version(&self) -> u32 { + self.version + } + + #[getter] + fn total_bytes(&self) -> usize { + self.buffer_lengths.iter().sum() + } + + fn __repr__(&self) -> String { + format!( + "TensorMessage(version={}, metadata_bytes={}, buffers={}, total_bytes={})", + self.version, + self.metadata.len(), + self.buffers.len(), + self.total_bytes() + ) + } +} + /// Descriptor object for optional zerocopy payload transport. #[pyclass(name = "ZeroCopyDescriptor")] #[derive(Clone)] @@ -754,6 +1086,11 @@ async fn reassemble_zerocopy_stream( /// Small zerocopy payloads → `Message::Single`; large ones → `Message::Stream` /// (descriptor-first + chunked data). Non-zerocopy objects → pickle. pub(crate) fn encode_python_payload(py: Python<'_>, obj: &PyObject) -> PyResult { + if obj.bind(py).is_instance_of::() { + let tensor = obj.bind(py).extract::>()?; + return tensor.to_rust_message(py); + } + match zerocopy_mode().as_str() { "off" => Ok(Message::single(SEALED_PY_MSG_TYPE, pickle_object(py, obj)?)), "force" => { @@ -833,6 +1170,11 @@ pub(crate) async fn decode_message_to_pyobject(msg: Message) -> PyResult Python::with_gil(|py| { + let message = PyTensorMessage::from_rust_message(py, tensor)?; + let obj = Py::new(py, message)?; + Ok(obj.into_pyobject(py)?.into_any().unbind()) + }), _ => Python::with_gil(|py| { Ok(PyMessage::from_rust_message(msg) .into_pyobject(py)? @@ -2120,6 +2462,21 @@ impl PyActorSystem { } } +#[pyfunction] +fn tensor_transport_stats(py: Python<'_>) -> PyResult> { + let stats = pulsing_actor::transport::raw_tensor_transport_stats(); + let result = PyDict::new(py); + result.set_item("raw_frames_sent", stats.frames_sent)?; + result.set_item("raw_frames_received", stats.frames_received)?; + result.set_item("raw_bytes_sent", stats.bytes_sent)?; + result.set_item("raw_bytes_received", stats.bytes_received)?; + result.set_item("http2_fallback_frames", stats.http2_fallback_frames)?; + result.set_item("http2_fallback_bytes", stats.http2_fallback_bytes)?; + result.set_item("active_copy_model", stats.active_copy_model)?; + result.set_item("raw_connections_accepted", stats.raw_connections_accepted)?; + Ok(result.unbind()) +} + pub fn add_to_module(m: &Bound<'_, pyo3::types::PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; @@ -2132,5 +2489,7 @@ pub fn add_to_module(m: &Bound<'_, pyo3::types::PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_function(wrap_pyfunction!(tensor_transport_stats, m)?)?; Ok(()) } diff --git a/crates/pulsing-rpymod/src/bindings/codec.rs b/crates/pulsing-rpymod/src/bindings/codec.rs index 201057ed8..6e2bfd1be 100644 --- a/crates/pulsing-rpymod/src/bindings/codec.rs +++ b/crates/pulsing-rpymod/src/bindings/codec.rs @@ -146,7 +146,7 @@ pub async fn decode_message_to_pyobject( Ok(desc.into_ref(&vm.ctx).into()) } other => { - let py_msg = PyMessage::from_rust_message(other); + let py_msg = PyMessage::from_rust_message(other, vm)?; Ok(py_msg.into_ref(&vm.ctx).into()) } } diff --git a/crates/pulsing-rpymod/src/bindings/message.rs b/crates/pulsing-rpymod/src/bindings/message.rs index e4d853d8f..6db4af5dd 100644 --- a/crates/pulsing-rpymod/src/bindings/message.rs +++ b/crates/pulsing-rpymod/src/bindings/message.rs @@ -218,21 +218,25 @@ impl PyMessage { } } - pub(crate) fn from_rust_message(msg: Message) -> Self { + pub(crate) fn from_rust_message(msg: Message, vm: &VirtualMachine) -> PyResult { match msg { - Message::Single { msg_type, data } => Self { + Message::Single { msg_type, data } => Ok(Self { msg_type, payload: Some(data), stream_reader: None, - }, + }), Message::Stream { default_msg_type, stream, - } => Self { + } => Ok(Self { msg_type: default_msg_type, payload: None, stream_reader: Some(Arc::new(TokioMutex::new(Some(stream)))), - }, + }), + Message::Tensor(_) => { + Err(vm + .new_runtime_error("TensorMessage is not supported by the RustPython binding")) + } } } } diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index 560edb0f3..85cf517e7 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -131,6 +131,7 @@ plugins: Cluster Networking: 集群组网 Actor Addressing: Actor 寻址 HTTP2 Transport: HTTP2 传输 + TensorMessage Transport: TensorMessage 快速传输 Load Sync: 负载同步 AS Actor Decorator: AS Actor 装饰器 Communication Evolution: 集群内通信技术演进 @@ -221,6 +222,7 @@ nav: - Load Sync: design/load_sync.md - Implementation: - AS Actor Decorator: design/as-actor-decorator.md + - TensorMessage Transport: design/tensor-message-transport.md - Communication Evolution: design/cluster-communication-evolution.md - Out-Cluster Connect: design/out-cluster-connect.md - Forge: diff --git a/docs/pyproject.toml b/docs/pyproject.toml index 870695020..f6af2639c 100644 --- a/docs/pyproject.toml +++ b/docs/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "pulsing-docs" -version = "0.1.2" +version = "0.1.3" description = "Documentation for Pulsing - Backbone for Distributed AI Systems" readme = "README.md" requires-python = ">=3.10" diff --git a/docs/src/design/tensor-message-transport.md b/docs/src/design/tensor-message-transport.md new file mode 100644 index 000000000..da90d9a2b --- /dev/null +++ b/docs/src/design/tensor-message-transport.md @@ -0,0 +1,336 @@ +# TensorMessage Fast Transport + +## Status and scope + +This document describes the implemented CPU-to-CPU transport for +`TensorMessage`. It is the low-level data plane used by callers such as +PulsingQueue to move a TensorDict without first serializing all tensor payloads +into one Python or Rust byte string. + +The responsibility boundary is deliberate: + +| Layer | Responsibility | +|---|---| +| PulsingQueue or another caller | Move CUDA tensors to CPU, make non-contiguous tensors contiguous, define the TensorDict schema, and encode dtype, shape, stride, names, and byte order in metadata. | +| Pulsing | Route an opaque metadata blob plus ordered, contiguous CPU buffers; preserve buffer ownership; select the transport; and enforce wire limits. | +| Receiving application | Decode the metadata, construct tensor views over the received buffers, and define business-level completion or durability semantics. | + +Pulsing does **not** inspect or convert tensor dtype, shape, stride, byte order, +or TensorDict structure. The `version` field is transported unchanged; it is +currently not capability-negotiated and unknown versions are not rejected by +the transport. + +## Goals + +- Avoid pickle for tensor payloads. +- Avoid a combined userspace payload allocation on the clear-text TCP path. +- Read each remote payload directly into the allocation later exposed to + Python and downstream tensor views. +- Keep the local mailbox path copy-free for payloads by retaining the original + Python buffer exporter. +- Support multiple ordered buffers and a small opaque metadata document. +- Preserve ordinary HTTP/2 operation as a compatibility path, including TLS. +- Give the receive allocation an explicit lifetime tied to Python objects. + +## Non-goals + +The current implementation does not provide: + +- GPU Direct, CUDA IPC, or GPU-to-GPU transport; +- same-host shared memory transport (the transport abstraction reserves it); +- tensor schema encoding, dtype conversion, or byte-order conversion; +- checksum, compression, encryption on raw TCP, or durable storage semantics; +- raw protocol capability negotiation, request IDs, transparent retry, or + automatic raw-to-HTTP/2 fallback; +- a streaming response made of `TensorMessage` frames. + +## Public API + +### Python + +```python +import pulsing as pul + +message = pul.TensorMessage( + metadata=b"opaque schema bytes", + buffers=[memoryview(first), memoryview(second)], + version=1, +) + +print(message.metadata) +print(message.buffers) +print(message.version) +print(message.total_bytes) # payload buffers only; metadata is excluded +``` + +Every item in `buffers` must implement the Python buffer protocol, be +C-contiguous, and be castable to a one-dimensional byte view. Pulsing rejects a +non-contiguous buffer instead of silently materializing a copy. An empty buffer +list is valid for a metadata-only control message. + +For an actor declared with `@pul.remote`, implement `receive_tensor()` and send +the value through the underlying actor reference: + +```python +@pul.remote +class Sink: + async def receive_tensor(self, message: pul.TensorMessage): + # Decode metadata and consume message.buffers here. + return {"ok": True, "received_bytes": message.total_bytes} + +sink = await Sink.resolve(...) +ack = await sink.ref.ask(message) +``` + +Calling a generated proxy method such as `sink.put(message)` uses Pulsing's +ordinary method-call/pickle envelope. `sink.ref.ask(message)` and +`sink.ref.tell(message)` select the dedicated tensor message path. + +A class derived directly from `pulsing.core.Actor` receives `TensorMessage` in +its normal `receive()` method. + +### Rust + +At the transport boundary, the representation is: + +```rust +pub struct TensorMessage { + pub version: u32, + pub metadata: bytes::Bytes, + pub buffers: Vec, +} +``` + +The public transport model distinguishes `DirectTcp`, +`PackedHttp2Compatibility`, and the reserved `SharedMemory` copy model. These +are Pulsing transport choices, not PulsingQueue storage backends. Both the +Python and Rust representations treat metadata as opaque. + +## Data paths + +```mermaid +flowchart LR + Q["Caller: CPU + contiguous buffers"] --> M["TensorMessage"] + M -->|same process| L["Actor mailbox: shared buffer owner"] + M -->|clear-text remote| R["Raw TCP: vectored write"] + R --> O["Final receive allocations"] + M -->|TLS or forced compatibility| H["Packed HTTP/2 body"] + H --> C["Final Python receive buffers"] + L --> V["Tensor views"] + O --> V + C --> V +``` + +### Same process + +The Python binding acquires a `Py_buffer` lease and creates Rust `Bytes` backed +by that owner. A local mailbox message retains the same storage; it does not +materialize a payload copy. The receive-side memoryview is read-only, but it +still aliases the source allocation. Mutating the source through another +writable reference remains visible to the receiver. Clone before constructing +the message when snapshot semantics are required. + +### Remote clear-text TCP + +Clear-text connections use the raw data plane by default. It shares the actor +system's listening address with HTTP/2; the server peeks for the `PTR1` magic +and dispatches the connection to the raw tensor handler. Connections are +long-lived, pooled per remote address, and configured with `TCP_NODELAY`. + +The sender uses vectored writes over the header, actor path, metadata, and the +original buffers. It does not first concatenate all payloads. A logical frame +may require multiple `write_vectored` calls, and TCP may split it into any +number of packets; vectored I/O does not imply one syscall or one packet. + +The receiver validates advertised lengths and then calls `read_exact` for each +buffer into its final `Vec`. Moving those vectors into Rust `Bytes` and +exposing them as writable Python memoryviews does not make another payload +copy. + +This is a **direct TCP payload copy model**, not kernel-bypass zero-copy. For +remote payload bytes, the intended model is: + +1. application buffer to the sender's TCP/kernel buffer; +2. receiver's TCP/kernel buffer to the final userspace receive allocation. + +Small headers and metadata have their own allocations/copies and are not part +of this two-payload-copy claim. In particular, Python-to-Rust and +Rust-to-Python metadata conversion copies the metadata. + +### HTTP/2 compatibility path + +TLS connections and `PULSING_TENSOR_TRANSPORT=http2` use the existing HTTP/2 +transport. The sender packs the compatibility wire representation into one +body. The HTTP stack aggregates the body, and the Python binding copies each +decoded buffer into its final writable receive allocation. This path is +compatible but has more payload copies and must not be described as the direct +or zero-copy path. + +## Buffer ownership and lifetime + +| Path | Receive storage | Aliases sender? | Python view | +|---|---|---|---| +| Same-process mailbox | Original exported Python storage retained through `Bytes` | Yes | Read-only alias | +| Remote raw TCP | One final allocation per received buffer | No | Writable | +| Remote HTTP/2 | Final writable allocation copied from the packed body | No | Writable | + +The sender-side `TensorMessage` retains each Python buffer export. The caller +must not mutate or resize the source while `ask()` or `tell()` is in progress. +After a remote call completes, the peer has independent storage. For a local +call, aliasing continues for as long as either side retains a reference. +Because local `tell()` returns after enqueue rather than actor completion, it +does not provide a safe point for reusing mutable source storage; keep the +source immutable or use `ask()` for an explicit application completion point. + +On receive, the memoryview owns a private Python object that owns the Rust +buffer. A downstream `torch.frombuffer()` view retains that owner, so the +allocation remains alive after the `TensorMessage` object itself is released. + +An upstream PyTorch preparation step can look like this: + +```python +cpu_tensor = tensor.detach().to("cpu").contiguous() +byte_view = memoryview(cpu_tensor.view(torch.uint8).numpy()).cast("B") +message = pul.TensorMessage(metadata, [byte_view]) +``` + +Pulsing deliberately does not call `.cpu()` or `.contiguous()` internally. + +## Request, response, and ACK semantics + +`ask(TensorMessage)` waits for actor handling and accepts either: + +- a `TensorMessage` response on the dedicated data plane; or +- a normal single response, useful for a small application ACK. + +It does not accept a streaming response on the raw path. + +Remote raw `tell(TensorMessage)` waits for a protocol ACK, but the server sends +that ACK after the message has been enqueued in the actor mailbox. It does not +mean the actor finished processing, wrote a storage bucket, replicated data, or +made data durable. Local `tell()` has the same enqueue-only semantic. + +When a producer must know that a storage actor has completed its write, use +`ask()` and have the actor return an empty or small application ACK only after +the write is complete. Pulsing supplies transport completion; transaction, +visibility, idempotency, and durability remain application responsibilities. + +The raw connection executes one request followed by one response while it is +checked out of the pool. The whole exchange uses the configured stream timeout. +A failed connection is discarded. Pulsing does not transparently replay the +message; retry and idempotency must be defined by the caller. + +## Wire formats + +All integer fields below use little-endian encoding. Tensor payload byte order +is not converted and must be declared by the caller's metadata schema. + +### Raw TCP frame (`PTR1`) + +```text +magic[4] = "PTR1" +kind: u8 +reserved[3] +version: u32 +actor_path_len: u32 +metadata_len: u64 +buffer_count: u32 +buffer_lengths[buffer_count]: u64[] +actor_path[actor_path_len] +metadata[metadata_len] +buffers... +``` + +Frame kinds are `Ask`, `Tell`, `TensorResponse`, `Ack`, `Error`, and +`SingleResponse`. A `SingleResponse` uses the path slot for the ordinary message +type and the metadata slot for its serialized data. + +### HTTP/2 compatibility body (`PTM1`) + +```text +magic[4] = "PTM1" +version: u32 +metadata_len: u64 +buffer_count: u32 +buffer_lengths[buffer_count]: u64[] +metadata[metadata_len] +buffers... +``` + +HTTP/2 carries the actor route and request mode separately, so they are not in +the compatibility body. + +## Selection and deployment compatibility + +| Configuration | Selected path | +|---|---| +| Clear-text connection; variable unset or `auto`/`raw` | Raw TCP | +| `PULSING_TENSOR_TRANSPORT=http2`, `legacy`, or `off` | HTTP/2 compatibility | +| TLS enabled | HTTP/2 compatibility | + +There is no raw protocol handshake and no automatic downgrade if a peer does +not understand `PTR1`. Deploy raw mode with compatible Pulsing versions on both +ends. Force HTTP/2 during a mixed-version migration. + +## Limits and validation + +The receiver validates counts and total sizes before allocating large payload +buffers: + +| Environment variable | Default | Meaning | +|---|---:|---| +| `PULSING_MAX_TENSOR_WIRE_BYTES` | 64 GiB | Maximum complete tensor frame/body | +| `PULSING_MAX_TENSOR_METADATA_BYTES` | 64 MiB | Maximum opaque metadata size | +| `PULSING_MAX_TENSOR_BUFFERS` | 65,536 | Maximum buffers per message | + +Raw actor paths also have a fixed 64 KiB limit. These controls bound allocation +but do not authenticate metadata or validate a caller-defined schema. + +## Observability + +`pulsing.tensor_transport_stats()` returns process-global cumulative counters: + +- `raw_frames_sent`, `raw_frames_received`; +- `raw_bytes_sent`, `raw_bytes_received`; +- `http2_fallback_frames`, `http2_fallback_bytes`; +- `raw_connections_accepted`; +- `active_copy_model`. + +`active_copy_model` is the last path recorded by the process, with values +`unused`, `direct_tcp`, or `packed_http2_compatibility`. It is not a per-message +label and can change when concurrent traffic uses different paths. + +## Performance guidance + +- Prepare CPU-contiguous buffers before creating `TensorMessage`. +- Keep metadata compact and version its schema explicitly. +- Preserve tensor order in metadata because Pulsing only preserves buffer order. +- Merge very many small tensors upstream when header processing and many small + I/O slices dominate payload transfer. +- Prefer a small normal ACK rather than echoing tensor payloads for storage PUT. +- Benchmark raw and HTTP/2 separately and inspect transport stats to confirm the + path under test. + +## Runnable example and tests + +Run the [TensorMessage fast-path example](../../../examples/python/tensor_message_fast_path.py): + +```bash +python examples/python/tensor_message_fast_path.py +python examples/python/tensor_message_fast_path.py --transport http2 +``` + +The example creates two ActorSystems in one process, forces actor resolution to +the remote node ID, sends one `float32-native` buffer over TCP, and waits for a +small application-level ACK. The implementation tests additionally cover local +aliasing, source lifetime, raw/HTTP2 selection, writable remote buffers, tell +ACK, pooled reconnect, size validation, and the `receive_tensor()` hook. + +## Future work + +The `TensorTransport` abstraction reserves a `SharedMemory` copy model. A +same-host transport can add shared-memory registration and lease/release +semantics without changing `TensorMessage` or the PulsingQueue schema. Capability +negotiation, checksums, request IDs, and explicit retry policy should be added +together so fallback and failure semantics remain observable rather than +silently replaying a non-idempotent storage operation. diff --git a/docs/src/design/tensor-message-transport.zh.md b/docs/src/design/tensor-message-transport.zh.md new file mode 100644 index 000000000..a44eab1e3 --- /dev/null +++ b/docs/src/design/tensor-message-transport.zh.md @@ -0,0 +1,304 @@ +# TensorMessage 快速传输设计 + +## 状态与范围 + +本文描述 Pulsing 已实现的 CPU 到 CPU `TensorMessage` 传输。它是供 +PulsingQueue 等上层调用者传输 TensorDict 的底层数据面，目标是避免先把所有 Tensor +payload 序列化并拼接成一个 Python 或 Rust 大字节串。 + +职责边界如下： + +| 层 | 职责 | +|---|---| +| PulsingQueue 或其他调用方 | 把 CUDA Tensor 搬到 CPU；将非连续 Tensor 转为连续；定义 TensorDict schema；把 dtype、shape、stride、名称和字节序写入 metadata。 | +| Pulsing | 路由不透明 metadata 和有序的连续 CPU buffers；维持 buffer 所有权；选择传输后端；执行 wire size 限制。 | +| 接收应用 | 解析 metadata；基于接收 buffer 创建 Tensor view；定义业务完成和持久化语义。 | + +Pulsing **不解析也不转换** Tensor 的 dtype、shape、stride、字节序或 TensorDict +结构。`version` 只会原样传递；当前协议不会协商版本，也不会因为未知版本而拒绝消息。 + +## 目标 + +- Tensor payload 不走 pickle。 +- 明文 TCP 主路径不创建合并后的 userspace payload 大包。 +- 远端接收时直接读入最终暴露给 Python 和下游 Tensor view 的内存。 +- 同进程 mailbox 路径通过持有原 Python buffer exporter，避免 payload copy。 +- 支持一段小型不透明 metadata 和多段有序 buffer。 +- TLS 等场景保留 HTTP/2 兼容路径。 +- 接收内存拥有明确且由 Python 引用链管理的生命周期。 + +## 非目标 + +当前实现不提供： + +- GPU Direct、CUDA IPC 或 GPU 到 GPU 传输； +- 同节点共享内存传输（传输抽象中已经预留）； +- Tensor schema 编解码、dtype 转换或字节序转换； +- checksum、压缩、raw TCP 加密或持久化存储语义； +- raw 协议能力协商、request ID、透明重试或 raw 到 HTTP/2 自动 fallback； +- 由 `TensorMessage` frame 组成的流式响应。 + +## 公开 API + +### Python + +```python +import pulsing as pul + +message = pul.TensorMessage( + metadata=b"opaque schema bytes", + buffers=[memoryview(first), memoryview(second)], + version=1, +) + +print(message.metadata) +print(message.buffers) +print(message.version) +print(message.total_bytes) # 只统计 payload buffers，不含 metadata +``` + +`buffers` 中每一项都必须实现 Python buffer protocol、保持 C-contiguous，并且能够 +cast 成一维 byte view。遇到非连续 buffer，Pulsing 会直接报错，不会静默补一次 copy。 +空 buffer list 是合法的，可用于只携带 metadata 的控制消息。 + +使用 `@pul.remote` 声明 actor 时，实现 `receive_tensor()`，并通过底层 actor +reference 发送： + +```python +@pul.remote +class Sink: + async def receive_tensor(self, message: pul.TensorMessage): + # 在这里解析 metadata 并消费 message.buffers。 + return {"ok": True, "received_bytes": message.total_bytes} + +sink = await Sink.resolve(...) +ack = await sink.ref.ask(message) +``` + +调用 `sink.put(message)` 之类的生成代理方法，会使用 Pulsing 普通方法调用的 pickle +envelope。只有 `sink.ref.ask(message)` 和 `sink.ref.tell(message)` 会选择专用 Tensor +消息路径。 + +直接继承 `pulsing.core.Actor` 的类，则在普通 `receive()` 中收到 `TensorMessage`。 + +### Rust + +传输边界的数据表示为： + +```rust +pub struct TensorMessage { + pub version: u32, + pub metadata: bytes::Bytes, + pub buffers: Vec, +} +``` + +公开传输模型区分 `DirectTcp`、`PackedHttp2Compatibility` 和预留的 +`SharedMemory` copy model。它们是 Pulsing 的物理传输选择，不是 PulsingQueue 的存储后端。 +Python 与 Rust 表示都把 metadata 当作不透明数据。 + +## 数据路径 + +```mermaid +flowchart LR + Q["调用方: CPU + 连续 buffers"] --> M["TensorMessage"] + M -->|同进程| L["Actor mailbox: 共享 buffer owner"] + M -->|明文远端| R["Raw TCP: vectored write"] + R --> O["最终接收内存"] + M -->|TLS 或强制兼容模式| H["打包后的 HTTP/2 body"] + H --> C["最终 Python 接收 buffers"] + L --> V["Tensor views"] + O --> V + C --> V +``` + +### 同进程 + +Python binding 获取 `Py_buffer` lease，并创建由该 owner 支撑的 Rust `Bytes`。本地 +mailbox 消息保留相同 storage，不会物化一份 payload copy。接收端 memoryview 是只读 +的，但仍与源内存 alias；通过其他可写引用修改源数据，接收方仍能看到变化。如果需要 +snapshot 语义，调用方必须在构造消息前 clone。 + +### 远端明文 TCP + +明文连接默认使用 raw 数据面。它与 HTTP/2 共用 ActorSystem 的监听地址；服务端通过 +peek `PTR1` magic 将连接分发给 raw Tensor handler。连接按远端地址长期池化，并启用 +`TCP_NODELAY`。 + +发送端对 header、actor path、metadata 和原始 buffers 使用 vectored write，不会先把 +所有 payload 拼接起来。一个逻辑 frame 可能需要多次 `write_vectored`，TCP 也可以把它 +拆成任意数量的 packet；vectored I/O 不等于一次 syscall 或一个 TCP packet。 + +接收端先验证声明的长度，再对每个 buffer 调用 `read_exact`，直接写入最终 `Vec`。 +随后把这些 vector move 给 Rust `Bytes` 并暴露为可写 Python memoryview，不会再产生一份 +payload copy。 + +这是一种 **direct TCP payload copy model**，不是绕过内核的 zero-copy。远端 payload +数据的目标模型是： + +1. 应用 buffer 到发送端 TCP/kernel buffer； +2. 接收端 TCP/kernel buffer 到最终 userspace 接收内存。 + +小型 header 和 metadata 有自己的分配/copy，不属于上述“两次 payload copy”的统计口径。 +尤其是 Python 到 Rust、Rust 到 Python 的 metadata 转换会复制 metadata。 + +### HTTP/2 兼容路径 + +TLS 连接以及 `PULSING_TENSOR_TRANSPORT=http2` 使用原有 HTTP/2 传输。发送方先把兼容 +wire 表示打包成一个 body；HTTP stack 会聚合 body；Python binding 再把每个解码后的 +buffer 复制进最终可写接收内存。该路径具有兼容性，但 payload copy 更多，不能称为 direct +或 zero-copy 路径。 + +## Buffer 所有权和生命周期 + +| 路径 | 接收 storage | 是否与发送端 alias | Python view | +|---|---|---|---| +| 同进程 mailbox | 通过 `Bytes` 保留的原 Python export | 是 | 只读 alias | +| 远端 raw TCP | 每个 buffer 一块最终接收内存 | 否 | 可写 | +| 远端 HTTP/2 | 从打包 body 复制出的最终可写内存 | 否 | 可写 | + +发送端 `TensorMessage` 会保留每一项 Python buffer export。在 `ask()` 或 `tell()` 尚未 +完成时，调用方不得修改或 resize 源 buffer。远端调用完成后，对端已经持有独立 storage； +本地调用则会持续 alias，直到双方都释放引用。 +本地 `tell()` 在 enqueue 后而不是 actor 处理完成后返回，因此它不能提供安全复用可变源内存 +的时刻；调用方应保持源数据不可变，或改用 `ask()` 获得明确的应用完成点。 + +接收端 memoryview 持有一个私有 Python owner，而该 owner 持有 Rust buffer。下游 +`torch.frombuffer()` view 会继续保留这个 owner，因此释放 `TensorMessage` 本身后，内存 +依然有效。 + +上层的 PyTorch 准备逻辑可以类似： + +```python +cpu_tensor = tensor.detach().to("cpu").contiguous() +byte_view = memoryview(cpu_tensor.view(torch.uint8).numpy()).cast("B") +message = pul.TensorMessage(metadata, [byte_view]) +``` + +Pulsing 内部不会调用 `.cpu()` 或 `.contiguous()`。 + +## Request、response 与 ACK 语义 + +`ask(TensorMessage)` 会等待 actor 处理完成，并允许 actor 返回： + +- 走专用数据面的 `TensorMessage`；或 +- 普通 single response，适合用作小型应用层 ACK。 + +raw 路径不允许返回 streaming response。 + +远端 raw `tell(TensorMessage)` 会等待一个协议 ACK，但服务端在消息进入 actor mailbox 后 +就会发送该 ACK。它不表示 actor 已经处理完成、storage bucket 已写完、数据已复制或已经 +持久化。本地 `tell()` 同样只有 enqueue 语义。 + +如果 producer 必须确认 storage actor 已经写完，应使用 `ask()`，并让 actor 在写入完成后 +才返回空 ACK 或小型应用 ACK。Pulsing 提供传输完成能力；事务性、可见性、幂等和持久化 +仍由应用定义。 + +raw 连接从连接池 checkout 后，严格执行一次 request 和一次 response。整个交换使用配置的 +stream timeout。失败连接会被丢弃；Pulsing 不会透明重放消息，重试和幂等必须由调用方定义。 + +## Wire format + +下面的所有整数都使用 little-endian 编码。Tensor payload 的字节序不会被转换，必须由调用方 +在 metadata schema 中声明。 + +### Raw TCP frame（`PTR1`） + +```text +magic[4] = "PTR1" +kind: u8 +reserved[3] +version: u32 +actor_path_len: u32 +metadata_len: u64 +buffer_count: u32 +buffer_lengths[buffer_count]: u64[] +actor_path[actor_path_len] +metadata[metadata_len] +buffers... +``` + +frame kind 包括 `Ask`、`Tell`、`TensorResponse`、`Ack`、`Error` 和 +`SingleResponse`。`SingleResponse` 使用 path 字段保存普通 message type，使用 metadata +字段保存其序列化数据。 + +### HTTP/2 兼容 body（`PTM1`） + +```text +magic[4] = "PTM1" +version: u32 +metadata_len: u64 +buffer_count: u32 +buffer_lengths[buffer_count]: u64[] +metadata[metadata_len] +buffers... +``` + +HTTP/2 会单独携带 actor route 和 request mode，因此兼容 body 不包含这两项。 + +## 路径选择与部署兼容性 + +| 配置 | 使用的路径 | +|---|---| +| 明文连接；环境变量未设置或为 `auto`/`raw` | Raw TCP | +| `PULSING_TENSOR_TRANSPORT=http2`、`legacy` 或 `off` | HTTP/2 兼容路径 | +| 启用 TLS | HTTP/2 兼容路径 | + +raw 协议没有握手；对端不识别 `PTR1` 时也不会自动降级。使用 raw 模式时，两端必须部署兼容 +版本；混合版本迁移期间应显式强制 HTTP/2。 + +## 限制与校验 + +接收端会在分配大 payload buffer 前校验数量和总长度： + +| 环境变量 | 默认值 | 含义 | +|---|---:|---| +| `PULSING_MAX_TENSOR_WIRE_BYTES` | 64 GiB | 单个完整 Tensor frame/body 的最大值 | +| `PULSING_MAX_TENSOR_METADATA_BYTES` | 64 MiB | 不透明 metadata 最大值 | +| `PULSING_MAX_TENSOR_BUFFERS` | 65,536 | 单消息最大 buffer 数量 | + +raw actor path 另有固定 64 KiB 上限。这些限制只约束内存分配，不负责认证 metadata 或验证 +调用方定义的 schema。 + +## 可观测性 + +`pulsing.tensor_transport_stats()` 返回进程级累计统计： + +- `raw_frames_sent`、`raw_frames_received`； +- `raw_bytes_sent`、`raw_bytes_received`； +- `http2_fallback_frames`、`http2_fallback_bytes`； +- `raw_connections_accepted`； +- `active_copy_model`。 + +`active_copy_model` 表示该进程最后一次记录的路径，可取 `unused`、`direct_tcp` 或 +`packed_http2_compatibility`。它不是每条消息独立的标签；并发流量走不同路径时会变化。 + +## 性能建议 + +- 创建 `TensorMessage` 前准备好 CPU-contiguous buffers。 +- metadata 应保持紧凑，并显式维护 schema 版本。 +- Pulsing 只保证 buffer 顺序，因此 Tensor 顺序必须记录在 metadata 中。 +- 小 Tensor 数量非常多时，在上层合并它们，避免 header 处理和大量小 I/O slice 成为瓶颈。 +- storage PUT 建议只返回小型普通 ACK，不要回传整份 Tensor payload。 +- raw 与 HTTP/2 应分别 benchmark，并通过 transport stats 确认实际使用的路径。 + +## 可运行示例与测试 + +运行 [TensorMessage fast-path 示例](../../../examples/python/tensor_message_fast_path.py)： + +```bash +python examples/python/tensor_message_fast_path.py +python examples/python/tensor_message_fast_path.py --transport http2 +``` + +示例在同一进程创建两个 ActorSystem，指定远端 node ID 完成 actor resolve，通过 TCP 发送一段 +`float32-native` buffer，并等待小型应用层 ACK。实现测试还覆盖本地 alias、源数据生命周期、 +raw/HTTP2 路径选择、远端可写 buffer、tell ACK、连接池重连、size 校验和 +`receive_tensor()` 分发入口。 + +## 后续工作 + +`TensorTransport` 抽象已经预留 `SharedMemory` copy model。后续可在不改变 +`TensorMessage` 和 PulsingQueue schema 的前提下，增加同节点共享内存注册及 lease/release +语义。能力协商、checksum、request ID 和明确的 retry 策略应配套设计，避免系统静默重放 +非幂等 storage 操作。 diff --git a/examples/README.md b/examples/README.md index b6a60cc14..753b0ed0c 100644 --- a/examples/README.md +++ b/examples/README.md @@ -65,11 +65,13 @@ cd examples/agent/langgraph && ./run_distributed.sh | `named_actors.py` | 服务发现 | | `cluster.py` | 集群通信 | | `remote_actor_example.py` | @remote 装饰器 | +| `tensor_message_fast_path.py` | TensorMessage raw TCP 快速传输与 HTTP/2 兼容路径 | | `subprocess_example.py` | `subprocess` 兼容 API | ```bash python examples/python/ping_pong.py python examples/python/cluster.py --port 8000 +python examples/python/tensor_message_fast_path.py python examples/python/subprocess_example.py USE_POLSING_SUBPROCESS=1 python examples/python/subprocess_example.py --resources ``` @@ -134,6 +136,7 @@ cargo run --example behavior_fsm -p pulsing-actor | AI 辩论/讨论 | `agent/pulsing/mbti_discussion.py` | | 并行任务竞争 | `agent/pulsing/parallel_ideas_async.py` | | 集群部署 | `python/cluster.py` | +| 传输连续 CPU Tensor | `python/tensor_message_fast_path.py` | | 子进程资源调度 | `python/subprocess_example.py` | | 学习 CLI 工具 | `inspect/demo_service.py` | | 接入 AutoGen | `agent/autogen/` | diff --git a/examples/python/README.md b/examples/python/README.md index 22636428a..ee863a697 100644 --- a/examples/python/README.md +++ b/examples/python/README.md @@ -35,6 +35,8 @@ python examples/python/ping_pong.py # Basic communication python examples/python/message_patterns.py # RPC and streaming python examples/python/named_actors.py # Service discovery python examples/python/cluster.py # Multi-node (see --help) +python examples/python/tensor_message_fast_path.py # Tensor buffers over raw TCP +python examples/python/tensor_message_fast_path.py --transport http2 # Compatibility path python examples/python/subprocess_example.py # Native subprocess-compatible API USE_POLSING_SUBPROCESS=1 python examples/python/subprocess_example.py --resources # Pulsing backend python examples/python/isolated_actor_spawn.py # Actor in child OS process; cluster sees parent bridge @@ -46,6 +48,11 @@ python examples/python/forge_agent_quickstart.py # ForgeAgent demo (no API python examples/python/workspace_minimal_demo.py # same, single Python process ``` +`tensor_message_fast_path.py` 会创建两个 ActorSystem，确保请求真实经过 TCP。示例使用 +标准库 `array` 构造连续 CPU buffer，通过 `TensorMessage` 发送，并在 storage-like actor +处理完数据后返回小型应用 ACK。完整语义见 +[TensorMessage 快速传输设计](../../docs/src/design/tensor-message-transport.zh.md)。 + 同步包装器说明： - `transfer_queue.get_client()`、`queue.sync()`、`pulsing.subprocess` 的 Pulsing 后端都依赖 Pulsing 自己维护的 event loop。 diff --git a/examples/python/tensor_message_fast_path.py b/examples/python/tensor_message_fast_path.py new file mode 100644 index 000000000..5e1a1b05b --- /dev/null +++ b/examples/python/tensor_message_fast_path.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python3 +"""Send contiguous CPU tensor buffers over Pulsing's TensorMessage fast path. + +The example uses two ActorSystems so the request always crosses TCP. It keeps +the schema in opaque JSON metadata and uses only the standard library; a real +TensorDict integration can replace ``array`` with contiguous CPU tensor views. + +Usage: + python examples/python/tensor_message_fast_path.py + python examples/python/tensor_message_fast_path.py --transport http2 +""" + +import argparse +import asyncio +import json +import os +import sys +from array import array + +import pulsing as pul + + +@pul.remote +class TensorSink: + """Consume a TensorMessage and return a small application-level ACK.""" + + async def receive_tensor(self, message: pul.TensorMessage) -> dict: + metadata = json.loads(message.metadata) + if metadata["dtype"] != "float32-native": + raise ValueError(f"unsupported dtype: {metadata['dtype']}") + + values = message.buffers[0].cast("f") + expected_values = metadata["shape"][0] * metadata["shape"][1] + if len(values) != expected_values: + raise ValueError( + f"shape expects {expected_values} values, received {len(values)}" + ) + + # Returning only after processing gives ask() application-level + # completion semantics. A storage actor would write its bucket here. + return { + "ok": True, + "received_bytes": message.total_bytes, + "sum": sum(values), + "version": message.version, + } + + +async def run(transport: str) -> None: + os.environ["PULSING_TENSOR_TRANSPORT"] = transport + + server = await pul.actor_system(addr="127.0.0.1:0") + client = None + try: + actor_name = "tensor-message-fast-path" + await TensorSink.spawn(system=server, name=actor_name, public=True) + + client = await pul.actor_system( + addr="127.0.0.1:0", + seeds=[server.addr], + ) + sink = await TensorSink.resolve( + name=actor_name, + system=client, + node_id=server.node_id.id, + timeout=5.0, + ) + if sink.ref.is_local(): + raise RuntimeError("example must resolve a remote actor") + + source = array("f", range(12)) + metadata = json.dumps( + { + "dtype": "float32-native", + "shape": [3, 4], + "byte_order": sys.byteorder, + }, + separators=(",", ":"), + ).encode("utf-8") + message = pul.TensorMessage( + metadata=metadata, + buffers=[memoryview(source).cast("B")], + version=1, + ) + + # Use .ref.ask() for TensorMessage. Calling a generated proxy method + # would put the value inside the ordinary pickle method envelope. + ack = await sink.ref.ask(message) + assert ack == { + "ok": True, + "received_bytes": 48, + "sum": 66.0, + "version": 1, + } + + stats = pul.tensor_transport_stats() + print(f"transport={transport}") + print(f"ack={ack}") + print(f"tensor_transport_stats={stats}") + finally: + if client is not None: + await client.shutdown() + await server.shutdown() + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--transport", + choices=("raw", "http2"), + default="raw", + help="raw uses the direct TCP path; http2 exercises compatibility mode", + ) + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + asyncio.run(run(args.transport)) diff --git a/pyproject.toml b/pyproject.toml index 4a0b78fa9..aed2644e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "pulsing" -version = "0.1.2" +version = "0.1.3" description = "Pulsing: Backbone for distributed AI systems. Actor runtime with streaming, zero dependencies, and built-in discovery." readme = "README.md" authors = [ diff --git a/python/pulsing/__init__.py b/python/pulsing/__init__.py index 23d0f5438..a6b5902fa 100644 --- a/python/pulsing/__init__.py +++ b/python/pulsing/__init__.py @@ -22,12 +22,13 @@ def incr(self): self.value += 1; return self.value import asyncio from typing import Any -__version__ = "0.1.2" +__version__ = "0.1.3" # Submodule first: ``@remote`` must be bound before anything that might # re-enter this package during ``pulsing.core`` import. from pulsing.core.remote import Actor, ActorClass, remote, resolve from pulsing.core import ( + _HAS_NATIVE_TENSOR_TRANSPORT, init, shutdown, get_system, @@ -55,6 +56,9 @@ def incr(self): self.value += 1; return self.value ) from . import transfer_queue +if _HAS_NATIVE_TENSOR_TRANSPORT: + from pulsing.core import TensorMessage, tensor_transport_stats + bootstrap.stop = bootstrap_stop @@ -434,3 +438,6 @@ async def read(self, topic, **kwargs): "PulsingTimeoutError", "PulsingUnsupportedError", ] + +if _HAS_NATIVE_TENSOR_TRANSPORT: + __all__.extend(["TensorMessage", "tensor_transport_stats"]) diff --git a/python/pulsing/core/__init__.py b/python/pulsing/core/__init__.py index 38b45b6d2..205aa9426 100644 --- a/python/pulsing/core/__init__.py +++ b/python/pulsing/core/__init__.py @@ -19,6 +19,7 @@ def incr(self): self.value += 1; return self.value import asyncio import os +import sys from pulsing._async_bridge import ( clear_pulsing_loop, @@ -43,6 +44,15 @@ def incr(self): self.value += 1; return self.value StreamMessage, ) # internal: used by service.py / integrations +_native_core = sys.modules["pulsing._core"] +_HAS_NATIVE_TENSOR_TRANSPORT = all( + hasattr(_native_core, name) + for name in ("TensorMessage", "tensor_transport_stats") +) +if _HAS_NATIVE_TENSOR_TRANSPORT: + TensorMessage = _native_core.TensorMessage + tensor_transport_stats = _native_core.tensor_transport_stats + # ============================================================================= # Global system for simple API # ============================================================================= @@ -247,3 +257,6 @@ def _cli_attach_from_native(system: ActorSystem) -> None: "PulsingRuntimeError", "PulsingActorError", ] + +if _HAS_NATIVE_TENSOR_TRANSPORT: + __all__.extend(["TensorMessage", "tensor_transport_stats"]) diff --git a/python/pulsing/core/remote.py b/python/pulsing/core/remote.py index 31aed5829..b004fab7f 100644 --- a/python/pulsing/core/remote.py +++ b/python/pulsing/core/remote.py @@ -6,6 +6,7 @@ import inspect import logging import random +import sys import uuid from abc import ABC, abstractmethod from typing import Any, TypeVar @@ -22,10 +23,18 @@ ) from .proxy import ActorProxy, _DelayedCallProxy +_native_core = sys.modules["pulsing._core"] +_NATIVE_TENSOR_MESSAGE = getattr(_native_core, "TensorMessage", None) + logger = logging.getLogger(__name__) T = TypeVar("T") +# Native data-plane hooks are invoked only by their dedicated Rust/PyO3 +# message branch. Exposing the same name through generic dict RPC would let a +# caller pickle a message-like payload and bypass TensorMessage transport. +_RESERVED_WIRE_METHODS = frozenset({"receive_tensor"}) + # Trace context ContextVars — set by _WrappedActor.receive (from Rust handler # attributes), read by PyO3 ask()/tell() to propagate W3C traceparent across # actor-to-actor calls within the same asyncio Task. @@ -122,7 +131,7 @@ def _extract_methods(cls: type) -> tuple[list[str], set[str]]: methods = [] async_methods = set() for name, method in inspect.getmembers(cls, predicate=inspect.isfunction): - if name.startswith("_"): + if name.startswith("_") or name in _RESERVED_WIRE_METHODS: continue methods.append(name) if inspect.iscoroutinefunction(method) or inspect.isasyncgenfunction(method): @@ -203,10 +212,32 @@ async def receive(self, msg) -> Any: _current_traceparent.set(getattr(self, "__pulsing_tp__", None)) _current_tracestate.set(getattr(self, "__pulsing_ts__", None)) + if _NATIVE_TENSOR_MESSAGE is not None and isinstance( + msg, _NATIVE_TENSOR_MESSAGE + ): + receive_tensor = getattr(self._instance, "receive_tensor", None) + if not callable(receive_tensor): + return _wrap_response( + error=( + f"{type(self._instance).__name__} does not implement " + "receive_tensor(message)" + ) + ) + result = receive_tensor(msg) + if inspect.isawaitable(result): + result = await result + # TensorMessage must stay outside the normal method-response pickle + # envelope so Rust can retain or transport its buffers directly. + return result + if isinstance(msg, dict): method, args, kwargs, is_async_call = _unwrap_call(msg) - if not method or method.startswith("_"): + if ( + not method + or method.startswith("_") + or method in _RESERVED_WIRE_METHODS + ): return _wrap_response(error=f"Invalid method: {method}") _MISSING = object() diff --git a/tests/python/test_tensor_message.py b/tests/python/test_tensor_message.py new file mode 100644 index 000000000..7f7b71240 --- /dev/null +++ b/tests/python/test_tensor_message.py @@ -0,0 +1,264 @@ +import asyncio +import ctypes +import gc + +import pytest +import pulsing as pul +from pulsing.core import Actor, TensorMessage +from pulsing.core.protocol import _wrap_call +from pulsing.core.remote import _WrappedActor + + +def test_tensor_message_public_api_and_exports(): + first = bytearray(b"abc") + second = memoryview(bytearray(b"defg")) + message = TensorMessage(b"opaque-metadata", [first, second]) + + assert pul.TensorMessage is TensorMessage + assert message.metadata == b"opaque-metadata" + assert message.version == 1 + assert isinstance(message.buffers, list) + assert [bytes(buffer) for buffer in message.buffers] == [b"abc", b"defg"] + assert message.total_bytes == 7 + + +def test_tensor_message_allows_control_only_message(): + message = TensorMessage(b"control", [], version=3) + + assert message.metadata == b"control" + assert message.buffers == [] + assert message.version == 3 + assert message.total_bytes == 0 + + +def test_tensor_message_rejects_non_contiguous_buffer_without_copying(): + non_contiguous = memoryview(bytearray(b"abcdef"))[::2] + + with pytest.raises(ValueError, match="C-contiguous"): + TensorMessage(b"meta", [non_contiguous]) + + +class _TensorEcho(Actor): + async def receive(self, message): + assert isinstance(message, TensorMessage) + return message + + +@pytest.mark.asyncio +async def test_local_actor_tensor_roundtrip_shares_payload_and_keeps_owner_alive(): + system = await pul.actor_system() + try: + ref = await system.spawn(_TensorEcho(), name="tensor-local-echo") + source = bytearray(b"payload") + request = TensorMessage(b"schema", [source]) + + response = await ref.ask(request) + assert isinstance(response, TensorMessage) + assert response.metadata == b"schema" + assert bytes(response.buffers[0]) == b"payload" + + # The local mailbox path retains the original buffer export instead of + # materializing a payload copy. + source[0] = ord("P") + assert bytes(response.buffers[0]) == b"Payload" + + torch = pytest.importorskip("torch") + source_ptr = ctypes.addressof(ctypes.c_ubyte.from_buffer(source)) + tensor = torch.frombuffer(response.buffers[0], dtype=torch.uint8) + assert tensor.data_ptr() == source_ptr + assert bytes(tensor.tolist()) == b"Payload" + + del request + del source + del response + gc.collect() + # torch.frombuffer retains the read-only exporter, which retains Rust + # Bytes and ultimately the original Python buffer lease. + assert bytes(tensor.tolist()) == b"Payload" + finally: + await system.shutdown() + + +@pytest.mark.parametrize("transport", ["raw", "http2"]) +@pytest.mark.asyncio +async def test_remote_actor_tensor_roundtrip_uses_selected_transport( + transport, monkeypatch +): + monkeypatch.setenv("PULSING_TENSOR_TRANSPORT", transport) + before = pul.tensor_transport_stats() + server = await pul.actor_system(addr="127.0.0.1:0") + client = None + try: + actor_name = f"tensor-remote-echo-{transport}" + await server.spawn(_TensorEcho(), name=actor_name, public=True) + client = await pul.actor_system(addr="127.0.0.1:0", seeds=[server.addr]) + remote_ref = await client.resolve_named( + actor_name, node_id=server.node_id.id, timeout=5.0 + ) + assert not remote_ref.is_local() + + source = bytearray(b"remote-payload") + response = await remote_ref.ask(TensorMessage(b"remote-schema", [source])) + + assert isinstance(response, TensorMessage) + assert response.metadata == b"remote-schema" + assert bytes(response.buffers[0]) == b"remote-payload" + assert not response.buffers[0].readonly + + # Crossing TCP materializes a receive allocation, so it must no longer + # alias the sender's Python bytearray. + source[0] = ord("R") + assert bytes(response.buffers[0]) == b"remote-payload" + + torch = pytest.importorskip("torch") + tensor = torch.frombuffer(response.buffers[0], dtype=torch.uint8) + tensor[0] = ord("R") + del response + gc.collect() + assert bytes(tensor.tolist()) == b"Remote-payload" + + after = pul.tensor_transport_stats() + if transport == "raw": + assert after["raw_frames_sent"] >= before["raw_frames_sent"] + 2 + assert after["active_copy_model"] == "direct_tcp" + else: + assert ( + after["http2_fallback_frames"] + >= before["http2_fallback_frames"] + 1 + ) + assert after["active_copy_model"] == "packed_http2_compatibility" + finally: + if client is not None: + await client.shutdown() + await server.shutdown() + + +class _TensorPutAck(Actor): + async def receive(self, message): + assert isinstance(message, TensorMessage) + return {"ok": True, "received_bytes": message.total_bytes} + + +@pytest.mark.parametrize("transport", ["raw", "http2"]) +@pytest.mark.asyncio +async def test_tensor_request_supports_pickled_single_ack(transport, monkeypatch): + monkeypatch.setenv("PULSING_TENSOR_TRANSPORT", transport) + server = await pul.actor_system(addr="127.0.0.1:0") + client = None + try: + actor_name = f"tensor-put-ack-{transport}" + await server.spawn(_TensorPutAck(), name=actor_name, public=True) + client = await pul.actor_system(addr="127.0.0.1:0", seeds=[server.addr]) + remote_ref = await client.resolve_named( + actor_name, node_id=server.node_id.id, timeout=5.0 + ) + + response = await remote_ref.ask(TensorMessage(b"put", [b"123", b"4567"])) + + assert response == {"ok": True, "received_bytes": 7} + expected_model = ( + "direct_tcp" if transport == "raw" else "packed_http2_compatibility" + ) + assert pul.tensor_transport_stats()["active_copy_model"] == expected_model + finally: + if client is not None: + await client.shutdown() + await server.shutdown() + + +@pytest.mark.asyncio +async def test_raw_tensor_tell_receives_protocol_ack(monkeypatch): + monkeypatch.setenv("PULSING_TENSOR_TRANSPORT", "raw") + before = pul.tensor_transport_stats() + server = await pul.actor_system(addr="127.0.0.1:0") + client = None + try: + await server.spawn(_TensorEcho(), name="tensor-tell-ack", public=True) + client = await pul.actor_system(addr="127.0.0.1:0", seeds=[server.addr]) + remote_ref = await client.resolve_named( + "tensor-tell-ack", node_id=server.node_id.id, timeout=5.0 + ) + + # tell() completes only after the raw peer returns its protocol ACK. + await remote_ref.tell(TensorMessage(b"tell", [b"payload"])) + + after = pul.tensor_transport_stats() + assert after["raw_frames_sent"] >= before["raw_frames_sent"] + 2 + assert after["active_copy_model"] == "direct_tcp" + finally: + if client is not None: + await client.shutdown() + await server.shutdown() + + +@pytest.mark.asyncio +async def test_raw_tensor_pool_reconnects_only_before_next_request(monkeypatch): + monkeypatch.setenv("PULSING_TENSOR_TRANSPORT", "raw") + monkeypatch.setenv("PULSING_TENSOR_MAX_REQUESTS_PER_CONNECTION", "1") + before = pul.tensor_transport_stats() + server = await pul.actor_system(addr="127.0.0.1:0") + client = None + try: + await server.spawn(_TensorEcho(), name="tensor-reconnect", public=True) + client = await pul.actor_system(addr="127.0.0.1:0", seeds=[server.addr]) + remote_ref = await client.resolve_named( + "tensor-reconnect", node_id=server.node_id.id, timeout=5.0 + ) + + first = await remote_ref.ask(TensorMessage(b"one", [b"first"])) + assert bytes(first.buffers[0]) == b"first" + # Let the peer FIN become visible before checking out the pooled socket. + await asyncio.sleep(0.05) + second = await remote_ref.ask(TensorMessage(b"two", [b"second"])) + assert bytes(second.buffers[0]) == b"second" + + after = pul.tensor_transport_stats() + assert ( + after["raw_connections_accepted"] + >= before["raw_connections_accepted"] + 2 + ) + finally: + if client is not None: + await client.shutdown() + await server.shutdown() + + +class _TensorService: + def __init__(self): + self.calls = 0 + + async def receive_tensor(self, message): + self.calls += 1 + return TensorMessage(message.metadata + b"-reply", message.buffers, message.version) + + +@pytest.mark.asyncio +async def test_wrapped_actor_dispatches_receive_tensor_and_returns_direct_message(): + service = _TensorService() + wrapped = _WrappedActor(service) + request = TensorMessage(b"request", [b"data"]) + + response = await wrapped.receive(request) + + assert isinstance(response, TensorMessage) + assert response.metadata == b"request-reply" + assert bytes(response.buffers[0]) == b"data" + assert service.calls == 1 + + +@pytest.mark.asyncio +async def test_receive_tensor_is_reserved_from_generic_rpc_dispatch(): + service = _TensorService() + wrapped = _WrappedActor(service) + + response = await wrapped.receive( + _wrap_call( + "receive_tensor", + (TensorMessage(b"pickled", [b"payload"]),), + {}, + False, + ) + ) + + assert response == {"__error__": "Invalid method: receive_tensor"} + assert service.calls == 0