feat: add nvidia MIG Settings

Author: piyush-jena

Cargo.toml

@@ -57,7 +57,7 @@ bottlerocket-template-helper = { path = "./bottlerocket-template-helper", versio
 
 # Settings Models
 bottlerocket-model-derive = { path = "./bottlerocket-settings-models/model-derive", version = "0.1" }
-bottlerocket-modeled-types = { path = "./bottlerocket-settings-models/modeled-types", version = "0.7" }
+bottlerocket-modeled-types = { path = "./bottlerocket-settings-models/modeled-types", version = "0.8" }
 bottlerocket-scalar = { path = "./bottlerocket-settings-models/scalar", version = "0.1" }
 bottlerocket-scalar-derive = { path = "./bottlerocket-settings-models/scalar-derive", version = "0.1" }
 bottlerocket-string-impls-for = { path = "./bottlerocket-settings-models/string-impls-for", version = "0.1" }

bottlerocket-settings-models/modeled-types/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "bottlerocket-modeled-types"
-version = "0.7.0"
+version = "0.8.0"
 authors = []
 license = "Apache-2.0 OR MIT"
 edition = "2021"

bottlerocket-settings-models/modeled-types/src/kubernetes.rs

@@ -1469,6 +1469,8 @@ pub struct NvidiaDevicePluginSettings {
     device_list_strategy: NvidiaDeviceListStrategy,
     device_sharing_strategy: NvidiaDeviceSharingStrategy,
     time_slicing: NvidiaTimeSlicingSettings,
+    device_partitioning_strategy: NvidiaDevicePartitioningStrategy,
+    mig: NvidiaMIGSettings,
 }
 
 #[derive(Debug, Clone, Eq, PartialEq, Hash, Serialize, Deserialize)]
@@ -1515,7 +1517,9 @@ mod tests {
                 device_id_strategy: Some(NvidiaDeviceIdStrategy::Uuid),
                 device_list_strategy: Some(NvidiaDeviceListStrategy::Envvar),
                 device_sharing_strategy: None,
-                time_slicing: None
+                time_slicing: None,
+                device_partitioning_strategy: None,
+                mig: None
             }
         );
         let results = serde_json::to_string(&nvidia_device_plugins).unwrap();
@@ -1534,7 +1538,9 @@ mod tests {
                 device_id_strategy: Some(NvidiaDeviceIdStrategy::Uuid),
                 device_list_strategy: Some(NvidiaDeviceListStrategy::Envvar),
                 device_sharing_strategy: Some(NvidiaDeviceSharingStrategy::TimeSlicing),
-                time_slicing: None
+                time_slicing: None,
+                device_partitioning_strategy: None,
+                mig: None
             }
         );
 
