"docs/pages/design-docs/disagg-serving.md" did not exist on "cea1902de7c54bd0a890b70559b1057e82fd7279"
pending.rs 17.5 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
//! # Transfer Managers
//!
//! Transfer managers are responsible for multiple things:
//! - Before the transfer:
//!     - Rate-limiting the number of transfers that can be initiated concurrently. This is implemented through bounded channels.
//!         - Due to the nature of the [`super::OffloadManager`], we only apply this rate-limiting to offloads.
//! - During the transfer:
//!     - Initiating the transfer
//!     - Holding strong references to blocks being transfered.
//! - After the transfer:
//!     - Dropping these references once the transfer is complete.
//!     - Registering the blocks with the target pool.
//!     - Returning the registered blocks to the caller.
//!
//! This is implemented through the [`TransferManager`] trait, which takes a single [`PendingTransfer`]
//! and initiates the transfer.
//!
//! Since CUDA and NIXL transfers use completely different semantics, we implement two separate transfer managers.
//!
//! ## Workflow
36
37
//! 1. A transfer request is made by calling [`TransferManager::enqueue_transfer`]
//! 2. [`TransferManager::enqueue_transfer`] performs the transfer, and enqueues relevant data into a bounded channel.
38
39
40
//! 3. A worker thread (consuming this bounded channel and enforcing rate limiting) awaits the incoming transfers.
//! 4. After a transfer is complete, the worker thread registers the blocks with the target pool, and returns the registered blocks to the caller.

41
use std::marker::PhantomData;
42
use std::pin::Pin;
43
use std::sync::Arc;
44
use tokio::runtime::Handle;
45
use tokio::sync::mpsc;
46
use tokio_util::sync::CancellationToken;
47

48
49
use crate::block_manager::block::{
    transfer::{WriteTo, WriteToStrategy},
50
51
    BlockError, BlockExt, BlockMetadata, BlockState, MutableBlock, ReadableBlock, TransferContext,
    WritableBlock,
52
};
53
use crate::block_manager::pool::BlockPoolError;
54
use crate::block_manager::storage::{Local, Storage};
55
use crate::block_manager::BlockPool;
56

57
use anyhow::Result;
58
use async_trait::async_trait;
59
use futures::{stream::FuturesUnordered, StreamExt};
60

61
use super::BlockResult;
62

63
64
use dynamo_runtime::utils::task::CriticalTaskExecutionHandle;

65
66
67
/// Manage a set of pending transfers.
pub struct PendingTransfer<Source: Storage, Target: Storage, Metadata: BlockMetadata> {
    /// The block being copied from.
68
    sources: Vec<Arc<MutableBlock<Source, Metadata>>>,
69
70
71
    /// The block being copied to.
    targets: Vec<MutableBlock<Target, Metadata>>,
    /// The oneshot sender that optionally returns the registered blocks once the transfer is complete.
72
    completion_indicator: Option<oneshot::Sender<BlockResult<Target, Metadata>>>,
73
    /// The target pool that will receive the registered block.
74
    target_pool: Arc<BlockPool<Target, Metadata>>,
75
76
77
78
79
80
81
82
}

impl<Source: Storage, Target: Storage, Metadata: BlockMetadata>
    PendingTransfer<Source, Target, Metadata>
{
    pub fn new(
        sources: Vec<Arc<MutableBlock<Source, Metadata>>>,
        targets: Vec<MutableBlock<Target, Metadata>>,
83
        completion_indicator: Option<oneshot::Sender<BlockResult<Target, Metadata>>>,
84
        target_pool: Arc<BlockPool<Target, Metadata>>,
85
    ) -> Self {
86
        assert_eq!(sources.len(), targets.len());
87
        Self {
88
89
90
            sources,
            targets,
            completion_indicator,
91
            target_pool,
92
93
94
95
96
        }
    }

    fn handle_complete(self) -> Result<()> {
        let Self {
97
98
99
            sources,
            mut targets,
            target_pool,
100
            completion_indicator,
101
102
103
            ..
        } = self;

104
105
106
107
108
        for (source, target) in sources.iter().zip(targets.iter_mut()) {
            transfer_metadata(source, target)?;
        }

        let blocks = target_pool.register_blocks_blocking(targets)?;
109

110
        if let Some(completion_indicator) = completion_indicator {
111
112
113
            completion_indicator
                .send(Ok(blocks))
                .map_err(|_| BlockPoolError::ProgressEngineShutdown)?;
114
        }
115
116
117
118
119
120
121
122
123
124

        Ok(())
    }
}

fn transfer_metadata<Source: Storage, Target: Storage, Metadata: BlockMetadata>(
    source: &Arc<MutableBlock<Source, Metadata>>,
    target: &mut MutableBlock<Target, Metadata>,
) -> Result<()> {
    // Only registered blocks can be transferred. There are upstream checks for this, so this shouldn't ever fail.
125
    if let BlockState::Registered(reg_handle, _) = source.state() {
126
127
128
129
130
131
132
133
134
135
        // Bring the block back to the 'Reset' state.
        target.reset();
        // Transfer metadata.
        target.update_metadata(source.metadata().clone());
        // Copy tokens
        target.apply_token_block(reg_handle.token_block().clone())?;
    } else {
        Err(BlockPoolError::BlockError(BlockError::InvalidState(
            "Block is not registered.".to_string(),
        )))?;
136
    }
137
138
139
140
141
142
143
144
145

    Ok(())
}

#[async_trait]
pub trait TransferManager<Source: Storage, Target: Storage, Metadata: BlockMetadata>:
    Send + Sync
{
    /// Begin a transfer. Blocks if the pending queue is full.
146
    async fn enqueue_transfer(
147
148
149
        &self,
        pending_transfer: PendingTransfer<Source, Target, Metadata>,
    ) -> Result<()>;
150
151
}

152
pub struct CudaTransferManager<Source: Storage, Target: Storage, Metadata: BlockMetadata> {
153
154
155
156
    pending_transfer_q: mpsc::Sender<(
        PendingTransfer<Source, Target, Metadata>,
        tokio::sync::oneshot::Receiver<()>,
    )>,
157
    transfer_ctx: Arc<TransferContext>,
158
159
160
}

impl<Source: Storage, Target: Storage, Metadata: BlockMetadata>
161
    CudaTransferManager<Source, Target, Metadata>
162
{
163
164
165
    pub fn new(
        transfer_ctx: Arc<TransferContext>,
        max_concurrent_transfers: usize,
166
        runtime: &Handle,
167
        cancellation_token: CancellationToken,
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
    ) -> Result<Self> {
        let (tx, mut rx) = mpsc::channel::<(
            PendingTransfer<Source, Target, Metadata>,
            tokio::sync::oneshot::Receiver<()>,
        )>(max_concurrent_transfers);

        CriticalTaskExecutionHandle::new_with_runtime(
            move |cancel_token| async move {
                loop {
                    tokio::select! {
                        Some((pending_transfer, notify)) = rx.recv() => {
                            // Wait for the event.
                            notify.await.map_err(|_| BlockPoolError::ProgressEngineShutdown)?;
                            // Only finalize the transfer after the event is signaled.
                            match pending_transfer.handle_complete() {
                                Ok(_) => {}
                                Err(e) => {
                                    // The only case where this can fail is if the progress engine is being shutdown.
                                    // This is not a problem, so we can just ignore it.
                                    tracing::warn!("Error handling transfer completion: {:?}", e);
                                }
                            }
                        }
191

192
193
194
195
                        _ = cancel_token.cancelled() => {
                            return Ok(());
                        }
                    }
196
                }
197
198
199
200
201
202
203
204
            },
            cancellation_token.clone(),
            "Cuda Transfer Manager",
            runtime,
        )?
        .detach();

        Ok(Self {
205
            pending_transfer_q: tx,
206
            transfer_ctx,
207
        })
208
    }
209
}
210

211
212
213
214
215
216
217
218
219
220
221
222
223
224
#[async_trait]
impl<Source, Target, Metadata> TransferManager<Source, Target, Metadata>
    for CudaTransferManager<Source, Target, Metadata>
