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

4
use dynamo_llm::local_model::LocalModel;
5
6
use dynamo_runtime::distributed::DistributedConfig;
use dynamo_runtime::storage::key_value_store::KeyValueStoreSelect;
7
8
use futures::StreamExt;
use once_cell::sync::OnceCell;
9
use pyo3::IntoPyObjectExt;
10
use pyo3::exceptions::PyStopAsyncIteration;
Richard Huo's avatar
Richard Huo committed
11
use pyo3::types::PyCapsule;
12
use pyo3::types::{PyDict, PyString};
13
use pyo3::{exceptions::PyException, prelude::*};
14
use rand::seq::IteratorRandom as _;
15
use rs::pipeline::network::Ingress;
Richard Huo's avatar
Richard Huo committed
16
use std::ffi::CString;
17
use std::fs;
18
use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4};
19
use std::path::PathBuf;
20
use std::time::Duration;
Richard Huo's avatar
Richard Huo committed
21
22
23
24
use std::{
    fmt::Display,
    sync::{Arc, Weak},
};
25
use tokio::sync::Mutex;
26
use tracing::Instrument;
27

Neelay Shah's avatar
Neelay Shah committed
28
use dynamo_runtime::{
Ryan Olson's avatar
Ryan Olson committed
29
    self as rs, logging,
30
    pipeline::{
31
        AsyncEngineContextProvider, EngineStream, ManyOut, SingleIn, context::Context as RsContext,
32
        network::egress::push_router::RouterMode as RsRouterMode,
33
    },
34
    protocols::annotated::Annotated as RsAnnotated,
35
    traits::DistributedRuntimeProvider,
36
37
};

Neelay Shah's avatar
Neelay Shah committed
38
use dynamo_llm::{self as llm_rs};
39
40
use dynamo_llm::{entrypoint::RouterConfig, kv_router::KvRouterConfig};

41
use crate::llm::local_model::ModelRuntimeConfig;
42
use crate::llm::preprocessor::{MediaDecoder, MediaFetcher};
43

44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
#[pyclass(eq, eq_int)]
#[derive(Clone, Debug, PartialEq)]
pub enum RouterMode {
    RoundRobin,
    Random,
    KV,
}

impl From<RouterMode> for RsRouterMode {
    fn from(mode: RouterMode) -> Self {
        match mode {
            RouterMode::RoundRobin => Self::RoundRobin,
            RouterMode::Random => Self::Random,
            RouterMode::KV => Self::KV,
        }
    }
}
61

62
mod context;
63
mod engine;
64
mod http;
65
mod kserve_grpc;
66
mod llm;
67
mod parsers;
68
mod planner;
69
mod prometheus_metrics;
70
71
72
73
74
75
76
77

type JsonServerStreamingIngress =
    Ingress<SingleIn<serde_json::Value>, ManyOut<RsAnnotated<serde_json::Value>>>;

static INIT: OnceCell<()> = OnceCell::new();

const DEFAULT_ANNOTATED_SETTING: Option<bool> = Some(true);

78
79
// Helper to get appropriate span for instrumentation - always emit spans
fn get_span_for_context(context: &context::Context, operation: &str) -> tracing::Span {
80
81
82
83
84
85
    logging::make_client_request_span(
        operation,
        context.inner().id(),
        context.trace_context(),
        None,
    )
86
87
88
89
90
91
92
93
}

// Helper to create span for direct method with instance_id
fn get_span_for_direct_context(
    context: &context::Context,
    operation: &str,
    instance_id: &str,
) -> tracing::Span {
94
95
96
97
98
99
    logging::make_client_request_span(
        operation,
        context.inner().id(),
        context.trace_context(),
        Some(instance_id),
    )
100
101
}

102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
// Helper to create request context with proper linking and cancellation handling
fn create_request_context(
    request: serde_json::Value,
    parent_ctx: &Option<context::Context>,
) -> RsContext<serde_json::Value> {
    match parent_ctx {
        // If there is a parent context, link the request as a child context of it
        Some(parent_ctx) => {
            let child_ctx = RsContext::with_id(request, parent_ctx.inner().id().to_string());
            parent_ctx.inner().link_child(child_ctx.context());
            if parent_ctx.inner().is_stopped() || parent_ctx.inner().is_killed() {
                // Let the server handle the cancellation for now since not all backends are
                // properly handling request exceptions
                // TODO: (DIS-830) Return an error if context is cancelled
                child_ctx.context().stop_generating();
            }
            child_ctx
        }
        // Otherwise if there is no parent context, use the request as-is
        _ => request.into(),
    }
}

125
126
127
128
129
/// A Python module implemented in Rust. The name of this function must match
/// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
/// import the module.
#[pymodule]
fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
130
    // Initialize logging early unless OTEL export is enabled (which requires tokio runtime)
131
    if rs::config::env_is_truthy("OTEL_EXPORT_ENABLED") {
132
        eprintln!(
133
            "Warning: OTEL_EXPORT_ENABLED detected. Logging initialization deferred until runtime is available. Early logs may be dropped."
134
135
136
137
138
        );
    } else {
        rs::logging::init();
    }

