Skip to main content

kvbm_engine/worker/group/
mod.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Distributed Workers
5//!
6//! This module provides the interface for how the leader will drive multiple workers.
7
8// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
9// SPDX-License-Identifier: Apache-2.0
10
11mod spmd;
12
13use std::sync::Arc;
14
15use super::{
16    ImportMetadataResponse, SerializedLayout, SerializedLayoutResponse, Worker, WorkerTransfers, *,
17};
18use crate::object::ObjectBlockOps;
19use anyhow::Result;
20
21pub use spmd::SpmdParallelWorkers;
22
23/// A cohort of parallel workers.
24///
25/// This trait is used to drive one or more parallel workers.
26pub trait ParallelWorkers: WorkerTransfers + ObjectBlockOps + Send + Sync {
27    /// Export the local metadata for a set of workers.
28    ///
29    /// Layouts will be returned in rank order.
30    ///
31    /// # Returns
32    /// A [`kvbm_physical::manager::SerializedLayout`] containing the local metadata
33    fn export_metadata(&self) -> Result<Vec<SerializedLayoutResponse>>;
34
35    /// Import the remote metadata for this worker.
36    ///
37    /// Handles will be returned in rank order.
38    ///
39    /// # Arguments
40    /// * `metadata` - A [`kvbm_physical::manager::SerializedLayout`] containing the remote metadata
41    ///
42    /// # Returns
43    /// A vector of [`kvbm_physical::manager::LayoutHandle`] for the imported remote layouts
44    fn import_metadata(
45        &self,
46        metadata: Vec<SerializedLayout>,
47    ) -> Result<Vec<ImportMetadataResponse>>;
48
49    /// Get the number of workers.
50    fn worker_count(&self) -> usize;
51
52    /// Get access to the underlying workers for metadata/handle queries.
53    ///
54    /// This is useful for operations that need to query individual workers
55    /// (e.g., collecting layout handles) without executing transfers.
56    fn workers(&self) -> &[Arc<dyn Worker>];
57}