worker.rs 26.3 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Ryan Olson's avatar
Ryan Olson committed
2
3
4
5
// SPDX-License-Identifier: Apache-2.0

use super::*;

6
use async_trait::async_trait;
Ryan Olson's avatar
Ryan Olson committed
7
8
9
10
11
use transfer::*;
use utils::*;
use zmq::*;

use crate::block_manager::{
12
    BasicMetadata, BlockMetadata, LayoutConfigBuilder, NixlLayout, Storage,
13
14
15
16
    block::{
        Block, layout_to_blocks, locality,
        transfer::{PoolConfig, TransferContext},
    },
Ryan Olson's avatar
Ryan Olson committed
17
18
    connector::scheduler::TransferSchedulerClient,
    layout::LayoutType,
19
    offload::{MAX_CONCURRENT_TRANSFERS, MAX_TRANSFER_BATCH_SIZE},
20
    storage::{DeviceAllocator, DeviceStorage, DiskAllocator, PinnedAllocator, torch::TorchTensor},
Ryan Olson's avatar
Ryan Olson committed
21
22
23
24
25
26
};

use derive_builder::Builder;
use nixl_sys::Agent as NixlAgent;
use std::collections::HashMap;
use std::sync::Arc;
27
use std::sync::atomic::{AtomicBool, Ordering};
Ryan Olson's avatar
Ryan Olson committed
28
29
30
31

use tokio::runtime::Handle;
use tokio_util::sync::CancellationToken;

Richard Huo's avatar
Richard Huo committed
32
use dynamo_runtime::utils::task::CriticalTaskExecutionHandle;
33
use tokio::sync::{Mutex, RwLock, oneshot};
Ryan Olson's avatar
Ryan Olson committed
34

35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
struct WorkerState {
    ready_for_ping: AtomicBool,
}

impl WorkerState {
    fn new() -> Self {
        Self {
            ready_for_ping: AtomicBool::new(false),
        }
    }
    fn mark_ready(&self) {
        self.ready_for_ping.store(true, Ordering::SeqCst);
    }
    fn is_ready(&self) -> bool {
        self.ready_for_ping.load(Ordering::SeqCst)
    }
Ryan Olson's avatar
Ryan Olson committed
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
}

pub fn load_and_validate_tensors(
    tensors: &[Arc<dyn TorchTensor>],
    device_id: usize,
) -> anyhow::Result<(Vec<DeviceStorage>, Vec<usize>)> {
    let mut shape = None;

    let mut device_tensors = Vec::with_capacity(tensors.len());
    let allocator = DeviceAllocator::new(device_id)?;

    for tensor in tensors {
        // Check the stride, and ensure our tensor is contiguous.
        // TODO: We eventually need to be able to handle this.
        let stride = tensor.stride();
66
67
68
        tracing::debug!("stride: {:?}", stride);
        tracing::debug!("stride is monotonically decreasing for NHD layout");
        tracing::debug!("stride is NOT monotonically decreasing for HND layout");
Ryan Olson's avatar
Ryan Olson committed
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92

        // Check that all layer tensors have the same shape.
        // TODO: We eventually need to support the weirder models with heterogenous layers.
        if let Some(shape) = shape.as_ref() {
            if *shape != tensor.shape() {
                return Err(anyhow::anyhow!(
                    "All tensors must have the same shape! Got {:?} and {:?}",
                    *shape,
                    tensor.shape()
                ));
            }
        } else {
            shape = Some(tensor.shape());
        }

        // Build the storage object from the tensor.
        let device_tensor = DeviceStorage::new_from_torch(allocator.ctx(), tensor.clone())?;

        device_tensors.push(device_tensor);
    }

    Ok((device_tensors, shape.unwrap()))
}

93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
fn build_agent(worker_id: usize, use_gds: bool) -> anyhow::Result<NixlAgent> {
    let agent = NixlAgent::new(&format!("kvbm-worker-{}", worker_id))?;
    if use_gds {
        let (_, gds_params) = agent.get_plugin_params("GDS_MT")?;
        agent.create_backend("GDS_MT", &gds_params)?;
    }
    let (_, posix_params) = agent.get_plugin_params("POSIX")?;
    agent.create_backend("POSIX", &posix_params)?;

    Ok(agent)
}

