mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
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 <yujong@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
7bdccd7371
commit
f1ef7fc0c2
17 changed files with 956 additions and 186 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<F, R>(&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);
|
||||
|
|
|
|||
|
|
@ -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<Py<PyAny>>);
|
||||
|
||||
impl PythonContext {
|
||||
pub fn capture(py: Python<'_>) -> PyResult<Self> {
|
||||
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<T, F>(&self, py: Python<'_>, f: F) -> PyResult<T>
|
||||
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::<T>));
|
||||
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<T, F>(context: PythonContext, f: F) -> PyResult<T>
|
||||
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<PyDict> {
|
||||
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<T: for<'a, 'py> FromPyObject<'a, 'py, Error: std::fmt::Debug>>(
|
||||
namespace: &Py<PyDict>,
|
||||
py: Python<'_>,
|
||||
expression: &str,
|
||||
) -> T {
|
||||
namespace
|
||||
.bind(py)
|
||||
.get_item("observe")
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.call1((expression,))
|
||||
.unwrap()
|
||||
.extract()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn work(namespace: &Py<PyDict>, py: Python<'_>, seconds: f64) {
|
||||
namespace
|
||||
.bind(py)
|
||||
.get_item("work")
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.call1((seconds,))
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
fn shared(namespace: &Py<PyDict>) -> Py<PyDict> {
|
||||
Python::attach(|py| namespace.clone_ref(py))
|
||||
}
|
||||
|
||||
fn context() -> PythonContext {
|
||||
Python::attach(|py| PythonContext::capture(py).unwrap())
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn request_context() -> (PythonContext, Py<PyAny>) {
|
||||
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<PyDict>,
|
||||
) {
|
||||
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<PyDict>) {
|
||||
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<PyDict>) {
|
||||
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<PyAny>),
|
||||
) {
|
||||
let (context, var) = request_context;
|
||||
|
||||
let seen: String = attach_blocking(context, move |py| {
|
||||
var.bind(py)
|
||||
.call_method0("get")
|
||||
.unwrap()
|
||||
.extract::<String>()
|
||||
.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<PyAny>),
|
||||
) {
|
||||
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::<String>()
|
||||
.unwrap()
|
||||
}),
|
||||
attach_blocking(context, move |py| {
|
||||
second
|
||||
.bind(py)
|
||||
.call_method0("get")
|
||||
.unwrap()
|
||||
.extract::<String>()
|
||||
.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<PyAny>),
|
||||
) {
|
||||
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");
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<F, R>(&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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<T>`, 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
|
||||
|
|
|
|||
|
|
@ -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<T>` 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
|
||||
|
|
@ -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<PythonClient>,
|
||||
context: PythonContext,
|
||||
}
|
||||
|
||||
struct PythonClient {
|
||||
client: Py<PyAny>,
|
||||
system: Option<KeyManagementSystem>,
|
||||
settings: Option<Py<PyAny>>,
|
||||
|
|
@ -29,14 +35,20 @@ impl PythonSecretManager {
|
|||
client: Py<PyAny>,
|
||||
system: Option<KeyManagementSystem>,
|
||||
settings: Option<Py<PyAny>>,
|
||||
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<Option<String>> {
|
||||
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<Box<dyn Future<Output = Result<Option<Secret>, 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::<PyException>(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<PyDict>) {
|
||||
) -> (SecretResolver, Py<PyDict>, 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<KeyManagementSystem>,
|
||||
#[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!(
|
||||
|
|
|
|||
|
|
@ -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<SecretManagerState> {
|
||||
pub(crate) fn into_state(self, context: PythonContext) -> Arc<SecretManagerState> {
|
||||
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,
|
||||
)),
|
||||
|
|
|
|||
|
|
@ -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<bool> = FieldSpec::new("readable", |field| field.schema_bool());
|
||||
const NATIVE: FieldSpec<bool> = 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<Arc<dyn SecretSource>> {
|
||||
select(&PythonSettings::SecretManager.read(py)?, || {
|
||||
Ok(Arc::new(ResolvedSecrets::new(config::read(py)?)))
|
||||
})
|
||||
}
|
||||
|
||||
fn select(
|
||||
manager: &Snapshot<'_>,
|
||||
resolved: impl FnOnce() -> PyResult<Arc<dyn SecretSource>>,
|
||||
) -> PyResult<Arc<dyn SecretSource>> {
|
||||
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<dyn SecretSource>)
|
||||
});
|
||||
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::<RustBridgeDeclined>(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)?))
|
||||
}
|
||||
|
|
|
|||
190
litellm-rust/crates/python-bridge/src/secrets/python.rs
Normal file
190
litellm-rust/crates/python-bridge/src/secrets/python.rs
Normal file
|
|
@ -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<Py<PyAny>>,
|
||||
context: PythonContext,
|
||||
}
|
||||
|
||||
impl PythonSecrets {
|
||||
pub(super) fn new(py: Python<'_>) -> PyResult<Self> {
|
||||
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<PyAny>, 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<Option<SecretValue>, 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::<Option<String>>())
|
||||
.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<PyDict> {
|
||||
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<PyDict>) -> (PythonSecrets, Py<PyDict>) {
|
||||
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<T: for<'a, 'py> FromPyObject<'a, 'py, Error: std::fmt::Debug>>(
|
||||
namespace: &Py<PyDict>,
|
||||
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<PyDict>),
|
||||
#[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<PyDict>),
|
||||
) {
|
||||
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<PyAny> = 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<PyDict>)) {
|
||||
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<PyDict>)) {
|
||||
secrets.0.get_secret_str("MISTRAL_API_KEY").await.unwrap();
|
||||
|
||||
let seen: Vec<String> = Python::attach(|py| global(&secrets.1, py, "seen_context_values"));
|
||||
assert_eq!(seen, vec!["request-value".to_owned()]);
|
||||
}
|
||||
}
|
||||
|
|
@ -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<SecretManagerState>) -> Self {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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 == ()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue