scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 182

fix(python-bridge): harden sync and async route boundaries

Browse source
This commit is contained in:
Yujong Lee 2026-09-01 21:01:28 -07:00 • committed by GitHub
parent 157c629b90
commit a5b9fba287
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 1095 additions and 177 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
.github/workflows/test-rust.yml vendored
View file

@ -9,6 +9,7 @@ on:
- "rust-toolchain.toml"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
- ".github/workflows/test-rust.yml"
pull_request:
branches:
@ -23,6 +24,7 @@ on:
- "rust-toolchain.toml"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
- ".github/workflows/test-rust.yml"
permissions:
@ -121,3 +123,6 @@ jobs:
env:
RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl
- name: Test native route wheel
run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl

23
litellm-rust/Cargo.lock generated
View file

@ -1442,13 +1442,16 @@ name = "litellm-python-bridge"
version = "0.1.0"
dependencies = [
"criterion",
"futures-util",
"litellm-ai-gateway",
"litellm-core",
"litellm-python-interop",
"pyo3",
"pyo3-async-runtimes",
"serde",
"serde_json",
"tokio",
"tokio-tungstenite",
]
[[package]]
@ -1669,9 +1672,9 @@ dependencies = [
[[package]]
name = "pyo3"
version = "0.29.0"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c"
checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b"
dependencies = [
"libc",
"once_cell",
@ -1697,18 +1700,18 @@ dependencies = [
[[package]]
name = "pyo3-build-config"
version = "0.29.0"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078"
checksum = "f41027e41b4bd03f6e60f9f417fe24a6341a6bb744edd62b6f709f2a52ea30e9"
dependencies = [
"target-lexicon",
]
[[package]]
name = "pyo3-ffi"
version = "0.29.0"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b"
checksum = "e591a95526fead067432c3b3a33fc74770b87b1e04e73671090d9c2055a2b327"
dependencies = [
"libc",
"pyo3-build-config",
@ -1716,9 +1719,9 @@ dependencies = [
[[package]]
name = "pyo3-macros"
version = "0.29.0"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771"
checksum = "73225868fc1cd84eef2c3c230ddb91273bf1de46aeb8a4248da76d32a0924a1c"
dependencies = [
"proc-macro2",
"pyo3-macros-backend",
@ -1728,9 +1731,9 @@ dependencies = [
[[package]]
name = "pyo3-macros-backend"
version = "0.29.0"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362"
checksum = "571575aa3749fa6216757dd47d2a3e7ef360f329a40f0666a9fbd14889024952"
dependencies = [
"heck",
"proc-macro2",

2
litellm-rust/Cargo.toml
View file

@ -18,7 +18,7 @@ litellm-core = { path = "crates/core" }
litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
litellm-python-interop = { path = "crates/python-interop" }
axum = "0.7"
pyo3 = "0.29.0"
pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
pythonize = "0.29.0"
rand = "0.8"

3
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -16,16 +16,19 @@ extension-module = ["pyo3/extension-module"]
panic-test = []
[dependencies]
futures-util.workspace = true
litellm-core = { workspace = true, features = ["bedrock-auth"] }
litellm-ai-gateway = { workspace = true, default-features = false }
litellm-python-interop.workspace = true
pyo3.workspace = true
pyo3-async-runtimes.workspace = true
serde.workspace = true
serde_json.workspace = true
tokio.workspace = true
[dev-dependencies]
criterion = "0.8.2"
tokio-tungstenite.workspace = true
[[bench]]
name = "serialization"

115
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -6,6 +6,7 @@ mod routes;
use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
use pyo3::prelude::*;
use pyo3::types::PyAny;
use serde_json::Value;
use crate::errors::core_error_to_pyerr;
use crate::marshal::{marshal_headers, optional_timeout};
@ -23,16 +24,16 @@ impl ResponsesWebSocketConnection {
_cls: &Bound<'py, pyo3::types::PyType>,
py: Python<'py>,
url: String,
headers: Option<Py<PyAny>>,
#[pyo3(from_py_with = litellm_python_interop::from_py)] headers: Option<Value>,
timeout_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
let headers = marshal_headers(py, headers)?;
let headers = marshal_headers(headers)?;
let timeout = optional_timeout(timeout_seconds);
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout)
.await
.map_err(core_error_to_pyerr)?;
Python::attach(|py| Py::new(py, ResponsesWebSocketConnection { inner }))
Ok(ResponsesWebSocketConnection { inner })
})
}
@ -58,24 +59,36 @@ impl ResponsesWebSocketConnection {
}
}
#[pymodule]
fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> {
errors::register(module)?;
routes::register(module)?;
module.add_class::<ResponsesWebSocketConnection>()?;
diagnostics::register(module)
#[pymodule(gil_used = false)]
mod _native {
use pyo3::prelude::*;
#[pymodule_init]
fn init(module: &Bound<'_, PyModule>) -> PyResult<()> {
super::errors::register(module)?;
super::routes::register(module)?;
module.add_class::<super::ResponsesWebSocketConnection>()?;
super::diagnostics::register(module)
}
}
#[cfg(test)]
mod tests {
use std::ffi::CString;
use std::time::Duration;
use futures_util::{SinkExt, StreamExt};
use pyo3::types::PyDict;
use tokio::net::TcpListener;
use tokio_tungstenite::{accept_async, tungstenite::Message};
use super::*;
#[test]
fn module_registration_preserves_the_public_surface() {
Python::initialize();
Python::attach(|py| {
let module = PyModule::new(py, "_native").expect("module should be created");
_native(&module).expect("module should register");
let module = pyo3::wrap_pymodule!(_native)(py).into_bound(py);
let expected = [
"RustBridgeDeclined",
@ -93,13 +106,79 @@ mod tests {
"gil_stats",
];
for name in expected {
assert!(
module
.hasattr(name)
.expect("attribute lookup should succeed")
);
}
let public_names: Vec<String> = module
.dict()
.keys()
.extract::<Vec<String>>()
.expect("module names should be strings")
.into_iter()
.filter(|name| !name.starts_with("__"))
.collect();
assert_eq!(public_names, expected);
});
}
#[test]
fn responses_websocket_connection_round_trips_through_python() {
Python::initialize();
let runtime = pyo3_async_runtimes::tokio::get_runtime();
let listener = runtime
.block_on(TcpListener::bind("127.0.0.1:0"))
.expect("listener should bind");
let address = listener
.local_addr()
.expect("listener should have an address");
let server = runtime.spawn(async move {
let (stream, _) = listener.accept().await.expect("server should accept");
let mut socket = accept_async(stream)
.await
.expect("handshake should succeed");
let message = socket
.next()
.await
.expect("client should send a frame")
.expect("client frame should be valid");
assert_eq!(message, Message::Text("from-python".into()));
socket
.send(Message::Text("from-server".into()))
.await
.expect("server should reply");
assert!(matches!(socket.next().await, Some(Ok(Message::Close(_)))));
});
Python::attach(|py| {
let module = pyo3::wrap_pymodule!(_native)(py).into_bound(py);
let locals = PyDict::new(py);
locals
.set_item("native", &module)
.expect("module should enter Python locals");
locals
.set_item("url", format!("ws://{address}"))
.expect("URL should enter Python locals");
let code = CString::new(
r#"
import asyncio
async def exercise():
connection = await native.ResponsesWebSocketConnection.connect(url)
assert type(connection) is native.ResponsesWebSocketConnection
await connection.send_text("from-python")
assert await connection.recv_text() == "from-server"
await connection.close()
assert await connection.recv_text() is None
asyncio.run(asyncio.wait_for(exercise(), timeout=5))
"#,
)
.expect("Python source should not contain null bytes");
py.run(&code, Some(&locals), Some(&locals))
.expect("Python WebSocket methods should round trip");
});
runtime
.block_on(async { tokio::time::timeout(Duration::from_secs(5), server).await })
.expect("server should finish")
.expect("server task should not panic");
}
}

7
litellm-rust/crates/python-bridge/src/marshal.rs
View file

@ -30,12 +30,9 @@ pub(crate) fn optional_timeout(timeout_seconds: Option<f64>) -> Option<Duration>
})
}
pub(crate) fn marshal_headers(
py: Python<'_>,
headers: Option<Py<PyAny>>,
) -> PyResult<HashMap<String, String>> {
pub(crate) fn marshal_headers(headers: Option<Value>) -> PyResult<HashMap<String, String>> {
let value = match headers {
Some(headers) => from_py(headers.bind(py))?,
Some(headers) => headers,
None => Value::Object(Map::new()),
};
let Value::Object(headers) = value else {

55
litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs
View file

@ -1,12 +1,14 @@
use litellm_core::audio_transcription::{
AudioTranscriptionRequest, audio_transcription as run_audio_transcription,
};
use litellm_python_interop::{from_py, release_gil, to_py};
use litellm_python_interop::from_py;
use pyo3::prelude::*;
use crate::errors::core_error_to_pyerr;
use crate::marshal::{optional_object_to_map, optional_timeout};
use super::runtime::{run_async, run_sync};
#[pyfunction]
#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
#[allow(clippy::too_many_arguments)]
@ -28,9 +30,10 @@ fn transcription(
};
let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
let timeout = optional_timeout(timeout_seconds);
let result = release_gil(py, || {
pyo3_async_runtimes::tokio::get_runtime().block_on(run_audio_transcription(
AudioTranscriptionRequest {
run_sync(
py,
async move {
run_audio_transcription(AudioTranscriptionRequest {
model: &model,
audio,
api_key: api_key.as_deref(),
@ -39,13 +42,11 @@ fn transcription(
extra_headers,
optional_params,
timeout,
},
))
});
match result {
Ok(value) => to_py(py, &value),
Err(err) => Err(core_error_to_pyerr(err)),
}
})
.await
},
core_error_to_pyerr,
)
}
#[pyfunction]
@ -69,21 +70,23 @@ fn atranscription(
};
let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
let timeout = optional_timeout(timeout_seconds);
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let value = run_audio_transcription(AudioTranscriptionRequest {
model: &model,
audio,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
optional_params,
timeout,
})
.await
.map_err(core_error_to_pyerr)?;
Python::attach(|py| to_py(py, &value))
})
run_async(
py,
async move {
run_audio_transcription(AudioTranscriptionRequest {
model: &model,
audio,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
optional_params,
timeout,
})
.await
},
core_error_to_pyerr,
)
}
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {

64
litellm-rust/crates/python-bridge/src/routes/chat_completions.rs
View file

@ -1,10 +1,10 @@
use std::time::Duration;
use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse};
use litellm_core::chat_completions::types::ChatCompletionsRequest;
use litellm_core::chat_completions::{
chat_completions as run_chat_completions, chat_completions_decline_reason,
};
use litellm_python_interop::{from_py, release_gil, to_py};
use litellm_python_interop::from_py;
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use serde_json::{Map, Value};
@ -12,12 +12,7 @@ use serde_json::{Map, Value};
use crate::errors::chat_completions_error_to_pyerr;
use crate::marshal::{optional_object_to_map, optional_timeout};
fn chat_completions_response_to_py(
py: Python<'_>,
response: ChatCompletionsResponse,
) -> PyResult<Py<PyAny>> {
to_py(py, &response)
}
use super::runtime::{run_async, run_sync};
type MarshaledChatCompletionsInputs = (
Value,
@ -95,9 +90,10 @@ fn chat_completions(
timeout_seconds,
)?;
let result = release_gil(py, || {
pyo3_async_runtimes::tokio::get_runtime().block_on(run_chat_completions(
ChatCompletionsRequest {
run_sync(
py,
async move {
run_chat_completions(ChatCompletionsRequest {
model: &model,
messages,
optional_params,
@ -106,14 +102,11 @@ fn chat_completions(
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
timeout,
},
))
});
match result {
Ok(response) => chat_completions_response_to_py(py, response),
Err(err) => Err(chat_completions_error_to_pyerr(err)),
}
})
.await
},
chat_completions_error_to_pyerr,
)
}
#[pyfunction]
@ -138,22 +131,23 @@ fn achat_completions(
timeout_seconds,
)?;
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let response = run_chat_completions(ChatCompletionsRequest {
model: &model,
messages,
optional_params,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
timeout,
})
.await
.map_err(chat_completions_error_to_pyerr)?;
Python::attach(|py| chat_completions_response_to_py(py, response))
})
run_async(
py,
async move {
run_chat_completions(ChatCompletionsRequest {
model: &model,
messages,
optional_params,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
timeout,
})
.await
},
chat_completions_error_to_pyerr,
)
}
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {

74
litellm-rust/crates/python-bridge/src/routes/messages.rs
View file

@ -1,8 +1,8 @@
use std::time::Duration;
use litellm_core::messages::messages as run_messages;
use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest};
use litellm_python_interop::{from_py, release_gil, to_py};
use litellm_core::messages::types::MessagesRequest;
use litellm_python_interop::from_py;
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use serde_json::{Map, Value};
@ -10,12 +10,7 @@ use serde_json::{Map, Value};
use crate::errors::core_error_to_pyerr;
use crate::marshal::{optional_object_to_map, optional_timeout};
fn messages_response_to_py(
py: Python<'_>,
response: AnthropicMessagesResponse,
) -> PyResult<Py<PyAny>> {
to_py(py, &response)
}
use super::runtime::{run_async, run_sync};
type MarshaledMessagesInputs = (Value, Option<Map<String, Value>>, Option<Duration>);
@ -52,22 +47,22 @@ fn messages(
let (body, extra_headers, timeout) =
marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?;
let result = release_gil(py, || {
pyo3_async_runtimes::tokio::get_runtime().block_on(run_messages(MessagesRequest {
model: &model,
body,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
timeout,
}))
});
match result {
Ok(response) => messages_response_to_py(py, response),
Err(err) => Err(core_error_to_pyerr(err)),
}
run_sync(
py,
async move {
run_messages(MessagesRequest {
model: &model,
body,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
timeout,
})
.await
},
core_error_to_pyerr,
)
}
#[pyfunction]
@ -86,21 +81,22 @@ fn amessages(
let (body, extra_headers, timeout) =
marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?;
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let response = run_messages(MessagesRequest {
model: &model,
body,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
timeout,
})
.await
.map_err(core_error_to_pyerr)?;
Python::attach(|py| messages_response_to_py(py, response))
})
run_async(
py,
async move {
run_messages(MessagesRequest {
model: &model,
body,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
timeout,
})
.await
},
core_error_to_pyerr,
)
}
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {

1
litellm-rust/crates/python-bridge/src/routes/mod.rs
View file

@ -4,6 +4,7 @@ mod audio_transcription;
mod chat_completions;
mod messages;
mod ocr;
mod runtime;
pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
ocr::register(module)?;

87
litellm-rust/crates/python-bridge/src/routes/ocr.rs
View file

@ -1,13 +1,15 @@
use std::time::Duration;
use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
use litellm_python_interop::{from_py, release_gil, to_py};
use litellm_python_interop::from_py;
use pyo3::prelude::*;
use serde_json::{Map, Value};
use crate::errors::core_error_to_pyerr;
use crate::marshal::{optional_object_to_map, optional_timeout};
use super::runtime::{run_async, run_sync};
type MarshaledOcrInputs = (
Value,
Option<Map<String, Value>>,
@ -55,27 +57,27 @@ fn ocr(
timeout_seconds,
)?;
let result = release_gil(py, || {
pyo3_async_runtimes::tokio::get_runtime().block_on(run_ocr(OcrRequest {
model: &model,
document,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
optional_params,
timeout,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: Default::default(),
litellm_call_id: None,
}))
});
match result {
Ok(value) => to_py(py, &value),
Err(err) => Err(core_error_to_pyerr(err)),
}
run_sync(
py,
async move {
run_ocr(OcrRequest {
model: &model,
document,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
optional_params,
timeout,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: Default::default(),
litellm_call_id: None,
})
.await
},
core_error_to_pyerr,
)
}
#[pyfunction]
@ -100,26 +102,27 @@ fn aocr(
timeout_seconds,
)?;
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let value = run_ocr(OcrRequest {
model: &model,
document,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
optional_params,
timeout,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: Default::default(),
litellm_call_id: None,
})
.await
.map_err(core_error_to_pyerr)?;
Python::attach(|py| to_py(py, &value))
})
run_async(
py,
async move {
run_ocr(OcrRequest {
model: &model,
document,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
optional_params,
timeout,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: Default::default(),
litellm_call_id: None,
})
.await
},
core_error_to_pyerr,
)
}
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {

423
litellm-rust/crates/python-bridge/src/routes/runtime.rs Normal file
View file

@ -0,0 +1,423 @@
use std::future::Future;
use std::panic::AssertUnwindSafe;
use std::time::Duration;
use futures_util::FutureExt;
use litellm_core::error::Error;
use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil};
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
use serde::Serialize;
use tokio::runtime::{Handle, Runtime};
use tokio::time::{self, MissedTickBehavior};
pub(super) fn run_sync<T, F>(
py: Python<'_>,
future: F,
map_error: fn(Error) -> PyErr,
) -> PyResult<Py<PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
{
run_sync_on(
py,
pyo3_async_runtimes::tokio::get_runtime(),
future,
map_error,
)
}
fn run_sync_on<T, F>(
py: Python<'_>,
runtime: &Runtime,
future: F,
map_error: fn(Error) -> PyErr,
) -> PyResult<Py<PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
{
if Handle::try_current().is_ok() {
return Err(PyRuntimeError::new_err(
"synchronous native routes cannot run from a Tokio context; use the async route",
));
}
let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?;
let result = map_core_result(result, map_error)?;
Pythonized(result).into_pyobject(py).map(Bound::unbind)
}
pub(super) fn run_async<T, F>(
py: Python<'_>,
future: F,
map_error: fn(Error) -> PyErr,
) -> PyResult<Bound<'_, PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
{
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let result = catch_route_panic(future).await?;
let result = map_core_result(result, map_error)?;
Ok(Pythonized(result))
})
}
fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -> PyResult<T> {
match result {
Ok(value) => Ok(value),
Err(error) => Err(
std::panic::catch_unwind(AssertUnwindSafe(|| map_error(error)))
.map_err(panic_to_pyerr)?,
),
}
}
async fn catch_route_panic<T, F>(future: F) -> PyResult<Result<T, Error>>
where
F: Future<Output = Result<T, Error>>,
{
AssertUnwindSafe(future)
.catch_unwind()
.await
.map_err(panic_to_pyerr)
}
async fn wait_for_sync_result<T, F>(future: F) -> PyResult<Result<T, Error>>
where
F: Future<Output = Result<T, Error>>,
{
let future = catch_route_panic(future);
tokio::pin!(future);
let signal_interval = Duration::from_millis(50);
let mut signal_checks =
time::interval_at(time::Instant::now() + signal_interval, signal_interval);
signal_checks.set_missed_tick_behavior(MissedTickBehavior::Delay);
loop {
tokio::select! {
result = &mut future => return result,
_ = signal_checks.tick() => Python::attach(|py| py.check_signals())?,
}
}
}
#[cfg(test)]
mod tests {
use std::ffi::CString;
use std::future::poll_fn;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, mpsc};
use std::task::Poll;
use std::thread;
use std::time::Instant;
use pyo3::panic::PanicException;
use pyo3::types::{PyDict, PyModule};
use serde::Serializer;
use tokio::runtime::Builder;
use super::*;
fn runtime_error(error: Error) -> PyErr {
PyRuntimeError::new_err(error.to_string())
}
fn panicking_error_mapper(_error: Error) -> PyErr {
panic!("error mapper panicked")
}
struct PanickingOutput;
static ASYNC_PROBE_COMPLETED: AtomicUsize = AtomicUsize::new(0);
impl Serialize for PanickingOutput {
fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
panic!("serializer panicked")
}
}
#[pyfunction]
fn async_serialization_panic(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
run_async(py, async { Ok(PanickingOutput) }, runtime_error)
}
#[pyfunction]
fn async_runtime_probe(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
run_async(
py,
async {
ASYNC_PROBE_COMPLETED.fetch_add(1, Ordering::SeqCst);
Ok(true)
},
runtime_error,
)
}
#[pyfunction]
fn runtime_worker_count() -> usize {
pyo3_async_runtimes::tokio::get_runtime()
.metrics()
.num_workers()
}
#[pyfunction]
fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool {
let completion_deadline = Instant::now() + Duration::from_secs(2);
while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions {
if Instant::now() >= completion_deadline {
return false;
}
thread::sleep(Duration::from_millis(1));
}
let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1);
pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
let _ = heartbeat_tx.send(());
});
heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok()
}
fn extract_bool(py: Python<'_>, result: PyResult<Py<PyAny>>) -> bool {
result
.expect("route should complete")
.bind(py)
.extract()
.expect("result should convert")
}
#[test]
fn sync_runner_polls_future_on_the_caller_thread() {
Python::initialize();
Python::attach(|py| {
let caller_thread = std::thread::current().id();
let result = run_sync(
py,
async move { Ok(std::thread::current().id() == caller_thread) },
runtime_error,
);
assert!(extract_bool(py, result));
});
}
#[test]
fn sync_runner_releases_gil_while_waiting() {
Python::initialize();
Python::attach(|py| {
let result = run_sync(
py,
async {
let gil_acquired = tokio::time::timeout(
Duration::from_secs(2),
tokio::task::spawn_blocking(|| Python::attach(|_| true)),
)
.await;
Ok(matches!(gil_acquired, Ok(Ok(true))))
},
runtime_error,
);
assert!(extract_bool(py, result));
});
}
#[test]
fn sync_runner_rejects_calls_from_a_tokio_context() {
Python::initialize();
let runtime = Builder::new_current_thread()
.enable_all()
.build()
.expect("runtime should build");
let error = runtime.block_on(async {
Python::attach(|py| {
run_sync::<bool, _>(py, async { Ok(true) }, runtime_error)
.expect_err("sync route should reject a nested Tokio runtime")
})
});
assert_eq!(
error.to_string(),
"RuntimeError: synchronous native routes cannot run from a Tokio context; use the async route"
);
}
#[test]
fn sync_runner_can_drive_a_current_thread_runtime() {
Python::initialize();
let runtime = Builder::new_current_thread()
.enable_all()
.build()
.expect("runtime should build");
Python::attach(|py| {
let result = run_sync_on(
py,
&runtime,
async {
tokio::task::yield_now().await;
Ok(true)
},
runtime_error,
);
assert!(extract_bool(py, result));
});
}
#[test]
fn sync_runner_maps_a_panicked_future() {
Python::initialize();
Python::attach(|py| {
let error = run_sync::<bool, _>(
py,
poll_fn(|_| -> Poll<Result<bool, Error>> { panic!("route future panicked") }),
runtime_error,
)
.expect_err("panicked route should become a Python exception");
assert!(error.is_instance_of::<PanicException>(py));
assert_eq!(error.to_string(), "PanicException: route future panicked");
});
}
#[test]
fn sync_runner_maps_a_panicked_error_mapper() {
Python::initialize();
Python::attach(|py| {
let error = run_sync::<bool, _>(
py,
async { Err(Error::InvalidRequest("invalid".to_string())) },
panicking_error_mapper,
)
.expect_err("panicked mapper should become a Python exception");
assert!(error.is_instance_of::<PanicException>(py));
assert_eq!(error.to_string(), "PanicException: error mapper panicked");
});
}
#[test]
fn sync_runner_surfaces_serializer_panics() {
Python::initialize();
Python::attach(|py| {
let error = run_sync(py, async { Ok(PanickingOutput) }, runtime_error)
.expect_err("serializer panic should become a Python exception");
assert!(error.is_instance_of::<PanicException>(py));
assert_eq!(error.to_string(), "PanicException: serializer panicked");
});
}
#[test]
fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() {
Python::initialize();
let barrier = Arc::new(tokio::sync::Barrier::new(2));
let callers: Vec<_> = (0..2)
.map(|_| {
let barrier = Arc::clone(&barrier);
thread::spawn(move || {
Python::attach(|py| {
extract_bool(
py,
run_sync(
py,
async move {
Ok(tokio::time::timeout(Duration::from_secs(2), barrier.wait())
.await
.is_ok())
},
runtime_error,
),
)
})
})
})
.collect();
let results: Vec<_> = callers
.into_iter()
.map(|caller| caller.join().expect("caller should not panic"))
.collect();
assert_eq!(results, vec![true, true]);
}
#[test]
fn async_runner_surfaces_serializer_panics() {
Python::initialize();
Python::attach(|py| {
let module = PyModule::new(py, "runtime").expect("module should be created");
module
.add_function(
wrap_pyfunction!(async_serialization_panic, &module)
.expect("function should wrap"),
)
.expect("function should register");
let locals = PyDict::new(py);
locals
.set_item("runtime", &module)
.expect("module should enter Python locals");
let code = CString::new(
r#"
import asyncio
async def exercise():
try:
await runtime.async_serialization_panic()
except BaseException as error:
assert type(error).__name__ == "PanicException"
assert str(error) == "serializer panicked"
else:
raise AssertionError("serializer panic was not raised")
asyncio.run(exercise())
"#,
)
.expect("Python source should not contain null bytes");
py.run(&code, Some(&locals), Some(&locals))
.expect("serializer panic should reach the Python awaiter");
});
}
#[test]
fn async_result_delivery_does_not_stall_tokio_workers() {
Python::initialize();
ASYNC_PROBE_COMPLETED.store(0, Ordering::SeqCst);
Python::attach(|py| {
let module = PyModule::new(py, "runtime").expect("module should be created");
for function in [
wrap_pyfunction!(async_runtime_probe, &module).expect("function should wrap"),
wrap_pyfunction!(runtime_worker_count, &module).expect("function should wrap"),
wrap_pyfunction!(runtime_is_responsive, &module).expect("function should wrap"),
] {
module
.add_function(function)
.expect("function should register");
}
let locals = PyDict::new(py);
locals
.set_item("runtime", &module)
.expect("module should enter Python locals");
let code = CString::new(
r#"
import asyncio
async def exercise():
worker_count = runtime.runtime_worker_count()
awaitables = [runtime.async_runtime_probe() for _ in range(worker_count)]
assert runtime.runtime_is_responsive(worker_count)
assert await asyncio.gather(*awaitables) == [True] * worker_count
asyncio.run(exercise())
"#,
)
.expect("Python source should not contain null bytes");
py.run(&code, Some(&locals), Some(&locals))
.expect("result delivery should leave Tokio workers responsive");
});
}
}

