diff --git a/crates/skippy-metrics/src/lib.rs b/crates/skippy-metrics/src/lib.rs index 65649d92e..e208751b3 100644 --- a/crates/skippy-metrics/src/lib.rs +++ b/crates/skippy-metrics/src/lib.rs @@ -6,6 +6,12 @@ pub mod attr { pub const TOPOLOGY_ID: &str = "skippy.topology_id"; pub const REQUEST_ID: &str = "skippy.request_id"; pub const SESSION_ID: &str = "skippy.session_id"; + pub const PROMPT_INDEX: &str = "skippy.prompt_index"; + pub const MESSAGE_KIND: &str = "skippy.message_kind"; + pub const TOKEN_COUNT: &str = "skippy.token_count"; + pub const CHECKPOINT_GENERATION: &str = "skippy.checkpoint_generation"; + pub const PROMPT_TOKEN_COUNT: &str = "skippy.prompt_token_count"; + pub const DECODE_STEP: &str = "skippy.decode_step"; pub const STAGE_ID: &str = "skippy.stage_id"; pub const STAGE_INDEX: &str = "skippy.stage_index"; pub const LAYER_START: &str = "skippy.layer_start"; diff --git a/crates/skippy-protocol/src/binary/mod.rs b/crates/skippy-protocol/src/binary/mod.rs index 610286890..28c099f7a 100644 --- a/crates/skippy-protocol/src/binary/mod.rs +++ b/crates/skippy-protocol/src/binary/mod.rs @@ -19,9 +19,9 @@ pub use types::{ MAX_STAGE_SIDEBAND_VALUES, MAX_STAGE_STATE_IMPORT_BYTES, READY_MAGIC, STAGE_LOGIT_BIAS_WIRE_BYTES, STAGE_SAMPLING_CONFIG_BASE_BYTES, STAGE_STATE_HEADER_BYTES, STAGE_STATE_VERSION, STAGE_WIRE_FIXED_HEADER_BYTES, StageLogitBias, StageReply, - StageReplyStats, StageSamplingConfig, StageStateHeader, StageWireMessage, WireActivationDType, - WireMessageKind, WireReplyKind, WireStagePhase, activation_frame_flags_from_state_flags, - activation_state_flags_from_frame_flags, state_flags, + StageReplyStats, StageRequestEpoch, StageSamplingConfig, StageStateHeader, StageWireMessage, + WireActivationDType, WireMessageKind, WireReplyKind, WireStagePhase, + activation_frame_flags_from_state_flags, activation_state_flags_from_frame_flags, state_flags, }; pub(crate) fn invalid_data(message: &'static str) -> std::io::Error { @@ -153,6 +153,7 @@ mod tests { fn stage_message_round_trips_f32() { let mut state = StageStateHeader::new(WireMessageKind::DecodeEmbd, WireActivationDType::F32); + state.checkpoint_generation = 3; state.prompt_token_count = 1; state.decode_step = 0; state.current_token = 11; @@ -192,6 +193,16 @@ mod tests { assert_eq!(decoded.state.source_stage_index, 0); assert_eq!(decoded.request_id, 7); assert_eq!(decoded.session_id, 11); + assert_eq!( + decoded.request_epoch(), + StageRequestEpoch { + request_id: 7, + session_id: 11, + checkpoint_generation: 3, + prompt_token_count: 1, + decode_step: 0, + } + ); assert_ne!(decoded.state.flags & state_flags::SAMPLING, 0); assert_eq!(decoded.state.flags & state_flags::CHAT_SAMPLING_METADATA, 0); assert_eq!(decoded.chat_sampling_metadata, None); @@ -203,6 +214,68 @@ mod tests { assert_eq!(sampling.logit_bias[0].bias, -50.0); } + #[test] + fn request_epoch_orders_only_matching_flows() { + let older = StageRequestEpoch { + request_id: 7, + session_id: 11, + checkpoint_generation: 1, + prompt_token_count: 8, + decode_step: 2, + }; + let newer = StageRequestEpoch { + request_id: 7, + session_id: 11, + checkpoint_generation: 1, + prompt_token_count: 8, + decode_step: 3, + }; + let different_session = StageRequestEpoch { + session_id: 12, + ..newer + }; + + assert!(older.same_flow(newer)); + assert!(older.is_stale_for(newer)); + assert!(!newer.is_stale_for(older)); + assert!(!older.same_flow(different_session)); + assert!(!older.is_stale_for(different_session)); + } + + #[test] + fn request_epoch_staleness_orders_generation_before_prompt_before_decode() { + let base = StageRequestEpoch { + request_id: 7, + session_id: 11, + checkpoint_generation: 1, + prompt_token_count: 8, + decode_step: 3, + }; + let newer_checkpoint = StageRequestEpoch { + checkpoint_generation: 2, + prompt_token_count: 0, + decode_step: 0, + ..base + }; + let newer_prompt = StageRequestEpoch { + prompt_token_count: 9, + decode_step: 0, + ..base + }; + let newer_decode = StageRequestEpoch { + decode_step: 4, + ..base + }; + + assert!(base.same_flow(newer_checkpoint)); + assert!(base.is_stale_for(newer_checkpoint)); + assert!(!newer_checkpoint.is_stale_for(base)); + assert!(base.is_stale_for(newer_prompt)); + assert!(!newer_prompt.is_stale_for(base)); + assert!(base.is_stale_for(newer_decode)); + assert!(!newer_decode.is_stale_for(base)); + } + #[test] fn generation_config_round_trips_sampling_metadata() { let message = StageWireMessage::configure_generation( diff --git a/crates/skippy-protocol/src/binary/types.rs b/crates/skippy-protocol/src/binary/types.rs index 28ce11029..29c11aeb7 100644 --- a/crates/skippy-protocol/src/binary/types.rs +++ b/crates/skippy-protocol/src/binary/types.rs @@ -343,6 +343,41 @@ impl Default for StageStateHeader { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct StageRequestEpoch { + pub request_id: u64, + pub session_id: u64, + pub checkpoint_generation: i32, + pub prompt_token_count: i32, + pub decode_step: i32, +} + +impl StageRequestEpoch { + pub fn same_flow(self, other: Self) -> bool { + self.request_id == other.request_id && self.session_id == other.session_id + } + + /// Returns true when this epoch is strictly older than `current` within the + /// same request/session flow. + /// + /// Epochs from different flows are never comparable. For matching flows, + /// staleness uses lexicographic ordering of checkpoint generation, prompt + /// token count, and decode step, so a newer checkpoint dominates prompt and + /// decode progress, and prompt progress dominates decode progress. + pub fn is_stale_for(self, current: Self) -> bool { + self.same_flow(current) + && ( + self.checkpoint_generation, + self.prompt_token_count, + self.decode_step, + ) < ( + current.checkpoint_generation, + current.prompt_token_count, + current.decode_step, + ) + } +} + #[derive(Debug, Clone, PartialEq)] pub struct StageWireMessage { pub kind: WireMessageKind, @@ -360,6 +395,16 @@ pub struct StageWireMessage { } impl StageWireMessage { + pub fn request_epoch(&self) -> StageRequestEpoch { + StageRequestEpoch { + request_id: self.request_id, + session_id: self.session_id, + checkpoint_generation: self.state.checkpoint_generation, + prompt_token_count: self.state.prompt_token_count, + decode_step: self.state.decode_step, + } + } + pub fn stop(dtype: WireActivationDType) -> Self { Self::stop_with_identity(dtype, 0, 0) } diff --git a/crates/skippy-server/src/binary_transport.rs b/crates/skippy-server/src/binary_transport.rs index ee398c0a4..314c317f3 100644 --- a/crates/skippy-server/src/binary_transport.rs +++ b/crates/skippy-server/src/binary_transport.rs @@ -1609,28 +1609,27 @@ fn binary_message_attrs( message: &StageWireMessage, ) -> std::collections::BTreeMap { let mut attrs = lifecycle_attrs(config); + let epoch = message.request_epoch(); attrs.insert(attr::SESSION_ID.to_string(), json!(session_id.to_string())); attrs.insert( attr::REQUEST_ID.to_string(), json!(binary_message_request_id(message)), ); + attrs.insert(attr::PROMPT_INDEX.to_string(), json!(message.state.seq_id)); attrs.insert( - "skippy.prompt_index".to_string(), - json!(message.state.seq_id), - ); - attrs.insert( - "skippy.message_kind".to_string(), + attr::MESSAGE_KIND.to_string(), json!(format!("{:?}", message.kind)), ); - attrs.insert("skippy.token_count".to_string(), json!(message.token_count)); + attrs.insert(attr::TOKEN_COUNT.to_string(), json!(message.token_count)); attrs.insert( - "skippy.prompt_token_count".to_string(), - json!(message.state.prompt_token_count), + attr::CHECKPOINT_GENERATION.to_string(), + json!(epoch.checkpoint_generation), ); attrs.insert( - "skippy.decode_step".to_string(), - json!(message.state.decode_step), + attr::PROMPT_TOKEN_COUNT.to_string(), + json!(epoch.prompt_token_count), ); + attrs.insert(attr::DECODE_STEP.to_string(), json!(epoch.decode_step)); let layer_count = i64::from(config.layer_end.saturating_sub(config.layer_start)); let kv_tokens_after = estimated_kv_tokens_after(message); attrs.insert("skippy.kv_tokens_after".to_string(), json!(kv_tokens_after)); @@ -3018,9 +3017,9 @@ impl BinaryRequestSummary { if let Some(request_id) = self.request_id.as_ref() { attrs.insert(attr::REQUEST_ID.to_string(), json!(request_id)); } - attrs.insert("skippy.prompt_index".to_string(), json!(self.prompt_index)); + attrs.insert(attr::PROMPT_INDEX.to_string(), json!(self.prompt_index)); attrs.insert( - "skippy.prompt_token_count".to_string(), + attr::PROMPT_TOKEN_COUNT.to_string(), json!(self.prompt_token_count), ); attrs.insert(