@@ -1549,3 +1555,179 @@ mod tests {
         assert!(result.is_err(), "The JSON should not be parsed successfully as it contains an invalid value for 'replicas'.");
     }
 }
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
+#[serde(rename_all = "lowercase")]
+pub enum NvidiaDevicePartitioningStrategy {
+    #[default]
+    None,
+    MIG,
+}
+
+#[cfg(test)]
+mod test_valid_device_partitioning_strategy {
+    use super::*;
+
+    #[test]
+    fn test_serde_nvidia_device_plugins_with_mig() {
+        let test_json = r#"{"pass-device-specs":false,"device-id-strategy":"uuid","device-list-strategy":"envvar","device-partitioning-strategy":"mig"}"#;
+        let nvidia_device_plugins: NvidiaDevicePluginSettings =
+            serde_json::from_str(test_json).unwrap();
+        assert_eq!(
+            nvidia_device_plugins,
+            NvidiaDevicePluginSettings {
+                pass_device_specs: Some(false),
+                device_id_strategy: Some(NvidiaDeviceIdStrategy::Uuid),
+                device_list_strategy: Some(NvidiaDeviceListStrategy::Envvar),
+                device_sharing_strategy: None,
+                time_slicing: None,
+                device_partitioning_strategy: Some(NvidiaDevicePartitioningStrategy::MIG),
+                mig: None
+            }
+        );
+
+        let results = serde_json::to_string(&nvidia_device_plugins).unwrap();
+        assert_eq!(results, test_json);
+    }
+}
+
+#[derive(Debug, Clone, Eq, PartialEq, Hash)]
+pub struct NvidiaGPUModel {
+    inner: String,
+}
+
+lazy_static! {
+    pub(crate) static ref NVIDIAGPU_NAME: Regex = Regex::new(r"^([a-z])(\d+)\.(\d+)gb$").unwrap();
+}
+
+impl TryFrom<&str> for NvidiaGPUModel {
+    type Error = error::Error;
+
+    fn try_from(input: &str) -> Result<Self, Self::Error> {
+        ensure!(
+            NVIDIAGPU_NAME.is_match(input),
+            error::PatternSnafu {
+                thing: "NVIDIA GPU Model",
+                pattern: NVIDIAGPU_NAME.clone(),
+                input
+            }
+        );
+
+        Ok(NvidiaGPUModel {
+            inner: input.to_string(),
+        })
+    }
+}
+
+string_impls_for!(NvidiaGPUModel, "NvidiaGPUModel");
+
+#[cfg(test)]
+mod test_valid_gpu_model {
+    use super::NvidiaGPUModel;
+    use std::convert::TryFrom;
+
+    #[test]
+    fn valid_gpu_model() {
+        for ok in &["a100.40gb", "a100.80gb", "h100.80gb", "h100.141gb"] {
+            assert!(NvidiaGPUModel::try_from(*ok).is_ok());
+        }
+    }
+
+    #[test]
+    fn invalid_gpu_model() {
+        assert!(NvidiaGPUModel::try_from("invalid").is_err());
+        assert!(NvidiaGPUModel::try_from("1000").is_err());
+    }
+}
+
+#[derive(Debug, Clone, Eq, PartialEq, Hash)]
+pub struct MIGProfile {
+    inner: String,
+}
+
+lazy_static! {
+    pub(crate) static ref MIGPROFILE_NAME: Regex = Regex::new(r"^[0-9]g\.\d+gb$").unwrap();
+}
+
+impl TryFrom<&str> for MIGProfile {
+    type Error = error::Error;
+
+    fn try_from(input: &str) -> Result<Self, Self::Error> {
+        let slice_format = matches!(input, "1" | "2" | "3" | "4" | "7");
+
+        ensure!(
+            slice_format | MIGPROFILE_NAME.is_match(input),
+            error::PatternSnafu {
+                thing: "MIG Profile",
+                pattern: MIGPROFILE_NAME.clone(),
+                input
+            }
+        );
+
+        Ok(MIGProfile {
+            inner: input.to_string(),
+        })
+    }
+}
+
+string_impls_for!(MIGProfile, "MIGProfile");
+
+#[cfg(test)]
+mod test_valid_mig_profile {
+    use super::MIGProfile;
+    use std::convert::TryFrom;
+
+    #[test]
+    fn valid_mig_profile() {
+        for ok in &[
+            "1g.5gb", "2g.10gb", "3g.20gb", "7g.40gb", "1g.10gb", "1g.20gb", "2g.20gb", "3g.40gb",
+            "7g.80gb", "1g.18gb", "1g.35gb", "2g.35gb", "3g.71gb", "7g.141gb", "1", "2", "3", "4",
+            "7",
+        ] {
+            assert!(MIGProfile::try_from(*ok).is_ok());
+        }
+    }
+
+    #[test]
+    fn invalid_mig_profile() {
+        assert!(MIGProfile::try_from("invalid").is_err());
+        assert!(MIGProfile::try_from("1000").is_err());
+    }
+}
+
+#[model(impl_default = true)]
+pub struct NvidiaMIGSettings {
+    profile: HashMap<NvidiaGPUModel, MIGProfile>,
+}
+
+#[cfg(test)]
+mod test_valid_nvidia_mig_settings {
+    use super::*;
+
+    #[test]
+    fn test_serde_nvidia_device_plugins_with_mig_profile() {
+        let test_json = r#"{"pass-device-specs":false,"device-id-strategy":"uuid","device-list-strategy":"envvar","device-partitioning-strategy":"mig","mig":{"profile":{"a100.40gb":"1g.5gb"}}}"#;
+        let nvidia_device_plugins: NvidiaDevicePluginSettings =
+            serde_json::from_str(test_json).unwrap();
+        assert_eq!(
+            nvidia_device_plugins,
+            NvidiaDevicePluginSettings {
+                pass_device_specs: Some(false),
+                device_id_strategy: Some(NvidiaDeviceIdStrategy::Uuid),
+                device_list_strategy: Some(NvidiaDeviceListStrategy::Envvar),
+                device_sharing_strategy: None,
+                time_slicing: None,
+                device_partitioning_strategy: Some(NvidiaDevicePartitioningStrategy::MIG),
+                mig: Some(NvidiaMIGSettings {
+                    profile: Some(HashMap::from([(
+                        NvidiaGPUModel::try_from("a100.40gb").unwrap(),
+                        MIGProfile::try_from("1g.5gb").unwrap()
+                    )]))
+                }),
+            }
+        );
+
+        let results = serde_json::to_string(&nvidia_device_plugins).unwrap();
+        assert_eq!(results, test_json);
+    }
+}

