diff --git a/src/runtime-rs/crates/resource/src/cdi_devices/container_device.rs b/src/runtime-rs/crates/resource/src/cdi_devices/container_device.rs new file mode 100644 index 0000000000..867d02f15a --- /dev/null +++ b/src/runtime-rs/crates/resource/src/cdi_devices/container_device.rs @@ -0,0 +1,108 @@ +// +// Copyright (c) 2024 Ant Group +// +// SPDX-License-Identifier: Apache-2.0 +// + +use std::collections::HashMap; +use std::path::Path; + +use anyhow::Result; +use oci_spec::runtime::Spec; + +use super::{resolve_cdi_device_kind, ContainerDevice}; +use agent::types::Device; + +const CDI_PREFIX: &str = "cdi.k8s.io"; + +// Sort the devices based on the first element's PCI_Guest_Path in the PCI bus according to options. +fn sort_devices_by_guest_pcipath(devices: &mut [ContainerDevice]) { + // Extract first guest_pcipath from device_options + let extract_first_guest_pcipath = |options: &[String]| -> Option { + options + .first() + .and_then(|option| option.split('=').nth(1)) + .map(|path| path.to_string()) + }; + + devices.sort_by(|a, b| { + let guest_path_a = extract_first_guest_pcipath(&a.device.options); + let guest_path_b = extract_first_guest_pcipath(&b.device.options); + + guest_path_a.cmp(&guest_path_b) + }); +} + +// Annotate container devices with CDI annotations in OCI Spec +pub fn annotate_container_devices( + spec: &mut Spec, + container_devices: Vec, +) -> Result> { + let mut devices_agent: Vec = Vec::new(); + // Make sure that annotations is Some(). + if spec.annotations().is_none() { + spec.set_annotations(Some(HashMap::new())); + } + + // Step 1: Extract all devices and filter out devices without device_info for vfio_devices + let vfio_devices: Vec = container_devices + .into_iter() + .map(|device| { + // push every device's Device to agent_devices + devices_agent.push(device.device.clone()); + device + }) + .filter(|device| device.device_info.is_some()) + .collect(); + + // Step 2: Group devices by vendor_id-class_id + let mut grouped_devices: HashMap> = HashMap::new(); + for device in vfio_devices { + // Extract the vendor/class key and insert into the map if both are present + if let Some(key) = device + .device_info + .as_ref() + .and_then(|info| resolve_cdi_device_kind(&info.vendor_id, &info.class_id)) + { + grouped_devices + .entry(key.to_owned()) + .or_default() + .push(device); + } + } + + // Step 3: Sort devices within each group by guest_pcipath + grouped_devices + .iter_mut() + .for_each(|(vendor_class, container_devices)| { + // The *offset* is a monotonically increasing counter that keeps track of the number of devices + // within an IOMMU group. It increments by total_of whenever a new IOMMU group is processed. + let offset: &mut usize = &mut 0; + + sort_devices_by_guest_pcipath(container_devices); + container_devices + .iter() + .enumerate() + .for_each(|(base, container_device)| { + let total_of = container_device.device.options.len(); + // annotate device with cdi information in OCI Spec. + for index in 0..total_of { + if let Some(iommu_grpid) = + Path::new(&container_device.device.container_path) + .file_name() + .and_then(|name| name.to_str()) + { + spec.annotations_mut().as_mut().unwrap().insert( + format!("{}/vfio{}.{}", CDI_PREFIX, iommu_grpid, index), // cdi.k8s.io/vfioX.y + format!("{}={}", vendor_class, base + *offset), // vendor/class=name + ); + } + } + + // update the offset with *total_of*. + *offset += total_of - 1; + }); + }); + + Ok(devices_agent) +} diff --git a/src/runtime-rs/crates/resource/src/cdi_devices/mod.rs b/src/runtime-rs/crates/resource/src/cdi_devices/mod.rs new file mode 100644 index 0000000000..60edf93129 --- /dev/null +++ b/src/runtime-rs/crates/resource/src/cdi_devices/mod.rs @@ -0,0 +1,63 @@ +// +// Copyright (c) 2024 Ant Group +// +// SPDX-License-Identifier: Apache-2.0 +// + +pub mod container_device; + +use agent::types::Device; +use std::collections::HashMap; +use std::path::PathBuf; + +#[derive(Clone, Default)] +pub struct DeviceInfo { + pub class_id: String, + pub vendor_id: String, + pub host_path: PathBuf, +} + +#[derive(Clone, Default)] +pub struct ContainerDevice { + pub device_info: Option, + pub device: Device, +} + +lazy_static! { + // *CDI_DEVICE_KIND_TABLE* is static hash map to store a mapping between device vendor and class + // identifiers and their corresponding CDI vendor and class strings. This mapping is essentially a + // lookup table that allows the system to determine the appropriate CDI for a given device based on + // its vendor and class information. + // Note: Our device mapping is designed to be flexible and responsive to user needs. The current list + // is not exhaustive and will be updated as required. + pub static ref CDI_DEVICE_KIND_TABLE: HashMap<&'static str, &'static str> = { + let mut m = HashMap::new(); + m.insert("0x10de-0x030", "nvidia.com/gpu"); + m.insert("0x8086-0x030", "intel.com/gpu"); + m.insert("0x1002-0x030", "amd.com/gpu"); + m.insert("0x15b3-0x020", "nvidia.com/nic"); + // TODO: it will be updated as required. + m + }; +} + +// Sort devices by guest_pcipath +pub fn sort_options_by_pcipath(mut device_options: Vec) -> Vec { + device_options.sort_by(|a, b| { + let extract_path = |s: &str| s.split('=').nth(1).map(|path| path.to_string()); + let guest_path_a = extract_path(a); + let guest_path_b = extract_path(b); + + guest_path_a.cmp(&guest_path_b) + }); + device_options +} + +// Resolve the CDI vendor ID/device Class by a lookup table based on the provided vendor and class. +pub fn resolve_cdi_device_kind<'a>(vendor_id: &'a str, class_id: &'a str) -> Option<&'a str> { + let vendor_class = format!("{}-{}", vendor_id, class_id); + // The first 12 characters of the string ("0x10de-0x030") provide a concise + // and clear identification of both the manufacturer and the device category. + // it returns "nvidia.com/gpu", "amd.com/gpu" or others. + CDI_DEVICE_KIND_TABLE.get(&vendor_class[..12]).copied() +} diff --git a/src/runtime-rs/crates/resource/src/lib.rs b/src/runtime-rs/crates/resource/src/lib.rs index 3a7524da9c..413b1c6512 100644 --- a/src/runtime-rs/crates/resource/src/lib.rs +++ b/src/runtime-rs/crates/resource/src/lib.rs @@ -23,6 +23,7 @@ pub mod rootfs; pub mod share_fs; pub mod volume; pub use manager::ResourceManager; +pub mod cdi_devices; pub mod cpu_mem; use kata_types::config::hypervisor::SharedFsInfo;