fix: DPU Extension Service status only include used DPUs

Signed-off-by: Felicity Xu <hanyux@nvidia.com>

Author: hanyux-nv

crates/api-model/src/instance/config/network.rs

@@ -240,6 +240,54 @@ impl InstanceNetworkConfig {
         }
     }
 
+    /// Returns the DPU machine IDs used by the instance network configuration.
+    pub fn get_used_dpus(
+        &self,
+        device_to_id_map: &HashMap<String, Vec<MachineId>>,
+    ) -> Vec<MachineId> {
+        let device_locators: Vec<DeviceLocator> = self
+            .interfaces
+            .iter()
+            .filter_map(|i| i.device_locator.clone())
+            .collect();
+
+        let legacy_physical_interface_count = self
+            .interfaces
+            .iter()
+            .filter(|iface| {
+                iface.function_id == InterfaceFunctionId::Physical {}
+                    && iface.device_locator.is_none()
+            })
+            .count();
+
+        let use_primary_dpu_only = legacy_physical_interface_count > 0
+            || device_locators.is_empty()
+            || device_to_id_map.is_empty();
+
+        if use_primary_dpu_only {
+            return device_to_id_map
+                .values()
+                .find_map(|ids| ids.first().copied())
+                .into_iter()
+                .collect();
+        }
+
+        let mut used_dpus = Vec::new();
+        for iface in &self.interfaces {
+            let Some(device_locator) = iface.device_locator.as_ref() else {
+                continue;
+            };
+            if let Some(dpu_ids) = device_to_id_map.get(&device_locator.device)
+                && let Some(dpu_id) = dpu_ids.get(device_locator.device_instance)
+                && !used_dpus.contains(dpu_id)
+            {
+                used_dpus.push(*dpu_id);
+            }
+        }
+
+        used_dpus
+    }
+
     /// Validates the network configuration.
     ///
     /// Note: this is also called on POST-resolution configs (i.e. after

crates/api-model/src/instance/status/extension_service.rs

