ghostdog: Add AMD GPU driver validation

Add match-driver support for AMD GPUs to validate that the amdgpu
kernel module and ROCm KFD driver are properly loaded before starting
the device plugin service.

Signed-off-by: Gaurav Sharma <[email protected]>

Author: mgsharm

packages/amd-k8s-device-plugin/amd-k8s-device-plugin.service

@@ -2,8 +2,15 @@
 Description=Start AMD kubernetes device plugin
 After=kubelet.service
 Wants=kubelet.service
+ConditionPathExists=/usr/bin/ghostdog
 
 [Service]
+# Verify AMD GPU is detected
+ExecStartPre=/usr/bin/ghostdog amd-gpu-present
+# Verify amdgpu driver is loaded
+ExecStartPre=/usr/bin/ghostdog match-driver amd rocm
+# Ensure that the kubelet device plugin socket exists before we start
+# A brief sleep is needed to avoid the `test` failing its first check
 ExecStartPre=/usr/bin/sleep 0.1
 ExecStartPre=/usr/bin/test -S /var/lib/kubelet/device-plugins/kubelet.sock
 ExecStart=/usr/bin/amd-device-plugin -logtostderr=true -stderrthreshold=INFO -v=5

sources/ghostdog/src/error.rs

@@ -32,10 +32,14 @@ pub(super) enum Error {
     CheckEfaFailure { source: pciclient::PciClientError },
     #[snafu(display("Failed to check if Neuron device is attached: {}", source))]
     CheckNeuronFailure { source: pciclient::PciClientError },
+    #[snafu(display("Failed to check if AMD GPU device is attached: {}", source))]
+    CheckAmdGpuFailure { source: pciclient::PciClientError },
     #[snafu(display("Did not detect EFA"))]
     NoEfaPresent,
     #[snafu(display("Did not detect Neuron"))]
     NoNeuronPresent,
+    #[snafu(display("Did not detect AMD GPU"))]
+    NoAmdGpuPresent,
     #[snafu(display("'{}' has no parent directory", path.display()))]
     NoParentDirectory { path: std::path::PathBuf },
     #[snafu(display("Failed to open '{}': {}", path.display(), source))]
@@ -78,6 +82,8 @@ pub(super) enum Error {
     NoInf1Present,
     #[snafu(display("Did not detect inf2+ hardware"))]
     NoInf2Present,
+    #[snafu(display("AMD GPU driver not loaded: {}", reason))]
+    AmdDriverNotLoaded { reason: String },
 }
 
 pub(crate) type Result<T> = std::result::Result<T, Error>;

sources/ghostdog/src/main.rs

@@ -58,6 +58,7 @@ enum SubCommand {
     EbsDeviceName(EbsDeviceNameArgs),
     EfaPresent(EfaPresentArgs),
     NeuronPresent(NeuronPresentArgs),
+    AmdGpuPresent(AmdGpuPresentArgs),
     MatchDriver(MatchDriverArgs),
     MatchNvidiaDriver(MatchNvidiaDriverArgs),
     WriteInfinibandGuid(WriteInfinibandGuidArgs),
@@ -73,6 +74,11 @@ struct EfaPresentArgs {}
 /// Detect if Neuron devices are attached.
 struct NeuronPresentArgs {}
 
+#[derive(FromArgs, PartialEq, Debug)]
+#[argh(subcommand, name = "amd-gpu-present")]
+/// Detect if AMD GPU devices are attached.
+struct AmdGpuPresentArgs {}
+
 #[derive(FromArgs, PartialEq, Debug)]
 #[argh(subcommand, name = "scan")]
 /// Scan a device to see if it is an ephemeral disk.
@@ -163,6 +169,9 @@ fn main() -> Result<()> {
         SubCommand::NeuronPresent(_) => {
             is_neuron_attached()?;
         }
+        SubCommand::AmdGpuPresent(_) => {
+            is_amd_gpu_attached()?;
+        }
         SubCommand::MatchNvidiaDriver(driver) => {
             let driver_name = driver.driver_name;
             nvidia_driver_supported(&driver_name)?;
@@ -173,6 +182,7 @@ fn main() -> Result<()> {
             match driver_name.as_str() {
                 "nvidia" => nvidia_driver_supported(&flavor_name)?,
                 "neuron" => match_neuron_driver(&flavor_name)?,
+                "amd" => match_amd_driver(&flavor_name)?,
                 _ => {
                     return Err(error::Error::UnsupportedDriver {
                         driver: driver_name.to_string(),
@@ -203,6 +213,14 @@ fn is_neuron_attached() -> Result<()> {
     }
 }
 
+fn is_amd_gpu_attached() -> Result<()> {
+    if pciclient::is_amd_gpu_attached().context(error::CheckAmdGpuFailureSnafu)? {
+        Ok(())
+    } else {
+        Err(error::Error::NoAmdGpuPresent)
+    }
+}
+
 // Returns true if this is an inf1 instance
 fn is_inf1_instance() -> Result<()> {
     if pciclient::is_inf1_instance().context(error::CheckInf1FailureSnafu)? {
@@ -412,6 +430,35 @@ fn match_neuron_driver(driver_flavor: &str) -> Result<()> {
     }
 }
 
+fn match_amd_driver(driver_flavor: &str) -> Result<()> {
+    match driver_flavor {
+        "rocm" => amd_driver_loaded(),
+        _ => Err(error::Error::UnsupportedDriverFlavor {
+            driver: "amd".to_string(),
+            flavor: driver_flavor.to_string(),
+        }),
+    }
+}
+
+fn amd_driver_loaded() -> Result<()> {
+    let kfd_path = "/sys/class/kfd";
+    let amdgpu_module_path = "/sys/module/amdgpu";
+    
+    if !std::path::Path::new(kfd_path).exists() {
+        return Err(error::Error::AmdDriverNotLoaded {
+            reason: format!("KFD device path {} not found - ROCm kernel driver not loaded", kfd_path),
+        });
+    }
+    
+    if !std::path::Path::new(amdgpu_module_path).exists() {
+        return Err(error::Error::AmdDriverNotLoaded {
+            reason: format!("amdgpu kernel module not loaded ({} not found)", amdgpu_module_path),
+        });
+    }
+    
+    Ok(())
+}
+
 // Known system partition types for Bottlerocket.
 lazy_static! {
     static ref SYSTEM_PARTITION_TYPES: HashSet<[u8; 16]> = [

sources/pciclient/src/lib.rs

@@ -9,7 +9,7 @@ mod private;
 
 use private::{
     call_list_devices, check_efa_attachment, check_inf1_attachment, check_inf2_attachment,
-    check_neuron_attachment, PciClient,
+    check_neuron_attachment, check_amd_gpu_attachment, PciClient,
 };
 
 use bon::Builder;
@@ -123,6 +123,11 @@ pub fn is_neuron_attached() -> Result<bool> {
     check_neuron_attachment(PciClient {})
 }
 
+/// Call `lspci` and check if there is any AMD GPU device attached.
+pub fn is_amd_gpu_attached() -> Result<bool> {
+    check_amd_gpu_attachment(PciClient {})
+}
+
 /// Call `lspci` and check if there are inf1 devices attached
 pub fn is_inf1_instance() -> Result<bool> {
     check_inf1_attachment(PciClient {})

sources/pciclient/src/private.rs

@@ -14,6 +14,7 @@ use crate::{
 
 const AMAZON_VENDOR_CODE: &str = "1d0f";
 const EFA_KEYWORD: &str = "efa";
+const AMD_VENDOR_CODE: &str = "1002";
 
 // Neuron devices specific to inf1 instance types
 lazy_static! {
@@ -34,6 +35,12 @@ lazy_static! {
     };
 }
 
+lazy_static! {
+    static ref AMD_GPU_DEVICES: HashSet<&'static str> = hashset! {
+        "75a3" // AMD Instinct MI355X
+    };
+}
+
 const LSPCI_PATH: &str = "/usr/bin/lspci";
 
 pub(crate) trait CommandExecutor {
@@ -183,6 +190,19 @@ pub(crate) fn check_neuron_attachment<T: CommandExecutor>(command_executor: T) -
     }))
 }
 
+/// Call `lspci` and check if there is any AMD GPU device attached.
+/// Internal usage, adding command_executor as parameter allows us to better unit test.
+/// For external usage, check [`crate::is_amd_gpu_attached`].
+pub(crate) fn check_amd_gpu_attachment<T: CommandExecutor>(command_executor: T) -> Result<bool> {
+    let list_devices_param = ListDevicesParam::builder()
+        .vendor(AMD_VENDOR_CODE.to_string())
+        .build();
+    let list_device_results = call_list_devices(command_executor, list_devices_param)?;
+    Ok(list_device_results
+        .iter()
+        .any(|device_info| AMD_GPU_DEVICES.contains(&device_info.device.as_str())))
+}
+
 /// Call `lspci` and check if there is any inf1-specific Neuron device attached.
 /// Internal usage, adding command_executor as parameter allows us to better unit test.
 /// For external usage, check [`crate::is_inf1_instance`].
@@ -229,7 +249,7 @@ mod test {
 
     use super::{
         call_list_devices, check_efa_attachment, check_inf1_attachment, check_neuron_attachment,
-        parse_list_devices_output, CommandExecutor,
+        check_amd_gpu_attachment, parse_list_devices_output, CommandExecutor,
     };
 
     struct MockPciClient {
@@ -436,6 +456,31 @@ mod test {
         assert!(check_inf1_attachment_result.unwrap());
     }
 
+    #[test]
+    fn test_is_amd_gpu_attached() {
+        let mock_pci_client = MockPciClient {
+            // AMD Instinct MI355X device has vendor 1002 for AMD, device code 75a3.
+            output: vec![
+                r#"00:1e.0 "0300" "1002" "75a3" -p00 "1002" "0123""#.to_string(),
+                r#"00:1f.0 "0302" "10de" "1eb8" -ra1 -p00 "10de" "12a2""#.to_string(),
+            ],
+        };
+        let check_amd_gpu_attachment_result = check_amd_gpu_attachment(mock_pci_client);
+        assert!(check_amd_gpu_attachment_result.is_ok());
+        assert!(check_amd_gpu_attachment_result.unwrap());
+    }
+   
+    #[test]
+    fn test_is_amd_gpu_attached_negative_case() {
+        let mock_pci_client = MockPciClient {
+            // Below is an actual output from lspci for ena device (not AMD GPU).
+            output: vec![r#"00:06.0 "0200" "1d0f" "ec20" -p00 "1d0f" "ec20""#.to_string()],
+        };
+        let check_amd_gpu_attachment_result = check_amd_gpu_attachment(mock_pci_client);
+        assert!(check_amd_gpu_attachment_result.is_ok());
+        assert!(!check_amd_gpu_attachment_result.unwrap());
+    }
+
     #[test]
     fn test_is_inf1_attached_negative_case() {
         let mock_pci_client = MockPciClient {