Yan Ru Pei's avatar
Yan Ru Pei committed
139
    m.add_function(wrap_pyfunction!(llm::kv::compute_block_hash_for_seq_py, m)?)?;
Ryan Olson's avatar
Ryan Olson committed
140
    m.add_function(wrap_pyfunction!(log_message, m)?)?;
141
    m.add_function(wrap_pyfunction!(register_llm, m)?)?;
142
    m.add_function(wrap_pyfunction!(fetch_llm, m)?)?;
143
144
    m.add_function(wrap_pyfunction!(llm::entrypoint::make_engine, m)?)?;
    m.add_function(wrap_pyfunction!(llm::entrypoint::run_input, m)?)?;
145

146
147
148
149
150
151
152
    m.add_class::<DistributedRuntime>()?;
    m.add_class::<CancellationToken>()?;
    m.add_class::<Namespace>()?;
    m.add_class::<Component>()?;
    m.add_class::<Endpoint>()?;
    m.add_class::<Client>()?;
    m.add_class::<AsyncResponseStream>()?;
153
154
155
    m.add_class::<llm::entrypoint::EntrypointArgs>()?;
    m.add_class::<llm::entrypoint::EngineConfig>()?;
    m.add_class::<llm::entrypoint::EngineType>()?;
156
157
    m.add_class::<llm::entrypoint::RouterConfig>()?;
    m.add_class::<llm::entrypoint::KvRouterConfig>()?;
158
    m.add_class::<llm::kv::WorkerMetricsPublisher>()?;
159
    m.add_class::<llm::model_card::ModelDeploymentCard>()?;
160
    m.add_class::<llm::local_model::ModelRuntimeConfig>()?;
161
    m.add_class::<llm::preprocessor::OAIChatPreprocessor>()?;
162
163
    m.add_class::<llm::preprocessor::MediaDecoder>()?;
    m.add_class::<llm::preprocessor::MediaFetcher>()?;
164
    m.add_class::<llm::backend::Backend>()?;
165
166
    m.add_class::<llm::kv::OverlapScores>()?;
    m.add_class::<llm::kv::KvIndexer>()?;
167
    m.add_class::<llm::kv::ApproxKvIndexer>()?;
168
    m.add_class::<llm::kv::KvEventPublisher>()?;
Yan Ru Pei's avatar
Yan Ru Pei committed
169
170
    m.add_class::<llm::kv::RadixTree>()?;
    m.add_class::<llm::kv::ZmqKvEventListener>()?;
171
172
    m.add_class::<llm::kv::ZmqKvEventPublisher>()?;
    m.add_class::<llm::kv::ZmqKvEventPublisherConfig>()?;
173
    m.add_class::<llm::kv::KvRecorder>()?;
174
175
    m.add_class::<http::HttpService>()?;
    m.add_class::<http::HttpAsyncEngine>()?;
176
    m.add_class::<context::Context>()?;
177
    m.add_class::<ModelType>()?;
178
    m.add_class::<ModelInput>()?;
179
180
181
182
    m.add_class::<llm::kv::ForwardPassMetrics>()?;
    m.add_class::<llm::kv::WorkerStats>()?;
    m.add_class::<llm::kv::KvStats>()?;
    m.add_class::<llm::kv::SpecDecodeStats>()?;
183
184
    m.add_class::<llm::kv::KvPushRouter>()?;
    m.add_class::<llm::kv::KvPushRouterStream>()?;
185
    m.add_class::<RouterMode>()?;
186
    m.add_class::<kserve_grpc::KserveGrpcService>()?;
187
    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
188
189
190
    m.add_class::<planner::VirtualConnectorCoordinator>()?;
    m.add_class::<planner::VirtualConnectorClient>()?;
    m.add_class::<planner::PlannerDecision>()?;
191
192

    engine::add_to_module(m)?;
193
    parsers::add_to_module(m)?;
194

195
    m.add_class::<prometheus_metrics::RuntimeMetrics>()?;
196
197
198
199
    let prometheus_metrics = PyModule::new(m.py(), "prometheus_metrics")?;
    prometheus_metrics::add_to_module(&prometheus_metrics)?;
    m.add_submodule(&prometheus_metrics)?;

200
201
202
203
204
205
206
207
208
209
    Ok(())
}

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

Ryan Olson's avatar
Ryan Olson committed
210
211
212
213
214
215
216
/// Log a message from Python with file and line info
#[pyfunction]
#[pyo3(text_signature = "(level, message, module, file, line)")]
fn log_message(level: &str, message: &str, module: &str, file: &str, line: u32) {
    logging::log_message(level, message, module, file, line);
}

