diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8c468eb2..4fc24272 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -142,6 +142,11 @@ jobs: persist-credentials: false - name: Run ruff check uses: astral-sh/ruff-action@v3 + with: + # Unpinned, this floats to the newest ruff and every new lint rule + # breaks CI on untouched files (0.16.0 did exactly that). Bump + # deliberately, together with the fixes it demands. + version: "0.15.4" python-unit-tests: name: python unit tests diff --git a/pegaflow-core/src/backing/transfer_lock_guard.rs b/pegaflow-core/src/backing/transfer_lock_guard.rs index e75f8b75..b7734ff7 100644 --- a/pegaflow-core/src/backing/transfer_lock_guard.rs +++ b/pegaflow-core/src/backing/transfer_lock_guard.rs @@ -81,11 +81,12 @@ mod tests { use pegaflow_proto::proto::engine::engine_server::{Engine, EngineServer}; use pegaflow_proto::proto::engine::{ - HealthRequest, HealthResponse, LoadRequest, LoadResponse, QueryBlocksForTransferRequest, - QueryBlocksForTransferResponse, QueryRequest, QueryResponse, RdmaHandshakeRequest, - RdmaHandshakeResponse, RegisterContextRequest, RegisterContextResponse, ReleaseRequest, - ReleaseResponse, ReleaseTransferLockResponse, SaveRequest, SaveResponse, SessionEvent, - SessionRequest, ShutdownRequest, ShutdownResponse, UnregisterRequest, UnregisterResponse, + FlushRequest, FlushResponse, HealthRequest, HealthResponse, LoadRequest, LoadResponse, + QueryBlocksForTransferRequest, QueryBlocksForTransferResponse, QueryRequest, QueryResponse, + RdmaHandshakeRequest, RdmaHandshakeResponse, RegisterContextRequest, + RegisterContextResponse, ReleaseRequest, ReleaseResponse, ReleaseTransferLockResponse, + SaveRequest, SaveResponse, SessionEvent, SessionRequest, ShutdownRequest, ShutdownResponse, + UnregisterRequest, UnregisterResponse, }; use tokio_stream::wrappers::TcpListenerStream; use tonic::transport::Endpoint; @@ -125,6 +126,12 @@ mod tests { ) -> Result, Status> { Err(Status::unimplemented("stub")) } + async fn flush( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("stub")) + } async fn register_context_batch( &self, _request: Request, diff --git a/pegaflow-core/src/internode/p2p_service.rs b/pegaflow-core/src/internode/p2p_service.rs index cb9dd294..923cd78e 100644 --- a/pegaflow-core/src/internode/p2p_service.rs +++ b/pegaflow-core/src/internode/p2p_service.rs @@ -20,12 +20,12 @@ use tonic::{Request, Response, Status, async_trait}; use pegaflow_proto::proto::engine::engine_server::{Engine, EngineServer}; use pegaflow_proto::proto::engine::{ - HealthRequest, HealthResponse, LoadRequest, LoadResponse, QueryBlocksForTransferRequest, - QueryBlocksForTransferResponse, QueryRequest, QueryResponse, RdmaHandshakeRequest, - RdmaHandshakeResponse, RegisterContextRequest, RegisterContextResponse, ReleaseRequest, - ReleaseResponse, ReleaseTransferLockRequest, ReleaseTransferLockResponse, ResponseStatus, - SaveRequest, SaveResponse, SessionEvent, SessionRequest, ShutdownRequest, ShutdownResponse, - TransferBlockInfo, TransferSlotInfo, UnregisterRequest, UnregisterResponse, + FlushRequest, FlushResponse, HealthRequest, HealthResponse, LoadRequest, LoadResponse, + QueryBlocksForTransferRequest, QueryBlocksForTransferResponse, QueryRequest, QueryResponse, + RdmaHandshakeRequest, RdmaHandshakeResponse, RegisterContextRequest, RegisterContextResponse, + ReleaseRequest, ReleaseResponse, ReleaseTransferLockRequest, ReleaseTransferLockResponse, + ResponseStatus, SaveRequest, SaveResponse, SessionEvent, SessionRequest, ShutdownRequest, + ShutdownResponse, TransferBlockInfo, TransferSlotInfo, UnregisterRequest, UnregisterResponse, }; use crate::{LayerBlock, PegaEngine}; @@ -246,6 +246,13 @@ impl Engine for P2pTransferService { Self::not_served("load") } + async fn flush( + &self, + _request: Request, + ) -> Result, Status> { + Self::not_served("flush") + } + async fn query_prefetch( &self, _request: Request, diff --git a/pegaflow-core/src/storage/mod.rs b/pegaflow-core/src/storage/mod.rs index 857ee8f2..5ac2a7d1 100644 --- a/pegaflow-core/src/storage/mod.rs +++ b/pegaflow-core/src/storage/mod.rs @@ -5,8 +5,6 @@ pub(crate) mod transfer_lock; mod write_path; use bytesize::ByteSize; -#[cfg(not(feature = "rdma"))] -use log::warn; use log::{debug, info, warn}; use std::collections::HashSet; use std::num::NonZeroU64; diff --git a/pegaflow-proto/proto/engine.proto b/pegaflow-proto/proto/engine.proto index 41b7818b..dbd58f90 100644 --- a/pegaflow-proto/proto/engine.proto +++ b/pegaflow-proto/proto/engine.proto @@ -46,10 +46,32 @@ message RegisterContextRequest { // Set by the connector for MLA; the engine derives per-layer page offsets // from the registered layouts. See spec.md. bool page_first = 18; + // Native (non-Python) registration: per-layer views into one KV arena the + // server allocates on `device_id` and shares back as a CUDA IPC handle. + // Empty for Python registrations. + repeated NativeKvTensor native_kv_tensors = 19; + // Total arena size the server allocates for a native registration; every + // layer view must fit inside it. Zero for Python registrations. + uint64 native_alloc_size = 20; +} + +// One layer's view into the server-allocated KV arena. +message NativeKvTensor { + uint64 offset_bytes = 1; + uint64 size_bytes = 2; + // Distance between consecutive blocks of this layer inside the arena + // (fused allocations interleave layers, so blocks are not contiguous). + uint64 block_stride_bytes = 3; } message RegisterContextResponse { ResponseStatus status = 1; + // CUDA IPC handle (CUipcMemHandle, 64 bytes) of the server-allocated arena. + // The client imports it with cuIpcOpenMemHandle and uses the mapping as its + // KV cache. Present only for native registrations; the arena is freed when + // the instance unregisters, so the client must not touch the mapping after + // unregister_context returns. + bytes arena_ipc_handle = 2; } message SaveLayer { @@ -77,6 +99,9 @@ message LoadRequest { string load_state_shm = 4; repeated string layer_names = 5; repeated LeaseLoad loads = 6; + // Return only after the H2D copies complete. Native clients must set this: + // they have no load_state_shm completion flag to poll. + bool wait_for_completion = 7; } message LeaseLoad { @@ -88,6 +113,12 @@ message LoadResponse { ResponseStatus status = 1; } +message FlushRequest {} + +message FlushResponse { + ResponseStatus status = 1; +} + message QueryRequest { string instance_id = 1; repeated bytes block_hashes = 2; @@ -207,6 +238,9 @@ service Engine { rpc RegisterContextBatch(RegisterContextRequest) returns (RegisterContextResponse); rpc Save(SaveRequest) returns (SaveResponse); rpc Load(LoadRequest) returns (LoadResponse); + // Durability barrier: returns once every accepted save is cache-visible and + // its MetaServer registration has been attempted. + rpc Flush(FlushRequest) returns (FlushResponse); // Prefetch query: check memory + trigger SSD prefetch or cross-node RDMA fetch for missing blocks rpc QueryPrefetch(QueryRequest) returns (QueryResponse); rpc Release(ReleaseRequest) returns (ReleaseResponse); diff --git a/pegaflow-proto/src/lib.rs b/pegaflow-proto/src/lib.rs index a9f53d2a..b575ece9 100644 --- a/pegaflow-proto/src/lib.rs +++ b/pegaflow-proto/src/lib.rs @@ -1,3 +1,8 @@ +/// Exact wire capability version. Registration is fenced on this string, so a +/// client built before the native-arena contract fails before any memory +/// changes hands. +pub const VERSION: &str = concat!(env!("CARGO_PKG_VERSION"), "+native-arena-v1"); + pub mod proto { #[allow( clippy::allow_attributes_without_reason, diff --git a/pegaflow-server/src/http_server.rs b/pegaflow-server/src/http_server.rs index dd0cba8e..e7de85d2 100644 --- a/pegaflow-server/src/http_server.rs +++ b/pegaflow-server/src/http_server.rs @@ -82,8 +82,13 @@ async fn cleanup_handler( ) -> impl IntoResponse { match query.id { None => { - let removed_tensors = state.registry.clear().await; + // Engine first, then a save-pipeline barrier, then the registry + // drop: the registry drop frees native arenas (cuMemFree), so the + // engine must forget its raw pointers and drain accepted saves + // before the memory behind them goes away. let removed_instances = state.engine.unregister_all_instances(); + state.engine.flush_saves().await; + let removed_tensors = state.registry.clear().await; if !removed_instances.is_empty() || removed_tensors > 0 { warn!( @@ -104,8 +109,11 @@ async fn cleanup_handler( ) } Some(instance_id) => { + // Same ordering as the clear-all branch above. + let engine_result = state.engine.unregister_instance(&instance_id); + state.engine.flush_saves().await; let removed_tensors = state.registry.drop_instance(instance_id.clone()).await; - match state.engine.unregister_instance(&instance_id) { + match engine_result { Ok(()) => { warn!( "Cleanup instance {}: {} CUDA tensor(s) released", diff --git a/pegaflow-server/src/lib.rs b/pegaflow-server/src/lib.rs index e5117d20..2c81ba2d 100644 --- a/pegaflow-server/src/lib.rs +++ b/pegaflow-server/src/lib.rs @@ -1,6 +1,7 @@ mod check_cuda_version; pub mod http_server; pub mod metric; +mod native_arena; pub mod proto; pub mod registry; pub mod service; @@ -61,6 +62,11 @@ pub struct Cli { #[arg(long, value_delimiter = ',')] pub devices: Vec, + /// Initialize the torch CUDA registry for Python (vLLM connector) clients. + /// `false` runs torch-free and serves native clients only. + #[arg(long, default_value_t = true, action = clap::ArgAction::Set)] + pub python_registry: bool, + /// Pinned memory pool size (supports units: kb, mb, gb, tb) /// Examples: "10gb", "500mb", "1tb" #[arg(long, default_value = "30gb", value_parser = parse_memory_size)] @@ -322,6 +328,30 @@ fn detect_cuda_devices() -> Result, std::io::Error> { }) } +/// Torch-free device enumeration for native-only deployments. +fn detect_native_cuda_devices() -> Result, std::io::Error> { + let count = cudarc::driver::CudaContext::device_count() + .map_err(|e| std::io::Error::other(format!("cudarc device_count: {e}")))?; + Ok((0..count).collect()) +} + +/// Torch-free CUDA context initialization for native-only deployments. +fn init_cudarc_cuda(device_ids: &[i32]) -> Result<(), std::io::Error> { + if device_ids.is_empty() { + return Err(std::io::Error::other("no CUDA devices to initialize")); + } + for &device_id in device_ids { + let ordinal = usize::try_from(device_id) + .map_err(|_| std::io::Error::other(format!("device_id {device_id} must be >= 0")))?; + let ctx = cudarc::driver::CudaContext::new(ordinal) + .map_err(|e| std::io::Error::other(format!("cudarc init device {device_id}: {e}")))?; + ctx.bind_to_thread() + .map_err(|e| std::io::Error::other(format!("cudarc bind device {device_id}: {e}")))?; + info!("Initialized CUDA context for device {device_id} (cudarc, torch-free)"); + } + Ok(()) +} + fn init_python_cuda(device_ids: &[i32]) -> Result<(), std::io::Error> { if device_ids.is_empty() { return Err(std::io::Error::other("no CUDA devices to initialize")); @@ -447,7 +477,11 @@ pub fn run() -> Result<(), Box> { // Determine which devices to initialize let devices = if cli.devices.is_empty() { // Auto-detect all available devices - let detected = detect_cuda_devices()?; + let detected = if cli.python_registry { + detect_cuda_devices()? + } else { + detect_native_cuda_devices()? + }; info!( "Auto-detected {} CUDA device(s): {:?}", detected.len(), @@ -463,17 +497,26 @@ pub fn run() -> Result<(), Box> { return Err("No CUDA devices available".into()); } - init_python_cuda(&devices)?; + if cli.python_registry { + init_python_cuda(&devices)?; + } else { + init_cudarc_cuda(&devices)?; + } info!( "CUDA runtime initialized for {} device(s): {:?}", devices.len(), devices ); - let registry = CudaTensorRegistry::new().map_err(|err| { - let msg = format_py_err(err); - std::io::Error::other(format!("failed to initialize torch CUDA context: {msg}")) - })?; + let registry = if cli.python_registry { + CudaTensorRegistry::new().map_err(|err| { + let msg = format_py_err(err); + std::io::Error::other(format!("failed to initialize torch CUDA context: {msg}")) + })? + } else { + info!("Python/torch registry disabled; serving native clients only"); + CudaTensorRegistry::empty() + }; // Confine the registry to its own thread: GIL + CUDA work now happens off // the async runtime, so a wedged `empty_cache` can't starve tokio workers. let registry = RegistryHandle::spawn(registry); diff --git a/pegaflow-server/src/native_arena.rs b/pegaflow-server/src/native_arena.rs new file mode 100644 index 00000000..e39b4faa --- /dev/null +++ b/pegaflow-server/src/native_arena.rs @@ -0,0 +1,218 @@ +//! Server-owned KV arenas for native (non-Python) clients. +//! +//! A native client does not export memory to us — we allocate the arena in +//! this process with `cuMemAlloc` and hand back a CUDA IPC handle in the +//! registration response. Owning the allocation is what makes the follow-up +//! RDMA work possible: `ibv_reg_mr`/dma-buf registration works on memory this +//! process allocated, while an IPC-*imported* pointer can never be registered +//! into a NIC. The client side only runs compute kernels on its imported +//! mapping, which CUDA IPC fully supports. +//! +//! One arena per context; layer views are plain `base + offset` arithmetic. +//! The arena is freed when its context leaves the registry, strictly after the +//! engine forgot the raw pointers derived from it. + +use cudarc::driver::{CudaContext, result::DriverError, sys}; +use std::collections::HashMap; +use std::sync::Arc; + +use crate::registry::TensorMetadata; + +/// One layer's requested view into the arena, straight from the wire. +pub(crate) struct NativeLayerView { + pub layer_name: String, + pub offset_bytes: u64, + pub size_bytes: u64, +} + +/// Result of a native registration: raw pointers for the engine plus the IPC +/// handle the client imports. +pub(crate) struct NativeRegistration { + pub metadatas: Vec, + pub arena_ipc_handle: Vec, +} + +/// A device allocation owned by this process, shared to one client via CUDA +/// IPC. Freed on drop, on the registry actor thread. +struct NativeArena { + context: Arc, + base_ptr: sys::CUdeviceptr, + size_bytes: usize, + ipc_handle: sys::CUipcMemHandle, +} + +// SAFETY: the raw device pointer and IPC handle are plain values; all CUDA +// calls bind the context to the calling thread first. +unsafe impl Send for NativeArena {} + +impl NativeArena { + fn allocate(device_id: i32, size_bytes: usize) -> Result { + let device = usize::try_from(device_id) + .map_err(|_| format!("device_id {device_id} must be >= 0"))?; + let context = CudaContext::new(device).map_err(|e| cuda_error("retain CUDA context", e))?; + context + .bind_to_thread() + .map_err(|e| cuda_error("bind CUDA context", e))?; + + let mut base_ptr: sys::CUdeviceptr = 0; + // SAFETY: base_ptr receives the allocation; size is non-zero + // (validated by the caller). + unsafe { sys::cuMemAlloc_v2(&mut base_ptr, size_bytes).result() } + .map_err(|e| cuda_error("allocate KV arena", e))?; + + // Zero the arena so a KV hit never reads stale device memory. The + // memset is async with respect to the host and CUDA IPC gives the + // importing process no cross-process stream ordering, so synchronize + // before the handle leaves this function. + // SAFETY: base_ptr..base_ptr+size_bytes was just allocated. + let zeroed = unsafe { sys::cuMemsetD8_v2(base_ptr, 0, size_bytes).result() } + .and_then(|_| unsafe { sys::cuCtxSynchronize().result() }); + if let Err(e) = zeroed { + unsafe { sys::cuMemFree_v2(base_ptr).result().ok() }; + return Err(cuda_error("zero KV arena", e)); + } + + let mut ipc_handle = sys::CUipcMemHandle { reserved: [0; 64] }; + // SAFETY: base_ptr is a live cuMemAlloc allocation, the only kind + // cuIpcGetMemHandle accepts. + if let Err(e) = unsafe { sys::cuIpcGetMemHandle(&mut ipc_handle, base_ptr).result() } { + unsafe { sys::cuMemFree_v2(base_ptr).result().ok() }; + return Err(cuda_error("export KV arena IPC handle", e)); + } + + Ok(Self { + context, + base_ptr, + size_bytes, + ipc_handle, + }) + } +} + +impl Drop for NativeArena { + fn drop(&mut self) { + self.context + .bind_to_thread() + .expect("bind CUDA context before freeing native KV arena"); + // SAFETY: base_ptr is the live allocation from `allocate`; the engine + // dropped its raw pointers before the registry released this arena. + unsafe { sys::cuMemFree_v2(self.base_ptr).result() }.expect("free native KV arena"); + } +} + +fn cuda_error(operation: &str, error: DriverError) -> String { + format!("{operation}: {error}") +} + +struct NativeContext { + #[allow( + dead_code, + reason = "owning the arena keeps the registered addresses alive; freed on drop" + )] + arena: NativeArena, + layer_count: usize, +} + +/// Native contexts keyed by the same `instance:tp:pp:dev` context key as the +/// Python registry, kept in a separate map so the Python bookkeeping stays +/// untouched. +#[derive(Default)] +pub(crate) struct NativeArenaMap { + contexts: HashMap, +} + +impl NativeArenaMap { + pub(crate) fn contains(&self, context_key: &str) -> bool { + self.contexts.contains_key(context_key) + } + + /// Allocate an arena on `device_id`, carve the requested layer views out + /// of it, and record the context. Runs on the registry actor thread. + pub(crate) fn register( + &mut self, + context_key: &str, + device_id: i32, + layers: &[NativeLayerView], + alloc_size: usize, + ) -> Result { + if self.contexts.contains_key(context_key) { + return Err(format!("context {context_key} is already registered")); + } + let mut seen = std::collections::HashSet::with_capacity(layers.len()); + for layer in layers { + if !seen.insert(layer.layer_name.as_str()) { + return Err(format!( + "layer {} appears more than once in context {context_key}", + layer.layer_name + )); + } + } + + let arena = NativeArena::allocate(device_id, alloc_size)?; + let mut metadatas = Vec::with_capacity(layers.len()); + for layer in layers { + let size = usize::try_from(layer.size_bytes) + .map_err(|_| format!("layer {} size does not fit usize", layer.layer_name))?; + layer + .offset_bytes + .checked_add(layer.size_bytes) + .filter(|end| *end <= arena.size_bytes as u64) + .ok_or_else(|| { + format!( + "layer {} view [{} +{}] is outside its {}-byte arena", + layer.layer_name, layer.offset_bytes, layer.size_bytes, arena.size_bytes + ) + })?; + metadatas.push(TensorMetadata { + data_ptr: arena.base_ptr + layer.offset_bytes, + size_bytes: size, + device_id, + }); + } + + let arena_ipc_handle: Vec = arena + .ipc_handle + .reserved + .iter() + .map(|&byte| byte as u8) + .collect(); + self.contexts.insert( + context_key.to_string(), + NativeContext { + arena, + layer_count: layers.len(), + }, + ); + Ok(NativeRegistration { + metadatas, + arena_ipc_handle, + }) + } + + /// Free the arena of one context; returns the number of layers dropped. + pub(crate) fn drop_context(&mut self, context_key: &str) -> usize { + self.contexts + .remove(context_key) + .map(|ctx| ctx.layer_count) + .unwrap_or(0) + } + + /// Free every arena belonging to `instance_id`; returns layers dropped. + pub(crate) fn drop_instance(&mut self, instance_id: &str) -> usize { + let prefix = format!("{instance_id}:"); + let keys: Vec = self + .contexts + .keys() + .filter(|key| key.starts_with(&prefix)) + .cloned() + .collect(); + keys.iter().map(|key| self.drop_context(key)).sum() + } + + /// Free every arena; returns layers dropped. + pub(crate) fn clear(&mut self) -> usize { + let count = self.contexts.values().map(|ctx| ctx.layer_count).sum(); + self.contexts.clear(); + count + } +} diff --git a/pegaflow-server/src/registry.rs b/pegaflow-server/src/registry.rs index 6dbd33b6..e89ac41f 100644 --- a/pegaflow-server/src/registry.rs +++ b/pegaflow-server/src/registry.rs @@ -1,3 +1,4 @@ +use crate::native_arena::{NativeArenaMap, NativeLayerView, NativeRegistration}; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::types::PyBytes; @@ -44,6 +45,9 @@ impl ContextState { pub struct CudaTensorRegistry { contexts: HashMap, + /// Server-allocated arenas for native clients; disjoint keyspace from + /// `contexts` (enforced at registration). + native: NativeArenaMap, } impl CudaTensorRegistry { @@ -54,6 +58,7 @@ impl CudaTensorRegistry { cuda.call_method0("init")?; Ok(Self { contexts: HashMap::new(), + native: NativeArenaMap::default(), }) }) } @@ -61,6 +66,7 @@ impl CudaTensorRegistry { pub fn empty() -> Self { Self { contexts: HashMap::new(), + native: NativeArenaMap::default(), } } @@ -70,7 +76,7 @@ impl CudaTensorRegistry { device_id: i32, layers: Vec<(String, Vec)>, ) -> PyResult> { - if self.contexts.contains_key(context_key) { + if self.contexts.contains_key(context_key) || self.native.contains(context_key) { return Err(PyValueError::new_err(format!( "context {context_key} is already registered" ))); @@ -106,8 +112,25 @@ impl CudaTensorRegistry { Ok(metadatas) } + /// Allocate a server-owned arena and register its layer views. Native + /// contexts share the registry keyspace with Python contexts but none of + /// the Python (GIL/torch) machinery. + fn register_native( + &mut self, + context_key: &str, + device_id: i32, + layers: &[NativeLayerView], + alloc_size: usize, + ) -> Result { + if self.contexts.contains_key(context_key) { + return Err(format!("context {context_key} is already registered")); + } + self.native + .register(context_key, device_id, layers, alloc_size) + } + fn drop_context(&mut self, context_key: &str) -> usize { - self.release_contexts(vec![context_key.to_string()]) + self.native.drop_context(context_key) + self.release_contexts(vec![context_key.to_string()]) } fn drop_instance(&mut self, instance_id: &str) -> usize { @@ -118,13 +141,13 @@ impl CudaTensorRegistry { .filter(|key| key.starts_with(&prefix)) .cloned() .collect(); - self.release_contexts(keys) + self.native.drop_instance(instance_id) + self.release_contexts(keys) } /// Clear all contexts and return the total number of tensors removed. fn clear_and_count(&mut self) -> usize { let keys: Vec = self.contexts.keys().cloned().collect(); - self.release_contexts(keys) + self.native.clear() + self.release_contexts(keys) } /// Remove `keys` from the registry, returning the number of CUDA IPC tensors @@ -215,6 +238,17 @@ enum RegistryCommand { // so callers never need to touch the GIL to read an error message. reply: oneshot::Sender, String>>, }, + RegisterNative { + context_key: String, + device_id: i32, + layers: Vec, + alloc_size: usize, + reply: oneshot::Sender>, + }, + ContainsContext { + context_key: String, + reply: oneshot::Sender, + }, DropInstance { instance_id: String, reply: oneshot::Sender, @@ -280,6 +314,36 @@ impl RegistryHandle { rx.await.expect("cuda-registry thread dropped reply") } + /// Allocate a server-owned arena on the registry thread and register its + /// layer views under `context_key`. Returns per-layer metadata plus the + /// CUDA IPC handle the client imports. + pub(crate) async fn register_native( + &self, + context_key: String, + device_id: i32, + layers: Vec, + alloc_size: usize, + ) -> Result { + let (reply, rx) = oneshot::channel(); + self.dispatch(RegistryCommand::RegisterNative { + context_key, + device_id, + layers, + alloc_size, + reply, + }) + .await; + rx.await.expect("cuda-registry thread dropped reply") + } + + /// Whether `context_key` is currently registered (Python or native). + pub(crate) async fn contains_context(&self, context_key: String) -> bool { + let (reply, rx) = oneshot::channel(); + self.dispatch(RegistryCommand::ContainsContext { context_key, reply }) + .await; + rx.await.expect("cuda-registry thread dropped reply") + } + /// Drop all CUDA tensors belonging to `instance_id`; returns the count /// released. pub async fn drop_instance(&self, instance_id: String) -> usize { @@ -330,6 +394,21 @@ fn registry_actor(mut registry: CudaTensorRegistry, mut rx: mpsc::Receiver { + let result = registry.register_native(&context_key, device_id, &layers, alloc_size); + let _ = reply.send(result); + } + RegistryCommand::ContainsContext { context_key, reply } => { + let present = registry.contexts.contains_key(&context_key) + || registry.native.contains(&context_key); + let _ = reply.send(present); + } RegistryCommand::DropInstance { instance_id, reply } => { let _ = reply.send(registry.drop_instance(&instance_id)); } diff --git a/pegaflow-server/src/service.rs b/pegaflow-server/src/service.rs index 832b9d7c..9044fa35 100644 --- a/pegaflow-server/src/service.rs +++ b/pegaflow-server/src/service.rs @@ -1,15 +1,17 @@ use pegaflow_core::{trace_in_span, trace_root}; use crate::metric::record_rpc_result; +use crate::native_arena::NativeLayerView; use crate::proto::engine::engine_server::Engine; use crate::proto::engine::{ - HealthRequest, HealthResponse, LoadRequest, LoadResponse, QueryBlocksForTransferRequest, - QueryBlocksForTransferResponse, QueryLoading, QueryReady, QueryRequest, QueryResponse, - RdmaHandshakeRequest, RdmaHandshakeResponse, RegisterContextRequest, RegisterContextResponse, - ReleaseRequest, ReleaseResponse, ReleaseTransferLockRequest, ReleaseTransferLockResponse, - ResponseStatus, SaveRequest, SaveResponse, SessionEvent, SessionRequest, ShutdownRequest, - ShutdownResponse, TransferBlockInfo, TransferMode as ProtoTransferMode, TransferSlotInfo, - UnregisterRequest, UnregisterResponse, query_response, + FlushRequest, FlushResponse, HealthRequest, HealthResponse, LoadRequest, LoadResponse, + QueryBlocksForTransferRequest, QueryBlocksForTransferResponse, QueryLoading, QueryReady, + QueryRequest, QueryResponse, RdmaHandshakeRequest, RdmaHandshakeResponse, + RegisterContextRequest, RegisterContextResponse, ReleaseRequest, ReleaseResponse, + ReleaseTransferLockRequest, ReleaseTransferLockResponse, ResponseStatus, SaveRequest, + SaveResponse, SessionEvent, SessionRequest, ShutdownRequest, ShutdownResponse, + TransferBlockInfo, TransferMode as ProtoTransferMode, TransferSlotInfo, UnregisterRequest, + UnregisterResponse, query_response, }; use crate::registry::RegistryHandle; use crate::session::SessionRegistry; @@ -54,13 +56,12 @@ impl GrpcEngineService { instance_id: &str, reason: &'static str, ) { - let removed = registry.drop_instance(instance_id.to_string()).await; - if removed > 0 { - info!( - "Session cleanup ({}): dropped {} CUDA tensors for instance {}", - reason, removed, instance_id - ); - } + // Engine first: it must forget raw device pointers before the registry + // drop frees the memory behind them (server-owned native arenas are + // cuMemFree'd right there; Python IPC imports are unmapped by GC). The + // flush barrier then drains saves already handed to the write pipeline. + // Saves still queued inside a GPU worker are NOT fenced — closing that + // residual window needs the engine-level drain-before-unregister work. if let Err(err) = engine.unregister_instance(instance_id) { // `InstanceMissing` is normal if the instance was never registered // (vllm died before any register_context_batch). Log at debug. @@ -69,6 +70,14 @@ impl GrpcEngineService { reason, instance_id, err ); } + engine.flush_saves().await; + let removed = registry.drop_instance(instance_id.to_string()).await; + if removed > 0 { + info!( + "Session cleanup ({}): dropped {} CUDA tensors for instance {}", + reason, removed, instance_id + ); + } } fn context_key(instance_id: &str, tp_rank: u32, pp_rank: u32, device_id: i32) -> String { @@ -116,8 +125,29 @@ impl GrpcEngineService { Ok(()) } + /// Per-layer arena views plus the engine's block-stride list, from the + /// wire representation. Lengths are already validated. + fn native_layer_views( + req: &RegisterContextRequest, + ) -> Result<(Vec, Vec), Status> { + let mut layers = Vec::with_capacity(req.native_kv_tensors.len()); + let mut block_stride_bytes = Vec::with_capacity(req.native_kv_tensors.len()); + for (layer_name, tensor) in req.layer_names.iter().zip(&req.native_kv_tensors) { + block_stride_bytes.push(Self::usize_from_u64( + tensor.block_stride_bytes, + "block_stride_bytes", + )?); + layers.push(NativeLayerView { + layer_name: layer_name.clone(), + offset_bytes: tensor.offset_bytes, + size_bytes: tensor.size_bytes, + }); + } + Ok((layers, block_stride_bytes)) + } + fn validate_register_context_request(req: &RegisterContextRequest) -> Result<(), Status> { - let server_version = env!("CARGO_PKG_VERSION"); + let server_version = pegaflow_proto::VERSION; if req.client_version != server_version { return Err(Status::failed_precondition(format!( "PegaFlow version mismatch: client={} server={server_version}", @@ -141,6 +171,60 @@ impl GrpcEngineService { req.tp_rank, req.tp_size ))); } + let batch_len = req.layer_names.len(); + if req.native_kv_tensors.is_empty() { + if req.native_alloc_size != 0 { + return Err(Status::invalid_argument( + "native_alloc_size requires native_kv_tensors", + )); + } + if req.wrapper_bytes.len() != batch_len { + return Err(Status::invalid_argument(format!( + "wrapper_bytes length {} does not match layer_names {batch_len}", + req.wrapper_bytes.len() + ))); + } + } else { + if !req.wrapper_bytes.is_empty() { + return Err(Status::invalid_argument( + "a registration is either native or Python, not both", + )); + } + if req.native_kv_tensors.len() != batch_len { + return Err(Status::invalid_argument(format!( + "native_kv_tensors length {} does not match layer_names {batch_len}", + req.native_kv_tensors.len() + ))); + } + if req.native_alloc_size == 0 { + return Err(Status::invalid_argument( + "native registration requires a non-zero native_alloc_size", + )); + } + // Native v1 scope: one process, one GPU, one arena. + if req.tp_size != 1 || req.world_size != 1 { + return Err(Status::invalid_argument( + "native registration supports tp_size=1 world_size=1 only", + )); + } + for tensor in &req.native_kv_tensors { + if tensor.size_bytes == 0 || tensor.block_stride_bytes == 0 { + return Err(Status::invalid_argument( + "native layer views need non-zero size_bytes and block_stride_bytes", + )); + } + let inside = tensor + .offset_bytes + .checked_add(tensor.size_bytes) + .is_some_and(|end| end <= req.native_alloc_size); + if !inside { + return Err(Status::invalid_argument(format!( + "native layer view [{} +{}] is outside the {}-byte arena", + tensor.offset_bytes, tensor.size_bytes, req.native_alloc_size + ))); + } + } + } Ok(()) } @@ -165,12 +249,6 @@ impl GrpcEngineService { Ok(()) } - fn build_register_context_response() -> RegisterContextResponse { - RegisterContextResponse { - status: Some(Self::ok_status()), - } - } - fn build_simple_response() -> ResponseStatus { Self::ok_status() } @@ -235,10 +313,11 @@ impl Engine for GrpcEngineService { ProtoTransferMode::Kernel => pegaflow_core::TransferMode::Kernel, }; - // Validate array lengths are consistent with each other. + // Validate array lengths are consistent with each other. The + // payload arrays (wrapper_bytes / native_kv_tensors) are already + // checked in validate_register_context_request. let batch_len = req.layer_names.len(); if batch_len == 0 - || req.wrapper_bytes.len() != batch_len || req.num_blocks.len() != batch_len || req.bytes_per_block.len() != batch_len || req.kv_stride_bytes.len() != batch_len @@ -249,6 +328,15 @@ impl Engine for GrpcEngineService { ))); } + let native = !req.native_kv_tensors.is_empty(); + // Extract native views before the conversions below consume `req` + // fields by value. + let native_views = if native { + Some(Self::native_layer_views(&req)?) + } else { + None + }; + let num_blocks_list: Vec = req .num_blocks .into_iter() @@ -275,24 +363,45 @@ impl Engine for GrpcEngineService { let tp_size = Self::usize_from_u32(req.tp_size, "tp_size")?; let world_size = Self::usize_from_u32(req.world_size, "world_size")?; - // Materialize tensors and collect data_ptr/size_bytes let context_key = Self::context_key(&req.instance_id, req.tp_rank, req.pp_rank, req.device_id); - // Materialize on the dedicated registry thread (GIL + CUDA IPC) and - // await the result, so this RPC never blocks an async worker. Move - // the (large) wrapper bytes over; clone the layer names since the - // engine call below still needs them. - let layers: Vec<(String, Vec)> = req - .layer_names - .iter() - .cloned() - .zip(req.wrapper_bytes) - .collect(); - let metadatas = self - .registry - .register_layers(context_key.clone(), req.device_id, layers) - .await - .map_err(|message| Status::internal(format!("register tensor failed: {message}")))?; + // Registry work (torch materialization or arena allocation) runs on + // the dedicated registry thread, so this RPC never blocks an async + // worker. + let (metadatas, arena_ipc_handle, block_stride_bytes) = if native { + let alloc_size = Self::usize_from_u64(req.native_alloc_size, "native_alloc_size")?; + let (layers, block_stride_bytes) = + native_views.expect("native_views populated when native"); + let registration = self + .registry + .register_native(context_key.clone(), req.device_id, layers, alloc_size) + .await + .map_err(|message| { + Status::internal(format!("register native arena failed: {message}")) + })?; + ( + registration.metadatas, + registration.arena_ipc_handle, + Some(block_stride_bytes), + ) + } else { + // Move the (large) wrapper bytes over; clone the layer names + // since the engine call below still needs them. + let layers: Vec<(String, Vec)> = req + .layer_names + .iter() + .cloned() + .zip(req.wrapper_bytes) + .collect(); + let metadatas = self + .registry + .register_layers(context_key.clone(), req.device_id, layers) + .await + .map_err(|message| { + Status::internal(format!("register tensor failed: {message}")) + })?; + (metadatas, Vec::new(), None) + }; let mut data_ptrs = Vec::with_capacity(batch_len); let mut size_bytes_list = Vec::with_capacity(batch_len); for metadata in &metadatas { @@ -301,7 +410,7 @@ impl Engine for GrpcEngineService { } // Call engine batch registration - if let Err(err) = self.engine.register_context_layer_batch( + if let Err(err) = self.engine.register_context_layer_batch_strided( &req.instance_id, &req.namespace, req.device_id, @@ -316,6 +425,7 @@ impl Engine for GrpcEngineService { &bytes_per_block_list, &kv_stride_bytes_list, &segments_list, + block_stride_bytes.as_deref(), transfer_mode, req.page_first, ) { @@ -330,7 +440,22 @@ impl Engine for GrpcEngineService { return Err(status); } - Ok(Response::new(Self::build_register_context_response())) + // A session/HTTP cleanup that raced this RPC between the registry + // step and the engine step has already freed the arena; the engine + // registration published above would point at freed memory. Confirm + // the context survived, or roll the engine back and fail. + if native && !self.registry.contains_context(context_key.clone()).await { + let _ = self.engine.unregister_instance(&req.instance_id); + return Err(Status::aborted(format!( + "instance {} was cleaned up while registering", + req.instance_id + ))); + } + + Ok(Response::new(RegisterContextResponse { + status: Some(Self::ok_status()), + arena_ipc_handle, + })) } .await; @@ -450,6 +575,7 @@ impl Engine for GrpcEngineService { layer_names, loads, load_state_shm, + wait_for_completion, .. } = req; Self::validate_device_id(device_id)?; @@ -475,16 +601,38 @@ impl Engine for GrpcEngineService { }) .collect::>()?; - self.engine - .batch_load_kv_blocks_multi_layer( - &instance_id, - tp_rank, - device_id, - &load_state_shm, - &layer_refs, - &loads, - ) - .map_err(Self::map_engine_error)?; + if wait_for_completion { + // Native clients have no load_state_shm completion flag to + // poll; run the in-process load and reply once DMA finished. + if !load_state_shm.is_empty() { + return Err(Status::invalid_argument( + "synchronous load must not include load_state_shm", + )); + } + self.engine + .batch_load_kv_blocks_multi_layer_inproc( + &instance_id, + tp_rank, + device_id, + &layer_refs, + &loads, + ) + .map_err(Self::map_engine_error)? + .await + .map_err(|_| Status::internal("load worker dropped completion"))? + .map_err(Self::map_engine_error)?; + } else { + self.engine + .batch_load_kv_blocks_multi_layer( + &instance_id, + tp_rank, + device_id, + &load_state_shm, + &layer_refs, + &loads, + ) + .map_err(Self::map_engine_error)?; + } Ok(Response::new(LoadResponse { status: Some(Self::build_simple_response()), @@ -510,6 +658,16 @@ impl Engine for GrpcEngineService { result } + async fn flush( + &self, + _request: Request, + ) -> Result, Status> { + self.engine.flush_saves_and_registrations().await; + Ok(Response::new(FlushResponse { + status: Some(Self::build_simple_response()), + })) + } + async fn query_prefetch( &self, request: Request, @@ -659,6 +817,15 @@ impl Engine for GrpcEngineService { let result: Result, Status> = async { let req = request.into_inner(); debug!("RPC [unregister_context]: instance_id={}", req.instance_id); + // Engine first: it must forget raw device pointers before the + // registry drop frees the memory behind them (native arenas are + // cuMemFree'd there); the flush barrier drains saves already in + // the write pipeline (see cleanup_instance for the residual + // GPU-worker-queue window). The registry drop still runs if the + // engine never saw the instance. + let engine_result = self.engine.unregister_instance(&req.instance_id); + self.engine.flush_saves().await; + let removed = self.registry.drop_instance(req.instance_id.clone()).await; if removed > 0 { info!( @@ -666,10 +833,7 @@ impl Engine for GrpcEngineService { removed, req.instance_id ); } - - self.engine - .unregister_instance(&req.instance_id) - .map_err(Self::map_engine_error)?; + engine_result.map_err(Self::map_engine_error)?; Ok(Response::new(UnregisterResponse { status: Some(Self::build_simple_response()), @@ -994,7 +1158,7 @@ mod tests { let err = GrpcEngineService::validate_register_context_request(&RegisterContextRequest { instance_id: "instance".to_string(), namespace: "namespace".to_string(), - client_version: env!("CARGO_PKG_VERSION").to_string(), + client_version: pegaflow_proto::VERSION.to_string(), tp_rank: 1, tp_size: 1, world_size: 1, @@ -1008,6 +1172,8 @@ mod tests { pp_rank: 0, transfer_mode: ProtoTransferMode::Direct as i32, page_first: false, + native_kv_tensors: Vec::new(), + native_alloc_size: 0, }) .expect_err("tp_rank outside tp_size must be rejected at RPC boundary"); @@ -1034,6 +1200,8 @@ mod tests { pp_rank: 0, transfer_mode: ProtoTransferMode::Direct as i32, page_first: false, + native_kv_tensors: Vec::new(), + native_alloc_size: 0, }) .expect_err("client/server version mismatch must be rejected before registration"); diff --git a/pegaflow-server/tests/common/mod.rs b/pegaflow-server/tests/common/mod.rs index 2a4669b8..e0d94ec5 100644 --- a/pegaflow-server/tests/common/mod.rs +++ b/pegaflow-server/tests/common/mod.rs @@ -378,6 +378,7 @@ impl MockVllmRpcHarness { lease, block_ids: (0..block_count as u32).collect(), }], + wait_for_completion: false, }; match self.worker.load(request.clone()).await { Ok(response) => Ok(LoadRpcExchange { diff --git a/pegaflow-server/tests/http_cleanup_hang_repro.rs b/pegaflow-server/tests/http_cleanup_hang_repro.rs index cdf3ad23..816a1cad 100644 --- a/pegaflow-server/tests/http_cleanup_hang_repro.rs +++ b/pegaflow-server/tests/http_cleanup_hang_repro.rs @@ -235,7 +235,7 @@ fn wedge_register_request(i: usize) -> RegisterContextRequest { RegisterContextRequest { instance_id: format!("wedge-{i}"), namespace: "wedge".to_string(), - client_version: env!("CARGO_PKG_VERSION").to_string(), + client_version: pegaflow_proto::VERSION.to_string(), tp_rank: 0, tp_size: 1, world_size: 1, @@ -249,6 +249,8 @@ fn wedge_register_request(i: usize) -> RegisterContextRequest { pp_rank: 0, transfer_mode: TransferMode::Direct as i32, page_first: false, + native_kv_tensors: Vec::new(), + native_alloc_size: 0, } } diff --git a/pegaflow-server/tests/native_arena_rpc_e2e.rs b/pegaflow-server/tests/native_arena_rpc_e2e.rs new file mode 100644 index 00000000..66861894 --- /dev/null +++ b/pegaflow-server/tests/native_arena_rpc_e2e.rs @@ -0,0 +1,345 @@ +//! Native arena registration + save/load over gRPC. +//! +//! Covers the torch-free path end to end: +//! `RegisterContextBatch(native_*)` → server allocates the arena and returns a +//! CUDA IPC handle → a child process imports it (CUDA forbids importing an IPC +//! handle in the exporting process) and writes a pattern → D2H save → child +//! wipes the arena → H2D load with `wait_for_completion` → child verifies the +//! restored bytes. +//! +//! Run: `cargo test -p pegaflow-server --test native_arena_rpc_e2e --features cuda-13,rdma` + +use std::ffi::c_void; +use std::net::{SocketAddr, TcpListener}; +use std::process::Command; +use std::sync::Arc; +use std::time::Duration; + +use cudarc::driver::sys; +use pegaflow_core::{PegaEngine, StorageConfig}; +use pegaflow_server::proto::engine::engine_client::EngineClient; +use pegaflow_server::proto::engine::engine_server::EngineServer; +use pegaflow_server::proto::engine::{ + LeaseLoad, LoadRequest, NativeKvTensor, QueryRequest, RegisterContextRequest, SaveLayer, + SaveRequest, TransferMode, UnregisterRequest, query_response, +}; +use pegaflow_server::{CudaTensorRegistry, GrpcEngineService, RegistryHandle}; +use tokio::sync::Notify; +use tonic::transport::Server; + +const INSTANCE_ID: &str = "native-arena-rpc-e2e"; +const NAMESPACE: &str = "native-arena"; +const LAYER_NAME: &str = "layer_0"; +const BLOCK_COUNT: usize = 4; +const BYTES_PER_BLOCK: usize = 1024; +const TOTAL_BYTES: usize = BLOCK_COUNT * BYTES_PER_BLOCK; + +const MODE_ENV: &str = "PEGAFLOW_ARENA_E2E_MODE"; +const HANDLE_ENV: &str = "PEGAFLOW_ARENA_E2E_HANDLE_HEX"; +const OUT_ENV: &str = "PEGAFLOW_ARENA_E2E_OUT"; + +#[tokio::test] +async fn native_arena_register_save_load_roundtrip() { + let engine = Arc::new( + PegaEngine::new_with_config( + 16 << 20, + false, + StorageConfig { + enable_lfu_admission: false, + ..StorageConfig::default() + }, + ) + .expect("engine"), + ); + // Torch-free registry: native registration only. + let registry = RegistryHandle::spawn(CudaTensorRegistry::empty()); + let port = unused_port(); + let addr: SocketAddr = ([127, 0, 0, 1], port).into(); + let shutdown = Arc::new(Notify::new()); + let hll = Arc::new(std::sync::Mutex::new( + pegaflow_common::hll::MultiWindowHllTracker::new( + vec![("24h".into(), Duration::from_secs(86400))], + 14, + ), + )); + let service = GrpcEngineService::new(Arc::clone(&engine), registry, Arc::clone(&shutdown), hll); + let server = tokio::spawn(async move { + Server::builder() + .add_service(EngineServer::new(service)) + .serve(addr) + .await + .expect("serve"); + }); + + let mut client = connect(&format!("http://127.0.0.1:{port}")).await; + + // 1) Native RegisterContextBatch: the server allocates the arena and + // returns its CUDA IPC handle. + let reg = client + .register_context_batch(RegisterContextRequest { + instance_id: INSTANCE_ID.to_string(), + namespace: NAMESPACE.to_string(), + client_version: pegaflow_proto::VERSION.to_string(), + tp_rank: 0, + tp_size: 1, + world_size: 1, + device_id: 0, + layer_names: vec![LAYER_NAME.to_string()], + wrapper_bytes: vec![], + num_blocks: vec![BLOCK_COUNT as u64], + bytes_per_block: vec![BYTES_PER_BLOCK as u64], + kv_stride_bytes: vec![0], + segments: vec![1], + pp_rank: 0, + transfer_mode: TransferMode::Direct as i32, + page_first: false, + native_kv_tensors: vec![NativeKvTensor { + offset_bytes: 0, + size_bytes: TOTAL_BYTES as u64, + block_stride_bytes: BYTES_PER_BLOCK as u64, + }], + native_alloc_size: TOTAL_BYTES as u64, + }) + .await + .expect("register_context_batch") + .into_inner(); + assert!( + reg.status.as_ref().is_some_and(|s| s.ok), + "register failed: {:?}", + reg.status + ); + assert_eq!( + reg.arena_ipc_handle.len(), + 64, + "expected a CUipcMemHandle in the response" + ); + let handle_hex = hex_encode(®.arena_ipc_handle); + + // 2) The "client" (a separate process, as CUDA IPC requires) writes the + // pattern into its imported mapping. + run_child("write", &handle_hex, None); + + // 3) D2H save reads through the server's own arena pointers. + let hashes: Vec> = (0..BLOCK_COUNT) + .map(|i| { + let mut h = vec![7u8]; + h.extend_from_slice(&(i as u32).to_le_bytes()); + h + }) + .collect(); + let save = client + .save(SaveRequest { + instance_id: INSTANCE_ID.to_string(), + tp_rank: 0, + device_id: 0, + pp_rank: 0, + saves: vec![SaveLayer { + layer_name: LAYER_NAME.to_string(), + block_ids: (0..BLOCK_COUNT as u32).collect(), + block_hashes: hashes.clone(), + }], + }) + .await + .expect("save") + .into_inner(); + assert!( + save.status.as_ref().is_some_and(|s| s.ok), + "{:?}", + save.status + ); + engine.flush_saves().await; + + // 4) Query hits. + let query = client + .query_prefetch(QueryRequest { + instance_id: INSTANCE_ID.to_string(), + block_hashes: hashes.clone(), + req_id: "native-arena-hit".into(), + wait_for_full_prefix: false, + }) + .await + .expect("query") + .into_inner(); + let ready = match query.outcome { + Some(query_response::Outcome::Ready(r)) => r, + other => panic!("expected Ready, got {other:?}"), + }; + assert_eq!(ready.num_hit_blocks as usize, BLOCK_COUNT); + assert!(!ready.lease.is_empty()); + + // 5) Client wipes the arena, then a synchronous load restores it. + run_child("wipe", &handle_hex, None); + let load = client + .load(LoadRequest { + instance_id: INSTANCE_ID.to_string(), + tp_rank: 0, + device_id: 0, + load_state_shm: String::new(), + layer_names: vec![LAYER_NAME.to_string()], + loads: vec![LeaseLoad { + lease: ready.lease, + block_ids: (0..BLOCK_COUNT as u32).collect(), + }], + wait_for_completion: true, + }) + .await + .expect("load") + .into_inner(); + assert!( + load.status.as_ref().is_some_and(|s| s.ok), + "{:?}", + load.status + ); + + // 6) Client reads its mapping back; must be bit-identical to the pattern. + let out = std::env::temp_dir().join(format!("pegaflow-arena-e2e-{}", std::process::id())); + run_child("dump", &handle_hex, Some(&out)); + let restored = std::fs::read(&out).expect("read child dump"); + let _ = std::fs::remove_file(&out); + let mut expected = vec![0u8; TOTAL_BYTES]; + fill_pattern(&mut expected); + assert_eq!( + restored, expected, + "restored arena must match saved pattern" + ); + + // 7) Unregister frees the arena. + let unreg = client + .unregister_context(UnregisterRequest { + instance_id: INSTANCE_ID.to_string(), + }) + .await + .expect("unregister") + .into_inner(); + assert!(unreg.status.as_ref().is_some_and(|s| s.ok)); + + server.abort(); +} + +/// Child-process entry point, driven by `run_child`. Ignored so plain +/// `cargo test` never runs it directly. +#[test] +#[ignore = "spawned as a helper process by native_arena_register_save_load_roundtrip"] +fn arena_ipc_child_helper() { + let mode = std::env::var(MODE_ENV).expect("child mode"); + let handle = hex_decode(&std::env::var(HANDLE_ENV).expect("child handle")); + assert_eq!(handle.len(), 64); + + let ctx = cudarc::driver::CudaContext::new(0).expect("CUDA device 0"); + ctx.bind_to_thread().expect("bind CUDA context"); + + let mut ipc_handle = sys::CUipcMemHandle { reserved: [0; 64] }; + for (dst, src) in ipc_handle.reserved.iter_mut().zip(&handle) { + *dst = *src as i8; + } + let mut base_ptr: sys::CUdeviceptr = 0; + // SAFETY: the handle references a live allocation owned by the parent + // process (the server keeps it registered while children run). + check_cuda( + unsafe { + sys::cuIpcOpenMemHandle_v2( + &mut base_ptr, + ipc_handle, + sys::CUipcMem_flags_enum::CU_IPC_MEM_LAZY_ENABLE_PEER_ACCESS as u32, + ) + }, + "cuIpcOpenMemHandle", + ); + + match mode.as_str() { + "write" => { + let mut pattern = vec![0u8; TOTAL_BYTES]; + fill_pattern(&mut pattern); + // SAFETY: the imported mapping spans TOTAL_BYTES. + check_cuda( + unsafe { + sys::cuMemcpyHtoD_v2(base_ptr, pattern.as_ptr() as *const c_void, TOTAL_BYTES) + }, + "cuMemcpyHtoD", + ); + } + "wipe" => { + // SAFETY: same mapping bounds as above. + check_cuda( + unsafe { sys::cuMemsetD8_v2(base_ptr, 0, TOTAL_BYTES) }, + "cuMemsetD8", + ); + } + "dump" => { + let out = std::env::var(OUT_ENV).expect("child out path"); + let mut bytes = vec![0u8; TOTAL_BYTES]; + // SAFETY: same mapping bounds as above. + check_cuda( + unsafe { + sys::cuMemcpyDtoH_v2(bytes.as_mut_ptr() as *mut c_void, base_ptr, TOTAL_BYTES) + }, + "cuMemcpyDtoH", + ); + std::fs::write(out, bytes).expect("write dump"); + } + other => panic!("unknown child mode {other}"), + } + + // SAFETY: base_ptr came from cuIpcOpenMemHandle above. + check_cuda( + unsafe { sys::cuIpcCloseMemHandle(base_ptr) }, + "cuIpcCloseMemHandle", + ); +} + +fn run_child(mode: &str, handle_hex: &str, out: Option<&std::path::Path>) { + let exe = std::env::current_exe().expect("current_exe"); + let mut cmd = Command::new(exe); + cmd.args(["arena_ipc_child_helper", "--exact", "--include-ignored"]) + .env(MODE_ENV, mode) + .env(HANDLE_ENV, handle_hex); + if let Some(out) = out { + cmd.env(OUT_ENV, out); + } + let output = cmd.output().expect("spawn child helper"); + assert!( + output.status.success(), + "child '{mode}' failed:\nstdout: {}\nstderr: {}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); +} + +fn fill_pattern(buf: &mut [u8]) { + for (i, byte) in buf.iter_mut().enumerate() { + *byte = ((i * 31 + 7) % 251) as u8; + } +} + +fn hex_encode(bytes: &[u8]) -> String { + bytes.iter().map(|b| format!("{b:02x}")).collect() +} + +fn hex_decode(hex: &str) -> Vec { + hex.as_bytes() + .chunks(2) + .map(|pair| u8::from_str_radix(std::str::from_utf8(pair).unwrap(), 16).unwrap()) + .collect() +} + +fn unused_port() -> u16 { + TcpListener::bind("127.0.0.1:0") + .expect("bind ephemeral port") + .local_addr() + .expect("local addr") + .port() +} + +async fn connect(endpoint: &str) -> EngineClient { + for _ in 0..50 { + if let Ok(client) = EngineClient::connect(endpoint.to_string()).await { + return client; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + panic!("gRPC server did not come up at {endpoint}"); +} + +fn check_cuda(result: sys::CUresult, op: &str) { + assert_eq!(result, sys::CUresult::CUDA_SUCCESS, "{op} failed"); +} diff --git a/python/src/lib.rs b/python/src/lib.rs index 60f024a4..0be636fa 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -308,7 +308,7 @@ impl EngineRpcClient { .register_context_batch(RegisterContextRequest { instance_id, namespace, - client_version: env!("CARGO_PKG_VERSION").to_string(), + client_version: pegaflow_proto::VERSION.to_string(), tp_rank, tp_size, world_size, @@ -322,6 +322,8 @@ impl EngineRpcClient { pp_rank, transfer_mode: transfer_mode as i32, page_first, + native_kv_tensors: Vec::new(), + native_alloc_size: 0, }) .await?; Ok(resp.into_inner()) @@ -419,6 +421,7 @@ impl EngineRpcClient { load_state_shm, layer_names, loads, + wait_for_completion: false, }) .await?; Ok(resp.into_inner())