feat: NIXL stub support

Signed-off-by: Alexandre Milesi <[email protected]>

Author: milesial

Cargo.lock

@@ -97,7 +97,7 @@ dialoguer = { version = "0.11", default-features = false, features = [
 
 # block_manager
 aligned-vec = { version = "0.6.4", optional = true }
-nixl-sys = { version = "=0.7.0", optional = true }
+nixl-sys = { git = "https://github.com/ai-dynamo/nixl", rev = "ae3f8af", optional = true }
 cudarc = { workspace = true, optional = true }
 nix = { version = "0.26", optional = true }
 

lib/bindings/python/Cargo.lock

@@ -47,8 +47,6 @@ pub struct ImageMetadata {
     #[allow(dead_code)] // used in followup MR
     pub(crate) format: Option<ImageFormat>,
     #[allow(dead_code)] // used in followup MR
-    pub(crate) color_type: ColorType,
-    #[allow(dead_code)] // used in followup MR
     pub(crate) layout: ImageLayout,
 }
 
@@ -82,7 +80,6 @@ impl Decoder for ImageDecoder {
         let mut decoded: DecodedMediaData = array.try_into()?;
         decoded.tensor_info.metadata = Some(DecodedMediaMetadata::Image(ImageMetadata {
             format,
-            color_type,
             layout: ImageLayout::HWC,
         }));
         Ok(decoded)

lib/llm/Cargo.toml

@@ -41,7 +41,7 @@ pub struct MediaLoader {
     media_decoder: MediaDecoder,
     http_client: reqwest::Client,
     media_fetcher: MediaFetcher,
-    nixl_agent: NixlAgent,
+    nixl_agent: Option<NixlAgent>,
 }
 
 impl MediaLoader {
@@ -55,7 +55,15 @@ impl MediaLoader {
 
         let http_client = http_client_builder.build()?;
 
-        let nixl_agent = get_nixl_agent()?;
+        let nixl_agent = match get_nixl_agent() {
+            Ok(agent) => Some(agent),
+            Err(e) => {
+                tracing::warn!(
+                    "Error when creating NIXL agent (will not be able to register media data): {e}"
+                );
+                None
+            }
+        };
 
         Ok(Self {
             media_decoder,
@@ -96,6 +104,10 @@ impl MediaLoader {
         oai_content_part: &ChatCompletionRequestUserMessageContentPart,
         // TODO: request-level options
     ) -> Result<RdmaMediaDataDescriptor> {
+        if self.nixl_agent.is_none() {
+            anyhow::bail!("NIXL agent is not available, cannot decode and register media data");
+        }
+
         // fetch the media, decode and NIXL-register
         let decoded = match oai_content_part {
             ChatCompletionRequestUserMessageContentPart::ImageUrl(image_part) => {
@@ -116,7 +128,7 @@ impl MediaLoader {
             _ => anyhow::bail!("Unsupported media type"),
         };
 
-        let rdma_descriptor = decoded.into_rdma_descriptor(&self.nixl_agent)?;
+        let rdma_descriptor = decoded.into_rdma_descriptor(self.nixl_agent.as_ref().unwrap())?;
         Ok(rdma_descriptor)
     }
 }
@@ -148,21 +160,22 @@ mod tests {
             ..Default::default()
         };
 
-        let loader = MediaLoader::new(media_decoder, fetcher).unwrap();
+        let loader: MediaLoader = MediaLoader::new(media_decoder, fetcher).unwrap();
 
         let image_url = ImageUrl::from(format!("{}/llm-optimize-deploy-graphic.png", server.url()));
         let content_part = ChatCompletionRequestUserMessageContentPart::ImageUrl(
             ChatCompletionRequestMessageContentPartImage { image_url },
         );
 
         let result = loader.fetch_and_decode_media_part(&content_part).await;
-        assert!(
-            result.is_ok(),
-            "Failed to fetch and decode image: {:?}",
-            result.err()
-        );
-
-        let descriptor = result.unwrap();
+        let descriptor = match result {
+            Ok(descriptor) => descriptor,
+            Err(e) if e.to_string().contains("NIXL agent is not available") => {
+                eprintln!("Skipping test: NIXL agent not available");
+                return;
+            }
+            Err(e) => panic!("Failed to fetch and decode image: {}", e),
+        };
         assert_eq!(descriptor.tensor_info.dtype, DataType::UINT8);
 
         // Verify image dimensions: 1,999px × 1,125px (width × height)

lib/llm/src/preprocessor/media/decoders/image.rs

@@ -26,7 +26,7 @@ dynamo-config = { workspace = true }
 
 anyhow = { workspace = true }
 cudarc = { workspace = true }
-nixl-sys = { version = "0.7" }
+nixl-sys = { git = "https://github.com/ai-dynamo/nixl", rev = "ae3f8af" }
 serde = { workspace = true}
 thiserror = { workspace = true }
 tracing = { workspace = true }

lib/llm/src/preprocessor/media/loader.rs

@@ -32,6 +32,10 @@ pub struct NixlAgent {
 impl NixlAgent {
     /// Create a NIXL agent without any backends.
     pub fn new(name: &str) -> Result<Self> {
+        if nixl_sys::is_stub() {
+            return Err(anyhow::anyhow!("NIXL is stubbed, cannot create agent"));
+        }
+
         let agent = Agent::new(name)?;
 
         Ok(Self {