217
218
/// Create an engine and attach it to an endpoint to make it visible to the frontend.
/// This is the main way you create a Dynamo worker / backend.
219
#[pyfunction]
220
#[pyo3(signature = (model_input, model_type, endpoint, model_path, model_name=None, context_length=None, kv_cache_block_size=None, router_mode=None, migration_limit=0, runtime_config=None, user_data=None, custom_template_path=None, media_decoder=None, media_fetcher=None))]
221
#[allow(clippy::too_many_arguments)]
222
223
fn register_llm<'p>(
    py: Python<'p>,
224
    model_input: ModelInput,
225
    model_type: ModelType,
226
227
228
    endpoint: Endpoint,
    model_path: &str,
    model_name: Option<&str>,
229
230
    context_length: Option<u32>,
    kv_cache_block_size: Option<u32>,
231
    router_mode: Option<RouterMode>,
232
    migration_limit: u32,
233
    runtime_config: Option<ModelRuntimeConfig>,
234
    user_data: Option<&Bound<'p, PyDict>>,
235
    custom_template_path: Option<&str>,
236
237
    media_decoder: Option<MediaDecoder>,
    media_fetcher: Option<MediaFetcher>,
238
) -> PyResult<Bound<'p, PyAny>> {
239
240
241
242
243
244
245
246
247
248
249
250
251
252
    // Validate Prefill model type requirements
    if model_type.inner == llm_rs::model_type::ModelType::Prefill {
        if !matches!(model_input, ModelInput::Tokens) {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "ModelType::Prefill requires model_input to be ModelInput::Tokens",
            ));
        }
        if migration_limit != 0 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "ModelType::Prefill requires migration_limit to be 0",
            ));
        }
    }

253
254
255
    let model_input = match model_input {
        ModelInput::Text => llm_rs::model_type::ModelInput::Text,
        ModelInput::Tokens => llm_rs::model_type::ModelInput::Tokens,
256
        ModelInput::Tensor => llm_rs::model_type::ModelInput::Tensor,
257
258
    };

259
260
    let model_type_obj = model_type.inner;

261
    let inner_path = model_path.to_string();
262
    let mut model_name = model_name.map(|n| n.to_string());
263
264
265
    let router_mode = router_mode.unwrap_or(RouterMode::RoundRobin);
    let router_config = RouterConfig::new(router_mode.into(), KvRouterConfig::default());

266
267
268
269
270
271
272
273
274
275
276
277
278
    // Early validation of custom template path
    let custom_template_path_owned = custom_template_path
        .map(|s| {
            let path = PathBuf::from(s);
            if !path.exists() {
                return Err(PyErr::new::<pyo3::exceptions::PyFileNotFoundError, _>(
                    format!("Custom template file does not exist: {}", path.display()),
                ));
            }
            Ok(path)
        })
        .transpose()?;

279
280
281
282
283
284
285
    let user_data_json = user_data
        .map(|dict| pythonize::depythonize(dict))
        .transpose()
        .map_err(|err| {
            PyErr::new::<PyException, _>(format!("Failed to convert user_data: {}", err))
        })?;

286
    pyo3_async_runtimes::tokio::future_into_py(py, async move {
287
288
289
290
291
292
293
294
295
296
297
298
299
        let model_path = if fs::exists(&inner_path)? {
            PathBuf::from(inner_path)
        } else {
            // Preserve the model name
            if model_name.is_none() {
                model_name = Some(inner_path.clone());
            }
            // Likely it's a Hugging Face repo, download it
            LocalModel::fetch(&inner_path, false)
                .await
                .map_err(to_pyerr)?
        };

300
301
        let mut builder = dynamo_llm::local_model::LocalModelBuilder::default();
        builder
302
            .model_path(model_path)
303
304
            .model_name(model_name)
            .context_length(context_length)
305
            .kv_cache_block_size(kv_cache_block_size)
306
            .router_config(Some(router_config))
307
            .migration_limit(Some(migration_limit))
308
            .runtime_config(runtime_config.unwrap_or_default().inner)
309
            .user_data(user_data_json)
310
311
312
            .custom_template_path(custom_template_path_owned)
            .media_decoder(media_decoder.map(|m| m.inner))
            .media_fetcher(media_fetcher.map(|m| m.inner));
313
        // Load the ModelDeploymentCard
314
        let mut local_model = builder.build().await.map_err(to_pyerr)?;
315
316
        // Advertise ourself on etcd so ingress can find us
        local_model
317
            .attach(&endpoint.inner, model_type_obj, model_input)
318
319
320
321
322
323
324
            .await
            .map_err(to_pyerr)?;

        Ok(())
    })
}

325
326
327
328
329
330
331
332
333
334
335
/// Download a model from Hugging Face, returning it's local path
/// Example: `model_path = await fetch_llm("Qwen/Qwen3-0.6B")`
#[pyfunction]
#[pyo3(signature = (remote_name))]
fn fetch_llm<'p>(py: Python<'p>, remote_name: &str) -> PyResult<Bound<'p, PyAny>> {
    let repo = remote_name.to_string();
    pyo3_async_runtimes::tokio::future_into_py(py, async move {
        LocalModel::fetch(&repo, false).await.map_err(to_pyerr)
    })
}

