chore: Cleanup restructure

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

Author: milesial

Cargo.lock

@@ -303,7 +303,7 @@ ENV CARGO_BUILD_JOBS=${CARGO_BUILD_JOBS:-16} \
 # Install system dependencies
 RUN dnf install -y https://download1.rpmfusion.org/free/el/rpmfusion-free-release-8.noarch.rpm && dnf install -y https://download1.rpmfusion.org/nonfree/el/rpmfusion-nonfree-release-8.noarch.rpm
 RUN dnf update -y \
-    && dnf install -y llvm-toolset protobuf-compiler wget unzip ffmpeg-devel \
+    && dnf install -y llvm-toolset protobuf-compiler wget unzip libavdevice-dev libavutil-dev libavcodec-dev libavformat-dev pkg-config \
     && dnf clean all \
     && rm -rf /var/cache/dnf
 

container/Dockerfile

@@ -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", "block-manager"]
+media-loading = ["dep:ndarray", "dep:video-rs", "dep:image", "dep:reqwest", "dep:base64", "dep:tokio-rayon", "dep:flate2", "block-manager"]
 
 [[bench]]
 name = "tokenizer"
@@ -145,6 +145,7 @@ image = { version = "0.25", optional = true }
 reqwest = { workspace = true, optional = true }
 base64 = { version = "0.22", optional = true }
 tokio-rayon = {version = "2", optional = true }
+flate2 = { version = "1.1.2", optional = true }
 
 # Publishers
 zeromq = "0.4.1"

lib/llm/Cargo.toml

@@ -312,7 +312,6 @@ impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<LLMEngineOutput>, Error>
     ) -> Result<ManyOut<LLMEngineOutput>, Error> {
         let (request, ctx) = input.into_parts();
 
-        println!("multi_modal_data: {:?}", request.multi_modal_data);
         // Extract dp_rank from request field (defaults to 0 if not set)
         let dp_rank = request.dp_rank.unwrap_or(0);
 

lib/llm/src/mocker/engine.rs

@@ -3,17 +3,6 @@
 
 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
 #[derive(Debug)]
@@ -22,29 +11,6 @@ pub struct EncodedMediaData {
     pub(crate) b64_encoded: bool,
 }
 
-// Decoded media data (image RGB, video frames pixels, ...)
-#[derive(Debug)]
-pub struct DecodedMediaData {
-    pub(crate) data: SystemStorage,
-    pub(crate) shape: Vec<usize>,
-    pub(crate) 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
@@ -86,120 +52,6 @@ impl EncodedMediaData {
     }
 }
 
-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> 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();
-        Ok(Self {
-            data: SystemStorage::try_from(data)?,
-            shape,
-            dtype: "uint8".to_string(),
-        })
-    }
-}
-
-#[async_trait::async_trait]
-pub trait Decoder: Clone + Send + 'static {
-    fn decode(&self, data: EncodedMediaData) -> Result<DecodedMediaData>;
-
-    async fn decode_async(&self, data: EncodedMediaData) -> Result<DecodedMediaData> {
-        // light clone (only config params)
-        let decoder = self.clone();
-        // compute heavy -> rayon
-        let result = tokio_rayon::spawn(move || decoder.decode(data)).await?;
-        Ok(result)
-    }
-}
-
-#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
-pub struct MediaDecoder {
-    #[serde(default)]
-    pub image_decoder: ImageDecoder,
-    #[serde(default)]
-    pub video_decoder: VideoDecoder,
-}
-
-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(
-                "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<RdmaMediaDataDescriptor> {
-        // TODO: request-level options
-        // 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?;
-                self.media_decoder.image_decoder.decode_async(data).await
-            }
-            ChatCompletionRequestUserMessageContentPart::VideoUrl(video_part) => {
-                let url = &video_part.video_url.url;
-                let data = EncodedMediaData::from_url(url, &self.http_client).await?;
-                self.media_decoder.video_decoder.decode_async(data).await
-            }
-            ChatCompletionRequestUserMessageContentPart::AudioUrl(_) => {
-                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)
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;

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

@@ -5,7 +5,9 @@ use anyhow::Result;
 use image::GenericImageView;
 use ndarray::Array3;
 
-use super::common::{DecodedMediaData, Decoder, EncodedMediaData};
+use super::super::common::EncodedMediaData;
+use super::super::rdma::DecodedMediaData;
+use super::Decoder;
 
 #[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
 #[serde(deny_unknown_fields)]

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

@@ -0,0 +1,34 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use anyhow::Result;
+
+use super::common::EncodedMediaData;
+use super::rdma::DecodedMediaData;
+
+mod image;
+mod video;
+
+pub use image::ImageDecoder;
+pub use video::VideoDecoder;
+
+#[async_trait::async_trait]
+pub trait Decoder: Clone + Send + 'static {
+    fn decode(&self, data: EncodedMediaData) -> Result<DecodedMediaData>;
+
+    async fn decode_async(&self, data: EncodedMediaData) -> Result<DecodedMediaData> {
+        // light clone (only config params)
+        let decoder = self.clone();
+        // compute heavy -> rayon
+        let result = tokio_rayon::spawn(move || decoder.decode(data)).await?;
+        Ok(result)
+    }
+}
+
+#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
+pub struct MediaDecoder {
+    #[serde(default)]
+    pub image_decoder: ImageDecoder,
+    #[serde(default)]
+    pub video_decoder: VideoDecoder,
+}

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

@@ -1,7 +1,9 @@
 // SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-use super::common::{DecodedMediaData, Decoder, EncodedMediaData};
+use super::super::common::EncodedMediaData;
+use super::super::rdma::DecodedMediaData;
+use super::Decoder;
 use anyhow::Result;
 use ndarray::Array4;
 use std::io::Write;

lib/llm/src/preprocessor/media/video.rs renamed to lib/llm/src/preprocessor/media/decoders/video.rs

@@ -0,0 +1,65 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use anyhow::Result;
+
+use dynamo_async_openai::types::ChatCompletionRequestUserMessageContentPart;
+
+use super::common::EncodedMediaData;
+use super::decoders::{Decoder, MediaDecoder};
+use super::rdma::RdmaMediaDataDescriptor;
+use nixl_sys::Agent as NixlAgent;
+
+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(
+                "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<RdmaMediaDataDescriptor> {
+        // TODO: request-level options
+        // 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?;
+                self.media_decoder.image_decoder.decode_async(data).await
+            }
+            ChatCompletionRequestUserMessageContentPart::VideoUrl(video_part) => {
+                let url = &video_part.video_url.url;
+                let data = EncodedMediaData::from_url(url, &self.http_client).await?;
+                self.media_decoder.video_decoder.decode_async(data).await
+            }
+            ChatCompletionRequestUserMessageContentPart::AudioUrl(_) => {
+                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/loader.rs

@@ -2,9 +2,11 @@
 // SPDX-License-Identifier: Apache-2.0
 
 mod common;
-mod image;
-mod video;
+mod decoders;
+mod loader;
+mod rdma;
 
-pub use common::{Decoder, EncodedMediaData, MediaDecoder, MediaLoader, RdmaMediaDataDescriptor};
-pub use image::ImageDecoder;
-pub use video::VideoDecoder;
+pub use common::EncodedMediaData;
+pub use decoders::{Decoder, ImageDecoder, MediaDecoder, VideoDecoder};
+pub use loader::MediaLoader;
+pub use rdma::{DecodedMediaData, RdmaMediaDataDescriptor};