entrypoint.rs 16.2 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
3
4
// SPDX-License-Identifier: Apache-2.0

use std::fmt::Display;
5
use std::future::Future;
6
use std::path::PathBuf;
7
8
use std::pin::Pin;
use std::sync::Arc;
9
10

use pyo3::{exceptions::PyException, prelude::*};
11
use pyo3_async_runtimes::TaskLocals;
12

13
use dynamo_llm::discovery::LoadThresholdConfig as RsLoadThresholdConfig;
14
use dynamo_llm::entrypoint::EngineConfig as RsEngineConfig;
15
use dynamo_llm::entrypoint::EngineFactoryCallback;
16
use dynamo_llm::entrypoint::RouterConfig as RsRouterConfig;
17
use dynamo_llm::entrypoint::input::Input;
18
use dynamo_llm::kv_router::KvRouterConfig as RsKvRouterConfig;
Graham King's avatar
Graham King committed
19
use dynamo_llm::local_model::DEFAULT_HTTP_PORT;
20
use dynamo_llm::local_model::{LocalModel, LocalModelBuilder};
21
use dynamo_llm::mocker::protocols::MockEngineArgs;
22
23
use dynamo_llm::model_card::ModelDeploymentCard as RsModelDeploymentCard;
use dynamo_llm::types::openai::chat_completions::OpenAIChatCompletionsStreamingEngine;
24
use dynamo_runtime::protocols::EndpointId;
25

26
use super::model_card::ModelDeploymentCard;
27
use crate::RouterMode;
28
use crate::engine::PythonAsyncEngine;
29

30
31
32
33
34
#[pyclass(eq, eq_int)]
#[derive(Clone, Debug, PartialEq)]
#[repr(i32)]
pub enum EngineType {
    Echo = 1,
35
36
    Dynamic = 2,
    Mocker = 3,
37
38
}

39
40
41
42
43
44
#[pyclass]
#[derive(Default, Clone, Debug, Copy)]
pub struct KvRouterConfig {
    inner: RsKvRouterConfig,
}

45
46
47
48
49
50
impl KvRouterConfig {
    pub fn inner(&self) -> RsKvRouterConfig {
        self.inner
    }
}

51
52
53
#[pymethods]
impl KvRouterConfig {
    #[new]
54
    #[pyo3(signature = (overlap_score_weight=1.0, router_temperature=0.0, use_kv_events=true, router_replica_sync=false, router_track_active_blocks=true, router_track_output_blocks=false, router_assume_kv_reuse=true, router_snapshot_threshold=1000000, router_reset_states=false, router_ttl_secs=120.0, router_max_tree_size=1048576, router_prune_target_ratio=0.8))]
55
    #[allow(clippy::too_many_arguments)]
56
57
58
59
60
    fn new(
        overlap_score_weight: f64,
        router_temperature: f64,
        use_kv_events: bool,
        router_replica_sync: bool,
61
        router_track_active_blocks: bool,
62
        router_track_output_blocks: bool,
63
        router_assume_kv_reuse: bool,
64
65
        router_snapshot_threshold: Option<u32>,
        router_reset_states: bool,
66
67
68
        router_ttl_secs: f64,
        router_max_tree_size: usize,
        router_prune_target_ratio: f64,
69
    ) -> Self {
70
71
72
73
74
        KvRouterConfig {
            inner: RsKvRouterConfig {
                overlap_score_weight,
                router_temperature,
                use_kv_events,
75
                router_replica_sync,
76
                router_track_active_blocks,
77
                router_track_output_blocks,
78
                router_assume_kv_reuse,
79
80
                router_snapshot_threshold,
                router_reset_states,
81
82
83
                router_ttl_secs,
                router_max_tree_size,
                router_prune_target_ratio,
84
85
86
87
88
89
90
91
92
93
            },
        }
    }
}

#[pyclass]
#[derive(Clone, Debug)]
pub struct RouterConfig {
    router_mode: RouterMode,
    kv_router_config: KvRouterConfig,
94
95
96
97
    /// Threshold for active decode blocks utilization (0.0-1.0)
    active_decode_blocks_threshold: Option<f64>,
    /// Threshold for active prefill tokens utilization (literal token count)
    active_prefill_tokens_threshold: Option<u64>,
98
99
    /// Threshold for active prefill tokens as fraction of max_num_batched_tokens
    active_prefill_tokens_threshold_frac: Option<f64>,
100
    enforce_disagg: bool,
101
102
103
104
105
}