336
337
#[pyclass]
#[derive(Clone)]
Ryan Olson's avatar
Ryan Olson committed
338
pub struct DistributedRuntime {
339
340
341
342
    inner: rs::DistributedRuntime,
    event_loop: PyObject,
}

Ryan Olson's avatar
Ryan Olson committed
343
344
impl DistributedRuntime {
    #[allow(dead_code)]
345
    pub(crate) fn inner(&self) -> &rs::DistributedRuntime {
Ryan Olson's avatar
Ryan Olson committed
346
347
348
349
        &self.inner
    }
}

350
#[pyclass]
351
#[derive(Clone)]
352
353
354
355
356
struct CancellationToken {
    inner: rs::CancellationToken,
}

#[pyclass]
357
#[derive(Clone)]
358
359
360
361
362
363
struct Namespace {
    inner: rs::component::Namespace,
    event_loop: PyObject,
}

#[pyclass]
364
#[derive(Clone)]
365
366
367
368
369
370
struct Component {
    inner: rs::component::Component,
    event_loop: PyObject,
}

#[pyclass]
371
#[derive(Clone)]
372
373
374
375
376
377
struct Endpoint {
    inner: rs::component::Endpoint,
    event_loop: PyObject,
}

#[pyclass]
378
#[derive(Clone)]
379
struct Client {
380
    router: rs::pipeline::PushRouter<serde_json::Value, RsAnnotated<serde_json::Value>>,
381
382
}

383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
#[pyclass]
#[derive(Clone, PartialEq)]
struct ModelType {
    inner: llm_rs::model_type::ModelType,
}

#[pymethods]
#[allow(non_upper_case_globals)]
impl ModelType {
    #[classattr]
    const Chat: Self = ModelType {
        inner: llm_rs::model_type::ModelType::Chat,
    };
    #[classattr]
    const Completions: Self = ModelType {
        inner: llm_rs::model_type::ModelType::Completions,
    };
    #[classattr]
    const Embedding: Self = ModelType {
        inner: llm_rs::model_type::ModelType::Embedding,
    };
404
405
406
407
    #[classattr]
    const TensorBased: Self = ModelType {
        inner: llm_rs::model_type::ModelType::TensorBased,
    };
408
409
410
411
    #[classattr]
    const Prefill: Self = ModelType {
        inner: llm_rs::model_type::ModelType::Prefill,
    };
412
413
414
415
416
417
418
419
420
421
422
423

    fn __or__(&self, other: &Self) -> Self {
        ModelType {
            inner: self.inner | other.inner,
        }
    }

    fn __str__(&self) -> String {
        self.inner.to_string()
    }
}

424
425
#[pyclass(eq, eq_int)]
#[derive(Clone, PartialEq)]
426
427
428
enum ModelInput {
    Text = 1,
    Tokens = 2,
429
    Tensor = 3,
430
431
}