// Helper: perform allocation and build transfer handler (factored from previous code)
async fn perform_allocation_and_build_handler(
    device_layout: Box<dyn NixlLayout<StorageType = DeviceStorage>>,
    mut layout_builder: LayoutConfigBuilder,
    worker_config: KvbmWorkerConfig,
    leader_meta: LeaderMetadata,
    worker_id: usize,
    device_id: usize,
    scheduler_client: Option<TransferSchedulerClient>,
114
115
116
117
118
119
120
121
122
) -> anyhow::Result<BlockTransferHandler> {
    let agent = build_agent(worker_id, leader_meta.num_disk_blocks > 0)?;
    let pool_config = PoolConfig {
        enable_pool: true,
        max_concurrent_transfers: MAX_CONCURRENT_TRANSFERS,
        max_transfer_batch_size: MAX_TRANSFER_BATCH_SIZE,
        num_outer_components: device_layout.config().outer_dim,
        num_layers: device_layout.config().num_layers,
    };
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
    let transfer_context = Arc::new(
        TransferContext::new(
            Arc::new(Some(agent)),
            DeviceAllocator::new(device_id)?.ctx().new_stream()?,
            Handle::current(),
            Some(pool_config),
        )
        .map_err(|e| {
            anyhow::anyhow!(
                "Failed to create transfer context for worker {} with CUDA memory pool: {}. \
                 This is a critical error - the worker cannot start without CUDA memory pools. \
                 Please ensure sufficient GPU memory is available on device {}.",
                worker_id,
                e,
                device_id
            )
        })?,
    );
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156

    // device
    let device_blocks = Some(KvbmWorker::make_layout::<_, BasicMetadata>(
        device_layout,
        transfer_context.nixl_agent().as_ref(),
        0,
        worker_id,
    )?);
    // host
    let host_blocks = if leader_meta.num_host_blocks > 0 {
        let host_allocator = Arc::new(PinnedAllocator::default());
        let host_layout = layout_builder
            .num_blocks(leader_meta.num_host_blocks)
            .build()?
            .allocate_layout(worker_config.host_layout_type, host_allocator)?;
        Some(KvbmWorker::make_layout::<_, BasicMetadata>(
157
            host_layout,
158
159
160
161
            transfer_context.nixl_agent().as_ref(),
            1,
            worker_id,
        )?)
jthomson04's avatar
jthomson04 committed
162
    } else {
163
164
165
166
167
168
169
170
171
172
173
        None
    };
    // disk
    let disk_blocks = if leader_meta.num_disk_blocks > 0 {
        let disk_allocator = Arc::new(DiskAllocator);
        let disk_layout = layout_builder
            .num_blocks(leader_meta.num_disk_blocks)
            .build()?
            .allocate_layout(worker_config.disk_layout_type, disk_allocator)?;
        Some(KvbmWorker::make_layout::<_, BasicMetadata>(
            disk_layout,
174
            transfer_context.nixl_agent().as_ref(),
175
            2,
176
            worker_id,
177
178
179
180
181
182
183
184
185
186
187
188
189
        )?)
    } else {
        None
    };

    let handler = BlockTransferHandler::new(
        device_blocks,
        host_blocks,
        disk_blocks,
        transfer_context,
        scheduler_client,
    )?;
    Ok(handler)
190
191
192
193
194
195
196
197
198
199
}

struct WorkerMetadataHandler {
    num_device_blocks: usize,
    bytes_per_block: usize,
}

#[async_trait]
impl Handler for WorkerMetadataHandler {
    async fn handle(&self, mut message: MessageHandle) -> anyhow::Result<()> {
200
201
202
203
204
205
206
        let payload = bincode::serde::encode_to_vec(
            &WorkerMetadata {
                num_device_blocks: self.num_device_blocks,
                bytes_per_block: self.bytes_per_block,
            },
            bincode::config::standard(),
        )?;
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
        message
            .reply(ZMQ_WORKER_METADATA_MESSAGE, &[payload])
            .await?;
        Ok(())
    }
}

