Skip to content

GitLab

".buildkite/vscode:/vscode.git/clone" did not exist on "1b7c7f5159484063af28cb47809d79e83d3301ec"
Unverified Commit 94aa2a7b authored by Richard Huo's avatar Richard Huo Committed by GitHub
Browse files

refactor: kvbm modularity DIS-657 Eliminate ETCD from the leader-worker initialization (#3202) 


Signed-off-by: default avatarrichardhuo-nv <rihuo@nvidia.com>
parent 5d90e530
Show whitespace changes
Inline Side-by-side
Showing with 738 additions and 715 deletions
Cargo.lock
View file @ 94aa2a7b
......@@ -2136,6 +2136,7 @@ dependencies = [
"async_zmq",
"axum 0.8.4",
"axum-server",
"bincode",
"bitflags 2.9.4",
"blake3",
"bs62",
......
components/backends/vllm/deploy/disagg_kvbm_2p2d.yaml
View file @ 94aa2a7b
......@@ -54,10 +54,6 @@ spec:
envs:
- name: DYN_KVBM_CPU_CACHE_GB
value: "100"
- name: DYN_KVBM_BARRIER_ID_PREFIX
valueFrom:
fieldRef:
fieldPath: metadata.name
extraPodSpec:
mainContainer:
image: nvcr.io/nvidia/ai-dynamo/vllm-runtime:my-tag
......
components/backends/vllm/deploy/disagg_kvbm_tp2.yaml
View file @ 94aa2a7b
......@@ -58,10 +58,6 @@ spec:
envs:
- name: DYN_KVBM_CPU_CACHE_GB
value: "100"
- name: DYN_KVBM_BARRIER_ID_PREFIX
valueFrom:
fieldRef:
fieldPath: metadata.name
extraPodSpec:
mainContainer:
image: nvcr.io/nvidia/ai-dynamo/vllm-runtime:my-tag
......
components/backends/vllm/launch/disagg_kvbm_2p2d.sh
View file @ 94aa2a7b
......@@ -15,7 +15,6 @@ CUDA_VISIBLE_DEVICES=1 python3 -m dynamo.vllm --model Qwen/Qwen3-0.6B --connecto
# run prefill workers on GPU 2 and 3 with KVBM enabled using 20GB of CPU cache
# NOTE: use different barrier id prefixes for each prefill worker to avoid conflicts
# NOTE: remove --enforce-eager for production use
DYN_KVBM_BARRIER_ID_PREFIX=kvbm_0 \
DYN_KVBM_CPU_CACHE_GB=20 \
CUDA_VISIBLE_DEVICES=2 \
python3 -m dynamo.vllm \
......@@ -24,7 +23,8 @@ CUDA_VISIBLE_DEVICES=2 \
--connector kvbm nixl \
--enforce-eager &
DYN_KVBM_BARRIER_ID_PREFIX=kvbm_1 \
DYN_KVBM_LEADER_ZMQ_PUB_PORT=56003 \
DYN_KVBM_LEADER_ZMQ_ACK_PORT=56004 \
DYN_KVBM_CPU_CACHE_GB=20 \
CUDA_VISIBLE_DEVICES=3 \
python3 -m dynamo.vllm \
......
lib/bindings/python/Cargo.lock
View file @ 94aa2a7b
......@@ -1466,6 +1466,7 @@ dependencies = [
"async_zmq",
"axum",
"axum-server",
"bincode",
"bitflags 2.9.3",
"blake3",
"bs62",
......
lib/bindings/python/rust/llm/block_manager/distributed.rs
View file @ 94aa2a7b
......@@ -8,5 +8,5 @@ mod utils;
mod worker;
pub use leader::KvbmLeader;
pub use utils::get_barrier_id_prefix;
pub use utils::{get_leader_zmq_ack_url, get_leader_zmq_pub_url};
pub use worker::{KvbmWorker, PyLayoutType, VllmTensor};
lib/bindings/python/rust/llm/block_manager/distributed/leader.rs
View file @ 94aa2a7b
......@@ -2,12 +2,11 @@
// SPDX-License-Identifier: Apache-2.0
use super::*;
use utils::get_barrier_id_prefix;
use derive_getters::Dissolve;
use llm_rs::block_manager::distributed::{
KvbmLeader as KvbmLeaderImpl, KvbmLeaderConfig, KvbmLeaderNumBlocksConfig,
};
use utils::{get_leader_zmq_ack_url, get_leader_zmq_pub_url};
const CPU_CACHE: &str = "DYN_KVBM_CPU_CACHE_GB";
const CPU_CACHE_OVERRIDE: &str = "DYN_KVBM_CPU_CACHE_OVERRIDE_NUM_BLOCKS";
......@@ -72,17 +71,16 @@ impl KvbmLeader {
#[new]
#[pyo3(signature = (world_size, drt))]
fn new(world_size: usize, drt: DistributedRuntime) -> PyResult<Self> {
let barrier_id_prefix = get_barrier_id_prefix();
let leader_init_timeout_sec: u64 =
get_leader_init_timeout_secs(LEADER_WORKER_INIT_TIMEOUT_SECS);
let config = KvbmLeaderConfig::builder()
.barrier_id_prefix(barrier_id_prefix)
.world_size(world_size)
.leader_init_timeout_secs(leader_init_timeout_sec)
.drt(drt.inner().clone())
.host_blocks_config(get_blocks_config(CPU_CACHE, CPU_CACHE_OVERRIDE))
.disk_blocks_config(get_blocks_config(DISK_CACHE, DISK_CACHE_OVERRIDE))
.leader_pub_url(get_leader_zmq_pub_url())
.leader_ack_url(get_leader_zmq_ack_url())
.build()
.map_err(to_pyerr)?;
......
lib/bindings/python/rust/llm/block_manager/distributed/utils.rs
View file @ 94aa2a7b
// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
use std::env;
pub fn get_barrier_id_prefix() -> String {
std::env::var("DYN_KVBM_BARRIER_ID_PREFIX")
const DEFAULT_LEADER_ZMQ_HOST: &str = "127.0.0.1";
const DEFAULT_LEADER_ZMQ_PUB_PORT: u16 = 56001;
const DEFAULT_LEADER_ZMQ_ACK_PORT: u16 = 56002;
fn read_env_trimmed(key: &str) -> Option<String> {
env::var(key)
.ok()
.filter(|s| !s.trim().is_empty())
.unwrap_or_else(|| "kvbm".to_string())
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
}
fn parse_port_u16(s: &str) -> Option<u16> {
match s.parse::<u32>() {
Ok(v) if (1..=65535).contains(&v) => Some(v as u16),
_ => None,
}
}
fn validated_port_from_env(key: &str, default_port: u16) -> u16 {
if let Some(val) = read_env_trimmed(key) {
if let Some(p) = parse_port_u16(&val) {
if p < 1024 {
tracing::warn!("{key} is a privileged port ({p}); binding may require extra caps");
}
return p;
} else {
tracing::warn!("{key} invalid value '{val}', falling back to default {default_port}");
}
}
default_port
}
fn get_leader_zmq_host() -> String {
read_env_trimmed("DYN_KVBM_LEADER_ZMQ_HOST")
.unwrap_or_else(|| DEFAULT_LEADER_ZMQ_HOST.to_string())
}
fn get_leader_zmq_pub_port() -> String {
validated_port_from_env("DYN_KVBM_LEADER_ZMQ_PUB_PORT", DEFAULT_LEADER_ZMQ_PUB_PORT).to_string()
}
fn get_leader_zmq_ack_port() -> String {
validated_port_from_env("DYN_KVBM_LEADER_ZMQ_ACK_PORT", DEFAULT_LEADER_ZMQ_ACK_PORT).to_string()
}
pub fn get_leader_zmq_pub_url() -> String {
format!(
"tcp://{}:{}",
get_leader_zmq_host(),
get_leader_zmq_pub_port()
)
}
pub fn get_leader_zmq_ack_url() -> String {
format!(
"tcp://{}:{}",
get_leader_zmq_host(),
get_leader_zmq_ack_port()
)
}
lib/bindings/python/rust/llm/block_manager/distributed/worker.rs
View file @ 94aa2a7b
// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
use utils::{get_leader_zmq_ack_url, get_leader_zmq_pub_url};
use super::*;
use std::sync::Arc;
use utils::get_barrier_id_prefix;
use llm_rs::block_manager::distributed::{
BlockTransferHandler as RustBlockTransferHandler, KvbmWorker as KvbmWorkerImpl,
......@@ -171,8 +172,6 @@ impl KvbmWorker {
vllm_tensors.push(Arc::new(vllm_tensor));
}
let barrier_id_prefix = get_barrier_id_prefix();
let config = KvbmWorkerConfig::builder()
.drt(drt)
.num_device_blocks(num_device_blocks)
......@@ -180,7 +179,6 @@ impl KvbmWorker {
.tensors(vllm_tensors)
.device_id(device_id)
.dtype_width_bytes(dtype_width_bytes)
.barrier_id_prefix(barrier_id_prefix)
.device_layout_type(
device_layout_type
.map(|py_layout| py_layout.into())
......@@ -196,6 +194,8 @@ impl KvbmWorker {
.map(|py_layout| py_layout.into())
.unwrap_or(LayoutType::FullyContiguous),
)
.leader_pub_url(get_leader_zmq_pub_url())
.leader_ack_url(get_leader_zmq_ack_url())
.build()
.map_err(to_pyerr)?;
......
lib/bindings/python/rust/llm/block_manager/vllm/connector/leader.rs
View file @ 94aa2a7b
......@@ -150,9 +150,6 @@ impl KvConnectorLeader {
let _ = slot_manager_cell.set(sm);
// another barrier sync to make sure worker init won't return before leader is ready
let _ = leader.run_leader_readiness_barrier_blocking(drt);
if leader_ready_tx.send("finished".to_string()).is_err() {
tracing::error!("main routine receiver dropped before result was sent");
}
......
lib/bindings/python/rust/llm/block_manager/vllm/connector/leader/recorder.rs
View file @ 94aa2a7b
......@@ -166,9 +166,6 @@ impl KvConnectorLeaderRecorder {
let _ = slot_manager_cell.set(sm);
// another barrier sync to make sure worker init won't return before leader is ready
leader.spawn_leader_readiness_barrier(drt);
if leader_ready_tx.send("finished".to_string()).is_err() {
tracing::error!("main routine receiver dropped before result was sent");
}
......
lib/bindings/python/rust/llm/block_manager/vllm/connector/trtllm_leader.rs
View file @ 94aa2a7b
......@@ -126,9 +126,6 @@ impl KvConnectorLeader {
let _ = slot_manager_cell.set(sm);
// another barrier sync to make sure worker init won't return before leader is ready
leader.spawn_leader_readiness_barrier(drt);
tracing::info!("KvConnectorLeader init complete.");
});
}
......
lib/bindings/python/rust/llm/block_manager/vllm/connector/trtllm_worker.rs
View file @ 94aa2a7b
......@@ -10,7 +10,7 @@ use std::collections::HashSet;
use std::sync::{Arc, OnceLock};
use super::*;
use crate::llm::block_manager::distributed::get_barrier_id_prefix;
use crate::llm::block_manager::distributed::{get_leader_zmq_ack_url, get_leader_zmq_pub_url};
use crate::llm::block_manager::vllm::connector::worker::event_sync_blocking;
use crate::{
DistributedRuntime as PyDistributedRuntime, llm::block_manager::distributed::VllmTensor,
......@@ -138,7 +138,8 @@ impl Worker for KvConnectorWorker {
.device_layout_type(LayoutType::FullyContiguous)
.host_layout_type(LayoutType::FullyContiguous)
.disk_layout_type(LayoutType::FullyContiguous)
.barrier_id_prefix(get_barrier_id_prefix())
.leader_pub_url(get_leader_zmq_pub_url())
.leader_ack_url(get_leader_zmq_ack_url())
.scheduler_client(Some(self.transfer_client.clone()))
.build()?;
......
lib/bindings/python/rust/llm/block_manager/vllm/connector/worker.rs
View file @ 94aa2a7b
......@@ -10,7 +10,7 @@ use std::collections::HashSet;
use std::sync::{Arc, OnceLock};
use super::*;
use crate::llm::block_manager::distributed::get_barrier_id_prefix;
use crate::llm::block_manager::distributed::{get_leader_zmq_ack_url, get_leader_zmq_pub_url};
use crate::{
DistributedRuntime as PyDistributedRuntime, llm::block_manager::distributed::VllmTensor,
to_pyerr,
......@@ -200,7 +200,8 @@ impl Worker for KvConnectorWorker {
.tensors(vllm_tensors)
.device_id(device_id)
.dtype_width_bytes(dtype_width_bytes)
.barrier_id_prefix(get_barrier_id_prefix())
.leader_pub_url(get_leader_zmq_pub_url())
.leader_ack_url(get_leader_zmq_ack_url())
.scheduler_client(Some(self.transfer_client.clone()))
.device_layout_type(detected_device_layout_type)
.host_layout_type(host_layout_type.unwrap_or(LayoutType::FullyContiguous))
......
lib/llm/Cargo.toml
View file @ 94aa2a7b
......@@ -85,6 +85,7 @@ offset-allocator = "0.2"
regex = "1"
rayon = "1"
dashmap = { version = "5.5.3" }
bincode = "1"
# input/text
dialoguer = { version = "0.11", default-features = false, features = [
......
lib/llm/src/block_manager/distributed.rs
View file @ 94aa2a7b
......@@ -123,14 +123,12 @@ mod tests {
async fn build_leader_and_workers(num_workers: usize) -> Result<(KvbmLeader, Vec<KvbmWorker>)> {
let mut workers = Vec::new();
let barrier_id = get_unique_barrier_id();
for i in 0..num_workers {
let tensors: Vec<Arc<dyn TorchTensor>> =
vec![Arc::new(MockTensor::new(vec![2, NUM_BLOCKS, 4096]))];
let config = KvbmWorkerConfig::builder()
.barrier_id_prefix(barrier_id.clone())
.num_device_blocks(NUM_BLOCKS)
.tensors(tensors)
.device_id(i)
......@@ -151,7 +149,6 @@ mod tests {
};
let leader_config = KvbmLeaderConfig::builder()
.barrier_id_prefix(barrier_id)
.world_size(num_workers)
.host_blocks_config(host_blocks)
.disk_blocks_config(disk_blocks)
......
lib/llm/src/block_manager/distributed/leader.rs
View file @ 94aa2a7b
......@@ -3,15 +3,10 @@
use super::*;
use dynamo_runtime::DistributedRuntime;
use utils::*;
use zmq::*;
use dynamo_runtime::utils::leader_worker_barrier::LeaderBarrier;
use anyhow::{Context, anyhow};
use derive_builder::Builder;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::time::Duration;
......@@ -20,15 +15,6 @@ use tokio::sync::OnceCell;
use tokio::sync::oneshot;
use tokio::time::sleep;
/// Data that is sent to workers over ETCD to establish a ZMQ connection.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KvbmLeaderData {
pub pub_url: String,
pub ack_url: String,
pub num_host_blocks: usize,
pub num_disk_blocks: usize,
}
#[derive(Builder, Clone, Debug, Default)]
pub struct KvbmLeaderNumBlocksConfig {
#[builder(default = "0.0")]
......@@ -51,10 +37,6 @@ fn compute_num_blocks(
#[derive(Builder, Clone, Debug)]
pub struct KvbmLeaderConfig {
/// The barrier id to use for syncing with workers.
#[builder(default = "String::from(\"kvbm\")")]
barrier_id_prefix: String,
/// The world size.
#[builder(default = "1")]
world_size: usize,
......@@ -63,14 +45,17 @@ pub struct KvbmLeaderConfig {
#[builder(default = "120")]
leader_init_timeout_secs: u64,
#[builder(setter(strip_option))]
drt: Option<DistributedRuntime>,
#[builder(default = "KvbmLeaderNumBlocksConfig::default()")]
host_blocks_config: KvbmLeaderNumBlocksConfig,
#[builder(default = "KvbmLeaderNumBlocksConfig::default()")]
disk_blocks_config: KvbmLeaderNumBlocksConfig,
#[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,
}
impl KvbmLeaderConfig {
......@@ -79,6 +64,11 @@ impl KvbmLeaderConfig {
}
pub fn sanity_check(&self) -> anyhow::Result<()> {
if self.leader_pub_url == self.leader_ack_url {
anyhow::bail!(
"leader_pub_url and leader_ack_url must differ (same endpoint would fail to bind)."
);
}
let cpu = &self.host_blocks_config;
let disk = &self.disk_blocks_config;
let cpu_configured = cpu.num_blocks_overriden > 0 || cpu.cache_size_in_gb > 0.0;
......@@ -121,166 +111,24 @@ pub struct KvbmLeader {
state: Arc<KvbmLeaderState>,
zmq_leader: Arc<OnceCell<ZmqActiveMessageLeader>>,
config: KvbmLeaderConfig,
//readiness flags
workers_sync_ready: Arc<AtomicBool>,
workers_sync_ready_notify: Arc<Notify>,
workers_sync_done: Arc<AtomicBool>,
}
impl KvbmLeader {
pub async fn new(mut config: KvbmLeaderConfig) -> anyhow::Result<Self> {
let drt = match config.drt.take() {
Some(dtr) => dtr,
None => {
anyhow::bail!("No distributed runtime provided");
}
};
let leader_sockets = new_leader_sockets("tcp://127.0.0.1")?;
pub async fn new(config: KvbmLeaderConfig) -> anyhow::Result<Self> {
let leader_sockets = new_leader_sockets(&config.leader_pub_url, &config.leader_ack_url)?;
let leader = Self {
state: Arc::new(KvbmLeaderState::default()),
zmq_leader: Arc::new(tokio::sync::OnceCell::new()),
config,
workers_sync_ready: Arc::new(AtomicBool::new(false)),
workers_sync_ready_notify: Arc::new(Notify::new()),
workers_sync_done: Arc::new(AtomicBool::new(false)),
};
let cancel_token = tokio_util::sync::CancellationToken::new();
// The leader_sockets struct cannot be cloned,
// so we use a tuple to "struct" the two urls
let leader_urls = (
leader_sockets.pub_url.clone(),
leader_sockets.ack_url.clone(),
);
leader.spawn_barrier_task(drt, leader_urls);
leader.spawn_zmq_task(leader_sockets, cancel_token);
Ok(leader)
}
fn spawn_barrier_task(&self, drt: DistributedRuntime, leader_urls: (String, String)) {
let state = self.state.clone();
let leader_config = self.config.clone();
let ready = Arc::clone(&self.workers_sync_ready);
let notify = Arc::clone(&self.workers_sync_ready_notify);
let done = Arc::clone(&self.workers_sync_done);
tokio::spawn(async move {
match KvbmLeader::run_barrier_sync(drt, leader_urls, leader_config).await {
Ok((num_device_blocks, num_host_blocks, num_disk_blocks)) => {
// write back results
state
.num_device_blocks
.store(num_device_blocks, Ordering::Release);
state
.num_host_blocks
.store(num_host_blocks, Ordering::Release);
state
.num_disk_blocks
.store(num_disk_blocks, Ordering::Release);
ready.store(true, Ordering::Release);
done.store(true, Ordering::Release);
notify.notify_waiters();
}
Err(e) => {
tracing::error!("Barrier sync failed: {e:?}");
done.store(true, Ordering::Release);
notify.notify_waiters();
}
}
});
}
async fn run_barrier_sync(
drt: DistributedRuntime,
leader_urls: (String, String),
leader_config: KvbmLeaderConfig,
) -> anyhow::Result<(usize, usize, usize)> {
let barrier_id_worker_to_leader =
format!("{}{}", leader_config.barrier_id_prefix, "-worker-to-leader");
tracing::info!(
"Syncing leader barrier with {} workers on barrier id {}",
leader_config.world_size,
barrier_id_worker_to_leader
);
// Build our leader barrier and publish the data.
// TODO: Use a separate timeout parameter from the ZMQ connection timeout
let worker_to_leader_barrier: LeaderBarrier<(), worker::KvbmWorkerData> =
LeaderBarrier::new(
barrier_id_worker_to_leader.clone(),
leader_config.world_size,
Some(Duration::from_secs(leader_config.leader_init_timeout_secs)),
);
let worker_data = worker_to_leader_barrier
.sync(&drt, &())
.await
.map_err(|e| anyhow::anyhow!("Failed to sync worker to leader barrier: {:?}", e))?;
let num_device_blocks = worker_data
.values()
.map(|data| data.num_device_blocks)
.min()
.unwrap();
// TODO: this works for TP, need to redefine bytes_per_block when we enable the DP/PP
let bytes_per_block: usize = worker_data.values().map(|d| d.bytes_per_block).sum();
assert!(
bytes_per_block > 0,
"bytes_per_block must be greater than 0"
);
tracing::info!(
"Worker to leader barrier synced with {} workers",
leader_config.world_size
);
tracing::debug!("Worker data: {:?}", worker_data);
let num_host_blocks =
compute_num_blocks(&leader_config.host_blocks_config, bytes_per_block);
let num_disk_blocks =
compute_num_blocks(&leader_config.disk_blocks_config, bytes_per_block);
// Start the second sync to transfer num_host_blocks and num_disk_blocks to worker
let barrier_id_leader_to_worker =
format!("{}{}", leader_config.barrier_id_prefix, "-leader-to-worker");
tracing::info!(
"Syncing leader barrier with {} workers on barrier id {}",
leader_config.world_size,
barrier_id_leader_to_worker
);
let (leader_pub_url, leader_ack_url) = leader_urls;
let zmq_data_leader_to_worker = Arc::new(KvbmLeaderData {
pub_url: leader_pub_url,
ack_url: leader_ack_url,
num_host_blocks,
num_disk_blocks,
});
let leader_to_worker_barrier: LeaderBarrier<KvbmLeaderData, ()> = LeaderBarrier::new(
barrier_id_leader_to_worker.clone(),
leader_config.world_size,
Some(Duration::from_secs(leader_config.leader_init_timeout_secs)),
);
let _worker_data = leader_to_worker_barrier
.sync(&drt, zmq_data_leader_to_worker.as_ref())
.await
.map_err(|e| anyhow::anyhow!("Failed to sync leader to worker barrier: {:?}", e))?;
tracing::info!(
"Worker to leader barrier synced with {} workers",
leader_config.world_size
);
Ok((num_device_blocks, num_host_blocks, num_disk_blocks))
}
fn spawn_zmq_task(
&self,
leader_sockets: LeaderSockets,
......@@ -290,133 +138,59 @@ impl KvbmLeader {
let state = self.state.clone();
let world_size = self.config.world_size;
let timeout = self.config.leader_init_timeout_secs;
let host_cfg = self.config.host_blocks_config.clone();
let disk_cfg = self.config.disk_blocks_config.clone();
// capture num_device_blocks so we can set it inside the closure
let num_device_blocks_cell = state.num_device_blocks.clone();
let num_host_blocks_cell = state.num_host_blocks.clone();
let num_disk_blocks_cell = state.num_disk_blocks.clone();
tokio::spawn(async move {
let res = ZmqActiveMessageLeader::new(
let res = ZmqActiveMessageLeader::new_with_handshake(
leader_sockets,
world_size,
std::time::Duration::from_secs(timeout),
cancel,
cancel.clone(),
move |workers: &[WorkerMetadata]| -> LeaderMetadata {
// Record device blocks: min across workers
if let Some(min_dev) = workers.iter().map(|w| w.num_device_blocks).min() {
num_device_blocks_cell.store(min_dev, Ordering::Release);
}
// For TP, sum bytes_per_block; adjust policy for DP/PP if needed.
let bytes_per_block: usize = workers.iter().map(|w| w.bytes_per_block).sum();
let num_host_blocks = compute_num_blocks(&host_cfg, bytes_per_block);
let num_disk_blocks = compute_num_blocks(&disk_cfg, bytes_per_block);
// store into leader state
num_host_blocks_cell.store(num_host_blocks, Ordering::Release);
num_disk_blocks_cell.store(num_disk_blocks, Ordering::Release);
LeaderMetadata {
num_host_blocks,
num_disk_blocks,
}
},
)
.await;
match res {
Ok(zmq) => {
let _ = cell.set(zmq);
// mark ready
state
.workers_allocation_ready
.store(true, Ordering::Release);
state.workers_ready_notify.notify_waiters();
tracing::info!("ZMQ handshake complete; workers allocation ready");
}
Err(e) => {
tracing::error!("ZMQ init failed: {e:?}");
tracing::error!("ZMQ init/handshake failed: {e:?}");
}
}
});
}
// This is supposed to be used in non-blocking leader initialization
pub fn spawn_leader_readiness_barrier(&self, drt: DistributedRuntime) {
let timeout_secs = self.config.leader_init_timeout_secs;
let state = self.state.clone();
let leader_config = self.config.clone();
let handle = drt.runtime().primary();
handle.spawn(async move {
if !state.workers_allocation_ready.load(Ordering::Acquire) {
// Wait until ZMQ marks ready or we time out.
let waited = tokio::time::timeout(
Duration::from_secs(timeout_secs),
state.workers_ready_notify.notified(),
)
.await;
if waited.is_err() {
tracing::error!(
"leader readiness barrier wait timed out after {timeout_secs} seconds"
);
return;
}
// Double-check the flag (Acquire) after wakeup.
if !state.workers_allocation_ready.load(Ordering::Acquire) {
tracing::error!("leader readiness notify fired but flag not set; aborting");
return;
}
}
match KvbmLeader::run_leader_readiness(drt, leader_config).await {
Ok(()) => {
tracing::info!("leader readiness barrier synced!");
}
Err(e) => {
tracing::error!("leader readiness barrier failed: {e:?}");
}
}
});
}
// This is supposed to be used in blocking leader initialization
pub fn run_leader_readiness_barrier_blocking(
&self,
drt: DistributedRuntime,
) -> anyhow::Result<()> {
let state = self.state.clone();
let timeout_secs = self.config.leader_init_timeout_secs;
let leader_config = self.config.clone();
tokio::task::block_in_place(|| {
tokio::runtime::Handle::current()
.block_on(async move {
// Create the future *before* checking the flag to avoid a lost-notify race.
let notified = state.workers_ready_notify.notified();
if !state.workers_allocation_ready.load(Ordering::Acquire) {
// Wait (with timeout) until ZMQ task marks ready.
tokio::time::timeout(Duration::from_secs(timeout_secs), notified)
.await
.map_err(|_| anyhow!("timed out waiting for workers_allocation_ready after {timeout_secs} seconds"))?;
// Double-check after wake to ensure the flag is actually set.
if !state.workers_allocation_ready.load(Ordering::Acquire) {
return Err(anyhow!(
"notified but workers_allocation_ready is still false"
));
}
}
KvbmLeader::run_leader_readiness(drt, leader_config).await
})
.context("leader readiness barrier failed")
})
}
async fn run_leader_readiness(
drt: DistributedRuntime,
leader_config: KvbmLeaderConfig,
) -> anyhow::Result<()> {
let barrier_id_leader_ready =
format!("{}{}", leader_config.barrier_id_prefix, "-leader-ready");
tracing::info!(
"Syncing leader readiness barrier with {} workers on barrier id {}",
leader_config.world_size,
barrier_id_leader_ready
);
let leader_readiness_barrier: LeaderBarrier<(), ()> = LeaderBarrier::new(
barrier_id_leader_ready.clone(),
leader_config.world_size,
Some(Duration::from_secs(leader_config.leader_init_timeout_secs)),
);
let _ = leader_readiness_barrier
.sync(&drt, &())
.await
.map_err(|e| {
anyhow::anyhow!("Failed to sync leader readiness barrier on leader: {:?}", e)
})?;
Ok(())
}
pub async fn transfer_blocks_request(
&self,
request: BlockTransferRequest,
......@@ -429,14 +203,6 @@ impl KvbmLeader {
zmq.broadcast(ZMQ_TRANSFER_BLOCKS_MESSAGE, data).await
}
pub fn is_worker_sync_ready(&self) -> bool {
self.workers_sync_ready.load(Ordering::Acquire)
}
pub fn is_worker_sync_done(&self) -> bool {
self.workers_sync_done.load(Ordering::Acquire)
}
pub fn num_device_blocks(&self) -> usize {
self.state.num_device_blocks.load(Ordering::Acquire)
}
......@@ -450,26 +216,12 @@ impl KvbmLeader {
}
pub async fn wait_worker_sync_ready(&self) -> bool {
if self.is_worker_sync_ready() {
return true;
}
if self.is_worker_sync_done() {
return false;
}
let notified = self.workers_sync_ready_notify.notified();
if self.is_worker_sync_ready() {
if self.state.workers_allocation_ready.load(Ordering::Acquire) {
return true;
}
if self.is_worker_sync_done() {
return false;
}
// bounded wait
let notified = self.state.workers_ready_notify.notified();
tokio::select! {
_ = notified => {
self.is_worker_sync_ready()
}
_ = notified => true,
_ = sleep(Duration::from_secs(self.config.leader_init_timeout_secs)) => false,
}
}
......
lib/llm/src/block_manager/distributed/utils.rs
View file @ 94aa2a7b
......@@ -7,8 +7,22 @@ use serde::{Deserialize, Serialize};
use crate::block_manager::connector::protocol::LeaderTransferRequest;
pub const ZMQ_PING_MESSAGE: &str = "ping";
pub const ZMQ_WORKER_METADATA_MESSAGE: &str = "worker_metadata";
pub const ZMQ_LEADER_METADATA_MESSAGE: &str = "leader_metadata";
pub const ZMQ_TRANSFER_BLOCKS_MESSAGE: &str = "transfer_blocks";
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkerMetadata {
pub num_device_blocks: usize,
pub bytes_per_block: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LeaderMetadata {
pub num_host_blocks: usize,
pub num_disk_blocks: usize,
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Copy)]
pub enum BlockTransferPool {
Device,
......
lib/llm/src/block_manager/distributed/worker.rs
View file @ 94aa2a7b
......@@ -3,8 +3,7 @@
use super::*;
use leader::KvbmLeaderData;
use async_trait::async_trait;
use transfer::*;
use utils::*;
use zmq::*;
......@@ -23,23 +22,32 @@ use crate::block_manager::{
use derive_builder::Builder;
use nixl_sys::Agent as NixlAgent;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use tokio::runtime::Handle;
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;
use dynamo_runtime::{
DistributedRuntime,
utils::{leader_worker_barrier::WorkerBarrier, task::CriticalTaskExecutionHandle},
};
use dynamo_runtime::{DistributedRuntime, utils::task::CriticalTaskExecutionHandle};
use tokio::sync::{Mutex, RwLock, oneshot};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KvbmWorkerData {
pub num_device_blocks: usize,
pub bytes_per_block: usize,
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)
}
}
pub fn load_and_validate_tensors(
......@@ -82,6 +90,269 @@ pub fn load_and_validate_tensors(
Ok((device_tensors, shape.unwrap()))
}
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>,
) -> 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,
};
let transfer_context = Arc::new(TransferContext::new(
Arc::new(Some(agent)),
DeviceAllocator::new(device_id)?.ctx().new_stream()?,
Handle::current(),
Some(pool_config),
));
// 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>(
host_layout,
transfer_context.nixl_agent().as_ref(),
1,
worker_id,
)?)
} else {
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,
transfer_context.nixl_agent().as_ref(),
2,
worker_id,
)?)
} else {
None
};
let handler = BlockTransferHandler::new(
device_blocks,
host_blocks,
disk_blocks,
transfer_context,
scheduler_client,
)?;
Ok(handler)
}
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<()> {
let payload = bincode::serialize(&WorkerMetadata {
num_device_blocks: self.num_device_blocks,
bytes_per_block: self.bytes_per_block,
})?;
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>,
handler_cell: Arc<RwLock<Option<BlockTransferHandler>>>,
handler_tx: Arc<Mutex<Option<oneshot::Sender<BlockTransferHandler>>>>,
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(());
}
let leader_meta: LeaderMetadata = match bincode::deserialize(&message.data[0]) {
Ok(m) => m,
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() {
tracing::info!("Ping received but worker not ready; deferring ACK");
// 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 {
cell: Arc<RwLock<Option<BlockTransferHandler>>>,
}
#[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"))
}
}
}
#[derive(Builder, Clone)]
#[builder(pattern = "owned")]
pub struct KvbmWorkerConfig {
......@@ -110,11 +381,14 @@ pub struct KvbmWorkerConfig {
#[builder(default = "LayoutType::FullyContiguous")]
disk_layout_type: LayoutType,
#[builder(default = "String::from(\"kvbm\")")]
barrier_id_prefix: String,
#[builder(default = "None")]
scheduler_client: Option<TransferSchedulerClient>,
#[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,
}
impl KvbmWorkerConfig {
......@@ -123,18 +397,6 @@ impl KvbmWorkerConfig {
}
}
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)
}
pub struct KvbmWorker {
task: Option<CriticalTaskExecutionHandle>,
block_transfer_handler_rx: Option<oneshot::Receiver<transfer::BlockTransferHandler>>,
......@@ -265,26 +527,13 @@ impl KvbmWorker {
)> {
let cancel_token = config.drt.primary_token().clone();
// barrier sync with leader to get the leader data
let leader_data = tokio::task::block_in_place(|| {
// This is now synchronous blocking code
// We need a separate current-thread runtime to block_on async calls here
let rt = tokio::runtime::Handle::current();
rt.block_on(async {
KvbmWorker::leader_barrier_sync(
config.clone(),
cancel_token.clone(),
bytes_per_block,
)
.await
})
})?;
// establish a oneshot channel to get back the raw BlockTransferHandler
let (handler_tx, handler_rx) = oneshot::channel();
let handler_tx_cell = Arc::new(Mutex::new(Some(handler_tx)));
// 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>();
let layout_ready_tx_cell = Mutex::new(Some(layout_ready_tx));
let scheduler_client = config.scheduler_client.clone();
......@@ -295,13 +544,13 @@ impl KvbmWorker {
KvbmWorker::worker_task(
device_layout,
layout_builder,
leader_data,
layout_type,
worker_config,
cancel_token,
handler_tx,
layout_ready_tx,
handler_tx_cell,
layout_ready_tx_cell,
scheduler_client,
bytes_per_block,
)
},
cancel_token.clone(),
......@@ -314,18 +563,6 @@ impl KvbmWorker {
Err(_) => tracing::error!("Worker layout dropped without sending"),
}
let worker_config = config.clone();
let cancel_for_barrier = cancel_token.clone();
// wait until the leader finished the initialization of all components
tokio::task::block_in_place(|| {
// This is now synchronous blocking code
// We need a separate current-thread runtime to block_on async calls here
let rt = tokio::runtime::Handle::current();
rt.block_on(async {
KvbmWorker::leader_readiness_sync(worker_config, cancel_for_barrier).await
})
})?;
Ok((task, handler_rx))
}
......@@ -344,9 +581,11 @@ impl KvbmWorker {
// channel to get BlockTransferHandler back to the caller
let (handler_tx, handler_rx) = oneshot::channel::<transfer::BlockTransferHandler>();
let handler_tx_cell = Arc::new(Mutex::new(Some(handler_tx)));
// channel that the worker will use to signal layout readiness
let (layout_ready_tx, layout_ready_rx) = oneshot::channel::<String>();
let layout_ready_tx_cell = Mutex::new(Some(layout_ready_tx));
// clone what we need inside the orchestrator
let worker_config = config.clone();
......@@ -359,13 +598,7 @@ impl KvbmWorker {
let scheduler = scheduler_client.clone();
async move {
// 1) barrier (must finish before worker_task starts)
let leader_data =
KvbmWorker::leader_barrier_sync(cfg.clone(), ct.clone(), bytes_per_block)
.await?;
// 2) start the long-running worker (after barrier)
// Spawn it so the orchestrator can continue with readiness + waiting.
// Start the long-running worker.
let dev_layout = device_layout; // moved in
let lb = layout_builder; // moved in
let lt = layout_type; // moved in
......@@ -373,13 +606,13 @@ impl KvbmWorker {
let worker_fut = KvbmWorker::worker_task(
dev_layout,
lb,
leader_data,
lt,
cfg.clone(),
ct.clone(),
handler_tx,
layout_ready_tx,
handler_tx_cell,
layout_ready_tx_cell,
scheduler,
bytes_per_block,
);
// If worker_task returns Result, handle/log it inside the spawned task.
......@@ -395,9 +628,6 @@ impl KvbmWorker {
Err(_) => tracing::warn!("worker layout readiness channel dropped"),
}
// 4) wait for leader to finish its side of initialization
KvbmWorker::leader_readiness_sync(cfg.clone(), ct.clone()).await?;
Ok::<(), anyhow::Error>(())
}
},
......@@ -407,6 +637,7 @@ impl KvbmWorker {
Ok((task, handler_rx))
}
/// 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.
......@@ -433,134 +664,17 @@ impl KvbmWorker {
Ok(blocks)
}
async fn leader_barrier_sync(
config: KvbmWorkerConfig,
cancel_token: CancellationToken,
bytes_per_block: usize,
) -> anyhow::Result<KvbmLeaderData> {
let drt = config.drt.clone();
let worker_id = drt
.primary_lease()
.ok_or(anyhow::anyhow!(
"unable to get primary lease; check that drt is not static"
))?
.id() as usize;
let barrier_id_worker_to_leader =
format!("{}{}", config.barrier_id_prefix, "-worker-to-leader");
tracing::info!(
"Worker {} waiting on barrier {}",
worker_id,
barrier_id_worker_to_leader
);
let worker_to_leader_barrier = WorkerBarrier::<(), KvbmWorkerData>::new(
barrier_id_worker_to_leader,
worker_id.to_string(),
);
let worker_data = KvbmWorkerData {
num_device_blocks: config.num_device_blocks,
bytes_per_block,
};
tokio::select! {
_ = cancel_token.cancelled() => {
return Err(anyhow::anyhow!("Cancelled"))
}
_leader_data = worker_to_leader_barrier.sync(&drt, &worker_data) => {
_leader_data
}
}
.map_err(|e| anyhow::anyhow!("Failed to sync worker to leader barrier: {:?}", e))?;
tracing::debug!(
"Worker {} sent the worker data in worker to leader phase",
worker_id
);
let barrier_id_leader_to_worker =
format!("{}{}", config.barrier_id_prefix, "-leader-to-worker");
tracing::info!(
"Worker {} waiting on barrier {}",
worker_id,
barrier_id_leader_to_worker
);
let leader_to_worker_barrier = WorkerBarrier::<KvbmLeaderData, ()>::new(
barrier_id_leader_to_worker,
worker_id.to_string(),
);
let leader_data = tokio::select! {
_ = cancel_token.cancelled() => {
return Err(anyhow::anyhow!("Cancelled"))
}
leader_data = leader_to_worker_barrier.sync(&drt, &()) => {
leader_data
}
}
.map_err(|e| anyhow::anyhow!("Failed to sync worker to leader barrier: {:?}", e))?;
tracing::info!(
"Worker {} received leader data: {:?}",
worker_id,
leader_data
);
Ok(leader_data)
}
async fn leader_readiness_sync(
config: KvbmWorkerConfig,
cancel_token: CancellationToken,
) -> anyhow::Result<()> {
let drt = config.drt.clone();
let worker_id = drt
.primary_lease()
.ok_or(anyhow::anyhow!(
"unable to get primary lease; check that drt is not static"
))?
.id() as usize;
let barrier_id_leader_readiness =
format!("{}{}", config.barrier_id_prefix, "-leader-ready");
tracing::info!(
"Worker {} waiting on barrier {}",
worker_id,
barrier_id_leader_readiness
);
let leader_readiness_barrier =
WorkerBarrier::<(), ()>::new(barrier_id_leader_readiness, worker_id.to_string());
// leader_data is not important in the leader readiness case
tokio::select! {
_ = cancel_token.cancelled() => {
return Err(anyhow::anyhow!("Cancelled"))
}
_leader_data = leader_readiness_barrier.sync(&drt, &()) => {
_leader_data
}
}
.map_err(|e| anyhow::anyhow!("Failed to sync leader readiness barrier: {:?}", e))?;
Ok(())
}
#[allow(clippy::too_many_arguments)]
async fn worker_task(
device_layout: Box<dyn NixlLayout<StorageType = DeviceStorage>>,
mut layout_builder: LayoutConfigBuilder,
leader_data: KvbmLeaderData,
_layout_type: LayoutType,
layout_builder: LayoutConfigBuilder,
_device_layout_type: LayoutType,
config: KvbmWorkerConfig,
cancel_token: CancellationToken,
handler_tx: oneshot::Sender<BlockTransferHandler>,
layout_ready_tx: oneshot::Sender<String>,
handler_tx: Arc<Mutex<Option<oneshot::Sender<BlockTransferHandler>>>>,
layout_ready_tx: tokio::sync::Mutex<Option<oneshot::Sender<String>>>,
scheduler_client: Option<TransferSchedulerClient>,
bytes_per_block: usize,
) -> anyhow::Result<()> {
let drt = config.drt.clone();
......@@ -571,101 +685,63 @@ impl KvbmWorker {
))?
.id() as usize;
let agent = build_agent(worker_id, leader_data.num_disk_blocks > 0)?;
// Readiness gating for ping
let state = Arc::new(WorkerState::new());
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,
};
let transfer_context = Arc::new(TransferContext::new(
Arc::new(Some(agent)),
DeviceAllocator::new(config.device_id)
.unwrap()
.ctx()
.new_stream()
.unwrap(),
Handle::current(),
Some(pool_config),
));
// Build our device, host, and disk block lists.
let device_blocks = Some(Self::make_layout::<_, BasicMetadata>(
device_layout,
transfer_context.nixl_agent().as_ref(),
0,
worker_id,
)?);
// Cell to publish the transfer handler
let transfer_handler_cell: Arc<RwLock<Option<BlockTransferHandler>>> =
Arc::new(RwLock::new(None));
let host_blocks = if leader_data.num_host_blocks > 0 {
let host_allocator = Arc::new(PinnedAllocator::default());
let host_layout = layout_builder
.num_blocks(leader_data.num_host_blocks)
.build()?
.allocate_layout(config.host_layout_type, host_allocator)?;
// Build handlers map
let mut handlers: HashMap<String, Arc<dyn Handler>> = HashMap::new();
Some(Self::make_layout::<_, BasicMetadata>(
host_layout,
transfer_context.nixl_agent().as_ref(),
1,
worker_id,
)?)
} else {
None
};
handlers.insert(
ZMQ_PING_MESSAGE.to_string(),
Arc::new(GatedPing {
state: state.clone(),
layout_ready_tx,
}) as Arc<dyn Handler>,
);
let disk_blocks = if leader_data.num_disk_blocks > 0 {
let disk_allocator = Arc::new(DiskAllocator);
let disk_layout = layout_builder
.num_blocks(leader_data.num_disk_blocks)
.build()?
.allocate_layout(config.disk_layout_type, disk_allocator)?;
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>,
);
Some(Self::make_layout::<_, BasicMetadata>(
disk_layout,
transfer_context.nixl_agent().as_ref(),
2,
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,
)?)
} else {
None
};
let block_transfer_handler = BlockTransferHandler::new(
device_blocks,
host_blocks,
disk_blocks,
transfer_context,
device_id: config.device_id,
scheduler_client,
)?;
tracing::debug!("sending block transfer handler to worker");
handler_tx
.send(block_transfer_handler.clone())
.map_err(|_| {
anyhow::anyhow!("Failed to send block transfer handler over oneshot channel")
})?;
tracing::debug!("sent block transfer handler to worker");
handler_cell: transfer_handler_cell.clone(),
handler_tx, // sends BlockTransferHandler to caller
started: AtomicBool::new(false),
}) as Arc<dyn Handler>,
);
let handlers = HashMap::from([(
// transfer requests get dispatched to built handler (after allocation)
handlers.insert(
ZMQ_TRANSFER_BLOCKS_MESSAGE.to_string(),
Arc::new(block_transfer_handler) as Arc<dyn Handler>,
)]);
Arc::new(BlockTransferDispatch {
cell: transfer_handler_cell.clone(),
}) as Arc<dyn Handler>,
);
let _zmq_worker = ZmqActiveMessageWorker::new(
&leader_data.pub_url,
&leader_data.ack_url,
&config.leader_pub_url,
&config.leader_ack_url,
handlers,
cancel_token.clone(),
)?;
if layout_ready_tx.send("finished".to_string()).is_err() {
tracing::error!("worker receiver dropped before result was sent");
}
// TODO: Some sort of fancy loop here.
// For now, just wait for cancellation.
cancel_token.cancelled().await;
......
lib/llm/src/block_manager/distributed/zmq.rs
View file @ 94aa2a7b
......@@ -22,11 +22,17 @@ use tmq::{
use tokio::sync::{Mutex, oneshot};
use tokio_util::sync::CancellationToken;
use bincode;
use futures_util::{SinkExt, StreamExt};
use std::cmp::min;
struct PendingMessage {
remaining_workers: usize,
completion_indicator: oneshot::Sender<()>,
completion_indicator: Option<oneshot::Sender<()>>,
// If true, collect one payload (bytes) from each worker reply.
want_payload: bool,
// Collected raw payloads (one per worker), if want_payload == true
payloads: Option<Vec<Vec<u8>>>,
}
pub struct LeaderSockets {
......@@ -36,18 +42,16 @@ pub struct LeaderSockets {
pub ack_url: String,
}
pub fn new_leader_sockets(url: &str) -> Result<LeaderSockets> {
let url = format!("{}:0", url);
pub fn new_leader_sockets(pub_url: &str, ack_url: &str) -> Result<LeaderSockets> {
let context = Context::new();
let pub_socket = publish(&context).bind(url.as_str())?;
let pub_socket = publish(&context).bind(pub_url)?;
let pub_url = pub_socket
.get_socket()
.get_last_endpoint()
.unwrap()
.unwrap();
let ack_socket = pull(&context).bind(url.as_str())?;
let ack_socket = pull(&context).bind(ack_url)?;
let ack_url = ack_socket
.get_socket()
.get_last_endpoint()
......@@ -78,12 +82,18 @@ pub struct ZmqActiveMessageLeader {
}
impl ZmqActiveMessageLeader {
pub async fn new(
/// Handshake-first constructor: collects WorkerMetaData, broadcasts LeaderMetadata,
/// waits for allocation ACKs, then runs the final ping loop.
pub async fn new_with_handshake<F>(
leader_sockets: LeaderSockets,
num_workers: usize,
timeout: Duration,
overall_timeout: Duration,
cancel_token: CancellationToken,
) -> Result<Self> {
make_leader_meta: F,
) -> Result<Self>
where
F: Fn(&[WorkerMetadata]) -> LeaderMetadata + Send + Sync + 'static,
{
let pub_socket = Arc::new(Mutex::new(leader_sockets.pub_socket));
let pull_socket = leader_sockets.ack_socket;
......@@ -94,48 +104,128 @@ impl ZmqActiveMessageLeader {
);
let pending_messages = Arc::new(Mutex::new(HashMap::new()));
let pending_messages_clone = pending_messages.clone();
CriticalTaskExecutionHandle::new(
|cancel_token| Self::pull_worker(pull_socket, pending_messages_clone, cancel_token),
cancel_token,
|ct| Self::pull_worker(pull_socket, pending_messages_clone, ct),
cancel_token.clone(),
"ZmqActiveMessageLeader: Pull worker",
)?
.detach();
let self_ = Self {
let this = Self {
pub_socket,
message_id: Arc::new(Mutex::new(0)),
pending_messages,
num_workers: Arc::new(num_workers),
};
// Ping our workers.
let start = Instant::now();
loop {
if start.elapsed() > timeout {
return Err(anyhow::anyhow!("Timed out waiting for workers."));
let deadline = Instant::now() + overall_timeout;
// 1) Collect KvbmWorkerData from ALL workers in a single round.
// Keep rebroadcasting until we get exactly `num_workers` replies to the SAME broadcast.
let workers_payloads: Vec<Vec<u8>> = loop {
if Instant::now() >= deadline {
return Err(anyhow::anyhow!(
"Handshake timed out (device-config collection)."
));
}
let remain = deadline.saturating_duration_since(Instant::now());
let round_to = min(Duration::from_secs(2), remain);
tracing::info!("Handshake: requesting worker device configs...");
match this
.broadcast_collect(
ZMQ_WORKER_METADATA_MESSAGE,
&[],
/* want_payload */ true,
round_to,
)
.await
{
Ok(payloads) if payloads.len() == num_workers => {
tracing::info!(
"Handshake: received {} worker metadata replies in this round.",
payloads.len()
);
break payloads;
}
Ok(payloads) => {
tracing::warn!(
"Handshake: got {} / {} worker metadata replies; rebroadcasting...",
payloads.len(),
num_workers
);
continue;
}
Err(e) => {
tracing::debug!(
"Handshake: worker metadata round timed out/failed: {e}; retrying..."
);
continue;
}
}
};
// Try to send a ping to all workers.
tracing::info!("ZmqActiveMessageLeader: Pinging workers...");
let ping_receiver = self_.broadcast(ZMQ_PING_MESSAGE, vec![]).await?;
let workers: Vec<WorkerMetadata> = workers_payloads
.into_iter()
.map(|b| bincode::deserialize::<WorkerMetadata>(&b))
.collect::<std::result::Result<_, _>>()?;
tokio::select! {
// If we receive an ACK from every worker, we're done.
_ = ping_receiver => {
tracing::info!("ZmqActiveMessageLeader: Worker ping successful. Startup complete.");
// 2) Compute & broadcast LeaderMetadata; wait for ALL acks in the SAME round.
let leader_meta = make_leader_meta(&workers);
let leader_meta_bytes = bincode::serialize(&leader_meta)?;
loop {
if Instant::now() >= deadline {
return Err(anyhow::anyhow!(
"Handshake timed out (allocation-config broadcast)."
));
}
let remain = deadline.saturating_duration_since(Instant::now());
let round_to = min(Duration::from_secs(2), remain);
tracing::info!("Handshake: broadcasting allocation config to workers...");
match this
.broadcast_collect(
ZMQ_LEADER_METADATA_MESSAGE,
std::slice::from_ref(&leader_meta_bytes),
/* want_payload */ false,
round_to,
)
.await
{
Ok(_) => {
// Success: all workers acked in this round.
tracing::info!("Handshake: all workers acked allocation config.");
break;
}
// Wait for 1 second before pinging again.
_ = tokio::time::sleep(Duration::from_millis(1000)) => {
tracing::info!("ZmqActiveMessageLeader: Ping timed out. Retrying...");
Err(e) => {
tracing::warn!(
"Handshake: allocation-config round incomplete: {e}; rebroadcasting..."
);
continue;
}
}
}
Ok(self_)
// 3) Final readiness ping loop (workers only ACK after allocation ready)
let ping_deadline = deadline;
loop {
if Instant::now() >= ping_deadline {
return Err(anyhow::anyhow!(
"Timed out waiting for ping readiness after handshake."
));
}
tracing::info!("Handshake: final readiness ping...");
let ping = this.broadcast(ZMQ_PING_MESSAGE, vec![]).await?;
tokio::select! {
_ = ping => break,
_ = tokio::time::sleep(Duration::from_millis(500)) => continue,
_ = cancel_token.cancelled() => return Err(anyhow::anyhow!("Startup canceled")),
}
}
Ok(this)
}
/// Broadcast a message to all workers.
......@@ -157,7 +247,9 @@ impl ZmqActiveMessageLeader {
let pending_message = PendingMessage {
// We start with the number of workers we're waiting for.
remaining_workers: *self.num_workers,
completion_indicator,
completion_indicator: Some(completion_indicator),
want_payload: false,
payloads: None,
};
// Add the message to the pending messages map.
......@@ -187,7 +279,68 @@ impl ZmqActiveMessageLeader {
Ok(completion_receiver)
}
/// Pull worker is responsible for receiving ACKs from workers.
/// Generic broadcast that can collect one reply payload from each worker.
/// - `function`: handler name on workers
/// - `data_frames`: optional extra frames after [id, function]
/// - `want_payload`: if true, expects replies shaped as [id, function, payload]
/// Returns payloads (empty if want_payload == false).
pub async fn broadcast_collect(
&self,
function: &str,
data_frames: &[Vec<u8>],
want_payload: bool,
timeout: Duration,
) -> Result<Vec<Vec<u8>>> {
// Generate a unique id.
let id = {
let mut id = self.message_id.lock().await;
*id += 1;
*id
};
let (completion_indicator, completion_receiver) = oneshot::channel();
let pending_message = PendingMessage {
remaining_workers: *self.num_workers,
completion_indicator: Some(completion_indicator),
want_payload,
payloads: want_payload.then(|| Vec::with_capacity(*self.num_workers)),
};
self.pending_messages
.lock()
.await
.insert(id, pending_message);
// Build message: [id, function, ...data]
let mut message: VecDeque<Message> = VecDeque::with_capacity(2 + data_frames.len());
message.push_back(id.to_be_bytes().as_slice().into());
message.push_back(function.into());
for df in data_frames {
message.push_back(df.clone().into());
}
self.pub_socket
.lock()
.await
.send(Multipart(message))
.await?;
// Await all replies or timeout.
tokio::select! {
_ = completion_receiver => { /* done */ }
_ = tokio::time::sleep(timeout) => {
let mut map = self.pending_messages.lock().await;
map.remove(&id);
return Err(anyhow::anyhow!("Timed out waiting for '{}' responses", function));
}
}
// Extract payloads (if any).
let mut map = self.pending_messages.lock().await;
let entry = map
.remove(&id)
.ok_or_else(|| anyhow::anyhow!("pending entry missing"))?;
Ok(entry.payloads.unwrap_or_default())
}
async fn pull_worker(
mut pull_socket: Pull,
pending_messages: Arc<Mutex<HashMap<usize, PendingMessage>>>,
......@@ -196,63 +349,40 @@ impl ZmqActiveMessageLeader {
loop {
tokio::select! {
Some(Ok(message)) = pull_socket.next() => {
// The leader should only ever receive ACKs.
// ACKs have no data.
if message.len() != 1 {
tracing::error!(
"Received message with unexpected length: {:?}",
message.len()
);
if message.is_empty() {
tracing::error!("Leader PULL: empty message");
continue;
}
// TODO: This looks ugly.
let arr: [u8; std::mem::size_of::<usize>()] = (*message[0]).try_into()?;
let id = usize::from_be_bytes(arr);
let mut pending_messages = pending_messages.lock().await;
// TODO: Should we error if we can't find the pending message?
// if let std::collections::hash_map::Entry::Occupied(mut entry) =
// pending_messages.entry(id)
// {
// entry.get_mut().remaining_workers -= 1;
// tracing::debug!(
// "ZmqActiveMessageLeader: Received ACK for message with id: {}. There are {} remaining workers.",
// id,
// entry.get().remaining_workers
// );
// // If all workers have ACKed, notify the completion indicator.
// if entry.get().remaining_workers == 0 {
// let e = entry.remove();
// tracing::debug!(
// "ZmqActiveMessageLeader: Message with id: {} completed.",
// id
// );
// // It's possible that the receiver has already been dropped,
// // so ignore any send error here.
// let _ = e.completion_indicator.send(());
// }
// }
match pending_messages.entry(id) {
std::collections::hash_map::Entry::Occupied(mut entry) => {
let pending_message = entry.get_mut();
debug_assert!(pending_message.remaining_workers > 0);
pending_message.remaining_workers -= 1;
tracing::debug!(
"ZmqActiveMessageLeader: Received ACK for message with id: {}. There are {} remaining workers.",
id,
pending_message.remaining_workers
);
if pending_message.remaining_workers == 0 {
let e = entry.remove();
tracing::debug!("ZmqActiveMessageLeader: Message with id: {} completed.", id);
let _ = e.completion_indicator.send(());
let mut map = pending_messages.lock().await;
if let Some(pm) = map.get_mut(&id) {
// payload reply or pure ACK?
if message.len() == 1 {
if pm.remaining_workers > 0 { pm.remaining_workers -= 1; }
} else {
if pm.want_payload && message.len() >= 3
&& let Some(bufs) = pm.payloads.as_mut() {
bufs.push((*message[2]).to_vec());
}
if pm.remaining_workers > 0 { pm.remaining_workers -= 1; }
}
std::collections::hash_map::Entry::Vacant(_) => {
tracing::error!("Received ACK for unknown message with id: {}", id);
tracing::debug!(
"Leader PULL: got {} for id {} (remaining={})",
if message.len()==1 { "ACK" } else { "REPLY" }, id, pm.remaining_workers
);
// IMPORTANT: do NOT remove here; just notify completion.
if pm.remaining_workers == 0
&& let Some(tx) = pm.completion_indicator.take() {
let _ = tx.send(());
}
} else {
// Late reply for a round we've already collected/removed.
tracing::debug!("Leader PULL: late/unknown id {}", id);
}
}
_ = cancel_token.cancelled() => {
......@@ -269,10 +399,10 @@ impl ZmqActiveMessageLeader {
/// A message handle is used to track a message.
/// It contains a way to ACK the message, as well as the data.
pub struct MessageHandle {
message_id: usize,
pub message_id: usize,
function: String,
pub data: Vec<Vec<u8>>,
push_handle: Arc<Mutex<Push>>,
pub push_handle: Arc<Mutex<Push>>,
acked: bool,
}
......@@ -321,6 +451,36 @@ impl MessageHandle {
tracing::debug!("ZmqActiveMessageWorker: ACKed message with id: {}", id);
Ok(())
}
/// Reply to the leader with arbitrary payload frames and mark as acked.
/// Frames shape: [id, function, payload_0, payload_1, ...]
pub async fn reply(
&mut self,
function: &str,
payload_frames: &[Vec<u8>],
) -> anyhow::Result<()> {
let mut frames: std::collections::VecDeque<tmq::Message> =
std::collections::VecDeque::with_capacity(2 + payload_frames.len());
frames.push_back(self.message_id.to_be_bytes().as_slice().into());
frames.push_back(function.into());
for p in payload_frames {
frames.push_back(p.clone().into());
}
self.push_handle
.lock()
.await
.send(tmq::Multipart(frames))
.await?;
// Mark as acked so Drop won't panic; leader treats the reply as the "ack".
self.acked = true;
Ok(())
}
/// Mark this message as handled locally without sending an ACK/reply.
/// Use when intentionally ignoring a message (e.g. ping before readiness).
pub fn mark_handled(&mut self) {
self.acked = true;
}
}
/// We must always ACK a message.
......@@ -340,21 +500,6 @@ pub trait Handler: Send + Sync {
async fn handle(&self, message: MessageHandle) -> Result<()>;
}
/// A super simple handler that responds to a ping.
/// This is used in the startup sequence to check worker liveness.
struct Ping;
#[async_trait]
impl Handler for Ping {
async fn handle(&self, mut message: MessageHandle) -> Result<()> {
if !message.data.is_empty() {
return Err(anyhow::anyhow!("Ping message should not have data."));
}
message.ack().await?;
Ok(())
}
}
type MessageHandlers = HashMap<String, Arc<dyn Handler>>;
/// The ActiveMessageWorker receives commands from the leader, and ACKs them.
......@@ -364,7 +509,7 @@ impl ZmqActiveMessageWorker {
pub fn new(
sub_url: &str,
push_url: &str,
mut message_handlers: MessageHandlers,
message_handlers: MessageHandlers,
cancel_token: CancellationToken,
) -> Result<Self> {
let context = Context::new();
......@@ -380,8 +525,6 @@ impl ZmqActiveMessageWorker {
push_url
);
// Add our ping handler.
message_handlers.insert(ZMQ_PING_MESSAGE.to_string(), Arc::new(Ping));
let message_handlers = Arc::new(message_handlers);
CriticalTaskExecutionHandle::new(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment