entrypoint.rs 21.9 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
use pythonize::pythonize;
13

14
use dynamo_kv_router::config::KvRouterConfig as RsKvRouterConfig;
15
use dynamo_llm::discovery::LoadThresholdConfig as RsLoadThresholdConfig;
16
use dynamo_llm::entrypoint::ChatEngineFactoryCallback;
17
use dynamo_llm::entrypoint::EngineConfig as RsEngineConfig;
18
use dynamo_llm::entrypoint::RouterConfig as RsRouterConfig;
19
use dynamo_llm::entrypoint::input::Input;
Graham King's avatar
Graham King committed
20
use dynamo_llm::local_model::DEFAULT_HTTP_PORT;
21
use dynamo_llm::local_model::{LocalModel, LocalModelBuilder};
22
use dynamo_llm::mocker::make_mocker_engine;
23
24
use dynamo_llm::model_card::ModelDeploymentCard as RsModelDeploymentCard;
use dynamo_llm::types::openai::chat_completions::OpenAIChatCompletionsStreamingEngine;
25
26
27
use dynamo_mocker::common::perf_model::PerfModel;

use super::aic_callback::create_aic_callback;
28
use dynamo_mocker::common::protocols::MockEngineArgs;
29
use dynamo_runtime::discovery::ModelCardInstanceId as RsModelCardInstanceId;
30
use dynamo_runtime::protocols::EndpointId;
31

32
use super::local_model::ModelRuntimeConfig;
33
use super::model_card::ModelDeploymentCard;
34
use crate::RouterMode;
35
use crate::engine::PythonAsyncEngine;
36

37
38
39
40
41
#[pyclass(eq, eq_int)]
#[derive(Clone, Debug, PartialEq)]
#[repr(i32)]
pub enum EngineType {
    Echo = 1,
42
43
    Dynamic = 2,
    Mocker = 3,
44
45
}

46
#[pyclass]
47
#[derive(Default, Clone, Debug)]
48
49
50
51
pub struct KvRouterConfig {
    inner: RsKvRouterConfig,
}

52
53
impl KvRouterConfig {
    pub fn inner(&self) -> RsKvRouterConfig {
54
        self.inner.clone()
55
56
57
    }
}

58
59
60
#[pymethods]
impl KvRouterConfig {
    #[new]
61
    #[pyo3(signature = (overlap_score_weight=1.0, router_temperature=0.0, use_kv_events=true, durable_kv_events=false, 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, router_queue_threshold=Some(2.0), router_event_threads=4, router_enable_cache_control=false, router_queue_policy="fcfs", remote_indexer_component=None))]
62
    #[allow(clippy::too_many_arguments)]
63
64
65
66
    fn new(
        overlap_score_weight: f64,
        router_temperature: f64,
        use_kv_events: bool,
67
        durable_kv_events: bool,
68
        router_replica_sync: bool,
69
        router_track_active_blocks: bool,
70
        router_track_output_blocks: bool,
71
        router_assume_kv_reuse: bool,
72
73
        router_snapshot_threshold: Option<u32>,
        router_reset_states: bool,
74
75
76
        router_ttl_secs: f64,
        router_max_tree_size: usize,
        router_prune_target_ratio: f64,
77
        router_queue_threshold: Option<f64>,
Yan Ru Pei's avatar
Yan Ru Pei committed
78
        router_event_threads: u32,
79
        router_enable_cache_control: bool,
80
        router_queue_policy: &str,
81
        remote_indexer_component: Option<String>,
82
    ) -> Self {
83
84
85
86
87
        KvRouterConfig {
            inner: RsKvRouterConfig {
                overlap_score_weight,
                router_temperature,
                use_kv_events,
88
                durable_kv_events,
89
                router_replica_sync,
90
                router_track_active_blocks,
91
                router_track_output_blocks,
92
                router_assume_kv_reuse,
93
94
                router_snapshot_threshold,
                router_reset_states,
95
96
97
                router_ttl_secs,
                router_max_tree_size,
                router_prune_target_ratio,
98
                router_queue_threshold,
Yan Ru Pei's avatar
Yan Ru Pei committed
99
                router_event_threads,
100
                router_enable_cache_control,
101
                skip_initial_worker_wait: false,
102
103
104
                router_queue_policy: router_queue_policy.parse().unwrap_or_else(|_| {
                    panic!("invalid router_queue_policy: {router_queue_policy:?}")
                }),
105
                remote_indexer_component,
106
107
108
109
110
111
112
113
            },
        }
    }
}

#[pyclass]
#[derive(Clone, Debug)]
pub struct RouterConfig {
114
115
116
117
118
119
    #[pyo3(get, set)]
    pub router_mode: RouterMode,