// Leader sends allocation config -> allocate -> publish handler -> mark ready -> ACK
struct LeaderMetadataHandler {
    state: Arc<WorkerState>,
    device_layout: Mutex<Option<Box<dyn NixlLayout<StorageType = DeviceStorage>>>>,
    layout_builder: LayoutConfigBuilder,
    worker_config: KvbmWorkerConfig,
    worker_id: usize,
    device_id: usize,
    scheduler_client: Option<TransferSchedulerClient>,
223
224
    handler_cell: Arc<RwLock<Option<BlockTransferHandler>>>,
    handler_tx: Arc<Mutex<Option<oneshot::Sender<BlockTransferHandler>>>>,
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
    started: AtomicBool,
}

#[async_trait]
impl Handler for LeaderMetadataHandler {
    async fn handle(&self, mut message: MessageHandle) -> anyhow::Result<()> {
        // Always ACK ASAP so Drop can't panic and leader can finish the round.
        if let Err(e) = message.ack().await {
            tracing::error!("leader_metadata: failed to ACK: {e:#}");
        }

        // Validate payload; if bad, ignore.
        if message.data.len() != 1 {
            tracing::error!(
                "leader_metadata expects 1 payload frame (got {})",
                message.data.len()
            );
            return Ok(());
        }
244
245
246
247
248
        let leader_meta: LeaderMetadata = match bincode::serde::decode_from_slice(
            &message.data[0],
            bincode::config::standard(),
        ) {
            Ok((m, _)) => m,
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
            Err(e) => {
                tracing::error!("leader_metadata: bad payload: {e:#}");
                return Ok(());
            }
        };

        // Single-flight: only the first message triggers allocation.
        if self
            .started
            .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
            .is_err()
        {
            tracing::debug!("leader_metadata: allocation already started; dropping duplicate");
            return Ok(());
        }

        // Take device_layout once.
        let dev_layout = {
            let mut guard = self.device_layout.lock().await;
            match guard.take() {
                Some(d) => d,
                None => {
                    tracing::warn!("leader_metadata: device_layout already consumed; dropping");
                    return Ok(());
                }
            }
        };

        // Capture what we need and run allocation in the background.
        let layout_builder = self.layout_builder.clone();
        let worker_config = self.worker_config.clone();
        let worker_id = self.worker_id;
        let device_id = self.device_id;
        let scheduler_client = self.scheduler_client.clone();
        let handler_cell = self.handler_cell.clone();
        let handler_tx = self.handler_tx.clone();
        let state = self.state.clone();

        tokio::spawn(async move {
            match perform_allocation_and_build_handler(
                dev_layout,
                layout_builder,
                worker_config,
                leader_meta,
                worker_id,
                device_id,
                scheduler_client,
            )
            .await
            {
                Ok(handler) => {
                    // Install transfer handler
                    {
                        let mut w = handler_cell.write().await;
                        *w = Some(handler.clone());
                    }
                    // Return handler to creator (once)
                    {
                        let mut g = handler_tx.lock().await;
                        if let Some(tx) = g.take() {
                            let _ = tx.send(handler);
                        }
                    }
                    // Now the worker can ACK pings
                    state.mark_ready();
                    tracing::info!("allocation finished; worker is ping-ACK-able");
                }
                Err(e) => {
                    tracing::error!("allocation failed: {e:#}");
                    // leave ready=false so pings keep being ignored
                }
            }
        });

        Ok(())
    }
}

// Gated ping: the worker can only response to ping after the state is ready
struct GatedPing {
    state: Arc<WorkerState>,
    // fired exactly once after the first successful ping ACK
    layout_ready_tx: Mutex<Option<oneshot::Sender<String>>>,
}

#[async_trait]
impl Handler for GatedPing {
    async fn handle(&self, mut message: MessageHandle) -> anyhow::Result<()> {
        if !self.state.is_ready() {
338
339
340
341
            tracing::info!(
                "KVBM worker is under initialization. It could take a while if set with large CPU or DISK cache size. Please wait..."
            );
            tracing::debug!("Ping received but worker not ready; deferring ACK");
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
            // Prevent Drop panic; leader won't get an ACK for this round and will retry.
            message.mark_handled();
            return Ok(());
        }

        message.ack().await?;

        // After a successful ACK, flip the readiness oneshot exactly once
        let mut guard = self.layout_ready_tx.lock().await;
        if let Some(tx) = guard.take() {
            let _ = tx.send("ping-acked".to_string());
            tracing::info!("Reported ping-ready after first ACK");
        }

        Ok(())
    }
}

