From d76a7d76c4ce99bdc13c3ca68eefa53438000f22 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Sat, 14 Nov 2026 22:27:55 -0800 Subject: [PATCH] perf(host-python): settle ready async values inline instead of spawning on tokio (#51906) run_async_value polls the future once on the caller's thread with the GIL released and the runtime entered. A future that completes returns an already settled asyncio future on the caller's loop, and only a pending one is handed to future_into_py. --- .../rust/crates/host-python/src/runtime.rs | 113 +++++++++++++++++- .../test_request_redis_batch_post_call.py | 2 + 2 files changed, 110 insertions(+), 5 deletions(-) diff --git a/litellm/rust/crates/host-python/src/runtime.rs b/litellm/rust/crates/host-python/src/runtime.rs index 2a8fc598fc4..bc16faae0da 100644 --- a/litellm/rust/crates/host-python/src/runtime.rs +++ b/litellm/rust/crates/host-python/src/runtime.rs @@ -8,6 +8,7 @@ use crate::fork_gate::{ForkGate, Refused, RuntimeAlreadyStarted}; use crate::{Error, Pythonized, panic_to_pyerr, release_gil}; use futures_util::FutureExt; use litellm_error::Report; +use pyo3::IntoPyObjectExt; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; use serde::Serialize; @@ -95,18 +96,32 @@ where pyo3_async_runtimes::tokio::future_into_py(py, future) } +fn loop_future(py: Python<'_>) -> PyResult> { + py.import("asyncio")? + .call_method0("get_running_loop")? + .call_method0("create_future") +} + pub fn ready_future<'py>( py: Python<'py>, value: &Bound<'py, PyAny>, ) -> PyResult> { - let future = py - .import("asyncio")? - .call_method0("get_running_loop")? - .call_method0("create_future")?; + let future = loop_future(py)?; future.call_method1("set_result", (value,))?; Ok(future) } +fn settle<'py>( + future: Bound<'py, PyAny>, + result: PyResult>, +) -> PyResult> { + match result { + Ok(value) => future.call_method1("set_result", (value,))?, + Err(error) => future.call_method1("set_exception", (error.into_value(future.py()),))?, + }; + Ok(future) +} + pub fn run_sync( py: Python<'_>, future: F, @@ -181,7 +196,18 @@ where T: for<'py> IntoPyObject<'py> + Send + 'static, F: Future> + Send + 'static, { - future_into_py(py, async move { catch_future_panic(future).await? }) + let settled = loop_future(py)?; + let mut future = Box::pin(async move { Ok(catch_future_panic(future).await.flatten()) }); + // If the future is already done (e.g. in-memory read), just poll once to + // return it instead of relying on a separate thread to wakeup. This saves + // ~100 microseconds per call + match poll_async_value(py, future.as_mut())? { + Poll::Ready(result) => settle( + settled, + result.and_then(|value| value.into_bound_py_any(py)), + ), + Poll::Pending => future_into_py(py, async move { future.await? }), + } } pub fn poll_async_value(py: Python<'_>, future: Pin<&mut F>) -> PyResult> @@ -520,6 +546,83 @@ asyncio.run(exercise()) }); } + #[pyfunction] + fn inline_value(py: Python<'_>, mode: String) -> PyResult> { + let caller = thread::current().id(); + run_async_value(py, async move { + match mode.as_str() { + "error" => Err(PyLookupError::new_err("inline failure")), + "pending" => { + tokio::task::yield_now().await; + Ok(thread::current().id() == caller) + } + _ => Ok(thread::current().id() == caller), + } + }) + } + + #[rstest] + #[case::ready("ready", "done and await future is True")] + #[case::error("error", "done and await future is LookupError")] + #[case::pending("pending", "not done and await future is False")] + fn an_async_value_settles_inline_only_when_its_first_poll_completes( + #[from(initialized_python)] python: &InitializedPython, + #[case] mode: &str, + #[case] expected: &str, + ) { + python.attach(|py| { + let locals = PyDict::new(py); + locals + .set_item("inline_value", wrap_pyfunction!(inline_value, py).unwrap()) + .unwrap(); + locals.set_item("mode", mode).unwrap(); + py.run( + c" +import asyncio + +async def exercise(): + future = inline_value(mode) + assert future.get_loop() is asyncio.get_running_loop() + state = 'done' if future.done() else 'not done' + try: + value = repr(await asyncio.wait_for(future, 2)) + except LookupError: + value = 'LookupError' + return f'{state} and await future is {value}' + +result = asyncio.run(exercise()) +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let result: String = locals + .get_item("result") + .unwrap() + .unwrap() + .extract() + .unwrap(); + assert_eq!(result, expected); + }); + } + + #[rstest] + fn an_async_value_without_a_running_loop_fails_before_running_the_future( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { + let ran = Arc::new(AtomicBool::new(false)); + let flag = Arc::clone(&ran); + let error = run_async_value(py, async move { + flag.store(true, Ordering::SeqCst); + Ok(()) + }) + .unwrap_err(); + assert!(error.is_instance_of::(py)); + assert!(!ran.load(Ordering::SeqCst)); + }); + } + #[rstest] fn sync_runner_polls_future_on_the_caller_thread( #[from(initialized_python)] python: &InitializedPython, diff --git a/tests/unit/caching/test_request_redis_batch_post_call.py b/tests/unit/caching/test_request_redis_batch_post_call.py index 85e3fd24cbe..eefadfdd853 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -127,6 +127,7 @@ async def test_post_call_owners_declare_on_the_native_post_call_batch_until_the_ await limiter.async_increment_tokens_with_ttl_preservation(_token_ops(_K1)) await limiter._release_stashed_parallel_slot(_slot_stash("slot-1", _K1), None, in_logging_callback=True) await flush_post_call_redis_batches() + await fakes.settle() assert fakes.post_call == ["token_increment", "parallel_release"] assert calls == [ @@ -228,6 +229,7 @@ async def test_a_released_slot_is_free_locally_at_once_and_the_older_redis_count assert limiter.rate_limits.local((_K1.gauge,)) == ({"slot-2": 1.0, "slot-3": 1.0},) limiter.rate_limits.set_local(_K1.gauge, {"slot-2": 1.0, "slot-3": 1.0, "slot-4": 1.0}, None) await flush_post_call_redis_batches() + await fakes.settle() assert calls == [("script", ["{api_key:k1}:max_parallel_requests"], ["slot-1"])] assert limiter.rate_limits.local((_K1.gauge,)) == ({"slot-2": 1.0, "slot-3": 1.0, "slot-4": 1.0},)