    #[pyo3(get, set)]
    pub kv_router_config: KvRouterConfig,

120
121
122
123
    /// 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>,
124
125
    /// Threshold for active prefill tokens as fraction of max_num_batched_tokens
    active_prefill_tokens_threshold_frac: Option<f64>,
126
    enforce_disagg: bool,
127
128
129
130
131
}

#[pymethods]
impl RouterConfig {
    #[new]
132
    #[pyo3(signature = (mode, config=None, active_decode_blocks_threshold=None, active_prefill_tokens_threshold=None, active_prefill_tokens_threshold_frac=None, enforce_disagg=false))]
133
134
135
    pub fn new(
        mode: RouterMode,
        config: Option<KvRouterConfig>,
136
137
        active_decode_blocks_threshold: Option<f64>,
        active_prefill_tokens_threshold: Option<u64>,
138
        active_prefill_tokens_threshold_frac: Option<f64>,
139
        enforce_disagg: bool,
140
    ) -> Self {
141
142
143
        Self {
            router_mode: mode,
            kv_router_config: config.unwrap_or_default(),
144
145
            active_decode_blocks_threshold,
            active_prefill_tokens_threshold,
146
            active_prefill_tokens_threshold_frac,
147
            enforce_disagg,
148
149
150
151
152
153
154
155
156
        }
    }
}

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,
157
158
159
160
161
            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,
            },
162
            enforce_disagg: rc.enforce_disagg,
163
164
165
166
        }
    }
}

167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
/// 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()
    }
}

182
183
184
185
186
187
188
189
190
#[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>,
191
    router_config: Option<RouterConfig>,
192
    kv_cache_block_size: Option<u32>,
193
    http_host: Option<String>,
Graham King's avatar
Graham King committed
194
    http_port: u16,
195
    http_metrics_port: Option<u16>,
Graham King's avatar
Graham King committed
196
197
    tls_cert_path: Option<PathBuf>,
    tls_key_path: Option<PathBuf>,
198
    extra_engine_args: Option<PathBuf>,
199
    runtime_config: Option<ModelRuntimeConfig>,
200
    namespace: Option<String>,
201
    namespace_prefix: Option<String>,
202
    is_prefill: bool,
203
    migration_limit: u32,
204
    chat_engine_factory: Option<PyEngineFactory>,
205
206
207
208
209
210
}

#[pymethods]
impl EntrypointArgs {
    #[allow(clippy::too_many_arguments)]
    #[new]
211
    #[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, runtime_config=None, namespace=None, namespace_prefix=None, is_prefill=false, migration_limit=0, chat_engine_factory=None))]
212
    pub fn new(
213
        py: Python<'_>,
214
215
216
217
218
219
        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>,
220
        router_config: Option<RouterConfig>,
221
        kv_cache_block_size: Option<u32>,
222
        http_host: Option<String>,
223
        http_port: Option<u16>,
224
        http_metrics_port: Option<u16>,
Graham King's avatar
Graham King committed
225
226
        tls_cert_path: Option<PathBuf>,
        tls_key_path: Option<PathBuf>,
227
        extra_engine_args: Option<PathBuf>,
228
        runtime_config: Option<ModelRuntimeConfig>,
229
        namespace: Option<String>,
230
        namespace_prefix: Option<String>,
231
        is_prefill: bool,
232
        migration_limit: u32,
233
        chat_engine_factory: Option<PyObject>,
234
    ) -> PyResult<Self> {
235
        let endpoint_id_obj: Option<EndpointId> = endpoint_id.as_deref().map(EndpointId::from);
Graham King's avatar
Graham King committed
236
237
238
239
240
241
242
        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",
            ));
        }
243

244
245
        // Capture TaskLocals at registration time for the chat engine factory callback
        let chat_engine_factory = chat_engine_factory
246
247
248
            .map(|callback| {
                let locals = pyo3_async_runtimes::tokio::get_current_locals(py).map_err(|e| {
                    pyo3::exceptions::PyRuntimeError::new_err(format!(
249
                        "Failed to get TaskLocals for chat_engine_factory: {}",
250
251
252
253
254
255
256
257
258
259
                        e
                    ))
                })?;
                Ok::<_, PyErr>(PyEngineFactory {
                    callback: Arc::new(callback),
                    locals: Arc::new(locals),
                })
            })
            .transpose()?;

260
261
262
263
264
265
266
        Ok(EntrypointArgs {
            engine_type,
            model_path,
            model_name,
            endpoint_id: endpoint_id_obj,
            context_length,
            template_file,
267
            router_config,
268
            kv_cache_block_size,
269
            http_host,
Graham King's avatar
Graham King committed
270
            http_port: http_port.unwrap_or(DEFAULT_HTTP_PORT),
271
            http_metrics_port,
Graham King's avatar
Graham King committed
272
273
            tls_cert_path,
            tls_key_path,
274
            extra_engine_args,
275
            runtime_config,
276
            namespace,
277
            namespace_prefix,
278
            is_prefill,
279
            migration_limit,
280
            chat_engine_factory,
281
282
283
284
285
286
287
288
289
290
        })
    }
}

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