// Transfer dispatcher that waits until block transfer handler exists
struct BlockTransferDispatch {
362
    cell: Arc<RwLock<Option<BlockTransferHandler>>>,
363
364
365
366
367
368
369
370
371
372
373
374
375
376
}

#[async_trait]
impl Handler for BlockTransferDispatch {
    async fn handle(&self, message: MessageHandle) -> anyhow::Result<()> {
        let maybe = { self.cell.read().await.clone() };
        if let Some(inner) = maybe {
            inner.handle(message).await
        } else {
            Err(anyhow::anyhow!("transfer handler not ready yet"))
        }
    }
}

377
#[derive(Builder, Clone)]
Ryan Olson's avatar
Ryan Olson committed
378
379
#[builder(pattern = "owned")]
pub struct KvbmWorkerConfig {
Richard Huo's avatar
Richard Huo committed
380
    cancel_token: CancellationToken,
Ryan Olson's avatar
Ryan Olson committed
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395

    num_device_blocks: usize,

    #[builder(default = "32")]
    page_size: usize,

    #[builder(default = "Vec::new()")]
    tensors: Vec<Arc<dyn TorchTensor>>,

    #[builder(default = "0")]
    device_id: usize,

    #[builder(default = "2")]
    dtype_width_bytes: usize,

396
397
398
399
400
401
402
403
    #[builder(default = "LayoutType::FullyContiguous")]
    device_layout_type: LayoutType,

    #[builder(default = "LayoutType::FullyContiguous")]
    host_layout_type: LayoutType,

    #[builder(default = "LayoutType::FullyContiguous")]
    disk_layout_type: LayoutType,
404

Ryan Olson's avatar
Ryan Olson committed
405
406
    #[builder(default = "None")]
    scheduler_client: Option<TransferSchedulerClient>,
407
408
409
410
411
412

    #[builder(default = "String::from(\"tcp://127.0.0.1:56001\")")]
    leader_pub_url: String,

    #[builder(default = "String::from(\"tcp://127.0.0.1:56002\")")]
    leader_ack_url: String,
Ryan Olson's avatar
Ryan Olson committed
413
414
415
416
417
418
419
420
421
422
}

impl KvbmWorkerConfig {
    pub fn builder() -> KvbmWorkerConfigBuilder {
        KvbmWorkerConfigBuilder::default()
    }
}

pub struct KvbmWorker {
    task: Option<CriticalTaskExecutionHandle>,
423
    block_transfer_handler_rx: Option<oneshot::Receiver<transfer::BlockTransferHandler>>,
Ryan Olson's avatar
Ryan Olson committed
424
425
426
}

