Graceful shutdown + batched nixl transfers + fix DiskStorage unregistration

Author: jthomson04

lib/llm/src/block_manager/block/transfer.rs

@@ -133,7 +133,7 @@ where
 pub trait WriteTo<Target> {
     fn write_to(
         &self,
-        dst: &mut Target,
+        dst: &mut Vec<Target>,
         notify: Option<String>,
         ctx: Arc<TransferContext>,
     ) -> Result<(), TransferError>;
@@ -143,49 +143,61 @@ pub trait WriteTo<Target> {
     /// Returns a future that will complete when the transfer is complete.
     fn nixl_write_to(
         &self,
-        dst: &mut Target,
+        dst: &mut Vec<Target>,
         notify: Option<String>,
         ctx: Arc<TransferContext>,
     ) -> Result<Box<dyn Future<Output = ()> + Send + Sync + Unpin>, TransferError>;
 }
 
-impl<RB: ReadableBlock, WB: WritableBlock> WriteTo<WB> for RB
+impl<RB: ReadableBlock, WB: WritableBlock> WriteTo<WB> for Vec<Arc<RB>>
 where
     RB: WriteToStrategy<WB> + Local,
 {
     fn write_to(
         &self,
-        dst: &mut WB,
+        dst: &mut Vec<WB>,
         notify: Option<String>,
         ctx: Arc<TransferContext>,
     ) -> Result<(), TransferError> {
-        match Self::write_to_strategy() {
-            TransferStrategy::Memcpy => memcpy::copy_block(self, dst),
+        match RB::write_to_strategy() {
+            TransferStrategy::Memcpy => {
+                for (src, dst) in self.iter().zip(dst.iter_mut()) {
+                    memcpy::copy_block(src.as_ref(), dst)?;
+                }
+                Ok(())
+            }
             TransferStrategy::CudaAsyncH2D
             | TransferStrategy::CudaAsyncD2H
             | TransferStrategy::CudaAsyncD2D => {
-                cuda::copy_block(self, dst, ctx.stream().as_ref(), RB::write_to_strategy())
+                for (src, dst) in self.iter().zip(dst.iter_mut()) {
+                    cuda::copy_block(
+                        src.as_ref(),
+                        dst,
+                        ctx.stream().as_ref(),
+                        RB::write_to_strategy(),
+                    )?;
+                }
+                Ok(())
             }
             TransferStrategy::NixlWrite => {
-                std::mem::drop(nixl::write_block_to(self, dst, ctx, notify)?);
+                std::mem::drop(nixl::write_blocks_to(self, dst, ctx, notify)?);
                 Ok(())
             }
             _ => Err(TransferError::IncompatibleTypes(format!(
                 "Unsupported copy strategy: {:?}",
                 RB::write_to_strategy()
             ))),
         }
-        // dispatch_copy_to(self, dst, self.transfer_context())
     }
 
     fn nixl_write_to(
         &self,
-        dst: &mut WB,
+        dst: &mut Vec<WB>,
         notify: Option<String>,
         ctx: Arc<TransferContext>,
     ) -> Result<Box<dyn Future<Output = ()> + Send + Sync + Unpin>, TransferError> {
         if let TransferStrategy::NixlWrite = RB::write_to_strategy() {
-            Ok(nixl::write_block_to(self, dst, ctx, notify)?)
+            Ok(nixl::write_blocks_to(self, dst, ctx, notify)?)
         } else {
             Err(TransferError::IncompatibleTypes(format!(
                 "Expected NIXL transfer strategy, got: {:?}",

lib/llm/src/block_manager/block/transfer/nixl.rs

@@ -18,16 +18,14 @@ use super::*;
 use anyhow::Result;
 use nixl_sys::{MemoryRegion, NixlDescriptor, OptArgs, XferDescList, XferOp};
 use std::future::{poll_fn, Future};
-use std::ops::Range;
 use std::task::Poll;
 
-/// Copy a block from a source to a destination using CUDA memcpy
-pub fn write_block_to<'a, Source, Destination>(
-    src: &'a Source,
-    dst: &'a mut Destination,
-    ctx: Arc<TransferContext>,
-    notify: Option<String>,
-) -> Result<Box<dyn Future<Output = ()> + Send + Sync + Unpin>>
+fn append_xfer_request<Source, Destination>(
+    src: &Arc<Source>,
+    dst: &mut Destination,
+    src_dl: &mut XferDescList,
+    dst_dl: &mut XferDescList,
+) -> Result<()>
 where
     Source: BlockDataProvider,
     Destination: BlockDataProviderMut,
@@ -36,17 +34,6 @@ where
     let dst_data = dst.block_data_mut(private::PrivateToken);
 
     if src_data.is_fully_contiguous() && dst_data.is_fully_contiguous() {
-        // Keep the arc to use in the returned future.
-        let nixl_agent_arc = ctx.as_ref().nixl_agent();
-
-        let nixl_agent = nixl_agent_arc
-            .as_ref()
-            .as_ref()
-            .expect("NIXL agent not found");
-
-        let mut src_dl = XferDescList::new(src_data.storage_type().nixl_mem_type())?;
-        let mut dst_dl = XferDescList::new(dst_data.storage_type().nixl_mem_type())?;
-
         let src_desc = src_data.block_view()?.as_nixl_descriptor();
         let dst_desc = dst_data.block_view_mut()?.as_nixl_descriptor_mut();
 
@@ -64,121 +51,99 @@ where
             )?;
         }
 
-        let xfer_req = nixl_agent
-            .create_xfer_req(XferOp::Write, &src_dl, &dst_dl, &nixl_agent.name(), None)
-            .unwrap();
-
-        let mut xfer_args = OptArgs::new()?;
-
-        if let Some(notify) = notify {
-            xfer_args.set_has_notification(true)?;
-            xfer_args.set_notification_message(notify.as_bytes())?;
-        }
-
-        let _ = nixl_agent.post_xfer_req(&xfer_req, Some(&xfer_args))?;
-
-        // Return a future that completes when the transfer is complete.
-        // TODO: How efficient is this? Can we do better?
-        Ok(Box::new(poll_fn(move |_cx| {
-            let nixl_agent = nixl_agent_arc
-                .as_ref()
-                .as_ref()
-                .expect("NIXL agent not found");
-
-            // The nixl agent returns true if the transfer is still in progress.
-            if !nixl_agent.get_xfer_status(&xfer_req).unwrap() {
-                Poll::Ready(())
-            } else {
-                Poll::Pending
-            }
-        })))
+        Ok(())
     } else {
         assert_eq!(src_data.num_layers(), dst_data.num_layers());
-        write_layers_to(0..src_data.num_layers(), src, dst, ctx, notify)
+        for layer_idx in 0..src_data.num_layers() {
+            for outer_idx in 0..src_data.num_outer_dims() {
+                let src_view = src_data.layer_view(layer_idx, outer_idx)?;
+                let mut dst_view = dst_data.layer_view_mut(layer_idx, outer_idx)?;
+
+                debug_assert_eq!(src_view.size(), dst_view.size());
+
+                let src_desc = src_view.as_nixl_descriptor();
+                let dst_desc = dst_view.as_nixl_descriptor_mut();
+
+                unsafe {
+                    src_dl.add_desc(
+                        src_desc.as_ptr() as usize,
+                        src_desc.size(),
+                        src_desc.device_id(),
+                    )?;
+
+                    dst_dl.add_desc(
+                        dst_desc.as_ptr() as usize,
+                        dst_desc.size(),
+                        dst_desc.device_id(),
+                    )?;
+                }
+            }
+        }
+        Ok(())
     }
 }
 
-/// Copy a range of layers from a source to a destination using CUDA memcpy
-pub fn write_layers_to<'a, Source, Destination>(
-    layer_range: Range<usize>,
-    src: &'a Source,
-    dst: &'a mut Destination,
+/// Copy a block from a source to a destination using CUDA memcpy
+pub fn write_blocks_to<Source, Destination>(
+    src: &[Arc<Source>],
+    dst: &mut [Destination],
     ctx: Arc<TransferContext>,
     notify: Option<String>,
 ) -> Result<Box<dyn Future<Output = ()> + Send + Sync + Unpin>>
 where
     Source: BlockDataProvider,
     Destination: BlockDataProviderMut,
 {
-    let src_data = src.block_data(private::PrivateToken);
-    let dst_data = dst.block_data_mut(private::PrivateToken);
+    if src.is_empty() || dst.is_empty() {
+        return Ok(Box::new(std::future::ready(())));
+    }
+    assert_eq!(src.len(), dst.len());
 
     let nixl_agent_arc = ctx.as_ref().nixl_agent();
     let nixl_agent = nixl_agent_arc
         .as_ref()
         .as_ref()
         .expect("NIXL agent not found");
 
-    let remote_worker_id = dst_data.worker_id.to_string();
-    let mut src_dl = XferDescList::new(src_data.storage_type().nixl_mem_type())?;
-    let mut dst_dl = XferDescList::new(dst_data.storage_type().nixl_mem_type())?;
-
-    // #[cfg(debug_assertions)]
-    // {
-    //     let expected_strategy = <<Source as BlockDataProvider>::StorageType as WriteToStrategy<
-    //         Destination::StorageType,
-    //     >>::write_to_strategy();
-    //     assert_eq!(strategy, expected_strategy);
-    // }
-
-    for layer_idx in layer_range {
-        for outer_idx in 0..src_data.num_outer_dims() {
-            let src_view = src_data.layer_view(layer_idx, outer_idx)?;
-            let mut dst_view = dst_data.layer_view_mut(layer_idx, outer_idx)?;
-
-            debug_assert_eq!(src_view.size(), dst_view.size());
-
-            let src_desc = src_view.as_nixl_descriptor();
-            let dst_desc = dst_view.as_nixl_descriptor_mut();
-
-            unsafe {
-                src_dl.add_desc(
-                    src_desc.as_ptr() as usize,
-                    src_desc.size(),
-                    src_desc.device_id(),
-                )?;
-
-                dst_dl.add_desc(
-                    dst_desc.as_ptr() as usize,
-                    dst_desc.size(),
-                    dst_desc.device_id(),
-                )?;
-            }
-        }
+    let src_mem_type = src
+        .first()
+        .unwrap()
+        .block_data(private::PrivateToken)
+        .storage_type()
+        .nixl_mem_type();
+    let dst_mem_type = dst
+        .first()
+        .unwrap()
+        .block_data(private::PrivateToken)
+        .storage_type()
+        .nixl_mem_type();
+
+    let mut src_dl = XferDescList::new(src_mem_type)?;
+    let mut dst_dl = XferDescList::new(dst_mem_type)?;
+
+    for (src, dst) in src.iter().zip(dst.iter_mut()) {
+        append_xfer_request(src, dst, &mut src_dl, &mut dst_dl)?;
     }
 
+    let xfer_req =
+        nixl_agent.create_xfer_req(XferOp::Write, &src_dl, &dst_dl, &nixl_agent.name(), None)?;
+
     let mut xfer_args = OptArgs::new()?;
 
     if let Some(notify) = notify {
         xfer_args.set_has_notification(true)?;
         xfer_args.set_notification_message(notify.as_bytes())?;
     }
 
-    let xfer_req = nixl_agent.create_xfer_req(
-        XferOp::Write,
-        &src_dl,
-        &dst_dl,
-        &remote_worker_id,
-        Some(&xfer_args),
-    )?;
-
     let _ = nixl_agent.post_xfer_req(&xfer_req, Some(&xfer_args))?;
 
     Ok(Box::new(poll_fn(move |_cx| {
         let nixl_agent = nixl_agent_arc
             .as_ref()
             .as_ref()
             .expect("NIXL agent not found");
+
+        // The nixl agent returns true if the transfer is still in progress.
         if !nixl_agent.get_xfer_status(&xfer_req).unwrap() {
             Poll::Ready(())
         } else {

lib/llm/src/block_manager/offload.rs

@@ -129,7 +129,7 @@ impl<Metadata: BlockMetadata> OffloadManager<Metadata> {
         let device_clone = this.device.clone();
         let host_clone = this.host.clone();
         async_rt_handle.spawn(async move {
-            OffloadManager::offload_worker(
+            let res = OffloadManager::offload_worker(
                 device_clone,
                 host_clone,
                 device_offload_rx,
@@ -138,8 +138,8 @@ impl<Metadata: BlockMetadata> OffloadManager<Metadata> {
                     MAX_OFFLOAD_STREAM_DEPTH,
                 )),
             )
-            .await
-            .unwrap()
+            .await;
+            tracing::warn!("Offload worker finished: {:?}", res);
         });
 
         let transfer_ctx = Arc::new(TransferContext::new(
@@ -152,7 +152,7 @@ impl<Metadata: BlockMetadata> OffloadManager<Metadata> {
         let disk_clone = this.disk.clone();
         let transfer_ctx_clone = transfer_ctx.clone();
         async_rt_handle.spawn(async move {
-            OffloadManager::offload_worker(
+            let res = OffloadManager::offload_worker(
                 host_clone,
                 disk_clone,
                 host_offload_rx,
@@ -161,38 +161,38 @@ impl<Metadata: BlockMetadata> OffloadManager<Metadata> {
                     MAX_OFFLOAD_STREAM_DEPTH,
                 )),
             )
-            .await
-            .unwrap()
+            .await;
+            tracing::warn!("Offload worker finished: {:?}", res);
         });
 
         // Host -> Device onboarding
         let host_clone = this.host.clone();
         let device_clone = this.device.clone();
         let transfer_ctx_clone = transfer_ctx.clone();
         async_rt_handle.spawn(async move {
-            OffloadManager::onboard_worker(
+            let res = OffloadManager::onboard_worker(
                 host_clone,
                 device_clone,
                 host_onboard_rx,
                 Arc::new(CudaTransferManager::new(transfer_ctx_clone, 16384)),
             )
-            .await
-            .unwrap()
+            .await;
+            tracing::warn!("Onboard worker finished: {:?}", res);
         });
 
         // Disk -> Device onboarding
         let disk_clone = this.disk.clone();
         let device_clone = this.device.clone();
         let transfer_ctx_clone = transfer_ctx.clone();
         async_rt_handle.spawn(async move {
-            OffloadManager::onboard_worker(
+            let res = OffloadManager::onboard_worker(
                 disk_clone,
                 device_clone,
                 disk_onboard_rx,
                 Arc::new(DiskTransferManager::new(transfer_ctx_clone, 16384)),
             )
-            .await
-            .unwrap()
+            .await;
+            tracing::warn!("Onboard worker terminated: {:?}", res);
         });
 
         Ok(this_clone)