Add AMD GPU detection to ghostdog

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

Author: mgsharm

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))]

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)?;
@@ -203,6 +212,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)? {

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 {