#[pymethods]
impl RouterConfig {
    #[new]
106
    #[pyo3(signature = (mode, config=None, active_decode_blocks_threshold=None, active_prefill_tokens_threshold=None, active_prefill_tokens_threshold_frac=None, enforce_disagg=false))]
107
108
109
    pub fn new(
        mode: RouterMode,
        config: Option<KvRouterConfig>,
110
111
        active_decode_blocks_threshold: Option<f64>,
        active_prefill_tokens_threshold: Option<u64>,
112
        active_prefill_tokens_threshold_frac: Option<f64>,
113
        enforce_disagg: bool,
114
    ) -> Self {
115
116
117
        Self {
            router_mode: mode,
            kv_router_config: config.unwrap_or_default(),
118
119
            active_decode_blocks_threshold,
            active_prefill_tokens_threshold,
120
            active_prefill_tokens_threshold_frac,
121
            enforce_disagg,
122
123
124
125
126
127
128
129
130
        }
    }
}

impl From<RouterConfig> for RsRouterConfig {
    fn from(rc: RouterConfig) -> RsRouterConfig {
        RsRouterConfig {
            router_mode: rc.router_mode.into(),
            kv_router_config: rc.kv_router_config.inner,
131
132
133
134
135
            load_threshold_config: RsLoadThresholdConfig {
                active_decode_blocks_threshold: rc.active_decode_blocks_threshold,
                active_prefill_tokens_threshold: rc.active_prefill_tokens_threshold,
                active_prefill_tokens_threshold_frac: rc.active_prefill_tokens_threshold_frac,
            },
136
            enforce_disagg: rc.enforce_disagg,
137
138
139
140
        }
    }
}

141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
/// Wrapper to hold Python callback and its TaskLocals for async execution
#[derive(Clone)]
struct PyEngineFactory {
    callback: Arc<PyObject>,
    locals: Arc<TaskLocals>,
}

impl std::fmt::Debug for PyEngineFactory {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PyEngineFactory")
            .field("callback", &"<PyObject>")
            .finish()
    }
}

156
157
158
159
160
161
162
163
164
#[pyclass]
#[derive(Clone, Debug)]
pub(crate) struct EntrypointArgs {
    engine_type: EngineType,
    model_path: Option<PathBuf>,
    model_name: Option<String>,
    endpoint_id: Option<EndpointId>,
    context_length: Option<u32>,
    template_file: Option<PathBuf>,
165
    router_config: Option<RouterConfig>,
166
    kv_cache_block_size: Option<u32>,
167
    http_host: Option<String>,
Graham King's avatar
Graham King committed
168
    http_port: u16,
169
    http_metrics_port: Option<u16>,
Graham King's avatar
Graham King committed
170
171
    tls_cert_path: Option<PathBuf>,
    tls_key_path: Option<PathBuf>,
172
    extra_engine_args: Option<PathBuf>,
173
    namespace: Option<String>,
174
175
    custom_backend_metrics_endpoint: Option<String>,
    custom_backend_metrics_polling_interval: Option<f64>,
176
    is_prefill: bool,
177
    engine_factory: Option<PyEngineFactory>,
178
179
180
181
182
183
}

#[pymethods]
impl EntrypointArgs {
    #[allow(clippy::too_many_arguments)]
    #[new]
184
    #[pyo3(signature = (engine_type, model_path=None, model_name=None, endpoint_id=None, context_length=None, template_file=None, router_config=None, kv_cache_block_size=None, http_host=None, http_port=None, http_metrics_port=None, tls_cert_path=None, tls_key_path=None, extra_engine_args=None, namespace=None, custom_backend_metrics_endpoint=None, custom_backend_metrics_polling_interval=None, is_prefill=false, engine_factory=None))]
185
    pub fn new(
186
        py: Python<'_>,
187
188
189
190
191
192
        engine_type: EngineType,
        model_path: Option<PathBuf>,
        model_name: Option<String>, // e.g. "dyn://namespace.component.endpoint"
        endpoint_id: Option<String>,
        context_length: Option<u32>,
        template_file: Option<PathBuf>,
193
        router_config: Option<RouterConfig>,
194
        kv_cache_block_size: Option<u32>,
195
        http_host: Option<String>,
196
        http_port: Option<u16>,
197
        http_metrics_port: Option<u16>,
Graham King's avatar
Graham King committed
198
199
        tls_cert_path: Option<PathBuf>,
        tls_key_path: Option<PathBuf>,
200
        extra_engine_args: Option<PathBuf>,
201
        namespace: Option<String>,
202
203
        custom_backend_metrics_endpoint: Option<String>,
        custom_backend_metrics_polling_interval: Option<f64>,
204
        is_prefill: bool,
205
        engine_factory: Option<PyObject>,
206
    ) -> PyResult<Self> {
207
        let endpoint_id_obj: Option<EndpointId> = endpoint_id.as_deref().map(EndpointId::from);
Graham King's avatar
Graham King committed
208
209
210
211
212
213
214
        if (tls_cert_path.is_some() && tls_key_path.is_none())
            || (tls_cert_path.is_none() && tls_key_path.is_some())
        {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "tls_cert_path and tls_key_path must be provided together",
            ));
        }
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231

        // Capture TaskLocals at registration time for the engine factory callback
        let engine_factory = engine_factory
            .map(|callback| {
                let locals = pyo3_async_runtimes::tokio::get_current_locals(py).map_err(|e| {
                    pyo3::exceptions::PyRuntimeError::new_err(format!(
                        "Failed to get TaskLocals for engine_factory: {}",
                        e
                    ))
                })?;
                Ok::<_, PyErr>(PyEngineFactory {
                    callback: Arc::new(callback),
                    locals: Arc::new(locals),
                })
            })
            .transpose()?;