impl KvbmWorker {
427
    pub async fn new(config: KvbmWorkerConfig, layout_blocking: bool) -> anyhow::Result<Self> {
Ryan Olson's avatar
Ryan Olson committed
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
        tracing::info!(
            "Initializing KvbmWorker with params: num_device_blocks={}, page_size={}, dtype_width_bytes={}",
            config.num_device_blocks,
            config.page_size,
            config.dtype_width_bytes
        );

        if config.num_device_blocks == 0 {
            return Err(anyhow::anyhow!("num_device_blocks must be greater than 0"));
        }

        let (device_tensors, shape) = load_and_validate_tensors(&config.tensors, config.device_id)?;

        if shape.len() < 3 {
            return Err(anyhow::anyhow!(format!(
                "Unsupported kv cache layout. Got shape: {:?}",
                shape
            )));
        }

448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
        let (layout_type, num_layers, outer_dim, inner_dim) = match config.device_layout_type {
            LayoutType::FullyContiguous => {
                let num_layers = shape[1];
                let outer_dim = shape[2];
                let inner_dim = shape[3..].iter().product::<usize>() / config.page_size;
                tracing::info!(
                    "Inferred layout: num_layers={}, outer_dim={}, page_size={}, inner_dim={}",
                    num_layers,
                    outer_dim,
                    config.page_size,
                    inner_dim
                );

                (
                    LayoutType::FullyContiguous,
                    num_layers,
                    outer_dim,
                    inner_dim,
                )
            }
            LayoutType::LayerSeparate { outer_contiguous } => {
                // Use the already-detected layout type from config (no re-detection needed)
                let layout_type = config.device_layout_type;

                // Extract outer_dim based on the provided outer_contiguous value
                let outer_dim = if outer_contiguous {
                    shape[0] // Outer contiguous: [outer_dim, n_blocks, ...]
                } else {
                    shape[1] // Block contiguous: [n_blocks, outer_dim, ...]
                };

                let num_layers = device_tensors.len();
                let inner_dim = shape[2..].iter().product::<usize>() / config.page_size;

                tracing::info!(
                    "Inferred layout: num_layers={}, outer_dim={}, outer_contiguous={}, page_size={}, inner_dim={}",
                    num_layers,
                    outer_dim,
                    outer_contiguous,
                    config.page_size,
                    inner_dim
                );

                (layout_type, num_layers, outer_dim, inner_dim)
            }
Ryan Olson's avatar
Ryan Olson committed
493
494
        };

495
496
        let bytes_per_block =
            num_layers * outer_dim * config.page_size * inner_dim * config.dtype_width_bytes;
Ryan Olson's avatar
Ryan Olson committed
497
498
499

        let mut layout_builder_instance = LayoutConfigBuilder::default();
        let layout_builder = layout_builder_instance
500
            .num_layers(num_layers)
Ryan Olson's avatar
Ryan Olson committed
501
502
503
504
505
506
507
508
509
510
            .outer_dim(outer_dim)
            .page_size(config.page_size)
            .inner_dim(inner_dim)
            .dtype_width_bytes(config.dtype_width_bytes);

        let device_layout = layout_builder
            .num_blocks(config.num_device_blocks)
            .build()?
            .create_layout(layout_type, device_tensors)?;

511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
        let layout_builder = layout_builder.clone();

        let (task, handler_rx) = if layout_blocking {
            Self::run_blocking_layout_initialization(
                config,
                bytes_per_block,
                device_layout,
                layout_builder,
                layout_type,
            )
            .await?
        } else {
            Self::run_non_blocking_layout_initialization(
                config,
                bytes_per_block,
                device_layout,
                layout_builder,
                layout_type,
            )
            .await?
        };
Ryan Olson's avatar
Ryan Olson committed
532

533
534
535
536
537
538
539
540
541
542
543
544
545
546
        Ok(Self {
            task: Some(task),
            block_transfer_handler_rx: Some(handler_rx),
        })
    }