2
litellm-rust/crates/python-interop/src/lib.rs
View file

@ -2,4 +2,4 @@ mod gil;
mod marshal;
pub use gil::{release_count, release_gil};
pub use marshal::{from_py, to_py};
pub use marshal::{Pythonized, from_py, panic_to_pyerr, to_py};

72
litellm-rust/crates/python-interop/src/marshal.rs
View file

@ -1,4 +1,8 @@
use std::any::Any;
use std::panic::{AssertUnwindSafe, catch_unwind};
use pyo3::exceptions::PyValueError;
use pyo3::panic::PanicException;
use pyo3::prelude::*;
use serde::Serialize;
use serde::de::DeserializeOwned;
@ -18,3 +22,71 @@ where
.map(Bound::unbind)
.map_err(|error| PyValueError::new_err(error.to_string()))
}
pub struct Pythonized<T>(pub T);
impl<'py, T> IntoPyObject<'py> for Pythonized<T>
where
T: Serialize,
{
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
catch_unwind(AssertUnwindSafe(|| pythonize::pythonize(py, &self.0)))
.map_err(panic_to_pyerr)?
.map_err(|error| PyValueError::new_err(error.to_string()))
}
}
pub fn panic_to_pyerr(payload: Box<dyn Any + Send>) -> PyErr {
let message = payload
.downcast_ref::<String>()
.map(String::as_str)
.or_else(|| payload.downcast_ref::<&str>().copied())
.unwrap_or("panic from Rust code");
PanicException::new_err(message.to_string())
}
#[cfg(test)]
mod tests {
use serde::Serializer;
use super::*;
struct PanickingSerializer;
impl Serialize for PanickingSerializer {
fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
panic!("serializer panicked")
}
}
#[test]
fn pythonized_converts_on_the_attached_thread() {
Python::initialize();
Python::attach(|py| {
let value: Vec<i32> = Pythonized(vec![1, 2, 3])
.into_pyobject(py)
.and_then(|value| value.extract())
.expect("value should convert");
assert_eq!(value, vec![1, 2, 3]);
});
}
#[test]
fn pythonized_maps_serializer_panics_to_a_base_exception() {
Python::initialize();
Python::attach(|py| {
let error = Pythonized(PanickingSerializer)
.into_pyobject(py)
.expect_err("serializer panic should become a Python exception");
assert!(error.is_instance_of::<PanicException>(py));
assert_eq!(error.to_string(), "PanicException: serializer panicked");
});
}
}

