lib.rs 39.4 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
42
use crate::llm::local_model::ModelRuntimeConfig;

43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
#[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,
        }
    }
}
60

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

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);

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

// 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 {
93
94
95
96
97
98
    logging::make_client_request_span(
        operation,
        context.inner().id(),
        context.trace_context(),
        Some(instance_id),
    )
99
100
}

101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
// 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(),
    }
}

124
125
126
127
128
/// 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<()> {
129
130
131
132
133
134
135
136
137
138
139
140
    // Initialize logging early unless OTEL export is enabled (which requires tokio runtime)
    if std::env::var("OTEL_EXPORT_ENABLED")
        .map(|v| v == "1")
        .unwrap_or(false)
    {
        eprintln!(
            "Warning: OTEL_EXPORT_ENABLED=1 detected. Logging initialization deferred until runtime is available. Early logs may be dropped."
        );
    } else {
        rs::logging::init();
    }

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

148
149
150
151
152
153
154
    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>()?;
155
156
157
    m.add_class::<llm::entrypoint::EntrypointArgs>()?;
    m.add_class::<llm::entrypoint::EngineConfig>()?;
    m.add_class::<llm::entrypoint::EngineType>()?;
158
159
    m.add_class::<llm::entrypoint::RouterConfig>()?;
    m.add_class::<llm::entrypoint::KvRouterConfig>()?;
160
    m.add_class::<llm::kv::WorkerMetricsPublisher>()?;
161
    m.add_class::<llm::model_card::ModelDeploymentCard>()?;
162
    m.add_class::<llm::local_model::ModelRuntimeConfig>()?;
163
164
    m.add_class::<llm::preprocessor::OAIChatPreprocessor>()?;
    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))]
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
) -> PyResult<Bound<'p, PyAny>> {
237
238
239
240
241
242
243
244
245
246
247
248
249
250
    // 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",
            ));
        }
    }

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

257
258
    let model_type_obj = model_type.inner;

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

264
265
266
267
268
269
270
271
272
273
274
275
276
    // 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()?;

277
278
279
280
281
282
283
    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))
        })?;

284
    pyo3_async_runtimes::tokio::future_into_py(py, async move {
285
286
287
288
289
290
291
292
293
294
295
296
297
        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)?
        };

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

        Ok(())
    })
}

321
322
323
324
325
326
327
328
329
330
331
/// 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)
    })
}

332
333
#[pyclass]
#[derive(Clone)]
Ryan Olson's avatar
Ryan Olson committed
334
pub struct DistributedRuntime {
335
336
337
338
    inner: rs::DistributedRuntime,
    event_loop: PyObject,
}

Ryan Olson's avatar
Ryan Olson committed
339
340
impl DistributedRuntime {
    #[allow(dead_code)]
341
    pub(crate) fn inner(&self) -> &rs::DistributedRuntime {
Ryan Olson's avatar
Ryan Olson committed
342
343
344
345
        &self.inner
    }
}

346
#[pyclass]
347
#[derive(Clone)]
348
349
350
351
352
struct CancellationToken {
    inner: rs::CancellationToken,
}

#[pyclass]
353
#[derive(Clone)]
354
355
356
357
358
359
struct Namespace {
    inner: rs::component::Namespace,
    event_loop: PyObject,
}

#[pyclass]
360
#[derive(Clone)]
361
362
363
364
365
366
struct Component {
    inner: rs::component::Component,
    event_loop: PyObject,
}

#[pyclass]
367
#[derive(Clone)]
368
369
370
371
372
373
struct Endpoint {
    inner: rs::component::Endpoint,
    event_loop: PyObject,
}

#[pyclass]
374
#[derive(Clone)]
375
struct Client {
376
    router: rs::pipeline::PushRouter<serde_json::Value, RsAnnotated<serde_json::Value>>,
377
378
}

379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
#[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,
    };
400
401
402
403
    #[classattr]
    const TensorBased: Self = ModelType {
        inner: llm_rs::model_type::ModelType::TensorBased,
    };
404
405
406
407
    #[classattr]
    const Prefill: Self = ModelType {
        inner: llm_rs::model_type::ModelType::Prefill,
    };
