dynamo_mocker/
engine.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Single-rank compatibility entry points for the grouped generalized engine.
5
6use anyhow::{Context, ensure};
7use tokio::sync::mpsc;
8use tokio::task::JoinHandle;
9use tokio_util::sync::CancellationToken;
10
11use crate::common::protocols::{FpmPublisher, KvEventPublishers, MockEngineArgs, OutputSignal};
12use crate::grouped_scheduler::{
13    GroupedSchedulerRankEventSinks, GroupedSchedulers,
14    create_single_rank_scheduler_with_event_sender,
15};
16use crate::scheduler::{SchedulerEventSender, SchedulerHandle};
17
18pub(crate) struct LiveEngineScheduler {
19    pub(crate) handle: Box<dyn SchedulerHandle>,
20    pub(crate) actor: JoinHandle<anyhow::Result<()>>,
21    pub(crate) completion_drain: crate::grouped_scheduler::CompletionBoundaryDrain,
22}
23
24/// Create a scheduler for the configured engine type.
25///
26/// Returns a boxed [`SchedulerHandle`] that the engine wrapper can use
27/// without knowing which backend is running underneath.
28pub fn create_engine(
29    args: MockEngineArgs,
30    dp_rank: u32,
31    output_tx: Option<mpsc::UnboundedSender<Vec<OutputSignal>>>,
32    kv_event_publishers: KvEventPublishers,
33    cancellation_token: Option<CancellationToken>,
34    fpm_publisher: FpmPublisher,
35) -> anyhow::Result<Box<dyn SchedulerHandle>> {
36    let LiveEngineScheduler {
37        handle,
38        actor,
39        completion_drain: _,
40    } = create_engine_with_event_sender(
41        args,
42        dp_rank,
43        output_tx.map(SchedulerEventSender::from),
44        kv_event_publishers,
45        cancellation_token,
46        fpm_publisher,
47    )?;
48    // Dropping a Tokio JoinHandle detaches the actor. The SchedulerHandle's
49    // cancellation guard remains the compatibility API's shutdown owner.
50    drop(actor);
51    Ok(handle)
52}
53
54pub(crate) fn create_engine_with_event_sender(
55    args: MockEngineArgs,
56    dp_rank: u32,
57    event_tx: Option<SchedulerEventSender>,
58    kv_event_publishers: KvEventPublishers,
59    cancellation_token: Option<CancellationToken>,
60    fpm_publisher: FpmPublisher,
61) -> anyhow::Result<LiveEngineScheduler> {
62    // This compatibility API cannot safely construct attention-DP ranks one at
63    // a time: each call would own a different generalized engine and bypass the
64    // group barrier. Production Live Mocker constructs the complete group once
65    // through `create_grouped_scheduler`.
66    ensure!(
67        args.dp_size == 1,
68        "single-rank create_engine does not support attention DP; use create_grouped_scheduler"
69    );
70    ensure!(
71        dp_rank == 0,
72        "single-rank create_engine requires dp_rank=0; use create_grouped_scheduler"
73    );
74    let GroupedSchedulers {
75        mut schedulers,
76        actor,
77        completion_drain,
78    } = create_single_rank_scheduler_with_event_sender(
79        args,
80        dp_rank,
81        GroupedSchedulerRankEventSinks {
82            event_tx,
83            kv_event_publishers,
84            fpm_publisher,
85        },
86        cancellation_token,
87    )?;
88    let handle = schedulers
89        .pop()
90        .context("single-rank generalized Mocker engine returned no scheduler handle")?;
91    ensure!(
92        schedulers.is_empty(),
93        "single-rank generalized Mocker engine returned extra scheduler handles"
94    );
95    Ok(LiveEngineScheduler {
96        handle,
97        actor,
98        completion_drain,
99    })
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105
106    #[tokio::test]
107    async fn compatibility_entrypoint_rejects_attention_dp() {
108        let args = MockEngineArgs {
109            dp_size: 4,
110            ..MockEngineArgs::default()
111        };
112        let cancel = CancellationToken::new();
113
114        let result = create_engine_with_event_sender(
115            args,
116            3,
117            None,
118            KvEventPublishers::default(),
119            Some(cancel.clone()),
120            FpmPublisher::default(),
121        );
122        let error = match result {
123            Ok(_) => panic!("attention DP must be rejected by the single-rank entrypoint"),
124            Err(error) => error,
125        };
126        assert!(error.to_string().contains("does not support attention DP"));
127    }
128}
dynamo_mocker/engine.rs

dynamo_mocker/
engine.rs