432
433
434
#[pymethods]
impl DistributedRuntime {
    #[new]
435
    fn new(event_loop: PyObject, store_kv: String) -> PyResult<Self> {
436
437
        let selected_kv_store: KeyValueStoreSelect = store_kv.parse().map_err(to_pyerr)?;

438
439
440
        // Try to get existing runtime first, create new Worker only if needed
        // This allows multiple DistributedRuntime instances to share the same tokio runtime
        let runtime = rs::Worker::runtime_from_existing()
441
            .or_else(|_| -> anyhow::Result<rs::Runtime> {
442
443
444
445
                // No existing Worker, create new one
                let worker = rs::Worker::from_settings()?;

                // Initialize pyo3 bridge (only happens once per process)
446
                INIT.get_or_try_init(|| -> anyhow::Result<()> {
447
448
                    let primary = worker.tokio_runtime()?;
                    pyo3_async_runtimes::tokio::init_with_runtime(primary).map_err(|e| {
449
                        anyhow::anyhow!("failed to initialize pyo3 static runtime: {:?}", e)
450
                    })?;
451
                    Ok(())
452
453
                })?;

454
                Ok(worker.runtime().clone())
455
456
            })
            .map_err(to_pyerr)?;
457

458
459
        // Initialize logging in context where tokio runtime is available
        // otel exporter requires it
460
        if rs::config::env_is_truthy("OTEL_EXPORT_ENABLED") {
461
462
463
464
            runtime.secondary().block_on(async {
                rs::logging::init();
            });
        }
465

466
467
468
469
470
471
472
473
        let runtime_config = DistributedConfig {
            store_backend: selected_kv_store,
            nats_config: dynamo_runtime::transports::nats::ClientOptions::default(),
        };
        let inner = runtime
            .secondary()
            .block_on(rs::DistributedRuntime::new(runtime, runtime_config))
            .map_err(to_pyerr)?;
474
475
476
477

        Ok(DistributedRuntime { inner, event_loop })
    }

Ryan Olson's avatar
Ryan Olson committed
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
    #[staticmethod]
    fn detached(py: Python) -> PyResult<Self> {
        let rt = rs::Worker::runtime_from_existing().map_err(to_pyerr)?;
        let handle = rt.primary();

        let inner = handle
            .block_on(rs::DistributedRuntime::from_settings(rt))
            .map_err(to_pyerr)?;

        Ok(DistributedRuntime {
            inner,
            event_loop: py.None(),
        })
    }

493
494
495
496
497
498
499
    fn namespace(&self, name: String) -> PyResult<Namespace> {
        Ok(Namespace {
            inner: self.inner.namespace(name).map_err(to_pyerr)?,
            event_loop: self.event_loop.clone(),
        })
    }

500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
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
562
563
564
565
566
567
568
    /// Allocate a contiguous block of ports from the specified range and atomically reserve them.
    /// Returns a list of all allocated ports in order.
    #[pyo3(signature = (namespace, port_min, port_max, block_size, context=None))]
    fn allocate_port_block<'p>(
        &self,
        py: Python<'p>,
        namespace: &str,
        port_min: u16,
        port_max: u16,
        block_size: u16,
        context: Option<String>, // Optional info to store alongside the reservation
    ) -> PyResult<Bound<'p, PyAny>> {
        const MAX_ALLOCATE_ATTEMPTS: usize = 100;
        if block_size == 0 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "Block size must be at least 1",
            ));
        }

        let Some(etcd_client) = self.inner.etcd_client() else {
            return Err(PyErr::new::<PyException, _>(
                "Static workers should not need to reserve ports",
            ));
        };

        let min = port_min;
        let max = port_max;

        // Compute maximum valid starting port (inclusive)
        let max_start_port = max.saturating_sub(block_size.saturating_sub(1));
        if max_start_port < min {
            return Err(PyErr::new::<PyException, _>(format!(
                "Port range {min}-{max} is too small for block size {block_size}",
            )));
        }

        // Randomize candidate starting ports to reduce contention/races
        let candidate_count =
            (max_start_port - port_min + 1).min(MAX_ALLOCATE_ATTEMPTS as u16) as usize;
        let mut rng = rand::rng();
        let candidate_ports: Vec<u16> =
            (port_min..=max_start_port).choose_multiple(&mut rng, candidate_count);

        let local_ip = match local_ip() {
            Ok(ip) => ip,
            Err(err) => {
                return Err(PyErr::new::<PyException, _>(format!(
                    "Failed fetching local IP address: {err}"
                )));
            }
        };

        let context_bytes = context.map(|s| s.as_bytes().to_vec()).unwrap_or_default();
        let namespace = namespace.to_owned();

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            for (attempt_idx, start_port) in candidate_ports.into_iter().enumerate() {
                let end_port_exclusive = start_port + block_size;
                let ports_to_reserve: Vec<u16> = (start_port..end_port_exclusive).collect();

                // Hold/bind all ports in the block
                let mut sockets = Vec::with_capacity(ports_to_reserve.len());
                let mut bind_failed = false;

                for &port in &ports_to_reserve {
                    match bind_tcp_port(port) {
                        Ok(sock) => sockets.push(sock),
                        Err(e) => {
                            tracing::error!(
569
570
571
                                "Failed to bind to port block starting at {start_port} (attempt {}): {e}",
                                attempt_idx + 1,
                            );
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
                            bind_failed = true;
                            break;
                        }
                    }
                }

                if bind_failed {
                    // Let previously bound sockets drop here
                    if attempt_idx < candidate_count - 1 {
                        tokio::time::sleep(Duration::from_millis(10)).await;
                    }
                    continue;
                }

                // With sockets held, reserve in ETCD
                let mut reserved_keys = Vec::with_capacity(ports_to_reserve.len());
                let mut reservation_failed = false;
                for port in &ports_to_reserve {
                    let key = make_port_key(&namespace, local_ip, *port).map_err(to_pyerr)?;

                    if let Err(e) = etcd_client
                        .kv_create(&key, context_bytes.clone(), None)
                        .await
                    {
                        tracing::error!(
                            "Failed to reserve port block starting at {start_port} (attempt {}): {e}",
                            attempt_idx + 1,
                        );
                        reservation_failed = true;
                        break;
                    }
                    reserved_keys.push(key);
                }

                if reservation_failed {
                    // Cleanup partial reservations
                    for key in reserved_keys {
                        if let Err(e) = etcd_client.kv_delete(key.as_str(), None).await {
                            tracing::warn!("Failed to cleanup reserved port {key}: {e}");
                        }
                    }

                    // Sockets automatically released via RAII
                    if attempt_idx < candidate_count - 1 {
                        tokio::time::sleep(Duration::from_millis(10)).await;
                    }
                    continue;
                }

                // Success - sockets will be released automatically
                tracing::debug!("Reserved port block {ports_to_reserve:?}");
                return Ok(ports_to_reserve);
            }

            Err(PyErr::new::<PyException, _>(format!(
627
628
                "Failed to allocate and reserve a port block of size {block_size} from range {min}-{max} after {candidate_count} attempts"
            )))
629
630
631
        })
    }