408
409
410
411
412
413
414
415
416
417
418
419

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

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

420
421
#[pyclass(eq, eq_int)]
#[derive(Clone, PartialEq)]
422
423
424
enum ModelInput {
    Text = 1,
    Tokens = 2,
425
    Tensor = 3,
426
427
}

428
429
430
#[pymethods]
impl DistributedRuntime {
    #[new]
431
432
433
    fn new(event_loop: PyObject, store_kv: String, is_static: bool) -> PyResult<Self> {
        let selected_kv_store: KeyValueStoreSelect = store_kv.parse().map_err(to_pyerr)?;

434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
        // 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()
            .or_else(|_| {
                // No existing Worker, create new one
                let worker = rs::Worker::from_settings()?;

                // Initialize pyo3 bridge (only happens once per process)
                INIT.get_or_try_init(|| {
                    let primary = worker.tokio_runtime()?;
                    pyo3_async_runtimes::tokio::init_with_runtime(primary).map_err(|e| {
                        rs::error!("failed to initialize pyo3 static runtime: {:?}", e)
                    })?;
                    rs::OK(())
                })?;

                rs::OK(worker.runtime().clone())
            })
            .map_err(to_pyerr)?;
453

454
455
        // Initialize logging in context where tokio runtime is available
        // otel exporter requires it
456
457
458
459
460
461
462
463
        if std::env::var("OTEL_EXPORT_ENABLED")
            .map(|v| v == "1")
            .unwrap_or(false)
        {
            runtime.secondary().block_on(async {
                rs::logging::init();
            });
        }
464

465
466
467
468
469
470
        let inner =
            if is_static {
                runtime.secondary().block_on(
                    rs::DistributedRuntime::from_settings_without_discovery(runtime),
                )
            } else {
471
472
473
474
475
                let config = DistributedConfig {
                    store_backend: selected_kv_store,
                    is_static: false,
                    nats_config: dynamo_runtime::transports::nats::ClientOptions::default(),
                };
476
477
                runtime
                    .secondary()
478
                    .block_on(rs::DistributedRuntime::new(runtime, config))
479
480
            };
        let inner = inner.map_err(to_pyerr)?;
481
482
483
484

        Ok(DistributedRuntime { inner, event_loop })
    }

Ryan Olson's avatar
Ryan Olson committed
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
    #[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(),
        })
    }

500
501
502
503
504
505
506
    fn namespace(&self, name: String) -> PyResult<Namespace> {
        Ok(Namespace {
            inner: self.inner.namespace(name).map_err(to_pyerr)?,
            event_loop: self.event_loop.clone(),
        })
    }

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
569
570
571
572
573
574
575
    /// 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!(
576
577
578
                                "Failed to bind to port block starting at {start_port} (attempt {}): {e}",
                                attempt_idx + 1,
                            );
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
627
628
629
630
631
632
633
                            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!(
634
635
                "Failed to allocate and reserve a port block of size {block_size} from range {min}-{max} after {candidate_count} attempts"
            )))
636
637
638
        })
    }

639
    fn shutdown(&self) {
640
        self.inner.shutdown();
641
642
643
644
645
    }

    fn event_loop(&self) -> PyObject {
        self.event_loop.clone()
    }
646
647
648
649
650

    fn child_token(&self) -> CancellationToken {
        let inner = self.inner.runtime().child_token();
        CancellationToken { inner }
    }
Richard Huo's avatar
Richard Huo committed
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665

    // 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))
    }
666
667
}

668
669
670
671
672
673
674
675
676
677
678
679
680
681
// 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> {
682
    Ok(format!("v1/{namespace}/ports/{node_ip}/{port}"))
683
684
685
686
687
688
689
690
691
692
693
694
}

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

695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
#[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(),
        })
    }

720
    /// NATS specific stats/metrics call
721
    fn create_service<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
722
        let mut inner = self.inner.clone();
723
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
724
            inner.add_stats_service().await.map_err(to_pyerr)?;
725
726
727
            Ok(())
        })
    }
728
729
730

    /// Get a RuntimeMetrics helper for creating Prometheus metrics
    #[getter]