    async fn run_blocking_layout_initialization(
        config: KvbmWorkerConfig,
        bytes_per_block: usize,
        device_layout: Box<dyn NixlLayout<StorageType = DeviceStorage>>,
        layout_builder: LayoutConfigBuilder,
        layout_type: LayoutType,
    ) -> anyhow::Result<(
        CriticalTaskExecutionHandle,
547
        oneshot::Receiver<transfer::BlockTransferHandler>,
548
    )> {
Richard Huo's avatar
Richard Huo committed
549
        let cancel_token = config.cancel_token.clone();
Ryan Olson's avatar
Ryan Olson committed
550
551
552

        // establish a oneshot channel to get back the raw BlockTransferHandler
        let (handler_tx, handler_rx) = oneshot::channel();
553
        let handler_tx_cell = Arc::new(Mutex::new(Some(handler_tx)));
Ryan Olson's avatar
Ryan Olson committed
554

555
556
        // establish a oneshot channel to block on the main routine to wait for layout allocation readiness
        let (layout_ready_tx, layout_ready_rx) = oneshot::channel::<String>();
557
        let layout_ready_tx_cell = Mutex::new(Some(layout_ready_tx));
558

Ryan Olson's avatar
Ryan Olson committed
559
560
        let scheduler_client = config.scheduler_client.clone();

561
562
        let worker_config = config.clone();
        // start background worker task to do layout allocation for host or disk
Ryan Olson's avatar
Ryan Olson committed
563
564
565
566
        let task = CriticalTaskExecutionHandle::new(
            move |cancel_token| {
                KvbmWorker::worker_task(
                    device_layout,
567
                    layout_builder,
Ryan Olson's avatar
Ryan Olson committed
568
                    layout_type,
569
                    worker_config,
Ryan Olson's avatar
Ryan Olson committed
570
                    cancel_token,
571
572
                    handler_tx_cell,
                    layout_ready_tx_cell,
Ryan Olson's avatar
Ryan Olson committed
573
                    scheduler_client,
574
                    bytes_per_block,
Ryan Olson's avatar
Ryan Olson committed
575
576
577
578
579
580
                )
            },
            cancel_token.clone(),
            "kvbm-worker-task",
        )?;

581
582
583
584
585
586
587
        // waiting for the worker layout allocation ready
        match layout_ready_rx.await {
            Ok(_) => tracing::info!("worker layout allocation finished."),
            Err(_) => tracing::error!("Worker layout dropped without sending"),
        }

        Ok((task, handler_rx))
Ryan Olson's avatar
Ryan Olson committed
588
589
    }

590
591
592
593
594
595
596
597
    async fn run_non_blocking_layout_initialization(
        config: KvbmWorkerConfig,
        bytes_per_block: usize,
        device_layout: Box<dyn NixlLayout<StorageType = DeviceStorage> + Send + 'static>,
        layout_builder: LayoutConfigBuilder,
        layout_type: LayoutType,
    ) -> anyhow::Result<(
        CriticalTaskExecutionHandle,
598
        oneshot::Receiver<transfer::BlockTransferHandler>,
599
    )> {
Richard Huo's avatar
Richard Huo committed
600
        let cancel_token = config.cancel_token.clone();
601
602
603
        let scheduler_client = config.scheduler_client.clone();

        // channel to get BlockTransferHandler back to the caller
604
        let (handler_tx, handler_rx) = oneshot::channel::<transfer::BlockTransferHandler>();
605
        let handler_tx_cell = Arc::new(Mutex::new(Some(handler_tx)));
606
607
608

        // channel that the worker will use to signal layout readiness
        let (layout_ready_tx, layout_ready_rx) = oneshot::channel::<String>();
609
        let layout_ready_tx_cell = Mutex::new(Some(layout_ready_tx));
610
611
612
613
614
615
616
617
618
619
620
621

        // clone what we need inside the orchestrator
        let worker_config = config.clone();
        let cancel_token_for_task = cancel_token.clone();

        // Single task that orchestrates everything in-order.
        let task = CriticalTaskExecutionHandle::new(
            move |ct| {
                let cfg = worker_config.clone();
                let scheduler = scheduler_client.clone();

                async move {
622
                    // Start the long-running worker.
623
624
625
626
627
628
629
630
631
632
                    let dev_layout = device_layout; // moved in
                    let lb = layout_builder; // moved in
                    let lt = layout_type; // moved in

                    let worker_fut = KvbmWorker::worker_task(
                        dev_layout,
                        lb,
                        lt,
                        cfg.clone(),
                        ct.clone(),
633
634
                        handler_tx_cell,
                        layout_ready_tx_cell,
635
                        scheduler,
636
                        bytes_per_block,
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
                    );

                    // If worker_task returns Result, handle/log it inside the spawned task.
                    tokio::spawn(async move {
                        if let Err(e) = worker_fut.await {
                            tracing::error!("worker_task exited with error: {e:#}");
                        }
                    });

                    // 3) wait for the worker’s layout allocation readiness
                    match layout_ready_rx.await {
                        Ok(_) => tracing::info!("worker layout allocation finished."),
                        Err(_) => tracing::warn!("worker layout readiness channel dropped"),
                    }

                    Ok::<(), anyhow::Error>(())
                }
            },
            cancel_token_for_task,
            "kvbm-worker-task",
        )?;

        Ok((task, handler_rx))
    }
661

Ryan Olson's avatar
Ryan Olson committed
662
663
664
665
666
    /// One-time use method to extract the block transfer handler from the worker.
    ///
    /// This is a bit of a hack. Improve the API design around this in the future.
    pub fn block_transfer_handler_rx(
        &mut self,
667
    ) -> Option<tokio::sync::oneshot::Receiver<BlockTransferHandler>> {
Ryan Olson's avatar
Ryan Olson committed
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
        self.block_transfer_handler_rx.take()
    }

