diff --git a/lib/llm/src/discovery/controller.rs b/lib/llm/src/discovery/controller.rs index 37ea1b2bd44c..550fca7427ca 100644 --- a/lib/llm/src/discovery/controller.rs +++ b/lib/llm/src/discovery/controller.rs @@ -47,6 +47,8 @@ pub(crate) struct DesiredInstance { pub(crate) group_key: GroupKey, pub(crate) mdc_checksum: String, pub(crate) projection_fingerprint: String, + /// Digest of this worker's Qwen video prompt-expansion contract. + pub(crate) video_contract: Option, } impl DesiredInstance { @@ -55,12 +57,43 @@ impl DesiredInstance { } } +/// Returns the cohort's shared contract, if every member publishes one. +/// +/// A WorkerSet has one video-routing processor, so exact video routing is +/// enabled only when all members agree. +fn cohort_video_contract(members: &[DesiredInstance]) -> Option { + let mut members = members.iter(); + let agreed = members.next()?.video_contract.clone()?; + members + .all(|member| member.video_contract.as_deref() == Some(agreed.as_str())) + .then_some(agreed) +} + +/// Identifies a cohort's active video contract without changing its MDC checksum. +fn cohort_fingerprint(mdc_checksum: &str, video_contract: Option<&str>) -> String { + match video_contract { + Some(contract) => format!("{mdc_checksum}\0video_contract\0{contract}"), + None => mdc_checksum.to_string(), + } +} + +/// Splits a fingerprint back into the MDC checksum and video contract it carries. +fn fingerprint_parts(fingerprint: &str) -> (&str, Option<&str>) { + match fingerprint.split_once("\0video_contract\0") { + Some((mdc_checksum, contract)) => (mdc_checksum, Some(contract)), + None => (fingerprint, None), + } +} + #[derive(Clone, Debug)] pub(crate) struct GroupSpec { pub(crate) key: GroupKey, pub(crate) mdc_checksum: String, + pub(crate) fingerprint: String, pub(crate) generation: u64, pub(crate) representative: DesiredInstance, + /// Contract shared by the cohort, if any. + pub(crate) video_contract: Option, } #[async_trait] @@ -95,6 +128,14 @@ pub(crate) trait ControllerHost: Send + Sync + 'static { adapters: &[DesiredInstance], ) -> anyhow::Result<()>; + fn replace_prepared_group( + &self, + spec: &GroupSpec, + prepared: Self::Prepared, + members: &[DesiredInstance], + adapters: &[DesiredInstance], + ) -> anyhow::Result<()>; + fn remove_group(&self, key: &GroupKey); fn discard_prepared(&self, prepared: Self::Prepared); @@ -107,11 +148,13 @@ enum GroupStatus { Idle, Queued { mdc_checksum: String, + committed_members: Option>, }, Building { mdc_checksum: String, generation: u64, cancellation: CancellationToken, + committed_members: Option>, }, Ready { mdc_checksum: String, @@ -120,10 +163,12 @@ enum GroupStatus { Retrying { mdc_checksum: String, deadline: Instant, + committed_members: Option>, }, Blocked { mdc_checksum: String, deadline: Instant, + committed_members: Option>, }, BlockedReady { mdc_checksum: String, @@ -132,12 +177,24 @@ enum GroupStatus { }, } +/// The committed pipeline that keeps serving while its replacement is prepared. +/// Retired once the replacement commits or the commit is withdrawn. +struct RetainedPipeline { + admission_tx: watch::Sender>, + /// Contract this pipeline expands video prompts with. Held here because the + /// group's status carries the *replacement's* fingerprint while a rebuild is + /// pending, so it is the only record of what the retained workers must publish. + /// `None` means it was built without exact video routing. + contract: Option, +} + struct DesiredGroup { retry_attempt: u32, cohorts: HashMap>, cohort_order: VecDeque, reported_rejections: HashSet, admission_tx: watch::Sender>, + retained: Option, status: GroupStatus, } @@ -150,6 +207,7 @@ impl DesiredGroup { cohort_order: VecDeque::new(), reported_rejections: HashSet::new(), admission_tx, + retained: None, status: GroupStatus::Idle, } } @@ -340,6 +398,7 @@ impl ModelDiscoveryController { if let Some(existing) = self.desired.get(&instance.key) { if existing.mdc_checksum == instance.mdc_checksum && existing.projection_fingerprint == instance.projection_fingerprint + && existing.video_contract == instance.video_contract { // A repeated registration still confirms liveness after an older snapshot. self.record_mutation(instance.key.clone()); @@ -424,6 +483,9 @@ impl ModelDiscoveryController { if group.cohorts.is_empty() { group.admission_tx.send_replace(Vec::new()); + if let Some(retained) = group.retained.take() { + retained.admission_tx.send_replace(Vec::new()); + } cancel_build(&old_status); if status_has_commit(&old_status) { self.host.remove_group(key); @@ -435,26 +497,92 @@ impl ModelDiscoveryController { .selected_checksum() .expect("non-empty group has a cohort") .to_string(); - if status_checksum(&old_status).is_some_and(|previous| previous != mdc_checksum) { - group.admission_tx.send_replace(Vec::new()); + let member_keys = group.cohorts[&mdc_checksum].clone(); + let members = self.members(&member_keys); + let fingerprint = + cohort_fingerprint(&mdc_checksum, cohort_video_contract(&members).as_deref()); + let fingerprint_changed = + status_checksum(&old_status).is_some_and(|previous| previous != fingerprint); + let mut retained_commit = status_committed_members(&old_status).cloned(); + // The contract the currently committed pipeline serves with. While a rebuild + // is pending, `old_status` names the replacement, so the retained pipeline is + // the authority on what its own workers must publish. + let committed_contract = match &group.retained { + Some(retained) => retained.contract.clone(), + None => fingerprint_parts(status_checksum(&old_status).unwrap_or_default()) + .1 + .map(str::to_string), + }; + if fingerprint_changed { + retained_commit = retained_commit.and_then(|committed_members| { + let (previous_checksum, _) = fingerprint_parts(status_checksum(&old_status)?); + if previous_checksum != mdc_checksum { + return None; + } + surviving_commit(&members, &committed_members, committed_contract.as_deref()) + }); cancel_build(&old_status); - if status_has_commit(&old_status) { - self.host.remove_group(key); + if retained_commit.is_some() { + // The retained pipeline keeps its own sender, or its client leaves the + // update loop and freezes on the worker list it last saw. Retained + // pipeline clients must never observe the successor's IDs, so the + // successor builds against a fresh channel. + let (successor_tx, _) = watch::channel(Vec::new()); + let previous_tx = std::mem::replace(&mut group.admission_tx, successor_tx); + group.retained.get_or_insert_with(|| RetainedPipeline { + admission_tx: previous_tx, + contract: committed_contract, + }); + } else { + group.admission_tx.send_replace(Vec::new()); + if let Some(retained) = group.retained.take() { + retained.admission_tx.send_replace(Vec::new()); + } + if status_has_commit(&old_status) { + self.host.remove_group(key); + } + // Retained pipeline clients must never observe the successor's IDs. + let (admission_tx, _) = watch::channel(Vec::new()); + group.admission_tx = admission_tx; } - // Retained pipeline clients must never observe the successor's IDs. - let (admission_tx, _) = watch::channel(Vec::new()); - group.admission_tx = admission_tx; group.reported_rejections.clear(); group.retry_attempt = 0; + } else if group.retained.is_some() { + // A rebuild is already pending, so the fingerprint tracks the replacement + // and can stay put while the retained workers change underneath it. Its + // contract is checked again here, or a worker that republished a different + // one keeps receiving requests expanded with the contract it dropped. + let had_commit = retained_commit.is_some(); + retained_commit = retained_commit.and_then(|committed_members| { + surviving_commit(&members, &committed_members, committed_contract.as_deref()) + }); + if had_commit && retained_commit.is_none() { + self.host.remove_group(key); + } + } + // Follow the retained pipeline's membership until it is retired, so a worker + // that leaves during a slow or failing rebuild stops receiving requests. + if let Some(retained_members) = &retained_commit { + if let Some(retained) = &group.retained { + let still_serving = members + .iter() + .filter(|member| retained_members.contains(&member.key)) + .cloned() + .collect::>(); + retained + .admission_tx + .send_replace(admitted_ids(&still_serving)); + } + } else if let Some(retained) = group.retained.take() { + retained.admission_tx.send_replace(Vec::new()); } group.report_rejections(key); - let member_keys = group.cohorts[&mdc_checksum].clone(); - let members = self.members(&member_keys); let admitted = admitted_ids(&members); if !matches!( &old_status, GroupStatus::Ready { .. } | GroupStatus::BlockedReady { .. } - ) { + ) || (fingerprint_changed && retained_commit.is_some()) + { group.admission_tx.send_replace(admitted); } @@ -467,7 +595,7 @@ impl ModelDiscoveryController { mdc_checksum: ready_checksum, committed_members, .. - } if ready_checksum == mdc_checksum => { + } if ready_checksum == fingerprint => { let current_members = member_keys; let old_admitted = group.admission_tx.borrow().clone(); let new_admitted = admitted_ids(&members); @@ -485,7 +613,7 @@ impl ModelDiscoveryController { group.admission_tx.send_replace(new_admitted); group.retry_attempt = 0; GroupStatus::Ready { - mdc_checksum, + mdc_checksum: fingerprint, committed_members: current_members, } } @@ -500,7 +628,7 @@ impl ModelDiscoveryController { "Discovery-group replacement blocked; retaining the last safe commit" ); GroupStatus::BlockedReady { - mdc_checksum, + mdc_checksum: fingerprint, committed_members, deadline: Instant::now() + delay, } @@ -515,43 +643,60 @@ impl ModelDiscoveryController { "Discovery-group membership replacement failed; withdrawing stale commit" ); GroupStatus::Blocked { - mdc_checksum, + mdc_checksum: fingerprint, deadline: Instant::now() + retry_delay(group.retry_attempt), + committed_members: None, } } } } + // These arms hold a replacement whose fingerprint has not moved. They carry + // `retained_commit` rather than the commit the status arrived with, so a + // commit withdrawn above is not claimed again by the pending rebuild. GroupStatus::Building { mdc_checksum: building_checksum, generation, cancellation, - } if building_checksum == mdc_checksum => GroupStatus::Building { - mdc_checksum, + .. + } if building_checksum == fingerprint => GroupStatus::Building { + mdc_checksum: fingerprint, generation, cancellation, + committed_members: retained_commit, }, GroupStatus::Queued { mdc_checksum: queued_checksum, - } if queued_checksum == mdc_checksum => GroupStatus::Queued { mdc_checksum }, + .. + } if queued_checksum == fingerprint => GroupStatus::Queued { + mdc_checksum: fingerprint, + committed_members: retained_commit, + }, GroupStatus::Retrying { mdc_checksum: retry_checksum, deadline, - } if retry_checksum == mdc_checksum && !desired_changed => GroupStatus::Retrying { - mdc_checksum, + .. + } if retry_checksum == fingerprint && !desired_changed => GroupStatus::Retrying { + mdc_checksum: fingerprint, deadline, + committed_members: retained_commit, }, GroupStatus::Blocked { mdc_checksum: blocked_checksum, deadline, - } if blocked_checksum == mdc_checksum && !desired_changed => GroupStatus::Blocked { - mdc_checksum, + .. + } if blocked_checksum == fingerprint && !desired_changed => GroupStatus::Blocked { + mdc_checksum: fingerprint, deadline, + committed_members: retained_commit, }, previous => { cancel_build(&previous); group.admission_tx.send_replace(admitted_ids(&members)); group.retry_attempt = 0; - GroupStatus::Queued { mdc_checksum } + GroupStatus::Queued { + mdc_checksum: fingerprint, + committed_members: retained_commit, + } } }; self.groups.insert(key.clone(), group); @@ -630,20 +775,30 @@ impl ModelDiscoveryController { let Some(group) = self.groups.get_mut(&key) else { continue; }; - let GroupStatus::Queued { mdc_checksum } = &group.status else { + let GroupStatus::Queued { .. } = &group.status else { continue; }; - let mdc_checksum = mdc_checksum.clone(); - let Some(member_key) = group - .cohorts - .get(&mdc_checksum) - .and_then(|members| members.first()) - else { + let committed_members = match &group.status { + GroupStatus::Queued { + committed_members, .. + } => committed_members.clone(), + _ => unreachable!("queued status was checked above"), + }; + let Some(mdc_checksum) = group.selected_checksum().map(str::to_string) else { continue; }; - let Some(representative) = self.desired.get(member_key).cloned() else { + let Some(member_keys) = group.cohorts.get(&mdc_checksum) else { continue; }; + let members = member_keys + .iter() + .filter_map(|member_key| self.desired.get(member_key).cloned()) + .collect::>(); + let Some(representative) = members.first().cloned() else { + continue; + }; + let video_contract = cohort_video_contract(&members); + let fingerprint = cohort_fingerprint(&mdc_checksum, video_contract.as_deref()); let cancellation = CancellationToken::new(); let generation = self.next_build_generation; @@ -651,13 +806,16 @@ impl ModelDiscoveryController { let spec = GroupSpec { key: key.clone(), mdc_checksum: mdc_checksum.clone(), + fingerprint: fingerprint.clone(), generation, representative, + video_contract, }; group.status = GroupStatus::Building { - mdc_checksum, + mdc_checksum: fingerprint, generation, cancellation: cancellation.clone(), + committed_members, }; let host = self.host.clone(); @@ -668,7 +826,11 @@ impl ModelDiscoveryController { tokio::select! { biased; _ = cancellation.cancelled() => BuildOutcome::Cancelled, - result = host.prepare(task_spec.clone(), admitted_ids, cancellation.clone()) => { + result = host.prepare( + task_spec.clone(), + admitted_ids, + cancellation.clone(), + ) => { match result { Ok(prepared) => BuildOutcome::Prepared(prepared), Err(error) => BuildOutcome::Failed(error), @@ -678,9 +840,9 @@ impl ModelDiscoveryController { }); let outcome = match future.catch_unwind().await { Ok(outcome) => outcome, - Err(_) => BuildOutcome::Failed(anyhow::anyhow!( - "model materialization panicked" - )), + Err(_) => { + BuildOutcome::Failed(anyhow::anyhow!("model materialization panicked")) + } }; BuildResult { spec: task_spec, @@ -704,9 +866,24 @@ impl ModelDiscoveryController { mdc_checksum, generation, .. - } if mdc_checksum == &result.spec.mdc_checksum && *generation == result.spec.generation + } if mdc_checksum == &result.spec.fingerprint && *generation == result.spec.generation ) && group.selected_checksum() == Some(result.spec.mdc_checksum.as_str()) - && group.cohorts.contains_key(&result.spec.mdc_checksum); + && group.cohorts.contains_key(&result.spec.mdc_checksum) + && group + .cohorts + .get(&result.spec.mdc_checksum) + .is_some_and(|member_keys| { + cohort_fingerprint( + &result.spec.mdc_checksum, + cohort_video_contract(&self.members(member_keys)).as_deref(), + ) == result.spec.fingerprint + }); + let committed_members = match &group.status { + GroupStatus::Building { + committed_members, .. + } => committed_members.clone(), + _ => None, + }; if !is_current { if let BuildOutcome::Prepared(prepared) = result.outcome { self.host.discard_prepared(prepared); @@ -725,14 +902,21 @@ impl ModelDiscoveryController { let members = self.members(&member_keys); let adapters = self.adapters_for_members(&member_keys); group.admission_tx.send_replace(admitted_ids(&members)); - match self - .host - .commit_group(&result.spec, prepared, &members, &adapters) - { + let commit_result = if committed_members.is_some() { + self.host + .replace_prepared_group(&result.spec, prepared, &members, &adapters) + } else { + self.host + .commit_group(&result.spec, prepared, &members, &adapters) + }; + match commit_result { Ok(()) => { + // The replacement has taken over, so retire the sender the + // superseded pipeline was watching. + group.retained = None; group.retry_attempt = 0; group.status = GroupStatus::Ready { - mdc_checksum: result.spec.mdc_checksum, + mdc_checksum: result.spec.fingerprint, committed_members: member_keys, }; } @@ -745,8 +929,9 @@ impl ModelDiscoveryController { ); group.retry_attempt = group.retry_attempt.saturating_add(1); group.status = GroupStatus::Blocked { - mdc_checksum: result.spec.mdc_checksum, + mdc_checksum: result.spec.fingerprint, deadline: Instant::now() + retry_delay(group.retry_attempt), + committed_members, }; } } @@ -762,13 +947,15 @@ impl ModelDiscoveryController { "Model materialization failed; scheduling retry" ); group.status = GroupStatus::Retrying { - mdc_checksum: result.spec.mdc_checksum, + mdc_checksum: result.spec.fingerprint, deadline: Instant::now() + delay, + committed_members, }; } BuildOutcome::Cancelled => { group.status = GroupStatus::Queued { - mdc_checksum: result.spec.mdc_checksum, + mdc_checksum: result.spec.fingerprint, + committed_members, }; } } @@ -791,15 +978,17 @@ impl ModelDiscoveryController { let now = Instant::now(); let mut retained_retries = Vec::new(); for (key, group) in &mut self.groups { - let (mdc_checksum, deadline) = match &group.status { + let (mdc_checksum, deadline, committed_members) = match &group.status { GroupStatus::Retrying { mdc_checksum, deadline, + committed_members, } | GroupStatus::Blocked { mdc_checksum, deadline, - } => (mdc_checksum, deadline), + committed_members, + } => (mdc_checksum, deadline, committed_members.clone()), GroupStatus::BlockedReady { mdc_checksum, committed_members, @@ -817,6 +1006,7 @@ impl ModelDiscoveryController { if *deadline <= now { group.status = GroupStatus::Queued { mdc_checksum: mdc_checksum.clone(), + committed_members, }; } } @@ -828,6 +1018,9 @@ impl ModelDiscoveryController { async fn shutdown_builds(&mut self) { for group in self.groups.values_mut() { group.admission_tx.send_replace(Vec::new()); + if let Some(retained) = group.retained.take() { + retained.admission_tx.send_replace(Vec::new()); + } cancel_build(&group.status); } self.builds.abort_all(); @@ -925,10 +1118,46 @@ fn cancel_build(status: &GroupStatus) { } } +/// The committed workers that may keep serving: those still present, provided they +/// all still publish `contract`. `None` withdraws the commit. +fn surviving_commit( + members: &[DesiredInstance], + committed_members: &BTreeSet, + contract: Option<&str>, +) -> Option> { + // Keep serving the committed workers that are still here rather than demanding + // the whole membership, so dropping the last worker of a rolling upgrade does + // not withdraw a healthy group. + let surviving = members + .iter() + .filter(|member| committed_members.contains(&member.key)) + .collect::>(); + if surviving.is_empty() { + return None; + } + // The committed pipeline expands video prompts with the contract it was built + // from, so every worker it still serves must publish that same contract; one + // that republished a different contract under the same ID would otherwise keep + // receiving requests against the old one. + if let Some(contract) = contract + && !surviving + .iter() + .all(|member| member.video_contract.as_deref() == Some(contract)) + { + return None; + } + Some( + surviving + .into_iter() + .map(|member| member.key.clone()) + .collect(), + ) +} + fn status_checksum(status: &GroupStatus) -> Option<&str> { match status { GroupStatus::Idle => None, - GroupStatus::Queued { mdc_checksum } + GroupStatus::Queued { mdc_checksum, .. } | GroupStatus::Building { mdc_checksum, .. } | GroupStatus::Ready { mdc_checksum, .. } | GroupStatus::Retrying { mdc_checksum, .. } @@ -938,10 +1167,31 @@ fn status_checksum(status: &GroupStatus) -> Option<&str> { } fn status_has_commit(status: &GroupStatus) -> bool { - matches!( - status, - GroupStatus::Ready { .. } | GroupStatus::BlockedReady { .. } - ) + status_committed_members(status).is_some() +} + +fn status_committed_members(status: &GroupStatus) -> Option<&BTreeSet> { + match status { + GroupStatus::Ready { + committed_members, .. + } + | GroupStatus::BlockedReady { + committed_members, .. + } => Some(committed_members), + GroupStatus::Queued { + committed_members, .. + } + | GroupStatus::Building { + committed_members, .. + } + | GroupStatus::Retrying { + committed_members, .. + } + | GroupStatus::Blocked { + committed_members, .. + } => committed_members.as_ref(), + GroupStatus::Idle => None, + } } async fn wait_for_deadline(deadline: Option) { @@ -972,6 +1222,8 @@ mod tests { }; use tokio::sync::{Semaphore, mpsc}; + use crate::local_model::runtime_config::VLLM_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY; + struct Prepared(u64); struct FakeHost { @@ -987,6 +1239,7 @@ mod tests { admissions: Mutex>>>, removed_groups: AtomicUsize, discarded: AtomicUsize, + prepared_replacements: AtomicUsize, } impl FakeHost { @@ -1006,6 +1259,7 @@ mod tests { admissions: Mutex::new(Vec::new()), removed_groups: AtomicUsize::new(0), discarded: AtomicUsize::new(0), + prepared_replacements: AtomicUsize::new(0), }), start_rx, ) @@ -1095,6 +1349,7 @@ mod tests { group_key: group_key(), mdc_checksum: card.mdcsum().to_string(), projection_fingerprint: card.source_path.clone().unwrap(), + video_contract: crate::discovery::watcher::qwen_video_contract_digest(&card), card, })) } @@ -1138,6 +1393,12 @@ mod tests { { anyhow::bail!("injected commit conflict"); } + // `ModelManager::commit_discovery_group` refuses a group ID it already holds. + anyhow::ensure!( + !self.committed.lock().unwrap().contains_key(&spec.key.id()), + "discovery group {:?} is already committed", + spec.key.id() + ); self.committed.lock().unwrap().insert( spec.key.id(), members.iter().map(|member| member.key.clone()).collect(), @@ -1161,6 +1422,12 @@ mod tests { { anyhow::bail!("injected replacement conflict"); } + // `ModelManager::replace_discovery_group` requires a committed group ID. + anyhow::ensure!( + self.committed.lock().unwrap().contains_key(&key.id()), + "committed discovery group {:?} not found", + key.id() + ); self.committed.lock().unwrap().insert( key.id(), members.iter().map(|member| member.key.clone()).collect(), @@ -1169,6 +1436,18 @@ mod tests { Ok(()) } + fn replace_prepared_group( + &self, + spec: &GroupSpec, + prepared: Self::Prepared, + members: &[DesiredInstance], + adapters: &[DesiredInstance], + ) -> anyhow::Result<()> { + let Prepared(_build) = prepared; + self.prepared_replacements.fetch_add(1, Ordering::SeqCst); + self.replace_group(&spec.key, members, adapters) + } + fn remove_group(&self, key: &GroupKey) { self.committed.lock().unwrap().remove(&key.id()); self.adapters.lock().unwrap().remove(&key.id()); @@ -1194,8 +1473,22 @@ mod tests { } fn instance(id: u64, mdc_checksum: &str) -> DesiredInstance { + build_instance(id, mdc_checksum, None) + } + + fn build_instance( + id: u64, + mdc_checksum: &str, + video_contract: Option, + ) -> DesiredInstance { let mut card = ModelDeploymentCard::with_name_only("model"); card.source_path = Some(mdc_checksum.to_string()); + if let Some(contract) = video_contract { + card.runtime_config.runtime_data.insert( + VLLM_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY.to_string(), + contract, + ); + } let mcid = ModelCardInstanceId { namespace: "namespace".to_string(), component: "worker".to_string(), @@ -1213,11 +1506,39 @@ mod tests { }, group_key: group_key(), mdc_checksum: card.mdcsum().to_string(), - card, projection_fingerprint: mdc_checksum.to_string(), + video_contract: crate::discovery::watcher::qwen_video_contract_digest(&card), + card, } } + /// Builds a worker that publishes a Qwen video prompt-expansion contract. + /// + /// The contract stays out of `mdcsum()`, so workers that differ only in + /// `resize_mode` share a cohort and disagree only on the contract. + fn instance_with_contract(id: u64, mdc_checksum: &str, resize_mode: &str) -> DesiredInstance { + let instance = build_instance( + id, + mdc_checksum, + Some(serde_json::json!({ + "placeholder_target": "bare_video_token", + "resize_mode": resize_mode, + })), + ); + assert!(instance.video_contract.is_some()); + assert_eq!( + instance.mdc_checksum, + self::instance(id, mdc_checksum).mdc_checksum + ); + instance + } + + fn contract_digest(resize_mode: &str) -> String { + instance_with_contract(0, "spec", resize_mode) + .video_contract + .unwrap() + } + fn discovery_instance(instance: &DesiredInstance) -> DiscoveryInstance { DiscoveryInstance::Model { namespace: instance.mcid.namespace.clone(), @@ -1235,6 +1556,235 @@ mod tests { controller.apply_build_result(result); } + #[tokio::test] + async fn contract_agreement_drives_the_fingerprint_and_survives_a_rolling_upgrade() { + let (host, mut starts) = FakeHost::new(); + let mut controller = ModelDiscoveryController::new(host.clone()); + // A legacy worker that predates the contract, alongside one that publishes it. + let legacy = instance(1, "same"); + let upgraded = instance_with_contract(2, "same", "round_ties_even"); + assert_eq!(legacy.mdc_checksum, upgraded.mdc_checksum); + + controller.apply_added(legacy.clone()); + controller.apply_added(upgraded.clone()); + controller.start_queued_builds(); + let mixed_spec = starts.recv().await.unwrap(); + // The cohort disagrees, so exact video routing stays off. + assert_eq!(mixed_spec.video_contract, None); + assert_eq!(mixed_spec.fingerprint, mixed_spec.mdc_checksum); + host.release.add_permits(1); + finish_build(&mut controller).await; + assert_eq!( + host.members(&group_key()), + BTreeSet::from([legacy.key.clone(), upgraded.key.clone()]) + ); + let committed_admissions = host.admissions.lock().unwrap()[0].clone(); + assert_eq!(*committed_admissions.borrow(), vec![1, 2]); + + // Removing the last legacy worker changes membership and contract agreement + // at once. Text service must keep running while the replacement is prepared. + controller.apply_removed(&legacy.key); + assert_eq!( + host.members(&group_key()), + BTreeSet::from([legacy.key.clone(), upgraded.key.clone()]), + "the committed group must survive the upgrade's last legacy worker leaving" + ); + assert_eq!(host.removed_groups.load(Ordering::SeqCst), 0); + // The retained pipeline keeps its channel and stops admitting the worker + // that left, so it never routes to a worker that is gone. + assert_eq!(*committed_admissions.borrow(), vec![2]); + assert!(committed_admissions.has_changed().is_ok()); + + controller.start_queued_builds(); + let upgraded_spec = starts.recv().await.unwrap(); + assert_eq!( + upgraded_spec.video_contract, + Some(contract_digest("round_ties_even")) + ); + assert_eq!( + upgraded_spec.fingerprint, + cohort_fingerprint( + &upgraded_spec.mdc_checksum, + Some(&contract_digest("round_ties_even")) + ) + ); + host.release.add_permits(1); + finish_build(&mut controller).await; + + // A retained commit is replaced in place, never committed a second time. + assert_eq!(host.prepared_replacements.load(Ordering::SeqCst), 1); + assert_eq!(host.members(&group_key()), BTreeSet::from([upgraded.key])); + assert_eq!(host.removed_groups.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn republished_contract_withdraws_the_commit_it_no_longer_matches() { + let (host, mut starts) = FakeHost::new(); + let mut controller = ModelDiscoveryController::new(host.clone()); + let worker = instance_with_contract(1, "same", "round_ties_even"); + let peer = instance_with_contract(2, "same", "round_ties_even"); + + controller.apply_added(worker.clone()); + controller.apply_added(peer.clone()); + controller.start_queued_builds(); + let spec = starts.recv().await.unwrap(); + assert_eq!( + spec.video_contract, + Some(contract_digest("round_ties_even")) + ); + host.release.add_permits(1); + finish_build(&mut controller).await; + assert_eq!(host.members(&group_key()).len(), 2); + + // The same instance ID republishes a different contract. The committed + // pipeline expands video prompts with the old one, so it cannot be kept. + let republished = instance_with_contract(1, "same", "legacy_ceil"); + assert_eq!(republished.key, worker.key); + assert_ne!(republished.video_contract, worker.video_contract); + assert!(controller.apply_added(republished)); + + assert!( + host.members(&group_key()).is_empty(), + "a worker serving a contract it no longer publishes must lose its commit" + ); + assert_eq!(host.removed_groups.load(Ordering::SeqCst), 1); + + // The cohort now disagrees, so the rebuild turns exact video routing off. + controller.start_queued_builds(); + let rebuilt = starts.recv().await.unwrap(); + assert_eq!(rebuilt.video_contract, None); + host.release.add_permits(1); + finish_build(&mut controller).await; + assert_eq!(host.prepared_replacements.load(Ordering::SeqCst), 0); + assert_eq!(host.members(&group_key()).len(), 2); + } + + #[tokio::test] + async fn a_republished_contract_during_a_pending_rebuild_withdraws_the_retained_commit() { + let (host, mut starts) = FakeHost::new(); + let mut controller = ModelDiscoveryController::new(host.clone()); + let first = instance_with_contract(1, "same", "round_ties_even"); + let second = instance_with_contract(2, "same", "round_ties_even"); + + controller.apply_added(first.clone()); + controller.apply_added(second.clone()); + controller.start_queued_builds(); + starts.recv().await.unwrap(); + host.release.add_permits(1); + finish_build(&mut controller).await; + let committed_admissions = host.admissions.lock().unwrap()[0].clone(); + assert_eq!(*committed_admissions.borrow(), vec![1, 2]); + + // A worker without a contract joins, so the cohort stops agreeing and a + // replacement with exact video routing off is queued. The commit built from + // `round_ties_even` keeps serving while that replacement is prepared. + let legacy = instance(3, "same"); + controller.apply_added(legacy.clone()); + assert_eq!(host.members(&group_key()).len(), 2); + assert_eq!(*committed_admissions.borrow(), vec![1, 2]); + assert_eq!(host.removed_groups.load(Ordering::SeqCst), 0); + + // Before that rebuild completes, a retained worker republishes a different + // contract. The cohort still has no common contract, so the replacement's + // fingerprint does not move -- but the retained pipeline still expands video + // prompts with `round_ties_even`, which worker 1 no longer publishes. + let republished = instance_with_contract(1, "same", "legacy_ceil"); + assert_eq!(republished.key, first.key); + assert!(controller.apply_added(republished)); + + assert!( + host.members(&group_key()).is_empty(), + "a retained pipeline must not keep serving a worker that republished a different contract" + ); + assert_eq!(host.removed_groups.load(Ordering::SeqCst), 1); + assert!(committed_admissions.borrow().is_empty()); + } + + #[tokio::test(start_paused = true)] + async fn a_republished_contract_during_a_failed_rebuild_withdraws_the_retained_commit() { + let (host, mut starts) = FakeHost::new(); + let mut controller = ModelDiscoveryController::new(host.clone()); + let first = instance_with_contract(1, "same", "round_ties_even"); + let second = instance_with_contract(2, "same", "round_ties_even"); + + controller.apply_added(first.clone()); + controller.apply_added(second.clone()); + controller.start_queued_builds(); + starts.recv().await.unwrap(); + host.release.add_permits(1); + finish_build(&mut controller).await; + let committed_admissions = host.admissions.lock().unwrap()[0].clone(); + assert_eq!(*committed_admissions.borrow(), vec![1, 2]); + + // A worker without a contract joins and the replacement build fails, so the + // commit built from `round_ties_even` keeps serving across the retry. + host.failures.store(1, Ordering::SeqCst); + controller.apply_added(instance(3, "same")); + controller.start_queued_builds(); + starts.recv().await.unwrap(); + host.release.add_permits(1); + finish_build(&mut controller).await; + assert_eq!(host.members(&group_key()).len(), 2); + assert_eq!(host.removed_groups.load(Ordering::SeqCst), 0); + assert_eq!(*committed_admissions.borrow(), vec![1, 2]); + + // A retained worker republishes a different contract while the failed rebuild + // waits to retry. The retained pipeline still expands video prompts with + // `round_ties_even`, so it cannot keep serving that worker. + let republished = instance_with_contract(1, "same", "legacy_ceil"); + assert_eq!(republished.key, first.key); + assert!(controller.apply_added(republished)); + + assert!( + host.members(&group_key()).is_empty(), + "a failed rebuild must not keep a commit whose worker republished a different contract" + ); + assert_eq!(host.removed_groups.load(Ordering::SeqCst), 1); + assert!(committed_admissions.borrow().is_empty()); + } + + #[tokio::test(start_paused = true)] + async fn a_retained_pipeline_keeps_receiving_worker_updates_while_rebuilding() { + let (host, mut starts) = FakeHost::new(); + let mut controller = ModelDiscoveryController::new(host.clone()); + let first = instance_with_contract(1, "same", "round_ties_even"); + let second = instance_with_contract(2, "same", "round_ties_even"); + let legacy = instance(3, "same"); + + controller.apply_added(first.clone()); + controller.apply_added(second.clone()); + controller.apply_added(legacy.clone()); + controller.start_queued_builds(); + starts.recv().await.unwrap(); + host.release.add_permits(1); + finish_build(&mut controller).await; + let committed_admissions = host.admissions.lock().unwrap()[0].clone(); + assert_eq!(*committed_admissions.borrow(), vec![1, 2, 3]); + + // Dropping the legacy worker retains the commit and starts a rebuild that fails. + host.failures.store(1, Ordering::SeqCst); + controller.apply_removed(&legacy.key); + assert_eq!(*committed_admissions.borrow(), vec![1, 2]); + controller.start_queued_builds(); + starts.recv().await.unwrap(); + host.release.add_permits(1); + finish_build(&mut controller).await; + assert_eq!(host.members(&group_key()).len(), 3); + + // A worker leaving during the failed rebuild must still reach the retained + // pipeline, or it keeps routing to a worker that is gone. + controller.apply_removed(&second.key); + assert_eq!(*committed_admissions.borrow(), vec![1]); + assert!(committed_admissions.has_changed().is_ok()); + + controller.start_queued_builds(); + starts.recv().await.unwrap(); + host.release.add_permits(1); + finish_build(&mut controller).await; + assert_eq!(host.members(&group_key()), BTreeSet::from([first.key])); + assert_eq!(host.prepared_replacements.load(Ordering::SeqCst), 1); + } + #[tokio::test] async fn membership_churn_keeps_one_build_and_commits_latest_members() { let (host, mut starts) = FakeHost::new(); diff --git a/lib/llm/src/discovery/model_manager.rs b/lib/llm/src/discovery/model_manager.rs index 260c9b48f6a6..0e2ffa5b4c4f 100644 --- a/lib/llm/src/discovery/model_manager.rs +++ b/lib/llm/src/discovery/model_manager.rs @@ -812,6 +812,7 @@ impl ModelManager { pub(crate) fn replace_discovery_group( &self, group_id: &str, + replacement_worker_set: Option, members: Vec<(String, ModelDeploymentCard)>, adapters: Vec<(String, ModelDeploymentCard)>, ) -> anyhow::Result<()> { @@ -821,8 +822,13 @@ impl ModelManager { .get(group_id) .ok_or_else(|| anyhow::anyhow!("committed discovery group {group_id:?} not found"))?; let primary = group.primary.clone(); + let namespace = group.namespace.clone(); let worker_set_key = group.worker_set_key.clone(); - let worker_set = group.worker_set.clone(); + let aliases = group.aliases.clone(); + let replacing_worker_set = replacement_worker_set.is_some(); + let worker_set = replacement_worker_set + .map(Arc::new) + .unwrap_or_else(|| group.worker_set.clone()); let previous_member_keys = group.cards.keys().cloned().collect::>(); let previous_adapter_keys = group.adapters.keys().cloned().collect::>(); let previous_adapter_names = group @@ -836,6 +842,10 @@ impl ModelManager { !members.is_empty(), "cannot replace with an empty discovery group" ); + anyhow::ensure!( + worker_set.namespace() == namespace, + "replacement WorkerSet namespace does not match committed discovery group" + ); self.validate_adapter_claims(&primary, adapters.iter().map(|(_, card)| card))?; let members = members.into_iter().collect::>(); let adapters = adapters.into_iter().collect::>(); @@ -856,6 +866,21 @@ impl ModelManager { .collect::>(); let lora_before = self.lora_projection_locked(); + if replacing_worker_set { + let primary_model = self.get_or_create_model(&primary); + if let Some(displaced_worker_set) = primary_model.get_worker_set(&worker_set_key) { + Self::clear_worker_set_targets(&displaced_worker_set); + } + primary_model.add_worker_set(worker_set_key.clone(), worker_set.clone()); + for alias in &aliases { + let alias_model = self.get_or_create_model(alias); + if let Some(displaced_worker_set) = alias_model.get_worker_set(&worker_set_key) { + Self::clear_worker_set_targets(&displaced_worker_set); + } + alias_model.add_worker_set(worker_set_key.clone(), worker_set.clone()); + } + } + for key in previous_member_keys.difference(&desired_member_keys) { self.cards.remove(key); } @@ -877,6 +902,9 @@ impl ModelManager { .expect("non-empty members checked above"); group.cards = members; group.adapters = adapters; + if replacing_worker_set { + group.worker_set = worker_set; + } drop(group); for (name, adapter_view) in adapter_views { @@ -891,6 +919,9 @@ impl ModelManager { } let lora_after = self.lora_projection_locked(); self.publish_lora_projection_locked(Self::union_lora_projection(&lora_before, &lora_after)); + if replacing_worker_set { + self.reconcile_discovery_topology(&primary, &namespace); + } self.publish_catalog_locked(); self.publish_lora_projection_locked(lora_after); Ok(()) diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index 293625f03f0b..38a068244192 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -36,6 +36,7 @@ use crate::{ }, local_model::runtime_config::{ ModelRuntimeConfig, TokenizerBackend, VLLM_INFERENCE_V1_GENERATE_CAPABILITY, + VLLM_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY, }, model_card::ModelDeploymentCard, model_type::{ModelInput, ModelType}, @@ -435,6 +436,23 @@ where validate_selector_worker_role(card, self.require_typed_worker_role)?; + // Prepare without exact video routing unless the cohort agreed on a contract. + if spec.video_contract.is_none() + && card + .runtime_config + .runtime_data + .remove(VLLM_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY) + .is_some() + { + tracing::warn!( + target: "mm_routing", + model_name = card.name(), + group = %spec.key.id(), + "WorkerSet members publish different Qwen video prompt-expansion contracts; \ + exact video routing disabled for this group" + ); + } + // Use per-worker-set router config if the worker provided one in its MDC, // otherwise fall back to the frontend-level global config. Policy selections // are process-local, so preserve them when the MDC supplies the base config. @@ -1058,6 +1076,7 @@ where }; let mdc_checksum = card.mdcsum().to_string(); let projection_fingerprint = lora_projection_fingerprint(&card)?; + let video_contract = qwen_video_contract_digest(&card); Ok(Some(DesiredInstance { key: mcid.to_path(), mcid, @@ -1066,6 +1085,7 @@ where group_key, mdc_checksum, projection_fingerprint, + video_contract, })) } @@ -1174,6 +1194,7 @@ where .collect::>(); self.manager.replace_discovery_group( &group_id, + None, members .iter() .map(|member| (member.key.clone(), member.card.clone())) @@ -1200,6 +1221,83 @@ where Ok(()) } + fn replace_prepared_group( + &self, + spec: &GroupSpec, + mut prepared: Self::Prepared, + members: &[DesiredInstance], + adapters: &[DesiredInstance], + ) -> anyhow::Result<()> { + let group_id = spec.key.id(); + let previous = self + .manager + .discovery_group_adapter_cards(&group_id) + .into_iter() + .map(|card| (card.name().to_string(), card)) + .collect::>(); + let adapter_was_available = previous + .keys() + .cloned() + .chain( + adapters + .iter() + .map(|adapter| adapter.card.name().to_string()), + ) + .map(|name| { + ( + name.clone(), + self.manager.get_committed_model(&name).is_some(), + ) + }) + .collect::>(); + let worker_set = prepared + .worker_set + .take() + .ok_or_else(|| anyhow::anyhow!("prepared WorkerSet was already consumed"))?; + let mut committed_members = members + .iter() + .map(|member| (member.key.clone(), member.card.clone())) + .collect::>(); + if let Some((_, card)) = committed_members + .iter_mut() + .find(|(key, _)| key == &spec.representative.key) + { + *card = prepared.card.clone(); + } + let desired = adapters + .iter() + .map(|adapter| (adapter.card.name().to_string(), adapter.card.clone())) + .collect::>(); + self.manager.replace_discovery_group( + &group_id, + Some(worker_set), + committed_members, + adapters + .iter() + .map(|adapter| (adapter.key.clone(), adapter.card.clone())) + .collect(), + )?; + self.emit_update(ModelUpdate::Added(prepared.card.clone())); + for (name, card) in &desired { + if !adapter_was_available.get(name).copied().unwrap_or(false) + && self.manager.get_committed_model(name).is_some() + { + self.emit_update(ModelUpdate::Added(card.clone())); + } + } + for (name, card) in previous { + if adapter_was_available.get(&name).copied().unwrap_or(false) + && self.manager.get_committed_model(&name).is_none() + { + self.emit_update(ModelUpdate::Removed(card)); + } + } + if prepared.card.model_type.supports_chat() { + self.notify_on_model.notify_waiters(); + } + Ok(()) + } + fn remove_group(&self, key: &GroupKey) { let Some(removed) = self.manager.remove_discovery_group(&key.id()) else { return; @@ -1316,6 +1414,17 @@ fn lora_projection_fingerprint(card: &ModelDeploymentCard) -> anyhow::Result Option { + let mut contract = card + .runtime_config + .runtime_data + .get(VLLM_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY)? + .clone(); + canonicalize_json(&mut contract); + Some(blake3::hash(contract.to_string().as_bytes()).to_string()) +} + fn canonicalize_json(value: &mut serde_json::Value) { match value { serde_json::Value::Object(object) => { @@ -2044,14 +2153,17 @@ mod tests { endpoint_id, mdc_checksum: card.mdcsum().to_string(), projection_fingerprint: lora_projection_fingerprint(&card).unwrap(), + video_contract: qwen_video_contract_digest(&card), card, group_key: key.clone(), }; let spec = GroupSpec { key, mdc_checksum: desired.mdc_checksum.clone(), + fingerprint: desired.mdc_checksum.clone(), generation: 1, representative: desired.clone(), + video_contract: desired.video_contract.clone(), }; let (_admission_tx, admission_rx) = tokio::sync::watch::channel(vec![1]); let prepared = watcher @@ -2140,14 +2252,17 @@ mod tests { endpoint_id, mdc_checksum: card.mdcsum().to_string(), projection_fingerprint: lora_projection_fingerprint(&card).unwrap(), + video_contract: qwen_video_contract_digest(&card), card, group_key: key.clone(), }; let spec = GroupSpec { key, mdc_checksum: desired.mdc_checksum.clone(), + fingerprint: desired.mdc_checksum.clone(), generation: 1, representative: desired, + video_contract: None, }; let (_admission_tx, admission_rx) = tokio::sync::watch::channel(vec![1]); diff --git a/lib/llm/src/model_card.rs b/lib/llm/src/model_card.rs index 6db5c6eaf2b2..9b22c5a9efd1 100644 --- a/lib/llm/src/model_card.rs +++ b/lib/llm/src/model_card.rs @@ -20,7 +20,6 @@ use crate::common::checked_file::CheckedFile; use crate::entrypoint::RouterConfig; use crate::local_model::runtime_config::{ ModelRuntimeConfig, TokenizerBackend, VLLM_NEMOTRON_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY, - VLLM_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY, }; use crate::model_type::{ModelInput, ModelType}; use crate::protocols::tensor::TensorModelConfig; @@ -1260,20 +1259,13 @@ impl ModelDeploymentCard { bytes_to_hash.extend_from_slice(b"\0vllm_enable_tower_connector_lora\0true"); } - // The frontend constructs one video routing preprocessor from - // the cohort's representative card. Partition workers when a - // processor contract is missing or differs so every admitted - // worker uses the same model-visible prompt expansion. - for key in [ - VLLM_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY, + // The Qwen video contract is resolved per cohort, not per card. + // Nemotron contracts still partition WorkerSets by checksum. + append_runtime_contract_checksum( + &mut bytes_to_hash, + &self.runtime_config, VLLM_NEMOTRON_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY, - ] { - append_runtime_contract_checksum( - &mut bytes_to_hash, - &self.runtime_config, - key, - ); - } + ); // TODO: Do we want any other user_data or runtime_config? @@ -3255,7 +3247,57 @@ mod ownership_tests { } #[test] - fn video_processor_runtime_contracts_isolate_worker_sets() { + fn qwen_video_processor_contract_stays_out_of_the_checksum() { + use crate::local_model::runtime_config::VLLM_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY; + + // `mdcsum()` caches via `OnceLock`, so each case uses a fresh card. + + fn card_with_contract(contract: serde_json::Value) -> ModelDeploymentCard { + let mut card = ModelDeploymentCard::with_name_only("model"); + card.runtime_config.runtime_data.insert( + VLLM_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY.to_string(), + contract, + ); + card + } + + fn legacy_ceil() -> serde_json::Value { + serde_json::json!({ + "placeholder_target": "bare_video_token", + "resize_mode": "legacy_ceil", + }) + } + + let withheld = ModelDeploymentCard::with_name_only("model"); + assert_eq!( + withheld.mdcsum(), + card_with_contract(legacy_ceil()).mdcsum(), + "a worker that predates the contract must still join the WorkerSet of one that publishes it" + ); + + assert_eq!( + card_with_contract(legacy_ceil()).mdcsum(), + card_with_contract(serde_json::json!({ + "placeholder_target": "bare_video_token", + "resize_mode": "round_ties_even", + })) + .mdcsum(), + "two workers publishing different contracts must still serve as one group" + ); + + // Negative control: a deployment that never carries this key is + // unaffected, so its cache directory and reported checksum do not move. + let bare = ModelDeploymentCard::with_name_only("model"); + let mut unrelated = ModelDeploymentCard::with_name_only("model"); + unrelated + .runtime_config + .runtime_data + .insert("some_unrelated_runtime_key".to_string(), true.into()); + assert_eq!(bare.mdcsum(), unrelated.mdcsum()); + } + + #[test] + fn video_processor_runtime_contract_checksum_boundaries() { use crate::local_model::runtime_config::{ VLLM_NEMOTRON_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY, VLLM_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY, @@ -3302,9 +3344,9 @@ mod ownership_tests { let unrelated = card_with_contract("unrelated_runtime_metadata", serde_json::json!(true)); assert_eq!(missing.mdcsum(), unrelated.mdcsum()); - assert_ne!(missing.mdcsum(), qwen.mdcsum()); + assert_eq!(missing.mdcsum(), qwen.mdcsum()); assert_eq!(qwen.mdcsum(), same_qwen.mdcsum()); - assert_ne!(qwen.mdcsum(), different_qwen.mdcsum()); + assert_eq!(qwen.mdcsum(), different_qwen.mdcsum()); assert_ne!(missing.mdcsum(), nemotron.mdcsum()); assert_ne!(nemotron.mdcsum(), different_nemotron.mdcsum()); assert_ne!(qwen.mdcsum(), nemotron.mdcsum());