731
732
    fn metrics(&self) -> prometheus_metrics::RuntimeMetrics {
        prometheus_metrics::RuntimeMetrics::from_component(self.inner.clone())
733
    }
734
735
736
737
}

#[pymethods]
impl Endpoint {
738
    #[pyo3(signature = (generator, graceful_shutdown = true, metrics_labels = None, health_check_payload = None))]
739
740
741
742
    fn serve_endpoint<'p>(
        &self,
        py: Python<'p>,
        generator: PyObject,
743
        graceful_shutdown: Option<bool>,
744
        metrics_labels: Option<Vec<(String, String)>>,
745
        health_check_payload: Option<&Bound<'p, PyDict>>,
746
747
748
749
750
751
    ) -> 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)?;
752
753
754
755
756
757
758
759
760
761
762
763
764

        // 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
765
766
767
768
769
770
        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)",
            ));
771
772
773
        }

        let mut builder = self
774
775
776
777
            .inner
            .endpoint_builder()
            .metrics_labels(metrics_labels)
            .handler(ingress);
778
779
780
781
782

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

783
        let graceful_shutdown = graceful_shutdown.unwrap_or(true);
784
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
785
786
787
788
789
            builder
                .graceful_shutdown(graceful_shutdown)
                .start()
                .await
                .map_err(to_pyerr)?;
790
791
792
793
            Ok(())
        })
    }

794
795
796
    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 {
797
            let client = inner.client().await.map_err(to_pyerr)?;
798
799
800
801
802
803
            let push_router = rs::pipeline::PushRouter::<
                serde_json::Value,
                RsAnnotated<serde_json::Value>,
            >::from_client(client, Default::default())
            .await
            .map_err(to_pyerr)?;
804
805
806
            Ok(Client {
                router: push_router,
            })
807
808
        })
    }
809

810
811
812
    // Opaque unique ID for this worker. May change over worker lifetime.
    fn connection_id(&self) -> u64 {
        self.inner.drt().connection_id()
813
    }
814
815
816

    /// Get a RuntimeMetrics helper for creating Prometheus metrics
    #[getter]
817
818
    fn metrics(&self) -> prometheus_metrics::RuntimeMetrics {
        prometheus_metrics::RuntimeMetrics::from_endpoint(self.inner.clone())
819
    }
820
821
822
823
824
825
826
827
828
829
830
}

#[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(),
        })
    }
831
832
833

    /// Get a RuntimeMetrics helper for creating Prometheus metrics
    #[getter]
834
835
    fn metrics(&self) -> prometheus_metrics::RuntimeMetrics {
        prometheus_metrics::RuntimeMetrics::from_namespace(self.inner.clone())
836
    }
837
838
839
840
}

#[pymethods]
impl Client {
841
842
    /// Get list of current instances.
    /// Replaces endpoint_ids.
843
    fn instance_ids(&self) -> Vec<u64> {
844
        self.router.client.instance_ids()
845
846
    }

847
848
849
    /// 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>> {
850
        let inner = self.router.client.clone();
851
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
852
            inner
853
                .wait_for_instances()
854
                .await
855
                .map(|v| v.into_iter().map(|cei| cei.id()).collect::<Vec<u64>>())
856
                .map_err(to_pyerr)
857
858
859
860
        })
    }

    /// Issue a request to the endpoint using the default routing strategy.
861
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
862
863
864
865
866
    fn generate<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
867
        context: Option<context::Context>,
868
    ) -> PyResult<Bound<'p, PyAny>> {
869
        if self.router.client.is_static() {
870
            self.r#static(py, request, annotated, context)
871
        } else {
872
            self.random(py, request, annotated, context)
873
        }
874
875
876
    }

    /// Send a request to the next endpoint in a round-robin fashion.
877
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
878
879
880
881
882
    fn round_robin<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
883
        context: Option<context::Context>,
884
885
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
886
        let request_ctx = create_request_context(request, &context);
887
888
889
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
890
        let client = self.router.clone();