    fn make_layout<S: Storage, M: BlockMetadata>(
        mut layout: Box<dyn NixlLayout<StorageType = S>>,
        agent: &Option<NixlAgent>,
        block_set_idx: usize,
        worker_id: usize,
    ) -> anyhow::Result<Vec<Block<S, locality::Local, M>>> {
        // Register with NIXL, if applicable.
        if let Some(agent) = agent {
            layout.nixl_register(agent, None)?;
        }

        // Convert the layout into blocks.
        let layout: Arc<dyn NixlLayout<StorageType = S>> = Arc::from(layout);
        let blocks = layout_to_blocks::<_, M>(layout, block_set_idx, worker_id as u64)?;
        Ok(blocks)
    }

688
689
690
    #[allow(clippy::too_many_arguments)]
    async fn worker_task(
        device_layout: Box<dyn NixlLayout<StorageType = DeviceStorage>>,
691
692
        layout_builder: LayoutConfigBuilder,
        _device_layout_type: LayoutType,
693
694
        config: KvbmWorkerConfig,
        cancel_token: CancellationToken,
695
        handler_tx: Arc<Mutex<Option<oneshot::Sender<BlockTransferHandler>>>>,
696
        layout_ready_tx: tokio::sync::Mutex<Option<oneshot::Sender<String>>>,
697
        scheduler_client: Option<TransferSchedulerClient>,
698
        bytes_per_block: usize,
699
    ) -> anyhow::Result<()> {
Richard Huo's avatar
Richard Huo committed
700
        let worker_id = config.device_id;
701
702
        // Readiness gating for ping
        let state = Arc::new(WorkerState::new());
Ryan Olson's avatar
Ryan Olson committed
703

704
        // Cell to publish the transfer handler
705
        let transfer_handler_cell: Arc<RwLock<Option<BlockTransferHandler>>> =
706
            Arc::new(RwLock::new(None));
707

708
709
        // Build handlers map
        let mut handlers: HashMap<String, Arc<dyn Handler>> = HashMap::new();
Ryan Olson's avatar
Ryan Olson committed
710

711
712
713
714
715
716
717
        handlers.insert(
            ZMQ_PING_MESSAGE.to_string(),
            Arc::new(GatedPing {
                state: state.clone(),
                layout_ready_tx,
            }) as Arc<dyn Handler>,
        );
Ryan Olson's avatar
Ryan Olson committed
718

719
720
721
722
723
724
725
        handlers.insert(
            ZMQ_WORKER_METADATA_MESSAGE.to_string(),
            Arc::new(WorkerMetadataHandler {
                num_device_blocks: config.num_device_blocks,
                bytes_per_block,
            }) as Arc<dyn Handler>,
        );
Ryan Olson's avatar
Ryan Olson committed
726

727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
        handlers.insert(
            ZMQ_LEADER_METADATA_MESSAGE.to_string(),
            Arc::new(LeaderMetadataHandler {
                state: state.clone(),
                device_layout: tokio::sync::Mutex::new(Some(device_layout)), // moved in
                layout_builder,                                              // moved
                worker_config: config.clone(),
                worker_id,
                device_id: config.device_id,
                scheduler_client,
                handler_cell: transfer_handler_cell.clone(),
                handler_tx, // sends BlockTransferHandler to caller
                started: AtomicBool::new(false),
            }) as Arc<dyn Handler>,
        );
Ryan Olson's avatar
Ryan Olson committed
742

743
744
        // transfer requests get dispatched to built handler (after allocation)
        handlers.insert(
Ryan Olson's avatar
Ryan Olson committed
745
            ZMQ_TRANSFER_BLOCKS_MESSAGE.to_string(),
746
747
748
749
            Arc::new(BlockTransferDispatch {
                cell: transfer_handler_cell.clone(),
            }) as Arc<dyn Handler>,
        );
Ryan Olson's avatar
Ryan Olson committed
750
751

        let _zmq_worker = ZmqActiveMessageWorker::new(
752
753
            &config.leader_pub_url,
            &config.leader_ack_url,
Ryan Olson's avatar
Ryan Olson committed
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
            handlers,
            cancel_token.clone(),
        )?;

        // TODO: Some sort of fancy loop here.
        // For now, just wait for cancellation.
        cancel_token.cancelled().await;

        Ok(())
    }
}

impl Drop for KvbmWorker {
    fn drop(&mut self) {
        if let Some(task) = self.task.take() {
            task.cancel();
            task.detach();
        }
    }
}