@@ -44,7 +44,7 @@ impl InstanceExtensionServicesStatus {
     /// For each extension service, we aggregate the statuses from all DPUs.
     /// The config passed must be from database (not rpc InstanceConfig), and must contain any terminating services.
     pub fn from_config_and_observations(
-        dpu_id_to_device_map: &HashMap<String, Vec<MachineId>>,
+        dpu_ids: &[MachineId],
         config: Versioned<&InstanceExtensionServicesConfig>,
         observations: &HashMap<MachineId, InstanceExtensionServiceStatusObservation>,
     ) -> Self {
@@ -57,15 +57,11 @@ impl InstanceExtensionServicesStatus {
             };
         }
 
-        // Extract all unique DPU machine IDs from the dpu_id_to_device_map
-        let all_dpu_ids: Vec<MachineId> =
-            dpu_id_to_device_map.values().flatten().copied().collect();
-
         // Instance allocation rejects non-empty service_configs on zero-DPU
         // hosts, so, in practice, we *shouldn't* reach here. BUT, if we do,
         // assume it's from something like a stale pre-validation instance,
         // and just report unsynced.
-        if all_dpu_ids.is_empty() {
+        if dpu_ids.is_empty() {
             return Self::unsynced_for_config(&config);
         }
 
@@ -76,7 +72,7 @@ impl InstanceExtensionServicesStatus {
         for service in config.service_configs.iter() {
             let mut dpu_statuses = vec![];
 
-            for dpu_id in &all_dpu_ids {
+            for dpu_id in dpu_ids {
                 match observations.get(dpu_id) {
                     // DPU has observation with matching config version
                     Some(obs) if obs.config_version == config.version => {
@@ -447,12 +443,6 @@ mod tests {
         ExtensionServiceId::from_str("00000000-0000-0000-0000-000000000000").unwrap()
     }
 
-    fn create_dpu_map_with_one_dpu() -> HashMap<String, Vec<MachineId>> {
-        let mut map = HashMap::new();
-        map.insert("device0".to_string(), vec![get_test_machine_id()]);
-        map
-    }
-
     fn create_service_config(version: ConfigVersion) -> InstanceExtensionServicesConfig {
         InstanceExtensionServicesConfig {
             service_configs: vec![InstanceExtensionServiceConfig {
@@ -497,10 +487,9 @@ mod tests {
         let config = create_service_config(service_version);
 
         let config_version = ConfigVersion::initial();
-        let dpu_map = create_dpu_map_with_one_dpu();
 
         let status = InstanceExtensionServicesStatus::from_config_and_observations(
-            &dpu_map,
+            &vec![get_test_machine_id()],
             Versioned::new(&config, config_version),
             &HashMap::new(),
         );
@@ -530,15 +519,14 @@ mod tests {
         let config = create_service_config(service_version);
         let config_version = ConfigVersion::initial();
 
-        let dpu_map = create_dpu_map_with_one_dpu();
         let observations = create_observation(
             config_version,
             service_version,
             ExtensionServiceDeploymentStatus::Running,
         );
 
         let status = InstanceExtensionServicesStatus::from_config_and_observations(
-            &dpu_map,
+            &vec![get_test_machine_id()],
             Versioned::new(&config, config_version),
             &observations,
         );
@@ -568,15 +556,14 @@ mod tests {
     fn extension_service_status_with_outdated_observation() {
         let config = create_service_config(ConfigVersion::initial());
         let version = ConfigVersion::initial();
-        let dpu_map = create_dpu_map_with_one_dpu();
         let observations = create_observation(
             ConfigVersion::initial(),
             ConfigVersion::initial(),
             ExtensionServiceDeploymentStatus::Running,
         );
 
         let status = InstanceExtensionServicesStatus::from_config_and_observations(
-            &dpu_map,
+            &vec![get_test_machine_id()],
             Versioned::new(&config, version.increment()), // Increment config version for extension service config
             &observations,                                // Observation has initial config version
         );
@@ -631,7 +618,7 @@ mod tests {
         );
 
         let status = InstanceExtensionServicesStatus::from_config_and_observations(
-            &dpu_map,
+            &vec![get_test_machine_id()],
             Versioned::new(&config, config_version),
             &observations,
         );
@@ -718,7 +705,7 @@ mod tests {
         );
 
         let status = InstanceExtensionServicesStatus::from_config_and_observations(
-            &dpu_map,
+            &vec![get_test_machine_id()],
             Versioned::new(&config, config_version),
             &observations,
         );
@@ -745,6 +732,77 @@ mod tests {
         assert_eq!(readiness, ExtensionServicesReadiness::Ready);
     }
 
+    #[test]
+    fn extension_service_status_scopes_to_target_dpus() {
+        let service_version = ConfigVersion::initial();
+        let config = create_service_config(service_version);
+        let config_version = ConfigVersion::initial();
+
+        let dpu1_id = get_test_machine_id();
+        let dpu2_id =
+            MachineId::from_str("fm100ds27v4uuq7sgs4gsjummskt0b3tedugtpevjrbfh6su081n9jufcq0")
+                .unwrap();
+        let mut observations = HashMap::new();
+        observations.insert(
+            dpu1_id,
+            InstanceExtensionServiceStatusObservation {
+                config_version,
+                instance_config_version: None,
+                extension_service_statuses: vec![ExtensionServiceStatusObservation {
+                    service_id: get_test_service_id(),
+                    service_type: ExtensionServiceType::KubernetesPod,
+                    service_name: "test-service".to_string(),
+                    version: service_version,
+                    removed: None,
+                    overall_state: ExtensionServiceDeploymentStatus::Running,
+                    components: vec![],
+                    message: String::new(),
+                }],
+                observed_at: chrono::Utc::now(),
+            },
+        );
+        observations.insert(
+            dpu2_id,
+            InstanceExtensionServiceStatusObservation {
+                config_version,
+                instance_config_version: None,
+                extension_service_statuses: vec![ExtensionServiceStatusObservation {
+                    service_id: get_test_service_id(),
+                    service_type: ExtensionServiceType::KubernetesPod,
+                    service_name: "test-service".to_string(),
+                    version: service_version,
+                    removed: None,
+                    overall_state: ExtensionServiceDeploymentStatus::Pending,
+                    components: vec![],
+                    message: "No pod sandbox found".to_string(),
+                }],
+                observed_at: chrono::Utc::now(),
+            },
+        );
+
+        let status = InstanceExtensionServicesStatus::from_config_and_observations(
+            &[dpu1_id],
+            Versioned::new(&config, config_version),
+            &observations,
+        );
+
+        assert_eq!(status.configs_synced, SyncState::Synced);
+        assert_eq!(status.extension_services.len(), 1);
+        assert_eq!(
+            status.extension_services[0].overall_status,
+            ExtensionServiceDeploymentStatus::Running
+        );
+        assert_eq!(status.extension_services[0].dpu_statuses.len(), 1);
+        assert_eq!(
+            status.extension_services[0].dpu_statuses[0].machine_id,
+            dpu1_id
+        );
+        assert_eq!(
+            compute_extension_services_readiness(&status),
+            ExtensionServicesReadiness::Ready
+        );
+    }
+
     #[test]
     fn extension_service_calculate_overall_status_all_running() {
         let dpu_statuses = vec![
@@ -880,7 +938,7 @@ mod tests {
         );
 
         let status = InstanceExtensionServicesStatus::from_config_and_observations(
-            &dpu_map,
+            &vec![get_test_machine_id()],
             Versioned::new(&config, config_version),
             &observations,
         );

crates/machine-controller/src/handler.rs

@@ -6377,19 +6377,24 @@ impl StateHandler for InstanceStateHandler {
 }
 
 // Gets extension services status from DB, checks if any removed services are fully terminated
-// across all DPUs, if so, remove them from the instance config in the DB(without updating the version).
+// across targeted DPUs, if so, remove them from the instance config in the DB(without updating the version).
 fn get_extension_services_status(
     mh_snapshot: &ManagedHostStateSnapshot,
     instance: &InstanceSnapshot,
 ) -> InstanceExtensionServicesStatus {
-    let (_, dpu_id_to_device_map) = mh_snapshot
+    let (_, device_to_id_map) = mh_snapshot
         .host_snapshot
         .get_dpu_device_and_id_mappings()
         .unwrap_or_else(|_| (HashMap::default(), HashMap::default()));
 
-    // Gather instance extension services status from all DPU observations
+    let mut target_dpus = instance.config.network.get_used_dpus(&device_to_id_map);
+    if target_dpus.is_empty() {
+        target_dpus = device_to_id_map.values().flatten().copied().collect();
+    }
+
+    // Gather instance extension services status from targeted DPUs.
     InstanceExtensionServicesStatus::from_config_and_observations(
-        &dpu_id_to_device_map,
+        &target_dpus,
         Versioned::new(
             &instance.config.extension_services,
             instance.extension_services_config_version,

crates/rpc/src/model/instance/status.rs

@@ -102,6 +102,8 @@ pub fn instance_status_from_config_and_observation(
         // If observations.network.instance_config_version was None, then "ignore"
     }
 
+    let used_dpu_ids = network_config.value.get_used_dpus(&dpu_id_to_device_map);
+
     let network =
         model::instance::status::network::InstanceNetworkStatus::from_config_and_observations(
             dpu_id_to_device_map.clone(),
@@ -117,7 +119,7 @@ pub fn instance_status_from_config_and_observation(
 
     let extension_services =
         model::instance::status::extension_service::InstanceExtensionServicesStatus::from_config_and_observations(
-            &dpu_id_to_device_map,
+            &used_dpu_ids,
             extension_services_config,
             &observations.extension_services,
         );