From f1ef7fc0c2987ff76dc5c67eaaac82263b8a203c Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:18:50 -0700 Subject: [PATCH] feat(rust_bridge): read secrets through Python from Rust routes and declare Rust-only routes with NO_PYTHON (#43057) * done * fix(rust_bridge): run Python secret reads under the caller's contextvars Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(rust_bridge): run every blocking Python call under the caller's contextvars Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Yujong Lee Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/host-python/Cargo.toml | 2 +- .../crates/host-python/src/execution.rs | 21 +- litellm-rust/crates/host-python/src/gil.rs | 322 +++++++++++++++++- litellm-rust/crates/host-python/src/lib.rs | 23 +- litellm-rust/crates/python-bridge/AGENTS.md | 28 ++ litellm-rust/crates/python-bridge/README.md | 10 - .../python-bridge/src/secrets/callback.rs | 85 ++++- .../python-bridge/src/secrets/config.rs | 5 +- .../crates/python-bridge/src/secrets/mod.rs | 98 +----- .../python-bridge/src/secrets/python.rs | 190 +++++++++++ .../python-bridge/src/secrets/resolved.rs | 5 +- .../bedrock/audio_transcription/__init__.py | 14 +- litellm/rust_bridge/dispatch.py | 35 +- litellm/rust_bridge/runtime.py | 84 +++-- .../rust_bridge/ocr/test_secrets.py | 59 +++- .../test_litellm/rust_bridge/test_dispatch.py | 84 ++++- .../test_litellm/rust_bridge/test_runtime.py | 77 +++++ 17 files changed, 956 insertions(+), 186 deletions(-) delete mode 100644 litellm-rust/crates/python-bridge/README.md create mode 100644 litellm-rust/crates/python-bridge/src/secrets/python.rs diff --git a/litellm-rust/crates/host-python/Cargo.toml b/litellm-rust/crates/host-python/Cargo.toml index e2c83fe1081..fb6379dc35a 100644 --- a/litellm-rust/crates/host-python/Cargo.toml +++ b/litellm-rust/crates/host-python/Cargo.toml @@ -12,7 +12,7 @@ pyo3.workspace = true pyo3-async-runtimes.workspace = true pythonize.workspace = true serde.workspace = true -tokio = { workspace = true, features = ["sync"] } +tokio = { workspace = true, features = ["rt", "sync"] } [dev-dependencies] rstest.workspace = true diff --git a/litellm-rust/crates/host-python/src/execution.rs b/litellm-rust/crates/host-python/src/execution.rs index b435bf241d2..c948c15d8eb 100644 --- a/litellm-rust/crates/host-python/src/execution.rs +++ b/litellm-rust/crates/host-python/src/execution.rs @@ -225,29 +225,12 @@ mod tests { use pyo3::exceptions::PyLookupError; use pyo3::panic::PanicException; use pyo3::types::{PyDict, PyModule}; - use rstest::{fixture, rstest}; + use rstest::rstest; use serde::Serializer; use tokio::runtime::Builder; use super::*; - - struct InitializedPython; - - impl InitializedPython { - fn attach(&self, f: F) -> R - where - F: for<'py> FnOnce(Python<'py>) -> R, - { - Python::attach(f) - } - } - - #[fixture] - #[once] - fn initialized_python() -> InitializedPython { - crate::initialize_python(); - InitializedPython - } + use crate::{InitializedPython, initialized_python}; #[derive(Debug)] struct Error(String); diff --git a/litellm-rust/crates/host-python/src/gil.rs b/litellm-rust/crates/host-python/src/gil.rs index 04b966a6002..dc91c8987b7 100644 --- a/litellm-rust/crates/host-python/src/gil.rs +++ b/litellm-rust/crates/host-python/src/gil.rs @@ -1,6 +1,57 @@ -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{ + Arc, Mutex, + atomic::{AtomicU64, Ordering}, +}; -use pyo3::prelude::*; +use pyo3::{exceptions::PyRuntimeError, prelude::*}; + +/// The caller's `contextvars` context, captured at the Python entry point so blocking Python +/// work started from Rust sees the same request-local values as the Python caller. +#[derive(Clone)] +pub struct PythonContext(Arc>); + +impl PythonContext { + pub fn capture(py: Python<'_>) -> PyResult { + Ok(Self(Arc::new( + py.import("contextvars")? + .call_method0("copy_context")? + .unbind(), + ))) + } + + /// Runs `f` inside a fresh copy of the captured context. The copy is what lets two blocking + /// calls run concurrently: a `contextvars.Context` cannot be entered twice at once. + pub fn enter(&self, py: Python<'_>, f: F) -> PyResult + where + F: FnOnce(Python<'_>) -> T + Send + Sync + 'static, + T: Send + 'static, + { + let copy = self.0.bind(py).call_method0("copy")?; + let body = Arc::new(Mutex::new(Some(f))); + let value = Arc::new(Mutex::new(None::)); + let callback = { + let body = Arc::clone(&body); + let value = Arc::clone(&value); + pyo3::types::PyCFunction::new_closure(py, None, None, move |_, _| { + Python::attach(|py| -> PyResult<()> { + let body = body + .lock() + .expect("context body slot poisoned") + .take() + .expect("the context body ran more than once"); + *value.lock().expect("context value slot poisoned") = Some(body(py)); + Ok(()) + }) + })? + }; + copy.call_method1("run", (callback,))?; + value + .lock() + .expect("context value slot poisoned") + .take() + .ok_or_else(|| PyRuntimeError::new_err("the context body produced no value")) + } +} static GIL_RELEASES: AtomicU64 = AtomicU64::new(0); @@ -19,3 +70,270 @@ where pub fn release_count() -> u64 { GIL_RELEASES.load(Ordering::Relaxed) } + +/// Runs Python work that may block, such as a secret manager read or a callback that does +/// I/O, on the runtime's blocking pool so the async workers stay free to poll other calls. +/// The work runs inside a copy of `context` so request-local `contextvars` survive the hop. +pub async fn attach_blocking(context: PythonContext, f: F) -> PyResult +where + F: for<'py> FnOnce(Python<'py>) -> T + Send + Sync + 'static, + T: Send + 'static, +{ + match tokio::task::spawn_blocking(move || Python::attach(|py| context.enter(py, f))).await { + Ok(value) => value, + Err(error) if error.is_panic() => std::panic::resume_unwind(error.into_panic()), + Err(error) => panic!("the blocking pool dropped a Python call: {error}"), + } +} + +#[cfg(test)] +mod tests { + use std::time::{Duration, Instant}; + + use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyDict}; + use rstest::{fixture, rstest}; + + use super::{PythonContext, attach_blocking}; + use crate::{InitializedPython, initialized_python, run_sync_value}; + + #[fixture] + fn namespace(#[from(initialized_python)] python: &InitializedPython) -> Py { + python.attach(|py| { + let namespace = PyDict::new(py); + py.run( + c" +import threading, time +finished = False +def work(seconds): + global finished + time.sleep(seconds) + finished = True +def observe(expression): + return eval(expression) +", + Some(&namespace), + None, + ) + .unwrap(); + namespace.unbind() + }) + } + + fn observe FromPyObject<'a, 'py, Error: std::fmt::Debug>>( + namespace: &Py, + py: Python<'_>, + expression: &str, + ) -> T { + namespace + .bind(py) + .get_item("observe") + .unwrap() + .unwrap() + .call1((expression,)) + .unwrap() + .extract() + .unwrap() + } + + fn work(namespace: &Py, py: Python<'_>, seconds: f64) { + namespace + .bind(py) + .get_item("work") + .unwrap() + .unwrap() + .call1((seconds,)) + .unwrap(); + } + + fn shared(namespace: &Py) -> Py { + Python::attach(|py| namespace.clone_ref(py)) + } + + fn context() -> PythonContext { + Python::attach(|py| PythonContext::capture(py).unwrap()) + } + + #[fixture] + fn request_context() -> (PythonContext, Py) { + Python::initialize(); + Python::attach(|py| { + let namespace = PyDict::new(py); + py.run( + c"import contextvars\nrequest_var = contextvars.ContextVar('request_var', default='unset')", + Some(&namespace), + None, + ) + .unwrap(); + let var = namespace.get_item("request_var").unwrap().unwrap(); + var.call_method1("set", ("request-value",)).unwrap(); + (PythonContext::capture(py).unwrap(), var.unbind()) + }) + } + + #[rstest] + #[tokio::test] + async fn blocking_python_work_leaves_the_runtime_free_to_run_other_tasks( + namespace: Py, + ) { + let (python_done, timer_done) = tokio::join!( + attach_blocking(context(), move |py| { + work(&namespace, py, 0.3); + Instant::now() + }), + async { + tokio::time::sleep(Duration::from_millis(30)).await; + Instant::now() + } + ); + assert!( + timer_done < python_done.unwrap(), + "the timer only finished after the Python call: the call ran inline on the worker" + ); + } + + #[rstest] + #[tokio::test] + async fn python_work_runs_off_the_thread_polling_the_future(namespace: Py) { + let polling: u64 = Python::attach(|py| observe(&namespace, py, "threading.get_ident()")); + + let worker: u64 = attach_blocking(context(), move |py| { + observe(&namespace, py, "threading.get_ident()") + }) + .await + .unwrap(); + + assert_ne!(worker, polling); + } + + #[rstest] + #[tokio::test] + async fn a_dropped_await_never_interrupts_the_python_call(namespace: Py) { + let handle = shared(&namespace); + let started = tokio::time::timeout( + Duration::from_millis(10), + attach_blocking(context(), move |py| work(&handle, py, 0.1)), + ) + .await; + assert!( + started.is_err(), + "the await was dropped before the call returned" + ); + + tokio::time::sleep(Duration::from_millis(300)).await; + let finished: bool = Python::attach(|py| observe(&namespace, py, "finished")); + assert!(finished); + } + + #[rstest] + #[tokio::test] + async fn a_panic_in_python_work_reaches_the_awaiting_task( + #[from(initialized_python)] _python: &InitializedPython, + ) { + let joined = tokio::spawn(attach_blocking(context(), |_| -> () { + panic!("python work failed") + })) + .await; + let error = joined.expect_err("the panic propagates instead of being swallowed"); + assert!(error.is_panic()); + } + + #[rstest] + fn a_sync_route_can_await_python_work_without_deadlocking_on_the_gil( + #[from(initialized_python)] python: &InitializedPython, + ) { + let value = python + .attach(|py| { + let context = PythonContext::capture(py).unwrap(); + run_sync_value(py, async move { + tokio::time::timeout( + Duration::from_secs(5), + attach_blocking(context, |_| Python::version_str().len()), + ) + .await + .map_err(|_| { + PyRuntimeError::new_err("the blocking call never re-acquired the GIL") + }) + }) + }) + .unwrap() + .unwrap(); + assert!(value > 0); + } + + #[rstest] + #[tokio::test] + async fn blocking_work_sees_the_callers_contextvars( + request_context: (PythonContext, Py), + ) { + let (context, var) = request_context; + + let seen: String = attach_blocking(context, move |py| { + var.bind(py) + .call_method0("get") + .unwrap() + .extract::() + .unwrap() + }) + .await + .unwrap(); + + assert_eq!(seen, "request-value"); + } + + #[rstest] + #[tokio::test] + async fn concurrent_blocking_calls_each_enter_a_context_copy( + request_context: (PythonContext, Py), + ) { + let (context, var) = request_context; + let first = Python::attach(|py| var.clone_ref(py)); + let second = var; + + let (first_seen, second_seen) = tokio::join!( + attach_blocking(context.clone(), move |py| { + first + .bind(py) + .call_method0("get") + .unwrap() + .extract::() + .unwrap() + }), + attach_blocking(context, move |py| { + second + .bind(py) + .call_method0("get") + .unwrap() + .extract::() + .unwrap() + }), + ); + + assert_eq!(first_seen.unwrap(), "request-value"); + assert_eq!(second_seen.unwrap(), "request-value"); + } + + #[rstest] + #[tokio::test] + async fn writes_inside_blocking_work_do_not_leak_back_to_the_caller( + request_context: (PythonContext, Py), + ) { + let (context, var) = request_context; + let leaked = Python::attach(|py| var.clone_ref(py)); + + attach_blocking(context, move |py| { + var.bind(py).call_method1("set", ("worker-value",)).unwrap(); + }) + .await + .unwrap(); + + let caller_value: String = Python::attach(|py| { + leaked + .bind(py) + .call_method0("get") + .unwrap() + .extract() + .unwrap() + }); + assert_ne!(caller_value, "worker-value"); + } +} diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 37f8a2c9e4b..55b27e34b46 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -25,7 +25,7 @@ pub use execution::{ runtime_started, }; pub use fork_gate::RuntimeAlreadyStarted; -pub use gil::{release_count, release_gil}; +pub use gil::{PythonContext, attach_blocking, release_count, release_gil}; pub use handle::{Execution, ExecutionBody, ExecutionStep}; pub use marshal::{ Pythonized, from_py, from_py_argument, json_loads, json_object_field, panic_to_pyerr, to_py, @@ -43,3 +43,24 @@ pub(crate) fn initialize_python() { }); }); } + +#[cfg(test)] +pub(crate) struct InitializedPython; + +#[cfg(test)] +impl InitializedPython { + pub(crate) fn attach(&self, f: F) -> R + where + F: for<'py> FnOnce(pyo3::Python<'py>) -> R, + { + pyo3::Python::attach(f) + } +} + +#[cfg(test)] +#[rstest::fixture] +#[once] +pub(crate) fn initialized_python() -> InitializedPython { + initialize_python(); + InitializedPython +} diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index 79e78d150e0..e2b56e390f3 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -7,6 +7,26 @@ - Python, Rust SDK and gateway use one lifecycle-bearing core route entrypoint; provider helpers stay private, never bridge-accessible transport drivers - Built-in provider/config/secret/auth/document preparation stays in Rust; caller-authored callbacks and focused Python-file reads run only at core-selected points - Target GIL-enabled CPython explicitly with `#[pymodule(gil_used = true)]`; detach Rust-only work + - GIL and tokio invariants, each pinned by a test in `host-python` (`execution.rs`, + `gil.rs`) so a regression fails there before it deadlocks a proxy: + - Never hold the GIL while waiting on the runtime. A sync entrypoint releases it with + `release_gil` around `block_on`, because every task that attaches would otherwise wait + on the thread that is waiting on them ([pyo3 parallelism](https://pyo3.rs/v0.29.2/parallelism.html)) + - Never `block_on` from a tokio worker; the sync entrypoints refuse with "cannot run from + a Tokio context" instead of panicking inside the runtime ([tokio `Runtime::block_on`](https://docs.rs/tokio/latest/tokio/runtime/struct.Runtime.html#method.block_on)) + - Inside a future, `Python::attach` only for GIL-cheap work: cloning a `Py`, building + a small value, reading a settings snapshot. Anything that can block (a secret manager + read, a callback that does I/O, an import, a network call) goes through + `litellm_host_python::attach_blocking`, which runs it on the blocking pool so the async + workers keep polling other calls ([tokio `spawn_blocking`](https://docs.rs/tokio/latest/tokio/task/fn.spawn_blocking.html)). + `block_in_place` is not an alternative: it needs a multi-thread worker and still steals it + - `attach_blocking` work runs on a thread the interpreter did not create (pinned by the + `threading.get_ident()` test). Like any foreign-thread attach it therefore has no running + asyncio loop and a fresh `contextvars` context: do not hand it a coroutine or anything + bound to the caller's loop + - Dropping the await (an asyncio cancel) does not interrupt the Python call; it runs to + completion and its result is discarded. A panic in it reaches the awaiting task as a panic + - Add a case to `gil.rs` when a new seam changes any of these; the tests are the spec - Free-threading requires separate runtime/concurrency validation; omitting the attribute does not opt out on PyO3 0.28+ - Preserve public argument binding and Python object provenance - Project only consumed fields at reference read points; no eager whole-graph serialization or equality-based alias reconstruction @@ -56,6 +76,14 @@ GIL handling to `litellm-host-python`. decides whether to raise or fall back. For a rust-only provider/route (no Python reference), the Python side is a thin dispatch that calls Rust and raises when the bridge is unavailable, with no fallback. + - Declare it by passing `python=NO_PYTHON` (`litellm.rust_bridge.runtime`) + to `PublicDispatch.run`/`arun` or `runtime.run`/`arun`, never a stand-in + callable that raises, and give every context of it a `RUST_REQUIRED` + catalog rule + - Any other decision, an unprojectable call, or a bypass raises + `NoPythonImplementationError` before native runs, so a misdeclared route + fails in tests instead of reaching deleted code. When deleting a route's + Python implementation, switch its dispatch to `NO_PYTHON` in the same change - Keep the Python interface minimal (well under 100 lines per route): it only marshals inputs and calls Rust. Do not add per-route feature flags, and do not put provider dispatch in `litellm/main.py`; it lives in a thin dispatch diff --git a/litellm-rust/crates/python-bridge/README.md b/litellm-rust/crates/python-bridge/README.md deleted file mode 100644 index fa8df8757ff..00000000000 --- a/litellm-rust/crates/python-bridge/README.md +++ /dev/null @@ -1,10 +0,0 @@ -Native OCR uses `litellm_secrets::source::SecretSource`. Built-in secret managers resolve to retained Rust backends. Custom Python managers and overrides keep the callback path. Readable managers still require the Rust secret-manager binding to be enabled - -The shared proxy initializer captures native configuration without loading the extension or doing native I/O. `_SecretManagerRuntime.from_client` constructs a backend on first use and keeps its handle on the Python client. The secret-manager dispatcher selects Python or Rust through `catalog.py`. Native reads call that handle; Rust routes extract the backend directly. Configuration changes replace the handle, while calls already bound to the previous backend keep using it. Handles cannot be reused after fork. Directly constructed LiteLLM managers are adapted on first native use. Manually supplied SDK clients keep their Python behavior because their credentials cannot be inferred safely. Provider implementations contain no bridge registration - -Retention describes ownership and lifetime. `callbacks-legacy-python::PublicCall` owns Python references for one call to preserve identity. A native cache or secret-manager handle owns shared Rust state across calls to preserve connection pools and caches. Both use existing `Py` and shared Rust ownership, with execution and GIL transitions handled by `litellm-host-python` - - -Cache and secret-manager catalog entries remain Python-only, including when `LITELLM_RUST=1`. This wiring does not change rollout policy - -OCR provider requests use the shared `litellm-http` pool. AWS and Google secret-manager SDK clients keep their SDK transports, which do not yet inherit the pool's proxy, TLS, certificate, timeout, or observability configuration. Preserve those SDK transports and configure them equivalently instead of forcing them through reqwest diff --git a/litellm-rust/crates/python-bridge/src/secrets/callback.rs b/litellm-rust/crates/python-bridge/src/secrets/callback.rs index df9fc86f177..82bb4443f98 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/callback.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/callback.rs @@ -1,6 +1,7 @@ -use std::{future::Future, pin::Pin}; +use std::{future::Future, pin::Pin, sync::Arc}; use litellm_core_utils::settings::Lookup; +use litellm_host_python::{PythonContext, attach_blocking}; use litellm_secrets::{ Error, ExternalSecretManager, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, }; @@ -19,6 +20,11 @@ const ENVIRONMENT_FALLBACK_LOG: &str = /// A secret manager whose reads execute in Python: a custom manager, a legacy compatible /// client, or a manually assigned SDK client. pub(crate) struct PythonSecretManager { + client: Arc, + context: PythonContext, +} + +struct PythonClient { client: Py, system: Option, settings: Option>, @@ -29,14 +35,20 @@ impl PythonSecretManager { client: Py, system: Option, settings: Option>, + context: PythonContext, ) -> Self { Self { - client, - system, - settings, + client: Arc::new(PythonClient { + client, + system, + settings, + }), + context, } } +} +impl PythonClient { fn read(&self, py: Python<'_>, name: &str) -> PyResult> { let client = self.client.bind(py); let kwargs = PyDict::new(py); @@ -76,7 +88,7 @@ fn python_name(system: KeyManagementSystem) -> &'static str { impl ExternalSecretManager for PythonSecretManager { fn system(&self) -> KeyManagementSystem { - self.system.unwrap_or(KeyManagementSystem::Custom) + self.client.system.unwrap_or(KeyManagementSystem::Custom) } fn read_secret<'a>( @@ -85,18 +97,26 @@ impl ExternalSecretManager for PythonSecretManager { _settings: &'a KeyManagementSettings, _environment: &'a (dyn Lookup + Send + Sync), ) -> Pin, Error>> + Send + 'a>> { + let client = Arc::clone(&self.client); + let context = self.context.clone(); + let name = name.to_owned(); Box::pin(async move { - Python::attach(|py| match self.read(py, name) { + match attach_blocking(context, move |py| match client.read(py, &name) { Ok(value) => Ok(value.map(SecretValue::new).map(Secret::String)), // `get_secret` answers a failed manager read from the process environment, but // only for `Exception`: cancellation and other `BaseException`s propagate. Err(error) if error.is_instance_of::(py) => { - log_environment_fallback(py, name, &error) + log_environment_fallback(py, &name, &error) .map_err(|error| external_error(py, error))?; Err(read_error(py, error)) } Err(error) => Err(external_error(py, error)), }) + .await + { + Ok(result) => result, + Err(error) => Python::attach(|py| Err(external_error(py, error))), + } }) } } @@ -116,8 +136,9 @@ fn log_environment_fallback(py: Python<'_>, name: &str, error: &PyErr) -> PyResu } #[cfg(test)] +#[allow(clippy::await_holding_lock)] mod tests { - use std::sync::Arc; + use std::sync::{Arc, Mutex, MutexGuard}; use litellm_secrets::{ FailurePolicy, KeyManagementSettings, KeyManagementSystem, OidcResolver, SecretManager, @@ -126,15 +147,27 @@ mod tests { use pyo3::{prelude::*, types::PyDict}; use rstest::rstest; + use litellm_host_python::PythonContext; + use super::{HANDLER_MODULE, PythonSecretManager, python_name}; use crate::secrets::python_error; + /// `sys.modules` is interpreter-global, so tests that install or rely on the handler module + /// cannot overlap with any other test on this list. + static HANDLER_LOCK: Mutex<()> = Mutex::new(()); + + fn handler_guard() -> MutexGuard<'static, ()> { + HANDLER_LOCK.lock().expect("handler lock poisoned") + } + /// A resolver over a Python manager whose reads raise `failure_type`, with the chained - /// exceptions Python attaches, and `fallback` as the process environment. + /// exceptions Python attaches, and `fallback` as the process environment. The returned guard + /// keeps other module-mutating tests out for the lifetime of the returned resolver. fn failing_resolver( failure_type: &str, fallback: Option<&'static str>, - ) -> (SecretResolver, Py) { + ) -> (SecretResolver, Py, MutexGuard<'static, ()>) { + let handler = handler_guard(); Python::initialize(); let (reader, locals) = Python::attach(|py| { let locals = PyDict::new(py); @@ -167,6 +200,7 @@ handler.get_secret_from_manager = get_secret_from_manager locals.get_item("manager").unwrap().unwrap().unbind(), None, None, + PythonContext::capture(py).unwrap(), ); (reader, locals.unbind()) }); @@ -179,7 +213,7 @@ handler.get_secret_from_manager = get_secret_from_manager OidcResolver::default(), ) .with_failure_policy(FailurePolicy::EnvironmentFallback); - (resolver, locals) + (resolver, locals, handler) } #[rstest] @@ -191,7 +225,7 @@ handler.get_secret_from_manager = get_secret_from_manager #[case] failure_type: &str, #[case] fallback: Option<&'static str>, ) { - let (resolver, locals) = failing_resolver(failure_type, fallback); + let (resolver, locals, _handler) = failing_resolver(failure_type, fallback); let error = resolver.get_secret("API_KEY", None).await.unwrap_err(); Python::attach(|py| { let original = python_error(py, &error).unwrap(); @@ -264,7 +298,7 @@ sys.modules.setdefault('litellm._logging', logging) #[case] fallback: Option<&'static str>, #[case] name: &str, ) { - let (resolver, _locals) = failing_resolver(failure_type, fallback); + let (resolver, _locals, _handler) = failing_resolver(failure_type, fallback); Python::attach(|py| assert!(logged_errors(py, name).is_empty())); let secret = resolver.get_secret(name, None).await.unwrap(); assert_eq!( @@ -290,6 +324,7 @@ sys.modules.setdefault('litellm._logging', logging) /// Installs a fake `get_secret_from_manager` that records its kwargs, runs `body`, and /// removes the fake handler again; parent package stubs persist for concurrent tests. + /// Callers hold `handler_guard` before attaching so the GIL is never held while waiting on it. fn with_fake_handler<'py>(py: Python<'py>, body: impl FnOnce(&Bound<'py, PyDict>)) { let locals = PyDict::new(py); py.run( @@ -330,6 +365,7 @@ else: #[case("123")] #[case("{'key': 'value'}")] fn nonstring_results_are_absent_without_a_read_failure(#[case] expression: &str) { + let _handler = handler_guard(); Python::initialize(); Python::attach(|py| { with_fake_handler(py, |locals| { @@ -340,8 +376,13 @@ else: Some(locals), ) .unwrap(); - let reader = PythonSecretManager::new(py.None(), None, None); - assert_eq!(reader.read(py, "KEY").unwrap(), None); + let reader = PythonSecretManager::new( + py.None(), + None, + None, + PythonContext::capture(py).unwrap(), + ); + assert_eq!(reader.client.read(py, "KEY").unwrap(), None); }); }); } @@ -365,6 +406,7 @@ else: #[test] fn configured_systems_dispatch_through_the_python_handler_with_the_original_settings() { + let _handler = handler_guard(); Python::initialize(); Python::attach(|py| { with_fake_handler(py, |locals| { @@ -374,9 +416,10 @@ else: client.clone().unbind(), Some(KeyManagementSystem::AzureKeyVault), Some(settings.clone().unbind()), + PythonContext::capture(py).unwrap(), ); assert_eq!( - reader.read(py, "API_KEY").unwrap().as_deref(), + reader.client.read(py, "API_KEY").unwrap().as_deref(), Some("handled-API_KEY") ); assert!(py.import(HANDLER_MODULE).is_ok()); @@ -416,6 +459,7 @@ else: #[case] system: Option, #[case] key_manager: &str, ) { + let _handler = handler_guard(); Python::initialize(); Python::attach(|py| { with_fake_handler(py, |locals| { @@ -434,9 +478,14 @@ manager = Manager() ) .unwrap(); let manager = locals.get_item("manager").unwrap().unwrap(); - let reader = PythonSecretManager::new(manager.clone().unbind(), system, None); + let reader = PythonSecretManager::new( + manager.clone().unbind(), + system, + None, + PythonContext::capture(py).unwrap(), + ); assert_eq!( - reader.read(py, "API_KEY").unwrap().as_deref(), + reader.client.read(py, "API_KEY").unwrap().as_deref(), Some("handled-API_KEY") ); assert_eq!( diff --git a/litellm-rust/crates/python-bridge/src/secrets/config.rs b/litellm-rust/crates/python-bridge/src/secrets/config.rs index 4a666ca1ded..2424d9d83bf 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/config.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/config.rs @@ -5,6 +5,8 @@ use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSyst use pyo3::prelude::*; use serde_json::Value; +use litellm_host_python::PythonContext; + use super::callback::PythonSecretManager; use crate::{ coercion::{Field, FieldSpec, ProjectionError}, @@ -87,7 +89,7 @@ pub(crate) struct SecretManagerSnapshot { } impl SecretManagerSnapshot { - pub(crate) fn into_state(self) -> Arc { + pub(crate) fn into_state(self, context: PythonContext) -> Arc { match self.client { SecretManagerClient::Native(backend) => { Arc::new(SecretManagerState::new(*backend, self.settings)) @@ -98,6 +100,7 @@ impl SecretManagerSnapshot { client, self.system, self.settings_object, + context, ))), self.settings, )), diff --git a/litellm-rust/crates/python-bridge/src/secrets/mod.rs b/litellm-rust/crates/python-bridge/src/secrets/mod.rs index f0ad98dbbd8..439f9ddddd1 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/mod.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/mod.rs @@ -4,102 +4,30 @@ mod error; mod mutation; mod operations; mod provider; +mod python; pub(crate) mod resolved; pub(crate) mod runtime; mod vault; use std::sync::Arc; -use litellm_secrets::source::{EnvironmentSecrets, SecretSource}; -use pyo3::prelude::*; - pub(crate) use error::python_error; +use litellm_secrets::source::SecretSource; +use pyo3::prelude::*; +use python::PythonSecrets; use resolved::ResolvedSecrets; -use crate::{ - coercion::FieldSpec, - errors::RustBridgeDeclined, - python_settings::{PythonSettings, Snapshot}, -}; +use crate::{coercion::FieldSpec, python_settings::PythonSettings}; -const READABLE: FieldSpec = FieldSpec::new("readable", |field| field.schema_bool()); const NATIVE: FieldSpec = FieldSpec::new("native", |field| field.schema_bool()); -/// Where a Rust route reads provider secrets from, as `litellm.get_secret` would. +/// Where a Rust route reads provider secrets from. Python's `get_secret_str` until a +/// `SecretManagerRule` in `catalog.py` moves the configured system off `PYTHON_ONLY`, then the +/// native secret manager. pub(crate) fn source(py: Python<'_>) -> PyResult> { - select(&PythonSettings::SecretManager.read(py)?, || { - Ok(Arc::new(ResolvedSecrets::new(config::read(py)?))) - }) -} - -fn select( - manager: &Snapshot<'_>, - resolved: impl FnOnce() -> PyResult>, -) -> PyResult> { - if !manager.read(&READABLE)? { - return Ok(Arc::new(EnvironmentSecrets::python_compatible())); - } - if !manager.read(&NATIVE)? { - return Err(RustBridgeDeclined::new_err( - "the configured secret manager is not enabled for the Rust bridge", - )); - } - resolved() -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use litellm_secrets::source::{EnvironmentSecrets, SecretSource}; - use pyo3::{prelude::*, types::PyDict}; - use rstest::rstest; - - use super::select; - use crate::{errors::RustBridgeDeclined, python_settings::PythonSettings}; - - enum Selected { - Environment, - Declined, - Resolved, - } - - #[rstest] - #[case::unreadable(false, false, Selected::Environment)] - #[case::unreadable_even_if_native(false, true, Selected::Environment)] - #[case::readable_python_only(true, false, Selected::Declined)] - #[case::readable_native(true, true, Selected::Resolved)] - fn readable_and_native_select_the_secret_source( - #[case] readable: bool, - #[case] native: bool, - #[case] expected: Selected, - ) { - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - locals.set_item("readable", readable).unwrap(); - locals.set_item("native", native).unwrap(); - let manager = py - .eval( - c"__import__('types').SimpleNamespace(readable=readable, native=native)", - None, - Some(&locals), - ) - .unwrap(); - let mut resolved_called = false; - let selected = select(&PythonSettings::SecretManager.snapshot(manager), || { - resolved_called = true; - Ok(Arc::new(EnvironmentSecrets::python_compatible()) as Arc) - }); - match expected { - Selected::Environment => assert!(selected.is_ok() && !resolved_called), - Selected::Resolved => assert!(selected.is_ok() && resolved_called), - Selected::Declined => { - let error = selected.err().expect("the Rust route declines"); - assert!(error.is_instance_of::(py)); - assert!(!resolved_called); - } - } - }); - } + if PythonSettings::SecretManager.read(py)?.read(&NATIVE)? { + let context = litellm_host_python::PythonContext::capture(py)?; + return Ok(Arc::new(ResolvedSecrets::new(config::read(py)?, context))); + } + Ok(Arc::new(PythonSecrets::new(py)?)) } diff --git a/litellm-rust/crates/python-bridge/src/secrets/python.rs b/litellm-rust/crates/python-bridge/src/secrets/python.rs new file mode 100644 index 00000000000..fb26b1e80fc --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/secrets/python.rs @@ -0,0 +1,190 @@ +use std::sync::Arc; + +use futures_util::future::BoxFuture; +use litellm_host_python::{PythonContext, attach_blocking}; +use litellm_secrets::{Error, SecretValue, source::SecretSource}; +use pyo3::prelude::*; + +use super::error::external_error; + +/// Reads each secret through Python's `get_secret_str`, so the configured manager, the key +/// management settings and the environment fallback behave exactly as they do in Python. +pub(super) struct PythonSecrets { + get_secret_str: Arc>, + context: PythonContext, +} + +impl PythonSecrets { + pub(super) fn new(py: Python<'_>) -> PyResult { + Ok(Self::reading_with( + py.import("litellm.secret_managers.main")? + .getattr("get_secret_str")? + .unbind(), + PythonContext::capture(py)?, + )) + } + + fn reading_with(get_secret_str: Py, context: PythonContext) -> Self { + Self { + get_secret_str: Arc::new(get_secret_str), + context, + } + } +} + +impl SecretSource for PythonSecrets { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result, Error>> { + let get_secret_str = Arc::clone(&self.get_secret_str); + let context = self.context.clone(); + let name = name.to_owned(); + Box::pin(async move { + match attach_blocking(context, move |py| { + get_secret_str + .bind(py) + .call1((name,)) + .and_then(|value| value.extract::>()) + .map(|value| value.map(SecretValue::new)) + .map_err(|error| external_error(py, error)) + }) + .await + { + Ok(result) => result, + Err(error) => Python::attach(|py| Err(external_error(py, error))), + } + }) + } +} + +#[cfg(test)] +mod tests { + use litellm_secrets::source::SecretSource; + use pyo3::{prelude::*, types::PyDict}; + use rstest::{fixture, rstest}; + + use super::PythonSecrets; + use crate::secrets::python_error; + use litellm_host_python::PythonContext; + + #[fixture] + fn namespace() -> Py { + Python::initialize(); + Python::attach(|py| { + let namespace = PyDict::new(py); + py.run( + c" +import contextvars +import threading +read_on = None +request_var = contextvars.ContextVar('request_var', default=None) +seen_context_values = [] +raised = KeyboardInterrupt('secret manager stopped') +def get_secret_str(name): + global read_on + read_on = threading.get_ident() + seen_context_values.append(request_var.get()) + if name == 'RAISING': + raise raised + return {'MISTRAL_API_KEY': 'vault-key'}.get(name) +", + Some(&namespace), + None, + ) + .unwrap(); + namespace.unbind() + }) + } + + #[fixture] + fn secrets(namespace: Py) -> (PythonSecrets, Py) { + let (reader, context) = Python::attach(|py| { + let namespace = namespace.bind(py); + namespace + .get_item("request_var") + .unwrap() + .unwrap() + .call_method1("set", ("request-value",)) + .unwrap(); + ( + namespace + .get_item("get_secret_str") + .unwrap() + .unwrap() + .unbind(), + PythonContext::capture(py).unwrap(), + ) + }); + (PythonSecrets::reading_with(reader, context), namespace) + } + + fn global FromPyObject<'a, 'py, Error: std::fmt::Debug>>( + namespace: &Py, + py: Python<'_>, + name: &str, + ) -> T { + namespace + .bind(py) + .get_item(name) + .unwrap() + .unwrap() + .extract() + .unwrap() + } + + #[rstest] + #[case::found("MISTRAL_API_KEY", Some("vault-key"))] + #[case::missing("OTHER", None)] + #[tokio::test] + async fn returns_what_get_secret_str_returns( + secrets: (PythonSecrets, Py), + #[case] name: &str, + #[case] expected: Option<&str>, + ) { + let value = secrets.0.get_secret_str(name).await.unwrap(); + + assert_eq!(value.as_ref().map(|value| value.expose()), expected); + } + + #[rstest] + #[tokio::test] + async fn exceptions_surface_as_the_original_python_object( + secrets: (PythonSecrets, Py), + ) { + let error = secrets.0.get_secret_str("RAISING").await.unwrap_err(); + + Python::attach(|py| { + let surfaced = python_error(py, &error).expect("the Python exception is preserved"); + let raised: Py = global(&secrets.1, py, "raised"); + assert!(surfaced.value(py).is(raised.bind(py))); + }); + } + + #[rstest] + #[tokio::test] + async fn reads_run_off_the_thread_polling_the_route(secrets: (PythonSecrets, Py)) { + let polling: u64 = Python::attach(|py| { + py.import("threading") + .unwrap() + .call_method0("get_ident") + .unwrap() + .extract() + .unwrap() + }); + + secrets.0.get_secret_str("MISTRAL_API_KEY").await.unwrap(); + + let read_on: u64 = Python::attach(|py| global(&secrets.1, py, "read_on")); + assert_ne!(read_on, polling); + } + + #[rstest] + #[tokio::test] + async fn reads_see_the_callers_contextvars(secrets: (PythonSecrets, Py)) { + secrets.0.get_secret_str("MISTRAL_API_KEY").await.unwrap(); + + let seen: Vec = Python::attach(|py| global(&secrets.1, py, "seen_context_values")); + assert_eq!(seen, vec!["request-value".to_owned()]); + } +} diff --git a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs index 84594c5a3dd..40b187c99de 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use futures_util::future::BoxFuture; use litellm_core_utils::settings::ProcessEnvironment; +use litellm_host_python::PythonContext; use litellm_secrets::source::SecretSource; use litellm_secrets::{ Error, FailurePolicy, OidcResolver, SecretManagerState, SecretResolver, SecretValue, @@ -14,8 +15,8 @@ pub(crate) struct ResolvedSecrets { } impl ResolvedSecrets { - pub(crate) fn new(snapshot: SecretManagerSnapshot) -> Self { - Self::from_state(snapshot.into_state()) + pub(crate) fn new(snapshot: SecretManagerSnapshot, context: PythonContext) -> Self { + Self::from_state(snapshot.into_state(context)) } fn from_state(state: Arc) -> Self { diff --git a/litellm/llms/bedrock/audio_transcription/__init__.py b/litellm/llms/bedrock/audio_transcription/__init__.py index 172316027a8..c62587566c0 100644 --- a/litellm/llms/bedrock/audio_transcription/__init__.py +++ b/litellm/llms/bedrock/audio_transcription/__init__.py @@ -1,5 +1,5 @@ import base64 -from typing import Final, NoReturn +from typing import Final import httpx @@ -16,14 +16,6 @@ from litellm.rust_bridge.transcription.native import ( from litellm.types.utils import FileTypes, TranscriptionResponse -def _no_python_implementation() -> NoReturn: - raise NotImplementedError("Bedrock audio transcription is implemented in Rust only") - - -async def _no_async_python_implementation() -> NoReturn: - _no_python_implementation() - - class BedrockAudioTranscriptionRustDispatch: @staticmethod def _audio_payload(audio_file: FileTypes) -> dict[str, object]: @@ -77,7 +69,7 @@ class BedrockAudioTranscriptionRustDispatch: RouteContext(Route.TRANSCRIPTION, provider=custom_llm_provider, model=model), binding=NATIVE_TRANSCRIPTION, native=native, - python=_no_python_implementation, + python=runtime.NO_PYTHON, ) async def async_audio_transcriptions( @@ -110,5 +102,5 @@ class BedrockAudioTranscriptionRustDispatch: RouteContext(Route.TRANSCRIPTION, provider=custom_llm_provider, model=model), binding=NATIVE_ATRANSCRIPTION, native=native, - python=_no_async_python_implementation, + python=runtime.NO_PYTHON, ) diff --git a/litellm/rust_bridge/dispatch.py b/litellm/rust_bridge/dispatch.py index 94ccddc92b7..d0a39b18c53 100644 --- a/litellm/rust_bridge/dispatch.py +++ b/litellm/rust_bridge/dispatch.py @@ -44,17 +44,39 @@ class PublicDispatch(Generic[RequestT]): return rollout_decision(rule.rollout) is not Decision.PYTHON return False + def _native_request(self, args: tuple[object, ...], kwargs: Mapping[str, object]) -> RequestT: + request: Final = self.request(args, kwargs) + if request is None: + raise runtime.NoPythonImplementationError( + f"{self.route.value} has no Python implementation, so every call must project to a native request" + ) + if self.bypass is not None and self.bypass(request): + raise runtime.NoPythonImplementationError( + f"{self.route.value} has no Python implementation, so a call its bypass predicate matches cannot " + "be served" + ) + return request + def run( self, args: tuple[object, ...], kwargs: Mapping[str, object], *, - python: Callable[..., ResultT], + python: Callable[..., ResultT] | runtime.NoPythonImplementation, binding: NativeBinding[NativeT], native: Callable[[NativeT, RequestT, tuple[object, ...], Mapping[str, object]], ResultT], rules: Rules | None = None, ) -> ResultT: selected_rules: Final = catalog.RULES if rules is None else rules + if isinstance(python, runtime.NoPythonImplementation): + native_request: Final = self._native_request(args, kwargs) + return runtime.run( + self.context(native_request), + binding=binding, + native=lambda hook: native(hook, native_request, args, kwargs), + python=python, + rules=selected_rules, + ) if not self._requires_projection(selected_rules): return python(*args, **kwargs) request: Final = self.request(args, kwargs) @@ -73,12 +95,21 @@ class PublicDispatch(Generic[RequestT]): args: tuple[object, ...], kwargs: Mapping[str, object], *, - python: Callable[..., Awaitable[ResultT]], + python: Callable[..., Awaitable[ResultT]] | runtime.NoPythonImplementation, binding: NativeBinding[NativeT], native: Callable[[NativeT, RequestT, tuple[object, ...], Mapping[str, object]], Awaitable[ResultT]], rules: Rules | None = None, ) -> ResultT: selected_rules: Final = catalog.RULES if rules is None else rules + if isinstance(python, runtime.NoPythonImplementation): + native_request: Final = self._native_request(args, kwargs) + return await runtime.arun( + self.context(native_request), + binding=binding, + native=lambda hook: native(hook, native_request, args, kwargs), + python=python, + rules=selected_rules, + ) if not self._requires_projection(selected_rules): return await python(*args, **kwargs) request: Final = self.request(args, kwargs) diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py index cf02a33eab6..cf632c09074 100644 --- a/litellm/rust_bridge/runtime.py +++ b/litellm/rust_bridge/runtime.py @@ -41,29 +41,37 @@ class BridgeErrorContext: model: str +@dataclass(frozen=True, slots=True) +class NoPythonImplementation: + pass + + +NO_PYTHON: Final = NoPythonImplementation() + + +class NoPythonImplementationError(RuntimeError): + pass + + def run( context: RouteContext, *, binding: NativeBinding[NativeT], native: Callable[[NativeT], ResultT], - python: Callable[[], ResultT], + python: Callable[[], ResultT] | NoPythonImplementation, rules: Rules | None = None, ) -> ResultT: selected: Final = decision(context, rules) + if isinstance(python, NoPythonImplementation): + _require_rust(context, selected) + return _required(_attempt_native(context, binding, native), context) match selected: case Decision.PYTHON: return python() case Decision.RUST_WITH_FALLBACK | Decision.RUST_REQUIRED: - loaded: Final = binding.load() - result: Final = attempt( - native_call=None if loaded is None else lambda: native(loaded), - adapt=_identity, - context=_error_context(context), - ) - if isinstance(result, RustHandled): - return mark_rust_response(result.value) - if selected is Decision.RUST_REQUIRED: - _raise_required(result, _error_context(context)) + result: Final = _attempt_native(context, binding, native) + if isinstance(result, RustHandled) or selected is Decision.RUST_REQUIRED: + return _required(result, context) return python() case _: assert_never(selected) @@ -74,29 +82,61 @@ async def arun( *, binding: NativeBinding[NativeT], native: Callable[[NativeT], Awaitable[ResultT]], - python: Callable[[], Awaitable[ResultT]], + python: Callable[[], Awaitable[ResultT]] | NoPythonImplementation, rules: Rules | None = None, ) -> ResultT: selected: Final = decision(context, rules) + if isinstance(python, NoPythonImplementation): + _require_rust(context, selected) + return _required(await _aattempt_native(context, binding, native), context) match selected: case Decision.PYTHON: return await python() case Decision.RUST_WITH_FALLBACK | Decision.RUST_REQUIRED: - loaded: Final = binding.load() - result: Final = await aattempt( - native_call=None if loaded is None else lambda: native(loaded), - adapt=_identity, - context=_error_context(context), - ) - if isinstance(result, RustHandled): - return mark_rust_response(result.value) - if selected is Decision.RUST_REQUIRED: - _raise_required(result, _error_context(context)) + result: Final = await _aattempt_native(context, binding, native) + if isinstance(result, RustHandled) or selected is Decision.RUST_REQUIRED: + return _required(result, context) return await python() case _: assert_never(selected) +def _require_rust(context: RouteContext, selected: Decision) -> None: + if selected is not Decision.RUST_REQUIRED: + raise NoPythonImplementationError( + f"{context.route.value} has no Python implementation, so its catalog rules must resolve to " + f"RUST_REQUIRED, but provider={context.provider!r} model={context.model!r} resolved to {selected.name}" + ) + + +def _attempt_native( + context: RouteContext, binding: NativeBinding[NativeT], native: Callable[[NativeT], ResultT] +) -> RustAttempt[ResultT]: + loaded: Final = binding.load() + return attempt( + native_call=None if loaded is None else lambda: native(loaded), + adapt=_identity, + context=_error_context(context), + ) + + +async def _aattempt_native( + context: RouteContext, binding: NativeBinding[NativeT], native: Callable[[NativeT], Awaitable[ResultT]] +) -> RustAttempt[ResultT]: + loaded: Final = binding.load() + return await aattempt( + native_call=None if loaded is None else lambda: native(loaded), + adapt=_identity, + context=_error_context(context), + ) + + +def _required(result: RustAttempt[ResultT], context: RouteContext) -> ResultT: + if isinstance(result, RustHandled): + return mark_rust_response(result.value) + _raise_required(result, _error_context(context)) + + def _identity(value: ResultT) -> ResultT: return value diff --git a/tests/test_litellm/rust_bridge/ocr/test_secrets.py b/tests/test_litellm/rust_bridge/ocr/test_secrets.py index a91c0ff5bc8..f6746b97127 100644 --- a/tests/test_litellm/rust_bridge/ocr/test_secrets.py +++ b/tests/test_litellm/rust_bridge/ocr/test_secrets.py @@ -222,22 +222,59 @@ async def test_custom_secret_manager_cancellation_propagates_without_provider_io assert raised.value is failure -async def test_rust_declines_a_readable_secret_manager_it_cannot_resolve( +async def test_rust_reads_a_python_only_secret_manager_through_python( monkeypatch: pytest.MonkeyPatch, rust_ocr: Ocr ) -> None: manager: Final = _VaultSecrets() - _configure( - monkeypatch, - manager=manager, - key_management=KeyManagementSettings(access_mode="read_only"), - native_secret_manager=False, - ) + key_management: Final = KeyManagementSettings(access_mode="read_only") + _configure(monkeypatch, manager=manager, key_management=key_management, native_secret_manager=False) - with _mistral_service(expected_requests=0) as server: - with pytest.raises(native.RustBridgeDeclined): - await rust_ocr(server.base_url) + with _mistral_service() as server: + await rust_ocr(server.base_url) - assert manager.key_reads() == () + assert server.requests[0].headers["authorization"] == "Bearer vault-key" + assert manager.key_reads(), "the Python manager was never asked for MISTRAL_API_KEY" + assert all(params == key_management.model_dump() for params in manager.key_reads()), manager.key_reads() + + +class _PythonManagerReads: + def __init__(self) -> None: + self.reads: tuple[tuple[object, str, str], ...] = () + + def __call__( + self, + client: object, + key_manager: str, + secret_name: str, + key_management_settings: KeyManagementSettings | None = None, + ) -> str | None: + self.reads = (*self.reads, (client, key_manager, secret_name)) + return "python-key" if secret_name == "MISTRAL_API_KEY" else None + + +async def test_python_only_builtin_manager_is_read_through_python_not_its_native_backend( + monkeypatch: pytest.MonkeyPatch, rust_ocr: Ocr +) -> None: + from litellm.secret_managers import main as secret_manager_main + from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 + + python_reads: Final = _PythonManagerReads() + monkeypatch.setattr(secret_manager_main, "get_secret_from_manager", python_reads) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "native-access") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "native-secret") + monkeypatch.delenv("MISTRAL_API_KEY", raising=False) + manager: Final = AWSSecretsManagerV2(aws_region_name="us-east-1") + monkeypatch.setattr(litellm, "secret_manager_client", manager) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.AWS_SECRET_MANAGER) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(hosted_keys=["MISTRAL_API_KEY"])) + monkeypatch.setattr(settings, "secret_manager", lambda: settings.SecretManager(readable=True, native=False)) + + with _mistral_service() as provider: + await rust_ocr(provider.base_url) + + assert provider.requests[0].headers["authorization"] == "Bearer python-key" + assert (manager, "aws_secret_manager", "MISTRAL_API_KEY") in python_reads.reads + assert getattr(manager, "_litellm_native_secret_manager", None) is None async def test_no_secret_client_leaves_dormant_binding_settings_unread( diff --git a/tests/test_litellm/rust_bridge/test_dispatch.py b/tests/test_litellm/rust_bridge/test_dispatch.py index 9a3793a772e..0dc06b2905f 100644 --- a/tests/test_litellm/rust_bridge/test_dispatch.py +++ b/tests/test_litellm/rust_bridge/test_dispatch.py @@ -1,6 +1,6 @@ from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator, Mapping from dataclasses import dataclass -from typing import Final +from typing import Final, TypeAlias import pytest @@ -9,6 +9,7 @@ from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.catalog import CacheRule, Delivery, Route, RouteContext, RouteRule, Rules, SecretManagerRule from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.dispatch import PublicDispatch +from litellm.rust_bridge.runtime import NO_PYTHON, NoPythonImplementationError @dataclass(frozen=True, slots=True) @@ -233,3 +234,84 @@ async def test_async_bypass_forwards_to_python_without_native() -> None: rules=rules, ) assert result is expected + + +NativeRoute: TypeAlias = Callable[[Request, tuple[object, ...], Mapping[str, object]], object] + + +def native_route(result: object) -> NativeBinding[NativeRoute]: + bound: Final[NativeBinding[NativeRoute]] = NativeBinding("no_python", validate=lambda _: None) + bound.override(lambda request, args, kwargs: (result, request, args, dict(kwargs))) + return bound + + +async def dispatch_without_python( + dispatch: PublicDispatch[Request], bound: NativeBinding[NativeRoute], rules: Rules, *, asynchronous: bool +) -> object: + if not asynchronous: + return dispatch.run( + ("model",), + {"page": 1}, + python=NO_PYTHON, + binding=bound, + native=lambda hook, value, args, kwargs: hook(value, args, kwargs), + rules=rules, + ) + + async def native( + hook: NativeRoute, value: Request, args: tuple[object, ...], kwargs: Mapping[str, object] + ) -> object: + return hook(value, args, kwargs) + + return await dispatch.arun(("model",), {"page": 1}, python=NO_PYTHON, binding=bound, native=native, rules=rules) + + +def ocr_dispatch(request: Request | None, *, bypass: bool = False) -> PublicDispatch[Request]: + return PublicDispatch( + route=Route.OCR, + request=lambda args, kwargs: request, + context=lambda value: RouteContext(Route.OCR, model=value.model), + bypass=lambda _: bypass, + ) + + +@pytest.mark.parametrize("asynchronous", (False, True)) +@pytest.mark.parametrize("switch", (None, False)) +async def test_dispatch_without_python_hands_every_call_to_native(asynchronous: bool, switch: bool | None) -> None: + request: Final = Request(model="model") + result: Final = object() + configuration.rust(switch) + try: + dispatched: Final = await dispatch_without_python( + ocr_dispatch(request), + native_route(result), + (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),), + asynchronous=asynchronous, + ) + finally: + configuration.rust(None) + + assert dispatched == (result, request, ("model",), {"page": 1}) + + +@pytest.mark.parametrize("asynchronous", (False, True)) +@pytest.mark.parametrize( + ("request_value", "bypass", "rules", "reason"), + ( + (Request(model="model"), False, (), "must resolve to RUST_REQUIRED"), + (Request(model="model"), False, (RouteRule(Route.OCR, Rollout.RUST_OPT_OUT),), "must resolve to RUST_REQUIRED"), + (None, False, (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),), "must project to a native request"), + (Request(model="model"), True, (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),), "bypass predicate matches"), + ), + ids=("no-rule", "opt-out-rule", "unprojectable-call", "bypassed-call"), +) +async def test_dispatch_without_python_never_falls_back( + asynchronous: bool, request_value: Request | None, bypass: bool, rules: Rules, reason: str +) -> None: + bound: Final[NativeBinding[NativeRoute]] = NativeBinding("no_python", validate=lambda _: None) + bound.override(lambda request, args, kwargs: pytest.fail("a misdeclared route must not reach native")) + + with pytest.raises(NoPythonImplementationError, match=f"ocr has no Python implementation, so .*{reason}"): + await dispatch_without_python( + ocr_dispatch(request_value, bypass=bypass), bound, rules, asynchronous=asynchronous + ) diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/test_litellm/rust_bridge/test_runtime.py index bbe25e0de13..33c8e6112a6 100644 --- a/tests/test_litellm/rust_bridge/test_runtime.py +++ b/tests/test_litellm/rust_bridge/test_runtime.py @@ -427,3 +427,80 @@ async def test_arun_upstream_error_maps_to_api_error_without_fallback() -> None: assert caught.value.status_code == 503 assert calls.calls == (RUST,) + + +async def run_without_python( + rollout: Rollout, + calls: Recorder, + *, + asynchronous: bool, + native_missing: bool = False, + context: RouteContext = CONTEXT, +) -> str: + bound: Final = binding(None if native_missing else calls.rust) + if not asynchronous: + return runtime.run( + context, binding=bound, native=lambda fn: fn(), python=runtime.NO_PYTHON, rules=rules(rollout) + ) + + async def native(fn: NativeFn) -> str: + return fn() + + return await runtime.arun(context, binding=bound, native=native, python=runtime.NO_PYTHON, rules=rules(rollout)) + + +@pytest.mark.parametrize("asynchronous", (False, True)) +@pytest.mark.parametrize("switch", (None, False, True)) +async def test_route_without_python_runs_native_whatever_the_rust_switch( + asynchronous: bool, switch: bool | None +) -> None: + calls: Final = recorder() + if switch is not None: + configuration.rust(switch) + + assert await run_without_python(Rollout.RUST_REQUIRED, calls, asynchronous=asynchronous) == RUST + assert calls.calls == (RUST,) + + +@pytest.mark.parametrize("asynchronous", (False, True)) +@pytest.mark.parametrize( + ("native_missing", "effect", "message"), + ( + (True, None, "Rust messages bridge is unavailable"), + (False, RustBridgeDeclined("unsupported"), "Rust messages bridge declined the request: unsupported"), + ), +) +async def test_route_without_python_raises_when_native_cannot_serve_the_call( + asynchronous: bool, native_missing: bool, effect: BaseException | None, message: str +) -> None: + calls: Final = recorder(effect) + + with pytest.raises(RuntimeError, match=message) as raised: + await run_without_python(Rollout.RUST_REQUIRED, calls, asynchronous=asynchronous, native_missing=native_missing) + + assert not isinstance(raised.value, runtime.NoPythonImplementationError) + + +@pytest.mark.parametrize("asynchronous", (False, True)) +@pytest.mark.parametrize("switch", (None, False, True)) +@pytest.mark.parametrize( + ("rollout", "context"), + ( + (Rollout.PYTHON_ONLY, CONTEXT), + (Rollout.RUST_OPT_IN, CONTEXT), + (Rollout.RUST_OPT_OUT, CONTEXT), + (Rollout.RUST_REQUIRED, RouteContext(Route.MESSAGES, provider="unmatched", model="model")), + ), + ids=("python-only", "opt-in", "opt-out", "no-matching-rule"), +) +async def test_route_without_python_rejects_rules_that_could_select_python( + asynchronous: bool, switch: bool | None, rollout: Rollout, context: RouteContext +) -> None: + calls: Final = recorder() + if switch is not None: + configuration.rust(switch) + + with pytest.raises(runtime.NoPythonImplementationError, match="messages has no Python implementation"): + await run_without_python(rollout, calls, asynchronous=asynchronous, context=context) + + assert calls.calls == ()