632
    fn shutdown(&self) {
633
        self.inner.shutdown();
634
635
636
637
638
    }

    fn event_loop(&self) -> PyObject {
        self.event_loop.clone()
    }
639
640
641
642
643

    fn child_token(&self) -> CancellationToken {
        let inner = self.inner.runtime().child_token();
        CancellationToken { inner }
    }
Richard Huo's avatar
Richard Huo committed
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658

    // This is used to pass the DistributedRuntime from the dynamo-runtime bindings
    // to the KVBM bindings, since KVBM cannot directly use the struct from this cdylib.
    // TODO: Create a separate crate "dynamo-python" so that all binding crates can import
    // from it and share the same crate path. This will allow PyO3 to automatically
    // recognize that both bindings use the same PyClass.
    #[pyo3(name = "to_capsule")]
    fn to_capsule<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyCapsule>> {
        let arc: Arc<rs::DistributedRuntime> = Arc::new(self.inner.clone());
        let weak: Weak<rs::DistributedRuntime> = Arc::downgrade(&arc);

        let name = CString::new("dynamo.runtime.weak").expect("valid capsule name");

        PyCapsule::new(py, weak, Some(name))
    }
659
660
}

661
662
663
664
665
666
667
668
669
670
671
672
673
674
// Bind a TCP port and return a socket held until dropped.
fn bind_tcp_port(port: u16) -> std::io::Result<socket2::Socket> {
    let sock = socket2::Socket::new(
        socket2::Domain::IPV4,
        socket2::Type::STREAM,
        Some(socket2::Protocol::TCP),
    )?;
    sock.set_reuse_address(true)?;
    let addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, port));
    sock.bind(&addr.into())?;
    Ok(sock)
}

fn make_port_key(namespace: &str, node_ip: IpAddr, port: u16) -> anyhow::Result<String> {
675
    Ok(format!("v1/{namespace}/ports/{node_ip}/{port}"))
676
677
678
679
680
681
682
683
684
685
686
687
}

fn local_ip() -> Result<IpAddr, local_ip_address::Error> {
    local_ip_address::local_ip().or_else(|err| match err {
        local_ip_address::Error::LocalIpAddressNotFound => {
            // Fall back to IPv6 if no IPv4 addresses are found
            local_ip_address::local_ipv6()
        }
        _ => Err(err),
    })
}

688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
#[pymethods]
impl CancellationToken {
    fn cancel(&self) {
        self.inner.cancel();
    }

    fn cancelled<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
        let token = self.inner.clone();
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            token.cancelled().await;
            Ok(())
        })
    }
}

#[pymethods]
impl Component {
    fn endpoint(&self, name: String) -> PyResult<Endpoint> {
        let inner = self.inner.endpoint(name);
        Ok(Endpoint {
            inner,
            event_loop: self.event_loop.clone(),
        })
    }

713
    /// NATS specific stats/metrics call
714
    fn create_service<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
715
        let mut inner = self.inner.clone();
716
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
717
            inner.add_stats_service().await.map_err(to_pyerr)?;
718
719
720
            Ok(())
        })
    }
721
722
723

    /// Get a RuntimeMetrics helper for creating Prometheus metrics
    #[getter]
724
725
    fn metrics(&self) -> prometheus_metrics::RuntimeMetrics {
        prometheus_metrics::RuntimeMetrics::from_component(self.inner.clone())
726
    }
727
728
729
730
}

#[pymethods]
impl Endpoint {
731
    #[pyo3(signature = (generator, graceful_shutdown = true, metrics_labels = None, health_check_payload = None))]
732
733
734
735
    fn serve_endpoint<'p>(
        &self,
        py: Python<'p>,
        generator: PyObject,
736
        graceful_shutdown: Option<bool>,
737
        metrics_labels: Option<Vec<(String, String)>>,
738
        health_check_payload: Option<&Bound<'p, PyDict>>,
739
740
741
742
743
744
    ) -> PyResult<Bound<'p, PyAny>> {
        let engine = Arc::new(engine::PythonAsyncEngine::new(
            generator,
            self.event_loop.clone(),
        )?);
        let ingress = JsonServerStreamingIngress::for_engine(engine).map_err(to_pyerr)?;
745
746
747
748
749
750
751
752
753
754
755
756
757

        // Convert Python dict to serde_json::Value if provided and validate it's an object
        let health_payload_json = health_check_payload
            .map(|dict| pythonize::depythonize::<serde_json::Value>(dict))
            .transpose()
            .map_err(|err| {
                pyo3::exceptions::PyTypeError::new_err(format!(
                    "Failed to convert health_check_payload: {}",
                    err
                ))
            })?;

        // Require an object/dict
758
759
760
761
762
763
        if let Some(ref payload) = health_payload_json
            && !payload.is_object()
        {
            return Err(pyo3::exceptions::PyTypeError::new_err(
                "health_check_payload must be a JSON object (dict)",
            ));
764
765
766
        }

        let mut builder = self
767
768
769
770
            .inner
            .endpoint_builder()
            .metrics_labels(metrics_labels)
            .handler(ingress);
