NIXL data passing, install ffmpeg

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

Author: milesial

container/Dockerfile

@@ -125,6 +125,11 @@ RUN apt-get update -y \
         clang \
         libclang-dev \
         protobuf-compiler \
+        # media-loading rust build+runtime dependencies
+        libavcodec-dev \
+        libavutil-dev \
+        libavformat-dev \
+        pkg-config \
     && apt-get clean \
     && rm -rf /var/lib/apt/lists/*
 

container/Dockerfile.vllm

@@ -305,6 +305,11 @@ ARG WORKSPACE_DIR=/workspace
 RUN apt-get update -y && \
     apt-get install -y --no-install-recommends  \
     # Install utilities
+    wget \
+    && rm -f /etc/apt/sources.list.d/cuda*.list \
+    && wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb && dpkg -i cuda-keyring_1.1-1_all.deb && \
+    apt-get update -y && \
+    apt-get install -y --no-install-recommends \
     nvtop \
     wget \
     tmux \

lib/bindings/python/src/dynamo/nixl_connect/__init__.py

@@ -16,6 +16,7 @@
 from __future__ import annotations
 
 import asyncio
+import base64
 import logging
 import socket
 import uuid
@@ -1211,7 +1212,7 @@ def metadata(self) -> RdmaMetadata:
 
             self._serialized_request = RdmaMetadata(
                 descriptors=descriptors,
-                nixl_metadata=nixl_metadata.hex(),
+                nixl_metadata=base64.b64encode(nixl_metadata).decode("utf-8"),
                 notification_key=self._notification_key,
                 operation_kind=int(self._operation_kind),
             )
@@ -1471,13 +1472,18 @@ def __init__(
         self._connector = connector
 
         # When `nixl_metadata` is a string, it is assumed to have come from a remote worker
-        # via a `RdmaMetadata` object and therefore can assumed be a hex-encoded, compressed
+        # via a `RdmaMetadata` object and therefore can assumed be a base64-encoded, compressed
         # representation of the NIXL metadata.
         if isinstance(nixl_metadata, str):
-            # Decode the hex-encoded string into bytes.
-            nixl_metadata = bytes.fromhex(nixl_metadata)
-            # Decompress the NIXL metadata.
-            nixl_metadata = zlib.decompress(nixl_metadata)
+            # Decode the base64-encoded string into bytes.
+            nixl_metadata = base64.b64decode(nixl_metadata)
+            try:
+                # Decompress the NIXL metadata.
+                nixl_metadata = zlib.decompress(nixl_metadata)
+            except zlib.error:
+                # we accept cases when the data might not be compressed
+                if nixl_metadata[0] == 0x78 and nixl_metadata[1] in (0x01, 0x5E, 0x9C, 0xDA):
+                    raise
 
         self._name = connector._nixl.add_remote_agent(nixl_metadata)
         if isinstance(self._name, bytes):

lib/llm/Cargo.toml

@@ -24,7 +24,7 @@ testing-etcd = []
 block-manager = ["dep:nixl-sys", "dep:cudarc", "dep:ndarray", "dep:nix"]
 cuda = ["dep:cudarc"]
 integration = ["dynamo-runtime/integration"]
-media-loading = ["dep:ndarray", "dep:video-rs", "dep:image", "dep:reqwest", "dep:base64", "dep:tokio-rayon"]
+media-loading = ["dep:ndarray", "dep:video-rs", "dep:image", "dep:reqwest", "dep:base64", "dep:tokio-rayon", "block-manager"]
 
 [[bench]]
 name = "tokenizer"

lib/llm/src/block_manager/storage.rs

@@ -364,6 +364,30 @@ impl SystemStorage {
     }
 }
 
+impl TryFrom<Vec<u8>> for SystemStorage {
+    type Error = StorageError;
+
+    /// Create SystemStorage from an existing Vec<u8>
+    /// Takes ownership of the Vec and uses its memory directly (zero-copy)
+    fn try_from(mut vec: Vec<u8>) -> Result<Self, Self::Error> {
+        let size = vec.len();
+        let layout =
+            Layout::array::<u8>(size).map_err(|e| StorageError::AllocationFailed(e.to_string()))?;
+        let ptr = NonNull::new(vec.as_mut_ptr())
+            .ok_or_else(|| StorageError::AllocationFailed("vec pointer is null".into()))?;
+
+        // prevents Vec from freeing the memory
+        std::mem::forget(vec);
+
+        Ok(Self {
+            ptr,
+            layout,
+            len: size,
+            handles: RegistrationHandles::new(),
+        })
+    }
+}
+
 impl Drop for SystemStorage {
     fn drop(&mut self) {
         self.handles.release();

lib/llm/src/preprocessor.rs

@@ -315,11 +315,11 @@ impl OpenAIPreprocessor {
                     _ => continue,
                 };
 
-                let map_item = media_map.entry(type_str.clone()).or_insert_with(Vec::new);
+                let map_item = media_map.entry(type_str.clone()).or_default();
 
                 if let Some(loader) = &self.media_loader {
-                    let decoded_data = loader.fetch_and_decode_media_part(content_part).await?;
-                    map_item.push(MultimodalData::Decoded(decoded_data));
+                    let rdma_descriptor = loader.fetch_and_decode_media_part(content_part).await?;
+                    map_item.push(MultimodalData::Decoded(rdma_descriptor));
                 } else {
                     map_item.push(MultimodalData::Url(url));
                 }

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

@@ -4,17 +4,45 @@
 use anyhow::Result;
 use base64::{Engine as _, engine::general_purpose};
 use ndarray::{ArrayBase, Dimension, OwnedRepr};
+use serde::{Deserialize, Serialize};
+use std::sync::Arc;
 
 use dynamo_async_openai::types::ChatCompletionRequestUserMessageContentPart;
 
+use crate::block_manager::storage::{
+    StorageError, SystemStorage, nixl::NixlRegisterableStorage, nixl::NixlStorage,
+};
 use crate::preprocessor::media::{ImageDecoder, VideoDecoder};
+use nixl_sys::Agent as NixlAgent;
 
 // Raw encoded media data (.png, .mp4, ...), optionally b64-encoded
 pub struct EncodedMediaData {
     bytes: Vec<u8>,
     b64_encoded: bool,
 }
 
+// Decoded media data (image RGB, video frames pixels, ...)
+pub struct DecodedMediaData {
+    data: SystemStorage,
+    shape: Vec<usize>,
+    dtype: String,
+}
+
+// Decoded media data NIXL descriptor (sent to the next step in the pipeline / NATS)
+#[derive(Serialize, Deserialize, Clone, Debug)]
+pub struct RdmaMediaDataDescriptor {
+    // b64 agent metadata
+    nixl_metadata: String,
+    // tensor descriptor
+    nixl_descriptor: NixlStorage,
+    shape: Vec<usize>,
+    dtype: String,
+    // reference to the actual data, kept alive while the rdma descriptor is alive
+    #[serde(skip, default)]
+    #[allow(dead_code)]
+    source_storage: Option<Arc<SystemStorage>>,
+}
+
 impl EncodedMediaData {
     // Handles both web URLs (will download the bytes) and data URLs (will keep b64-encoded)
     // This function is kept in tokio runtime so we do not want any expensive operations
@@ -56,25 +84,42 @@ impl EncodedMediaData {
     }
 }
 
-// Decoded media data (image RGB, video frames pixels, ...)
-#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
-pub struct DecodedMediaData {
-    data: Vec<u8>,
-    shape: Vec<usize>,
-    dtype: String,
+impl DecodedMediaData {
+    pub fn into_rdma_descriptor(self, nixl_agent: &NixlAgent) -> Result<RdmaMediaDataDescriptor> {
+        // get NIXL metadata and descriptor
+        let mut source_storage = self.data;
+        source_storage.nixl_register(nixl_agent, None)?;
+        let nixl_descriptor = unsafe { source_storage.as_nixl_descriptor() }
+            .ok_or_else(|| anyhow::anyhow!("Cannot convert storage to NIXL descriptor"))?;
+
+        // TODO: cache this if this is constant across the worker lifetime?
+        let nixl_local_md = nixl_agent.get_local_md()?;
+        let nixl_metadata = general_purpose::STANDARD.encode(&nixl_local_md);
+
+        Ok(RdmaMediaDataDescriptor {
+            nixl_metadata,
+            nixl_descriptor,
+            shape: self.shape,
+            dtype: self.dtype,
+            // do not drop / free the storage yet
+            source_storage: Some(Arc::new(source_storage)),
+        })
+    }
 }
 
 // convert Array{N}<u8> to DecodedMediaData
 // TODO: Array1<f32> for audio
-impl<D: Dimension> From<ArrayBase<OwnedRepr<u8>, D>> for DecodedMediaData {
-    fn from(array: ArrayBase<OwnedRepr<u8>, D>) -> Self {
+impl<D: Dimension> TryFrom<ArrayBase<OwnedRepr<u8>, D>> for DecodedMediaData {
+    type Error = StorageError;
+
+    fn try_from(array: ArrayBase<OwnedRepr<u8>, D>) -> Result<Self, Self::Error> {
         let shape = array.shape().to_vec();
         let (data, _) = array.into_raw_vec_and_offset();
-        Self {
-            data,
+        Ok(Self {
+            data: SystemStorage::try_from(data)?,
             shape,
             dtype: "uint8".to_string(),
-        }
+        })
     }
 }
 
@@ -102,28 +147,36 @@ pub struct MediaDecoder {
 pub struct MediaLoader {
     media_decoder: MediaDecoder,
     http_client: reqwest::Client,
+    nixl_agent: NixlAgent,
 }
 
 impl MediaLoader {
     pub fn new(media_decoder: MediaDecoder) -> Result<Self> {
         let http_client = reqwest::Client::builder()
             .user_agent(
-                "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:143.0) Gecko/20100101 Firefox/143.0",
+                "dynamo-ai/dynamo", // TODO: use a proper user agent
             )
             .build()?;
 
+        let uuid = uuid::Uuid::new_v4();
+        let nixl_agent = NixlAgent::new(&format!("media-loader-{}", uuid))?;
+        let (_, ucx_params) = nixl_agent.get_plugin_params("UCX")?;
+        nixl_agent.create_backend("UCX", &ucx_params)?;
+
         Ok(Self {
             media_decoder,
             http_client,
+            nixl_agent,
         })
     }
 
     pub async fn fetch_and_decode_media_part(
         &self,
         oai_content_part: &ChatCompletionRequestUserMessageContentPart,
-    ) -> Result<DecodedMediaData> {
+    ) -> Result<RdmaMediaDataDescriptor> {
         // TODO: request-level options
-        match oai_content_part {
+        // fetch and decode the media
+        let decoded = match oai_content_part {
             ChatCompletionRequestUserMessageContentPart::ImageUrl(image_part) => {
                 let url = &image_part.image_url.url;
                 let data = EncodedMediaData::from_url(url, &self.http_client).await?;
@@ -138,6 +191,9 @@ impl MediaLoader {
                 anyhow::bail!("Audio decoding is not supported yet");
             }
             _ => anyhow::bail!("Unsupported media type"),
-        }
+        }?;
+
+        let rdma_descriptor = decoded.into_rdma_descriptor(&self.nixl_agent)?;
+        Ok(rdma_descriptor)
     }
 }

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

@@ -36,6 +36,6 @@ impl Decoder for ImageDecoder {
         };
         let shape = (height as usize, width as usize, n_channels as usize);
         let array = Array3::from_shape_vec(shape, data)?;
-        Ok(array.into())
+        Ok(array.try_into()?)
     }
 }

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

@@ -5,6 +5,6 @@ mod common;
 mod image;
 mod video;
 
-pub use common::{DecodedMediaData, Decoder, EncodedMediaData, MediaDecoder, MediaLoader};
+pub use common::{Decoder, EncodedMediaData, MediaDecoder, MediaLoader, RdmaMediaDataDescriptor};
 pub use image::ImageDecoder;
 pub use video::VideoDecoder;

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

@@ -111,6 +111,6 @@ impl Decoder for VideoDecoder {
             3,
         );
         let array = Array4::from_shape_vec(shape, all_frames)?;
-        Ok(array.into())
+        Ok(array.try_into()?)
     }
 }