fix(bindings): drain bridge tasks before interpreter finalization

lib/bindings/python/rust/backend.rs

@@ -591,7 +591,7 @@ impl Worker {
         // to the worker's DistributedRuntime, not the process-wide executor, and
         // are resolved directly by RsWorker without changing environment vars.
         let primary = rs::Worker::ensure_process_runtime().map_err(to_pyerr)?;
-        let _ = pyo3_async_runtimes::tokio::init_with_runtime(primary);
+        crate::adopt_bridge_runtime(primary);
 
         Ok(Self {
             engine: Arc::new(engine),
@@ -625,7 +625,7 @@ impl Worker {
         let config = self.config.clone();
         let raw = self.raw;
 
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             // No fallback: `runtime_from_existing` creates the process runtime when there isn't
             // one, so it only fails when the settings themselves are bad. Retrying through
             // `Worker::from_settings` would read those same settings and fail the same way.

lib/bindings/python/rust/context.rs

@@ -312,7 +312,7 @@ impl Context {
     fn async_killed_or_stopped<'a>(&self, py: Python<'a>) -> PyResult<Bound<'a, PyAny>> {
         let inner = self.inner.clone();
 
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             tokio::select! {
                 _ = inner.killed() => {
                     Ok(true)

lib/bindings/python/rust/http.rs

@@ -121,7 +121,7 @@ impl HttpService {
         // multi-minute streams are still draining.
         let guard = runtime.inner().register_graceful_task();
 
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let _guard = guard;
             service.run(token).await.map_err(to_pyerr)?;
             Ok(())

lib/bindings/python/rust/kserve_grpc.rs

@@ -166,7 +166,7 @@ impl KserveGrpcService {
                 )
             })?;
 
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             service.run(token).await.map_err(to_pyerr)?;
             Ok(())
         })

lib/bindings/python/rust/llm/entrypoint.rs

@@ -778,7 +778,7 @@ pub fn make_engine<'p>(
         .runtime_config(args.runtime_config.clone().inner)
         .namespace(args.namespace.clone())
         .namespace_prefix(args.namespace_prefix.clone());
-    pyo3_async_runtimes::tokio::future_into_py(py, async move {
+    crate::future_into_py(py, async move {
         if let Some(model_path) = args.model_path.clone() {
             let local_path = if model_path.exists() {
                 model_path
@@ -1045,7 +1045,7 @@ pub fn run_input<'p>(
             "linked worker-selection policies require HTTP frontend input",
         ));
     }
-    pyo3_async_runtimes::tokio::future_into_py(py, async move {
+    crate::future_into_py(py, async move {
         if let Some(factory) = worker_selection_policy_factory {
             HttpFrontend::default()
                 .frontend_route_extensions(frontend_route_extensions)

lib/bindings/python/rust/llm/kv.rs

@@ -673,7 +673,7 @@ impl SelectionService {
             builder = builder.replica_sync(port, replica_sync_peers);
         }
         let inner = py
-            .allow_threads(|| pyo3_async_runtimes::tokio::get_runtime().block_on(builder.build()))
+            .allow_threads(|| crate::bridge_runtime().block_on(builder.build()))
             .map_err(to_pyerr)?;
         Ok(Self {
             inner: Arc::new(inner),
@@ -687,13 +687,13 @@ impl SelectionService {
     /// dropped. Idempotent, and also run automatically on drop.
     fn shutdown(&self, py: Python<'_>) {
         let service = Arc::clone(&self.inner);
-        py.allow_threads(|| pyo3_async_runtimes::tokio::get_runtime().block_on(service.shutdown()));
+        py.allow_threads(|| crate::bridge_runtime().block_on(service.shutdown()));
     }
 
     /// Await service shutdown without blocking the Python event loop.
     fn shutdown_async<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
         let service = Arc::clone(&self.inner);
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             service.shutdown().await;
             Ok(())
         })
@@ -704,7 +704,7 @@ impl SelectionService {
         let req: WorkerRequest =
             depythonize(worker.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?;
         let core = self.inner.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let record = core.upsert_worker(req).await.map_err(selection_to_pyerr)?;
             Python::with_gil(|py| pythonize(py, &record).map(|o| o.unbind()).map_err(to_pyerr))
         })
@@ -720,7 +720,7 @@ impl SelectionService {
         let patch: WorkerPatchRequest =
             depythonize(patch.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?;
         let service = Arc::clone(&self.inner);
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let record = service
                 .patch_worker(worker_id, patch)
                 .await
@@ -732,7 +732,7 @@ impl SelectionService {
     /// Remove a worker and tear down its KV-event listener.
     fn delete_worker<'p>(&self, py: Python<'p>, worker_id: u64) -> PyResult<Bound<'p, PyAny>> {
         let core = self.inner.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let record = core
                 .delete_worker(worker_id)
                 .await
@@ -769,7 +769,7 @@ impl SelectionService {
         let req: OverlapScoresRequest =
             depythonize(request.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?;
         let core = self.inner.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let resp = core.overlap_scores(req).await.map_err(selection_to_pyerr)?;
             Python::with_gil(|py| pythonize(py, &resp).map(|o| o.unbind()).map_err(to_pyerr))
         })
@@ -780,7 +780,7 @@ impl SelectionService {
         let req: SelectRequest =
             depythonize(request.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?;
         let core = self.inner.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let resp = core.select(req).await.map_err(selection_to_pyerr)?;
             Python::with_gil(|py| pythonize(py, &resp).map(|o| o.unbind()).map_err(to_pyerr))
         })
@@ -795,7 +795,7 @@ impl SelectionService {
         let req: SelectAndReserveRequest =
             depythonize(request.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?;
         let core = self.inner.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let resp = core
                 .select_and_reserve(req)
                 .await
@@ -814,7 +814,7 @@ impl SelectionService {
         let req: ReservationRequest =
             depythonize(request.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?;
         let core = self.inner.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let resp = core
                 .create_reservation(req)
                 .await
@@ -830,7 +830,7 @@ impl SelectionService {
         selection_id: String,
     ) -> PyResult<Bound<'p, PyAny>> {
         let core = self.inner.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             core.prefill_complete(&selection_id)
                 .await
                 .map_err(selection_to_pyerr)?;
@@ -853,7 +853,7 @@ impl SelectionService {
         selection_id: String,
     ) -> PyResult<Bound<'p, PyAny>> {
         let core = self.inner.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             core.free_reservation(&selection_id)
                 .await
                 .map_err(selection_to_pyerr)?;
@@ -880,7 +880,7 @@ impl SelectionService {
         let req: PotentialLoadsRequest =
             depythonize(request.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?;
         let core = self.inner.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let resp = core
                 .potential_loads(req)
                 .await
@@ -896,7 +896,7 @@ impl SelectionService {
         endpoint: String,
     ) -> PyResult<Bound<'p, PyAny>> {
         let service = Arc::clone(&self.inner);
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             service
                 .register_replica_peer(endpoint)
                 .await
@@ -911,7 +911,7 @@ impl SelectionService {
         endpoint: String,
     ) -> PyResult<Bound<'p, PyAny>> {
         let service = Arc::clone(&self.inner);
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             service
                 .deregister_replica_peer(endpoint)
                 .await
@@ -927,7 +927,7 @@ impl SelectionService {
     /// Export the current indexer state in the standalone `/dump` shape.
     fn indexer_snapshot<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
         let service = Arc::clone(&self.inner);
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let snapshot = service.indexer_snapshot().await;
             Python::with_gil(|py| {
                 pythonize(py, &snapshot)
@@ -944,7 +944,7 @@ impl SelectionService {
         peers: Vec<String>,
     ) -> PyResult<Bound<'p, PyAny>> {
         let service = Arc::clone(&self.inner);
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             service
                 .recover_indexer_from_peers(&peers)
                 .await
@@ -1043,7 +1043,7 @@ impl WorkerMetricsPublisher {
     ) -> PyResult<Bound<'p, PyAny>> {
         let rs_publisher = self.inner.clone();
         let rs_endpoint = endpoint.inner;
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             rs_publisher
                 .create_endpoint(rs_endpoint)
                 .await
@@ -1094,7 +1094,7 @@ impl MultimodalEmbeddingCachePublisher {
     ) -> PyResult<Bound<'p, PyAny>> {
         let rs_publisher = self.inner.clone();
         let rs_endpoint = endpoint.inner;
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             rs_publisher
                 .create_endpoint(rs_endpoint)
                 .await
@@ -2106,7 +2106,7 @@ impl KvRouter {
         tracker: Option<Arc<RequestTracker>>,
         response_buffer_size: usize,
     ) -> PyResult<Bound<'p, PyAny>> {
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let single_in = SingleIn::new(request);
             let stream = inner.generate(single_in).await.map_err(to_pyerr)?;
             let (tx, rx) =
@@ -2256,7 +2256,7 @@ impl KvRouter {
 
         // The initial-worker wait can be unbounded. Releasing the GIL makes it
         // supervisable from another thread, but not cancellable.
-        let runtime = pyo3_async_runtimes::tokio::get_runtime();
+        let runtime = crate::bridge_runtime();
         py.allow_threads(|| {
             runtime.block_on(async move {
                 let client = endpoint.inner.client().await.map_err(to_pyerr)?;
@@ -2483,7 +2483,7 @@ impl KvRouter {
         let chooser = Arc::clone(self.inner.kv_router());
         let update_states = request_id.is_some();
 
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let admitted = chooser
                 .find_best_match_details_with_policy_class_admitted(
                     request_id.as_deref(),
@@ -2575,7 +2575,7 @@ impl KvRouter {
     ) -> PyResult<Bound<'p, PyAny>> {
         let chooser = Arc::clone(self.inner.kv_router());
 
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             chooser
                 .mark_prefill_completed(&request_id)
                 .await
@@ -2588,7 +2588,7 @@ impl KvRouter {
     fn free<'p>(&self, py: Python<'p>, request_id: String) -> PyResult<Bound<'p, PyAny>> {
         let chooser = Arc::clone(self.inner.kv_router());
 
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             chooser.free(&request_id).await.map_err(to_pyerr)?;
             Ok(())
         })
@@ -2608,7 +2608,7 @@ impl KvRouter {
             .transpose()?;
         let chooser = Arc::clone(self.inner.kv_router());
 
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let loads = chooser
                 .get_potential_loads(
                     &token_ids,
@@ -2654,7 +2654,7 @@ impl KvRouter {
             .transpose()?;
         let chooser = Arc::clone(self.inner.kv_router());
 
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let scores = chooser
                 .get_overlap_scores(
                     &token_ids,
@@ -2679,7 +2679,7 @@ impl KvRouter {
     fn dump_events<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
         let chooser = Arc::clone(self.inner.kv_router());
 
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let events = chooser.dump_events().await.map_err(to_pyerr)?;
             // Serialize to JSON string
             let json_str = serde_json::to_string(&events).map_err(to_pyerr)?;

lib/bindings/python/rust/llm/kv/demand_driven.rs

@@ -189,7 +189,7 @@ impl<T: 'static> Drop for DemandDrivenOwner<T> {
 
         // Python can release the stream outside Tokio. Keep the final Arc alive
         // until the PyO3 runtime drops it so nested router guards can spawn cleanup.
-        drop(pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
+        drop(crate::bridge_runtime().spawn(async move {
             drop(state);
         }));
     }
@@ -218,7 +218,7 @@ impl DemandDrivenResponseStream {
     #[pyo3(name = "__anext__")]
     fn next<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
         let source = self.source.state();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             loop {
                 let response = source
                     .next()
@@ -239,7 +239,7 @@ pub(super) fn process_request_to_stream<'p>(
     request: llm_rs::protocols::common::preprocessor::PreprocessedRequest,
     tracker: Option<Arc<RequestTracker>>,
 ) -> PyResult<Bound<'p, PyAny>> {
-    pyo3_async_runtimes::tokio::future_into_py(py, async move {
+    crate::future_into_py(py, async move {
         let single_in = SingleIn::new(request);
         let stream = inner.generate(single_in).await.map_err(to_pyerr)?;
         // Zero capacity is genuinely demand-driven only for direct Python

lib/bindings/python/rust/llm/kv_dc_relay.rs

@@ -198,7 +198,7 @@ impl KvDcRelay {
         let dc_id = self.dc_id.clone();
         let config = self.config.clone();
         let inner = self.inner.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             inner
                 .get_or_try_init(|| async move {
                     llm_rs::kv_dc_relay::KvDcRelay::start(
@@ -218,7 +218,7 @@ impl KvDcRelay {
     #[cfg(feature = "ckf-diagnostics")]
     fn stats<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         let inner = self.started()?;
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let stats = inner.stats().await.map_err(to_pyerr)?;
             Python::with_gil(|py| {
                 pythonize::pythonize(py, &stats)
@@ -230,7 +230,7 @@ impl KvDcRelay {
 
     fn health<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         let inner = self.started()?;
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let health = inner.health().await;
             Python::with_gil(|py| {
                 pythonize::pythonize(py, &health)
@@ -242,9 +242,7 @@ impl KvDcRelay {
 
     fn flush<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         let inner = self.started()?;
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            inner.flush().await.map_err(to_pyerr)
-        })
+        crate::future_into_py(py, async move { inner.flush().await.map_err(to_pyerr) })
     }
 
     #[cfg(feature = "ckf-diagnostics")]
@@ -254,7 +252,7 @@ impl KvDcRelay {
         serving_endpoint: String,
     ) -> PyResult<Bound<'py, PyAny>> {
         let inner = self.started()?;
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             let endpoint = dynamo_runtime::protocols::EndpointId::from(serving_endpoint.as_str());
             let diagnostic = inner
                 .diagnostic_snapshot(&endpoint)
@@ -270,14 +268,12 @@ impl KvDcRelay {
 
     fn shutdown<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         let inner = self.started()?;
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            inner.shutdown().await.map_err(to_pyerr)
-        })
+        crate::future_into_py(py, async move { inner.shutdown().await.map_err(to_pyerr) })
     }
 
     fn wait_for_shutdown<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         let inner = self.started()?;
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::future_into_py(py, async move {
             inner.wait_for_shutdown().await;
             Ok(())
         })