771
772
773
774
775

        if let Some(payload) = health_payload_json {
            builder = builder.health_check_payload(payload);
        }

776
        let graceful_shutdown = graceful_shutdown.unwrap_or(true);
777
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
778
779
780
781
782
            builder
                .graceful_shutdown(graceful_shutdown)
                .start()
                .await
                .map_err(to_pyerr)?;
783
784
785
786
            Ok(())
        })
    }

787
788
789
    fn client<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
790
            let client = inner.client().await.map_err(to_pyerr)?;
791
792
793
794
795
796
            let push_router = rs::pipeline::PushRouter::<
                serde_json::Value,
                RsAnnotated<serde_json::Value>,
            >::from_client(client, Default::default())
            .await
            .map_err(to_pyerr)?;
797
798
799
            Ok(Client {
                router: push_router,
            })
800
801
        })
    }
802

803
804
805
    // Opaque unique ID for this worker. May change over worker lifetime.
    fn connection_id(&self) -> u64 {
        self.inner.drt().connection_id()
806
    }
807
808
809

    /// Get a RuntimeMetrics helper for creating Prometheus metrics
    #[getter]
810
811
    fn metrics(&self) -> prometheus_metrics::RuntimeMetrics {
        prometheus_metrics::RuntimeMetrics::from_endpoint(self.inner.clone())
812
    }
813
814
815
816
817
818
819
820
821
822
823
}

#[pymethods]
impl Namespace {
    fn component(&self, name: String) -> PyResult<Component> {
        let inner = self.inner.component(name).map_err(to_pyerr)?;
        Ok(Component {
            inner,
            event_loop: self.event_loop.clone(),
        })
    }
824
825
826

    /// Get a RuntimeMetrics helper for creating Prometheus metrics
    #[getter]
827
828
    fn metrics(&self) -> prometheus_metrics::RuntimeMetrics {
        prometheus_metrics::RuntimeMetrics::from_namespace(self.inner.clone())
829
    }
830
831
832
833
}

#[pymethods]
impl Client {
834
835
    /// Get list of current instances.
    /// Replaces endpoint_ids.
836
    fn instance_ids(&self) -> Vec<u64> {
837
        self.router.client.instance_ids()
838
839
    }

840
841
842
    /// Wait for an instance to be available for work.
    /// Replaces wait_for_endpoints.
    fn wait_for_instances<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
843
        let inner = self.router.client.clone();
844
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
845
            inner
846
                .wait_for_instances()
847
                .await
848
                .map(|v| v.into_iter().map(|cei| cei.id()).collect::<Vec<u64>>())
849
                .map_err(to_pyerr)
850
851
852
853
        })
    }

    /// Issue a request to the endpoint using the default routing strategy.
854
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
855
856
857
858
859
    fn generate<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
860
        context: Option<context::Context>,
861
    ) -> PyResult<Bound<'p, PyAny>> {
862
        self.random(py, request, annotated, context)
863
864
865
    }

    /// Send a request to the next endpoint in a round-robin fashion.
866
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
867
868
869
870
871
    fn round_robin<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
872
        context: Option<context::Context>,
873
874
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
875
        let request_ctx = create_request_context(request, &context);
876
877
878
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
879
        let client = self.router.clone();
880
881

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
882
883
            let stream = match context {
                Some(context) => {
884
885
                    // Always instrument with appropriate span (none if no trace context)
                    let span = get_span_for_context(&context, "round_robin");
886
887
888
889
890
                    client
                        .round_robin(request_ctx)
                        .instrument(span)
                        .await
                        .map_err(to_pyerr)?
891
                }
892
                _ => client.round_robin(request_ctx).await.map_err(to_pyerr)?,
893
            };
894
895
896
897
898
899
900
901
902
            tokio::spawn(process_stream(stream, tx));
            Ok(AsyncResponseStream {
                rx: Arc::new(Mutex::new(rx)),
                annotated,
            })
        })
    }

    /// Send a request to a random endpoint.
903
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
904
905
906
907
908
    fn random<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
909
        context: Option<context::Context>,
910
911
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
912
        let request_ctx = create_request_context(request, &context);
913
914
915
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
916
        let client = self.router.clone();
917
918

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
919
920
            let stream = match context {
                Some(context) => {
921
                    // Always instrument with appropriate span (none if no trace context)
922
                    let span = get_span_for_context(&context, "random");
923
924
925
926
927
                    client
                        .random(request_ctx)
                        .instrument(span)
                        .await
                        .map_err(to_pyerr)?
928
                }
929
                _ => client.random(request_ctx).await.map_err(to_pyerr)?,
930
            };
931
932
933
934
935
936
937
938
939
            tokio::spawn(process_stream(stream, tx));
            Ok(AsyncResponseStream {
                rx: Arc::new(Mutex::new(rx)),
                annotated,
            })
        })
    }

    /// Directly send a request to a specific endpoint.
940
    #[pyo3(signature = (request, instance_id, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
