From ec3508956813cac9ad555ccc133eb0369613bbab Mon Sep 17 00:00:00 2001 From: Jay Zhu Date: Mon, 3 Aug 2026 22:13:42 -0600 Subject: [PATCH] fix(nvlink): refresh NVL domain UUIDs from NMX-C Hello for machines Signed-off-by: Jay Zhu --- crates/nvlink-manager/src/lib.rs | 154 ++++++++++++++++++++++++++++--- 1 file changed, 140 insertions(+), 14 deletions(-) diff --git a/crates/nvlink-manager/src/lib.rs b/crates/nvlink-manager/src/lib.rs index e07cd1a0fe..2526447dc3 100644 --- a/crates/nvlink-manager/src/lib.rs +++ b/crates/nvlink-manager/src/lib.rs @@ -145,6 +145,7 @@ fn domain_uuid_from_nmx_c_hello( .as_ref() .and_then(|header| uuid::Uuid::parse_str(&header.domain_uuid).ok()) .map(NvLinkDomainId::from) + .filter(|domain_uuid| *domain_uuid != NvLinkDomainId::nil()) .ok_or_else(|| { NvLinkManagerError::internal(format!( "Failed to parse domain UUID from NMX-C hello response: {hello:?}" @@ -183,9 +184,8 @@ fn build_machine_nvlink_info_from_nmx_c_hello( ) -> MachineNvLinkInfo { if let Some(existing) = existing { let mut info = existing.clone(); - if info.domain_uuid == NvLinkDomainId::nil() { - info.domain_uuid = domain_uuid; - } + info.domain_uuid = domain_uuid; + if info.chassis_serial.trim().is_empty() { info.chassis_serial = chassis_serial.to_string(); } @@ -218,11 +218,11 @@ fn build_machine_nvlink_info_from_nmx_c_hello( } } -/// Populates missing `machines.status.nvlink_info` entries (or nil `domain_uuid`) using NMX-C hello. +/// Updates missing or stale machine NVLink domains using a validated NMX-C Hello. fn populate_machine_nvlink_info_if_needed( machine_nvlink_info: &mut HashMap>, managed_host_snapshots: &HashMap, - chassis_serial: &str, + snapshot_chassis_serial: Option<&str>, machine_ids: &[MachineId], domain_uuid: NvLinkDomainId, ) -> Vec<(MachineId, MachineNvLinkInfo)> { @@ -233,14 +233,22 @@ fn populate_machine_nvlink_info_if_needed( .and_then(|info| info.as_ref()); let needs_update = match existing { None => true, - Some(info) => { - info.domain_uuid == NvLinkDomainId::nil() || info.chassis_serial.trim().is_empty() - } + Some(info) => info.domain_uuid != domain_uuid || info.chassis_serial.trim().is_empty(), }; if !needs_update { continue; } + // Fall back to persisted chassis inventory so a missing snapshot serial does + // not block refreshing the machine's NVLink domain UUID. + let Some(chassis_serial) = snapshot_chassis_serial + .or_else(|| existing.map(|info| info.chassis_serial.as_str())) + .map(str::trim) + .filter(|serial| !serial.is_empty()) + else { + continue; + }; + let snapshot = managed_host_snapshots.get(machine_id); let updated = build_machine_nvlink_info_from_nmx_c_hello( existing, @@ -1548,21 +1556,26 @@ impl NvlPartitionMonitor { managed_host_snapshots_domain.keys().copied().collect(); let mut nvlink_info_db_updates = Vec::new(); for snapshot in snapshots { - // The chassis serial number should always be available from hardware_info. - let Some(chassis_serial) = managed_host_chassis_serial(snapshot) else { + let machine_id = snapshot.host_snapshot.id; + let snapshot_chassis_serial = managed_host_chassis_serial(snapshot); + let has_persisted_chassis_serial = machine_nvlink_info + .get(&machine_id) + .and_then(|info| info.as_ref()) + .is_some_and(|info| !info.chassis_serial.trim().is_empty()); + if snapshot_chassis_serial.is_none() && !has_persisted_chassis_serial { tracing::warn!( group_id, group_type = group_type_label, - machine_id = %snapshot.host_snapshot.id, + %machine_id, "Skipping nvlink_info population; chassis serial unavailable" ); continue; - }; + } nvlink_info_db_updates.extend(populate_machine_nvlink_info_if_needed( machine_nvlink_info, all_managed_host_snapshots, - &chassis_serial, - &[snapshot.host_snapshot.id], + snapshot_chassis_serial.as_deref(), + &[machine_id], domain_uuid, )); } @@ -2950,6 +2963,119 @@ fn record_domain_health( } } +#[cfg(test)] +mod domain_uuid_observation_tests { + use carbide_test_support::{Check, check_values, value_scenarios}; + use carbide_uuid::machine::{MachineIdSource, MachineType}; + use libnmxc::nmxc_model::{ServerHeader, ServerHello}; + + use super::*; + + #[test] + fn hello_domain_uuid_accepts_only_valid_non_nil_values() { + value_scenarios!( + run = |domain_uuid| { + domain_uuid_from_nmx_c_hello(&ServerHello { + server_header: Some(ServerHeader { + domain_uuid: domain_uuid.to_string(), + ..Default::default() + }), + ..Default::default() + }) + .is_ok() + }; + + "accepted" { + "11111111-1111-1111-1111-111111111111" => true, + } + + "rejected" { + "not-a-uuid" => false, + "00000000-0000-0000-0000-000000000000" => false, + } + + ); + } + + #[test] + fn valid_domain_observations_update_only_changed_machine_domains() + -> Result<(), Box> { + let machine_id = MachineId::new(MachineIdSource::Tpm, [1; 32], MachineType::Host); + let old_domain: NvLinkDomainId = "11111111-1111-1111-1111-111111111111".parse()?; + let observed_domain: NvLinkDomainId = "22222222-2222-2222-2222-222222222222".parse()?; + + let gpu = NvLinkGpu { + tray_index: 1, + slot_id: 2, + device_id: 3, + guid: 4, + }; + + let old_info = MachineNvLinkInfo { + domain_uuid: old_domain, + chassis_serial: "CHASSIS-A".to_string(), + gpus: vec![gpu], + }; + + let observed_info = MachineNvLinkInfo { + domain_uuid: observed_domain, + ..old_info.clone() + }; + + let initial_info = MachineNvLinkInfo { + domain_uuid: observed_domain, + chassis_serial: "CHASSIS-A".to_string(), + gpus: Vec::new(), + }; + + check_values( + [ + Check { + scenario: "initial population", + input: (None, Some("CHASSIS-A")), + expect: (1, initial_info), + }, + Check { + scenario: "changed UUID", + input: (Some(old_info.clone()), Some("CHASSIS-A")), + expect: (1, observed_info.clone()), + }, + Check { + scenario: "changed UUID with persisted chassis serial only", + input: (Some(old_info), None), + expect: (1, observed_info.clone()), + }, + Check { + scenario: "identical persisted UUID after restart", + input: (Some(observed_info.clone()), Some("CHASSIS-A")), + expect: (0, observed_info), + }, + ], + |(existing, snapshot_chassis_serial)| { + let mut machine_nvlink_info = HashMap::from([(machine_id, existing)]); + + let updates = populate_machine_nvlink_info_if_needed( + &mut machine_nvlink_info, + &HashMap::new(), + snapshot_chassis_serial, + &[machine_id], + observed_domain, + ); + + ( + updates.len(), + machine_nvlink_info + .remove(&machine_id) + .flatten() + .expect("observation populates machine NVLink info"), + ) + }, + ); + + Ok(()) + } +} + #[cfg(test)] mod partition_classification_tests { use libnmxc::nmxc_model::PartitionInfo;