Suffix usage slot to annotation key name (#560)

* Suffix usage slot to annotation key name Signed-off-by: Johnson Shih <jshih@microsoft.com> * Update patch version Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> * Revert "Update patch version" Signed-off-by: Johnson Shih <jshih@microsoft.com> This reverts commit d4b5338. --------- Signed-off-by: Johnson Shih <jshih@microsoft.com> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
project-akri · Apr 5, 2023 · b999938 · b999938
1 parent 6fe39ec
commit b999938
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 16 deletions.
diff --git 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<String, String>,
 }
 
-/// 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<String> {
     match serde_json::from_str::<CriCtlOutput>(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,15 +114,15 @@ 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
         assert_eq!(
             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
@@ -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
@@ -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;