232
233
234
235
236
237
238
        Ok(EntrypointArgs {
            engine_type,
            model_path,
            model_name,
            endpoint_id: endpoint_id_obj,
            context_length,
            template_file,
239
            router_config,
240
            kv_cache_block_size,
241
            http_host,
Graham King's avatar
Graham King committed
242
            http_port: http_port.unwrap_or(DEFAULT_HTTP_PORT),
243
            http_metrics_port,
Graham King's avatar
Graham King committed
244
245
            tls_cert_path,
            tls_key_path,
246
            extra_engine_args,
247
            namespace,
248
249
            custom_backend_metrics_endpoint,
            custom_backend_metrics_polling_interval,
250
            is_prefill,
251
            engine_factory,
252
253
254
255
256
257
258
259
260
261
        })
    }
}

#[pyclass]
#[derive(Clone)]
pub(crate) struct EngineConfig {
    inner: RsEngineConfig,
}

262
263
/// Create the backend engine wrapper to run the model.
/// Download the model if necessary.
264
265
266
267
268
269
270
271
272
#[pyfunction]
#[pyo3(signature = (distributed_runtime, args))]
pub fn make_engine<'p>(
    py: Python<'p>,
    distributed_runtime: super::DistributedRuntime,
    args: EntrypointArgs,
) -> PyResult<Bound<'p, PyAny>> {
    let mut builder = LocalModelBuilder::default();
    builder
273
274
275
276
277
        .model_name(
            args.model_name
                .clone()
                .or_else(|| args.model_path.clone().map(|p| p.display().to_string())),
        )
278
        .endpoint_id(args.endpoint_id.clone())
279
        .context_length(args.context_length)
280
        .request_template(args.template_file.clone())
281
        .kv_cache_block_size(args.kv_cache_block_size)
282
        .router_config(args.router_config.clone().map(|rc| rc.into()))
283
        .http_host(args.http_host.clone())
284
        .http_port(args.http_port)
285
        .http_metrics_port(args.http_metrics_port)
Graham King's avatar
Graham King committed
286
287
        .tls_cert_path(args.tls_cert_path.clone())
        .tls_key_path(args.tls_key_path.clone())
288
        .is_mocker(matches!(args.engine_type, EngineType::Mocker))
289
        .extra_engine_args(args.extra_engine_args.clone())
290
291
292
        .namespace(args.namespace.clone())
        .custom_backend_metrics_endpoint(args.custom_backend_metrics_endpoint.clone())
        .custom_backend_metrics_polling_interval(args.custom_backend_metrics_polling_interval);
293
    pyo3_async_runtimes::tokio::future_into_py(py, async move {
294
295
296
297
        if let Some(model_path) = args.model_path.clone() {
            let local_path = if model_path.exists() {
                model_path
            } else {
298
299
300
                // Mocker only needs tokenizer, not weights
                let ignore_weights = matches!(args.engine_type, EngineType::Mocker);
                LocalModel::fetch(&model_path.display().to_string(), ignore_weights)
301
302
303
304
305
306
                    .await
                    .map_err(to_pyerr)?
            };
            builder.model_path(local_path);
        }

307
        let local_model = builder.build().await.map_err(to_pyerr)?;
308
        let inner = select_engine(distributed_runtime, args, local_model)
309
310
311
312
313
314
            .await
            .map_err(to_pyerr)?;
        Ok(EngineConfig { inner })
    })
}