891
892

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
893
894
            let stream = match context {
                Some(context) => {
895
896
                    // Always instrument with appropriate span (none if no trace context)
                    let span = get_span_for_context(&context, "round_robin");
897
898
899
900
901
                    client
                        .round_robin(request_ctx)
                        .instrument(span)
                        .await
                        .map_err(to_pyerr)?
902
                }
903
                _ => client.round_robin(request_ctx).await.map_err(to_pyerr)?,
904
            };
905
906
907
908
909
910
911
912
913
            tokio::spawn(process_stream(stream, tx));
            Ok(AsyncResponseStream {
                rx: Arc::new(Mutex::new(rx)),
                annotated,
            })
        })
    }

    /// Send a request to a random endpoint.
914
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
915
916
917
918
919
    fn random<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
920
        context: Option<context::Context>,
921
922
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
923
        let request_ctx = create_request_context(request, &context);
924
925
926
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
927
        let client = self.router.clone();
928
929

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
930
931
            let stream = match context {
                Some(context) => {
932
                    // Always instrument with appropriate span (none if no trace context)
933
                    let span = get_span_for_context(&context, "random");
934
935
936
937
938
                    client
                        .random(request_ctx)
                        .instrument(span)
                        .await
                        .map_err(to_pyerr)?
939
                }
940
                _ => client.random(request_ctx).await.map_err(to_pyerr)?,
941
            };
942
943
944
945
946
947
948
949
950
            tokio::spawn(process_stream(stream, tx));
            Ok(AsyncResponseStream {
                rx: Arc::new(Mutex::new(rx)),
                annotated,
            })
        })
    }

    /// Directly send a request to a specific endpoint.
951
    #[pyo3(signature = (request, instance_id, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
952
953
954
955
    fn direct<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
956
        instance_id: u64,
957
        annotated: Option<bool>,
958
        context: Option<context::Context>,
959
960
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
961
        let request_ctx = create_request_context(request, &context);
962
963
964
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
965
        let client = self.router.clone();
966
967

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
968
969
            let stream = match context {
                Some(context) => {
970
                    // Always instrument with appropriate span (none if no trace context)
971
972
                    let span =
                        get_span_for_direct_context(&context, "direct", &instance_id.to_string());
973
974
975
976
977
                    client
                        .direct(request_ctx, instance_id)
                        .instrument(span)
                        .await
                        .map_err(to_pyerr)?
978
979
                }
                _ => client
980
                    .direct(request_ctx, instance_id)
981
982
983
                    .await
                    .map_err(to_pyerr)?,
            };
984
985
986
987
988
989
990
991
992

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

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

    /// Directly send a request to a pre-defined static worker
995
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
996
997
998
999
1000
    fn r#static<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
1001
        context: Option<context::Context>,
1002
1003
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
1004
        let request_ctx = create_request_context(request, &context);
1005
1006
1007
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
1008
        let client = self.router.clone();
1009
1010

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
1011
1012
            let stream = match context {
                Some(context) => {
1013
                    // Always instrument with appropriate span (none if no trace context)
1014
                    let span = get_span_for_context(&context, "static");
1015
1016
1017
1018
1019
                    client
                        .r#static(request_ctx)
                        .instrument(span)
                        .await
                        .map_err(to_pyerr)?
1020
                }
1021
                _ => client.r#static(request_ctx).await.map_err(to_pyerr)?,
1022
            };
1023
1024
1025
1026
1027
1028
1029
1030
1031

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

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

async fn process_stream(
1035
    stream: EngineStream<RsAnnotated<serde_json::Value>>,
1036
1037
1038
1039
1040
    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
1041
        let annotated: RsAnnotated<serde_json::Value> = response;
1042
        let annotated: RsAnnotated<PyObject> = annotated.map_data(|data| {
1043
            Python::with_gil(|py| match pythonize::pythonize(py, &data) {
1044
1045
                Ok(pyobj) => Ok(pyobj.into()),
                Err(e) => Err(e.to_string()),
1046
            })
1047
1048
1049
1050
1051
1052
        });

        let is_error = annotated.is_error();

        // Send the PyObject through the channel or log an error
        if let Err(e) = tx.send(annotated).await {
1053
            tracing::error!("Failed to send response: {:?}", e);
1054
            break;
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
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
        }

        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")
        )
    }
}