where
    Source: Storage,
    Target: Storage,
    Metadata: BlockMetadata,
    // Check that the source block is readable, local, and writable to the target block.
    MutableBlock<Source, Metadata>: ReadableBlock<StorageType = Source>
        + Local
        + WriteToStrategy<MutableBlock<Target, Metadata>>,
    // Check that the target block is writable.
    MutableBlock<Target, Metadata>: WritableBlock<StorageType = Target>,
{
225
    async fn enqueue_transfer(
226
        &self,
227
        mut pending_transfer: PendingTransfer<Source, Target, Metadata>,
228
    ) -> Result<()> {
229
230
231
232
233
234
235
236
237
238
239
240
        let notify = pending_transfer
            .sources
            .write_to(
                &mut pending_transfer.targets,
                true,
                self.transfer_ctx.clone(),
            )?
            .ok_or_else(|| {
                anyhow::anyhow!(
                    "write_to returned None when notify was true. This should never happen!"
                )
            })?;
241
242
243
244

        // Send the pending transfer and event to the worker thread.
        // If the queue is full, we block the worker until space becomes available.
        self.pending_transfer_q
245
            .send((pending_transfer, notify))
246
247
248
249
250
251
252
253
254
255
256
257
            .await?;

        Ok(())
    }
}

pub struct DiskTransferManager {
    futures_tx: mpsc::Sender<Pin<Box<dyn std::future::Future<Output = ()> + Send + Sync>>>,
    transfer_ctx: Arc<TransferContext>,
}

impl DiskTransferManager {
258
259
260
261
262
    pub fn new(
        transfer_ctx: Arc<TransferContext>,
        max_concurrent_transfers: usize,
        runtime: &Handle,
        cancellation_token: CancellationToken,
263
    ) -> Result<Self> {
264
265
        let (futures_tx, mut futures_rx) = mpsc::channel(1);

266
267
268
269
        CriticalTaskExecutionHandle::new_with_runtime(
            move |cancel_token| async move {
                // Keep track of our pending transfers.
                // Consume the futures as they complete, while also receiving new ones.
270

271
272
273
                let mut pending_transfers = FuturesUnordered::new();
                loop {
                    tokio::select! {
274

275
276
277
                        _ = cancel_token.cancelled() => {
                            return Ok(());
                        }
278

279
280
281
282
283
284
285
286
287
288
                        Some(future) = futures_rx.recv() => {
                            // If we're at max size, block the worker thread on the next() call until we have capacity.
                            while pending_transfers.len() >= max_concurrent_transfers {
                                pending_transfers.next().await;
                            }
                            // Once we have capacity, push the new future onto the queue.
                            pending_transfers.push(future);
                        }
                        Some(_) = pending_transfers.next(), if !pending_transfers.is_empty() => {
                            // A transfer completed, just continue to process more
289
290
291
                        }
                    }
                }
292
293
294
295
296
297
298
299
            },
            cancellation_token.clone(),
            "Disk Transfer Manager",
            runtime,
        )?
        .detach();

        Ok(Self {
300
301
            futures_tx,
            transfer_ctx,
302
        })
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
    }
}

#[async_trait]
impl<Source, Target, Metadata> TransferManager<Source, Target, Metadata> for DiskTransferManager
where
    Source: Storage,
    Target: Storage,
    Metadata: BlockMetadata,
    // Check that the source block is readable, local, and writable to the target block.
    MutableBlock<Source, Metadata>: ReadableBlock<StorageType = Source>
        + Local
        + WriteToStrategy<MutableBlock<Target, Metadata>>,
    // Check that the target block is writable.
    MutableBlock<Target, Metadata>: WritableBlock<StorageType = Target>,
{
319
    async fn enqueue_transfer(
320
321
322
        &self,
        mut pending_transfer: PendingTransfer<Source, Target, Metadata>,
    ) -> Result<()> {
323
324
325
326
327
328
329
330
331
332
333
334
        let notify = pending_transfer
            .sources
            .write_to(
                &mut pending_transfer.targets,
                true,
                self.transfer_ctx.clone(),
            )?
            .ok_or_else(|| {
                anyhow::anyhow!(
                    "write_to returned None when notify was true. This should never happen!"
                )
            })?;
335
336

        let completion_future = async move {
337
            let _ = notify.await;
338
339
340
341
342
343
344
345
            match pending_transfer.handle_complete() {
                Ok(_) => {}
                Err(e) => {
                    // The only case where this can fail is if the progress engine is being shutdown.
                    // This is not a problem, so we can just ignore it.
                    tracing::warn!("Error handling transfer completion: {:?}", e);
                }
            }
346
347
348
349
350
        };

        // Futures_(tx/rx) has a capacity of 1. If the queue worker has received another future and is awaiting next() due to a full `FuturesUnordered`,
        // this call will block until the worker has processed the prior future.
        self.futures_tx.send(Box::pin(completion_future)).await?;
351
352
353
354

        Ok(())
    }
}
355
356
357
358
359
360
361
362
363
364
365