291
292
/// Create the backend engine wrapper to run the model.
/// Download the model if necessary.
293
294
295
296
297
298
299
300
301
#[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
302
303
304
305
306
        .model_name(
            args.model_name
                .clone()
                .or_else(|| args.model_path.clone().map(|p| p.display().to_string())),
        )
307
        .endpoint_id(args.endpoint_id.clone())
308
        .context_length(args.context_length)
309
        .request_template(args.template_file.clone())
310
        .kv_cache_block_size(args.kv_cache_block_size)
311
        .router_config(args.router_config.clone().map(|rc| rc.into()))
312
        .migration_limit(Some(args.migration_limit))
313
        .http_host(args.http_host.clone())
314
        .http_port(args.http_port)
315
        .http_metrics_port(args.http_metrics_port)
Graham King's avatar
Graham King committed
316
317
        .tls_cert_path(args.tls_cert_path.clone())
        .tls_key_path(args.tls_key_path.clone())
318
        .is_mocker(matches!(args.engine_type, EngineType::Mocker))
319
        .extra_engine_args(args.extra_engine_args.clone())
320
        .runtime_config(args.runtime_config.clone().unwrap_or_default().inner)
321
322
        .namespace(args.namespace.clone())
        .namespace_prefix(args.namespace_prefix.clone());
323
    pyo3_async_runtimes::tokio::future_into_py(py, async move {
324
325
326
327
        if let Some(model_path) = args.model_path.clone() {
            let local_path = if model_path.exists() {
                model_path
            } else {
328
329
330
                // Mocker only needs tokenizer, not weights
                let ignore_weights = matches!(args.engine_type, EngineType::Mocker);
                LocalModel::fetch(&model_path.display().to_string(), ignore_weights)
331
332
333
334
335
336
                    .await
                    .map_err(to_pyerr)?
            };
            builder.model_path(local_path);
        }

337
        let local_model = builder.build().await.map_err(to_pyerr)?;
338
        let inner = select_engine(distributed_runtime, args, local_model)
339
340
341
342
343
344
            .await
            .map_err(to_pyerr)?;
        Ok(EngineConfig { inner })
    })
}

345
346
/// Convert a PyEngineFactory to a Rust ChatEngineFactoryCallback
fn py_engine_factory_to_callback(factory: PyEngineFactory) -> ChatEngineFactoryCallback {
347
348
349
350
    let callback = factory.callback;
    let locals = factory.locals;

    Arc::new(
351
352
353
        move |instance_id: RsModelCardInstanceId,
              card: RsModelDeploymentCard|
              -> Pin<
354
355
356
357
358
359
360
361
            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| {
362
363
364
365
                    let py_instance_id =
                        Py::new(py, crate::ModelCardInstanceId { inner: instance_id }).map_err(
                            |e| anyhow::anyhow!("Failed to create Python ModelCardInstanceId: {e}"),
                        )?;
366
367
368
                    // Create Python ModelDeploymentCard wrapper
                    let py_card = ModelDeploymentCard { inner: card };
                    let py_card_obj = Py::new(py, py_card)
369
                        .map_err(|e| anyhow::anyhow!("Failed to create Python MDC: {e}"))?;
370
371
372

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

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

                // Await the Python coroutine (GIL is released during await)
                let py_result = py_future
                    .await
384
                    .map_err(|e| anyhow::anyhow!("chat_engine_factory callback failed: {}", e))?;
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399

                // 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)
            })
        },
    )
}

