Merge pull request #137 from Dstack-TEE/gpu-include

teepod: Add gpu whitelist

Author: kvinwang

teepod/src/config.rs

@@ -111,12 +111,23 @@ pub struct GpuConfig {
     pub listing: Vec<String>,
     /// The PCI addresses to exclude from passthrough
     pub exclude: Vec<String>,
+    /// The PCI addresses to include in passthrough
+    pub include: Vec<String>,
 }
 
 impl GpuConfig {
     pub(crate) fn list_devices(&self) -> Result<Vec<Device>> {
         let devices = lspci_filtered(|dev| {
-            self.listing.contains(&dev.full_product_id()) && !self.exclude.contains(&dev.slot)
+            if !self.listing.contains(&dev.full_product_id()) {
+                return false;
+            }
+            if self.exclude.contains(&dev.slot) {
+                return false;
+            }
+            if !self.include.is_empty() && !self.include.contains(&dev.slot) {
+                return false;
+            }
+            true
         })
         .context("Failed to list GPU devices")?;
 

teepod/teepod.toml

@@ -45,6 +45,8 @@ enabled = false
 listing = ["10de:2335"]
 # The PCI addresses of the cards to exclude
 exclude = []
+# The PCI addresses of the cards to include
+include = []
 
 [gateway]
 base_domain = "localhost"