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//! Engine factory — creates the appropriate scheduler based on [`EngineType`].
5
6use tokio::sync::mpsc;
7use tokio::task::JoinHandle;
8use tokio_util::sync::CancellationToken;
9
10use crate::common::protocols::{
11    EngineType, FpmPublisher, KvEventPublishers, MockEngineArgs, OutputSignal,
12};
13use crate::scheduler::{
14    Scheduler, SchedulerEventSender, SchedulerHandle, SchedulerOutputSender, SglangScheduler,
15};
16
17pub(crate) struct LiveEngineScheduler {
18    pub(crate) handle: Box<dyn SchedulerHandle>,
19    pub(crate) actor: JoinHandle<anyhow::Result<()>>,
20}
21
22/// Create a scheduler for the configured engine type.
23///
24/// Returns a boxed [`SchedulerHandle`] that the engine wrapper can use
25/// without knowing which backend is running underneath.
26pub fn create_engine(
27    args: MockEngineArgs,
28    dp_rank: u32,
29    output_tx: Option<mpsc::UnboundedSender<Vec<OutputSignal>>>,
30    kv_event_publishers: KvEventPublishers,
31    cancellation_token: Option<CancellationToken>,
32    fpm_publisher: FpmPublisher,
33) -> Box<dyn SchedulerHandle> {
34    create_engine_with_output_sender(
35        args,
36        dp_rank,
37        output_tx.map(SchedulerOutputSender::from),
38        kv_event_publishers,
39        cancellation_token,
40        fpm_publisher,
41    )
42}
43
44pub(crate) fn create_engine_with_output_sender(
45    args: MockEngineArgs,
46    dp_rank: u32,
47    output_tx: Option<SchedulerOutputSender>,
48    kv_event_publishers: KvEventPublishers,
49    cancellation_token: Option<CancellationToken>,
50    fpm_publisher: FpmPublisher,
51) -> Box<dyn SchedulerHandle> {
52    let LiveEngineScheduler { handle, actor } = create_engine_with_event_sender(
53        args,
54        dp_rank,
55        output_tx.map(SchedulerEventSender::from),
56        kv_event_publishers,
57        cancellation_token,
58        fpm_publisher,
59    );
60    drop(actor);
61    handle
62}
63
64pub(crate) fn create_engine_with_event_sender(
65    args: MockEngineArgs,
66    dp_rank: u32,
67    event_tx: Option<SchedulerEventSender>,
68    kv_event_publishers: KvEventPublishers,
69    cancellation_token: Option<CancellationToken>,
70    fpm_publisher: FpmPublisher,
71) -> LiveEngineScheduler {
72    let (handle, actor): (Box<dyn SchedulerHandle>, _) = match args.engine_type {
73        // TRT-LLM reuses the vLLM scheduler core; the GUARANTEED_NO_EVICT
74        // policy is carried in `args` and read by the core per pass.
75        EngineType::Vllm | EngineType::Trtllm => {
76            let (scheduler, actor) = Scheduler::spawn_with_event_sender(
77                args,
78                dp_rank,
79                event_tx,
80                kv_event_publishers,
81                cancellation_token,
82                fpm_publisher,
83            );
84            (Box::new(scheduler), actor)
85        }
86        EngineType::Sglang => {
87            let (scheduler, actor) = SglangScheduler::spawn_with_event_sender(
88                args,
89                dp_rank,
90                event_tx,
91                kv_event_publishers,
92                cancellation_token,
93                fpm_publisher,
94            );
95            (Box::new(scheduler), actor)
96        }
97    };
98    LiveEngineScheduler { handle, actor }
99}
dynamo_mocker/engine.rs

dynamo_mocker/
engine.rs