bottlerocket-settings-models/settings-extensions/kubelet-device-plugins/src/lib.rs

@@ -43,10 +43,12 @@ impl SettingsModel for KubeletDevicePluginsV1 {
 mod test {
     use super::*;
     use bottlerocket_modeled_types::{
-        NvidiaDeviceIdStrategy, NvidiaDeviceListStrategy, NvidiaDeviceSharingStrategy,
-        NvidiaTimeSlicingSettings,
+        MIGProfile, NvidiaDeviceIdStrategy, NvidiaDeviceListStrategy,
+        NvidiaDevicePartitioningStrategy, NvidiaDeviceSharingStrategy, NvidiaGPUModel,
+        NvidiaMIGSettings, NvidiaTimeSlicingSettings,
     };
     use bounded_integer::BoundedI32;
+    use std::collections::HashMap;
 
     #[test]
     fn test_generate_kubelet_device_plugins() {
@@ -59,7 +61,7 @@ mod test {
 
     #[test]
     fn test_serde_kubelet_device_plugins() {
-        let test_json = r#"{"nvidia":{"pass-device-specs":true,"device-id-strategy":"index","device-list-strategy":"volume-mounts","device-sharing-strategy":"time-slicing","time-slicing":{"replicas":2,"rename-by-default":true,"fail-requests-greater-than-one":true}}}"#;
+        let test_json = r#"{"nvidia":{"pass-device-specs":true,"device-id-strategy":"index","device-list-strategy":"volume-mounts","device-sharing-strategy":"time-slicing","time-slicing":{"replicas":2,"rename-by-default":true,"fail-requests-greater-than-one":true},"device-partitioning-strategy":"mig","mig":{"profile":{"a100.40gb":"1g.5gb"}}}}"#;
 
         let device_plugins: KubeletDevicePluginsV1 = serde_json::from_str(test_json).unwrap();
         assert_eq!(
@@ -75,6 +77,13 @@ mod test {
                         rename_by_default: Some(true),
                         fail_requests_greater_than_one: Some(true),
                     }),
+                    device_partitioning_strategy: Some(NvidiaDevicePartitioningStrategy::MIG),
+                    mig: Some(NvidiaMIGSettings {
+                        profile: Some(HashMap::from([(
+                            NvidiaGPUModel::try_from("a100.40gb").unwrap(),
+                            MIGProfile::try_from("1g.5gb").unwrap()
+                        )]))
+                    }),
                 }),
             }
         );

bottlerocket-settings-models/settings-models/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "bottlerocket-settings-models"
-version = "0.7.0"
+version = "0.8.0"
 authors = ["Tom Kirchner <tjk@amazon.com>"]
 license = "Apache-2.0 OR MIT"
 edition = "2021"