400
401
async fn select_engine(
    #[allow(unused_variables)] distributed_runtime: super::DistributedRuntime,
402
    args: EntrypointArgs,
403
404
    local_model: LocalModel,
) -> anyhow::Result<RsEngineConfig> {
405
    let inner = match args.engine_type {
406
407
        EngineType::Echo => {
            // There is no validation for the echo engine
408
            RsEngineConfig::InProcessText {
409
                model: Box::new(local_model),
410
                engine: dynamo_llm::engines::make_echo_engine(),
411
412
            }
        }
413
        EngineType::Dynamic => {
414
415
            //  Convert Python chat engine factory to Rust callback
            let chat_engine_factory = args.chat_engine_factory.map(py_engine_factory_to_callback);
416
417
            RsEngineConfig::Dynamic {
                model: Box::new(local_model),
418
                chat_engine_factory,
419
420
            }
        }
421
        EngineType::Mocker => {
422
            let mut mocker_args = if let Some(extra_args_path) = args.extra_engine_args {
423
424
425
426
427
428
429
430
431
432
433
434
435
436
                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()
            };

437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
            // If aic_backend is set, create Python AIC callback and override perf_model
            if let Some(ref backend_name) = mocker_args.aic_backend {
                let backend = backend_name.clone();
                let system = mocker_args.aic_system.as_deref().unwrap_or("h200_sxm");
                let model_name = mocker_args
                    .aic_model_path
                    .as_deref()
                    .unwrap_or_else(|| local_model.card().source_path());
                let backend_version = mocker_args.aic_backend_version.as_deref();
                let tp_size = mocker_args.aic_tp_size.unwrap_or(1);
                match Python::with_gil(|py| {
                    create_aic_callback(py, &backend, system, model_name, tp_size, backend_version)
                }) {
                    Ok(callback) => {
                        tracing::info!(
                            "AIC perf model: backend={}, gpu={}, model={}, version={:?}",
                            backend,
                            system,
                            model_name,
                            backend_version
                        );
                        mocker_args.perf_model = Arc::new(PerfModel::from_aic_callback(callback));
                    }
                    Err(e) => {
                        return Err(anyhow::anyhow!(
                            "Failed to create AIC callback (--aic-perf-model was requested): {}",
                            e
                        ));
                    }
                }
            }

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

471
            let engine =
472
                make_mocker_engine(distributed_runtime.inner, endpoint, mocker_args).await?;
473

474
            RsEngineConfig::InProcessTokens {
475
476
                engine,
                model: Box::new(local_model),
477
                is_prefill: args.is_prefill,
478
479
            }
        }
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
    };

    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(
496
            distributed_runtime.inner.clone(),
497
498
499
500
501
502
503
504
505
            input_enum,
            engine_config.inner,
        )
        .await
        .map_err(to_pyerr)?;
        Ok(())
    })
}

506
507
508
509
510
511
512
513
514
#[pyfunction]
#[pyo3(signature = (trace_file, extra_engine_args=None, num_workers=1, replay_concurrency=None))]
pub fn run_mocker_trace_replay(
    py: Python<'_>,
    trace_file: PathBuf,
    extra_engine_args: Option<PathBuf>,
    num_workers: usize,
    replay_concurrency: Option<isize>,
) -> PyResult<PyObject> {
515
516
517
518
519
520
521
522
523
524
525
    // Load args before allow_threads so we can use the GIL for AIC callback creation.
    let mut args = if let Some(ref extra_args_path) = extra_engine_args {
        MockEngineArgs::from_json_file(extra_args_path).map_err(|e| {
            PyException::new_err(format!(
                "Failed to load mocker args from {:?}: {}",
                extra_args_path, e
            ))
        })?
    } else {
        MockEngineArgs::default()
    };
526

527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
    // Create AIC callback if requested (requires GIL, must be done before allow_threads).
    if let Some(ref backend_name) = args.aic_backend.clone() {
        let backend = backend_name.clone();
        let system = args.aic_system.as_deref().unwrap_or("h200_sxm").to_string();
        let model_name = args
            .aic_model_path
            .clone()
            .ok_or_else(|| PyException::new_err("--aic-perf-model requires --model-path"))?;
        let backend_version = args.aic_backend_version.clone();
        let tp_size = args.aic_tp_size.unwrap_or(1);
        let callback = create_aic_callback(
            py,
            &backend,
            &system,
            &model_name,
            tp_size,
            backend_version.as_deref(),
        )
        .map_err(|e| {
            PyException::new_err(format!(
                "Failed to create AIC callback (--aic-perf-model was requested): {}",
                e
            ))
        })?;
        tracing::info!(
            "AIC perf model: backend={}, gpu={}, model={}, version={:?}",
            backend,
            system,
            model_name,
            backend_version
        );
        args.perf_model = Arc::new(PerfModel::from_aic_callback(callback));
    }

    let report = py.allow_threads(move || {
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
        let replay_concurrency = replay_concurrency
            .map(usize::try_from)
            .transpose()
            .map_err(|_| anyhow::anyhow!("replay_concurrency must be at least 1"))?;

        if let Some(max_in_flight) = replay_concurrency {
            dynamo_mocker::simulation::simulate_concurrency_file(
                args,
                &trace_file,
                max_in_flight,
                num_workers,
            )
        } else {
            dynamo_mocker::simulation::simulate_trace_file(args, &trace_file, num_workers)
        }
    });
    let report = report.map_err(to_pyerr)?;
    pythonize(py, &report)
        .map_err(to_pyerr)
        .map(|obj| obj.unbind())
}

584
585
586
587
588
589
pub fn to_pyerr<E>(err: E) -> PyErr
where
    E: Display,
{
    PyException::new_err(format!("{}", err))
}