941
942
943
944
    fn direct<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
945
        instance_id: u64,
946
        annotated: Option<bool>,
947
        context: Option<context::Context>,
948
949
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
950
        let request_ctx = create_request_context(request, &context);
951
952
953
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
954
        let client = self.router.clone();
955
956

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
957
958
            let stream = match context {
                Some(context) => {
959
                    // Always instrument with appropriate span (none if no trace context)
960
961
                    let span =
                        get_span_for_direct_context(&context, "direct", &instance_id.to_string());
962
963
964
965
966
                    client
                        .direct(request_ctx, instance_id)
                        .instrument(span)
                        .await
                        .map_err(to_pyerr)?
967
968
                }
                _ => client
969
                    .direct(request_ctx, instance_id)
970
971
972
                    .await
                    .map_err(to_pyerr)?,
            };
973
974
975
976
977
978
979
980
981
982
983
984

            tokio::spawn(process_stream(stream, tx));

            Ok(AsyncResponseStream {
                rx: Arc::new(Mutex::new(rx)),
                annotated,
            })
        })
    }
}

async fn process_stream(
985
    stream: EngineStream<RsAnnotated<serde_json::Value>>,
986
987
988
989
990
    tx: tokio::sync::mpsc::Sender<RsAnnotated<PyObject>>,
) {
    let mut stream = stream;
    while let Some(response) = stream.next().await {
        // Convert the response to a PyObject using Python's GIL
991
        let annotated: RsAnnotated<serde_json::Value> = response;
992
        let annotated: RsAnnotated<PyObject> = annotated.map_data(|data| {
993
            Python::with_gil(|py| match pythonize::pythonize(py, &data) {
994
995
                Ok(pyobj) => Ok(pyobj.into()),
                Err(e) => Err(e.to_string()),
996
            })
997
998
999
1000
1001
1002
        });

        let is_error = annotated.is_error();

        // Send the PyObject through the channel or log an error
        if let Err(e) = tx.send(annotated).await {
1003
            tracing::error!("Failed to send response: {:?}", e);
1004
            break;
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
        }

        if is_error {
            break;
        }
    }
}

#[pyclass]
struct AsyncResponseStream {
    rx: Arc<Mutex<tokio::sync::mpsc::Receiver<RsAnnotated<PyObject>>>>,
    annotated: bool,
}

#[pymethods]
impl AsyncResponseStream {
    /// This method is required to implement the `AsyncIterator` protocol.
    #[pyo3(name = "__aiter__")]
    fn aiter(slf: PyRef<Self>, py: Python) -> PyResult<Py<PyAny>> {
        slf.into_py_any(py)
    }
    /// This method is required to implement the `AsyncIterator` protocol.
    #[pyo3(name = "__anext__")]
    fn next<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
        let rx = self.rx.clone();
        let annotated = self.annotated;

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            loop {
                let value = rx.lock().await.recv().await;
                match value {
                    Some(pyobj) => {
                        let pyobj = match pyobj.ok() {
                            Ok(pyobj) => pyobj,
                            Err(e) => {
                                return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(e));
                            }
                        };

                        if annotated {
                            let object = Annotated { inner: pyobj };
                            #[allow(deprecated)]
                            let object = Python::with_gil(|py| object.into_py(py));
                            return Ok(object);
                        } else {
                            match pyobj.data {
                                Some(data) => return Ok(data),
                                None => continue,
                            }
                        }
                    }
                    None => return Err(PyStopAsyncIteration::new_err("Stream exhausted")),
                }
            }
        })
    }
}

#[pyclass]
struct Annotated {
    inner: RsAnnotated<PyObject>,
}

#[pymethods]
impl Annotated {
    #[new]
    fn new(data: PyObject) -> Self {
        Annotated {
            inner: RsAnnotated::from_data(data),
        }
    }

    fn is_error(&self) -> bool {
        self.inner.is_error()
    }

    fn data(&self) -> Option<PyObject> {
        self.inner.data.clone()
    }

    fn event(&self) -> Option<String> {
        self.inner.event.clone()
    }

    fn comments(&self) -> Option<Vec<String>> {
        self.inner.comment.clone()
    }

    fn id(&self) -> Option<String> {
        self.inner.id.clone()
    }

    #[pyo3(name = "__repr__")]
    fn _repr(&self, py: Python) -> String {
        let data = self.inner.data.clone().map(|obj| {
            obj.call_method0(py, "__repr__")
                .and_then(|repr_obj| repr_obj.extract::<Py<PyString>>(py))
                .map(|py_str| py_str.to_string_lossy(py).into_owned())
                .unwrap_or_else(|_| "<failed_repr>".to_string())
        });

        format!(
            "Annotated(data={}, event={}, comment={:?}, id={})",
            data.unwrap_or_else(|| "<no_data>".to_string()),
            self.inner.event.as_deref().unwrap_or("None"),
            self.inner.comment.as_deref().unwrap_or(&[]),
            self.inner.id.as_deref().unwrap_or("None")
        )
    }
}