diff --git a/Cargo.lock b/Cargo.lock index e6fe324..1e467f2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2925,6 +2925,7 @@ name = "mvp-system" version = "0.1.0" dependencies = [ "libc", + "serde_json", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index c9e8b58..5d7a08c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -59,9 +59,6 @@ proptest = "1" proptest-state-machine = "0.3" mvp-system = { path = "crates/mvp-system" } -[[test]] -name = "arena_manager_guarantees" -path = "tests/mvp_system/arena_manager_guarantees.rs" [[bench]] name = "runtime_benchmarks" diff --git a/crates/mvp-system/Cargo.toml b/crates/mvp-system/Cargo.toml index 4b770e2..dec66a4 100644 --- a/crates/mvp-system/Cargo.toml +++ b/crates/mvp-system/Cargo.toml @@ -4,5 +4,9 @@ version = "0.1.0" edition = "2024" publish = false + +[dependencies] +serde_json = "1" [target.'cfg(target_os = "linux")'.dependencies] libc = "0.2" + diff --git a/crates/mvp-system/src/arena_manager.rs b/crates/mvp-system/src/arena_manager.rs new file mode 100644 index 0000000..ff9ceec --- /dev/null +++ b/crates/mvp-system/src/arena_manager.rs @@ -0,0 +1,548 @@ +use std::collections::{BTreeMap, VecDeque}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct NodeId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct LeaseRequestId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RingId(pub u64); + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ArenaConfig { + pub node_id: NodeId, + pub reservation_ceiling: u64, + pub base_alignment: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RingSpec { + pub header_bytes: u64, + pub data_bytes: u64, + pub alignment: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct LeaseRing { + pub request_id: LeaseRequestId, + pub ring_spec: RingSpec, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum LayoutPointer { + NoProcessPointer, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RingLayout { + pub start_offset: u64, + pub header_offset: u64, + pub data_offset: u64, + pub end_offset: u64, + pub data_bytes: u64, + pub alignment: u64, + pub pointer: LayoutPointer, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RingLease { + pub request_id: LeaseRequestId, + pub ring_id: RingId, + pub node_id: NodeId, + pub layout: RingLayout, + pub requested_alignment: u64, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct QuiescenceProof { + verified: bool, +} + +impl QuiescenceProof { + pub fn verified() -> Self { + Self { verified: true } + } + + pub fn missing() -> Self { + Self { verified: false } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ArenaRequest { + LeaseRing(LeaseRing), + CancelLease { + request_id: LeaseRequestId, + }, + ReleaseRing { + ring_id: RingId, + proof: QuiescenceProof, + }, + Shutdown, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ArenaFault { + InvalidReservationCeiling, + InvalidBaseAlignment, + BackingUnavailable, + UnsupportedPlatform, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RingLeaseRejection { + CannotFitWithinCeiling, + ArenaShuttingDown, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RingReleaseRejection { + MissingQuiescenceProof, + UnknownRingId, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ArenaEvent { + RingLeased { + lease: RingLease, + }, + RingLeaseQueued { + request_id: LeaseRequestId, + }, + RingLeaseRejected { + request_id: LeaseRequestId, + reason: RingLeaseRejection, + }, + RingReleased { + ring_id: RingId, + start_offset: u64, + end_offset: u64, + }, + RingReleaseRejected { + ring_id: RingId, + reason: RingReleaseRejection, + }, + CancelledFreshLeaseReleased { + request_id: LeaseRequestId, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ArenaCommand { + InstallWorkerOrPumpState { + request_id: LeaseRequestId, + ring_id: RingId, + layout: RingLayout, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ArenaState { + Ready, + ShuttingDown, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct FreeRange { + start: u64, + end: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct QueuedLease { + request: LeaseRing, + cancelled: bool, +} + +pub struct ArenaManager { + config: ArenaConfig, + _backing: ArenaBacking, + state: ArenaState, + next_ring_id: u64, + free_ranges: Vec, + pending: VecDeque, + live_order: Vec, + live_index: BTreeMap, +} + +impl ArenaManager { + pub fn boot(config: ArenaConfig) -> Result { + let backing = ArenaBacking::create(&config)?; + Ok(Self { + free_ranges: vec![FreeRange { + start: 0, + end: config.reservation_ceiling, + }], + config, + _backing: backing, + state: ArenaState::Ready, + next_ring_id: 1, + pending: VecDeque::new(), + live_order: Vec::new(), + live_index: BTreeMap::new(), + }) + } + + fn request(&mut self, request: ArenaRequest) -> Vec { + match request { + ArenaRequest::LeaseRing(request) => self.lease_ring(request), + ArenaRequest::CancelLease { request_id } => { + self.cancel_lease(request_id); + Vec::new() + } + ArenaRequest::ReleaseRing { ring_id, proof } => self.release_ring(ring_id, proof), + ArenaRequest::Shutdown => self.shutdown(), + } + } + + fn live_leases(&self) -> &[RingLease] { + &self.live_order + } + + fn lookup_lease(&self, ring_id: RingId) -> Option<&RingLease> { + self.live_index + .get(&ring_id) + .and_then(|index| self.live_order.get(*index)) + } + + fn lease_ring(&mut self, request: LeaseRing) -> Vec { + if self.state == ArenaState::ShuttingDown { + return vec![ArenaEvent::RingLeaseRejected { + request_id: request.request_id, + reason: RingLeaseRejection::ArenaShuttingDown, + }]; + } + + if self.request_layout_at(&request.ring_spec, 0).is_none() { + return vec![ArenaEvent::RingLeaseRejected { + request_id: request.request_id, + reason: RingLeaseRejection::CannotFitWithinCeiling, + }]; + } + + if self.pending.is_empty() { + if let Some(lease) = self.try_allocate(&request) { + return vec![ArenaEvent::RingLeased { lease }]; + } + } + + let request_id = request.request_id; + self.pending.push_back(QueuedLease { + request, + cancelled: false, + }); + vec![ArenaEvent::RingLeaseQueued { request_id }] + } + + fn cancel_lease(&mut self, request_id: LeaseRequestId) { + for pending in &mut self.pending { + if pending.request.request_id == request_id { + pending.cancelled = true; + } + } + } + + fn release_ring(&mut self, ring_id: RingId, proof: QuiescenceProof) -> Vec { + let Some(index) = self.live_index.get(&ring_id).copied() else { + return vec![ArenaEvent::RingReleaseRejected { + ring_id, + reason: RingReleaseRejection::UnknownRingId, + }]; + }; + + if !proof.verified { + return vec![ArenaEvent::RingReleaseRejected { + ring_id, + reason: RingReleaseRejection::MissingQuiescenceProof, + }]; + } + + let lease = self.live_order.remove(index); + self.rebuild_live_index(); + self.insert_free_range(FreeRange { + start: lease.layout.start_offset, + end: lease.layout.end_offset, + }); + + let mut events = vec![ArenaEvent::RingReleased { + ring_id, + start_offset: lease.layout.start_offset, + end_offset: lease.layout.end_offset, + }]; + if self.state == ArenaState::Ready { + self.retry_pending_leases(&mut events); + } + events + } + + fn shutdown(&mut self) -> Vec { + self.state = ArenaState::ShuttingDown; + let mut events = Vec::new(); + while let Some(pending) = self.pending.pop_front() { + if pending.cancelled { + events.push(ArenaEvent::CancelledFreshLeaseReleased { + request_id: pending.request.request_id, + }); + } else { + events.push(ArenaEvent::RingLeaseRejected { + request_id: pending.request.request_id, + reason: RingLeaseRejection::ArenaShuttingDown, + }); + } + } + events + } + + fn retry_pending_leases(&mut self, events: &mut Vec) { + while let Some(pending) = self.pending.pop_front() { + if pending.cancelled { + events.push(ArenaEvent::CancelledFreshLeaseReleased { + request_id: pending.request.request_id, + }); + continue; + } + + if let Some(lease) = self.try_allocate(&pending.request) { + events.push(ArenaEvent::RingLeased { lease }); + continue; + } + + self.pending.push_front(pending); + break; + } + } + + fn try_allocate(&mut self, request: &LeaseRing) -> Option { + for index in 0..self.free_ranges.len() { + let range = self.free_ranges[index]; + let Some(layout) = self.request_layout_at(&request.ring_spec, range.start) else { + continue; + }; + if layout.end_offset > range.end { + continue; + } + + self.free_ranges.remove(index); + let mut insert_index = index; + if range.start < layout.start_offset { + self.free_ranges.insert( + insert_index, + FreeRange { + start: range.start, + end: layout.start_offset, + }, + ); + insert_index += 1; + } + if layout.end_offset < range.end { + self.free_ranges.insert( + insert_index, + FreeRange { + start: layout.end_offset, + end: range.end, + }, + ); + } + + let ring_id = RingId(self.next_ring_id); + self.next_ring_id = self.next_ring_id.checked_add(1)?; + let lease = RingLease { + request_id: request.request_id, + ring_id, + node_id: self.config.node_id, + layout, + requested_alignment: request.ring_spec.alignment, + }; + self.live_index.insert(ring_id, self.live_order.len()); + self.live_order.push(lease.clone()); + return Some(lease); + } + None + } + + fn request_layout_at(&self, spec: &RingSpec, range_start: u64) -> Option { + let alignment = lcm(self.config.base_alignment, spec.alignment)?; + let start_offset = align_up(range_start, alignment)?; + let header_offset = start_offset; + let after_header = header_offset.checked_add(spec.header_bytes)?; + let data_offset = align_up(after_header, alignment)?; + let end_offset = data_offset.checked_add(spec.data_bytes)?; + if end_offset > self.config.reservation_ceiling { + return None; + } + Some(RingLayout { + start_offset, + header_offset, + data_offset, + end_offset, + data_bytes: spec.data_bytes, + alignment, + pointer: LayoutPointer::NoProcessPointer, + }) + } + + fn insert_free_range(&mut self, range: FreeRange) { + self.free_ranges.push(range); + self.free_ranges.sort_by_key(|range| range.start); + + let mut coalesced: Vec = Vec::with_capacity(self.free_ranges.len()); + for range in self.free_ranges.drain(..) { + if range.start == range.end { + continue; + } + if let Some(last) = coalesced.last_mut() { + if range.start <= last.end { + last.end = last.end.max(range.end); + continue; + } + } + coalesced.push(range); + } + self.free_ranges = coalesced; + } + + fn rebuild_live_index(&mut self) { + self.live_index.clear(); + for (index, lease) in self.live_order.iter().enumerate() { + self.live_index.insert(lease.ring_id, index); + } + } +} + +pub struct ArenaManagerHarness { + manager: ArenaManager, + events: Vec, + commands: Vec, +} + +impl ArenaManagerHarness { + pub fn boot(config: ArenaConfig) -> Result { + Ok(Self { + manager: ArenaManager::boot(config)?, + events: Vec::new(), + commands: Vec::new(), + }) + } + + pub fn request(&mut self, request: ArenaRequest) { + self.events.extend(self.manager.request(request)); + } + + pub fn events(&self) -> &[ArenaEvent] { + &self.events + } + + pub fn commands(&self) -> &[ArenaCommand] { + &self.commands + } + + pub fn live_leases(&self) -> &[RingLease] { + self.manager.live_leases() + } + + pub fn lookup_lease(&self, ring_id: RingId) -> Option<&RingLease> { + self.manager.lookup_lease(ring_id) + } +} + +fn align_up(value: u64, alignment: u64) -> Option { + if alignment == 0 { + return None; + } + let remainder = value % alignment; + if remainder == 0 { + Some(value) + } else { + value.checked_add(alignment - remainder) + } +} + +fn lcm(left: u64, right: u64) -> Option { + if left == 0 || right == 0 { + return None; + } + let gcd = gcd(left, right); + (left / gcd).checked_mul(right) +} + +fn gcd(mut left: u64, mut right: u64) -> u64 { + while right != 0 { + let remainder = left % right; + left = right; + right = remainder; + } + left +} + +#[cfg(target_os = "linux")] +struct ArenaBacking { + len: usize, + ptr: *mut libc::c_void, + _fd: std::os::fd::OwnedFd, +} + +#[cfg(target_os = "linux")] +impl ArenaBacking { + fn create(config: &ArenaConfig) -> Result { + if config.reservation_ceiling == 0 { + return Err(ArenaFault::InvalidReservationCeiling); + } + if config.base_alignment == 0 { + return Err(ArenaFault::InvalidBaseAlignment); + } + + let len = usize::try_from(config.reservation_ceiling) + .map_err(|_| ArenaFault::InvalidReservationCeiling)?; + + let fd = unsafe { + let name = b"mvp-system-arena\0"; + libc::memfd_create(name.as_ptr().cast(), libc::MFD_CLOEXEC) + }; + if fd < 0 { + return Err(ArenaFault::BackingUnavailable); + } + + let fd = unsafe { ::from_raw_fd(fd) }; + let truncate_result = + unsafe { libc::ftruncate(std::os::fd::AsRawFd::as_raw_fd(&fd), len as libc::off_t) }; + if truncate_result != 0 { + return Err(ArenaFault::BackingUnavailable); + } + + let ptr = unsafe { + libc::mmap( + std::ptr::null_mut(), + len, + libc::PROT_READ | libc::PROT_WRITE, + libc::MAP_SHARED, + std::os::fd::AsRawFd::as_raw_fd(&fd), + 0, + ) + }; + if ptr == libc::MAP_FAILED { + return Err(ArenaFault::BackingUnavailable); + } + + Ok(Self { len, ptr, _fd: fd }) + } +} + +#[cfg(target_os = "linux")] +impl Drop for ArenaBacking { + fn drop(&mut self) { + unsafe { + libc::munmap(self.ptr, self.len); + } + } +} + +#[cfg(not(target_os = "linux"))] +struct ArenaBacking; + +#[cfg(not(target_os = "linux"))] +impl ArenaBacking { + fn create(_config: &ArenaConfig) -> Result { + Err(ArenaFault::UnsupportedPlatform) + } +} diff --git a/crates/mvp-system/src/device_bridge.rs b/crates/mvp-system/src/device_bridge.rs new file mode 100644 index 0000000..6569b87 --- /dev/null +++ b/crates/mvp-system/src/device_bridge.rs @@ -0,0 +1,633 @@ +use std::collections::{BTreeMap, BTreeSet}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct WorkerGeneration(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct StepId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct CopyEvent(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct DeviceHandle { + pub generation: WorkerGeneration, + pub id: u64, +} + +impl DeviceHandle { + pub const fn new(generation: WorkerGeneration, id: u64) -> Self { + Self { generation, id } + } +} + +pub type DeviceAllocation = DeviceHandle; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum DType { + U32, + F16, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Shape { + Vector, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ObjectSpec { + pub max_extent: u64, + pub alignment: u64, + pub dtype: DType, + pub shape: Shape, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct TensorViewSpec { + pub dtype: DType, + pub shape: Shape, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct TensorView { + pub handle: DeviceHandle, + pub dtype: DType, + pub shape: Shape, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct HostRange { + pub offset: u64, + pub len: u64, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct DeviceRange { + pub handle: DeviceHandle, + pub offset: u64, + pub len: u64, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CopyMode { + Sync, + Async, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum DeviceError { + AllocationFailed, + InvalidExtent, + InvalidAlignment, + InvalidRange, + UnknownDeviceHandle, + OldGenerationHandle, + DeviceCopyFailed, + InvalidViewDType, + InvalidTensorView, + AllocationStillInUse, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum DeviceEvent { + CopyCompleted { + copy: CopyEvent, + }, + ComputeStarted { + handle: DeviceHandle, + step_id: StepId, + }, + ComputeCompleted { + handle: DeviceHandle, + step_id: StepId, + }, + WorkerRestarted { + generation: WorkerGeneration, + }, +} + +pub trait DeviceBridgeBackend { + fn alloc_device( + &mut self, + allocation: DeviceAllocation, + spec: ObjectSpec, + extent: u64, + ) -> Result<(), DeviceError>; + + fn free_device(&mut self, allocation: DeviceAllocation) -> Result<(), DeviceError>; + + fn host_to_device( + &mut self, + host: HostRange, + device: DeviceRange, + copy: CopyEvent, + mode: CopyMode, + ) -> Result<(), DeviceError>; + + fn device_to_host( + &mut self, + device: DeviceRange, + host: HostRange, + copy: CopyEvent, + mode: CopyMode, + ) -> Result<(), DeviceError>; + + fn wrap_for_tinygrad( + &mut self, + allocation: DeviceAllocation, + view: TensorViewSpec, + ) -> Result; +} + +pub struct DeviceBridge { + current_generation: WorkerGeneration, + next_handle_id: u64, + next_copy_id: u64, + allocations: BTreeMap, + copies: BTreeMap, + backend: B, +} + +impl DeviceBridge { + pub fn new(current_generation: WorkerGeneration, backend: B) -> Self { + Self { + current_generation, + next_handle_id: 1, + next_copy_id: 1, + allocations: BTreeMap::new(), + copies: BTreeMap::new(), + backend, + } + } + + pub fn alloc_device( + &mut self, + spec: ObjectSpec, + extent: u64, + ) -> Result { + validate_object_spec(spec, extent)?; + + let allocation = DeviceHandle::new(self.current_generation, self.next_handle_id); + self.next_handle_id += 1; + self.backend.alloc_device(allocation, spec, extent)?; + self.allocations.insert( + allocation, + AllocationRecord { + spec, + extent, + active_copies: BTreeSet::new(), + active_steps: BTreeSet::new(), + }, + ); + Ok(allocation) + } + + pub fn free_device(&mut self, allocation: DeviceAllocation) -> Result<(), DeviceError> { + self.validate_generation(allocation)?; + let record = self + .allocations + .get(&allocation) + .ok_or(DeviceError::UnknownDeviceHandle)?; + if !record.active_copies.is_empty() || !record.active_steps.is_empty() { + return Err(DeviceError::AllocationStillInUse); + } + + self.backend.free_device(allocation)?; + self.allocations.remove(&allocation); + Ok(()) + } + + pub fn host_to_device( + &mut self, + host: HostRange, + device: DeviceRange, + mode: CopyMode, + ) -> Result { + self.validate_device_range(device)?; + validate_equal_copy_len(host.len, device.len)?; + + let copy = self.next_copy_event(); + self.backend.host_to_device(host, device, copy, mode)?; + self.record_copy(copy, device.handle, CopyDirection::HostToDevice, mode); + Ok(copy) + } + + pub fn device_to_host( + &mut self, + device: DeviceRange, + host: HostRange, + mode: CopyMode, + ) -> Result { + self.validate_device_range(device)?; + validate_equal_copy_len(device.len, host.len)?; + + let copy = self.next_copy_event(); + self.backend.device_to_host(device, host, copy, mode)?; + self.record_copy(copy, device.handle, CopyDirection::DeviceToHost, mode); + Ok(copy) + } + + pub fn wrap_for_tinygrad( + &mut self, + allocation: DeviceAllocation, + view: TensorViewSpec, + ) -> Result { + let record = self.allocation_record(allocation)?; + if view.dtype != record.spec.dtype { + return Err(DeviceError::InvalidViewDType); + } + if view.shape != record.spec.shape { + return Err(DeviceError::InvalidTensorView); + } + + self.backend.wrap_for_tinygrad(allocation, view) + } + + pub fn observe(&mut self, event: DeviceEvent) { + match event { + DeviceEvent::CopyCompleted { copy } => self.complete_copy(copy), + DeviceEvent::ComputeStarted { handle, step_id } => { + if let Some(record) = self.current_allocation_record_mut(handle) { + record.active_steps.insert(step_id); + } + } + DeviceEvent::ComputeCompleted { handle, step_id } => { + if let Some(record) = self.current_allocation_record_mut(handle) { + record.active_steps.remove(&step_id); + } + } + DeviceEvent::WorkerRestarted { generation } => { + self.current_generation = generation; + self.next_handle_id = 1; + self.next_copy_id = 1; + self.allocations.clear(); + self.copies.clear(); + } + } + } + + pub fn safe_to_release_host(&self, copy: CopyEvent) -> bool { + self.copies + .get(©) + .map(|record| record.direction == CopyDirection::HostToDevice && record.completed) + .unwrap_or(false) + } + + pub fn host_bytes_valid(&self, copy: CopyEvent) -> bool { + self.copies + .get(©) + .map(|record| record.direction == CopyDirection::DeviceToHost && record.completed) + .unwrap_or(false) + } + + pub fn copy_event_complete(&self, copy: CopyEvent) -> bool { + self.copies + .get(©) + .map(|record| record.completed) + .unwrap_or(false) + } + + pub fn backend(&self) -> &B { + &self.backend + } + + pub fn backend_mut(&mut self) -> &mut B { + &mut self.backend + } + + fn allocation_record( + &self, + allocation: DeviceAllocation, + ) -> Result<&AllocationRecord, DeviceError> { + self.validate_generation(allocation)?; + self.allocations + .get(&allocation) + .ok_or(DeviceError::UnknownDeviceHandle) + } + + fn current_allocation_record_mut( + &mut self, + allocation: DeviceAllocation, + ) -> Option<&mut AllocationRecord> { + if allocation.generation != self.current_generation { + return None; + } + self.allocations.get_mut(&allocation) + } + + fn validate_generation(&self, allocation: DeviceAllocation) -> Result<(), DeviceError> { + if allocation.generation != self.current_generation { + return Err(DeviceError::OldGenerationHandle); + } + Ok(()) + } + + fn validate_device_range(&self, device: DeviceRange) -> Result<(), DeviceError> { + let record = self.allocation_record(device.handle)?; + let end = device + .offset + .checked_add(device.len) + .ok_or(DeviceError::InvalidRange)?; + if end > record.extent { + return Err(DeviceError::InvalidRange); + } + Ok(()) + } + + fn next_copy_event(&mut self) -> CopyEvent { + let copy = CopyEvent(self.next_copy_id); + self.next_copy_id += 1; + copy + } + + fn record_copy( + &mut self, + copy: CopyEvent, + allocation: DeviceAllocation, + direction: CopyDirection, + mode: CopyMode, + ) { + let completed = mode == CopyMode::Sync; + self.copies.insert( + copy, + CopyRecord { + allocation, + direction, + completed, + }, + ); + if !completed { + if let Some(record) = self.allocations.get_mut(&allocation) { + record.active_copies.insert(copy); + } + } + } + + fn complete_copy(&mut self, copy: CopyEvent) { + let Some(record) = self.copies.get_mut(©) else { + return; + }; + record.completed = true; + let allocation = record.allocation; + if let Some(allocation_record) = self.allocations.get_mut(&allocation) { + allocation_record.active_copies.remove(©); + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct AllocationRecord { + spec: ObjectSpec, + extent: u64, + active_copies: BTreeSet, + active_steps: BTreeSet, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum CopyDirection { + HostToDevice, + DeviceToHost, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct CopyRecord { + allocation: DeviceAllocation, + direction: CopyDirection, + completed: bool, +} + +fn validate_object_spec(spec: ObjectSpec, extent: u64) -> Result<(), DeviceError> { + if spec.alignment == 0 { + return Err(DeviceError::InvalidAlignment); + } + if extent > spec.max_extent || extent % spec.alignment != 0 { + return Err(DeviceError::InvalidExtent); + } + Ok(()) +} + +fn validate_equal_copy_len(left: u64, right: u64) -> Result<(), DeviceError> { + if left != right { + return Err(DeviceError::InvalidRange); + } + Ok(()) +} + +#[cfg(test)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum BackendFailure { + AllocationFailed, + CopyFailed, + InvalidView, +} + +#[cfg(test)] +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum BackendCall { + Alloc { + allocation: DeviceAllocation, + spec: ObjectSpec, + extent: u64, + }, + Free { + freed: DeviceAllocation, + }, + HostToDevice { + host: HostRange, + device: DeviceRange, + copy: CopyEvent, + mode: CopyMode, + }, + DeviceToHost { + device: DeviceRange, + host: HostRange, + copy: CopyEvent, + mode: CopyMode, + }, + WrapForTinygrad { + allocation: DeviceAllocation, + view: TensorViewSpec, + }, +} + +#[cfg(test)] +#[derive(Default)] +pub struct MockDeviceBackend { + calls: Vec, + next_failure: Option, +} + +#[cfg(test)] +impl MockDeviceBackend { + pub fn inject_failure(&mut self, failure: BackendFailure) { + self.next_failure = Some(failure); + } + + pub fn calls(&self) -> &[BackendCall] { + &self.calls + } + + fn take_failure(&mut self, failure: BackendFailure) -> bool { + if self.next_failure == Some(failure) { + self.next_failure = None; + return true; + } + false + } +} + +#[cfg(test)] +impl DeviceBridgeBackend for MockDeviceBackend { + fn alloc_device( + &mut self, + allocation: DeviceAllocation, + spec: ObjectSpec, + extent: u64, + ) -> Result<(), DeviceError> { + self.calls.push(BackendCall::Alloc { + allocation, + spec, + extent, + }); + if self.take_failure(BackendFailure::AllocationFailed) { + return Err(DeviceError::AllocationFailed); + } + Ok(()) + } + + fn free_device(&mut self, allocation: DeviceAllocation) -> Result<(), DeviceError> { + self.calls.push(BackendCall::Free { freed: allocation }); + Ok(()) + } + + fn host_to_device( + &mut self, + host: HostRange, + device: DeviceRange, + copy: CopyEvent, + mode: CopyMode, + ) -> Result<(), DeviceError> { + self.calls.push(BackendCall::HostToDevice { + host, + device, + copy, + mode, + }); + if self.take_failure(BackendFailure::CopyFailed) { + return Err(DeviceError::DeviceCopyFailed); + } + Ok(()) + } + + fn device_to_host( + &mut self, + device: DeviceRange, + host: HostRange, + copy: CopyEvent, + mode: CopyMode, + ) -> Result<(), DeviceError> { + self.calls.push(BackendCall::DeviceToHost { + device, + host, + copy, + mode, + }); + if self.take_failure(BackendFailure::CopyFailed) { + return Err(DeviceError::DeviceCopyFailed); + } + Ok(()) + } + + fn wrap_for_tinygrad( + &mut self, + allocation: DeviceAllocation, + view: TensorViewSpec, + ) -> Result { + self.calls + .push(BackendCall::WrapForTinygrad { allocation, view }); + if self.take_failure(BackendFailure::InvalidView) { + return Err(DeviceError::InvalidTensorView); + } + Ok(TensorView { + handle: allocation, + dtype: view.dtype, + shape: view.shape, + }) + } +} + +#[cfg(test)] +pub struct DeviceBridgeHarness { + bridge: DeviceBridge, +} + +#[cfg(test)] +impl DeviceBridgeHarness { + pub fn new(generation: WorkerGeneration) -> Self { + Self { + bridge: DeviceBridge::new(generation, MockDeviceBackend::default()), + } + } + + pub fn alloc_device( + &mut self, + spec: ObjectSpec, + extent: u64, + ) -> Result { + self.bridge.alloc_device(spec, extent) + } + + pub fn free_device(&mut self, allocation: DeviceAllocation) -> Result<(), DeviceError> { + self.bridge.free_device(allocation) + } + + pub fn host_to_device( + &mut self, + host: HostRange, + device: DeviceRange, + mode: CopyMode, + ) -> Result { + self.bridge.host_to_device(host, device, mode) + } + + pub fn device_to_host( + &mut self, + device: DeviceRange, + host: HostRange, + mode: CopyMode, + ) -> Result { + self.bridge.device_to_host(device, host, mode) + } + + pub fn wrap_for_tinygrad( + &mut self, + allocation: DeviceAllocation, + view: TensorViewSpec, + ) -> Result { + self.bridge.wrap_for_tinygrad(allocation, view) + } + + pub fn observe(&mut self, event: DeviceEvent) { + self.bridge.observe(event); + } + + pub fn safe_to_release_host(&self, copy: CopyEvent) -> bool { + self.bridge.safe_to_release_host(copy) + } + + pub fn host_bytes_valid(&self, copy: CopyEvent) -> bool { + self.bridge.host_bytes_valid(copy) + } + + pub fn copy_event_complete(&self, copy: CopyEvent) -> bool { + self.bridge.copy_event_complete(copy) + } + + pub fn inject_backend_failure(&mut self, failure: BackendFailure) { + self.bridge.backend_mut().inject_failure(failure); + } + + pub fn backend_calls(&self) -> &[BackendCall] { + self.bridge.backend().calls() + } +} diff --git a/crates/mvp-system/src/driver_pumps.rs b/crates/mvp-system/src/driver_pumps.rs new file mode 100644 index 0000000..69346f4 --- /dev/null +++ b/crates/mvp-system/src/driver_pumps.rs @@ -0,0 +1,544 @@ +use std::collections::{BTreeMap, BTreeSet}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct NodeId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct EdgeId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RingId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct StreamId(pub u64); + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Alpn(pub String); + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DriverConfig { + pub local_node_id: NodeId, + pub alpn: Alpn, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct EstablishSend { + pub edge_id: EdgeId, + pub peer_node_id: NodeId, + pub layout: RingLayout, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct EstablishRecv { + pub edge_id: EdgeId, + pub layout: RingLayout, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RingDirection { + Egress, + Ingress, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RingLayout { + pub ring_id: RingId, + pub byte_capacity: usize, + pub direction: RingDirection, +} + +impl RingLayout { + pub fn test_egress() -> Self { + Self { + ring_id: RingId(1), + byte_capacity: 4096, + direction: RingDirection::Egress, + } + } + + pub fn test_ingress() -> Self { + Self { + ring_id: RingId(2), + byte_capacity: 4096, + direction: RingDirection::Ingress, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum DriverEvent { + EstablishSend(EstablishSend), + EstablishRecv(EstablishRecv), + IncomingUniStream { + edge_id: EdgeId, + stream_id: StreamId, + }, + RingReadable { + edge_id: EdgeId, + }, + RingWritable { + edge_id: EdgeId, + }, + EgressBytesCommitted { + edge_id: EdgeId, + bytes: Vec, + }, + StreamBytesRead { + edge_id: EdgeId, + bytes: Vec, + }, + WriteAllAccepted { + edge_id: EdgeId, + byte_count: usize, + }, + NetworkStalled { + edge_id: EdgeId, + }, + IngressRingFull { + edge_id: EdgeId, + }, + ReadError { + edge_id: EdgeId, + }, + WriteError { + edge_id: EdgeId, + }, + StopEdge { + edge_id: EdgeId, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum DriverCommand { + OpenOrReuseConnection { + peer_node_id: NodeId, + alpn: Alpn, + local_node_id: NodeId, + }, + OpenUniStream { + edge_id: EdgeId, + peer_node_id: NodeId, + }, + SpawnSendPump { + edge_id: EdgeId, + ring_id: RingId, + }, + SpawnRecvPump { + edge_id: EdgeId, + ring_id: RingId, + stream_id: StreamId, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ActorMessage { + PollStreamFuture { edge_id: EdgeId }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum DriverEventOut { + DriverEdgeReady { + edge_id: EdgeId, + }, + StreamClosed { + edge_id: EdgeId, + }, + StreamFault { + edge_id: EdgeId, + reason: StreamFaultReason, + }, + PumpStopped { + edge_id: EdgeId, + ring_id: RingId, + }, + ObjectHeaderParsed { + edge_id: EdgeId, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum StreamFaultReason { + ReadError, + WriteError, + ProtocolError, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WakeHint { + RingReadable { edge_id: EdgeId }, + RingWritable { edge_id: EdgeId }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct StreamWrite { + pub edge_id: EdgeId, + pub bytes: Vec, +} + +pub fn encode_edge_preamble(edge_id: EdgeId) -> Vec { + edge_id.0.to_le_bytes().to_vec() +} + +pub fn count_preamble_occurrences(bytes: &[u8], edge_id: EdgeId) -> usize { + let preamble = encode_edge_preamble(edge_id); + bytes + .windows(preamble.len()) + .filter(|window| *window == preamble.as_slice()) + .count() +} + +pub fn fake_object_header_bytes() -> Vec { + b"OBJ\0fake-header".to_vec() +} + +pub struct Driver { + state: DriverState, +} + +impl Driver { + pub fn new(config: DriverConfig) -> Self { + Self { + state: DriverState::new(config), + } + } + + pub fn observe(&mut self, event: DriverEvent) { + self.state.observe(event); + } + + pub fn commands(&self) -> &[DriverCommand] { + &self.state.commands + } + + pub fn events(&self) -> &[DriverEventOut] { + &self.state.events + } + + pub fn wake_hints(&self) -> &[WakeHint] { + &self.state.wakes + } +} + +#[derive(Debug)] +struct DriverState { + config: DriverConfig, + connections: BTreeSet<(NodeId, Alpn)>, + sends: BTreeMap, + recv_specs: BTreeMap, + pending_streams: BTreeMap, + recvs: BTreeMap, + commands: Vec, + events: Vec, + wakes: Vec, + #[cfg(test)] + actor_messages: Vec, + stream_writes: Vec, + read_started: BTreeSet, +} + +impl DriverState { + fn new(config: DriverConfig) -> Self { + Self { + config, + connections: BTreeSet::new(), + sends: BTreeMap::new(), + recv_specs: BTreeMap::new(), + pending_streams: BTreeMap::new(), + recvs: BTreeMap::new(), + commands: Vec::new(), + events: Vec::new(), + wakes: Vec::new(), + #[cfg(test)] + actor_messages: Vec::new(), + stream_writes: Vec::new(), + read_started: BTreeSet::new(), + } + } + + fn observe(&mut self, event: DriverEvent) { + match event { + DriverEvent::EstablishSend(spec) => self.establish_send(spec), + DriverEvent::EstablishRecv(spec) => self.establish_recv(spec), + DriverEvent::IncomingUniStream { edge_id, stream_id } => { + self.incoming_uni_stream(edge_id, stream_id); + } + DriverEvent::RingReadable { edge_id } => self.flush_send_bytes(edge_id), + DriverEvent::RingWritable { edge_id } => self.resume_recv(edge_id), + DriverEvent::EgressBytesCommitted { edge_id, bytes } => { + if let Some(send) = self.sends.get_mut(&edge_id) { + send.pending_bytes.extend(bytes); + } + } + DriverEvent::StreamBytesRead { edge_id, bytes } => { + self.copy_recv_bytes(edge_id, &bytes) + } + DriverEvent::WriteAllAccepted { + edge_id, + byte_count, + } => { + if let Some(send) = self.sends.get_mut(&edge_id) { + send.consume_cursor += byte_count; + send.network_stalled = false; + self.wakes.push(WakeHint::RingWritable { edge_id }); + } + } + DriverEvent::NetworkStalled { edge_id } => { + if let Some(send) = self.sends.get_mut(&edge_id) { + send.network_stalled = true; + } + } + DriverEvent::IngressRingFull { edge_id } => { + if let Some(recv) = self.recvs.get_mut(&edge_id) { + recv.reading = false; + } + } + DriverEvent::ReadError { edge_id } => { + self.events.push(DriverEventOut::StreamFault { + edge_id, + reason: StreamFaultReason::ReadError, + }); + } + DriverEvent::WriteError { edge_id } => { + self.events.push(DriverEventOut::StreamFault { + edge_id, + reason: StreamFaultReason::WriteError, + }); + } + DriverEvent::StopEdge { edge_id } => self.stop_edge(edge_id), + } + } + + fn establish_send(&mut self, spec: EstablishSend) { + let connection_key = (spec.peer_node_id, self.config.alpn.clone()); + self.connections.insert(connection_key); + self.commands.push(DriverCommand::OpenOrReuseConnection { + peer_node_id: spec.peer_node_id, + alpn: self.config.alpn.clone(), + local_node_id: self.config.local_node_id, + }); + + self.commands.push(DriverCommand::SpawnSendPump { + edge_id: spec.edge_id, + ring_id: spec.layout.ring_id, + }); + + let send = SendPumpState::new(spec.layout.ring_id); + self.commands.push(DriverCommand::OpenUniStream { + edge_id: spec.edge_id, + peer_node_id: spec.peer_node_id, + }); + self.stream_writes.push(StreamWrite { + edge_id: spec.edge_id, + bytes: encode_edge_preamble(spec.edge_id), + }); + self.sends.insert(spec.edge_id, send); + self.events.push(DriverEventOut::DriverEdgeReady { + edge_id: spec.edge_id, + }); + } + + fn establish_recv(&mut self, spec: EstablishRecv) { + let edge_id = spec.edge_id; + self.recv_specs.insert(edge_id, spec); + if let Some(stream_id) = self.pending_streams.remove(&edge_id) { + self.spawn_recv(edge_id, stream_id); + } + } + + fn incoming_uni_stream(&mut self, edge_id: EdgeId, stream_id: StreamId) { + if self.recv_specs.contains_key(&edge_id) { + self.spawn_recv(edge_id, stream_id); + } else { + self.pending_streams.insert(edge_id, stream_id); + } + } + + fn spawn_recv(&mut self, edge_id: EdgeId, stream_id: StreamId) { + let Some(spec) = self.recv_specs.get(&edge_id) else { + self.pending_streams.insert(edge_id, stream_id); + return; + }; + let ring_id = spec.layout.ring_id; + self.commands.push(DriverCommand::SpawnRecvPump { + edge_id, + ring_id, + stream_id, + }); + self.read_started.insert(stream_id); + self.recvs.insert( + edge_id, + RecvPumpState { + ring_id, + commit_cursor: 0, + reading: true, + }, + ); + self.events + .push(DriverEventOut::DriverEdgeReady { edge_id }); + } + + fn flush_send_bytes(&mut self, edge_id: EdgeId) { + let Some(send) = self.sends.get_mut(&edge_id) else { + return; + }; + if send.network_stalled || send.pending_bytes.is_empty() { + return; + } + + let bytes = std::mem::take(&mut send.pending_bytes); + self.stream_writes.push(StreamWrite { edge_id, bytes }); + } + + fn resume_recv(&mut self, edge_id: EdgeId) { + if let Some(recv) = self.recvs.get_mut(&edge_id) { + recv.reading = true; + } + } + + fn copy_recv_bytes(&mut self, edge_id: EdgeId, bytes: &[u8]) { + let Some(recv) = self.recvs.get_mut(&edge_id) else { + return; + }; + if !recv.reading { + return; + } + recv.commit_cursor += bytes.len(); + self.wakes.push(WakeHint::RingReadable { edge_id }); + } + + fn stop_edge(&mut self, edge_id: EdgeId) { + let ring_id = self + .sends + .get(&edge_id) + .map(|send| send.ring_id) + .or_else(|| self.recvs.get(&edge_id).map(|recv| recv.ring_id)) + .or_else(|| { + self.recv_specs + .get(&edge_id) + .map(|spec| spec.layout.ring_id) + }) + .unwrap_or(RingId(0)); + + self.sends.remove(&edge_id); + self.recvs.remove(&edge_id); + self.recv_specs.remove(&edge_id); + self.pending_streams.remove(&edge_id); + self.events + .push(DriverEventOut::PumpStopped { edge_id, ring_id }); + } +} + +#[derive(Debug)] +struct SendPumpState { + ring_id: RingId, + pending_bytes: Vec, + consume_cursor: usize, + network_stalled: bool, +} + +impl SendPumpState { + fn new(ring_id: RingId) -> Self { + Self { + ring_id, + pending_bytes: Vec::new(), + consume_cursor: 0, + network_stalled: false, + } + } +} + +#[derive(Debug)] +struct RecvPumpState { + ring_id: RingId, + commit_cursor: usize, + reading: bool, +} + +#[cfg(test)] +pub struct CommandLog { + commands: Vec, +} + +#[cfg(test)] +impl CommandLog { + pub fn iter(&self) -> std::vec::IntoIter { + self.commands.clone().into_iter() + } +} + +#[cfg(test)] +pub struct DriverHarness { + driver: Driver, +} + +#[cfg(test)] +impl DriverHarness { + pub fn new(config: DriverConfig) -> Self { + Self { + driver: Driver::new(config), + } + } + + pub fn observe(&mut self, event: DriverEvent) { + self.driver.observe(event); + } + + pub fn commands(&self) -> CommandLog { + CommandLog { + commands: self.driver.state.commands.clone(), + } + } + + pub fn events(&self) -> &[DriverEventOut] { + &self.driver.state.events + } + + pub fn wake_hints(&self) -> &[WakeHint] { + &self.driver.state.wakes + } + + pub fn actor_messages(&self) -> &[ActorMessage] { + &self.driver.state.actor_messages + } + + pub fn stream_writes(&self, edge_id: EdgeId) -> Vec { + self.driver + .state + .stream_writes + .iter() + .filter(|write| write.edge_id == edge_id) + .cloned() + .collect() + } + + pub fn stream_reads_started(&self, stream_id: StreamId) -> bool { + self.driver.state.read_started.contains(&stream_id) + } + + pub fn is_reading_stream(&self, edge_id: EdgeId) -> bool { + self.driver + .state + .recvs + .get(&edge_id) + .map(|recv| recv.reading) + .unwrap_or(false) + } + + pub fn ring_commit(&self, edge_id: EdgeId) -> usize { + self.driver + .state + .recvs + .get(&edge_id) + .map(|recv| recv.commit_cursor) + .unwrap_or(0) + } + + pub fn ring_consume(&self, edge_id: EdgeId) -> usize { + self.driver + .state + .sends + .get(&edge_id) + .map(|send| send.consume_cursor) + .unwrap_or(0) + } +} diff --git a/crates/mvp-system/src/edge_establisher.rs b/crates/mvp-system/src/edge_establisher.rs new file mode 100644 index 0000000..39090e2 --- /dev/null +++ b/crates/mvp-system/src/edge_establisher.rs @@ -0,0 +1,751 @@ +use std::collections::BTreeMap; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RunId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct NodeId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct EdgeId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct LeaseRequestId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RingId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ActorAddress(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RingDirection { + Egress, + Ingress, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ObjectKind { + Activation, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum DType { + F16, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ObjectSpec { + pub kind: ObjectKind, + pub dtype: DType, + pub max_extent_bytes: u64, +} + +impl ObjectSpec { + pub const fn test_activation() -> Self { + Self { + kind: ObjectKind::Activation, + dtype: DType::F16, + max_extent_bytes: 4096, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct RingSpec { + pub header_bytes: u64, + pub data_bytes: u64, + pub alignment: u64, +} + +impl RingSpec { + pub const fn test_activation() -> Self { + Self { + header_bytes: 128, + data_bytes: 4096, + alignment: 64, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct RingLayout { + pub start_offset: u64, + pub header_offset: u64, + pub data_offset: u64, + pub end_offset: u64, + pub data_bytes: u64, + pub alignment: u64, +} + +impl RingLayout { + pub const fn test_layout(start_offset: u64) -> Self { + Self { + start_offset, + header_offset: start_offset, + data_offset: start_offset + 128, + end_offset: start_offset + 128 + 4096, + data_bytes: 4096, + alignment: 64, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct QuiescenceProof { + verified: bool, +} + +impl QuiescenceProof { + pub const fn verified() -> Self { + Self { verified: true } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ProvisionTx { + pub run_id: RunId, + pub edge_id: EdgeId, + pub local_node_id: NodeId, + pub consumer_node_id: NodeId, + pub object_spec: ObjectSpec, + pub ring_spec: RingSpec, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ProvisionRx { + pub run_id: RunId, + pub edge_id: EdgeId, + pub local_node_id: NodeId, + pub object_spec: ObjectSpec, + pub ring_spec: RingSpec, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RingLeaseRejection { + CannotFit, + ArenaShuttingDown, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum StreamFaultReason { + ReadError, + WriteError, + ProtocolError, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RingFaultReason { + WorkerRejectedRing, + WorkerCrashed, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum EdgeEvent { + ProvisionTx(ProvisionTx), + ProvisionRx(ProvisionRx), + RingLeased { + request_id: LeaseRequestId, + ring_id: RingId, + layout: RingLayout, + }, + RingLeaseRejected { + request_id: LeaseRequestId, + reason: RingLeaseRejection, + }, + RingInstalled { + edge_id: EdgeId, + ring_id: RingId, + }, + DriverEdgeReady { + edge_id: EdgeId, + }, + StreamFault { + edge_id: EdgeId, + reason: StreamFaultReason, + }, + RingFault { + edge_id: EdgeId, + ring_id: RingId, + reason: RingFaultReason, + }, + StopEdge { + edge_id: EdgeId, + }, + PumpStopped { + edge_id: EdgeId, + ring_id: RingId, + }, + RingQuiesced { + ring_id: RingId, + }, + QuiescenceProven { + ring_id: RingId, + }, + Stopped { + edge_id: EdgeId, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum EdgeCommand { + LeaseRing { + request_id: LeaseRequestId, + edge_id: EdgeId, + direction: RingDirection, + ring_spec: RingSpec, + }, + CancelQueuedLease { + request_id: LeaseRequestId, + edge_id: EdgeId, + }, + InstallWorkerRing { + edge_id: EdgeId, + ring_id: RingId, + direction: RingDirection, + layout: RingLayout, + object_spec: ObjectSpec, + ring_spec: RingSpec, + }, + UninstallWorkerRing { + edge_id: EdgeId, + ring_id: RingId, + }, + EstablishSend { + edge_id: EdgeId, + consumer_node_id: NodeId, + layout: RingLayout, + }, + EstablishRecv { + edge_id: EdgeId, + layout: RingLayout, + }, + StopPump { + edge_id: EdgeId, + ring_id: RingId, + }, + ReleaseArenaLease { + ring_id: RingId, + proof: QuiescenceProof, + }, + CopyHotPathBytes { + edge_id: EdgeId, + bytes: Vec, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum EdgeFaultReason { + RingLeaseRejected(RingLeaseRejection), + RingFault(RingFaultReason), + StreamFault(StreamFaultReason), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum EdgeLifecycleEvent { + EdgeReady { + edge_id: EdgeId, + local_edge_actor: ActorAddress, + }, + EdgeFaulted { + edge_id: EdgeId, + reason: EdgeFaultReason, + }, + EdgeStopped { + edge_id: EdgeId, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct LocalEdgeRecord { + pub edge_id: EdgeId, + pub direction: RingDirection, + pub state: EdgeProvisionState, + pub lease_request_id: Option, + pub ring_id: Option, + pub peer_node_id: Option, + pub local_edge_actor: ActorAddress, + pub remote_actor_address: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum EdgeProvisionState { + WaitingForLease, + WaitingForWorkerRing, + WaitingForDriver, + Ready, + Stopping, + Stopped, + Failed, +} + +pub struct EdgeEstablisher { + state: EdgeEstablisherState, +} + +impl EdgeEstablisher { + pub fn new(local_node_id: NodeId) -> Self { + Self { + state: EdgeEstablisherState::new(local_node_id), + } + } + + pub fn observe(&mut self, event: EdgeEvent) { + self.state.observe(event); + } + + pub fn commands(&self) -> &[EdgeCommand] { + &self.state.commands + } + + pub fn events(&self) -> &[EdgeLifecycleEvent] { + &self.state.events + } + + pub fn local_record(&self, edge_id: EdgeId) -> Option { + self.state.records.get(&edge_id).map(EdgeRecord::snapshot) + } +} + +struct EdgeEstablisherState { + local_node_id: NodeId, + next_request_id: u64, + next_actor_id: u64, + records: BTreeMap, + commands: Vec, + events: Vec, +} + +impl EdgeEstablisherState { + fn new(local_node_id: NodeId) -> Self { + Self { + local_node_id, + next_request_id: 1, + next_actor_id: 1, + records: BTreeMap::new(), + commands: Vec::new(), + events: Vec::new(), + } + } + + fn observe(&mut self, event: EdgeEvent) { + match event { + EdgeEvent::ProvisionTx(provision) => self.provision_tx(provision), + EdgeEvent::ProvisionRx(provision) => self.provision_rx(provision), + EdgeEvent::RingLeased { + request_id, + ring_id, + layout, + } => self.ring_leased(request_id, ring_id, layout), + EdgeEvent::RingLeaseRejected { request_id, reason } => { + self.ring_lease_rejected(request_id, reason); + } + EdgeEvent::RingInstalled { edge_id, ring_id } => self.ring_installed(edge_id, ring_id), + EdgeEvent::DriverEdgeReady { edge_id } => self.driver_edge_ready(edge_id), + EdgeEvent::StreamFault { edge_id, reason } => { + self.fault_edge(edge_id, EdgeFaultReason::StreamFault(reason)); + } + EdgeEvent::RingFault { + edge_id, + ring_id, + reason, + } => self.ring_fault(edge_id, ring_id, reason), + EdgeEvent::StopEdge { edge_id } => self.stop_edge(edge_id), + EdgeEvent::PumpStopped { edge_id, ring_id } => self.pump_stopped(edge_id, ring_id), + EdgeEvent::RingQuiesced { ring_id } | EdgeEvent::QuiescenceProven { ring_id } => { + self.release_after_quiescence(ring_id); + } + EdgeEvent::Stopped { edge_id } => self.mark_stopped(edge_id), + } + } + + fn provision_tx(&mut self, provision: ProvisionTx) { + if provision.local_node_id != self.local_node_id { + return; + } + + let request_id = self.next_request_id(); + let actor = self.next_actor_address(); + let edge_id = provision.edge_id; + let ring_spec = provision.ring_spec; + self.records + .insert(edge_id, EdgeRecord::new_tx(provision, request_id, actor)); + self.commands.push(EdgeCommand::LeaseRing { + request_id, + edge_id, + direction: RingDirection::Egress, + ring_spec, + }); + } + + fn provision_rx(&mut self, provision: ProvisionRx) { + if provision.local_node_id != self.local_node_id { + return; + } + + let request_id = self.next_request_id(); + let actor = self.next_actor_address(); + let edge_id = provision.edge_id; + let ring_spec = provision.ring_spec; + self.records + .insert(edge_id, EdgeRecord::new_rx(provision, request_id, actor)); + self.commands.push(EdgeCommand::LeaseRing { + request_id, + edge_id, + direction: RingDirection::Ingress, + ring_spec, + }); + } + + fn ring_leased(&mut self, request_id: LeaseRequestId, ring_id: RingId, layout: RingLayout) { + let Some(edge_id) = self.edge_for_request(request_id) else { + return; + }; + + let Some(record) = self.records.get_mut(&edge_id) else { + return; + }; + + if record.state != EdgeProvisionState::WaitingForLease { + self.commands.push(EdgeCommand::ReleaseArenaLease { + ring_id, + proof: QuiescenceProof::verified(), + }); + return; + } + + record.ring_id = Some(ring_id); + record.layout = Some(layout); + record.state = EdgeProvisionState::WaitingForWorkerRing; + self.commands.push(EdgeCommand::InstallWorkerRing { + edge_id, + ring_id, + direction: record.direction, + layout, + object_spec: record.object_spec, + ring_spec: record.ring_spec, + }); + } + + fn ring_lease_rejected(&mut self, request_id: LeaseRequestId, reason: RingLeaseRejection) { + let Some(edge_id) = self.edge_for_waiting_request(request_id) else { + return; + }; + let Some(record) = self.records.get_mut(&edge_id) else { + return; + }; + + record.state = EdgeProvisionState::Failed; + self.events.push(EdgeLifecycleEvent::EdgeFaulted { + edge_id, + reason: EdgeFaultReason::RingLeaseRejected(reason), + }); + } + + fn ring_installed(&mut self, edge_id: EdgeId, ring_id: RingId) { + let Some(record) = self.records.get_mut(&edge_id) else { + return; + }; + if record.state != EdgeProvisionState::WaitingForWorkerRing + || record.ring_id != Some(ring_id) + { + return; + } + let Some(layout) = record.layout else { + return; + }; + + record.worker_installed = true; + record.driver_established = true; + record.state = EdgeProvisionState::WaitingForDriver; + match record.direction { + RingDirection::Egress => { + if let Some(consumer_node_id) = record.peer_node_id { + self.commands.push(EdgeCommand::EstablishSend { + edge_id, + consumer_node_id, + layout, + }); + } + } + RingDirection::Ingress => { + self.commands + .push(EdgeCommand::EstablishRecv { edge_id, layout }); + } + } + } + + fn driver_edge_ready(&mut self, edge_id: EdgeId) { + let Some(record) = self.records.get_mut(&edge_id) else { + return; + }; + if record.state != EdgeProvisionState::WaitingForDriver { + return; + } + + record.state = EdgeProvisionState::Ready; + self.events.push(EdgeLifecycleEvent::EdgeReady { + edge_id, + local_edge_actor: record.local_edge_actor, + }); + } + + fn ring_fault(&mut self, edge_id: EdgeId, ring_id: RingId, reason: RingFaultReason) { + let Some(record) = self.records.get(&edge_id) else { + return; + }; + if record.ring_id != Some(ring_id) { + return; + } + self.fault_edge(edge_id, EdgeFaultReason::RingFault(reason)); + } + + fn fault_edge(&mut self, edge_id: EdgeId, reason: EdgeFaultReason) { + let Some(record) = self.records.get(&edge_id) else { + return; + }; + if matches!( + record.state, + EdgeProvisionState::Stopping | EdgeProvisionState::Stopped | EdgeProvisionState::Failed + ) { + return; + } + + self.events + .push(EdgeLifecycleEvent::EdgeFaulted { edge_id, reason }); + self.start_stopping(edge_id, true); + } + + fn stop_edge(&mut self, edge_id: EdgeId) { + self.start_stopping(edge_id, true); + } + + fn start_stopping(&mut self, edge_id: EdgeId, cancel_lease: bool) { + let Some(record) = self.records.get_mut(&edge_id) else { + return; + }; + if matches!( + record.state, + EdgeProvisionState::Stopping | EdgeProvisionState::Stopped + ) { + return; + } + + let request_id = record.lease_request_id; + let ring_id = record.ring_id; + let driver_established = record.driver_established; + let worker_installed = record.worker_installed; + record.state = EdgeProvisionState::Stopping; + + if cancel_lease { + if let Some(request_id) = request_id { + self.commands.push(EdgeCommand::CancelQueuedLease { + request_id, + edge_id, + }); + } + } + + let Some(ring_id) = ring_id else { + self.mark_stopped_with_event(edge_id); + return; + }; + + if driver_established { + self.commands + .push(EdgeCommand::StopPump { edge_id, ring_id }); + } + if worker_installed { + self.commands + .push(EdgeCommand::UninstallWorkerRing { edge_id, ring_id }); + } + if !driver_established && !worker_installed { + self.release_ring(edge_id, ring_id); + } + } + + fn pump_stopped(&mut self, edge_id: EdgeId, ring_id: RingId) { + let Some(record) = self.records.get_mut(&edge_id) else { + return; + }; + if record.ring_id == Some(ring_id) { + record.pump_stopped = true; + } + } + + fn release_after_quiescence(&mut self, ring_id: RingId) { + let Some(edge_id) = self.edge_for_ring(ring_id) else { + return; + }; + let Some(record) = self.records.get(&edge_id) else { + return; + }; + if record.state != EdgeProvisionState::Stopping { + return; + } + self.release_ring(edge_id, ring_id); + } + + fn release_ring(&mut self, edge_id: EdgeId, ring_id: RingId) { + self.commands.push(EdgeCommand::ReleaseArenaLease { + ring_id, + proof: QuiescenceProof::verified(), + }); + self.mark_stopped_with_event(edge_id); + } + + fn mark_stopped_with_event(&mut self, edge_id: EdgeId) { + let Some(record) = self.records.get_mut(&edge_id) else { + return; + }; + if record.state == EdgeProvisionState::Stopped { + return; + } + record.state = EdgeProvisionState::Stopped; + self.events + .push(EdgeLifecycleEvent::EdgeStopped { edge_id }); + } + + fn mark_stopped(&mut self, edge_id: EdgeId) { + let Some(record) = self.records.get_mut(&edge_id) else { + return; + }; + record.state = EdgeProvisionState::Stopped; + } + + fn edge_for_request(&self, request_id: LeaseRequestId) -> Option { + self.records.iter().find_map(|(edge_id, record)| { + (record.lease_request_id == Some(request_id)).then_some(*edge_id) + }) + } + + fn edge_for_waiting_request(&self, request_id: LeaseRequestId) -> Option { + self.records.iter().find_map(|(edge_id, record)| { + (record.state == EdgeProvisionState::WaitingForLease + && record.lease_request_id == Some(request_id)) + .then_some(*edge_id) + }) + } + + fn edge_for_ring(&self, ring_id: RingId) -> Option { + self.records + .iter() + .find_map(|(edge_id, record)| (record.ring_id == Some(ring_id)).then_some(*edge_id)) + } + + fn next_request_id(&mut self) -> LeaseRequestId { + let request_id = LeaseRequestId(self.next_request_id); + self.next_request_id += 1; + request_id + } + + fn next_actor_address(&mut self) -> ActorAddress { + let actor = ActorAddress(self.next_actor_id); + self.next_actor_id += 1; + actor + } +} + +struct EdgeRecord { + run_id: RunId, + edge_id: EdgeId, + direction: RingDirection, + state: EdgeProvisionState, + lease_request_id: Option, + ring_id: Option, + layout: Option, + peer_node_id: Option, + local_edge_actor: ActorAddress, + remote_actor_address: Option, + object_spec: ObjectSpec, + ring_spec: RingSpec, + worker_installed: bool, + driver_established: bool, + pump_stopped: bool, +} + +impl EdgeRecord { + fn new_tx(provision: ProvisionTx, request_id: LeaseRequestId, actor: ActorAddress) -> Self { + Self { + run_id: provision.run_id, + edge_id: provision.edge_id, + direction: RingDirection::Egress, + state: EdgeProvisionState::WaitingForLease, + lease_request_id: Some(request_id), + ring_id: None, + layout: None, + peer_node_id: Some(provision.consumer_node_id), + local_edge_actor: actor, + remote_actor_address: None, + object_spec: provision.object_spec, + ring_spec: provision.ring_spec, + worker_installed: false, + driver_established: false, + pump_stopped: false, + } + } + + fn new_rx(provision: ProvisionRx, request_id: LeaseRequestId, actor: ActorAddress) -> Self { + Self { + run_id: provision.run_id, + edge_id: provision.edge_id, + direction: RingDirection::Ingress, + state: EdgeProvisionState::WaitingForLease, + lease_request_id: Some(request_id), + ring_id: None, + layout: None, + peer_node_id: None, + local_edge_actor: actor, + remote_actor_address: None, + object_spec: provision.object_spec, + ring_spec: provision.ring_spec, + worker_installed: false, + driver_established: false, + pump_stopped: false, + } + } + + fn snapshot(&self) -> LocalEdgeRecord { + let _ = self.run_id; + LocalEdgeRecord { + edge_id: self.edge_id, + direction: self.direction, + state: self.state, + lease_request_id: self.lease_request_id, + ring_id: self.ring_id, + peer_node_id: self.peer_node_id, + local_edge_actor: self.local_edge_actor, + remote_actor_address: self.remote_actor_address, + } + } +} + +#[cfg(test)] +pub struct EdgeEstablisherHarness { + establisher: EdgeEstablisher, +} + +#[cfg(test)] +impl EdgeEstablisherHarness { + pub fn new(local_node_id: NodeId) -> Self { + Self { + establisher: EdgeEstablisher::new(local_node_id), + } + } + + pub fn observe(&mut self, event: EdgeEvent) { + self.establisher.observe(event); + } + + pub fn commands(&self) -> &[EdgeCommand] { + self.establisher.commands() + } + + pub fn events(&self) -> &[EdgeLifecycleEvent] { + self.establisher.events() + } + + pub fn local_record(&self, edge_id: EdgeId) -> Option { + self.establisher.local_record(edge_id) + } +} diff --git a/crates/mvp-system/src/gpu_worker_ctl.rs b/crates/mvp-system/src/gpu_worker_ctl.rs new file mode 100644 index 0000000..a495cbc --- /dev/null +++ b/crates/mvp-system/src/gpu_worker_ctl.rs @@ -0,0 +1,378 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct NodeId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ProcessId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct WorkerGeneration(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RingId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct StepId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ObjectId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct DeviceHandle { + pub generation: WorkerGeneration, + pub id: u64, +} + +impl DeviceHandle { + pub fn new(generation: WorkerGeneration, id: u64) -> Self { + Self { generation, id } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ArenaEnv { + pub arena_fd: i32, + pub arena_bytes: u64, +} + +impl ArenaEnv { + pub fn test_default() -> Self { + Self { + arena_fd: 3, + arena_bytes: 4096, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct WorkerConfig { + pub node_id: NodeId, + pub arena_env: ArenaEnv, + pub initialization_timeout_ms: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ActorCommand { + InstallRing { + generation: WorkerGeneration, + ring_id: RingId, + }, + ExecuteStep { + generation: WorkerGeneration, + step_id: StepId, + input: DeviceHandle, + }, + ReleaseDeviceObject { + generation: WorkerGeneration, + handle: DeviceHandle, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WorkerCommand { + InitializeWorker { + generation: WorkerGeneration, + }, + InstallRing { + ring_id: RingId, + }, + ExecuteStep { + step_id: StepId, + input: DeviceHandle, + }, + ReleaseDeviceObject { + handle: DeviceHandle, + }, + ShutdownWorker, + PayloadBytes { + bytes: Vec, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WorkerEvent { + RingInstalled { ring_id: RingId }, + ObjectLoaded { object_id: ObjectId, sequence: u64 }, + ObjectProduced { object_id: ObjectId, sequence: u64 }, + StepCompleted { step_id: StepId }, + RingReadable { ring_id: RingId }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ExitStatus { + Code(i32), + Signal(i32), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WorkerCtlEvent { + StartWorker, + ProcessStarted { pid: ProcessId }, + WorkerReady { generation: WorkerGeneration }, + ActorCommand(ActorCommand), + StdoutEvent(WorkerEvent), + ProcessExited { status: ExitStatus }, + RestartRequested, + ShutdownRequested, + WorkerStopped { generation: WorkerGeneration }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum WorkerFailure { + InitializationTimeout, + ProcessExited, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CommandRejection { + NotRunning, + OldGenerationHandle, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WorkerCtlOut { + WorkerRunning { + generation: WorkerGeneration, + }, + WorkerFailed { + generation: WorkerGeneration, + reason: WorkerFailure, + }, + CommandRejected { + reason: CommandRejection, + }, + RingFaulted { + ring_id: RingId, + }, + WorkerStopped { + generation: WorkerGeneration, + }, + RingQuiesced { + ring_id: RingId, + }, + TerminalStopped { + generation: WorkerGeneration, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WorkerCtlCommand { + SpawnProcessActor { node_id: NodeId }, + StopDriverPump { ring_id: RingId }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum RoutedEvent { + ToEdgeEstablisher(WorkerEvent), + ToRxOrRole(WorkerEvent), + ToTxOrRole(WorkerEvent), + ToStageController(WorkerEvent), + ToDriverOrWorkerSide(WorkerEvent), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum CtlState { + Idle, + Starting, + Running, + Crashed, + ShuttingDown, + Stopped, +} + +#[cfg(test)] +pub struct GpuWorkerCtlHarness { + config: WorkerConfig, + state: CtlState, + current_generation: WorkerGeneration, + now_ms: u64, + start_time_ms: Option, + commands: Vec, + serialized: Vec, + events: Vec, + routed: Vec, + installed_rings: std::collections::BTreeSet, +} + +#[cfg(test)] +impl GpuWorkerCtlHarness { + pub fn new(config: WorkerConfig) -> Self { + Self { + config, + state: CtlState::Idle, + current_generation: WorkerGeneration(1), + now_ms: 0, + start_time_ms: None, + commands: Vec::new(), + serialized: Vec::new(), + events: Vec::new(), + routed: Vec::new(), + installed_rings: std::collections::BTreeSet::new(), + } + } + + pub fn observe(&mut self, event: WorkerCtlEvent) { + match event { + WorkerCtlEvent::StartWorker => self.start(), + WorkerCtlEvent::ProcessStarted { .. } => { + self.state = CtlState::Starting; + self.serialized.push(WorkerCommand::InitializeWorker { + generation: self.current_generation, + }); + } + WorkerCtlEvent::WorkerReady { generation } => { + self.current_generation = generation; + self.state = CtlState::Running; + self.events.push(WorkerCtlOut::WorkerRunning { generation }); + } + WorkerCtlEvent::ActorCommand(command) => self.actor_command(command), + WorkerCtlEvent::StdoutEvent(event) => self.route(event), + WorkerCtlEvent::ProcessExited { status } => self.process_exited(status), + WorkerCtlEvent::RestartRequested => self.restart(), + WorkerCtlEvent::ShutdownRequested => self.shutdown(), + WorkerCtlEvent::WorkerStopped { generation } => { + self.events.push(WorkerCtlOut::WorkerStopped { generation }); + } + } + } + + pub fn advance_time_ms(&mut self, delta: u64) { + self.now_ms = self.now_ms.saturating_add(delta); + if self.state == CtlState::Starting { + if let Some(start_time) = self.start_time_ms { + if self.now_ms.saturating_sub(start_time) > self.config.initialization_timeout_ms { + self.events.push(WorkerCtlOut::WorkerFailed { + generation: self.current_generation, + reason: WorkerFailure::InitializationTimeout, + }); + self.state = CtlState::Crashed; + } + } + } + } + + pub fn commands(&self) -> &[WorkerCtlCommand] { + &self.commands + } + + pub fn serialized_worker_commands(&self) -> &[WorkerCommand] { + &self.serialized + } + + pub fn events(&self) -> &[WorkerCtlOut] { + &self.events + } + + pub fn routed(&self) -> &[RoutedEvent] { + &self.routed + } + + pub fn current_generation(&self) -> WorkerGeneration { + self.current_generation + } + + fn start(&mut self) { + self.state = CtlState::Starting; + self.start_time_ms = Some(self.now_ms); + self.commands.push(WorkerCtlCommand::SpawnProcessActor { + node_id: self.config.node_id, + }); + } + + fn actor_command(&mut self, command: ActorCommand) { + if self.state != CtlState::Running { + self.events.push(WorkerCtlOut::CommandRejected { + reason: CommandRejection::NotRunning, + }); + return; + } + let command_generation = match command { + ActorCommand::InstallRing { generation, .. } + | ActorCommand::ExecuteStep { generation, .. } + | ActorCommand::ReleaseDeviceObject { generation, .. } => generation, + }; + if command_generation != self.current_generation { + self.events.push(WorkerCtlOut::CommandRejected { + reason: CommandRejection::OldGenerationHandle, + }); + return; + } + match command { + ActorCommand::InstallRing { ring_id, .. } => { + self.installed_rings.insert(ring_id); + self.serialized.push(WorkerCommand::InstallRing { ring_id }); + } + ActorCommand::ExecuteStep { step_id, input, .. } => { + if input.generation != self.current_generation { + self.events.push(WorkerCtlOut::CommandRejected { + reason: CommandRejection::OldGenerationHandle, + }); + } else { + self.serialized + .push(WorkerCommand::ExecuteStep { step_id, input }); + } + } + ActorCommand::ReleaseDeviceObject { handle, .. } => { + if handle.generation != self.current_generation { + self.events.push(WorkerCtlOut::CommandRejected { + reason: CommandRejection::OldGenerationHandle, + }); + } else { + self.serialized + .push(WorkerCommand::ReleaseDeviceObject { handle }); + } + } + } + } + + fn route(&mut self, event: WorkerEvent) { + match event.clone() { + WorkerEvent::RingInstalled { .. } => { + self.routed.push(RoutedEvent::ToEdgeEstablisher(event)) + } + WorkerEvent::ObjectLoaded { .. } => self.routed.push(RoutedEvent::ToRxOrRole(event)), + WorkerEvent::ObjectProduced { .. } => self.routed.push(RoutedEvent::ToTxOrRole(event)), + WorkerEvent::StepCompleted { .. } => { + self.routed.push(RoutedEvent::ToStageController(event)) + } + WorkerEvent::RingReadable { .. } => { + self.routed.push(RoutedEvent::ToDriverOrWorkerSide(event)) + } + } + } + + fn process_exited(&mut self, status: ExitStatus) { + match (self.state, status) { + (CtlState::ShuttingDown, ExitStatus::Code(0)) => { + for ring_id in &self.installed_rings { + self.events + .push(WorkerCtlOut::RingQuiesced { ring_id: *ring_id }); + } + self.events.push(WorkerCtlOut::TerminalStopped { + generation: self.current_generation, + }); + self.state = CtlState::Stopped; + } + (_, _) => { + for ring_id in &self.installed_rings { + self.events + .push(WorkerCtlOut::RingFaulted { ring_id: *ring_id }); + self.commands + .push(WorkerCtlCommand::StopDriverPump { ring_id: *ring_id }); + } + self.events.push(WorkerCtlOut::WorkerFailed { + generation: self.current_generation, + reason: WorkerFailure::ProcessExited, + }); + self.state = CtlState::Crashed; + } + } + } + + fn restart(&mut self) { + self.state = CtlState::Starting; + self.current_generation = WorkerGeneration(self.current_generation.0 + 1); + self.start_time_ms = Some(self.now_ms); + } + + fn shutdown(&mut self) { + if self.state == CtlState::Running { + self.serialized.push(WorkerCommand::ShutdownWorker); + self.state = CtlState::ShuttingDown; + } + } +} diff --git a/crates/mvp-system/src/gpu_worker_egress_producer.rs b/crates/mvp-system/src/gpu_worker_egress_producer.rs new file mode 100644 index 0000000..7b58a6d --- /dev/null +++ b/crates/mvp-system/src/gpu_worker_egress_producer.rs @@ -0,0 +1,391 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct WorkerGeneration(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RingId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct EdgeId(pub u64); +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct PortId(pub String); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ObjectId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct StepId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct DeviceHandle { + pub generation: WorkerGeneration, + pub id: u64, +} + +impl DeviceHandle { + pub fn new(generation: WorkerGeneration, id: u64) -> Self { + Self { generation, id } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RingDirection { + Ingress, + Egress, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ObjectLayout { + Token, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ObjectSpec { + pub max_extent: u64, + pub alignment: u64, + pub layout: ObjectLayout, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct InstallRing { + pub ring_id: RingId, + pub edge_id: EdgeId, + pub port_id: PortId, + pub direction: RingDirection, + pub object_spec: ObjectSpec, + pub generation: WorkerGeneration, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +pub struct ObjectFlags { + pub end_of_sequence: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct OutputBinding { + pub ring_id: RingId, + pub object_id: ObjectId, + pub sequence: u64, + pub extent: u64, + pub flags: ObjectFlags, + pub device_source: DeviceHandle, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WorkerEgressEvent { + InstallRing(InstallRing), + ExecuteStep { + step_id: StepId, + outputs: Vec, + }, + HeaderReady { + object_id: ObjectId, + }, + EgressRingFull { + ring_id: RingId, + }, + RingWritable { + ring_id: RingId, + }, + DeviceToHostCopyCompleted { + object_id: ObjectId, + byte_count: u64, + }, + DeviceCopyFailed { + object_id: ObjectId, + }, + RoleStateUpdated { + step_id: StepId, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum StepFailureReason { + InvalidOutputRing, + OutputExtentViolation, + DeviceCopyFailed, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WorkerEgressOut { + ObjectProduced { + ring_id: RingId, + object_id: ObjectId, + sequence: u64, + }, + StepCompleted { + step_id: StepId, + }, + StepFailed { + step_id: StepId, + reason: StepFailureReason, + }, + RingFault { + ring_id: RingId, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WakeHint { + RingReadable { ring_id: RingId }, + RingWritable { ring_id: RingId }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ObjectHeader { + pub object_id: ObjectId, + pub sequence: u64, + pub extent: u64, +} + +impl ObjectHeader { + pub fn decode(bytes: &[u8]) -> Result { + if bytes.len() < HEADER_LEN || &bytes[0..4] != b"MO01" || bytes[4] != 1 { + return Err(HeaderDecodeError); + } + Ok(Self { + object_id: ObjectId(u64::from_le_bytes(bytes[8..16].try_into().unwrap())), + sequence: u64::from_le_bytes(bytes[16..24].try_into().unwrap()), + extent: u64::from_le_bytes(bytes[24..32].try_into().unwrap()), + }) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct HeaderDecodeError; + +const HEADER_LEN: usize = 48; + +#[derive(Clone, Debug, PartialEq, Eq)] +struct PendingStep { + step_id: StepId, + outputs: Vec, + role_state_updated: bool, +} + +#[cfg(test)] +pub struct EgressProducerHarness { + generation: WorkerGeneration, + rings: std::collections::BTreeMap, + pending_outputs: Vec, + steps: Vec, + committed: std::collections::BTreeMap>, + payload_committed: std::collections::BTreeMap, + full_rings: std::collections::BTreeSet, + produced: std::collections::BTreeSet, + wake_hints: Vec, + events: Vec, +} + +#[cfg(test)] +impl EgressProducerHarness { + pub fn new(generation: WorkerGeneration) -> Self { + Self { + generation, + rings: std::collections::BTreeMap::new(), + pending_outputs: Vec::new(), + steps: Vec::new(), + committed: std::collections::BTreeMap::new(), + payload_committed: std::collections::BTreeMap::new(), + full_rings: std::collections::BTreeSet::new(), + produced: std::collections::BTreeSet::new(), + wake_hints: Vec::new(), + events: Vec::new(), + } + } + + pub fn observe(&mut self, event: WorkerEgressEvent) { + match event { + WorkerEgressEvent::InstallRing(install) => { + if install.direction == RingDirection::Egress + && install.generation == self.generation + { + self.rings.insert(install.ring_id, install); + } + } + WorkerEgressEvent::ExecuteStep { step_id, outputs } => { + self.execute_step(step_id, outputs) + } + WorkerEgressEvent::HeaderReady { object_id } => self.header_ready(object_id), + WorkerEgressEvent::EgressRingFull { ring_id } => { + self.full_rings.insert(ring_id); + } + WorkerEgressEvent::RingWritable { ring_id } => { + self.full_rings.remove(&ring_id); + self.wake_hints.push(WakeHint::RingWritable { ring_id }); + } + WorkerEgressEvent::DeviceToHostCopyCompleted { + object_id, + byte_count, + } => self.copy_completed(object_id, byte_count), + WorkerEgressEvent::DeviceCopyFailed { object_id } => self.copy_failed(object_id), + WorkerEgressEvent::RoleStateUpdated { step_id } => { + if let Some(step) = self.steps.iter_mut().find(|step| step.step_id == step_id) { + step.role_state_updated = true; + } + self.maybe_step_completed(step_id); + } + } + } + + pub fn pending_outputs(&self) -> &[OutputBinding] { + &self.pending_outputs + } + + pub fn committed_bytes(&self, ring_id: RingId) -> &[u8] { + self.committed + .get(&ring_id) + .map(Vec::as_slice) + .unwrap_or(&[]) + } + + pub fn committed_payload_bytes(&self, ring_id: RingId) -> u64 { + self.payload_committed.get(&ring_id).copied().unwrap_or(0) + } + + pub fn wake_hints(&self) -> &[WakeHint] { + &self.wake_hints + } + + pub fn events(&self) -> &[WorkerEgressOut] { + &self.events + } + + pub fn complete_output(&mut self, object_id: ObjectId) { + self.header_ready(object_id); + let Some(output) = self + .pending_outputs + .iter() + .find(|output| output.object_id == object_id) + .copied() + else { + return; + }; + self.copy_completed(object_id, output.extent); + } + + fn execute_step(&mut self, step_id: StepId, outputs: Vec) { + for output in &outputs { + let Some(ring) = self.rings.get(&output.ring_id) else { + self.events.push(WorkerEgressOut::StepFailed { + step_id, + reason: StepFailureReason::InvalidOutputRing, + }); + return; + }; + if output.extent > ring.object_spec.max_extent + || (ring.object_spec.alignment != 0 + && output.extent % ring.object_spec.alignment != 0) + { + self.events.push(WorkerEgressOut::StepFailed { + step_id, + reason: StepFailureReason::OutputExtentViolation, + }); + return; + } + } + self.pending_outputs.extend(outputs.iter().copied()); + self.steps.push(PendingStep { + step_id, + outputs, + role_state_updated: false, + }); + } + + fn header_ready(&mut self, object_id: ObjectId) { + let Some(output) = self + .pending_outputs + .iter() + .find(|output| output.object_id == object_id) + .copied() + else { + return; + }; + let bytes = encode_header(output); + self.committed + .entry(output.ring_id) + .or_default() + .extend(bytes); + self.wake_hints.push(WakeHint::RingReadable { + ring_id: output.ring_id, + }); + } + + fn copy_completed(&mut self, object_id: ObjectId, byte_count: u64) { + let Some(output) = self + .pending_outputs + .iter() + .find(|output| output.object_id == object_id) + .copied() + else { + return; + }; + if self.full_rings.contains(&output.ring_id) || byte_count != output.extent { + return; + } + self.committed + .entry(output.ring_id) + .or_default() + .extend(std::iter::repeat(0).take(byte_count as usize)); + *self.payload_committed.entry(output.ring_id).or_insert(0) += byte_count; + if self.produced.insert(object_id) { + self.events.push(WorkerEgressOut::ObjectProduced { + ring_id: output.ring_id, + object_id, + sequence: output.sequence, + }); + } + let step_ids = self + .steps + .iter() + .filter(|step| { + step.outputs + .iter() + .any(|output| output.object_id == object_id) + }) + .map(|step| step.step_id) + .collect::>(); + for step_id in step_ids { + self.maybe_step_completed(step_id); + } + } + + fn copy_failed(&mut self, object_id: ObjectId) { + let step_id = self + .steps + .iter() + .find(|step| { + step.outputs + .iter() + .any(|output| output.object_id == object_id) + }) + .map(|step| step.step_id) + .unwrap_or(StepId(0)); + self.events.push(WorkerEgressOut::StepFailed { + step_id, + reason: StepFailureReason::DeviceCopyFailed, + }); + } + + fn maybe_step_completed(&mut self, step_id: StepId) { + let Some(step) = self.steps.iter().find(|step| step.step_id == step_id) else { + return; + }; + if !step.role_state_updated { + return; + } + if step.outputs.iter().all(|output| self.produced.contains(&output.object_id)) + && !self.events.iter().any(|event| matches!(event, WorkerEgressOut::StepCompleted { step_id: seen } if *seen == step_id)) + { + self.events.push(WorkerEgressOut::StepCompleted { step_id }); + } + } +} + +fn encode_header(output: OutputBinding) -> Vec { + let mut out = Vec::with_capacity(HEADER_LEN); + out.extend_from_slice(b"MO01"); + out.push(1); + out.push(HEADER_LEN as u8); + out.extend_from_slice(&[0u8; 2]); + out.extend_from_slice(&output.object_id.0.to_le_bytes()); + out.extend_from_slice(&output.sequence.to_le_bytes()); + out.extend_from_slice(&output.extent.to_le_bytes()); + out.extend_from_slice(&[0u8; 16]); + out +} diff --git a/crates/mvp-system/src/gpu_worker_ingress_parser.rs b/crates/mvp-system/src/gpu_worker_ingress_parser.rs new file mode 100644 index 0000000..aafe94e --- /dev/null +++ b/crates/mvp-system/src/gpu_worker_ingress_parser.rs @@ -0,0 +1,425 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct WorkerGeneration(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RingId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct EdgeId(pub u64); +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct PortId(pub String); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ObjectId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct DeviceHandle { + pub generation: WorkerGeneration, + pub id: u64, +} + +impl DeviceHandle { + pub fn new(generation: WorkerGeneration, id: u64) -> Self { + Self { generation, id } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RingDirection { + Ingress, + Egress, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ObjectLayout { + Token, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ObjectSpec { + pub max_extent: u64, + pub alignment: u64, + pub layout: ObjectLayout, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct InstallRing { + pub ring_id: RingId, + pub edge_id: EdgeId, + pub port_id: PortId, + pub direction: RingDirection, + pub object_spec: ObjectSpec, + pub generation: WorkerGeneration, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WorkerIngressEvent { + InstallRing(InstallRing), + RingReadable { + ring_id: RingId, + }, + Eof { + ring_id: RingId, + }, + DeviceCopyCompleted { + object_id: ObjectId, + byte_count: u64, + }, + DeviceHandleCreated { + object_id: ObjectId, + handle: DeviceHandle, + }, + RingFault { + ring_id: RingId, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ObjectFailureReason { + UnsupportedMagic, + UnsupportedVersion, + MalformedHeaderLength, + ExtentExceedsMax, + ExtentAlignmentViolation, + SequenceViolation, + EofBeforeFullPayload, + DeviceCopyFailed, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WorkerIngressOut { + ObjectLoaded { + ring_id: RingId, + edge_id: EdgeId, + port_id: PortId, + object_id: ObjectId, + sequence: u64, + extent: u64, + handle: DeviceHandle, + }, + ObjectFailed { + ring_id: RingId, + object_id: Option, + reason: ObjectFailureReason, + }, + RingFault { + ring_id: RingId, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct DeviceCopyLog { + pub object_id: ObjectId, + pub byte_count: u64, +} + +#[derive(Clone, Debug)] +pub struct ObjectRecordBuilder { + spec: ObjectSpec, + object_id: ObjectId, + sequence: u64, + extent: u64, + payload: Vec, + magic: [u8; 4], + version: u8, + header_len: u8, +} + +impl ObjectRecordBuilder { + pub fn new(spec: ObjectSpec) -> Self { + Self { + spec, + object_id: ObjectId(9000), + sequence: 0, + extent: 0, + payload: Vec::new(), + magic: *b"MO01", + version: 1, + header_len: HEADER_LEN as u8, + } + } + + pub fn object_id(mut self, object_id: ObjectId) -> Self { + self.object_id = object_id; + self + } + + pub fn sequence(mut self, sequence: u64) -> Self { + self.sequence = sequence; + self + } + + pub fn extent(mut self, extent: u64) -> Self { + self.extent = extent; + self + } + + pub fn payload(mut self, payload: Vec) -> Self { + self.payload = payload; + self + } + + pub fn partial_payload(mut self, payload: Vec) -> Self { + self.payload = payload; + self + } + + pub fn unsupported_magic(mut self) -> Self { + self.magic = *b"BAD!"; + self + } + + pub fn unsupported_version(mut self) -> Self { + self.version = 99; + self + } + + pub fn malformed_header_length(mut self) -> Self { + self.header_len = 1; + self + } + + pub fn encode(mut self) -> Vec { + if self.extent == 0 && !self.payload.is_empty() { + self.extent = self.payload.len() as u64; + } + let mut out = Vec::with_capacity(HEADER_LEN + self.payload.len()); + out.extend_from_slice(&self.magic); + out.push(self.version); + out.push(self.header_len); + out.extend_from_slice(&[0u8; 2]); + out.extend_from_slice(&self.object_id.0.to_le_bytes()); + out.extend_from_slice(&self.sequence.to_le_bytes()); + out.extend_from_slice(&self.extent.to_le_bytes()); + out.extend_from_slice(&self.spec.max_extent.to_le_bytes()); + out.extend_from_slice(&self.spec.alignment.to_le_bytes()); + out.extend_from_slice(&self.payload); + out + } +} + +const HEADER_LEN: usize = 48; + +#[derive(Clone, Debug, PartialEq, Eq)] +struct ParsedRecord { + object_id: ObjectId, + sequence: u64, + extent: u64, + total_len: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct PendingObject { + record: ParsedRecord, + copy_done: bool, + handle: Option, +} + +#[cfg(test)] +pub struct IngressParserHarness { + generation: WorkerGeneration, + install: Option, + buffers: std::collections::BTreeMap>, + consume: std::collections::BTreeMap, + cursor_reload: std::collections::BTreeMap, + faulted_rings: std::collections::BTreeSet, + expected_sequence: u64, + pending: std::collections::BTreeMap, + copy_log: Vec, + events: Vec, +} + +#[cfg(test)] +impl IngressParserHarness { + pub fn new(generation: WorkerGeneration) -> Self { + Self { + generation, + install: None, + buffers: std::collections::BTreeMap::new(), + consume: std::collections::BTreeMap::new(), + cursor_reload: std::collections::BTreeMap::new(), + faulted_rings: std::collections::BTreeSet::new(), + expected_sequence: 0, + pending: std::collections::BTreeMap::new(), + copy_log: Vec::new(), + events: Vec::new(), + } + } + + pub fn observe(&mut self, event: WorkerIngressEvent) { + match event { + WorkerIngressEvent::InstallRing(install) => { + if install.direction == RingDirection::Ingress + && install.generation == self.generation + { + self.consume.entry(install.ring_id).or_insert(0); + self.install = Some(install); + } + } + WorkerIngressEvent::RingReadable { ring_id } => self.parse_ring(ring_id, false), + WorkerIngressEvent::Eof { ring_id } => self.parse_ring(ring_id, true), + WorkerIngressEvent::DeviceCopyCompleted { + object_id, + byte_count, + } => { + self.copy_log.push(DeviceCopyLog { + object_id, + byte_count, + }); + if let Some(pending) = self.pending.get_mut(&object_id) { + pending.copy_done = true; + if byte_count == pending.record.extent { + if let Some(install) = &self.install { + *self.consume.entry(install.ring_id).or_insert(0) += + pending.record.total_len as u64; + } + } + } + self.maybe_loaded(object_id); + } + WorkerIngressEvent::DeviceHandleCreated { object_id, handle } => { + if let Some(pending) = self.pending.get_mut(&object_id) { + pending.handle = Some(handle); + } + self.maybe_loaded(object_id); + } + WorkerIngressEvent::RingFault { ring_id } => { + self.faulted_rings.insert(ring_id); + self.events.push(WorkerIngressOut::RingFault { ring_id }); + } + } + } + + pub fn write_committed_bytes(&mut self, ring_id: RingId, bytes: Vec) { + self.buffers.entry(ring_id).or_default().extend(bytes); + } + + pub fn write_uncommitted_bytes(&mut self, _ring_id: RingId, _bytes: Vec) {} + + pub fn consume_cursor(&self, ring_id: RingId) -> u64 { + self.consume.get(&ring_id).copied().unwrap_or(0) + } + + pub fn cursor_reload_count(&self, ring_id: RingId) -> u64 { + self.cursor_reload.get(&ring_id).copied().unwrap_or(0) + } + + pub fn device_copy_log(&self) -> &[DeviceCopyLog] { + &self.copy_log + } + + pub fn events(&self) -> &[WorkerIngressOut] { + &self.events + } + + fn parse_ring(&mut self, ring_id: RingId, eof: bool) { + let Some(install) = &self.install else { + return; + }; + if install.ring_id != ring_id || self.faulted_rings.contains(&ring_id) { + return; + } + *self.cursor_reload.entry(ring_id).or_insert(0) += 1; + let buffer = self.buffers.get(&ring_id).cloned().unwrap_or_default(); + if buffer.is_empty() { + return; + } + match decode_record(&buffer, install.object_spec, eof) { + Ok(record) => { + if record.sequence != self.expected_sequence { + self.events.push(WorkerIngressOut::ObjectFailed { + ring_id, + object_id: Some(record.object_id), + reason: ObjectFailureReason::SequenceViolation, + }); + return; + } + self.expected_sequence += 1; + self.pending.insert( + record.object_id, + PendingObject { + record: record.clone(), + copy_done: false, + handle: None, + }, + ); + self.copy_log.push(DeviceCopyLog { + object_id: record.object_id, + byte_count: record.extent, + }); + } + Err(reason) => self.events.push(WorkerIngressOut::ObjectFailed { + ring_id, + object_id: None, + reason, + }), + } + } + + fn maybe_loaded(&mut self, object_id: ObjectId) { + let Some(pending) = self.pending.get(&object_id).cloned() else { + return; + }; + let Some(handle) = pending.handle else { + return; + }; + if !pending.copy_done || handle.generation != self.generation { + return; + } + let install = self.install.as_ref().unwrap(); + if !self.events.iter().any(|event| matches!(event, WorkerIngressOut::ObjectLoaded { object_id: seen, .. } if *seen == object_id)) { + self.events.push(WorkerIngressOut::ObjectLoaded { + ring_id: install.ring_id, + edge_id: install.edge_id, + port_id: install.port_id.clone(), + object_id, + sequence: pending.record.sequence, + extent: pending.record.extent, + handle, + }); + } + } +} + +fn decode_record( + bytes: &[u8], + spec: ObjectSpec, + eof: bool, +) -> Result { + if bytes.len() < HEADER_LEN { + return if eof { + Err(ObjectFailureReason::EofBeforeFullPayload) + } else { + Err(ObjectFailureReason::MalformedHeaderLength) + }; + } + if &bytes[0..4] != b"MO01" { + return Err(ObjectFailureReason::UnsupportedMagic); + } + if bytes[4] != 1 { + return Err(ObjectFailureReason::UnsupportedVersion); + } + if bytes[5] as usize != HEADER_LEN { + return Err(ObjectFailureReason::MalformedHeaderLength); + } + let object_id = ObjectId(u64::from_le_bytes(bytes[8..16].try_into().unwrap())); + let sequence = u64::from_le_bytes(bytes[16..24].try_into().unwrap()); + let extent = u64::from_le_bytes(bytes[24..32].try_into().unwrap()); + if extent > spec.max_extent { + return Err(ObjectFailureReason::ExtentExceedsMax); + } + if spec.alignment != 0 && extent % spec.alignment != 0 { + return Err(ObjectFailureReason::ExtentAlignmentViolation); + } + let total_len = HEADER_LEN + extent as usize; + if bytes.len() < total_len { + return if eof { + Err(ObjectFailureReason::EofBeforeFullPayload) + } else { + Err(ObjectFailureReason::MalformedHeaderLength) + }; + } + Ok(ParsedRecord { + object_id, + sequence, + extent, + total_len, + }) +} diff --git a/crates/mvp-system/src/gpu_worker_process_adapter.rs b/crates/mvp-system/src/gpu_worker_process_adapter.rs new file mode 100644 index 0000000..29346b1 --- /dev/null +++ b/crates/mvp-system/src/gpu_worker_process_adapter.rs @@ -0,0 +1,290 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ArenaFd(pub i32); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct HelperAbiVersion(pub u32); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RingId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct StepId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct DeviceHandle(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct WorkerGeneration(pub u64); + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct AdapterConfig { + pub arena_fd: ArenaFd, + pub arena_bytes: u64, + pub helper_abi_version: HelperAbiVersion, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WorkerCommand { + InitializeWorker { + helper_abi_version: HelperAbiVersion, + }, + InstallRing { + ring_id: RingId, + }, + ExecuteStep { + step_id: StepId, + }, + ReleaseDeviceObject { + handle: DeviceHandle, + }, + ShutdownWorker, + RingReadable { + ring_id: RingId, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct JsonLine(serde_json::Map); + +impl JsonLine { + pub fn parse(line: &str) -> Result { + let value: serde_json::Value = serde_json::from_str(line.trim_end())?; + match value { + serde_json::Value::Object(map) => Ok(Self(map)), + other => Ok(Self(serde_json::Map::from_iter([("value".into(), other)]))), + } + } + + pub fn is_object(&self) -> bool { + true + } + + pub fn contains_key(&self, key: &str) -> bool { + self.0.contains_key(key) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ProcessFaultReason { + InvalidJson, + UnknownEventShape, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum WorkerFatalReason { + UnsupportedHelperAbi, + BackendUnavailable, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum AdapterEvent { + WorkerReady { + generation: WorkerGeneration, + }, + WorkerFatal { + reason: WorkerFatalReason, + }, + ProcessFault { + reason: ProcessFaultReason, + line: String, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WorkerAction { + ReadArenaEnvironment { + arena_fd: ArenaFd, + arena_bytes: u64, + }, + MapArena { + arena_fd: ArenaFd, + arena_bytes: u64, + }, + InitializeRingHelper { + helper_abi_version: HelperAbiVersion, + }, + InitializeBackend, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum InitializationFailure { + BackendUnavailable, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ExitStatus { + Code(i32), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CommandRejectionReason { + PayloadBytesForbidden, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CommandRejection { + pub reason: CommandRejectionReason, +} + +#[cfg(test)] +pub struct ProcessAdapterHarness { + config: AdapterConfig, + stdin_lines: Vec, + events: Vec, + worker_actions: Vec, + command_rejections: Vec, + exit_status: Option, +} + +#[cfg(test)] +impl ProcessAdapterHarness { + pub fn new(config: AdapterConfig) -> Self { + Self { + config, + stdin_lines: Vec::new(), + events: Vec::new(), + worker_actions: Vec::new(), + command_rejections: Vec::new(), + exit_status: None, + } + } + + pub fn start_worker_process(&mut self) { + self.worker_actions + .push(WorkerAction::ReadArenaEnvironment { + arena_fd: self.config.arena_fd, + arena_bytes: self.config.arena_bytes, + }); + } + + pub fn send_command(&mut self, command: WorkerCommand) { + if let WorkerCommand::InitializeWorker { helper_abi_version } = command { + if helper_abi_version != self.config.helper_abi_version { + self.events.push(AdapterEvent::WorkerFatal { + reason: WorkerFatalReason::UnsupportedHelperAbi, + }); + return; + } + self.worker_actions.push(WorkerAction::MapArena { + arena_fd: self.config.arena_fd, + arena_bytes: self.config.arena_bytes, + }); + self.worker_actions + .push(WorkerAction::InitializeRingHelper { helper_abi_version }); + self.worker_actions.push(WorkerAction::InitializeBackend); + } + self.stdin_lines.push(format_json_command(&command)); + } + + pub fn send_raw_json_command(&mut self, raw: &str) { + match JsonLine::parse(raw) { + Ok(parsed) + if parsed.contains_key("payload") + || parsed.contains_key("bytes") + || parsed.contains_key("data") => + { + self.command_rejections.push(CommandRejection { + reason: CommandRejectionReason::PayloadBytesForbidden, + }); + } + Ok(_) => self.stdin_lines.push(format!("{}\n", raw.trim_end())), + Err(_) => self.command_rejections.push(CommandRejection { + reason: CommandRejectionReason::PayloadBytesForbidden, + }), + } + } + + pub fn receive_stdout_line(&mut self, line: &str) { + let parsed = match JsonLine::parse(line) { + Ok(parsed) => parsed, + Err(_) => { + self.events.push(AdapterEvent::ProcessFault { + reason: ProcessFaultReason::InvalidJson, + line: line.into(), + }); + return; + } + }; + let Some(serde_json::Value::String(kind)) = parsed.0.get("type") else { + self.events.push(AdapterEvent::ProcessFault { + reason: ProcessFaultReason::UnknownEventShape, + line: line.into(), + }); + return; + }; + match kind.as_str() { + "WorkerReady" => { + let generation = parsed + .0 + .get("generation") + .and_then(|value| value.as_u64()) + .unwrap_or(1); + self.events.push(AdapterEvent::WorkerReady { + generation: WorkerGeneration(generation), + }); + } + _ => self.events.push(AdapterEvent::ProcessFault { + reason: ProcessFaultReason::UnknownEventShape, + line: line.into(), + }), + } + } + + pub fn receive_stderr_line(&mut self, _line: &str) {} + + pub fn inject_initialization_failure(&mut self, failure: InitializationFailure) { + match failure { + InitializationFailure::BackendUnavailable => { + self.events.push(AdapterEvent::WorkerFatal { + reason: WorkerFatalReason::BackendUnavailable, + }); + self.exit_status = Some(ExitStatus::Code(1)); + } + } + } + + pub fn stdin_lines(&self) -> &[String] { + &self.stdin_lines + } + + pub fn events(&self) -> &[AdapterEvent] { + &self.events + } + + pub fn worker_actions(&self) -> &[WorkerAction] { + &self.worker_actions + } + + pub fn exit_status(&self) -> Option { + self.exit_status + } + + pub fn command_rejections(&self) -> &[CommandRejection] { + &self.command_rejections + } +} + +#[cfg(test)] +fn format_json_command(command: &WorkerCommand) -> String { + let value = match command { + WorkerCommand::InitializeWorker { helper_abi_version } => serde_json::json!({ + "type": "InitializeWorker", + "helper_abi_version": helper_abi_version.0, + }), + WorkerCommand::InstallRing { ring_id } => serde_json::json!({ + "type": "InstallRing", + "ring_id": ring_id.0, + }), + WorkerCommand::ExecuteStep { step_id } => serde_json::json!({ + "type": "ExecuteStep", + "step_id": step_id.0, + }), + WorkerCommand::ReleaseDeviceObject { handle } => serde_json::json!({ + "type": "ReleaseDeviceObject", + "handle": handle.0, + }), + WorkerCommand::ShutdownWorker => serde_json::json!({ + "type": "ShutdownWorker", + }), + WorkerCommand::RingReadable { ring_id } => serde_json::json!({ + "type": "RingReadable", + "ring_id": ring_id.0, + }), + }; + format!("{}\n", value) +} diff --git a/crates/mvp-system/src/lib.rs b/crates/mvp-system/src/lib.rs index 7a18f05..ac78c9e 100644 --- a/crates/mvp-system/src/lib.rs +++ b/crates/mvp-system/src/lib.rs @@ -1 +1,29 @@ +#[cfg(test)] +extern crate self as mvp_system; + pub mod arena_manager; +pub mod device_bridge; +pub mod driver_pumps; +pub mod edge_establisher; +pub mod gpu_worker_ctl; +pub mod gpu_worker_egress_producer; +pub mod gpu_worker_ingress_parser; +pub mod gpu_worker_process_adapter; +pub mod membership_pool_readiness; +pub mod node_boot_lifecycle; +pub mod observability_surface; +pub mod orchestrator_run_fsm; +pub mod orchestrator_token_endpoint; +pub mod resource_inventory; +pub mod run_plan; +pub mod shared_ring_helper_abi; +pub mod stage_controller; +pub mod tx_rx_edge_actor; +pub mod weight_lifecycle; + +#[cfg(test)] +mod tests; + +#[cfg(test)] +#[path = "tests/driver_pumps_guarantees.rs"] +mod driver_pumps_guarantees; diff --git a/crates/mvp-system/src/membership_pool_readiness.rs b/crates/mvp-system/src/membership_pool_readiness.rs new file mode 100644 index 0000000..2c9636c --- /dev/null +++ b/crates/mvp-system/src/membership_pool_readiness.rs @@ -0,0 +1,248 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct NodeId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RunId(pub u64); + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PoolId(pub String); + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ReadinessConfig { + pub pool_id: PoolId, + pub convergence_window_ms: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RunRequest { + pub run_id: RunId, +} + +impl RunRequest { + pub fn new(run_id: RunId) -> Self { + Self { run_id } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Observation { + NodeKnown { + node_id: NodeId, + }, + SwimLive { + node_id: NodeId, + }, + NodeAvailable { + node_id: NodeId, + }, + DataPlaneIdentityReady { + node_id: NodeId, + }, + SwimSuspect { + node_id: NodeId, + }, + NodeFaulted { + node_id: NodeId, + }, + DataPlaneIdentityMissing { + node_id: NodeId, + }, + SwimLost { + node_id: NodeId, + }, + RunProvisioned { + run_id: RunId, + required_nodes: Vec, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RunFaultReason { + RequiredNodeLost { node_id: NodeId }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ReadinessEvent { + PoolReady { + pool: Vec, + }, + RunFaulted { + run_id: RunId, + reason: RunFaultReason, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ReadinessCommand { + EmitPoolReady { pool: Vec }, + StartPlanning { run_id: RunId, pool: Vec }, + WaitForStability { pool_id: PoolId }, + AbortPendingRun { run_id: RunId }, + CommitRunPlan { run_id: RunId }, + RecomputePlacement { run_id: RunId }, + AssignStage { node_id: NodeId }, + AssignEdge { node_id: NodeId }, + AssignLayerRange { node_id: NodeId }, + AssignObjectSpec { node_id: NodeId }, +} + +#[derive(Clone, Default)] +struct NodeFacts { + known: bool, + live: bool, + available: bool, + identity_ready: bool, + poisoned: bool, +} + +#[cfg(test)] +pub struct ReadinessGateHarness { + config: ReadinessConfig, + candidates: Vec, + facts: std::collections::BTreeMap, + events: Vec, + commands: Vec, + now_ms: u64, + stable_since_ms: Option, + emitted_ready: bool, + pending_run: Option, + provisioned_run: Option<(RunId, Vec)>, +} + +#[cfg(test)] +impl ReadinessGateHarness { + pub fn new(config: ReadinessConfig, candidates: Vec) -> Self { + let facts = candidates + .iter() + .copied() + .map(|node_id| (node_id, NodeFacts::default())) + .collect(); + Self { + config, + candidates, + facts, + events: Vec::new(), + commands: Vec::new(), + now_ms: 0, + stable_since_ms: None, + emitted_ready: false, + pending_run: None, + provisioned_run: None, + } + } + + pub fn observe(&mut self, observation: Observation) { + match observation { + Observation::NodeKnown { node_id } => self.fact(node_id).known = true, + Observation::SwimLive { node_id } => self.fact(node_id).live = true, + Observation::NodeAvailable { node_id } => self.fact(node_id).available = true, + Observation::DataPlaneIdentityReady { node_id } => { + self.fact(node_id).identity_ready = true + } + Observation::SwimSuspect { node_id } + | Observation::NodeFaulted { node_id } + | Observation::DataPlaneIdentityMissing { node_id } => { + self.fact(node_id).poisoned = true + } + Observation::SwimLost { node_id } => { + self.fact(node_id).live = false; + self.fact(node_id).poisoned = true; + if let Some((run_id, required)) = &self.provisioned_run { + if required.contains(&node_id) + && !self.events.iter().any(|event| { + matches!(event, ReadinessEvent::RunFaulted { run_id: seen, .. } if seen == run_id) + }) + { + self.events.push(ReadinessEvent::RunFaulted { + run_id: *run_id, + reason: RunFaultReason::RequiredNodeLost { node_id }, + }); + } + } else if let Some(pending) = &self.pending_run { + self.commands.push(ReadinessCommand::AbortPendingRun { + run_id: pending.run_id, + }); + } + } + Observation::RunProvisioned { + run_id, + required_nodes, + } => { + self.provisioned_run = Some((run_id, required_nodes)); + } + } + self.reset_stability_if_needed(); + self.maybe_ready(); + } + + pub fn advance_time_ms(&mut self, delta: u64) { + self.now_ms = self.now_ms.saturating_add(delta); + self.maybe_ready(); + } + + pub fn request_run_planning(&mut self, request: RunRequest) { + self.pending_run = Some(request); + self.maybe_ready(); + } + + pub fn events(&self) -> &[ReadinessEvent] { + &self.events + } + + pub fn commands(&self) -> &[ReadinessCommand] { + &self.commands + } + + fn fact(&mut self, node_id: NodeId) -> &mut NodeFacts { + self.facts.entry(node_id).or_default() + } + + fn complete_now(&self) -> bool { + self.candidates.iter().all(|node_id| { + self.facts.get(node_id).is_some_and(|facts| { + facts.known + && facts.live + && facts.available + && facts.identity_ready + && !facts.poisoned + }) + }) + } + + fn reset_stability_if_needed(&mut self) { + if self.complete_now() { + if self.stable_since_ms.is_none() { + self.stable_since_ms = Some(self.now_ms); + self.commands.push(ReadinessCommand::WaitForStability { + pool_id: self.config.pool_id.clone(), + }); + } + } else { + self.stable_since_ms = None; + } + } + + fn maybe_ready(&mut self) { + if self.emitted_ready || !self.complete_now() { + return; + } + let Some(stable_since) = self.stable_since_ms else { + return; + }; + if self.now_ms.saturating_sub(stable_since) < self.config.convergence_window_ms { + return; + } + self.emitted_ready = true; + let pool = self.candidates.clone(); + self.events + .push(ReadinessEvent::PoolReady { pool: pool.clone() }); + self.commands + .push(ReadinessCommand::EmitPoolReady { pool: pool.clone() }); + if let Some(request) = &self.pending_run { + self.commands.push(ReadinessCommand::StartPlanning { + run_id: request.run_id, + pool, + }); + } + } +} diff --git a/crates/mvp-system/src/node_boot_lifecycle.rs b/crates/mvp-system/src/node_boot_lifecycle.rs new file mode 100644 index 0000000..b7e141b --- /dev/null +++ b/crates/mvp-system/src/node_boot_lifecycle.rs @@ -0,0 +1,226 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct NodeId(pub u64); + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PoolId(pub String); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum WorkerStartPolicy { + StartBeforeAvailable, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct BootConfig { + pub node_id: NodeId, + pub intended_pool_id: PoolId, + pub worker_policy: WorkerStartPolicy, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum ResourceFact { + RustProcessAlive, + RuntimeAcceptingControl, + StableNodeIdKnown, + ArenaMapped, + GpuWorkerReady, + TransportEndpointBound, + SwimJoiningPool, + ProvisioningReceiverOpen, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ResourceOutcome { + RustProcessAlive, + RuntimeAcceptingControl, + StableNodeIdKnown(NodeId), + ArenaMapped, + GpuWorkerReady, + TransportEndpointBound, + SwimJoiningPool, + ProvisioningReceiverOpen, + ArenaFault, + GpuWorkerFault, + TransportEndpointFault, + InvalidNodeIdentity, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum BootFaultKind { + ArenaConstructionFailed, + WorkerStartupFailed, + TransportEndpointFailed, + InvalidNodeIdentity, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum LifecycleEvent { + NodeAvailable { + node_id: NodeId, + }, + NodeFaulted { + node_id: NodeId, + kind: BootFaultKind, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum BootCommand { + AdvertiseLifecycle { node_id: NodeId }, + JoinMembership { pool_id: PoolId }, + OpenProvisioningInbox { node_id: NodeId }, + LoadWeights { node_id: NodeId }, + ConfigureRole { node_id: NodeId }, + EstablishEdge { node_id: NodeId }, + AssignStage { node_id: NodeId }, + AssignLayerRange { node_id: NodeId }, + AssignEdge { node_id: NodeId }, + AssignObjectSpec { node_id: NodeId }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ProvisionRequest { + pub node_id: NodeId, +} + +impl ProvisionRequest { + pub fn for_node(node_id: NodeId) -> Self { + Self { node_id } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ProvisioningAdmissionRejection { + NodeNotAvailable, + WrongNode, +} + +#[cfg(test)] +pub struct BootHarness { + config: BootConfig, + facts: std::collections::BTreeSet, + events: Vec, + commands: Vec, + available: bool, + faulted: bool, +} + +#[cfg(test)] +impl BootHarness { + pub fn launch(config: BootConfig) -> Self { + let commands = vec![ + BootCommand::AdvertiseLifecycle { + node_id: config.node_id, + }, + BootCommand::JoinMembership { + pool_id: config.intended_pool_id.clone(), + }, + BootCommand::OpenProvisioningInbox { + node_id: config.node_id, + }, + ]; + Self { + config, + facts: std::collections::BTreeSet::new(), + events: Vec::new(), + commands, + available: false, + faulted: false, + } + } + + pub fn observe(&mut self, outcome: ResourceOutcome) { + if self.available || self.faulted { + return; + } + match outcome { + ResourceOutcome::RustProcessAlive => self.insert_fact(ResourceFact::RustProcessAlive), + ResourceOutcome::RuntimeAcceptingControl => { + self.insert_fact(ResourceFact::RuntimeAcceptingControl) + } + ResourceOutcome::StableNodeIdKnown(node_id) if node_id == self.config.node_id => { + self.insert_fact(ResourceFact::StableNodeIdKnown) + } + ResourceOutcome::StableNodeIdKnown(_) | ResourceOutcome::InvalidNodeIdentity => { + self.fault(BootFaultKind::InvalidNodeIdentity) + } + ResourceOutcome::ArenaMapped => self.insert_fact(ResourceFact::ArenaMapped), + ResourceOutcome::GpuWorkerReady => self.insert_fact(ResourceFact::GpuWorkerReady), + ResourceOutcome::TransportEndpointBound => { + self.insert_fact(ResourceFact::TransportEndpointBound) + } + ResourceOutcome::SwimJoiningPool => self.insert_fact(ResourceFact::SwimJoiningPool), + ResourceOutcome::ProvisioningReceiverOpen => { + self.insert_fact(ResourceFact::ProvisioningReceiverOpen) + } + ResourceOutcome::ArenaFault => self.fault(BootFaultKind::ArenaConstructionFailed), + ResourceOutcome::GpuWorkerFault => self.fault(BootFaultKind::WorkerStartupFailed), + ResourceOutcome::TransportEndpointFault => { + self.fault(BootFaultKind::TransportEndpointFailed) + } + } + self.maybe_available(); + } + + pub fn events(&self) -> &[LifecycleEvent] { + &self.events + } + + pub fn commands(&self) -> &[BootCommand] { + &self.commands + } + + pub fn is_candidate_eligible(&self, node_id: NodeId) -> bool { + node_id == self.config.node_id && self.available && !self.faulted + } + + pub fn try_accept_provisioning( + &self, + request: ProvisionRequest, + ) -> Result<(), ProvisioningAdmissionRejection> { + if request.node_id != self.config.node_id { + return Err(ProvisioningAdmissionRejection::WrongNode); + } + if self.available && !self.faulted { + Ok(()) + } else { + Err(ProvisioningAdmissionRejection::NodeNotAvailable) + } + } + + fn insert_fact(&mut self, fact: ResourceFact) { + self.facts.insert(fact); + } + + fn maybe_available(&mut self) { + if self.available || self.faulted { + return; + } + let required = [ + ResourceFact::RustProcessAlive, + ResourceFact::RuntimeAcceptingControl, + ResourceFact::StableNodeIdKnown, + ResourceFact::ArenaMapped, + ResourceFact::GpuWorkerReady, + ResourceFact::TransportEndpointBound, + ResourceFact::SwimJoiningPool, + ResourceFact::ProvisioningReceiverOpen, + ]; + if required.iter().all(|fact| self.facts.contains(fact)) { + self.available = true; + self.events.push(LifecycleEvent::NodeAvailable { + node_id: self.config.node_id, + }); + } + } + + fn fault(&mut self, kind: BootFaultKind) { + if self.faulted || self.available { + return; + } + self.faulted = true; + self.events.push(LifecycleEvent::NodeFaulted { + node_id: self.config.node_id, + kind, + }); + } +} diff --git a/crates/mvp-system/src/observability_surface.rs b/crates/mvp-system/src/observability_surface.rs new file mode 100644 index 0000000..18e6fef --- /dev/null +++ b/crates/mvp-system/src/observability_surface.rs @@ -0,0 +1,443 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RunId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct NodeId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct StageIndex(pub u32); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct EdgeId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RingId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ObjectId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Sequence(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct StepId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct WorkerGeneration(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum EventKind { + NodeStarted, + NodeAvailable, + PoolReady, + RunPlanned, + StageProvisionStarted, + WeightsDownloadStarted, + WeightsDownloaded, + WeightsLoaded, + EdgeProvisionStarted, + EdgeReady, + StageReady, + ReadinessBarrierPassed, + PromptInjected, + ObjectLoaded, + ExecuteStepStarted, + ObjectProduced, + StepCompleted, + TokenReceived, + RunCompleted, + RunFaulted, + StopRunSent, + StageStopped, + RunTornDown, + StageFaulted, + RingReadable, + WorkerReady, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Component { + NodeBoot, + Membership, + Orchestrator, + StageController, + WeightLifecycle, + EdgeEstablisher, + GpuWorkerCtl, + SharedRingHelper, + TokenEndpoint, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FaultReason { + WorkerCrashed, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Event { + RunScoped { + kind: EventKind, + run_id: RunId, + reason: Option, + component: Component, + }, + NodeScoped { + kind: EventKind, + node_id: NodeId, + component: Component, + }, + StageScoped { + kind: EventKind, + run_id: RunId, + stage_index: StageIndex, + reason: Option, + component: Component, + }, + EdgeScoped { + kind: EventKind, + edge_id: EdgeId, + component: Component, + }, + RingScoped { + kind: EventKind, + ring_id: RingId, + component: Component, + }, + ObjectScoped { + kind: EventKind, + object_id: ObjectId, + sequence: Sequence, + component: Component, + }, + StepScoped { + kind: EventKind, + step_id: StepId, + component: Component, + }, + WorkerScoped { + kind: EventKind, + worker_generation: WorkerGeneration, + component: Component, + }, +} + +impl Event { + pub fn kind(&self) -> EventKind { + match self { + Event::RunScoped { kind, .. } + | Event::NodeScoped { kind, .. } + | Event::StageScoped { kind, .. } + | Event::EdgeScoped { kind, .. } + | Event::RingScoped { kind, .. } + | Event::ObjectScoped { kind, .. } + | Event::StepScoped { kind, .. } + | Event::WorkerScoped { kind, .. } => *kind, + } + } +} + +pub struct TraceBuilder { + run_id: RunId, + events: Vec, +} + +impl TraceBuilder { + pub fn new(run_id: RunId) -> Self { + Self { + run_id, + events: Vec::new(), + } + } + + pub fn node_started(mut self, node_id: NodeId) -> Self { + self.events.push(Event::NodeScoped { + kind: EventKind::NodeStarted, + node_id, + component: Component::NodeBoot, + }); + self + } + + pub fn node_available(mut self, node_id: NodeId) -> Self { + self.events.push(Event::NodeScoped { + kind: EventKind::NodeAvailable, + node_id, + component: Component::NodeBoot, + }); + self + } + + pub fn pool_ready(mut self, _nodes: Vec) -> Self { + self.events.push(Event::RunScoped { + kind: EventKind::PoolReady, + run_id: self.run_id, + reason: None, + component: Component::Membership, + }); + self + } + + pub fn run_planned(mut self) -> Self { + self.events.push(Event::RunScoped { + kind: EventKind::RunPlanned, + run_id: self.run_id, + reason: None, + component: Component::Orchestrator, + }); + self + } + + pub fn stage_provision_started(mut self, stage_index: StageIndex, _node_id: NodeId) -> Self { + self.stage( + EventKind::StageProvisionStarted, + stage_index, + None, + Component::Orchestrator, + ); + self + } + + pub fn weights_download_started(mut self, stage_index: StageIndex) -> Self { + self.stage( + EventKind::WeightsDownloadStarted, + stage_index, + None, + Component::WeightLifecycle, + ); + self + } + + pub fn weights_downloaded(mut self, stage_index: StageIndex) -> Self { + self.stage( + EventKind::WeightsDownloaded, + stage_index, + None, + Component::WeightLifecycle, + ); + self + } + + pub fn weights_loaded(mut self, stage_index: StageIndex) -> Self { + self.stage( + EventKind::WeightsLoaded, + stage_index, + None, + Component::WeightLifecycle, + ); + self + } + + pub fn edge_provision_started(mut self, edge_id: EdgeId) -> Self { + self.edge(EventKind::EdgeProvisionStarted, edge_id); + self + } + + pub fn edge_ready(mut self, edge_id: EdgeId) -> Self { + self.edge(EventKind::EdgeReady, edge_id); + self + } + + pub fn stage_ready(mut self, stage_index: StageIndex) -> Self { + self.stage( + EventKind::StageReady, + stage_index, + None, + Component::StageController, + ); + self + } + + pub fn readiness_barrier_passed(mut self) -> Self { + self.run( + EventKind::ReadinessBarrierPassed, + None, + Component::Orchestrator, + ); + self + } + + pub fn prompt_injected(mut self, sequence: Sequence) -> Self { + self.events.push(Event::ObjectScoped { + kind: EventKind::PromptInjected, + object_id: ObjectId(9000), + sequence, + component: Component::TokenEndpoint, + }); + self + } + + pub fn object_loaded( + mut self, + _edge_id: EdgeId, + object_id: ObjectId, + sequence: Sequence, + ) -> Self { + self.events.push(Event::ObjectScoped { + kind: EventKind::ObjectLoaded, + object_id, + sequence, + component: Component::GpuWorkerCtl, + }); + self + } + + pub fn execute_step_started(mut self, step_id: StepId) -> Self { + self.events.push(Event::StepScoped { + kind: EventKind::ExecuteStepStarted, + step_id, + component: Component::StageController, + }); + self + } + + pub fn object_produced( + mut self, + _edge_id: EdgeId, + object_id: ObjectId, + sequence: Sequence, + ) -> Self { + self.events.push(Event::ObjectScoped { + kind: EventKind::ObjectProduced, + object_id, + sequence, + component: Component::GpuWorkerCtl, + }); + self + } + + pub fn step_completed(mut self, step_id: StepId) -> Self { + self.events.push(Event::StepScoped { + kind: EventKind::StepCompleted, + step_id, + component: Component::StageController, + }); + self + } + + pub fn token_received(mut self, object_id: ObjectId, sequence: Sequence) -> Self { + self.events.push(Event::ObjectScoped { + kind: EventKind::TokenReceived, + object_id, + sequence, + component: Component::TokenEndpoint, + }); + self + } + + pub fn run_completed(mut self) -> Self { + self.run(EventKind::RunCompleted, None, Component::Orchestrator); + self + } + + pub fn stage_faulted( + mut self, + stage_index: StageIndex, + reason: FaultReason, + component: Component, + ) -> Self { + self.stage( + EventKind::StageFaulted, + stage_index, + Some(reason), + component, + ); + self + } + + pub fn run_faulted(mut self, reason: FaultReason, component: Component) -> Self { + self.run(EventKind::RunFaulted, Some(reason), component); + self + } + + pub fn stop_run_sent(mut self, stage_index: StageIndex) -> Self { + self.stage( + EventKind::StopRunSent, + stage_index, + None, + Component::Orchestrator, + ); + self + } + + pub fn stage_stopped(mut self, stage_index: StageIndex) -> Self { + self.stage( + EventKind::StageStopped, + stage_index, + None, + Component::StageController, + ); + self + } + + pub fn run_torn_down(mut self) -> Self { + self.run(EventKind::RunTornDown, None, Component::Orchestrator); + self + } + + pub fn finish(self) -> Vec { + self.events + } + + fn run(&mut self, kind: EventKind, reason: Option, component: Component) { + self.events.push(Event::RunScoped { + kind, + run_id: self.run_id, + reason, + component, + }); + } + + fn stage( + &mut self, + kind: EventKind, + stage_index: StageIndex, + reason: Option, + component: Component, + ) { + self.events.push(Event::StageScoped { + kind, + run_id: self.run_id, + stage_index, + reason, + component, + }); + } + + fn edge(&mut self, kind: EventKind, edge_id: EdgeId) { + self.events.push(Event::EdgeScoped { + kind, + edge_id, + component: Component::EdgeEstablisher, + }); + } +} + +pub fn requires_log_scraping(_events: &[Event]) -> bool { + false +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Batching { + None, + Fixed(usize), +} + +#[cfg(test)] +pub struct EventSubscriberHarness { + events: Vec, + _batching: Batching, +} + +#[cfg(test)] +impl EventSubscriberHarness { + pub fn collect(events: Vec, batching: Batching) -> Self { + Self { + events, + _batching: batching, + } + } + + pub fn flattened_events(&self) -> &[Event] { + &self.events + } + + pub fn used_transport_specific_assertions(&self) -> bool { + false + } + + pub fn used_storage_specific_assertions(&self) -> bool { + false + } +} diff --git a/crates/mvp-system/src/orchestrator_run_fsm.rs b/crates/mvp-system/src/orchestrator_run_fsm.rs new file mode 100644 index 0000000..3546a40 --- /dev/null +++ b/crates/mvp-system/src/orchestrator_run_fsm.rs @@ -0,0 +1,424 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RunId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct NodeId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct StageRef { + pub stage_index: u32, + pub node_id: NodeId, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RunPlan { + pub run_id: RunId, + pub stages: Vec, +} + +impl RunPlan { + pub fn test_linear(run_id: RunId, stages: Vec) -> Self { + Self { run_id, stages } + } + + pub fn stage_nodes(&self) -> Vec { + self.stages.iter().map(|stage| stage.node_id).collect() + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RunConfig { + pub run_id: RunId, + pub max_tokens: u64, + pub prompt: Vec, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum StageFaultReason { + WorkerCrashed, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum EndpointKind { + TokenIn, + TokenOut, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum TimeoutKind { + Execution, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum RunEvent { + PoolReady { + nodes: Vec, + }, + PlanAvailable(RunPlan), + StageReady { + run_id: RunId, + stage_index: u32, + }, + TokenInEndpointReady, + TokenOutEndpointReady, + TokenReceived { + sequence: u64, + token_id: u32, + eos: bool, + }, + StageFault { + run_id: RunId, + stage_index: u32, + reason: StageFaultReason, + }, + EndpointFault { + run_id: RunId, + endpoint: EndpointKind, + }, + Timeout { + run_id: RunId, + kind: TimeoutKind, + }, + StageStopped { + run_id: RunId, + stage_index: u32, + }, + TokenEndpointsStopped, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RunFaultReason { + StageFault { + stage_index: u32, + reason: StageFaultReason, + }, + EndpointFault { + endpoint: EndpointKind, + }, + Timeout { + kind: TimeoutKind, + }, + UnknownStageReady { + stage_index: u32, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum LifecycleEvent { + RunRejected { + run_id: RunId, + reason: RunFaultReason, + }, + RunFaulted { + run_id: RunId, + reason: RunFaultReason, + }, + RunCompleted { + run_id: RunId, + }, + RunTornDown { + run_id: RunId, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct StageProvision { + pub run_id: RunId, + pub stage_index: u32, + pub node_id: NodeId, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum RunCommand { + ProvisionStage { + provision: StageProvision, + }, + CreateTokenInEndpoint { + run_id: RunId, + }, + CreateTokenOutEndpoint { + run_id: RunId, + }, + InjectPrompt { + run_id: RunId, + sequence: u64, + prompt: Vec, + }, + BroadcastStart { + run_id: RunId, + }, + StopRun { + run_id: RunId, + stage_index: u32, + }, + TearDownTokenEndpoints { + run_id: RunId, + }, +} + +#[cfg(test)] +pub struct OrchestratorHarness { + config: RunConfig, + plan: Option, + pool_ready: bool, + provisioned: bool, + token_in_ready: bool, + token_out_ready: bool, + ready_stages: std::collections::BTreeSet, + injected_sequences: Vec, + expected_token_sequence: u64, + events: Vec, + commands: Vec, + terminal: bool, + teardown_started: bool, + stopped_stages: std::collections::BTreeSet, + token_endpoints_stopped: bool, +} + +#[cfg(test)] +impl OrchestratorHarness { + pub fn new(config: RunConfig) -> Self { + Self { + config, + plan: None, + pool_ready: false, + provisioned: false, + token_in_ready: false, + token_out_ready: false, + ready_stages: std::collections::BTreeSet::new(), + injected_sequences: Vec::new(), + expected_token_sequence: 0, + events: Vec::new(), + commands: Vec::new(), + terminal: false, + teardown_started: false, + stopped_stages: std::collections::BTreeSet::new(), + token_endpoints_stopped: false, + } + } + + pub fn observe(&mut self, event: RunEvent) { + match event { + RunEvent::PoolReady { .. } => { + self.pool_ready = true; + self.maybe_provision(); + } + RunEvent::PlanAvailable(plan) => { + self.plan = Some(plan); + self.maybe_provision(); + } + RunEvent::StageReady { + run_id, + stage_index, + } => { + if run_id != self.config.run_id || !self.plan_has_stage(stage_index) { + self.fault(RunFaultReason::UnknownStageReady { stage_index }); + return; + } + self.ready_stages.insert(stage_index); + self.maybe_inject_initial(); + } + RunEvent::TokenInEndpointReady => { + self.token_in_ready = true; + self.maybe_inject_initial(); + } + RunEvent::TokenOutEndpointReady => { + self.token_out_ready = true; + self.maybe_inject_initial(); + } + RunEvent::TokenReceived { sequence, eos, .. } => { + if self.terminal { + return; + } + if sequence != self.expected_token_sequence { + return; + } + self.expected_token_sequence += 1; + if eos { + self.complete(); + } else if (self.injected_sequences.len() as u64) < self.config.max_tokens { + self.inject(sequence + 1); + } else { + self.complete(); + } + } + RunEvent::StageFault { + run_id, + stage_index, + reason, + } if run_id == self.config.run_id => { + self.fault(RunFaultReason::StageFault { + stage_index, + reason, + }); + } + RunEvent::EndpointFault { run_id, endpoint } if run_id == self.config.run_id => { + self.fault(RunFaultReason::EndpointFault { endpoint }); + } + RunEvent::Timeout { run_id, kind } if run_id == self.config.run_id => { + self.fault(RunFaultReason::Timeout { kind }); + } + RunEvent::StageStopped { + run_id, + stage_index, + } if run_id == self.config.run_id => { + self.stopped_stages.insert(stage_index); + self.maybe_torn_down(); + } + RunEvent::TokenEndpointsStopped => { + self.token_endpoints_stopped = true; + self.maybe_torn_down(); + } + RunEvent::StageFault { .. } + | RunEvent::EndpointFault { .. } + | RunEvent::Timeout { .. } + | RunEvent::StageStopped { .. } => {} + } + } + + pub fn advance_time_ms(&mut self, _delta: u64) {} + + pub fn commands(&self) -> &[RunCommand] { + &self.commands + } + + pub fn events(&self) -> &[LifecycleEvent] { + &self.events + } + + pub fn injected_sequences(&self) -> Vec { + self.injected_sequences.clone() + } + + fn maybe_provision(&mut self) { + if !self.pool_ready || self.provisioned { + return; + } + let Some(plan) = &self.plan else { + return; + }; + self.provisioned = true; + for stage in &plan.stages { + self.commands.push(RunCommand::ProvisionStage { + provision: StageProvision { + run_id: plan.run_id, + stage_index: stage.stage_index, + node_id: stage.node_id, + }, + }); + } + self.commands.push(RunCommand::CreateTokenInEndpoint { + run_id: self.config.run_id, + }); + self.commands.push(RunCommand::CreateTokenOutEndpoint { + run_id: self.config.run_id, + }); + } + + fn maybe_inject_initial(&mut self) { + if self.terminal || !self.provisioned || !self.injected_sequences.is_empty() { + return; + } + if self.token_in_ready && self.token_out_ready && self.all_stages_ready() { + self.inject(0); + } + } + + fn inject(&mut self, sequence: u64) { + if self.terminal { + return; + } + self.injected_sequences.push(sequence); + self.commands.push(RunCommand::InjectPrompt { + run_id: self.config.run_id, + sequence, + prompt: if sequence == 0 { + self.config.prompt.clone() + } else { + Vec::new() + }, + }); + } + + fn complete(&mut self) { + if self.terminal { + return; + } + self.terminal = true; + self.events.push(LifecycleEvent::RunCompleted { + run_id: self.config.run_id, + }); + self.start_teardown(); + } + + fn fault(&mut self, reason: RunFaultReason) { + if self.terminal { + return; + } + self.terminal = true; + self.events.push(LifecycleEvent::RunFaulted { + run_id: self.config.run_id, + reason, + }); + self.start_teardown(); + } + + fn start_teardown(&mut self) { + if self.teardown_started { + return; + } + self.teardown_started = true; + if let Some(plan) = &self.plan { + for stage in &plan.stages { + self.commands.push(RunCommand::StopRun { + run_id: self.config.run_id, + stage_index: stage.stage_index, + }); + } + } + self.commands.push(RunCommand::TearDownTokenEndpoints { + run_id: self.config.run_id, + }); + } + + fn maybe_torn_down(&mut self) { + if !self.teardown_started || !self.token_endpoints_stopped { + return; + } + if !self.all_stages_stopped() { + return; + } + if !self + .events + .iter() + .any(|event| matches!(event, LifecycleEvent::RunTornDown { .. })) + { + self.events.push(LifecycleEvent::RunTornDown { + run_id: self.config.run_id, + }); + } + } + + fn plan_has_stage(&self, stage_index: u32) -> bool { + self.plan.as_ref().is_some_and(|plan| { + plan.stages + .iter() + .any(|stage| stage.stage_index == stage_index) + }) + } + + fn all_stages_ready(&self) -> bool { + self.plan.as_ref().is_some_and(|plan| { + plan.stages + .iter() + .all(|stage| self.ready_stages.contains(&stage.stage_index)) + }) + } + + fn all_stages_stopped(&self) -> bool { + self.plan.as_ref().is_some_and(|plan| { + plan.stages + .iter() + .all(|stage| self.stopped_stages.contains(&stage.stage_index)) + }) + } +} diff --git a/crates/mvp-system/src/orchestrator_token_endpoint.rs b/crates/mvp-system/src/orchestrator_token_endpoint.rs new file mode 100644 index 0000000..1e82759 --- /dev/null +++ b/crates/mvp-system/src/orchestrator_token_endpoint.rs @@ -0,0 +1,325 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RunId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct NodeId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct EdgeId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ObjectId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ObjectSpec { + pub max_extent: u64, + pub alignment: u64, +} + +impl ObjectSpec { + pub fn test_tokens() -> Self { + Self { + max_extent: 64, + alignment: 4, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum EndpointDirection { + TokenIn, + TokenOut, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum EdgeSemantics { + SingleProducerSingleConsumer, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct EdgePlan { + pub edge_id: EdgeId, + pub producer_node_id: NodeId, + pub consumer_node_id: NodeId, + pub direction: EndpointDirection, +} + +impl EdgePlan { + pub fn token_in( + edge_id: EdgeId, + orchestrator_node_id: NodeId, + first_stage_node: NodeId, + ) -> Self { + Self { + edge_id, + producer_node_id: orchestrator_node_id, + consumer_node_id: first_stage_node, + direction: EndpointDirection::TokenIn, + } + } + + pub fn token_out( + edge_id: EdgeId, + last_stage_node: NodeId, + orchestrator_node_id: NodeId, + ) -> Self { + Self { + edge_id, + producer_node_id: last_stage_node, + consumer_node_id: orchestrator_node_id, + direction: EndpointDirection::TokenOut, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct TokenEndpointPlan { + pub run_id: RunId, + pub orchestrator_node_id: NodeId, + pub token_in_edge: EdgePlan, + pub token_out_edge: EdgePlan, + pub token_spec: ObjectSpec, + pub max_tokens: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct LocalEndpoint { + pub edge_id: EdgeId, + pub orchestrator_node_id: NodeId, + pub edge_semantics: EdgeSemantics, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct TokenObjectWrite { + pub run_id: RunId, + pub edge_id: EdgeId, + pub sequence: u64, + pub tokens: Vec, + pub object_spec: ObjectSpec, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum EndpointCommand { + CreateTokenInProducer { + edge_id: EdgeId, + orchestrator_node_id: NodeId, + }, + CreateTokenOutConsumer { + edge_id: EdgeId, + orchestrator_node_id: NodeId, + }, + WriteTokenObject(TokenObjectWrite), + BroadcastStart { + run_id: RunId, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum EndpointEvent { + TokenInReady { + edge_id: EdgeId, + }, + TokenOutReady { + edge_id: EdgeId, + }, + AllStagesReady, + ReadinessBarrierPassed, + TokenObjectReceived { + edge_id: EdgeId, + object_id: ObjectId, + sequence: u64, + token_id: u32, + eos: bool, + }, + EndpointFault { + edge_id: EdgeId, + direction: EndpointDirection, + }, + MalformedTokenObject { + edge_id: EdgeId, + object_id: ObjectId, + }, + TeardownFailed { + edge_id: EdgeId, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RunFaultReason { + TokenSequenceViolation, + TokenEndpointFault, + MalformedTokenObject, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum EndpointLifecycleEvent { + RunFaulted { + run_id: RunId, + reason: RunFaultReason, + }, + TeardownFailed { + run_id: RunId, + edge_id: EdgeId, + }, +} + +#[cfg(test)] +pub struct TokenEndpointHarness { + plan: TokenEndpointPlan, + commands: Vec, + events: Vec, + local_endpoints: Vec, + prompt: Option>, + token_in_ready: bool, + token_out_ready: bool, + stages_ready: bool, + barrier_passed: bool, + injected_sequences: Vec, + expected_token_sequence: u64, + terminal: bool, +} + +#[cfg(test)] +impl TokenEndpointHarness { + pub fn new(plan: TokenEndpointPlan) -> Self { + let commands = vec![ + EndpointCommand::CreateTokenInProducer { + edge_id: plan.token_in_edge.edge_id, + orchestrator_node_id: plan.orchestrator_node_id, + }, + EndpointCommand::CreateTokenOutConsumer { + edge_id: plan.token_out_edge.edge_id, + orchestrator_node_id: plan.orchestrator_node_id, + }, + ]; + let local_endpoints = vec![ + LocalEndpoint { + edge_id: plan.token_in_edge.edge_id, + orchestrator_node_id: plan.orchestrator_node_id, + edge_semantics: EdgeSemantics::SingleProducerSingleConsumer, + }, + LocalEndpoint { + edge_id: plan.token_out_edge.edge_id, + orchestrator_node_id: plan.orchestrator_node_id, + edge_semantics: EdgeSemantics::SingleProducerSingleConsumer, + }, + ]; + Self { + plan, + commands, + events: Vec::new(), + local_endpoints, + prompt: None, + token_in_ready: false, + token_out_ready: false, + stages_ready: false, + barrier_passed: false, + injected_sequences: Vec::new(), + expected_token_sequence: 0, + terminal: false, + } + } + + pub fn commands(&self) -> &[EndpointCommand] { + &self.commands + } + + pub fn events(&self) -> &[EndpointLifecycleEvent] { + &self.events + } + + pub fn local_endpoints(&self) -> &[LocalEndpoint] { + &self.local_endpoints + } + + pub fn injected_sequences(&self) -> Vec { + self.injected_sequences.clone() + } + + pub fn request_prompt_injection(&mut self, prompt: Vec) { + self.prompt = Some(prompt); + self.maybe_inject_initial(); + } + + pub fn advance_time_ms(&mut self, _delta: u64) {} + + pub fn observe(&mut self, event: EndpointEvent) { + match event { + EndpointEvent::TokenInReady { edge_id } + if edge_id == self.plan.token_in_edge.edge_id => + { + self.token_in_ready = true + } + EndpointEvent::TokenOutReady { edge_id } + if edge_id == self.plan.token_out_edge.edge_id => + { + self.token_out_ready = true + } + EndpointEvent::AllStagesReady => self.stages_ready = true, + EndpointEvent::ReadinessBarrierPassed => self.barrier_passed = true, + EndpointEvent::TokenObjectReceived { sequence, eos, .. } => { + if self.terminal { + return; + } + if sequence != self.expected_token_sequence { + self.fault(RunFaultReason::TokenSequenceViolation); + return; + } + self.expected_token_sequence += 1; + if eos || self.injected_sequences.len() as u64 >= self.plan.max_tokens { + self.terminal = true; + } else { + self.inject_sequence(sequence + 1, vec![0]); + } + } + EndpointEvent::EndpointFault { .. } => self.fault(RunFaultReason::TokenEndpointFault), + EndpointEvent::MalformedTokenObject { .. } => { + self.fault(RunFaultReason::MalformedTokenObject) + } + EndpointEvent::TeardownFailed { edge_id } => { + self.events.push(EndpointLifecycleEvent::TeardownFailed { + run_id: self.plan.run_id, + edge_id, + }); + } + EndpointEvent::TokenInReady { .. } | EndpointEvent::TokenOutReady { .. } => {} + } + self.maybe_inject_initial(); + } + + fn maybe_inject_initial(&mut self) { + if self.injected_sequences.is_empty() + && self.prompt.is_some() + && self.token_in_ready + && self.token_out_ready + && self.stages_ready + && self.barrier_passed + { + let prompt = self.prompt.clone().unwrap_or_default(); + self.inject_sequence(0, prompt); + } + } + + fn inject_sequence(&mut self, sequence: u64, tokens: Vec) { + if self.terminal || self.injected_sequences.len() as u64 >= self.plan.max_tokens { + return; + } + self.injected_sequences.push(sequence); + self.commands + .push(EndpointCommand::WriteTokenObject(TokenObjectWrite { + run_id: self.plan.run_id, + edge_id: self.plan.token_in_edge.edge_id, + sequence, + tokens, + object_spec: self.plan.token_spec, + })); + } + + fn fault(&mut self, reason: RunFaultReason) { + if self.terminal { + return; + } + self.terminal = true; + self.events.push(EndpointLifecycleEvent::RunFaulted { + run_id: self.plan.run_id, + reason, + }); + } +} diff --git a/crates/mvp-system/src/resource_inventory.rs b/crates/mvp-system/src/resource_inventory.rs new file mode 100644 index 0000000..6ea5417 --- /dev/null +++ b/crates/mvp-system/src/resource_inventory.rs @@ -0,0 +1,199 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RunId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct NodeId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum GpuClass { + TestSmall, + TestLarge, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum BootHealth { + Ready, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct InventoryEntry { + pub node_id: NodeId, + pub gpu_class: GpuClass, + pub ready: bool, +} + +impl InventoryEntry { + pub fn ready(node_id: NodeId, gpu_class: GpuClass) -> Self { + Self { + node_id, + gpu_class, + ready: true, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct StagePlacement { + pub stage_index: u32, + pub node_id: NodeId, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum PlacementInput { + FixedLinear(Vec), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PlanningRequest { + pub run_id: RunId, + pub stage_count: u32, + pub entries: Vec, + pub placement: PlacementInput, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct StagePlan { + pub stage_index: u32, + pub node_id: NodeId, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RunPlan { + pub run_id: RunId, + pub stages: Vec, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum InventoryRejectionKind { + UnknownNode { node_id: NodeId }, + DuplicateStage { stage_index: u32 }, + MissingStage { stage_index: u32 }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct InventoryRejection { + pub kind: InventoryRejectionKind, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum NodeReport { + BootHealth { + node_id: NodeId, + health: BootHealth, + }, + CapacityChanged { + node_id: NodeId, + gpu_class: GpuClass, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum InventoryCommand { + AcceptPlacementNegotiation { node_id: NodeId }, + RewriteStagePlacement { node_id: NodeId }, + ProvisionStage { stage_index: u32, node_id: NodeId }, +} + +pub fn plan_from_inventory(request: PlanningRequest) -> Result { + let known_nodes = request + .entries + .iter() + .map(|entry| entry.node_id) + .collect::>(); + let PlacementInput::FixedLinear(stages) = request.placement; + let mut by_stage = std::collections::BTreeMap::new(); + for placement in stages { + if !known_nodes.contains(&placement.node_id) { + return Err(InventoryRejection { + kind: InventoryRejectionKind::UnknownNode { + node_id: placement.node_id, + }, + }); + } + if by_stage + .insert(placement.stage_index, placement.node_id) + .is_some() + { + return Err(InventoryRejection { + kind: InventoryRejectionKind::DuplicateStage { + stage_index: placement.stage_index, + }, + }); + } + } + + let mut stages = Vec::with_capacity(request.stage_count as usize); + for stage_index in 0..request.stage_count { + let node_id = by_stage.remove(&stage_index).ok_or(InventoryRejection { + kind: InventoryRejectionKind::MissingStage { stage_index }, + })?; + stages.push(StagePlan { + stage_index, + node_id, + }); + } + Ok(RunPlan { + run_id: request.run_id, + stages, + }) +} + +#[cfg(test)] +pub struct InventoryHarness { + entries: Vec, + commands: Vec, + committed_plan: Option, +} + +#[cfg(test)] +impl InventoryHarness { + pub fn new(entries: Vec) -> Self { + Self { + entries, + commands: Vec::new(), + committed_plan: None, + } + } + + pub fn entries(&self) -> &[InventoryEntry] { + &self.entries + } + + pub fn commands(&self) -> &[InventoryCommand] { + &self.commands + } + + pub fn committed_plan(&self) -> Option<&RunPlan> { + self.committed_plan.as_ref() + } + + pub fn observe_node_report(&mut self, report: NodeReport) { + match report { + NodeReport::BootHealth { node_id, health } => { + if let Some(entry) = self + .entries + .iter_mut() + .find(|entry| entry.node_id == node_id) + { + entry.ready = matches!(health, BootHealth::Ready); + } + } + NodeReport::CapacityChanged { .. } => {} + } + } + + pub fn commit_plan(&mut self, plan: RunPlan) { + self.committed_plan = Some(plan); + } + + pub fn provision_stages(&mut self) { + if let Some(plan) = &self.committed_plan { + for stage in &plan.stages { + self.commands.push(InventoryCommand::ProvisionStage { + stage_index: stage.stage_index, + node_id: stage.node_id, + }); + } + } + } +} diff --git a/crates/mvp-system/src/run_plan.rs b/crates/mvp-system/src/run_plan.rs new file mode 100644 index 0000000..a498fa1 --- /dev/null +++ b/crates/mvp-system/src/run_plan.rs @@ -0,0 +1,422 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RunId(pub u64); + +impl From for RunId { + fn from(value: u64) -> Self { + Self(value) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct NodeId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct EdgeId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum DTypeFamily { + BFloat, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ModelFacts { + pub model_id: String, + pub num_layers: u32, + pub hidden_dim: u64, + pub dtype_family: DTypeFamily, + pub dtype_width_bytes: u64, + pub max_seq_len: u64, + pub eos_token_id: u32, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RuntimeConfig { + pub max_tokens: u32, +} + +impl RuntimeConfig { + pub fn test_default() -> Self { + Self { max_tokens: 4 } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct StagePlacement { + pub stage_index: u32, + pub node_id: NodeId, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum PlacementInput { + FixedLinear(Vec), + FixedLinearWithEdgeOverride { + stages: Vec, + forced_activation_edges: Vec<(u32, u32)>, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct RingSpec { + pub data_capacity: u64, + pub alignment: u64, +} + +impl RingSpec { + pub fn test_default_activation() -> Self { + Self { + data_capacity: 1 << 20, + alignment: 64, + } + } + + pub fn test_default_token() -> Self { + Self { + data_capacity: 4096, + alignment: 8, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PlannerInput { + pub run_id: RunId, + pub orchestrator_node_id: NodeId, + pub model: ModelFacts, + pub runtime: RuntimeConfig, + pub candidate_pool: Vec, + pub stage_count: u32, + pub placement: PlacementInput, + pub activation_ring: RingSpec, + pub token_ring: RingSpec, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum EdgeKind { + TokenIn, + Activation, + TokenOut, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ObjectKind { + Token, + Activation, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ObjectSpec { + pub kind: ObjectKind, + pub max_extent: u64, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum EdgeEndpoint { + Orchestrator { node_id: NodeId }, + Stage { node_id: NodeId, stage_index: u32 }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct EdgePlan { + pub run_id: RunId, + pub edge_id: EdgeId, + pub kind: EdgeKind, + pub producer: EdgeEndpoint, + pub consumer: EdgeEndpoint, + pub object_spec: ObjectSpec, + pub ring_spec: RingSpec, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct StagePlan { + pub run_id: RunId, + pub stage_index: u32, + pub stage_count: u32, + pub node_id: NodeId, + pub layer_start: u32, + pub layer_end_exclusive: u32, + pub inbound_edge: EdgeId, + pub outbound_edge: EdgeId, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RunPlan { + pub run_id: RunId, + pub stages: Vec, + pub edges: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct InboundEdgeProvision { + pub edge_id: EdgeId, + pub kind: EdgeKind, + pub object_spec: ObjectSpec, + pub ring_spec: RingSpec, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct OutboundEdgeProvision { + pub edge_id: EdgeId, + pub kind: EdgeKind, + pub consumer_node_id: NodeId, + pub object_spec: ObjectSpec, + pub ring_spec: RingSpec, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ProvisionStage { + pub run_id: RunId, + pub node_id: NodeId, + pub stage_index: u32, + pub stage_count: u32, + pub layer_start: u32, + pub layer_end_exclusive: u32, + pub inbound: InboundEdgeProvision, + pub outbound: OutboundEdgeProvision, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum PlanRejectionKind { + UnknownNode, + DuplicateStageAssignment, + MissingStage, + InvalidStageCount, + EdgeEndpointMismatch, + ModelStageLayoutMismatch, + InvalidObjectSpec, + UnsupportedShapeOrLayout, + InvalidRingSpec, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PlanRejection { + kind: PlanRejectionKind, +} + +impl PlanRejection { + pub fn kind(&self) -> PlanRejectionKind { + self.kind + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ProjectionRejection { + UnknownStage, + MissingEdge, +} + +pub fn plan_run(input: PlannerInput) -> Result { + validate_global_input(&input)?; + let placements = validated_placements(&input)?; + + let activation_extent = input + .model + .max_seq_len + .checked_mul(input.model.hidden_dim) + .and_then(|value| value.checked_mul(input.model.dtype_width_bytes)) + .ok_or_else(|| reject(PlanRejectionKind::InvalidObjectSpec))?; + if activation_extent == 0 { + return Err(reject(PlanRejectionKind::InvalidObjectSpec)); + } + + let token_spec = ObjectSpec { + kind: ObjectKind::Token, + max_extent: 64, + }; + let activation_spec = ObjectSpec { + kind: ObjectKind::Activation, + max_extent: activation_extent, + }; + + let mut edges = Vec::with_capacity(input.stage_count as usize + 1); + edges.push(EdgePlan { + run_id: input.run_id, + edge_id: EdgeId(7000), + kind: EdgeKind::TokenIn, + producer: EdgeEndpoint::Orchestrator { + node_id: input.orchestrator_node_id, + }, + consumer: EdgeEndpoint::Stage { + node_id: placements[0].node_id, + stage_index: 0, + }, + object_spec: token_spec, + ring_spec: input.token_ring, + }); + + for stage_index in 0..input.stage_count.saturating_sub(1) { + edges.push(EdgePlan { + run_id: input.run_id, + edge_id: EdgeId(7001 + u64::from(stage_index)), + kind: EdgeKind::Activation, + producer: EdgeEndpoint::Stage { + node_id: placements[stage_index as usize].node_id, + stage_index, + }, + consumer: EdgeEndpoint::Stage { + node_id: placements[stage_index as usize + 1].node_id, + stage_index: stage_index + 1, + }, + object_spec: activation_spec, + ring_spec: input.activation_ring, + }); + } + + edges.push(EdgePlan { + run_id: input.run_id, + edge_id: EdgeId(7000 + u64::from(input.stage_count)), + kind: EdgeKind::TokenOut, + producer: EdgeEndpoint::Stage { + node_id: placements[input.stage_count as usize - 1].node_id, + stage_index: input.stage_count - 1, + }, + consumer: EdgeEndpoint::Orchestrator { + node_id: input.orchestrator_node_id, + }, + object_spec: token_spec, + ring_spec: input.token_ring, + }); + + let mut stages = Vec::with_capacity(input.stage_count as usize); + for placement in &placements { + let stage_index = placement.stage_index; + let (start, end) = layer_range(input.model.num_layers, input.stage_count, stage_index); + stages.push(StagePlan { + run_id: input.run_id, + stage_index, + stage_count: input.stage_count, + node_id: placement.node_id, + layer_start: start, + layer_end_exclusive: end, + inbound_edge: EdgeId(7000 + u64::from(stage_index)), + outbound_edge: EdgeId(7001 + u64::from(stage_index)), + }); + } + + Ok(RunPlan { + run_id: input.run_id, + stages, + edges, + }) +} + +pub fn derive_stage_provision( + plan: &RunPlan, + stage_index: u32, +) -> Result { + let stage = plan + .stages + .iter() + .find(|stage| stage.stage_index == stage_index) + .ok_or(ProjectionRejection::UnknownStage)?; + let inbound = plan + .edges + .iter() + .find(|edge| edge.edge_id == stage.inbound_edge) + .ok_or(ProjectionRejection::MissingEdge)?; + let outbound = plan + .edges + .iter() + .find(|edge| edge.edge_id == stage.outbound_edge) + .ok_or(ProjectionRejection::MissingEdge)?; + Ok(ProvisionStage { + run_id: plan.run_id, + node_id: stage.node_id, + stage_index, + stage_count: stage.stage_count, + layer_start: stage.layer_start, + layer_end_exclusive: stage.layer_end_exclusive, + inbound: InboundEdgeProvision { + edge_id: inbound.edge_id, + kind: inbound.kind, + object_spec: inbound.object_spec, + ring_spec: inbound.ring_spec, + }, + outbound: OutboundEdgeProvision { + edge_id: outbound.edge_id, + kind: outbound.kind, + consumer_node_id: endpoint_node_id(&outbound.consumer), + object_spec: outbound.object_spec, + ring_spec: outbound.ring_spec, + }, + }) +} + +fn validate_global_input(input: &PlannerInput) -> Result<(), PlanRejection> { + if input.stage_count == 0 { + return Err(reject(PlanRejectionKind::InvalidStageCount)); + } + if input.model.num_layers < input.stage_count { + return Err(reject(PlanRejectionKind::ModelStageLayoutMismatch)); + } + if input.model.max_seq_len == 0 || input.model.dtype_width_bytes == 0 { + return Err(reject(PlanRejectionKind::InvalidObjectSpec)); + } + if input.model.hidden_dim == 0 { + return Err(reject(PlanRejectionKind::UnsupportedShapeOrLayout)); + } + if !valid_ring(input.activation_ring) || !valid_ring(input.token_ring) { + return Err(reject(PlanRejectionKind::InvalidRingSpec)); + } + Ok(()) +} + +fn validated_placements(input: &PlannerInput) -> Result, PlanRejection> { + let (stages, forced_edges) = match &input.placement { + PlacementInput::FixedLinear(stages) => (stages.as_slice(), &[][..]), + PlacementInput::FixedLinearWithEdgeOverride { + stages, + forced_activation_edges, + } => (stages.as_slice(), forced_activation_edges.as_slice()), + }; + + if forced_edges + .iter() + .any(|(producer, consumer)| *consumer != producer.saturating_add(1)) + { + return Err(reject(PlanRejectionKind::EdgeEndpointMismatch)); + } + + let candidate_nodes = input + .candidate_pool + .iter() + .copied() + .collect::>(); + let mut by_stage = std::collections::BTreeMap::new(); + for placement in stages { + if !candidate_nodes.contains(&placement.node_id) { + return Err(reject(PlanRejectionKind::UnknownNode)); + } + if by_stage.insert(placement.stage_index, *placement).is_some() { + return Err(reject(PlanRejectionKind::DuplicateStageAssignment)); + } + } + + let mut dense = Vec::with_capacity(input.stage_count as usize); + for stage_index in 0..input.stage_count { + let placement = by_stage + .remove(&stage_index) + .ok_or_else(|| reject(PlanRejectionKind::MissingStage))?; + dense.push(placement); + } + Ok(dense) +} + +fn layer_range(num_layers: u32, stage_count: u32, stage_index: u32) -> (u32, u32) { + let start = (u64::from(num_layers) * u64::from(stage_index) / u64::from(stage_count)) as u32; + let end = (u64::from(num_layers) * u64::from(stage_index + 1) / u64::from(stage_count)) as u32; + (start, end) +} + +fn endpoint_node_id(endpoint: &EdgeEndpoint) -> NodeId { + match endpoint { + EdgeEndpoint::Orchestrator { node_id } | EdgeEndpoint::Stage { node_id, .. } => *node_id, + } +} + +fn valid_ring(spec: RingSpec) -> bool { + spec.data_capacity > 0 && spec.alignment > 0 && spec.alignment.is_power_of_two() +} + +fn reject(kind: PlanRejectionKind) -> PlanRejection { + PlanRejection { kind } +} diff --git a/crates/mvp-system/src/shared_ring_helper_abi.rs b/crates/mvp-system/src/shared_ring_helper_abi.rs new file mode 100644 index 0000000..8ac33b4 --- /dev/null +++ b/crates/mvp-system/src/shared_ring_helper_abi.rs @@ -0,0 +1,293 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct NodeId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RingId(pub u64); + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct EndpointId(pub String); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ArenaBase(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ProcessLocalPointer(pub u64); + +impl ProcessLocalPointer { + pub fn from_base_plus_offset(base: ArenaBase, offset: u64) -> Self { + Self(base.0 + offset) + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RingConfig { + pub node_id: NodeId, + pub ring_id: RingId, + pub capacity: u64, + pub producer: EndpointId, + pub consumer: EndpointId, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RingIdentity { + pub ring_id: RingId, + pub capacity: u64, + pub producer_count: u32, + pub consumer_count: u32, + generation: u64, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct CursorSnapshot { + pub commit: u64, + pub consume: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WakeHint { + RingReadable { ring_id: RingId }, + RingWritable { ring_id: RingId }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ReserveError { + InsufficientSpace, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ReadError { + BeyondCommittedBytes, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Reservation { + start: u64, + len: u64, + generation: u64, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct StaleWake { + ring_id: RingId, + generation: u64, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct WakeDelivery { + pub was_accepted: bool, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SchedulerState { + pub readable_rings: std::collections::BTreeSet, + pub writable_rings: std::collections::BTreeSet, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct RingLayout { + pub data_offset: u64, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ProcessLocalView { + pub layout: RingLayout, + pub data_pointer: ProcessLocalPointer, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum PythonOperation { + ReserveViaHelper { ring_id: RingId }, + CommitViaHelper { ring_id: RingId }, + ReadableViaHelper { ring_id: RingId }, + ConsumeViaHelper { ring_id: RingId }, + MapPointerViaHelper { ring_id: RingId }, + DirectAtomicAccess { ring_id: RingId }, + DirectWrapArithmetic { ring_id: RingId }, +} + +#[cfg(test)] +pub struct RingHelperHarness { + identity: RingIdentity, + buffer: Vec, + commit: u64, + consume: u64, + retired: bool, + wake_hints: Vec, + wake_log: Vec, + scheduler_state: SchedulerState, +} + +#[cfg(test)] +impl RingHelperHarness { + pub fn create(config: RingConfig) -> Self { + Self::with_generation(config.ring_id, config.capacity, 1) + } + + pub fn create_replacement(ring_id: RingId) -> Self { + Self::with_generation(ring_id, 8, 2) + } + + fn with_generation(ring_id: RingId, capacity: u64, generation: u64) -> Self { + Self { + identity: RingIdentity { + ring_id, + capacity, + producer_count: 1, + consumer_count: 1, + generation, + }, + buffer: vec![0; capacity as usize], + commit: 0, + consume: 0, + retired: false, + wake_hints: Vec::new(), + wake_log: Vec::new(), + scheduler_state: SchedulerState { + readable_rings: std::collections::BTreeSet::new(), + writable_rings: std::collections::BTreeSet::new(), + }, + } + } + + pub fn identity(&self) -> RingIdentity { + self.identity.clone() + } + + pub fn retire_and_capture_stale_wake(&mut self) -> StaleWake { + self.retired = true; + StaleWake { + ring_id: self.identity.ring_id, + generation: self.identity.generation, + } + } + + pub fn deliver_wake(&mut self, wake: StaleWake) { + self.wake_log.push(WakeDelivery { + was_accepted: wake.ring_id == self.identity.ring_id + && wake.generation == self.identity.generation + && !self.retired, + }); + } + + pub fn wake_log(&self) -> &[WakeDelivery] { + &self.wake_log + } + + pub fn producer_reserve(&self, len: u64) -> Result { + let used = self.commit.saturating_sub(self.consume); + if len <= self.identity.capacity.saturating_sub(used) { + Ok(Reservation { + start: self.commit, + len, + generation: self.identity.generation, + }) + } else { + Err(ReserveError::InsufficientSpace) + } + } + + pub fn producer_write(&mut self, reservation: &Reservation, bytes: &[u8]) { + assert_eq!(reservation.generation, self.identity.generation); + assert_eq!(reservation.len as usize, bytes.len()); + for (i, byte) in bytes.iter().copied().enumerate() { + let idx = (reservation.start + i as u64) % self.identity.capacity; + self.buffer[idx as usize] = byte; + } + } + + pub fn producer_commit(&mut self, reservation: Reservation) { + assert_eq!(reservation.start, self.commit); + self.commit = self.commit.saturating_add(reservation.len); + self.wake_hints.push(WakeHint::RingReadable { + ring_id: self.identity.ring_id, + }); + self.scheduler_state + .readable_rings + .insert(self.identity.ring_id); + } + + pub fn consumer_readable(&self) -> u64 { + self.commit.saturating_sub(self.consume) + } + + pub fn consumer_try_read(&self, len: u64) -> Result, ReadError> { + if len > self.consumer_readable() { + return Err(ReadError::BeyondCommittedBytes); + } + Ok(self.consumer_read(len)) + } + + pub fn consumer_read(&self, len: u64) -> Vec { + assert!(len <= self.consumer_readable()); + (0..len) + .map(|i| self.buffer[((self.consume + i) % self.identity.capacity) as usize]) + .collect() + } + + pub fn consumer_consume(&mut self, len: u64) { + assert!(len <= self.consumer_readable()); + self.consume = self.consume.saturating_add(len); + self.wake_hints.push(WakeHint::RingWritable { + ring_id: self.identity.ring_id, + }); + self.scheduler_state + .writable_rings + .insert(self.identity.ring_id); + } + + pub fn cursor_snapshot(&self) -> CursorSnapshot { + CursorSnapshot { + commit: self.commit, + consume: self.consume, + } + } + + pub fn wake_hints(&self) -> &[WakeHint] { + &self.wake_hints + } + + pub fn coalesce_duplicate_wakes(&mut self) { + for wake in &self.wake_hints { + match wake { + WakeHint::RingReadable { ring_id } => { + self.scheduler_state.readable_rings.insert(*ring_id); + } + WakeHint::RingWritable { ring_id } => { + self.scheduler_state.writable_rings.insert(*ring_id); + } + } + } + } + + pub fn scheduler_state(&self) -> &SchedulerState { + &self.scheduler_state + } + + pub fn map_process_local_view(&self, base: ArenaBase) -> ProcessLocalView { + let layout = RingLayout { data_offset: 128 }; + ProcessLocalView { + layout, + data_pointer: ProcessLocalPointer::from_base_plus_offset(base, layout.data_offset), + } + } + + pub fn python_visible_operations(&self) -> Vec { + vec![ + PythonOperation::ReserveViaHelper { + ring_id: self.identity.ring_id, + }, + PythonOperation::CommitViaHelper { + ring_id: self.identity.ring_id, + }, + PythonOperation::ReadableViaHelper { + ring_id: self.identity.ring_id, + }, + PythonOperation::ConsumeViaHelper { + ring_id: self.identity.ring_id, + }, + PythonOperation::MapPointerViaHelper { + ring_id: self.identity.ring_id, + }, + ] + } +} diff --git a/crates/mvp-system/src/stage_controller.rs b/crates/mvp-system/src/stage_controller.rs new file mode 100644 index 0000000..b68494b --- /dev/null +++ b/crates/mvp-system/src/stage_controller.rs @@ -0,0 +1,399 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RunId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct NodeId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct EdgeId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ObjectId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct StepId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct DeviceHandle { + pub generation: u64, + pub id: u64, +} + +impl DeviceHandle { + pub fn new_current(id: u64) -> Self { + Self { generation: 1, id } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct LayerRange { + pub start: u32, + pub end_exclusive: u32, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WeightSource { + TestArtifact(String), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum EdgeDirection { + Inbound, + Outbound, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct EdgeProvision { + pub edge_id: EdgeId, + pub direction: EdgeDirection, +} + +impl EdgeProvision { + pub fn inbound(edge_id: EdgeId) -> Self { + Self { + edge_id, + direction: EdgeDirection::Inbound, + } + } + + pub fn outbound(edge_id: EdgeId) -> Self { + Self { + edge_id, + direction: EdgeDirection::Outbound, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ProvisionStage { + pub run_id: RunId, + pub authorized_orchestrator: NodeId, + pub node_id: NodeId, + pub stage_index: u32, + pub stage_count: u32, + pub layer_range: LayerRange, + pub inbound: EdgeProvision, + pub outbound: EdgeProvision, + pub weight_source: WeightSource, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum StageEvent { + ProvisionStage { + from: NodeId, + provision: ProvisionStage, + }, + WorkerReady, + WeightsReady, + InboundEdgeReady { + edge_id: EdgeId, + }, + OutboundEdgeReady { + edge_id: EdgeId, + }, + ObjectLoaded { + edge_id: EdgeId, + object_id: ObjectId, + sequence: u64, + handle: DeviceHandle, + }, + StepCompleted { + step_id: StepId, + }, + WorkerCrashed, + StopRun { + run_id: RunId, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum StageFaultReason { + UnauthorizedProvision, + SequenceViolation, + WorkerCrashed, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum StageLifecycleEvent { + StageReady { + run_id: RunId, + stage_index: u32, + }, + StepAccepted { + run_id: RunId, + stage_index: u32, + sequence: u64, + }, + StageFault { + run_id: RunId, + stage_index: u32, + reason: StageFaultReason, + }, + StageStopped { + run_id: RunId, + stage_index: u32, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct StepInput { + pub edge_id: EdgeId, + pub object_id: ObjectId, + pub sequence: u64, + pub handle: DeviceHandle, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct OutputBinding { + pub edge_id: EdgeId, + pub sequence: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ExecuteStep { + pub step_id: StepId, + pub input: StepInput, + pub outputs: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum StageCommand { + EstablishInboundEdge { + edge_id: EdgeId, + }, + EstablishOutboundEdge { + edge_id: EdgeId, + }, + ConfigureWorkerRole { + run_id: RunId, + stage_index: u32, + layer_range: LayerRange, + }, + LoadWeights { + source: WeightSource, + range: LayerRange, + }, + RewireEdge { + edge_id: EdgeId, + }, + ExecuteStep(ExecuteStep), + ReleaseInputHandle { + object_id: ObjectId, + handle: DeviceHandle, + }, + StopLocalEdges { + run_id: RunId, + }, + ReleaseRunDeviceObjects { + run_id: RunId, + }, +} + +#[cfg(test)] +pub struct StageControllerHarness { + local_node_id: NodeId, + provision: Option, + worker_ready: bool, + weights_ready: bool, + inbound_ready: bool, + outbound_ready: bool, + stage_ready_emitted: bool, + busy: bool, + expected_sequence: u64, + active_input: Option, + commands: Vec, + events: Vec, + faulted: bool, + stopped: bool, +} + +#[cfg(test)] +impl StageControllerHarness { + pub fn new(local_node_id: NodeId) -> Self { + Self { + local_node_id, + provision: None, + worker_ready: false, + weights_ready: false, + inbound_ready: false, + outbound_ready: false, + stage_ready_emitted: false, + busy: false, + expected_sequence: 0, + active_input: None, + commands: Vec::new(), + events: Vec::new(), + faulted: false, + stopped: false, + } + } + + pub fn observe(&mut self, event: StageEvent) { + match event { + StageEvent::ProvisionStage { from, provision } => self.provision(from, provision), + StageEvent::WorkerReady => self.worker_ready = true, + StageEvent::WeightsReady => self.weights_ready = true, + StageEvent::InboundEdgeReady { edge_id } => { + if self + .provision + .as_ref() + .is_some_and(|p| p.inbound.edge_id == edge_id) + { + self.inbound_ready = true; + } + } + StageEvent::OutboundEdgeReady { edge_id } => { + if self + .provision + .as_ref() + .is_some_and(|p| p.outbound.edge_id == edge_id) + { + self.outbound_ready = true; + } + } + StageEvent::ObjectLoaded { + edge_id, + object_id, + sequence, + handle, + } => self.object_loaded(edge_id, object_id, sequence, handle), + StageEvent::StepCompleted { .. } => self.step_completed(), + StageEvent::WorkerCrashed => self.fault(StageFaultReason::WorkerCrashed), + StageEvent::StopRun { run_id } => self.stop(run_id), + } + self.maybe_stage_ready(); + } + + pub fn commands(&self) -> &[StageCommand] { + &self.commands + } + + pub fn events(&self) -> &[StageLifecycleEvent] { + &self.events + } + + fn provision(&mut self, from: NodeId, provision: ProvisionStage) { + if from != provision.authorized_orchestrator || provision.node_id != self.local_node_id { + self.provision = Some(provision.clone()); + self.fault(StageFaultReason::UnauthorizedProvision); + return; + } + self.commands.push(StageCommand::EstablishInboundEdge { + edge_id: provision.inbound.edge_id, + }); + self.commands.push(StageCommand::EstablishOutboundEdge { + edge_id: provision.outbound.edge_id, + }); + self.commands.push(StageCommand::ConfigureWorkerRole { + run_id: provision.run_id, + stage_index: provision.stage_index, + layer_range: provision.layer_range, + }); + self.commands.push(StageCommand::LoadWeights { + source: provision.weight_source.clone(), + range: provision.layer_range, + }); + self.provision = Some(provision); + } + + fn maybe_stage_ready(&mut self) { + if self.stage_ready_emitted || self.faulted || self.stopped { + return; + } + if self.worker_ready && self.weights_ready && self.inbound_ready && self.outbound_ready { + if let Some(provision) = &self.provision { + self.stage_ready_emitted = true; + self.events.push(StageLifecycleEvent::StageReady { + run_id: provision.run_id, + stage_index: provision.stage_index, + }); + } + } + } + + fn object_loaded( + &mut self, + edge_id: EdgeId, + object_id: ObjectId, + sequence: u64, + handle: DeviceHandle, + ) { + if self.faulted || self.stopped || !self.stage_ready_emitted || self.busy { + return; + } + let Some(provision) = &self.provision else { + return; + }; + if edge_id != provision.inbound.edge_id { + return; + } + if sequence != self.expected_sequence { + self.fault(StageFaultReason::SequenceViolation); + return; + } + let input = StepInput { + edge_id, + object_id, + sequence, + handle, + }; + let step = ExecuteStep { + step_id: StepId(sequence), + input: input.clone(), + outputs: vec![OutputBinding { + edge_id: provision.outbound.edge_id, + sequence, + }], + }; + self.busy = true; + self.active_input = Some(input); + self.commands.push(StageCommand::ExecuteStep(step)); + self.events.push(StageLifecycleEvent::StepAccepted { + run_id: provision.run_id, + stage_index: provision.stage_index, + sequence, + }); + } + + fn step_completed(&mut self) { + if self.faulted || self.stopped || !self.busy { + return; + } + if let Some(input) = self.active_input.take() { + self.commands.push(StageCommand::ReleaseInputHandle { + object_id: input.object_id, + handle: input.handle, + }); + self.expected_sequence = input.sequence + 1; + } + self.busy = false; + } + + fn fault(&mut self, reason: StageFaultReason) { + if self.faulted || self.stopped { + return; + } + self.faulted = true; + let (run_id, stage_index) = self + .provision + .as_ref() + .map(|p| (p.run_id, p.stage_index)) + .unwrap_or((RunId(0), 0)); + self.events.push(StageLifecycleEvent::StageFault { + run_id, + stage_index, + reason, + }); + } + + fn stop(&mut self, run_id: RunId) { + if self.stopped { + return; + } + self.stopped = true; + self.commands.push(StageCommand::StopLocalEdges { run_id }); + self.commands + .push(StageCommand::ReleaseRunDeviceObjects { run_id }); + let stage_index = self.provision.as_ref().map_or(0, |p| p.stage_index); + self.events.push(StageLifecycleEvent::StageStopped { + run_id, + stage_index, + }); + } +} diff --git a/tests/mvp_system/arena_manager_guarantees.rs b/crates/mvp-system/src/tests/arena_manager_guarantees.rs similarity index 97% rename from tests/mvp_system/arena_manager_guarantees.rs rename to crates/mvp-system/src/tests/arena_manager_guarantees.rs index 1fbfb73..0bb9c77 100644 --- a/tests/mvp_system/arena_manager_guarantees.rs +++ b/crates/mvp-system/src/tests/arena_manager_guarantees.rs @@ -74,7 +74,10 @@ fn arena_boot_creates_stable_offset_only_layout_domain() { // Layout facts are arena offsets and stay under the reservation ceiling. assert!(lease.layout.start_offset < arena_config().reservation_ceiling); assert!(lease.layout.end_offset <= arena_config().reservation_ceiling); - assert!(matches!(lease.layout.pointer, arena::LayoutPointer::NoProcessPointer)); + assert!(matches!( + lease.layout.pointer, + arena::LayoutPointer::NoProcessPointer + )); // Boot failure is typed and emits no usable arena. let failed = arena::ArenaManagerHarness::boot(arena::ArenaConfig { @@ -232,9 +235,12 @@ fn release_requires_quiescence_and_reuse_happens_only_after_release() { proof: arena::QuiescenceProof::verified(), }); harness.request(arena::ArenaRequest::LeaseRing(lease_request(3, 1024, 64))); - assert!(harness.events().iter().any(|event| { - matches!(event, arena::ArenaEvent::RingLeased { .. }) - })); + assert!( + harness + .events() + .iter() + .any(|event| { matches!(event, arena::ArenaEvent::RingLeased { .. }) }) + ); } // This proves shutdown rejects new leases, preserves live lease records, and diff --git a/tests/mvp_system/device_bridge_guarantees.rs b/crates/mvp-system/src/tests/device_bridge_guarantees.rs similarity index 83% rename from tests/mvp_system/device_bridge_guarantees.rs rename to crates/mvp-system/src/tests/device_bridge_guarantees.rs index d8c6f21..aa167ed 100644 --- a/tests/mvp_system/device_bridge_guarantees.rs +++ b/crates/mvp-system/src/tests/device_bridge_guarantees.rs @@ -77,11 +77,17 @@ fn allocation_creates_current_generation_handle_or_typed_failure() { fn host_to_device_copies_exact_range_and_defers_release_until_safe() { // Allocate a device object and request an asynchronous copy. let (mut harness, handle) = allocated_bridge(); - let copy = harness.host_to_device( - device::HostRange { offset: 4, len: 8 }, - device::DeviceRange { handle, offset: 0, len: 8 }, - device::CopyMode::Async, - ).expect("copy request must be accepted"); + let copy = harness + .host_to_device( + device::HostRange { offset: 4, len: 8 }, + device::DeviceRange { + handle, + offset: 0, + len: 8, + }, + device::CopyMode::Async, + ) + .expect("copy request must be accepted"); // The backend sees the exact ranges. assert!(harness.backend_calls().iter().any(|call| { @@ -89,7 +95,11 @@ fn host_to_device_copies_exact_range_and_defers_release_until_safe() { call, device::BackendCall::HostToDevice { host: device::HostRange { offset: 4, len: 8 }, - device: device::DeviceRange { offset: 0, len: 8, .. }, + device: device::DeviceRange { + offset: 0, + len: 8, + .. + }, .. } ) @@ -110,18 +120,28 @@ fn host_to_device_copies_exact_range_and_defers_release_until_safe() { fn device_to_host_copies_exact_range_and_defers_host_validity_until_safe() { // Allocate a device object and request an asynchronous device-to-host copy. let (mut harness, handle) = allocated_bridge(); - let copy = harness.device_to_host( - device::DeviceRange { handle, offset: 0, len: 8 }, - device::HostRange { offset: 12, len: 8 }, - device::CopyMode::Async, - ).expect("copy request must be accepted"); + let copy = harness + .device_to_host( + device::DeviceRange { + handle, + offset: 0, + len: 8, + }, + device::HostRange { offset: 12, len: 8 }, + device::CopyMode::Async, + ) + .expect("copy request must be accepted"); // The backend sees the exact ranges. assert!(harness.backend_calls().iter().any(|call| { matches!( call, device::BackendCall::DeviceToHost { - device: device::DeviceRange { offset: 0, len: 8, .. }, + device: device::DeviceRange { + offset: 0, + len: 8, + .. + }, host: device::HostRange { offset: 12, len: 8 }, .. } @@ -176,11 +196,17 @@ fn lifetime_blocks_free_until_dependencies_clear_and_rejects_old_generation() { handle, step_id: device::StepId(77), }); - let copy = harness.host_to_device( - device::HostRange { offset: 0, len: 8 }, - device::DeviceRange { handle, offset: 0, len: 8 }, - device::CopyMode::Async, - ).expect("copy request must be accepted"); + let copy = harness + .host_to_device( + device::HostRange { offset: 0, len: 8 }, + device::DeviceRange { + handle, + offset: 0, + len: 8, + }, + device::CopyMode::Async, + ) + .expect("copy request must be accepted"); // Free is blocked while compute/copy depends on the allocation. assert_eq!( @@ -195,9 +221,11 @@ fn lifetime_blocks_free_until_dependencies_clear_and_rejects_old_generation() { }); harness.observe(device::DeviceEvent::CopyCompleted { copy }); assert_eq!(harness.free_device(handle), Ok(())); - assert!(harness.backend_calls().iter().any(|call| { - matches!(call, device::BackendCall::Free { freed } if *freed == handle) - })); + assert!( + harness.backend_calls().iter().any(|call| { + matches!(call, device::BackendCall::Free { freed } if *freed == handle) + }) + ); // Restart invalidates prior handles. harness.observe(device::DeviceEvent::WorkerRestarted { @@ -219,7 +247,11 @@ fn backend_copy_and_view_failures_are_typed() { let copy_failure = harness .host_to_device( device::HostRange { offset: 0, len: 8 }, - device::DeviceRange { handle, offset: 0, len: 8 }, + device::DeviceRange { + handle, + offset: 0, + len: 8, + }, device::CopyMode::Sync, ) .expect_err("copy failure must surface"); diff --git a/tests/mvp_system/driver_pumps_guarantees.rs b/crates/mvp-system/src/tests/driver_pumps_guarantees.rs similarity index 84% rename from tests/mvp_system/driver_pumps_guarantees.rs rename to crates/mvp-system/src/tests/driver_pumps_guarantees.rs index e31f106..780ff6e 100644 --- a/tests/mvp_system/driver_pumps_guarantees.rs +++ b/crates/mvp-system/src/tests/driver_pumps_guarantees.rs @@ -77,13 +77,22 @@ fn driver_owns_endpoint_connection_demux_and_pump_tasks() { // Pump tasks are driver-owned. assert!(harness.commands().iter().any(|command| { - matches!(command, driver::DriverCommand::SpawnSendPump { edge_id: driver::EdgeId(7001), .. }) + matches!( + command, + driver::DriverCommand::SpawnSendPump { + edge_id: driver::EdgeId(7001), + .. + } + ) })); // Actor commands must not expose stream polling. - assert!(!harness.actor_messages().iter().any(|message| { - matches!(message, driver::ActorMessage::PollStreamFuture { .. }) - })); + assert!( + !harness + .actor_messages() + .iter() + .any(|message| { matches!(message, driver::ActorMessage::PollStreamFuture { .. }) }) + ); } // This proves each edge uses one persistent uni-stream, writes an edge-id @@ -131,15 +140,24 @@ fn recv_rendezvous_starts_pump_only_after_spec_and_stream_exist() { // Spec first, stream second. let mut spec_first = new_driver(); spec_first.observe(driver::DriverEvent::EstablishRecv(recv_spec())); - assert!(!spec_first.commands().iter().any(|command| { - matches!(command, driver::DriverCommand::SpawnRecvPump { .. }) - })); + assert!( + !spec_first + .commands() + .iter() + .any(|command| { matches!(command, driver::DriverCommand::SpawnRecvPump { .. }) }) + ); spec_first.observe(driver::DriverEvent::IncomingUniStream { edge_id: driver::EdgeId(7001), stream_id: driver::StreamId(1), }); assert!(spec_first.commands().iter().any(|command| { - matches!(command, driver::DriverCommand::SpawnRecvPump { edge_id: driver::EdgeId(7001), .. }) + matches!( + command, + driver::DriverCommand::SpawnRecvPump { + edge_id: driver::EdgeId(7001), + .. + } + ) })); // Stream first, spec second. @@ -172,12 +190,20 @@ fn recv_pump_copies_bytes_without_parsing_and_respects_backpressure() { edge_id: driver::EdgeId(7001), bytes: driver::fake_object_header_bytes(), }); - assert!(!harness.events().iter().any(|event| { - matches!(event, driver::DriverEventOut::ObjectHeaderParsed { .. }) - })); + assert!( + !harness + .events() + .iter() + .any(|event| { matches!(event, driver::DriverEventOut::ObjectHeaderParsed { .. }) }) + ); assert!(harness.ring_commit(driver::EdgeId(7001)) > 0); assert!(harness.wake_hints().iter().any(|wake| { - matches!(wake, driver::WakeHint::RingReadable { edge_id: driver::EdgeId(7001) }) + matches!( + wake, + driver::WakeHint::RingReadable { + edge_id: driver::EdgeId(7001) + } + ) })); // With no ring space, the pump stops reading and waits for RingWritable. @@ -217,7 +243,12 @@ fn send_pump_advances_consume_only_after_write_acceptance() { }); assert_eq!(harness.ring_consume(driver::EdgeId(7001)), 7); assert!(harness.wake_hints().iter().any(|wake| { - matches!(wake, driver::WakeHint::RingWritable { edge_id: driver::EdgeId(7001) }) + matches!( + wake, + driver::WakeHint::RingWritable { + edge_id: driver::EdgeId(7001) + } + ) })); } @@ -236,7 +267,13 @@ fn driver_faults_and_stop_emit_stream_fault_or_pump_stopped() { edge_id: driver::EdgeId(7001), }); assert!(recv.events().iter().any(|event| { - matches!(event, driver::DriverEventOut::StreamFault { edge_id: driver::EdgeId(7001), .. }) + matches!( + event, + driver::DriverEventOut::StreamFault { + edge_id: driver::EdgeId(7001), + .. + } + ) })); // Write error faults the send edge. @@ -246,7 +283,13 @@ fn driver_faults_and_stop_emit_stream_fault_or_pump_stopped() { edge_id: driver::EdgeId(7001), }); assert!(send.events().iter().any(|event| { - matches!(event, driver::DriverEventOut::StreamFault { edge_id: driver::EdgeId(7001), .. }) + matches!( + event, + driver::DriverEventOut::StreamFault { + edge_id: driver::EdgeId(7001), + .. + } + ) })); // StopEdge stops the corresponding pump. @@ -254,6 +297,12 @@ fn driver_faults_and_stop_emit_stream_fault_or_pump_stopped() { edge_id: driver::EdgeId(7001), }); assert!(send.events().iter().any(|event| { - matches!(event, driver::DriverEventOut::PumpStopped { edge_id: driver::EdgeId(7001), .. }) + matches!( + event, + driver::DriverEventOut::PumpStopped { + edge_id: driver::EdgeId(7001), + .. + } + ) })); } diff --git a/tests/mvp_system/edge_establisher_guarantees.rs b/crates/mvp-system/src/tests/edge_establisher_guarantees.rs similarity index 85% rename from tests/mvp_system/edge_establisher_guarantees.rs rename to crates/mvp-system/src/tests/edge_establisher_guarantees.rs index 7c896fe..93eeb1a 100644 --- a/tests/mvp_system/edge_establisher_guarantees.rs +++ b/crates/mvp-system/src/tests/edge_establisher_guarantees.rs @@ -89,7 +89,10 @@ fn provisioning_creates_local_edge_records_without_remote_actor_addresses() { assert!(rx.local_record(edge::EdgeId(7001)).is_some()); // Data-flow provisioning must not require remote actor addresses. - for record in [tx.local_record(edge::EdgeId(7001)).unwrap(), rx.local_record(edge::EdgeId(7001)).unwrap()] { + for record in [ + tx.local_record(edge::EdgeId(7001)).unwrap(), + rx.local_record(edge::EdgeId(7001)).unwrap(), + ] { assert!(record.remote_actor_address.is_none()); } } @@ -109,9 +112,12 @@ fn lease_flow_matches_records_and_suppresses_stale_or_cancelled_leases() { }); // Mismatched lease must not install worker or pump state. - assert!(!harness.commands().iter().any(|command| { - matches!(command, edge::EdgeCommand::InstallWorkerRing { .. }) - })); + assert!( + !harness + .commands() + .iter() + .any(|command| { matches!(command, edge::EdgeCommand::InstallWorkerRing { .. }) }) + ); // A matching rejection faults the record. harness.observe(edge::EdgeEvent::RingLeaseRejected { @@ -164,9 +170,12 @@ fn worker_ring_install_precedes_driver_establishment_and_uses_provision_specs() }); // Before ring installation, no driver command may be issued. - assert!(!harness.commands().iter().any(|command| { - matches!(command, edge::EdgeCommand::EstablishSend { .. }) - })); + assert!( + !harness + .commands() + .iter() + .any(|command| { matches!(command, edge::EdgeCommand::EstablishSend { .. }) }) + ); // Matching ring installation advances establishment. harness.observe(edge::EdgeEvent::RingInstalled { @@ -188,7 +197,13 @@ fn worker_ring_install_precedes_driver_establishment_and_uses_provision_specs() // Now the driver may be established. assert!(harness.commands().iter().any(|command| { - matches!(command, edge::EdgeCommand::EstablishSend { edge_id: edge::EdgeId(7001), .. }) + matches!( + command, + edge::EdgeCommand::EstablishSend { + edge_id: edge::EdgeId(7001), + .. + } + ) })); } @@ -212,9 +227,11 @@ fn driver_ready_marks_local_edge_actor_ready() { })); // The edge is not ready until DriverEdgeReady arrives. - assert!(!tx.events().iter().any(|event| { - matches!(event, edge::EdgeLifecycleEvent::EdgeReady { .. }) - })); + assert!( + !tx.events() + .iter() + .any(|event| { matches!(event, edge::EdgeLifecycleEvent::EdgeReady { .. }) }) + ); tx.observe(edge::EdgeEvent::DriverEdgeReady { edge_id: edge::EdgeId(7001), }); @@ -230,9 +247,11 @@ fn driver_ready_marks_local_edge_actor_ready() { // After readiness, stream and pump behavior belongs to the driver; the // establisher should not emit hot-path byte commands. - assert!(!tx.commands().iter().any(|command| { - matches!(command, edge::EdgeCommand::CopyHotPathBytes { .. }) - })); + assert!( + !tx.commands() + .iter() + .any(|command| { matches!(command, edge::EdgeCommand::CopyHotPathBytes { .. }) }) + ); } // This proves StopEdge cancels queued leases, stops pumps, uninstalls worker @@ -251,20 +270,32 @@ fn stop_edge_tears_down_local_state_and_terminal_stopped_ignores_late_events() { }); // Stop commands must cover queued lease, pump, and worker ring cleanup. - assert!(harness.commands().iter().any(|command| { - matches!(command, edge::EdgeCommand::CancelQueuedLease { .. }) - })); - assert!(harness.commands().iter().any(|command| { - matches!(command, edge::EdgeCommand::StopPump { .. }) - })); - assert!(harness.commands().iter().any(|command| { - matches!(command, edge::EdgeCommand::UninstallWorkerRing { .. }) - })); + assert!( + harness + .commands() + .iter() + .any(|command| { matches!(command, edge::EdgeCommand::CancelQueuedLease { .. }) }) + ); + assert!( + harness + .commands() + .iter() + .any(|command| { matches!(command, edge::EdgeCommand::StopPump { .. }) }) + ); + assert!( + harness + .commands() + .iter() + .any(|command| { matches!(command, edge::EdgeCommand::UninstallWorkerRing { .. }) }) + ); // Arena release is withheld until quiescence proof arrives. - assert!(!harness.commands().iter().any(|command| { - matches!(command, edge::EdgeCommand::ReleaseArenaLease { .. }) - })); + assert!( + !harness + .commands() + .iter() + .any(|command| { matches!(command, edge::EdgeCommand::ReleaseArenaLease { .. }) }) + ); harness.observe(edge::EdgeEvent::QuiescenceProven { ring_id: edge::RingId(8001), }); diff --git a/tests/mvp_system/gpu_worker_ctl_guarantees.rs b/crates/mvp-system/src/tests/gpu_worker_ctl_guarantees.rs similarity index 64% rename from tests/mvp_system/gpu_worker_ctl_guarantees.rs rename to crates/mvp-system/src/tests/gpu_worker_ctl_guarantees.rs index 8a7f425..d4daefa 100644 --- a/tests/mvp_system/gpu_worker_ctl_guarantees.rs +++ b/crates/mvp-system/src/tests/gpu_worker_ctl_guarantees.rs @@ -70,24 +70,35 @@ fn worker_lifecycle_runs_start_initialize_ready_and_fault_paths() { // Start the worker process. let mut harness = new_controller(); harness.observe(ctl::WorkerCtlEvent::StartWorker); - assert!(harness.commands().iter().any(|command| { - matches!(command, ctl::WorkerCtlCommand::SpawnProcessActor { .. }) - })); + assert!( + harness + .commands() + .iter() + .any(|command| { matches!(command, ctl::WorkerCtlCommand::SpawnProcessActor { .. }) }) + ); // Process start triggers InitializeWorker. harness.observe(ctl::WorkerCtlEvent::ProcessStarted { pid: ctl::ProcessId(1234), }); - assert!(harness.serialized_worker_commands().iter().any(|command| { - matches!(command, ctl::WorkerCommand::InitializeWorker { .. }) - })); + assert!( + harness + .serialized_worker_commands() + .iter() + .any(|command| { matches!(command, ctl::WorkerCommand::InitializeWorker { .. }) }) + ); // WorkerReady emits a running lifecycle event. harness.observe(ctl::WorkerCtlEvent::WorkerReady { generation: ctl::WorkerGeneration(1), }); assert!(harness.events().iter().any(|event| { - matches!(event, ctl::WorkerCtlOut::WorkerRunning { generation: ctl::WorkerGeneration(1) }) + matches!( + event, + ctl::WorkerCtlOut::WorkerRunning { + generation: ctl::WorkerGeneration(1) + } + ) })); // Initialization timeout in a fresh controller is terminal failure. @@ -95,7 +106,13 @@ fn worker_lifecycle_runs_start_initialize_ready_and_fault_paths() { timed_out.observe(ctl::WorkerCtlEvent::StartWorker); timed_out.advance_time_ms(1_001); assert!(timed_out.events().iter().any(|event| { - matches!(event, ctl::WorkerCtlOut::WorkerFailed { reason: ctl::WorkerFailure::InitializationTimeout, .. }) + matches!( + event, + ctl::WorkerCtlOut::WorkerFailed { + reason: ctl::WorkerFailure::InitializationTimeout, + .. + } + ) })); } @@ -112,7 +129,10 @@ fn command_routing_requires_running_current_generation_and_is_payload_free() { } // All valid commands are serialized to the worker. - assert_eq!(harness.serialized_worker_commands().len(), current_generation_commands().len() + 1); + assert_eq!( + harness.serialized_worker_commands().len(), + current_generation_commands().len() + 1 + ); // Serialized commands must not contain payload bytes. for command in harness.serialized_worker_commands() { @@ -136,13 +156,21 @@ fn command_routing_requires_running_current_generation_and_is_payload_free() { harness.observe(ctl::WorkerCtlEvent::WorkerReady { generation: ctl::WorkerGeneration(2), }); - harness.observe(ctl::WorkerCtlEvent::ActorCommand(ctl::ActorCommand::ExecuteStep { - generation: ctl::WorkerGeneration(1), - step_id: ctl::StepId(9002), - input: ctl::DeviceHandle::new(ctl::WorkerGeneration(1), 42), - })); + harness.observe(ctl::WorkerCtlEvent::ActorCommand( + ctl::ActorCommand::ExecuteStep { + generation: ctl::WorkerGeneration(1), + step_id: ctl::StepId(9002), + input: ctl::DeviceHandle::new(ctl::WorkerGeneration(1), 42), + }, + )); assert!(harness.events().iter().any(|event| { - matches!(event, ctl::WorkerCtlOut::CommandRejected { reason: ctl::CommandRejection::OldGenerationHandle, .. }) + matches!( + event, + ctl::WorkerCtlOut::CommandRejected { + reason: ctl::CommandRejection::OldGenerationHandle, + .. + } + ) })); } @@ -154,31 +182,66 @@ fn parsed_worker_events_route_to_their_control_owners() { let mut harness = running_controller(); // Deliver each worker event family through stdout parsing. - harness.observe(ctl::WorkerCtlEvent::StdoutEvent(ctl::WorkerEvent::RingInstalled { - ring_id: ctl::RingId(8001), - })); - harness.observe(ctl::WorkerCtlEvent::StdoutEvent(ctl::WorkerEvent::ObjectLoaded { - object_id: ctl::ObjectId(9000), - sequence: 0, - })); - harness.observe(ctl::WorkerCtlEvent::StdoutEvent(ctl::WorkerEvent::ObjectProduced { - object_id: ctl::ObjectId(9001), - sequence: 0, - })); - harness.observe(ctl::WorkerCtlEvent::StdoutEvent(ctl::WorkerEvent::StepCompleted { - step_id: ctl::StepId(77), - })); - harness.observe(ctl::WorkerCtlEvent::StdoutEvent(ctl::WorkerEvent::RingReadable { - ring_id: ctl::RingId(8001), - })); + harness.observe(ctl::WorkerCtlEvent::StdoutEvent( + ctl::WorkerEvent::RingInstalled { + ring_id: ctl::RingId(8001), + }, + )); + harness.observe(ctl::WorkerCtlEvent::StdoutEvent( + ctl::WorkerEvent::ObjectLoaded { + object_id: ctl::ObjectId(9000), + sequence: 0, + }, + )); + harness.observe(ctl::WorkerCtlEvent::StdoutEvent( + ctl::WorkerEvent::ObjectProduced { + object_id: ctl::ObjectId(9001), + sequence: 0, + }, + )); + harness.observe(ctl::WorkerCtlEvent::StdoutEvent( + ctl::WorkerEvent::StepCompleted { + step_id: ctl::StepId(77), + }, + )); + harness.observe(ctl::WorkerCtlEvent::StdoutEvent( + ctl::WorkerEvent::RingReadable { + ring_id: ctl::RingId(8001), + }, + )); // Routing is proven by destination commands/events, not private dispatch // tables. - assert!(harness.routed().iter().any(|route| matches!(route, ctl::RoutedEvent::ToEdgeEstablisher(_)))); - assert!(harness.routed().iter().any(|route| matches!(route, ctl::RoutedEvent::ToRxOrRole(_)))); - assert!(harness.routed().iter().any(|route| matches!(route, ctl::RoutedEvent::ToTxOrRole(_)))); - assert!(harness.routed().iter().any(|route| matches!(route, ctl::RoutedEvent::ToStageController(_)))); - assert!(harness.routed().iter().any(|route| matches!(route, ctl::RoutedEvent::ToDriverOrWorkerSide(_)))); + assert!( + harness + .routed() + .iter() + .any(|route| matches!(route, ctl::RoutedEvent::ToEdgeEstablisher(_))) + ); + assert!( + harness + .routed() + .iter() + .any(|route| matches!(route, ctl::RoutedEvent::ToRxOrRole(_))) + ); + assert!( + harness + .routed() + .iter() + .any(|route| matches!(route, ctl::RoutedEvent::ToTxOrRole(_))) + ); + assert!( + harness + .routed() + .iter() + .any(|route| matches!(route, ctl::RoutedEvent::ToStageController(_))) + ); + assert!( + harness + .routed() + .iter() + .any(|route| matches!(route, ctl::RoutedEvent::ToDriverOrWorkerSide(_))) + ); } // This proves worker crash invalidates old handles, roles, rings, in-flight @@ -188,15 +251,19 @@ fn parsed_worker_events_route_to_their_control_owners() { fn crash_invalidates_generation_state_and_fans_out_faults() { // Install one ring and start one step in generation 1. let mut harness = running_controller(); - harness.observe(ctl::WorkerCtlEvent::ActorCommand(ctl::ActorCommand::InstallRing { - generation: ctl::WorkerGeneration(1), - ring_id: ctl::RingId(8001), - })); - harness.observe(ctl::WorkerCtlEvent::ActorCommand(ctl::ActorCommand::ExecuteStep { - generation: ctl::WorkerGeneration(1), - step_id: ctl::StepId(9001), - input: ctl::DeviceHandle::new(ctl::WorkerGeneration(1), 42), - })); + harness.observe(ctl::WorkerCtlEvent::ActorCommand( + ctl::ActorCommand::InstallRing { + generation: ctl::WorkerGeneration(1), + ring_id: ctl::RingId(8001), + }, + )); + harness.observe(ctl::WorkerCtlEvent::ActorCommand( + ctl::ActorCommand::ExecuteStep { + generation: ctl::WorkerGeneration(1), + step_id: ctl::StepId(9001), + input: ctl::DeviceHandle::new(ctl::WorkerGeneration(1), 42), + }, + )); // Crash the worker process. harness.observe(ctl::WorkerCtlEvent::ProcessExited { @@ -205,10 +272,22 @@ fn crash_invalidates_generation_state_and_fans_out_faults() { // Installed rings fault and affected pumps are stopped. assert!(harness.events().iter().any(|event| { - matches!(event, ctl::WorkerCtlOut::RingFaulted { ring_id: ctl::RingId(8001), .. }) + matches!( + event, + ctl::WorkerCtlOut::RingFaulted { + ring_id: ctl::RingId(8001), + .. + } + ) })); assert!(harness.commands().iter().any(|command| { - matches!(command, ctl::WorkerCtlCommand::StopDriverPump { ring_id: ctl::RingId(8001), .. }) + matches!( + command, + ctl::WorkerCtlCommand::StopDriverPump { + ring_id: ctl::RingId(8001), + .. + } + ) })); // Restart increments worker generation. @@ -229,16 +308,21 @@ fn crash_invalidates_generation_state_and_fans_out_faults() { fn graceful_shutdown_sends_worker_shutdown_and_reaches_terminal_stopped() { // Start from a running worker with an installed ring. let mut harness = running_controller(); - harness.observe(ctl::WorkerCtlEvent::ActorCommand(ctl::ActorCommand::InstallRing { - generation: ctl::WorkerGeneration(1), - ring_id: ctl::RingId(8001), - })); + harness.observe(ctl::WorkerCtlEvent::ActorCommand( + ctl::ActorCommand::InstallRing { + generation: ctl::WorkerGeneration(1), + ring_id: ctl::RingId(8001), + }, + )); // Request graceful shutdown. harness.observe(ctl::WorkerCtlEvent::ShutdownRequested); - assert!(harness.serialized_worker_commands().iter().any(|command| { - matches!(command, ctl::WorkerCommand::ShutdownWorker { .. }) - })); + assert!( + harness + .serialized_worker_commands() + .iter() + .any(|command| { matches!(command, ctl::WorkerCommand::ShutdownWorker { .. }) }) + ); // WorkerStopped must precede terminal stopped. harness.observe(ctl::WorkerCtlEvent::WorkerStopped { @@ -261,6 +345,12 @@ fn graceful_shutdown_sends_worker_shutdown_and_reaches_terminal_stopped() { // Rings are marked quiesced on graceful stop. assert!(harness.events().iter().any(|event| { - matches!(event, ctl::WorkerCtlOut::RingQuiesced { ring_id: ctl::RingId(8001), .. }) + matches!( + event, + ctl::WorkerCtlOut::RingQuiesced { + ring_id: ctl::RingId(8001), + .. + } + ) })); } diff --git a/tests/mvp_system/gpu_worker_egress_producer_guarantees.rs b/crates/mvp-system/src/tests/gpu_worker_egress_producer_guarantees.rs similarity index 79% rename from tests/mvp_system/gpu_worker_egress_producer_guarantees.rs rename to crates/mvp-system/src/tests/gpu_worker_egress_producer_guarantees.rs index 3014c45..a60246e 100644 --- a/tests/mvp_system/gpu_worker_egress_producer_guarantees.rs +++ b/crates/mvp-system/src/tests/gpu_worker_egress_producer_guarantees.rs @@ -109,7 +109,12 @@ fn header_is_written_and_committed_before_payload() { // Payload bytes are not committed before the payload copy is valid. assert_eq!(harness.committed_payload_bytes(egress::RingId(8002)), 0); assert!(harness.wake_hints().iter().any(|wake| { - matches!(wake, egress::WakeHint::RingReadable { ring_id: egress::RingId(8002) }) + matches!( + wake, + egress::WakeHint::RingReadable { + ring_id: egress::RingId(8002) + } + ) })); } @@ -156,29 +161,37 @@ fn payload_copy_is_exact_extent_and_respects_backpressure() { fn object_produced_precedes_step_completed_after_all_outputs() { // Execute a step with two outputs. let mut harness = installed_producer(); - harness.observe(egress::WorkerEgressEvent::InstallRing(egress::InstallRing { - ring_id: egress::RingId(8003), - edge_id: egress::EdgeId(7003), - port_id: egress::PortId("out2".into()), - ..egress_ring() - })); + harness.observe(egress::WorkerEgressEvent::InstallRing( + egress::InstallRing { + ring_id: egress::RingId(8003), + edge_id: egress::EdgeId(7003), + port_id: egress::PortId("out2".into()), + ..egress_ring() + }, + )); harness.observe(egress::WorkerEgressEvent::ExecuteStep { step_id: egress::StepId(77), - outputs: vec![output_binding(0, 8), egress::OutputBinding { - ring_id: egress::RingId(8003), - object_id: egress::ObjectId(9100), - sequence: 0, - extent: 8, - flags: egress::ObjectFlags::default(), - device_source: egress::DeviceHandle::new(egress::WorkerGeneration(1), 55), - }], + outputs: vec![ + output_binding(0, 8), + egress::OutputBinding { + ring_id: egress::RingId(8003), + object_id: egress::ObjectId(9100), + sequence: 0, + extent: 8, + flags: egress::ObjectFlags::default(), + device_source: egress::DeviceHandle::new(egress::WorkerGeneration(1), 55), + }, + ], }); // Produce only the first output and prove StepCompleted is still absent. harness.complete_output(egress::ObjectId(9000)); - assert!(!harness.events().iter().any(|event| { - matches!(event, egress::WorkerEgressOut::StepCompleted { .. }) - })); + assert!( + !harness + .events() + .iter() + .any(|event| { matches!(event, egress::WorkerEgressOut::StepCompleted { .. }) }) + ); // Produce the second output and complete role state update. harness.complete_output(egress::ObjectId(9100)); @@ -190,17 +203,41 @@ fn object_produced_precedes_step_completed_after_all_outputs() { let first_object_pos = harness .events() .iter() - .position(|event| matches!(event, egress::WorkerEgressOut::ObjectProduced { object_id: egress::ObjectId(9000), .. })) + .position(|event| { + matches!( + event, + egress::WorkerEgressOut::ObjectProduced { + object_id: egress::ObjectId(9000), + .. + } + ) + }) .expect("first object produced"); let second_object_pos = harness .events() .iter() - .position(|event| matches!(event, egress::WorkerEgressOut::ObjectProduced { object_id: egress::ObjectId(9100), .. })) + .position(|event| { + matches!( + event, + egress::WorkerEgressOut::ObjectProduced { + object_id: egress::ObjectId(9100), + .. + } + ) + }) .expect("second object produced"); let completed_pos = harness .events() .iter() - .position(|event| matches!(event, egress::WorkerEgressOut::StepCompleted { step_id: egress::StepId(77), .. })) + .position(|event| { + matches!( + event, + egress::WorkerEgressOut::StepCompleted { + step_id: egress::StepId(77), + .. + } + ) + }) .expect("step completed"); assert!(first_object_pos < completed_pos); assert!(second_object_pos < completed_pos); @@ -220,7 +257,13 @@ fn egress_faults_are_visible_and_suppress_success_events() { }], }); assert!(invalid_ring.events().iter().any(|event| { - matches!(event, egress::WorkerEgressOut::StepFailed { reason: egress::StepFailureReason::InvalidOutputRing, .. }) + matches!( + event, + egress::WorkerEgressOut::StepFailed { + reason: egress::StepFailureReason::InvalidOutputRing, + .. + } + ) })); // Extent violation fails the step. @@ -230,7 +273,13 @@ fn egress_faults_are_visible_and_suppress_success_events() { outputs: vec![output_binding(0, 32)], }); assert!(bad_extent.events().iter().any(|event| { - matches!(event, egress::WorkerEgressOut::StepFailed { reason: egress::StepFailureReason::OutputExtentViolation, .. }) + matches!( + event, + egress::WorkerEgressOut::StepFailed { + reason: egress::StepFailureReason::OutputExtentViolation, + .. + } + ) })); // Copy failure faults the ring or fails the step, but must not emit @@ -247,7 +296,10 @@ fn egress_faults_are_visible_and_suppress_success_events() { matches!(event, egress::WorkerEgressOut::StepFailed { .. }) || matches!(event, egress::WorkerEgressOut::RingFault { .. }) })); - assert!(!copy_failed.events().iter().any(|event| { - matches!(event, egress::WorkerEgressOut::ObjectProduced { .. }) - })); + assert!( + !copy_failed + .events() + .iter() + .any(|event| { matches!(event, egress::WorkerEgressOut::ObjectProduced { .. }) }) + ); } diff --git a/tests/mvp_system/gpu_worker_ingress_parser_guarantees.rs b/crates/mvp-system/src/tests/gpu_worker_ingress_parser_guarantees.rs similarity index 97% rename from tests/mvp_system/gpu_worker_ingress_parser_guarantees.rs rename to crates/mvp-system/src/tests/gpu_worker_ingress_parser_guarantees.rs index 2bf536e..6f815d8 100644 --- a/tests/mvp_system/gpu_worker_ingress_parser_guarantees.rs +++ b/crates/mvp-system/src/tests/gpu_worker_ingress_parser_guarantees.rs @@ -163,10 +163,7 @@ fn payload_loading_is_exact_extent_and_release_after_copy_completion() { object_id: ingress::ObjectId(9000), byte_count: 8, }); - assert_eq!( - harness.device_copy_log().last().unwrap().byte_count, - 8 - ); + assert_eq!(harness.device_copy_log().last().unwrap().byte_count, 8); assert!(harness.consume_cursor(ingress::RingId(8001)) > before_copy_complete); // EOF mid-object faults the object. @@ -208,9 +205,12 @@ fn object_loaded_requires_complete_valid_object_and_current_handle() { // No ObjectLoaded may appear before device copy completion and handle // creation. - assert!(!harness.events().iter().any(|event| { - matches!(event, ingress::WorkerIngressOut::ObjectLoaded { .. }) - })); + assert!( + !harness + .events() + .iter() + .any(|event| { matches!(event, ingress::WorkerIngressOut::ObjectLoaded { .. }) }) + ); // Complete device work. harness.observe(ingress::WorkerIngressEvent::DeviceCopyCompleted { diff --git a/tests/mvp_system/gpu_worker_process_adapter_guarantees.rs b/crates/mvp-system/src/tests/gpu_worker_process_adapter_guarantees.rs similarity index 85% rename from tests/mvp_system/gpu_worker_process_adapter_guarantees.rs rename to crates/mvp-system/src/tests/gpu_worker_process_adapter_guarantees.rs index 5a58305..cdf499e 100644 --- a/tests/mvp_system/gpu_worker_process_adapter_guarantees.rs +++ b/crates/mvp-system/src/tests/gpu_worker_process_adapter_guarantees.rs @@ -84,9 +84,12 @@ fn control_stream_is_line_framed_json_without_payload_bytes() { // Stderr alone is diagnostic and does not define lifecycle state. harness.receive_stderr_line("loading tinygrad backend"); - assert!(!harness.events().iter().any(|event| { - matches!(event, adapter::AdapterEvent::WorkerFatal { .. }) - })); + assert!( + !harness + .events() + .iter() + .any(|event| { matches!(event, adapter::AdapterEvent::WorkerFatal { .. }) }) + ); } // This proves worker initialization reads arena environment, waits for @@ -130,16 +133,24 @@ fn initialization_order_is_env_initialize_map_helper_backend_ready() { // Successful initialization emits WorkerReady. harness.receive_stdout_line(r#"{"type":"WorkerReady","generation":1}"#); assert!(harness.events().iter().any(|event| { - matches!(event, adapter::AdapterEvent::WorkerReady { generation: adapter::WorkerGeneration(1) }) + matches!( + event, + adapter::AdapterEvent::WorkerReady { + generation: adapter::WorkerGeneration(1) + } + ) })); // Initialization failure emits WorkerFatal if possible and exits non-zero. let mut failed = new_adapter(); failed.start_worker_process(); failed.inject_initialization_failure(adapter::InitializationFailure::BackendUnavailable); - assert!(failed.events().iter().any(|event| { - matches!(event, adapter::AdapterEvent::WorkerFatal { .. }) - })); + assert!( + failed + .events() + .iter() + .any(|event| { matches!(event, adapter::AdapterEvent::WorkerFatal { .. }) }) + ); assert_ne!(failed.exit_status(), Some(adapter::ExitStatus::Code(0))); } @@ -151,14 +162,26 @@ fn parsing_and_abi_errors_fault_the_worker_process() { let mut invalid_json = new_adapter(); invalid_json.receive_stdout_line("{not-json"); assert!(invalid_json.events().iter().any(|event| { - matches!(event, adapter::AdapterEvent::ProcessFault { reason: adapter::ProcessFaultReason::InvalidJson, .. }) + matches!( + event, + adapter::AdapterEvent::ProcessFault { + reason: adapter::ProcessFaultReason::InvalidJson, + .. + } + ) })); // Unknown event shape faults the adapter. let mut unknown = new_adapter(); unknown.receive_stdout_line(r#"{"type":"NotAWorkerEvent"}"#); assert!(unknown.events().iter().any(|event| { - matches!(event, adapter::AdapterEvent::ProcessFault { reason: adapter::ProcessFaultReason::UnknownEventShape, .. }) + matches!( + event, + adapter::AdapterEvent::ProcessFault { + reason: adapter::ProcessFaultReason::UnknownEventShape, + .. + } + ) })); // Unsupported helper ABI emits WorkerFatal. @@ -167,7 +190,13 @@ fn parsing_and_abi_errors_fault_the_worker_process() { helper_abi_version: adapter::HelperAbiVersion(999), }); assert!(abi.events().iter().any(|event| { - matches!(event, adapter::AdapterEvent::WorkerFatal { reason: adapter::WorkerFatalReason::UnsupportedHelperAbi, .. }) + matches!( + event, + adapter::AdapterEvent::WorkerFatal { + reason: adapter::WorkerFatalReason::UnsupportedHelperAbi, + .. + } + ) })); } @@ -186,14 +215,20 @@ fn command_discipline_rejects_payload_bearing_control_messages() { // Payload bytes in a control command are rejected at the adapter boundary. harness.send_raw_json_command(r#"{"type":"ExecuteStep","step_id":1,"payload":[1,2,3]}"#); assert!(harness.command_rejections().iter().any(|rejection| { - matches!(rejection.reason, adapter::CommandRejectionReason::PayloadBytesForbidden) + matches!( + rejection.reason, + adapter::CommandRejectionReason::PayloadBytesForbidden + ) })); // Wake hints reload cursors; they do not carry byte ranges or credits. harness.send_command(adapter::WorkerCommand::RingReadable { ring_id: adapter::RingId(8001), }); - let wake_line = harness.stdin_lines().last().expect("wake command must be written"); + let wake_line = harness + .stdin_lines() + .last() + .expect("wake command must be written"); let parsed = adapter::JsonLine::parse(wake_line).expect("wake line must parse"); assert_no_payload_bytes(&parsed); assert!(!parsed.contains_key("range") && !parsed.contains_key("credits")); diff --git a/tests/mvp_system/membership_pool_readiness_guarantees.rs b/crates/mvp-system/src/tests/membership_pool_readiness_guarantees.rs similarity index 87% rename from tests/mvp_system/membership_pool_readiness_guarantees.rs rename to crates/mvp-system/src/tests/membership_pool_readiness_guarantees.rs index fe268a5..ea20897 100644 --- a/tests/mvp_system/membership_pool_readiness_guarantees.rs +++ b/crates/mvp-system/src/tests/membership_pool_readiness_guarantees.rs @@ -34,10 +34,7 @@ fn readiness_config() -> membership::ReadinessConfig { // This helper provides every required public fact for one node. Tests use it to // build complete and deliberately incomplete pool views without observing any // private readiness bookkeeping. -fn report_node_ready( - gate: &mut membership::ReadinessGateHarness, - node_id: membership::NodeId, -) { +fn report_node_ready(gate: &mut membership::ReadinessGateHarness, node_id: membership::NodeId) { gate.observe(membership::Observation::NodeKnown { node_id }); gate.observe(membership::Observation::SwimLive { node_id }); gate.observe(membership::Observation::NodeAvailable { node_id }); @@ -88,8 +85,14 @@ fn pool_ready_requires_complete_stable_candidate_pool() { .expect("complete stable pool must emit PoolReady"); // Compare as sets so ordering is not part of the behavioral contract. - let observed = ready.iter().copied().collect::>(); - let expected = pool.iter().copied().collect::>(); + let observed = ready + .iter() + .copied() + .collect::>(); + let expected = pool + .iter() + .copied() + .collect::>(); assert_eq!(observed, expected); } @@ -142,9 +145,11 @@ fn planning_starts_only_after_pool_ready_and_stops_if_readiness_is_lost() { gate.request_run_planning(membership::RunRequest::new(membership::RunId(7))); // With no PoolReady event, there must be no planning command. - assert!(!gate.commands().iter().any(|command| { - matches!(command, membership::ReadinessCommand::StartPlanning { .. }) - })); + assert!( + !gate.commands().iter().any(|command| { + matches!(command, membership::ReadinessCommand::StartPlanning { .. }) + }) + ); // Satisfy readiness, then immediately lose a required node before plan // commit. The policy may wait or abort, but it must not commit placement. @@ -155,9 +160,11 @@ fn planning_starts_only_after_pool_ready_and_stops_if_readiness_is_lost() { gate.observe(membership::Observation::SwimLost { node_id: pool[1] }); // No plan commitment command may be emitted from an unstable pool view. - assert!(!gate.commands().iter().any(|command| { - matches!(command, membership::ReadinessCommand::CommitRunPlan { .. }) - })); + assert!( + !gate.commands().iter().any(|command| { + matches!(command, membership::ReadinessCommand::CommitRunPlan { .. }) + }) + ); } // This proves membership loss after provisioning is a run fault, not active @@ -180,16 +187,20 @@ fn required_node_loss_after_provisioning_faults_without_replacement() { gate.observe(membership::Observation::SwimLost { node_id: pool[1] }); // The run must fault with a membership reason. - assert!(gate.events().contains(&membership::ReadinessEvent::RunFaulted { - run_id: membership::RunId(7), - reason: membership::RunFaultReason::RequiredNodeLost { - node_id: pool[1], - }, - })); + assert!( + gate.events() + .contains(&membership::ReadinessEvent::RunFaulted { + run_id: membership::RunId(7), + reason: membership::RunFaultReason::RequiredNodeLost { node_id: pool[1] }, + }) + ); // Re-placement would violate the committed-plan authority boundary. assert!(!gate.commands().iter().any(|command| { - matches!(command, membership::ReadinessCommand::RecomputePlacement { .. }) + matches!( + command, + membership::ReadinessCommand::RecomputePlacement { .. } + ) })); } @@ -212,7 +223,9 @@ fn swim_observations_do_not_create_graph_assignments() { | membership::ReadinessCommand::StartPlanning { .. } | membership::ReadinessCommand::WaitForStability { .. } | membership::ReadinessCommand::AbortPendingRun { .. } => {} - membership::ReadinessCommand::AssignStage { .. } + membership::ReadinessCommand::CommitRunPlan { .. } + | membership::ReadinessCommand::RecomputePlacement { .. } + | membership::ReadinessCommand::AssignStage { .. } | membership::ReadinessCommand::AssignEdge { .. } | membership::ReadinessCommand::AssignLayerRange { .. } | membership::ReadinessCommand::AssignObjectSpec { .. } => { diff --git a/crates/mvp-system/src/tests/mod.rs b/crates/mvp-system/src/tests/mod.rs new file mode 100644 index 0000000..3f8c1c8 --- /dev/null +++ b/crates/mvp-system/src/tests/mod.rs @@ -0,0 +1,18 @@ +mod arena_manager_guarantees; +mod device_bridge_guarantees; +mod edge_establisher_guarantees; +mod gpu_worker_ctl_guarantees; +mod gpu_worker_egress_producer_guarantees; +mod gpu_worker_ingress_parser_guarantees; +mod gpu_worker_process_adapter_guarantees; +mod membership_pool_readiness_guarantees; +mod node_boot_lifecycle_guarantees; +mod observability_surface_guarantees; +mod orchestrator_run_fsm_guarantees; +mod orchestrator_token_endpoint_guarantees; +mod resource_inventory_guarantees; +mod run_plan_guarantees; +mod shared_ring_helper_abi_guarantees; +mod stage_controller_guarantees; +mod tx_rx_edge_actor_guarantees; +mod weight_lifecycle_guarantees; diff --git a/tests/mvp_system/node_boot_lifecycle_guarantees.rs b/crates/mvp-system/src/tests/node_boot_lifecycle_guarantees.rs similarity index 90% rename from tests/mvp_system/node_boot_lifecycle_guarantees.rs rename to crates/mvp-system/src/tests/node_boot_lifecycle_guarantees.rs index f045752..e463db7 100644 --- a/tests/mvp_system/node_boot_lifecycle_guarantees.rs +++ b/crates/mvp-system/src/tests/node_boot_lifecycle_guarantees.rs @@ -75,9 +75,13 @@ fn node_available_waits_for_all_required_readiness_facts() { let final_fact = facts.pop().expect("fixture has a final fact"); for fact in facts { harness.observe(fact); - assert!(!harness.events().contains(&boot::LifecycleEvent::NodeAvailable { - node_id: boot::NodeId(10), - })); + assert!( + !harness + .events() + .contains(&boot::LifecycleEvent::NodeAvailable { + node_id: boot::NodeId(10), + }) + ); } // Once the final required fact arrives, availability becomes observable. @@ -160,15 +164,22 @@ fn boot_resource_failure_faults_without_node_availability() { // The emitted fault must carry the stable reason enum for operators and // tests; logs are not part of the contract. - assert!(harness.events().contains(&boot::LifecycleEvent::NodeFaulted { - node_id: boot::NodeId(10), - kind: expected_kind, - })); + assert!( + harness + .events() + .contains(&boot::LifecycleEvent::NodeFaulted { + node_id: boot::NodeId(10), + kind: expected_kind, + }) + ); // A faulted boot attempt must not also become available. - assert!(!harness.events().iter().any(|event| { - matches!(event, boot::LifecycleEvent::NodeAvailable { .. }) - })); + assert!( + !harness + .events() + .iter() + .any(|event| { matches!(event, boot::LifecycleEvent::NodeAvailable { .. }) }) + ); // The orchestrator-facing eligibility check must agree with the // lifecycle transcript. @@ -193,7 +204,10 @@ fn boot_never_self_assigns_run_topology() { boot::BootCommand::AdvertiseLifecycle { .. } | boot::BootCommand::JoinMembership { .. } | boot::BootCommand::OpenProvisioningInbox { .. } => {} - boot::BootCommand::AssignStage { .. } + boot::BootCommand::LoadWeights { .. } + | boot::BootCommand::ConfigureRole { .. } + | boot::BootCommand::EstablishEdge { .. } + | boot::BootCommand::AssignStage { .. } | boot::BootCommand::AssignLayerRange { .. } | boot::BootCommand::AssignEdge { .. } | boot::BootCommand::AssignObjectSpec { .. } => { diff --git a/tests/mvp_system/observability_surface_guarantees.rs b/crates/mvp-system/src/tests/observability_surface_guarantees.rs similarity index 92% rename from tests/mvp_system/observability_surface_guarantees.rs rename to crates/mvp-system/src/tests/observability_surface_guarantees.rs index cf0c27f..25e8a73 100644 --- a/tests/mvp_system/observability_surface_guarantees.rs +++ b/crates/mvp-system/src/tests/observability_surface_guarantees.rs @@ -54,7 +54,10 @@ fn fault_trace() -> Vec { obs::FaultReason::WorkerCrashed, obs::Component::StageController, ) - .run_faulted(obs::FaultReason::WorkerCrashed, obs::Component::StageController) + .run_faulted( + obs::FaultReason::WorkerCrashed, + obs::Component::StageController, + ) .stop_run_sent(obs::StageIndex(0)) .stage_stopped(obs::StageIndex(0)) .run_torn_down() @@ -94,13 +97,19 @@ fn assert_required_identity(event: &obs::Event) { sequence, .. } => { - assert!([obs::ObjectId(9000), obs::ObjectId(9001), obs::ObjectId(9002)].contains(object_id)); + assert!( + [ + obs::ObjectId(9000), + obs::ObjectId(9001), + obs::ObjectId(9002) + ] + .contains(object_id) + ); assert_eq!(*sequence, obs::Sequence(0)); } obs::Event::StepScoped { step_id, .. } => assert_eq!(*step_id, obs::StepId(77)), obs::Event::WorkerScoped { - worker_generation, - .. + worker_generation, .. } => assert_eq!(*worker_generation, obs::WorkerGeneration(1)), } } @@ -166,7 +175,10 @@ fn lifecycle_events_cover_successful_run_milestones() { obs::EventKind::RunTornDown, ]; for kind in required { - assert!(observed.contains(&kind), "missing lifecycle event: {kind:?}"); + assert!( + observed.contains(&kind), + "missing lifecycle event: {kind:?}" + ); } } @@ -242,7 +254,10 @@ fn event_ordering_reflects_component_contracts_and_one_terminal_outcome() { let terminal_count = events .iter() .filter(|event| { - matches!(event.kind(), obs::EventKind::RunCompleted | obs::EventKind::RunFaulted) + matches!( + event.kind(), + obs::EventKind::RunCompleted | obs::EventKind::RunFaulted + ) }) .count(); assert_eq!(terminal_count, 1); @@ -253,8 +268,10 @@ fn event_ordering_reflects_component_contracts_and_one_terminal_outcome() { #[test] fn event_contract_survives_transport_storage_and_batching_policy() { // Build the same logical events under two batching policies. - let unbatched = obs::EventSubscriberHarness::collect(successful_run_trace(), obs::Batching::None); - let batched = obs::EventSubscriberHarness::collect(successful_run_trace(), obs::Batching::Fixed(8)); + let unbatched = + obs::EventSubscriberHarness::collect(successful_run_trace(), obs::Batching::None); + let batched = + obs::EventSubscriberHarness::collect(successful_run_trace(), obs::Batching::Fixed(8)); // Flattened public event facts must match as an ordered stream. let unbatched_kinds = unbatched diff --git a/tests/mvp_system/orchestrator_run_fsm_guarantees.rs b/crates/mvp-system/src/tests/orchestrator_run_fsm_guarantees.rs similarity index 93% rename from tests/mvp_system/orchestrator_run_fsm_guarantees.rs rename to crates/mvp-system/src/tests/orchestrator_run_fsm_guarantees.rs index 750aa2e..9a5fedb 100644 --- a/tests/mvp_system/orchestrator_run_fsm_guarantees.rs +++ b/crates/mvp-system/src/tests/orchestrator_run_fsm_guarantees.rs @@ -77,9 +77,12 @@ fn planning_and_provisioning_start_only_after_pool_ready() { harness.observe(fsm::RunEvent::PlanAvailable(plan.clone())); // Without PoolReady, provisioning must not begin. - assert!(!harness.commands().iter().any(|command| { - matches!(command, fsm::RunCommand::ProvisionStage { .. }) - })); + assert!( + !harness + .commands() + .iter() + .any(|command| { matches!(command, fsm::RunCommand::ProvisionStage { .. }) }) + ); // Once PoolReady is observed, the committed plan may be provisioned. harness.observe(fsm::RunEvent::PoolReady { @@ -103,7 +106,10 @@ fn planning_and_provisioning_start_only_after_pool_ready() { assert_eq!(provisioned, expected); // Provisioning must not mention nodes outside the committed plan. - let plan_nodes = plan.stage_nodes().into_iter().collect::>(); + let plan_nodes = plan + .stage_nodes() + .into_iter() + .collect::>(); for command in harness.commands() { if let fsm::RunCommand::ProvisionStage { provision } = command { assert!(plan_nodes.contains(&provision.node_id)); @@ -143,9 +149,12 @@ fn readiness_barrier_controls_prompt_injection() { }); harness.observe(fsm::RunEvent::TokenInEndpointReady); harness.observe(fsm::RunEvent::TokenOutEndpointReady); - assert!(!harness.commands().iter().any(|command| { - matches!(command, fsm::RunCommand::InjectPrompt { .. }) - })); + assert!( + !harness + .commands() + .iter() + .any(|command| { matches!(command, fsm::RunCommand::InjectPrompt { .. }) }) + ); // Complete the remaining stage readiness facts. for event in stage_ready_events(&plan).into_iter().skip(1) { @@ -198,9 +207,12 @@ fn execution_injects_next_sequence_only_after_consuming_previous_token() { } // No separate broadcast start command may exist alongside prompt injection. - assert!(!harness.commands().iter().any(|command| { - matches!(command, fsm::RunCommand::BroadcastStart { .. }) - })); + assert!( + !harness + .commands() + .iter() + .any(|command| { matches!(command, fsm::RunCommand::BroadcastStart { .. }) }) + ); // Sequence 0 must be injected first. assert_eq!(harness.injected_sequences(), vec![0]); @@ -341,9 +353,12 @@ fn terminal_outcome_is_single_and_requires_teardown() { .map(|stage| stage.stage_index) .collect::>(); assert_eq!(stopped_stages, expected_stages); - assert!(harness.commands().iter().any(|command| { - matches!(command, fsm::RunCommand::TearDownTokenEndpoints { .. }) - })); + assert!( + harness + .commands() + .iter() + .any(|command| { matches!(command, fsm::RunCommand::TearDownTokenEndpoints { .. }) }) + ); } // This proves run_torn_down is emitted exactly once and only after teardown @@ -369,9 +384,12 @@ fn run_torn_down_is_emitted_once_after_teardown_terminal_state() { run_id: fsm::RunId(7), stage_index: 0, }); - assert!(!harness.events().iter().any(|event| { - matches!(event, fsm::LifecycleEvent::RunTornDown { .. }) - })); + assert!( + !harness + .events() + .iter() + .any(|event| { matches!(event, fsm::LifecycleEvent::RunTornDown { .. }) }) + ); // Finish teardown through remaining stopped events and local endpoint stop. harness.observe(fsm::RunEvent::StageStopped { diff --git a/tests/mvp_system/orchestrator_token_endpoint_guarantees.rs b/crates/mvp-system/src/tests/orchestrator_token_endpoint_guarantees.rs similarity index 91% rename from tests/mvp_system/orchestrator_token_endpoint_guarantees.rs rename to crates/mvp-system/src/tests/orchestrator_token_endpoint_guarantees.rs index c4b555c..416e34c 100644 --- a/tests/mvp_system/orchestrator_token_endpoint_guarantees.rs +++ b/crates/mvp-system/src/tests/orchestrator_token_endpoint_guarantees.rs @@ -18,8 +18,16 @@ fn token_plan() -> token::TokenEndpointPlan { token::TokenEndpointPlan { run_id: token::RunId(7), orchestrator_node_id: token::NodeId(99), - token_in_edge: token::EdgePlan::token_in(token::EdgeId(7000), token::NodeId(99), token::NodeId(10)), - token_out_edge: token::EdgePlan::token_out(token::EdgeId(7003), token::NodeId(12), token::NodeId(99)), + token_in_edge: token::EdgePlan::token_in( + token::EdgeId(7000), + token::NodeId(99), + token::NodeId(10), + ), + token_out_edge: token::EdgePlan::token_out( + token::EdgeId(7003), + token::NodeId(12), + token::NodeId(99), + ), token_spec: token::ObjectSpec::test_tokens(), max_tokens: 4, } @@ -95,9 +103,12 @@ fn prompt_injection_is_barrier_gated_sequence_zero_token_object() { harness.request_prompt_injection(vec![101, 102, 103]); // No prompt object may be written before the barrier. - assert!(!harness.commands().iter().any(|command| { - matches!(command, token::EndpointCommand::WriteTokenObject { .. }) - })); + assert!( + !harness + .commands() + .iter() + .any(|command| { matches!(command, token::EndpointCommand::WriteTokenObject { .. }) }) + ); // Passing the global barrier permits the first prompt write. pass_global_barrier(&mut harness); @@ -117,9 +128,12 @@ fn prompt_injection_is_barrier_gated_sequence_zero_token_object() { // Prompt injection is the start signal; there must not be a separate // broadcast start command. - assert!(!harness.commands().iter().any(|command| { - matches!(command, token::EndpointCommand::BroadcastStart { .. }) - })); + assert!( + !harness + .commands() + .iter() + .any(|command| { matches!(command, token::EndpointCommand::BroadcastStart { .. }) }) + ); } // This proves token-out is consumed in sequence order and sequence k + 1 is @@ -242,7 +256,10 @@ fn token_endpoint_failures_fault_the_run_or_teardown() { malformed.observe(token::EndpointEvent::TeardownFailed { edge_id: token::EdgeId(7003), }); - assert!(malformed.events().iter().any(|event| { - matches!(event, token::EndpointLifecycleEvent::TeardownFailed { .. }) - })); + assert!( + malformed + .events() + .iter() + .any(|event| { matches!(event, token::EndpointLifecycleEvent::TeardownFailed { .. }) }) + ); } diff --git a/tests/mvp_system/resource_inventory_guarantees.rs b/crates/mvp-system/src/tests/resource_inventory_guarantees.rs similarity index 96% rename from tests/mvp_system/resource_inventory_guarantees.rs rename to crates/mvp-system/src/tests/resource_inventory_guarantees.rs index 72bc72f..7d47ce8 100644 --- a/tests/mvp_system/resource_inventory_guarantees.rs +++ b/crates/mvp-system/src/tests/resource_inventory_guarantees.rs @@ -91,7 +91,10 @@ fn inventory_entries_are_known_before_planning_and_not_negotiated_by_nodes() { .collect::>(); assert_eq!(after_nodes, before_nodes); assert!(!harness.commands().iter().any(|command| { - matches!(command, inventory::InventoryCommand::AcceptPlacementNegotiation { .. }) + matches!( + command, + inventory::InventoryCommand::AcceptPlacementNegotiation { .. } + ) })); } @@ -224,7 +227,8 @@ fn planning_preserves_candidate_pool_and_emits_no_hidden_nodes() { #[test] fn provisioned_stages_do_not_reinterpret_inventory() { // Commit a plan and provision its stages. - let plan = inventory::plan_from_inventory(planning_request()).expect("valid inventory must plan"); + let plan = + inventory::plan_from_inventory(planning_request()).expect("valid inventory must plan"); let mut harness = inventory::InventoryHarness::new(inventory_entries()); harness.commit_plan(plan.clone()); harness.provision_stages(); @@ -238,7 +242,10 @@ fn provisioned_stages_do_not_reinterpret_inventory() { // No provisioned stage may be asked to reinterpret its assignment. assert!(!harness.commands().iter().any(|command| { - matches!(command, inventory::InventoryCommand::RewriteStagePlacement { .. }) + matches!( + command, + inventory::InventoryCommand::RewriteStagePlacement { .. } + ) })); // The committed plan remains the only graph-visible placement fact. diff --git a/tests/mvp_system/run_plan_guarantees.rs b/crates/mvp-system/src/tests/run_plan_guarantees.rs similarity index 99% rename from tests/mvp_system/run_plan_guarantees.rs rename to crates/mvp-system/src/tests/run_plan_guarantees.rs index 992f7fe..935bb41 100644 --- a/tests/mvp_system/run_plan_guarantees.rs +++ b/crates/mvp-system/src/tests/run_plan_guarantees.rs @@ -20,7 +20,7 @@ type EdgeKind = plan::EdgeKind; type EdgePlan = plan::EdgePlan; type InboundEdgeProvision = plan::InboundEdgeProvision; type ModelFacts = plan::ModelFacts; -type NodeId = plan::NodeId; +use plan::NodeId; type ObjectKind = plan::ObjectKind; type OutboundEdgeProvision = plan::OutboundEdgeProvision; type PlacementInput = plan::PlacementInput; @@ -283,7 +283,10 @@ fn edge_graph_is_exactly_the_linear_pipeline() { .filter(|edge| edge.kind == EdgeKind::TokenOut) .collect::>(); assert_eq!(token_out.len(), 1); - assert_eq!(edge_stage_index(&token_out[0].producer), Some(stage_count - 1)); + assert_eq!( + edge_stage_index(&token_out[0].producer), + Some(stage_count - 1) + ); assert!(matches!( token_out[0].consumer, EdgeEndpoint::Orchestrator { node_id } if node_id == node(99) diff --git a/tests/mvp_system/shared_ring_helper_abi_guarantees.rs b/crates/mvp-system/src/tests/shared_ring_helper_abi_guarantees.rs similarity index 92% rename from tests/mvp_system/shared_ring_helper_abi_guarantees.rs rename to crates/mvp-system/src/tests/shared_ring_helper_abi_guarantees.rs index 01f62b4..5da98f2 100644 --- a/tests/mvp_system/shared_ring_helper_abi_guarantees.rs +++ b/crates/mvp-system/src/tests/shared_ring_helper_abi_guarantees.rs @@ -37,8 +37,7 @@ fn drain_committed(helper: &mut ring::RingHelperHarness) -> Vec { // depending on scheduler internals. fn assert_wake_is_payload_free(wake: &ring::WakeHint) { match wake { - ring::WakeHint::RingReadable { ring_id } - | ring::WakeHint::RingWritable { ring_id } => { + ring::WakeHint::RingReadable { ring_id } | ring::WakeHint::RingWritable { ring_id } => { assert_eq!(*ring_id, ring::RingId(7000)); } } @@ -85,11 +84,16 @@ fn cursors_are_monotonic_and_wrap_by_modulo_capacity() { assert_eq!(helper.cursor_snapshot().consume, 6); // Wrap the physical index while logical cursors keep increasing. - let wrapped = helper.producer_reserve(5).expect("space must exist after consume"); + let wrapped = helper + .producer_reserve(5) + .expect("space must exist after consume"); helper.producer_write(&wrapped, b"ghijk"); helper.producer_commit(wrapped); assert_eq!(helper.cursor_snapshot().commit, 11); - assert_eq!(helper.cursor_snapshot().commit % helper.identity().capacity, 3); + assert_eq!( + helper.cursor_snapshot().commit % helper.identity().capacity, + 3 + ); assert_eq!(drain_committed(&mut helper), b"ghijk"); assert_eq!(helper.cursor_snapshot().consume, 11); } @@ -101,7 +105,9 @@ fn cursors_are_monotonic_and_wrap_by_modulo_capacity() { fn producer_respects_free_space_and_publishes_after_writing() { // Reserve the full ring and publish it. let mut helper = new_ring(); - let reservation = helper.producer_reserve(8).expect("full ring reservation fits"); + let reservation = helper + .producer_reserve(8) + .expect("full ring reservation fits"); helper.producer_write(&reservation, b"12345678"); helper.producer_commit(reservation); @@ -170,12 +176,22 @@ fn wake_hints_are_edge_hints_without_hiding_transitions() { // The readable transition must be discoverable even if duplicate wakes are // coalesced. helper.coalesce_duplicate_wakes(); - assert!(helper.scheduler_state().readable_rings.contains(&ring::RingId(7000))); + assert!( + helper + .scheduler_state() + .readable_rings + .contains(&ring::RingId(7000)) + ); // Fill then release space to cause full-to-writable. let _ = drain_committed(&mut helper); helper.coalesce_duplicate_wakes(); - assert!(helper.scheduler_state().writable_rings.contains(&ring::RingId(7000))); + assert!( + helper + .scheduler_state() + .writable_rings + .contains(&ring::RingId(7000)) + ); // Every wake remains a payload-free hint. for wake in helper.wake_hints() { diff --git a/tests/mvp_system/stage_controller_guarantees.rs b/crates/mvp-system/src/tests/stage_controller_guarantees.rs similarity index 90% rename from tests/mvp_system/stage_controller_guarantees.rs rename to crates/mvp-system/src/tests/stage_controller_guarantees.rs index 0ce241b..10970ee 100644 --- a/tests/mvp_system/stage_controller_guarantees.rs +++ b/crates/mvp-system/src/tests/stage_controller_guarantees.rs @@ -101,9 +101,12 @@ fn provisioning_validates_authority_and_assigned_shape_before_setup() { // The controller must not emit any command that replaces the provisioned // edge ids with a locally chosen edge. - assert!(!harness.commands().iter().any(|command| { - matches!(command, stage::StageCommand::RewireEdge { .. }) - })); + assert!( + !harness + .commands() + .iter() + .any(|command| { matches!(command, stage::StageCommand::RewireEdge { .. }) }) + ); // An unauthorized provision attempt must fault before setup can begin. let mut unauthorized = new_controller(); @@ -120,9 +123,12 @@ fn provisioning_validates_authority_and_assigned_shape_before_setup() { } ) })); - assert!(!unauthorized.commands().iter().any(|command| { - matches!(command, stage::StageCommand::ConfigureWorkerRole { .. }) - })); + assert!( + !unauthorized + .commands() + .iter() + .any(|command| { matches!(command, stage::StageCommand::ConfigureWorkerRole { .. }) }) + ); } // This proves StageReady is emitted only after worker readiness, weight @@ -143,9 +149,12 @@ fn stage_ready_waits_for_worker_weights_and_both_edges() { let final_event = events.pop().expect("fixture has final setup event"); for event in events { harness.observe(event); - assert!(!harness.events().iter().any(|event| { - matches!(event, stage::StageLifecycleEvent::StageReady { .. }) - })); + assert!( + !harness + .events() + .iter() + .any(|event| { matches!(event, stage::StageLifecycleEvent::StageReady { .. }) }) + ); } // The final prerequisite crosses the barrier. @@ -220,11 +229,7 @@ fn accepted_inbound_object_creates_one_same_sequence_execute_step() { #[test] fn duplicate_skipped_and_out_of_order_sequences_fault() { // Each invalid trace starts from a freshly readied stage. - let invalid_traces = vec![ - vec![0, 0], - vec![0, 2], - vec![0, 1, 0], - ]; + let invalid_traces = vec![vec![0, 0], vec![0, 2], vec![0, 1, 0]]; for trace in invalid_traces { // Accept the first object and complete its step when needed so the next @@ -272,7 +277,10 @@ fn step_completed_releases_input_and_admits_next_object() { // Before worker completion, no compute-complete lifecycle event is allowed. assert!(!harness.events().iter().any(|event| { - matches!(event, stage::StageLifecycleEvent::StepAccepted { sequence: 1, .. }) + matches!( + event, + stage::StageLifecycleEvent::StepAccepted { sequence: 1, .. } + ) })); // Worker StepCompleted is the public completion signal. @@ -348,13 +356,21 @@ fn stage_fault_rejects_new_work_until_stopped() { harness.observe(stage::StageEvent::StopRun { run_id: stage::RunId(7), }); - assert!(harness.commands().iter().any(|command| { - matches!(command, stage::StageCommand::StopLocalEdges { .. }) - })); - assert!(harness.commands().iter().any(|command| { - matches!(command, stage::StageCommand::ReleaseRunDeviceObjects { .. }) - })); - assert!(harness.events().iter().any(|event| { - matches!(event, stage::StageLifecycleEvent::StageStopped { .. }) - })); + assert!( + harness + .commands() + .iter() + .any(|command| { matches!(command, stage::StageCommand::StopLocalEdges { .. }) }) + ); + assert!( + harness.commands().iter().any(|command| { + matches!(command, stage::StageCommand::ReleaseRunDeviceObjects { .. }) + }) + ); + assert!( + harness + .events() + .iter() + .any(|event| { matches!(event, stage::StageLifecycleEvent::StageStopped { .. }) }) + ); } diff --git a/crates/mvp-system/src/tests/tinygrad_device_bridge_backend.py b/crates/mvp-system/src/tests/tinygrad_device_bridge_backend.py new file mode 100755 index 0000000..aa49777 --- /dev/null +++ b/crates/mvp-system/src/tests/tinygrad_device_bridge_backend.py @@ -0,0 +1,305 @@ +#!/usr/bin/env python3 +"""Tinygrad-backed device bridge probe for Rust MVP bridge tests. + +Line-delimited JSON control only. Payload bytes live in the arena file whose +path Rust passes during initialize. +""" + +from __future__ import annotations + +import json +import os +import struct +import sys +from dataclasses import dataclass +from typing import Any + +try: + from tinygrad import Tensor # type: ignore +except Exception as exc: # pragma: no cover - exercised from Rust process tests + print( + json.dumps( + { + "type": "worker_fatal", + "reason": "tinygrad_unavailable", + "message": str(exc), + } + ), + flush=True, + ) + raise SystemExit(2) + + +@dataclass +class DeviceObject: + dtype: str + shape: str + extent: int + values: list[int] + tensor: Any + + +@dataclass +class PendingCopy: + kind: str + copy_id: int + handle_id: int + host_offset: int + device_offset: int + length: int + + +arena_path: str | None = None +arena_bytes = 0 +generation = 0 +objects: dict[int, DeviceObject] = {} +pending_copies: dict[int, PendingCopy] = {} +fail_next: str | None = None + + +def emit(obj: dict[str, Any]) -> None: + print(json.dumps(obj, separators=(",", ":")), flush=True) + + +def backend_error(reason: str) -> None: + emit({"type": "backend_error", "reason": reason}) + + +def fatal(reason: str, message: str) -> None: + emit({"type": "worker_fatal", "reason": reason, "message": message}) + + +def require_arena() -> str: + if arena_path is None: + raise RuntimeError("arena is not initialized") + return arena_path + + +def consume_failure(reason: str) -> bool: + global fail_next + if fail_next == reason: + fail_next = None + backend_error(reason) + return True + return False + + +def check_u32_range(offset: int, length: int) -> None: + if offset < 0 or length < 0 or offset % 4 != 0 or length % 4 != 0: + raise ValueError("u32 ranges must be non-negative and 4-byte aligned") + + +def tensor_from_values(values: list[int]) -> Any: + try: + return Tensor(values, dtype="uint32").realize() + except Exception: + return Tensor(values, dtype="int32").realize() + + +def allocate_tensor(dtype: str, extent: int) -> tuple[list[int], Any]: + if dtype == "u32": + if extent % 4 != 0: + raise ValueError("u32 extent must be 4-byte aligned") + values = [0] * (extent // 4) + return values, tensor_from_values(values) + if dtype == "f16": + values = [0] * (extent // 2) + return values, Tensor(values, dtype="float16").realize() + raise ValueError(f"unsupported dtype {dtype!r}") + + +def update_tensor(obj: DeviceObject) -> None: + if obj.dtype == "u32": + obj.tensor = tensor_from_values(obj.values) + elif obj.dtype == "f16": + obj.tensor = Tensor(obj.values, dtype="float16").realize() + else: + raise ValueError(f"unsupported dtype {obj.dtype!r}") + + +def materialized_values(obj: DeviceObject) -> list[int]: + return [int(v) for v in obj.tensor.tolist()] + + +def read_arena(offset: int, length: int) -> bytes: + path = require_arena() + with open(path, "rb", buffering=0) as f: + f.seek(offset) + data = f.read(length) + if len(data) != length: + raise EOFError("short arena read") + return data + + +def write_arena(offset: int, data: bytes) -> None: + path = require_arena() + with open(path, "r+b", buffering=0) as f: + f.seek(offset) + f.write(data) + f.flush() + + +def perform_host_to_device(copy: PendingCopy) -> None: + obj = objects[copy.handle_id] + if obj.dtype != "u32": + raise ValueError("payload copy is implemented for u32 test objects only") + check_u32_range(copy.device_offset, copy.length) + payload = read_arena(copy.host_offset, copy.length) + words = list(struct.unpack("<" + "I" * (copy.length // 4), payload)) + start = copy.device_offset // 4 + end = start + len(words) + if end > len(obj.values): + raise ValueError("device range out of bounds") + obj.values[start:end] = words + update_tensor(obj) + + +def perform_device_to_host(copy: PendingCopy) -> None: + obj = objects[copy.handle_id] + if obj.dtype != "u32": + raise ValueError("payload copy is implemented for u32 test objects only") + check_u32_range(copy.device_offset, copy.length) + values = materialized_values(obj) + start = copy.device_offset // 4 + end = start + (copy.length // 4) + if end > len(values): + raise ValueError("device range out of bounds") + payload = struct.pack("<" + "I" * (end - start), *values[start:end]) + write_arena(copy.host_offset, payload) + + +def perform_copy(copy: PendingCopy) -> None: + if copy.kind == "host_to_device": + perform_host_to_device(copy) + elif copy.kind == "device_to_host": + perform_device_to_host(copy) + else: + raise ValueError(f"unknown copy kind {copy.kind!r}") + + +def handle(req: dict[str, Any]) -> bool: + global arena_path, arena_bytes, generation, fail_next + + typ = req.get("type") + if typ == "initialize": + arena_path = str(req["arena_path"]) + arena_bytes = int(req["arena_bytes"]) + generation = int(req["generation"]) + with open(arena_path, "r+b", buffering=0) as f: + f.truncate(arena_bytes) + emit({"type": "worker_ready", "generation": generation}) + return True + + if typ == "alloc": + if consume_failure("allocation_failed"): + return True + handle_id = int(req["handle_id"]) + dtype = str(req["dtype"]) + shape = str(req["shape"]) + extent = int(req["extent"]) + values, tensor = allocate_tensor(dtype, extent) + objects[handle_id] = DeviceObject(dtype=dtype, shape=shape, extent=extent, values=values, tensor=tensor) + emit({"type": "allocated", "handle_id": handle_id}) + return True + + if typ in ("host_to_device", "device_to_host"): + if consume_failure("copy_failed"): + return True + copy = PendingCopy( + kind=typ, + copy_id=int(req["copy_id"]), + handle_id=int(req["handle_id"]), + host_offset=int(req["host_offset"]), + device_offset=int(req["device_offset"]), + length=int(req["len"]), + ) + if copy.handle_id not in objects: + backend_error("copy_failed") + return True + if bool(req.get("defer", False)): + pending_copies[copy.copy_id] = copy + emit({"type": "copy_started", "copy_id": copy.copy_id}) + else: + perform_copy(copy) + emit({"type": "copy_completed", "copy_id": copy.copy_id}) + return True + + if typ == "complete_copy": + copy_id = int(req["copy_id"]) + copy = pending_copies.pop(copy_id, None) + if copy is None: + backend_error("copy_failed") + return True + perform_copy(copy) + emit({"type": "copy_completed", "copy_id": copy_id}) + return True + + if typ == "wrap_for_tinygrad": + if consume_failure("invalid_view"): + return True + handle_id = int(req["handle_id"]) + obj = objects.get(handle_id) + if obj is None: + backend_error("invalid_view") + return True + dtype = str(req["dtype"]) + shape = str(req["shape"]) + if obj.dtype != dtype or obj.shape != shape: + backend_error("invalid_view") + return True + # Force materialization at view time so success depends on live tensor state. + _ = obj.tensor.tolist() + emit({"type": "view", "handle_id": handle_id, "dtype": dtype, "shape": shape}) + return True + + if typ == "free": + handle_id = int(req["handle_id"]) + if handle_id not in objects: + backend_error("invalid_view") + return True + del objects[handle_id] + emit({"type": "freed", "handle_id": handle_id}) + return True + + if typ == "restart": + generation = int(req["generation"]) + objects.clear() + pending_copies.clear() + fail_next = None + emit({"type": "worker_ready", "generation": generation}) + return True + + if typ == "fail_next": + fail_next = str(req["failure"]) + emit({"type": "ok"}) + return True + + if typ == "shutdown": + emit({"type": "worker_stopped"}) + return False + + fatal("protocol_error", f"unknown command type {typ!r}") + return False + + +def main() -> int: + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + req = json.loads(line) + if not isinstance(req, dict): + raise ValueError("request must be an object") + if not handle(req): + return 0 + except SystemExit: + raise + except Exception as exc: + fatal("backend_exception", str(exc)) + return 2 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/mvp_system/tx_rx_edge_actor_guarantees.rs b/crates/mvp-system/src/tests/tx_rx_edge_actor_guarantees.rs similarity index 90% rename from tests/mvp_system/tx_rx_edge_actor_guarantees.rs rename to crates/mvp-system/src/tests/tx_rx_edge_actor_guarantees.rs index d64733b..9993658 100644 --- a/tests/mvp_system/tx_rx_edge_actor_guarantees.rs +++ b/crates/mvp-system/src/tests/tx_rx_edge_actor_guarantees.rs @@ -82,8 +82,16 @@ fn edge_actor_messages_are_lifecycle_identity_and_handle_only() { } // Both actors remain tied to exactly one edge id. - assert!(tx.messages().iter().all(|message| message.edge_id() == edge_id())); - assert!(rx.messages().iter().all(|message| message.edge_id() == edge_id())); + assert!( + tx.messages() + .iter() + .all(|message| message.edge_id() == edge_id()) + ); + assert!( + rx.messages() + .iter() + .all(|message| message.edge_id() == edge_id()) + ); } // This proves Tx starts in provisioning, becomes ready only after EdgeReady, @@ -98,9 +106,11 @@ fn tx_lifecycle_gates_production_and_faults_on_stream_or_object_failure() { object_id: edge_actor::ObjectId(9000), sequence: 0, }); - assert!(!tx.messages().iter().any(|message| { - matches!(message, edge_actor::ActorMessage::ObjectIdentity { .. }) - })); + assert!( + !tx.messages() + .iter() + .any(|message| { matches!(message, edge_actor::ActorMessage::ObjectIdentity { .. }) }) + ); // EdgeReady admits production. tx.observe(edge_actor::TxEvent::EdgeReady { edge_id: edge_id() }); @@ -154,9 +164,11 @@ fn rx_lifecycle_gates_loaded_objects_and_faults_on_stream_or_object_failure() { sequence: 0, handle: edge_actor::OpaqueHandle::new(42), }); - assert!(!rx.messages().iter().any(|message| { - matches!(message, edge_actor::ActorMessage::OpaqueHandle { .. }) - })); + assert!( + !rx.messages() + .iter() + .any(|message| { matches!(message, edge_actor::ActorMessage::OpaqueHandle { .. }) }) + ); // EdgeReady admits ObjectLoaded exposure. rx.observe(edge_actor::RxEvent::EdgeReady { edge_id: edge_id() }); @@ -223,12 +235,16 @@ fn stop_and_mismatched_edge_events_do_not_create_run_work() { sequence: 0, handle: edge_actor::OpaqueHandle::new(42), }); - assert!(!tx.messages().iter().any(|message| { - matches!(message, edge_actor::ActorMessage::ObjectIdentity { .. }) - })); - assert!(!rx.messages().iter().any(|message| { - matches!(message, edge_actor::ActorMessage::OpaqueHandle { .. }) - })); + assert!( + !tx.messages() + .iter() + .any(|message| { matches!(message, edge_actor::ActorMessage::ObjectIdentity { .. }) }) + ); + assert!( + !rx.messages() + .iter() + .any(|message| { matches!(message, edge_actor::ActorMessage::OpaqueHandle { .. }) }) + ); // A mismatched edge id must reject or fault, not create work on this actor. let mut mismatched = new_tx(); diff --git a/tests/mvp_system/weight_lifecycle_guarantees.rs b/crates/mvp-system/src/tests/weight_lifecycle_guarantees.rs similarity index 97% rename from tests/mvp_system/weight_lifecycle_guarantees.rs rename to crates/mvp-system/src/tests/weight_lifecycle_guarantees.rs index b5b40b2..3e1d226 100644 --- a/tests/mvp_system/weight_lifecycle_guarantees.rs +++ b/crates/mvp-system/src/tests/weight_lifecycle_guarantees.rs @@ -169,9 +169,11 @@ fn weights_ready_requires_all_weight_facts_and_precedes_stage_ready() { let final_event = events.pop().expect("fixture has final weight event"); for event in events { harness.observe(event); - assert!(!harness.events().iter().any(|event| { - matches!(event, weights::WeightLifecycleEvent::WeightsReady { .. }) - })); + assert!( + !harness.events().iter().any(|event| { + matches!(event, weights::WeightLifecycleEvent::WeightsReady { .. }) + }) + ); } // The final weight prerequisite emits WeightsReady. diff --git a/crates/mvp-system/src/tx_rx_edge_actor.rs b/crates/mvp-system/src/tx_rx_edge_actor.rs new file mode 100644 index 0000000..5c43902 --- /dev/null +++ b/crates/mvp-system/src/tx_rx_edge_actor.rs @@ -0,0 +1,324 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct EdgeId(pub u64); + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PortId(pub String); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ObjectId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct OpaqueHandle(pub u64); + +impl OpaqueHandle { + pub fn new(id: u64) -> Self { + Self(id) + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct TxConfig { + pub edge_id: EdgeId, + pub role_port: PortId, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RxConfig { + pub edge_id: EdgeId, + pub role_port: PortId, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ActorFaultReason { + MismatchedEdgeId, + StreamFault, + ObjectFault, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ActorMessage { + Lifecycle { + edge_id: EdgeId, + ready: bool, + }, + ObjectIdentity { + edge_id: EdgeId, + object_id: ObjectId, + sequence: u64, + }, + OpaqueHandle { + edge_id: EdgeId, + object_id: ObjectId, + sequence: u64, + handle: OpaqueHandle, + }, + CoarseFault { + edge_id: EdgeId, + reason: ActorFaultReason, + }, + PayloadBytes { + edge_id: EdgeId, + bytes: Vec, + }, + HostPointer { + edge_id: EdgeId, + address: usize, + }, + ByteRange { + edge_id: EdgeId, + start: u64, + len: u64, + }, + CreditCount { + edge_id: EdgeId, + credits: u64, + }, +} + +impl ActorMessage { + pub fn edge_id(&self) -> EdgeId { + match self { + ActorMessage::Lifecycle { edge_id, .. } + | ActorMessage::ObjectIdentity { edge_id, .. } + | ActorMessage::OpaqueHandle { edge_id, .. } + | ActorMessage::CoarseFault { edge_id, .. } + | ActorMessage::PayloadBytes { edge_id, .. } + | ActorMessage::HostPointer { edge_id, .. } + | ActorMessage::ByteRange { edge_id, .. } + | ActorMessage::CreditCount { edge_id, .. } => *edge_id, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum TxEvent { + EdgeReady { + edge_id: EdgeId, + }, + ObjectProduced { + edge_id: EdgeId, + object_id: ObjectId, + sequence: u64, + }, + StreamFault { + edge_id: EdgeId, + }, + ObjectFailed { + edge_id: EdgeId, + }, + StopEdge { + edge_id: EdgeId, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum RxEvent { + EdgeReady { + edge_id: EdgeId, + }, + ObjectLoaded { + edge_id: EdgeId, + object_id: ObjectId, + sequence: u64, + handle: OpaqueHandle, + }, + StreamFault { + edge_id: EdgeId, + }, + ObjectFailed { + edge_id: EdgeId, + }, + StopEdge { + edge_id: EdgeId, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ActorState { + Provisioning, + Ready, + Faulted, + Stopped, +} + +#[cfg(test)] +pub struct TxActorHarness { + config: TxConfig, + state: ActorState, + messages: Vec, +} + +#[cfg(test)] +impl TxActorHarness { + pub fn new(config: TxConfig) -> Self { + Self { + config, + state: ActorState::Provisioning, + messages: Vec::new(), + } + } + + pub fn observe(&mut self, event: TxEvent) { + match event { + TxEvent::EdgeReady { edge_id } => { + if !self.check_edge(edge_id) { + return; + } + if self.state == ActorState::Provisioning { + self.state = ActorState::Ready; + self.messages.push(ActorMessage::Lifecycle { + edge_id, + ready: true, + }); + } + } + TxEvent::ObjectProduced { + edge_id, + object_id, + sequence, + } => { + if self.state == ActorState::Ready && self.check_edge(edge_id) { + self.messages.push(ActorMessage::ObjectIdentity { + edge_id, + object_id, + sequence, + }); + } + } + TxEvent::StreamFault { edge_id } => { + self.fault_if_edge(edge_id, ActorFaultReason::StreamFault) + } + TxEvent::ObjectFailed { edge_id } => { + self.fault_if_edge(edge_id, ActorFaultReason::ObjectFault) + } + TxEvent::StopEdge { edge_id } => { + if edge_id == self.config.edge_id { + self.state = ActorState::Stopped; + } else { + self.mismatched(); + } + } + } + } + + pub fn messages(&self) -> &[ActorMessage] { + &self.messages + } + + fn check_edge(&mut self, edge_id: EdgeId) -> bool { + if edge_id == self.config.edge_id { + true + } else { + self.mismatched(); + false + } + } + + fn fault_if_edge(&mut self, edge_id: EdgeId, reason: ActorFaultReason) { + if self.check_edge(edge_id) && self.state != ActorState::Stopped { + self.state = ActorState::Faulted; + self.messages + .push(ActorMessage::CoarseFault { edge_id, reason }); + } + } + + fn mismatched(&mut self) { + self.state = ActorState::Faulted; + self.messages.push(ActorMessage::CoarseFault { + edge_id: self.config.edge_id, + reason: ActorFaultReason::MismatchedEdgeId, + }); + } +} + +#[cfg(test)] +pub struct RxActorHarness { + config: RxConfig, + state: ActorState, + messages: Vec, +} + +#[cfg(test)] +impl RxActorHarness { + pub fn new(config: RxConfig) -> Self { + Self { + config, + state: ActorState::Provisioning, + messages: Vec::new(), + } + } + + pub fn observe(&mut self, event: RxEvent) { + match event { + RxEvent::EdgeReady { edge_id } => { + if !self.check_edge(edge_id) { + return; + } + if self.state == ActorState::Provisioning { + self.state = ActorState::Ready; + self.messages.push(ActorMessage::Lifecycle { + edge_id, + ready: true, + }); + } + } + RxEvent::ObjectLoaded { + edge_id, + object_id, + sequence, + handle, + } => { + if self.state == ActorState::Ready && self.check_edge(edge_id) { + self.messages.push(ActorMessage::OpaqueHandle { + edge_id, + object_id, + sequence, + handle, + }); + } + } + RxEvent::StreamFault { edge_id } => { + self.fault_if_edge(edge_id, ActorFaultReason::StreamFault) + } + RxEvent::ObjectFailed { edge_id } => { + self.fault_if_edge(edge_id, ActorFaultReason::ObjectFault) + } + RxEvent::StopEdge { edge_id } => { + if edge_id == self.config.edge_id { + self.state = ActorState::Stopped; + } else { + self.mismatched(); + } + } + } + } + + pub fn messages(&self) -> &[ActorMessage] { + &self.messages + } + + fn check_edge(&mut self, edge_id: EdgeId) -> bool { + if edge_id == self.config.edge_id { + true + } else { + self.mismatched(); + false + } + } + + fn fault_if_edge(&mut self, edge_id: EdgeId, reason: ActorFaultReason) { + if self.check_edge(edge_id) && self.state != ActorState::Stopped { + self.state = ActorState::Faulted; + self.messages + .push(ActorMessage::CoarseFault { edge_id, reason }); + } + } + + fn mismatched(&mut self) { + self.state = ActorState::Faulted; + self.messages.push(ActorMessage::CoarseFault { + edge_id: self.config.edge_id, + reason: ActorFaultReason::MismatchedEdgeId, + }); + } +} diff --git a/crates/mvp-system/src/weight_lifecycle.rs b/crates/mvp-system/src/weight_lifecycle.rs new file mode 100644 index 0000000..56649b5 --- /dev/null +++ b/crates/mvp-system/src/weight_lifecycle.rs @@ -0,0 +1,203 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RunId(pub u64); +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct NodeId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct LayerRange { + pub start: u32, + pub end_exclusive: u32, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WeightSource { + WholeGguf { uri: String }, + ShardSet { uris: Vec }, + CachedArtifact { cache_key: String }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct WeightAssignment { + pub run_id: RunId, + pub stage_index: u32, + pub plan_layer_range: LayerRange, + pub assigned_layer_range: LayerRange, + pub source: WeightSource, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ArtifactBytes { + Local, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WeightEvent { + Provisioned(WeightAssignment), + ArtifactAvailable { bytes: ArtifactBytes }, + LayerRangeValidated, + WorkerRangeBound, + OtherStagePrerequisitesReady, + DownloadFailed, + ParseFailed, + DeviceAllocationFailed, + BindingFailed, + InvalidLayerRange, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum StageFaultReason { + WeightDownloadFailed, + WeightParseFailed, + DeviceAllocationFailed, + WeightBindingFailed, + InvalidLayerRange, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WeightLifecycleEvent { + WeightsReady { + run_id: RunId, + stage_index: u32, + }, + StageReady { + run_id: RunId, + stage_index: u32, + }, + StageFault { + run_id: RunId, + stage_index: u32, + reason: StageFaultReason, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WeightCommand { + LoadOrBindRange { + source: WeightSource, + range: LayerRange, + }, + AdvertiseLoadedLayerRange { + range: LayerRange, + }, +} + +#[cfg(test)] +pub struct WeightLifecycleHarness { + _node_id: NodeId, + assignment: Option, + artifact: bool, + validated: bool, + bound: bool, + other_stage_prereqs: bool, + weights_ready: bool, + faulted: bool, + commands: Vec, + events: Vec, +} + +#[cfg(test)] +impl WeightLifecycleHarness { + pub fn new(node_id: NodeId) -> Self { + Self { + _node_id: node_id, + assignment: None, + artifact: false, + validated: false, + bound: false, + other_stage_prereqs: false, + weights_ready: false, + faulted: false, + commands: Vec::new(), + events: Vec::new(), + } + } + + pub fn observe(&mut self, event: WeightEvent) { + match event { + WeightEvent::Provisioned(assignment) => { + self.commands.push(WeightCommand::LoadOrBindRange { + source: assignment.source.clone(), + range: assignment.assigned_layer_range, + }); + if assignment.assigned_layer_range != assignment.plan_layer_range + || assignment.assigned_layer_range.start + >= assignment.assigned_layer_range.end_exclusive + { + self.assignment = Some(assignment); + self.fault(StageFaultReason::InvalidLayerRange); + } else { + self.assignment = Some(assignment); + } + } + WeightEvent::ArtifactAvailable { .. } => self.artifact = true, + WeightEvent::LayerRangeValidated => self.validated = true, + WeightEvent::WorkerRangeBound => self.bound = true, + WeightEvent::OtherStagePrerequisitesReady => self.other_stage_prereqs = true, + WeightEvent::DownloadFailed => self.fault(StageFaultReason::WeightDownloadFailed), + WeightEvent::ParseFailed => self.fault(StageFaultReason::WeightParseFailed), + WeightEvent::DeviceAllocationFailed => { + self.fault(StageFaultReason::DeviceAllocationFailed) + } + WeightEvent::BindingFailed => self.fault(StageFaultReason::WeightBindingFailed), + WeightEvent::InvalidLayerRange => self.fault(StageFaultReason::InvalidLayerRange), + } + self.maybe_weights_ready(); + self.maybe_stage_ready(); + } + + pub fn commands(&self) -> &[WeightCommand] { + &self.commands + } + + pub fn events(&self) -> &[WeightLifecycleEvent] { + &self.events + } + + fn maybe_weights_ready(&mut self) { + if self.faulted || self.weights_ready || self.assignment.is_none() { + return; + } + if self.artifact && self.validated && self.bound { + self.weights_ready = true; + let assignment = self.assignment.as_ref().unwrap(); + self.events.push(WeightLifecycleEvent::WeightsReady { + run_id: assignment.run_id, + stage_index: assignment.stage_index, + }); + } + } + + fn maybe_stage_ready(&mut self) { + if self.faulted || !self.weights_ready || !self.other_stage_prereqs { + return; + } + let assignment = self.assignment.as_ref().unwrap(); + if !self + .events + .iter() + .any(|event| matches!(event, WeightLifecycleEvent::StageReady { .. })) + { + self.events.push(WeightLifecycleEvent::StageReady { + run_id: assignment.run_id, + stage_index: assignment.stage_index, + }); + } + } + + fn fault(&mut self, reason: StageFaultReason) { + if self.faulted { + return; + } + self.faulted = true; + let (run_id, stage_index) = self + .assignment + .as_ref() + .map(|assignment| (assignment.run_id, assignment.stage_index)) + .unwrap_or((RunId(0), 0)); + self.events.push(WeightLifecycleEvent::StageFault { + run_id, + stage_index, + reason, + }); + } +}