339
tests/test_litellm/rust_bridge/native_route_wheel_test.py Normal file
View file

@ -0,0 +1,339 @@
from __future__ import annotations
import asyncio
import importlib.util
import json
import os
import signal
import subprocess
import sys
import tempfile
import threading
import zipfile
from http.client import HTTPMessage
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from socket import socket as Socket
from typing import Final
REQUEST_STARTED: Final = threading.Event()
REQUEST_CANCELLED: Final = threading.Event()
ANTHROPIC_RESPONSE: Final = (
b'{"id":"msg_native","type":"message","role":"assistant",'
b'"model":"claude-sonnet-4-5","content":[{"type":"text","text":"native-message"}],'
b'"stop_reason":"end_turn","stop_sequence":null,'
b'"usage":{"input_tokens":2,"output_tokens":3}}'
)
class NativeRouteHandler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def do_POST(self) -> None:
content_length: Final = int(self.headers.get("content-length", "0"))
body: Final = json.loads(self.rfile.read(content_length))
route: Final = self.headers.get("x-test-route")
outcome: Final = self.headers.get("x-test-outcome")
assert_native_request(route, outcome, self.path, self.headers, body)
if outcome == "hang":
REQUEST_STARTED.set()
self.connection.settimeout(5)
if connection_was_cancelled(self.connection):
REQUEST_CANCELLED.set()
return
status: Final = 429 if outcome == "429" else 200
response_body: Final = native_response(status, route)
self.send_response(status)
self.send_header("content-type", "application/json")
self.send_header("content-length", str(len(response_body)))
self.send_header("connection", "close")
self.end_headers()
self.wfile.write(response_body)
def log_message(self, _message_format: str, *_args: object) -> None:
pass
def connection_was_cancelled(connection: Socket) -> bool:
try:
return connection.recv(1) == b""
except TimeoutError:
return False
except OSError:
return True
def assert_native_request(
route: str | None,
outcome: str | None,
path: str,
headers: HTTPMessage,
body: object,
) -> None:
if route not in {"ocr", "transcription", "messages", "chat_completions"}:
raise AssertionError(f"unexpected route marker: {route!r}")
if outcome not in {"success", "429", "hang"}:
raise AssertionError(f"unexpected outcome marker: {outcome!r}")
if not isinstance(body, dict):
raise TypeError(f"{route} sent {type(body).__name__}, expected a JSON object")
if route == "ocr":
assert path == "/v1/ocr"
assert headers.get("authorization") == "Bearer sk-native"
assert body["model"] == "mistral-ocr-latest"
assert body["document"]["document_url"] == "https://example.com/document.pdf"
assert body["include_image_base64"] is True
return
if route == "transcription":
assert path == "/model/mistral.voxtral-mini-3b-2507/converse"
assert headers.get("authorization", "").startswith("AWS4-HMAC-SHA256 ")
assert headers.get("x-amz-date")
assert body["messages"][0]["content"][0]["audio"]["source"]["bytes"] == "AQI="
assert "The audio language is en" in body["messages"][0]["content"][1]["text"]
return
assert path == "/v1/messages"
assert headers.get("x-api-key") == "sk-native"
assert body["model"] == "claude-sonnet-4-5"
if route == "messages":
assert body["max_tokens"] == 16
assert body["messages"][0]["content"] == "hello-from-messages"
return
assert body["max_tokens"] == 17
assert body["messages"][0]["content"] == [{"type": "text", "text": "hello-from-chat"}]
def native_response(status: int, route: str | None) -> bytes:
if status == 429:
return b'{"error":"native-rate-limit"}'
if route == "ocr":
return b'{"pages":[{"index":0,"markdown":"native-ocr"}]}'
if route == "transcription":
return b'{"output":{"message":{"content":[{"text":"native-transcription"}]}}}'
return ANTHROPIC_RESPONSE
def load_native(native_path: Path) -> object:
module_spec: Final = importlib.util.spec_from_file_location("litellm.rust_bridge._native", native_path)
if module_spec is None or module_spec.loader is None:
raise RuntimeError("cannot create native extension import specification")
native_module: Final = importlib.util.module_from_spec(module_spec)
module_spec.loader.exec_module(native_module)
return native_module
def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]:
common: Final = {
"api_base": api_base,
"extra_headers": {"x-test-outcome": outcome, "x-test-route": route},
"timeout_seconds": 3.0,
}
if route == "ocr":
return common | {
"model": "mistral-ocr-latest",
"document": {"type": "document_url", "document_url": "https://example.com/document.pdf"},
"api_key": "sk-native",
"custom_llm_provider": "mistral",
"optional_params": {"include_image_base64": True},
}
if route == "transcription":
return common | {
"model": "mistral.voxtral-mini-3b-2507",
"audio": {"data": "AQI=", "format": "wav", "filename": "audio.wav"},
"custom_llm_provider": "bedrock",
"optional_params": {
"aws_access_key_id": "native-access-key",
"aws_secret_access_key": "native-secret-key",
"aws_region_name": "us-east-1",
"language": "en",
},
}
if route == "messages":
return common | {
"model": "claude-sonnet-4-5",
"body": {
"model": "claude-sonnet-4-5",
"max_tokens": 16,
"messages": [{"role": "user", "content": "hello-from-messages"}],
},
"api_key": "sk-native",
"custom_llm_provider": "anthropic",
}
if route == "chat_completions":
return common | {
"model": "anthropic/claude-sonnet-4-5",
"messages": [{"role": "user", "content": "hello-from-chat"}],
"optional_params": {"max_tokens": 17},
"api_key": "sk-native",
}
raise AssertionError(f"unknown route: {route}")
def assert_success(route: str, response: object) -> None:
if not isinstance(response, dict):
raise TypeError(f"{route} returned {type(response).__name__}, expected dict")
actual: Final = success_value(route, response)
expected: Final = (
"native-ocr" if route == "ocr" else "native-transcription" if route == "transcription" else "native-message"
)
if actual != expected:
raise AssertionError(f"{route} returned {actual!r}, expected {expected!r}")
def success_value(route: str, response: dict[object, object]) -> object:
if route == "ocr":
return response["pages"][0]["markdown"]
if route == "transcription":
return response["text"]
if route == "messages":
return response["content"][0]["text"]
return response["choices"][0]["message"]["content"]
def assert_rate_limit(native: object, route: str, error: BaseException) -> None:
if route == "chat_completions":
upstream_error: Final = native.RustUpstreamError
if not isinstance(error, upstream_error) or error.args[0] != 429:
raise AssertionError(f"{route} returned the wrong 429 error: {error!r}")
return
if not isinstance(error, RuntimeError) or "429" not in str(error):
raise AssertionError(f"{route} returned the wrong 429 error: {error!r}")
def exercise_sync(native: object, api_base: str) -> None:
for route in ("ocr", "transcription", "messages", "chat_completions"):
function: Final = getattr(native, route)
assert_success(route, function(**route_kwargs(route, api_base, "success")))
try:
function(**route_kwargs(route, api_base, "429"))
except (RuntimeError, native.RustUpstreamError) as error:
assert_rate_limit(native, route, error)
else:
raise AssertionError(f"{route} accepted a 429 response")
async def exercise_async(native: object, api_base: str) -> None:
for route in ("ocr", "transcription", "messages", "chat_completions"):
function: Final = getattr(native, f"a{route}")
assert_success(route, await function(**route_kwargs(route, api_base, "success")))
try:
await function(**route_kwargs(route, api_base, "429"))
except (RuntimeError, native.RustUpstreamError) as error:
assert_rate_limit(native, route, error)
else:
raise AssertionError(f"a{route} accepted a 429 response")
async def exercise_async_concurrency(native: object, api_base: str) -> None:
responses: Final = await asyncio.wait_for(
asyncio.gather(
*(
native.amessages(**route_kwargs("messages", api_base, "success"))
for _ in range(32)
)
),
timeout=15,
)
for response in responses:
assert_success("messages", response)
def exercise_routes(native_path: Path, api_base: str) -> object:
native: Final = load_native(native_path)
exercise_sync(native, api_base)
asyncio.run(exercise_async(native, api_base))
asyncio.run(exercise_async_concurrency(native, api_base))
return native
def exercise_signal(native: object, api_base: str) -> int:
try:
native.messages(
**route_kwargs("messages", api_base, "hang"),
)
except KeyboardInterrupt:
sys.stdout.write("KeyboardInterrupt\n")
sys.stdout.flush()
sys.stdin.read(1)
return 0
raise AssertionError("sync native route ignored SIGINT")
def verify_sigint(native_path: Path, api_base: str) -> None:
REQUEST_STARTED.clear()
REQUEST_CANCELLED.clear()
process: Final = subprocess.Popen(
(sys.executable, __file__, "child", str(native_path), api_base),
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
)
try:
if not REQUEST_STARTED.wait(30):
process.kill()
stdout, stderr = process.communicate(timeout=5)
raise AssertionError(
f"native route matrix did not reach the hanging upstream\nstdout:\n{stdout}\nstderr:\n{stderr}"
)
os.kill(process.pid, signal.SIGINT)
if not REQUEST_CANCELLED.wait(5):
raise AssertionError("interrupted native route did not cancel its upstream future")
if process.poll() is not None:
raise AssertionError("signal child exited before cancellation was observed")
stdout, stderr = process.communicate(input="\n", timeout=5)
if process.returncode != 0 or stdout != "KeyboardInterrupt\n":
raise AssertionError(
f"signal child exited with status {process.returncode}\nstdout:\n{stdout}\nstderr:\n{stderr}"
)
finally:
if process.poll() is None:
process.kill()
process.wait(timeout=5)
def verify_wheel(wheel: Path) -> int:
with tempfile.TemporaryDirectory() as temporary_directory, zipfile.ZipFile(wheel) as archive:
wheel_root: Final = Path(temporary_directory)
for member in archive.infolist():
target: Final = wheel_root / member.filename
if member.is_dir():
target.mkdir(parents=True, exist_ok=True)
else:
target.parent.mkdir(parents=True, exist_ok=True)
target.write_bytes(archive.read(member))
native_members: Final = tuple(
member
for member in archive.infolist()
if member.filename.startswith("litellm/rust_bridge/_native.") and member.filename.endswith(".so")
)
if len(native_members) != 1:
raise AssertionError(f"expected one native extension, found {len(native_members)}")
native_path: Final = wheel_root / native_members[0].filename
server: Final = ThreadingHTTPServer(("127.0.0.1", 0), NativeRouteHandler)
server_thread: Final = threading.Thread(target=server.serve_forever, daemon=True)
server_thread.start()
api_base: Final = f"http://127.0.0.1:{server.server_address[1]}"
try:
verify_sigint(native_path, api_base)
finally:
server.shutdown()
server.server_close()
server_thread.join(timeout=5)
return 0
def main() -> int:
if len(sys.argv) == 2:
return verify_wheel(Path(sys.argv[1]))
if len(sys.argv) == 4 and sys.argv[1] == "child":
native: Final = exercise_routes(Path(sys.argv[2]), sys.argv[3])
return exercise_signal(native, sys.argv[3])
sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n")
return 2
if __name__ == "__main__":
sys.exit(main())