Merge pull request #323 from arnaldo2792/p6-b300-qualification

Extend supported MIG profiles

Author: arnaldo2792

sources/nvidia-migmanager/src/main.rs

@@ -97,6 +97,7 @@ enum NvidiaGpu {
     H100_80GB,
     H200_141GB,
     B200_180GB,
+    B300_269GB,
     RtxPro6000_96GB,
     Other,
 }
@@ -214,6 +215,9 @@ fn get_gpu_model(pci_device_id: &str) -> Result<NvidiaGpu> {
     } else if pci_device_id.starts_with("0x2BB5") {
         info!("Found NVIDIA RTX PRO 6000 Blackwell GPU.");
         Ok(NvidiaGpu::RtxPro6000_96GB)
+    } else if pci_device_id.starts_with("0x3182") {
+        info!("Found NVIDIA B300-269GB GPU.");
+        Ok(NvidiaGpu::B300_269GB)
     } else {
         warn!("Found NVIDIA Device but couldn't confirm variant.");
         Ok(NvidiaGpu::Other)
@@ -409,6 +413,9 @@ fn enable_mig(mig_settings: NvidiaMigConfig, gpu_info: &[MigGpu]) -> Result<()>
         Ok(NvidiaGpu::RtxPro6000_96GB) => {
             process_mig_config::<NvidiaRtxPro6000_96gbMigProfile>("rtxpro6000.96gb", &mig_settings)
         }
+        Ok(NvidiaGpu::B300_269GB) => {
+            process_mig_config::<NvidiaB300_269gbMigProfile>("b300.269gb", &mig_settings)
+        }
         _ => {
             let known_gpus: Vec<&str> = vec![
                 "a100.40gb",

sources/nvidia-migmanager/src/mig_profile.rs

@@ -1,5 +1,9 @@
 use serde::Deserialize;
 
+pub(crate) trait MigGpuProfile: for<'de> Deserialize<'de> {
+    fn get_mig_profile(&self) -> &str;
+}
+
 #[derive(Deserialize)]
 pub(crate) enum NvidiaA100_40gbMigProfile {
     #[serde(alias = "1g.5gb")]
@@ -216,6 +220,42 @@ impl MigGpuProfile for NvidiaRtxPro6000_96gbMigProfile {
     }
 }
 
-pub(crate) trait MigGpuProfile: for<'de> Deserialize<'de> {
-    fn get_mig_profile(&self) -> &str;
+#[derive(Deserialize)]
+pub(crate) enum NvidiaB300_269gbMigProfile {
+    #[serde(alias = "1g.34gb")]
+    #[serde(alias = "7")]
+    Mig1g34gb,
+
+    #[serde(alias = "1g.67gb")]
+    #[serde(alias = "4")]
+    Mig1g67gb,
+
+    #[serde(alias = "2g.67gb")]
+    #[serde(alias = "3")]
+    Mig2g67gb,
+
+    #[serde(alias = "3g.135gb")]
+    #[serde(alias = "2")]
+    Mig3g135gb,
+
+    #[serde(alias = "4g.135gb")]
+    Mig4g135gb,
+
+    #[serde(alias = "7g.269gb")]
+    #[serde(alias = "1")]
+    #[serde(other)]
+    Mig7g269gb,
+}
+
+impl MigGpuProfile for NvidiaB300_269gbMigProfile {
+    fn get_mig_profile(&self) -> &str {
+        match self {
+            Self::Mig7g269gb => "7g.269gb",
+            Self::Mig4g135gb => "4g.135gb",
+            Self::Mig3g135gb => "3g.135gb,3g.135gb",
+            Self::Mig2g67gb => "2g.67gb,2g.67gb,2g.67gb",
+            Self::Mig1g67gb => "1g.67gb,1g.67gb,1g.67gb,1g.67gb",
+            Self::Mig1g34gb => "1g.34gb,1g.34gb,1g.34gb,1g.34gb,1g.34gb,1g.34gb,1g.34gb",
+        }
+    }
 }