diff --git a/agent/src/util/crictl_containers.rs b/agent/src/util/crictl_containers.rs index e3dea472a..a6d257176 100644 --- a/agent/src/util/crictl_containers.rs +++ b/agent/src/util/crictl_containers.rs @@ -1,4 +1,4 @@ -use akri_shared::akri::AKRI_SLOT_ANNOTATION_NAME; +use akri_shared::akri::AKRI_SLOT_ANNOTATION_NAME_PREFIX; use std::collections::{HashMap, HashSet}; /// Output from crictl query @@ -15,18 +15,24 @@ struct CriCtlContainer { annotations: HashMap, } -/// This gets the usage slots for an instance by getting the annotations that were stored at id `AKRI_SLOT_ANNOTATION_NAME` during allocate. +/// This gets the usage slots for an instance by getting the annotations that were stored at id `AKRI_SLOT_ANNOTATION_NAME_PREFIX` during allocate. pub fn get_container_slot_usage(crictl_output: &str) -> HashSet { match serde_json::from_str::(crictl_output) { Ok(crictl_output_parsed) => crictl_output_parsed .containers .iter() - .filter_map(|container| { - container - .annotations - .get(&AKRI_SLOT_ANNOTATION_NAME.to_string()) + .flat_map(|container| &container.annotations) + .filter_map(|(key, value)| { + if key.starts_with(AKRI_SLOT_ANNOTATION_NAME_PREFIX) + && value.eq(key + .strip_prefix(AKRI_SLOT_ANNOTATION_NAME_PREFIX) + .unwrap_or_default()) + { + Some(value.clone()) + } else { + None + } }) - .map(|string_ref| string_ref.to_string()) .collect(), Err(e) => { trace!( @@ -108,7 +114,7 @@ mod tests { expected, get_container_slot_usage(&format!( "{{ \"ddd\": \"\", \"containers\": [ {} ] }}", - &get_container_str("\"akri.agent.slot\": \"foo\",") + &get_container_str("\"akri.agent.slot-foo\": \"foo\",") )) ); // Expected output with slot @@ -116,7 +122,7 @@ mod tests { expected, get_container_slot_usage(&format!( "{{ \"containers\": [ {} ] }}", - &get_container_str("\"akri.agent.slot\": \"foo\",") + &get_container_str("\"akri.agent.slot-foo\": \"foo\",") )) ); // Expected output with multiple containers @@ -127,8 +133,8 @@ mod tests { expected_2, get_container_slot_usage(&format!( "{{ \"containers\": [ {}, {} ] }}", - &get_container_str("\"akri.agent.slot\": \"foo1\","), - &get_container_str("\"akri.agent.slot\": \"foo2\","), + &get_container_str("\"akri.agent.slot-foo1\": \"foo1\","), + &get_container_str("\"akri.agent.slot-foo2\": \"foo2\","), )) ); } diff --git a/agent/src/util/device_plugin_service.rs b/agent/src/util/device_plugin_service.rs index c1f7ff5d5..3e2e5ffe5 100644 --- a/agent/src/util/device_plugin_service.rs +++ b/agent/src/util/device_plugin_service.rs @@ -13,7 +13,7 @@ use akri_shared::{ configuration::ConfigurationSpec, instance::InstanceSpec, retry::{random_delay, MAX_INSTANCE_UPDATE_TRIES}, - AKRI_SLOT_ANNOTATION_NAME, + AKRI_SLOT_ANNOTATION_NAME_PREFIX, }, k8s, k8s::KubeInterface, @@ -287,7 +287,7 @@ impl DevicePluginService { &self.instance_name, request, ); - let mut akri_annotations = std::collections::HashMap::new(); + let mut akri_annotations = HashMap::new(); for device_usage_id in request.devices_i_ds { trace!( "internal_allocate - for Instance {} processing request for device usage slot id {}", @@ -296,7 +296,7 @@ impl DevicePluginService { ); akri_annotations.insert( - AKRI_SLOT_ANNOTATION_NAME.to_string(), + format!("{}{}", AKRI_SLOT_ANNOTATION_NAME_PREFIX, &device_usage_id), device_usage_id.clone(), ); diff --git a/shared/src/akri/mod.rs b/shared/src/akri/mod.rs index 7a8ba22a3..aa4ca7484 100644 --- a/shared/src/akri/mod.rs +++ b/shared/src/akri/mod.rs @@ -11,8 +11,8 @@ pub const API_CONFIGURATIONS: &str = "configurations"; pub const API_INSTANCES: &str = "instances"; /// Akri prefix pub const AKRI_PREFIX: &str = "akri.sh"; -/// Container Annotation name used to store slot name -pub const AKRI_SLOT_ANNOTATION_NAME: &str = "akri.agent.slot"; +/// Container Annotation name prefix used to store slot name +pub const AKRI_SLOT_ANNOTATION_NAME_PREFIX: &str = "akri.agent.slot-"; pub mod configuration; pub mod instance;