lib.rs 31.7 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
use dynamo_runtime::distributed::{DistributedConfig, RequestPlaneMode};
6
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
14
use pyo3::{exceptions::PyException, prelude::*};
use rs::pipeline::network::Ingress;
Richard Huo's avatar
Richard Huo committed
15
use std::ffi::CString;
16
use std::fs;
17
use std::path::PathBuf;
Richard Huo's avatar
Richard Huo committed
18
19
20
21
use std::{
    fmt::Display,
    sync::{Arc, Weak},
};
22
use tokio::sync::Mutex;
23
use tracing::Instrument;
24

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

Neelay Shah's avatar
Neelay Shah committed
35
use dynamo_llm::{self as llm_rs};
36
37
use dynamo_llm::{entrypoint::RouterConfig, kv_router::KvRouterConfig};

38
use crate::llm::local_model::ModelRuntimeConfig;
39
use crate::llm::preprocessor::{MediaDecoder, MediaFetcher};
40

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

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

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

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

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

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

122
123
124
125
126
/// 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<()> {
127
    // Initialize logging early unless OTEL export is enabled (which requires tokio runtime)
128
    if rs::config::env_is_truthy("OTEL_EXPORT_ENABLED") {
129
        eprintln!(
130
            "Warning: OTEL_EXPORT_ENABLED detected. Logging initialization deferred until runtime is available. Early logs may be dropped."
131
132
133
134
135
        );
    } else {
        rs::logging::init();
    }

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

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

    engine::add_to_module(m)?;
190
    parsers::add_to_module(m)?;
191

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

197
198
199
200
201
202
203
204
205
206
    Ok(())
}

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

Ryan Olson's avatar
Ryan Olson committed
207
208
209
210
211
212
213
/// 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);
}

214
215
/// 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.
216
#[pyfunction]
217
#[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))]
218
#[allow(clippy::too_many_arguments)]
219
220
fn register_llm<'p>(
    py: Python<'p>,
221
    model_input: ModelInput,
222
    model_type: ModelType,
223
224
225
    endpoint: Endpoint,
    model_path: &str,
    model_name: Option<&str>,
226
227
    context_length: Option<u32>,
    kv_cache_block_size: Option<u32>,
228
    router_mode: Option<RouterMode>,
229
    migration_limit: u32,
230
    runtime_config: Option<ModelRuntimeConfig>,
231
    user_data: Option<&Bound<'p, PyDict>>,
232
    custom_template_path: Option<&str>,
233
234
    media_decoder: Option<MediaDecoder>,
    media_fetcher: Option<MediaFetcher>,
235
) -> PyResult<Bound<'p, PyAny>> {
236
237
238
239
240
241
242
243
244
245
246
247
248
249
    // 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",
            ));
        }
    }

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

256
257
    let model_type_obj = model_type.inner;

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

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

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

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

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

        Ok(())
    })
}

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

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

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

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

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

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

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

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

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

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

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

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

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

436
437
438
        // 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()
439
            .or_else(|_| -> anyhow::Result<rs::Runtime> {
440
441
442
443
                // No existing Worker, create new one
                let worker = rs::Worker::from_settings()?;

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

452
                Ok(worker.runtime().clone())
453
454
            })
            .map_err(to_pyerr)?;
455

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

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

        Ok(DistributedRuntime { inner, event_loop })
    }

Ryan Olson's avatar
Ryan Olson committed
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
    #[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(),
        })
    }

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

    fn shutdown(&self) {
500
        self.inner.shutdown();
501
502
503
504
505
    }

    fn event_loop(&self) -> PyObject {
        self.event_loop.clone()
    }
506
507
508
509
510

    fn child_token(&self) -> CancellationToken {
        let inner = self.inner.runtime().child_token();
        CancellationToken { inner }
    }
Richard Huo's avatar
Richard Huo committed
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525

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

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

553
    /// NATS specific stats/metrics call
554
    fn create_service<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
555
        let mut inner = self.inner.clone();
556
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
557
            inner.add_stats_service().await.map_err(to_pyerr)?;
558
559
560
            Ok(())
        })
    }
561
562
563

    /// Get a RuntimeMetrics helper for creating Prometheus metrics
    #[getter]
564
565
    fn metrics(&self) -> prometheus_metrics::RuntimeMetrics {
        prometheus_metrics::RuntimeMetrics::from_component(self.inner.clone())
566
    }
567
568
569
570
}

#[pymethods]
impl Endpoint {
571
    #[pyo3(signature = (generator, graceful_shutdown = true, metrics_labels = None, health_check_payload = None))]
572
573
574
575
    fn serve_endpoint<'p>(
        &self,
        py: Python<'p>,
        generator: PyObject,
576
        graceful_shutdown: Option<bool>,
577
        metrics_labels: Option<Vec<(String, String)>>,
578
        health_check_payload: Option<&Bound<'p, PyDict>>,
579
580
581
582
583
584
    ) -> 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)?;
585
586
587
588
589
590
591
592
593
594
595
596
597

        // 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
598
599
600
601
602
603
        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)",
            ));
604
605
606
        }

        let mut builder = self
607
608
609
610
            .inner
            .endpoint_builder()
            .metrics_labels(metrics_labels)
            .handler(ingress);
611
612
613
614
615

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

616
        let graceful_shutdown = graceful_shutdown.unwrap_or(true);
617
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
618
619
620
621
622
            builder
                .graceful_shutdown(graceful_shutdown)
                .start()
                .await
                .map_err(to_pyerr)?;
623
624
625
626
            Ok(())
        })
    }