/// A transfer manager that enforces a max batch size for transfers.
pub struct TransferBatcher<Source, Target, Metadata, Manager>
where
    Source: Storage,
    Target: Storage,
    Metadata: BlockMetadata,
    Manager: TransferManager<Source, Target, Metadata>,
{
    transfer_manager: Manager,
    max_transfer_batch_size: usize,
366
367
    runtime: Handle,
    cancellation_token: CancellationToken,
368
369
370
371
372
373
374
375
376
377
    _phantom: PhantomData<(Source, Target, Metadata)>,
}

impl<Source, Target, Metadata, Manager> TransferBatcher<Source, Target, Metadata, Manager>
where
    Source: Storage,
    Target: Storage,
    Metadata: BlockMetadata,
    Manager: TransferManager<Source, Target, Metadata>,
{
378
379
380
381
382
383
    pub fn new(
        transfer_manager: Manager,
        max_transfer_batch_size: usize,
        runtime: &Handle,
        cancellation_token: CancellationToken,
    ) -> Self {
384
385
386
        Self {
            transfer_manager,
            max_transfer_batch_size,
387
388
            runtime: runtime.clone(),
            cancellation_token,
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
            _phantom: PhantomData,
        }
    }
}

#[async_trait]
impl<Source, Target, Metadata, Manager> TransferManager<Source, Target, Metadata>
    for TransferBatcher<Source, Target, Metadata, Manager>
where
    Source: Storage,
    Target: Storage,
    Metadata: BlockMetadata,
    Manager: TransferManager<Source, Target, Metadata>,
{
    async fn enqueue_transfer(
        &self,
        pending_transfer: PendingTransfer<Source, Target, Metadata>,
    ) -> Result<()> {
        // If it's smaller than the max batch size, just enqueue it.
        if pending_transfer.sources.len() < self.max_transfer_batch_size {
            return self
                .transfer_manager
                .enqueue_transfer(pending_transfer)
                .await;
        }

        // Otherwise, we need to split the transfer into multiple smaller transfers.

        let PendingTransfer {
            mut sources,
            mut targets,
            completion_indicator,
            target_pool,
        } = pending_transfer;

        let mut indicators = Vec::new();

        while !sources.is_empty() {
            let sources = sources
                .drain(..std::cmp::min(self.max_transfer_batch_size, sources.len()))
                .collect();
            let targets = targets
                .drain(..std::cmp::min(self.max_transfer_batch_size, targets.len()))
                .collect();

            // If we have a completion indicator, we need to create a new one for each sub-transfer.
            let indicator = if completion_indicator.is_some() {
                let (batch_tx, batch_rx) = oneshot::channel();
                indicators.push(batch_rx);
                Some(batch_tx)
            } else {
                None
            };

            let request = PendingTransfer::new(sources, targets, indicator, target_pool.clone());
            // Enqueue our reduced transfer. This may block if the queue is full.
            self.transfer_manager.enqueue_transfer(request).await?;
        }

        if let Some(completion_indicator) = completion_indicator {
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
            CriticalTaskExecutionHandle::new_with_runtime(
                move |cancel_token| async move {
                    let mut results = Vec::new();

                    for indicator in indicators.into_iter() {
                        // Await each sub-transfer, and append the results to our final results.
                        tokio::select! {
                            _ = cancel_token.cancelled() => {
                                return Ok(());
                            }

                            Ok(indicator) = indicator => {
                                let result = match indicator {
                                    Ok(result) => result,
                                    Err(e) => {
                                        tracing::error!("Error receiving transfer results: {:?}", e);
                                        completion_indicator.send(Err(e)).unwrap();
                                        return Ok(());
                                    }
                                };
                                results.extend(result);
                            }
471
                        }
472
473
474
475
                    }

                    // Send the final results to the top-level completion indicator.
                    completion_indicator.send(Ok(results))?;
476

477
478
479
480
481
482
                    Ok(())
                },
                self.cancellation_token.clone(),
                "Transfer Batcher",
                &self.runtime,
            )?.detach();
483
484
485
486
487
        }

        Ok(())
    }
}