lib.rs 39.8 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
436
437
    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)?;

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
        let inner =
            if is_static {
                runtime.secondary().block_on(
                    rs::DistributedRuntime::from_settings_without_discovery(runtime),
                )
            } else {
472
473
474
475
476
                let config = DistributedConfig {
                    store_backend: selected_kv_store,
                    is_static: false,
                    nats_config: dynamo_runtime::transports::nats::ClientOptions::default(),
                };
477
478
                runtime
                    .secondary()
479
                    .block_on(rs::DistributedRuntime::new(runtime, config))
480
481
            };
        let inner = inner.map_err(to_pyerr)?;
482
483
484
485

        Ok(DistributedRuntime { inner, event_loop })
    }

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

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

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
576
    /// 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!(
577
578
579
                                "Failed to bind to port block starting at {start_port} (attempt {}): {e}",
                                attempt_idx + 1,
                            );
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
634
                            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!(
635
636
                "Failed to allocate and reserve a port block of size {block_size} from range {min}-{max} after {candidate_count} attempts"
            )))
637
638
639
        })
    }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        let is_error = annotated.is_error();

        // Send the PyObject through the channel or log an error
        if let Err(e) = tx.send(annotated).await {
1054
            tracing::error!("Failed to send response: {:?}", e);
1055
            break;
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
1165
        }

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