627
628
629
    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 {
630
            let client = inner.client().await.map_err(to_pyerr)?;
631
632
633
634
635
636
            let push_router = rs::pipeline::PushRouter::<
                serde_json::Value,
                RsAnnotated<serde_json::Value>,
            >::from_client(client, Default::default())
            .await
            .map_err(to_pyerr)?;
637
638
639
            Ok(Client {
                router: push_router,
            })
640
641
        })
    }
642

643
644
645
    // Opaque unique ID for this worker. May change over worker lifetime.
    fn connection_id(&self) -> u64 {
        self.inner.drt().connection_id()
646
    }
647
648
649

    /// Get a RuntimeMetrics helper for creating Prometheus metrics
    #[getter]
650
651
    fn metrics(&self) -> prometheus_metrics::RuntimeMetrics {
        prometheus_metrics::RuntimeMetrics::from_endpoint(self.inner.clone())
652
    }
653
654
655
656
657
658
659
660
661
662
663
}

#[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(),
        })
    }
664
665
666

    /// Get a RuntimeMetrics helper for creating Prometheus metrics
    #[getter]
667
668
    fn metrics(&self) -> prometheus_metrics::RuntimeMetrics {
        prometheus_metrics::RuntimeMetrics::from_namespace(self.inner.clone())
669
    }
670
671
672
673
}

#[pymethods]
impl Client {
674
675
    /// Get list of current instances.
    /// Replaces endpoint_ids.
676
    fn instance_ids(&self) -> Vec<u64> {
677
        self.router.client.instance_ids()
678
679
    }

680
681
682
    /// 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>> {
683
        let inner = self.router.client.clone();
684
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
685
            inner
686
                .wait_for_instances()
687
                .await
688
                .map(|v| v.into_iter().map(|cei| cei.id()).collect::<Vec<u64>>())
689
                .map_err(to_pyerr)
690
691
692
693
        })
    }

    /// Issue a request to the endpoint using the default routing strategy.
694
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
695
696
697
698
699
    fn generate<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
700
        context: Option<context::Context>,
701
    ) -> PyResult<Bound<'p, PyAny>> {
702
        self.random(py, request, annotated, context)
703
704
705
    }

    /// Send a request to the next endpoint in a round-robin fashion.
706
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
707
708
709
710
711
    fn round_robin<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
712
        context: Option<context::Context>,
713
714
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
715
        let request_ctx = create_request_context(request, &context);
716
717
718
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
719
        let client = self.router.clone();
720
721

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
722
723
            let stream = match context {
                Some(context) => {
724
725
                    // Always instrument with appropriate span (none if no trace context)
                    let span = get_span_for_context(&context, "round_robin");
726
727
728
729
730
                    client
                        .round_robin(request_ctx)
                        .instrument(span)
                        .await
                        .map_err(to_pyerr)?
731
                }
732
                _ => client.round_robin(request_ctx).await.map_err(to_pyerr)?,
733
            };
734
735
736
737
738
739
740
741
742
            tokio::spawn(process_stream(stream, tx));
            Ok(AsyncResponseStream {
                rx: Arc::new(Mutex::new(rx)),
                annotated,
            })
        })
    }

    /// Send a request to a random endpoint.
743
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
744
745
746
747
748
    fn random<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
749
        context: Option<context::Context>,
750
751
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
752
        let request_ctx = create_request_context(request, &context);
753
754
755
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
756
        let client = self.router.clone();
757
758

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
759
760
            let stream = match context {
                Some(context) => {
761
                    // Always instrument with appropriate span (none if no trace context)
762
                    let span = get_span_for_context(&context, "random");
763
764
765
766
767
                    client
                        .random(request_ctx)
                        .instrument(span)
                        .await
                        .map_err(to_pyerr)?
768
                }
769
                _ => client.random(request_ctx).await.map_err(to_pyerr)?,
770
            };
771
772
773
774
775
776
777
778
779
            tokio::spawn(process_stream(stream, tx));
            Ok(AsyncResponseStream {
                rx: Arc::new(Mutex::new(rx)),
                annotated,
            })
        })
    }

    /// Directly send a request to a specific endpoint.
780
    #[pyo3(signature = (request, instance_id, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
781
782
783
784
    fn direct<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
785
        instance_id: u64,
786
        annotated: Option<bool>,
787
        context: Option<context::Context>,
788
789
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
790
        let request_ctx = create_request_context(request, &context);
791
792
793
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
794
        let client = self.router.clone();
795
796

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
797
798
            let stream = match context {
                Some(context) => {
799
                    // Always instrument with appropriate span (none if no trace context)
800
801
                    let span =
                        get_span_for_direct_context(&context, "direct", &instance_id.to_string());
802
803
804
805
806
                    client
                        .direct(request_ctx, instance_id)
                        .instrument(span)
                        .await
                        .map_err(to_pyerr)?
807
808
                }
                _ => client
809
                    .direct(request_ctx, instance_id)
810
811
812
                    .await
                    .map_err(to_pyerr)?,
            };
813
814
815
816
817
818
819
820
821
822
823
824

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

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

async fn process_stream(
825
    stream: EngineStream<RsAnnotated<serde_json::Value>>,
826
827
828
829
830
    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
831
        let annotated: RsAnnotated<serde_json::Value> = response;
832
        let annotated: RsAnnotated<PyObject> = annotated.map_data(|data| {
833
            Python::with_gil(|py| match pythonize::pythonize(py, &data) {
834
835
                Ok(pyobj) => Ok(pyobj.into()),
                Err(e) => Err(e.to_string()),
836
            })
837
838
839
840
841
842
        });

        let is_error = annotated.is_error();

        // Send the PyObject through the channel or log an error
        if let Err(e) = tx.send(annotated).await {
843
            tracing::error!("Failed to send response: {:?}", e);
844
            break;
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
        }

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