315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
/// Convert a PyEngineFactory to a Rust EngineFactoryCallback
fn py_engine_factory_to_callback(factory: PyEngineFactory) -> EngineFactoryCallback {
    let callback = factory.callback;
    let locals = factory.locals;

    Arc::new(
        move |card: RsModelDeploymentCard| -> Pin<
            Box<dyn Future<Output = anyhow::Result<OpenAIChatCompletionsStreamingEngine>> + Send>,
        > {
            let callback = callback.clone();
            let locals = locals.clone();

            Box::pin(async move {
                // Acquire GIL to call Python callback and convert coroutine to future
                let py_future = Python::with_gil(|py| {
                    // Create Python ModelDeploymentCard wrapper
                    let py_card = ModelDeploymentCard { inner: card };
                    let py_card_obj = Py::new(py, py_card)
                        .map_err(|e| anyhow::anyhow!("Failed to create Python MDC: {}", e))?;

                    // Call Python async function to get a coroutine
                    let coroutine = callback
                        .call1(py, (py_card_obj,))
                        .map_err(|e| anyhow::anyhow!("Failed to call engine_factory: {}", e))?;

                    // Use the TaskLocals captured at registration time
                    pyo3_async_runtimes::into_future_with_locals(&locals, coroutine.into_bound(py))
                        .map_err(|e| {
                            anyhow::anyhow!("Failed to convert coroutine to future: {}", e)
                        })
                })?;

                // Await the Python coroutine (GIL is released during await)
                let py_result = py_future
                    .await
                    .map_err(|e| anyhow::anyhow!("engine_factory callback failed: {}", e))?;

                // Extract PythonAsyncEngine from the Python result and wrap in Arc
                let engine: OpenAIChatCompletionsStreamingEngine = Python::with_gil(|py| {
                    let engine: PythonAsyncEngine = py_result.extract(py).map_err(|e| {
                        anyhow::anyhow!("Failed to extract PythonAsyncEngine: {}", e)
                    })?;
                    Ok::<_, anyhow::Error>(Arc::new(engine))
                })?;

                Ok(engine)
            })
        },
    )
}

366
367
async fn select_engine(
    #[allow(unused_variables)] distributed_runtime: super::DistributedRuntime,
368
    args: EntrypointArgs,
369
370
    local_model: LocalModel,
) -> anyhow::Result<RsEngineConfig> {
371
    let inner = match args.engine_type {
372
373
        EngineType::Echo => {
            // There is no validation for the echo engine
374
            RsEngineConfig::InProcessText {
375
                model: Box::new(local_model),
376
                engine: dynamo_llm::engines::make_echo_engine(),
377
378
            }
        }
379
380
381
382
383
384
385
386
        EngineType::Dynamic => {
            //  Convert Python engine factory to Rust callback
            let engine_factory = args.engine_factory.map(py_engine_factory_to_callback);
            RsEngineConfig::Dynamic {
                model: Box::new(local_model),
                engine_factory,
            }
        }
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
        EngineType::Mocker => {
            let mocker_args = if let Some(extra_args_path) = args.extra_engine_args {
                MockEngineArgs::from_json_file(&extra_args_path).map_err(|e| {
                    anyhow::anyhow!(
                        "Failed to load mocker args from {:?}: {}",
                        extra_args_path,
                        e
                    )
                })?
            } else {
                tracing::warn!(
                    "No extra_engine_args specified for mocker engine. Using default mocker args."
                );
                MockEngineArgs::default()
            };

            let endpoint = local_model.endpoint_id().clone();

            let engine = dynamo_llm::mocker::engine::make_mocker_engine(
                distributed_runtime.inner,
                endpoint,
                mocker_args,
            )
            .await?;

412
            RsEngineConfig::InProcessTokens {
413
414
                engine,
                model: Box::new(local_model),
415
                is_prefill: args.is_prefill,
416
417
            }
        }
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
    };

    Ok(inner)
}

#[pyfunction]
#[pyo3(signature = (distributed_runtime, input, engine_config))]
pub fn run_input<'p>(
    py: Python<'p>,
    distributed_runtime: super::DistributedRuntime,
    input: &str,
    engine_config: EngineConfig,
) -> PyResult<Bound<'p, PyAny>> {
    let input_enum: Input = input.parse().map_err(to_pyerr)?;
    pyo3_async_runtimes::tokio::future_into_py(py, async move {
        dynamo_llm::entrypoint::input::run_input(
434
            distributed_runtime.inner.clone(),
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
            input_enum,
            engine_config.inner,
        )
        .await
        .map_err(to_pyerr)?;
        Ok(())
    })
}

pub fn to_pyerr<E>(err: E) -> PyErr
where
    E: Display,
{
    PyException::new_err(format!("{}", err))
}