scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-16 04:29:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 225

refactor(host-python): run the Python driver as one async body (#46235)

Browse source 
* refactor(host-python): run the Python driver as one async body

Co-Authored-By: Yujong Lee <yujong@berri.ai>

* refactor(host-python): give the driver a PythonCaller with named suspensions

PythonTask becomes PythonCaller with await_python, open_stream and yield_chunk, so the call body no longer builds Suspension values or maps stream replies itself

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
yujonglee 2026-10-14 17:19:06 -07:00 • committed by GitHub
parent ab0875b2da
commit 7d4be1e217
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 911 additions and 853 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm-rust/Cargo.lock generated
View file

@ -3922,7 +3922,9 @@ version = "0.1.0"
dependencies = [
"bytes",
"futures-util",
"litellm-coroutine",
"litellm-host",
"parking_lot",
"pyo3",
"pyo3-async-runtimes",
"pythonize",

1
litellm-rust/Cargo.toml
View file

@ -80,6 +80,7 @@ http = "1"
google-cloud-auth = { version = "1.16.0", default-features = false }
jsonwebtoken = { version = "11.1.0", default-features = false }
hyper-util = { version = "0.1.20", default-features = false, features = ["client-proxy"] }
parking_lot = "0.12.5"
proptest = "1.7.0"
pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }

2
litellm-rust/crates/host-python/Cargo.toml
View file

@ -6,10 +6,12 @@ license.workspace = true
repository.workspace = true
[dependencies]
litellm-coroutine.workspace = true
litellm-host.workspace = true
bytes.workspace = true
futures-util.workspace = true
parking_lot.workspace = true
serde.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["rt", "sync"] }

58
litellm-rust/crates/host-python/src/caller.rs Normal file
View file

@ -0,0 +1,58 @@
use std::ops::ControlFlow;
use litellm_coroutine::{Abandoned, Co, Reply};
use pyo3::prelude::*;
use crate::Error;
use crate::handle::Suspension;
pub(crate) type Answer = PyResult<Py<PyAny>>;
pub(crate) struct Yielded {
pub(crate) suspension: Suspension,
pub(crate) reply: Reply<Answer>,
}
/// Hands suspensions to the Python caller and receives its answers.
#[derive(Clone)]
pub(crate) struct PythonCaller {
co: Co<Yielded>,
}
impl PythonCaller {
/// Connects the call body to its coroutine's yield handle.
pub(crate) fn new(co: Co<Yielded>) -> Self {
Self { co }
}
/// Asks the caller to await Python work and return its value or exception.
pub(crate) async fn await_python(&self, awaitable: Py<PyAny>) -> Answer {
self.suspend(Suspension::Await(awaitable)).await
}
/// Hands the caller a stream head and waits for demand or closure.
pub(crate) async fn open_stream(&self, head: Py<PyAny>) -> ControlFlow<()> {
self.stream_boundary(Suspension::Open(head)).await
}
/// Delivers a chunk and waits for the caller to request another or close.
pub(crate) async fn yield_chunk(&self, chunk: Py<PyAny>) -> ControlFlow<()> {
self.stream_boundary(Suspension::Yield(chunk)).await
}
/// Maps stream demand to continuation and any error to detachment.
async fn stream_boundary(&self, suspension: Suspension) -> ControlFlow<()> {
match self.suspend(suspension).await {
Ok(_) => ControlFlow::Continue(()),
Err(_) => ControlFlow::Break(()),
}
}
/// Yields a suspension with its reply handle and waits for the answer.
async fn suspend(&self, suspension: Suspension) -> Answer {
match self.co.yield_(|reply| Yielded { suspension, reply }).await {
Ok(answer) => answer,
Err(Abandoned) => Err(Error::MissingState.into()),
}
}
}

1162
litellm-rust/crates/host-python/src/driver.rs
View file

File diff suppressed because it is too large Load diff

16
litellm-rust/crates/host-python/src/handle.rs
View file

@ -7,8 +7,9 @@ use pyo3::prelude::*;
pub type PythonLifecycle = for<'py> fn(Python<'py>) -> PyResult<Bound<'py, PyModule>>;
pub enum ExecutionStep {
Return(Py<PyAny>),
/// What a suspended execution waits on: an awaitable the caller awaits, or a stream
/// boundary the caller's stream reads past.
pub enum Suspension {
Await(Py<PyAny>),
/// The call streams: the caller gets a stream over this execution carrying this head,
/// and the execution stays suspended until the stream asks for a chunk.
@ -16,6 +17,11 @@ pub enum ExecutionStep {
Yield(Py<PyAny>),
}
pub enum ExecutionStep {
Suspended(Suspension),
Return(Py<PyAny>),
}
pub trait ExecutionBody: Send + Sync {
fn resume(&mut self, result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep>;
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>;
@ -94,9 +100,9 @@ impl Execution {
let outcome = catch_unwind(AssertUnwindSafe(|| {
let step = body.resume(result)?;
let (tag, value, suspended) = match step {
ExecutionStep::Await(value) => ("Await", value, true),
ExecutionStep::Open(head) => ("Open", head, true),
ExecutionStep::Yield(value) => ("Yield", value, true),
ExecutionStep::Suspended(Suspension::Await(value)) => ("Await", value, true),
ExecutionStep::Suspended(Suspension::Open(head)) => ("Open", head, true),
ExecutionStep::Suspended(Suspension::Yield(value)) => ("Yield", value, true),
ExecutionStep::Return(value) => ("Complete", value, false),
};
let step = lifecycle(py)?.getattr(tag)?.call1((value,))?.unbind();

3
litellm-rust/crates/host-python/src/lib.rs
View file

@ -5,6 +5,7 @@
mod argument;
mod binding;
mod caller;
mod conversion_cache;
mod driver;
mod error;
@ -27,7 +28,7 @@ pub use error::{Error, InvokeError};
pub use file_reader::{FileContent, PythonFileReader, py_bytes};
pub use fork_gate::RuntimeAlreadyStarted;
pub use gil::{PythonContext, attach_blocking, release_count, release_gil};
pub use handle::{Execution, ExecutionBody, ExecutionStep, PythonLifecycle};
pub use handle::{Execution, ExecutionBody, ExecutionStep, PythonLifecycle, Suspension};
pub use hooks::{HookChain, HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonRuntime};
pub use marshal::{
Pythonized, from_py, from_py_argument, json_loads, json_object_field, panic_to_pyerr, to_py,

316
litellm-rust/crates/host-python/tests/driver.rs
View file

@ -9,12 +9,8 @@ use litellm_host::{
lifecycle::{CallEvent, ExecutionEvent, FailureOrigin},
machine::CallMachine,
observation::observation_channel,
protocol::Reply,
};
use litellm_host_python::{
CallOptions, Error as HostError, HookChain, InvokeError, PythonBinding, PythonHostCalls,
PythonOwned, PythonRuntime, run_call,
};
use litellm_host_python::{CallOptions, HookChain, PythonOwned, PythonRuntime, run_call};
use pyo3::{
exceptions::{
PyBaseException, PyException, PyLookupError, PyRuntimeError, PyStopIteration, PyTypeError,
@ -26,9 +22,9 @@ use pyo3::{
};
use rstest::rstest;
use support::{
Classified, Error, HookScript, Interpreter, Log, MachineFactory, OpScript, Synthetic, binding,
call_options, failing_machine, hooks, interpreter, lifecycle_binding, open_stream, read_all,
run_composed, run_scripted, settled, success_machine,
HookScript, Interpreter, Log, MachineFactory, OpScript, Synthetic, binding, call_options,
failing_machine, hooks, interpreter, lifecycle_binding, open_stream, read_all, run_composed,
run_scripted, settled, success_machine,
};
fn observed(capacity: usize) -> (CallOptions, impl FnMut() -> Vec<CallEvent>) {
@ -261,7 +257,6 @@ fn startup_failures_release_the_factory_without_constructing_a_machine(
self.0.store(true, Ordering::SeqCst);
}
}
Python::attach(|py| {
let constructed = Arc::new(AtomicBool::new(false));
let did_construct = constructed.clone();
@ -404,131 +399,103 @@ fn a_stream_carries_its_head_before_the_first_chunk(
});
}
struct Classification {
provider_fails: bool,
op: OpScript,
classifier_fails: bool,
message: &'static str,
native_context: Option<&'static str>,
steps: &'static [&'static str],
}
#[rstest]
fn a_native_failure_is_classified_once_and_reported_classified(
#[case::provider_failure(Classification {
provider_fails: true,
op: OpScript::Answer,
classifier_fails: false,
message: "classified: provider exploded",
native_context: None,
steps: &["project", "classify:provider exploded"],
})]
#[case::host_op_rejection(Classification {
provider_fails: false,
op: OpScript::RejectNatively,
classifier_fails: false,
message: "classified: op rejected",
native_context: None,
steps: &["project", "op:sign", "classify:op rejected"],
})]
#[case::host_op_python_exception(Classification {
provider_fails: false,
op: OpScript::RaisePython,
classifier_fails: false,
message: "op failed",
native_context: None,
steps: &["project", "op:sign"],
})]
#[case::failing_classifier(Classification {
provider_fails: true,
op: OpScript::Answer,
classifier_fails: true,
message: "classifier failed",
native_context: Some("provider exploded"),
steps: &["project", "classify:provider exploded"],
})]
fn a_native_failure_is_classified_once_and_a_python_exception_is_raised_as_is(
#[from(interpreter)] _python: Interpreter,
#[case] case: Classification,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
failing_machine(),
OpScript::Answer,
HookScript::Plain,
asynchronous,
);
let error = result.unwrap_err();
assert!(error.is_instance_of::<PyValueError>(py));
assert_eq!(error.value(py).to_string(), "classified: provider exploded");
assert_eq!(
log,
[
"started",
"begin",
"project",
"classify:provider exploded",
"failed:Call:classified: provider exploded",
"adapter.close",
"host.close",
]
);
});
}
#[rstest]
fn a_native_rejection_from_a_host_operation_is_classified_once(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::RejectNatively,
HookScript::Plain,
false,
);
assert_eq!(
result.unwrap_err().value(py).to_string(),
"classified: op rejected"
);
assert_eq!(
log,
[
"started",
"begin",
"project",
"op:sign",
"classify:op rejected",
"failed:Call:classified: op rejected",
"adapter.close",
"host.close",
]
);
});
}
#[rstest]
fn a_python_exception_from_a_host_operation_is_reported_as_raised(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::RaisePython,
HookScript::Plain,
false,
);
let error = result.unwrap_err();
assert!(error.is_instance_of::<PyValueError>(py));
assert_eq!(error.value(py).to_string(), "op failed");
assert_eq!(
log,
[
"started",
"begin",
"project",
"op:sign",
"failed:Call:op failed",
"adapter.close",
"host.close",
]
);
});
}
#[rstest]
fn a_failing_classifier_surfaces_with_the_native_error_as_context(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
let log_handle = Log::default();
let Classification {
provider_fails,
op,
classifier_fails,
message,
native_context,
steps,
} = case;
let log = Log::default();
let binding = if classifier_fails {
binding(&log, op).with_failing_classifier()
} else {
binding(&log, op)
};
let machine = if provider_fails {
failing_machine()
} else {
success_machine()
};
let (result, log) = run_composed(
py,
failing_machine(),
binding(&log_handle, OpScript::Answer).with_failing_classifier(),
hooks(&log_handle, HookScript::Plain),
machine,
binding,
hooks(&log, HookScript::Plain),
std::convert::identity,
call_options(false),
call_options(asynchronous),
);
let error = result.unwrap_err();
assert!(error.is_instance_of::<PyTypeError>(py));
assert_eq!(error.value(py).to_string(), "classifier failed");
let context = error.context(py).unwrap();
assert!(context.is_instance_of::<PyRuntimeError>(py));
assert_eq!(context.value(py).to_string(), "provider exploded");
assert_eq!(
log,
[
"started",
"begin",
"project",
"classify:provider exploded",
"failed:Call:classifier failed",
"adapter.close",
"host.close",
]
);
assert_eq!(error.value(py).to_string(), message);
match native_context {
Some(context) => {
assert!(error.is_instance_of::<PyTypeError>(py));
let chained = error.context(py).unwrap();
assert!(chained.is_instance_of::<PyRuntimeError>(py));
assert_eq!(chained.value(py).to_string(), context);
}
None => assert!(error.is_instance_of::<PyValueError>(py)),
}
let expected: Vec<String> = ["started", "begin"]
.iter()
.chain(steps)
.map(|step| step.to_string())
.chain([
format!("failed:Call:{message}"),
"adapter.close".into(),
"host.close".into(),
])
.collect();
assert_eq!(log, expected);
});
}
@ -738,25 +705,6 @@ fn the_host_projects_from_the_keyword_view_the_argument_policy_rewrote(
});
}
#[rstest]
fn the_adapters_finalized_response_is_what_the_call_returns_and_reports(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::ReplaceResponse,
asynchronous,
);
assert_eq!(result.unwrap().extract::<String>(py).unwrap(), "replaced");
assert!(log.contains(&"succeeded:replaced".to_string()));
assert!(!log.contains(&"succeeded:project:1|rewritten".to_string()));
});
}
#[rstest]
fn a_failure_while_finalizing_fails_the_call_instead_of_succeeding(
#[from(interpreter)] _python: Interpreter,
@ -788,79 +736,23 @@ fn a_failure_while_finalizing_fails_the_call_instead_of_succeeding(
}
#[rstest]
fn cancellation_ends_the_call_without_terminal_dispatch(#[from(interpreter)] _python: Interpreter) {
fn cancellation_ends_the_call_without_terminal_dispatch(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
struct Cancelling(Log);
impl PythonBinding for Cancelling {
type Protocol = Synthetic;
type Failure = Classified;
fn decode_request(
&mut self,
_: Python<'_>,
_: &Bound<'_, PyDict>,
) -> Result<String, InvokeError<Error>> {
self.0.push("project");
Err(pyo3::exceptions::asyncio::CancelledError::new_err(()).into())
}
fn encode_stream_head(
&mut self,
_: Python<'_>,
head: std::convert::Infallible,
) -> PyResult<Py<PyAny>> {
match head {}
}
fn encode_chunk(
&mut self,
_: Python<'_>,
chunk: std::convert::Infallible,
) -> PyResult<Py<PyAny>> {
match chunk {}
}
fn encode_response(&mut self, _: Python<'_>, _: String) -> PyResult<Py<PyAny>> {
Err(HostError::MissingState.into())
}
fn map_error(&self, _: Python<'_>, error: Error) -> PyResult<Classified> {
self.0.push("classify");
Ok(Classified(error.0))
}
fn host_error(error: &PyErr) -> Error {
Error(error.to_string())
}
}
impl PythonHostCalls<Synthetic> for Cancelling {
fn handle_host_call(
&mut self,
_: Python<'_>,
_: (&'static str, Reply<String>),
) -> Result<(), InvokeError<Error>> {
Err(PyErr::from(HostError::MissingState).into())
}
}
impl PythonOwned for Cancelling {
fn close(&mut self, _: Python<'_>) {}
fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
Ok(())
}
}
let log = Log::default();
let host = Cancelling(log.clone());
let adapter = hooks(&log, HookScript::Plain);
let error = run_call(
let (result, log) = run_scripted(
py,
|_, _, request| Ok(success_machine()(request)),
host,
adapter,
PyDict::new(py).unbind(),
call_options(false),
)
.unwrap_err();
success_machine(),
OpScript::CancelRequest,
HookScript::Plain,
asynchronous,
);
let error = result.unwrap_err();
assert!(!error.is_instance_of::<PyException>(py));
assert_eq!(
log.entries(),
["started", "begin", "project", "adapter.close"]
log,
["started", "begin", "project", "adapter.close", "host.close"]
);
});
}

200
litellm-rust/crates/host-python/tests/handle.rs
View file

@ -1,42 +1,22 @@
mod support;
use litellm_host_python::{Execution, ExecutionBody, ExecutionStep};
use litellm_host::{hooks::CallHooks, lifecycle::Timing};
use litellm_host_python::{
Execution, ExecutionBody, ExecutionStep, HookStep, PythonCallEvent, PythonOwned, PythonRuntime,
Suspension, run_call,
};
use pyo3::{
exceptions::PyBaseException,
gc::{PyTraverseError, PyVisit},
prelude::*,
types::PyDict,
types::{PyCFunction, PyDict},
};
use rstest::rstest;
use support::{Interpreter, install_lifecycle_module, interpreter, lifecycle_binding};
use support::{
Interpreter, Log, OpScript, binding, call_options, install_lifecycle_module, interpreter,
lifecycle_binding, success_machine,
};
#[rstest]
fn python_driver_preserves_inline_await_and_native_ownership(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
py.import("asyncio").unwrap();
let module = install_lifecycle_module(py);
let locals = PyDict::new(py);
locals
.set_item("drive", module.getattr("drive").unwrap())
.unwrap();
locals
.set_item(
"await_execution",
wrap_pyfunction!(await_execution, py).unwrap(),
)
.unwrap();
locals
.set_item(
"calling_execution",
wrap_pyfunction!(calling_execution, py).unwrap(),
)
.unwrap();
let probe = std::ffi::CString::new(include_str!("lifecycle.py")).unwrap();
py.run(&probe, Some(&locals), Some(&locals)).unwrap();
});
}
struct RetainingHost {
retained: Option<Py<PyAny>>,
}
@ -69,7 +49,7 @@ struct AwaitBody(Option<Py<PyAny>>);
impl ExecutionBody for AwaitBody {
fn resume(&mut self, result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
match self.0.take() {
Some(awaitable) => Ok(ExecutionStep::Await(awaitable)),
Some(awaitable) => Ok(ExecutionStep::Suspended(Suspension::Await(awaitable))),
None => result
.expect("selected await completed")
.map(ExecutionStep::Return),
@ -124,6 +104,116 @@ fn error_execution(error: Bound<'_, PyBaseException>) -> Execution {
Execution::new(ErrorBody(Some(error.unbind())), lifecycle_binding)
}
/// Where a paused call keeps the test's object: in its hooks, or selected as the call's
/// response or exception, with the call paused in the hook that is told about it.
#[derive(Clone, Copy, Debug)]
enum Holder {
Hooks,
Response,
Exception,
}
struct PausingHooks {
retained: Option<Py<PyAny>>,
pause: Option<Py<PyAny>>,
holder: Holder,
}
impl CallHooks<PythonRuntime> for PausingHooks {
fn transform_response(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
_: Timing,
) -> PyResult<HookStep<Self, Py<PyAny>>> {
match self.holder {
Holder::Hooks => Ok(HookStep::Ready(response)),
Holder::Response => Ok(HookStep::Ready(self.retained.take().unwrap())),
Holder::Exception => Err(PyErr::from_value(
self.retained.take().unwrap().into_bound(py),
)),
}
}
fn on_event(
&mut self,
_: Python<'_>,
event: PythonCallEvent<'_>,
) -> PyResult<HookStep<Self, ()>> {
let pauses_here = match self.holder {
Holder::Hooks => true,
Holder::Response => matches!(event, PythonCallEvent::Succeeded { .. }),
Holder::Exception => matches!(event, PythonCallEvent::Failed { .. }),
};
Ok(match self.pause.take_if(|_| pauses_here) {
Some(pause) => {
HookStep::Await(pause, |_, _, answer| answer.map(|_| HookStep::Ready(())))
}
None => HookStep::Ready(()),
})
}
}
impl PythonOwned for PausingHooks {
fn close(&mut self, _: Python<'_>) {
self.retained = None;
self.pause = None;
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.retained)?;
visit.call(&self.pause)
}
}
fn paused_call(
py: Python<'_>,
retained: Py<PyAny>,
pause: Py<PyAny>,
holder: Holder,
) -> PyResult<Py<PyAny>> {
let machine = success_machine();
run_call(
py,
move |_, _, request| Ok(machine(request)),
binding(&Log::default(), OpScript::Answer),
PausingHooks {
retained: Some(retained),
pause: Some(pause),
holder,
},
PyDict::new(py).unbind(),
call_options(true),
)
}
#[rstest]
fn every_suspension_is_awaited_inline_on_the_callers_task(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
let module = install_lifecycle_module(py);
let locals = PyDict::new(py);
locals
.set_item("drive", module.getattr("drive").unwrap())
.unwrap();
locals
.set_item(
"await_execution",
wrap_pyfunction!(await_execution, py).unwrap(),
)
.unwrap();
locals
.set_item(
"calling_execution",
wrap_pyfunction!(calling_execution, py).unwrap(),
)
.unwrap();
let probe = std::ffi::CString::new(include_str!("lifecycle.py")).unwrap();
py.run(&probe, Some(&locals), Some(&locals)).unwrap();
});
}
#[rstest]
fn retained_exception_frames_are_collectable(#[from(interpreter)] _python: Interpreter) {
Python::attach(|py| {
@ -199,3 +289,51 @@ assert retained_ref() is None
.unwrap();
});
}
#[rstest]
#[case::held_by_the_hooks(Holder::Hooks)]
#[case::selected_as_the_response(Holder::Response)]
#[case::selected_as_the_exception(Holder::Exception)]
fn cycles_through_a_call_suspended_in_a_hook_are_collectable(
#[from(interpreter)] _python: Interpreter,
#[case] holder: Holder,
) {
Python::attach(|py| {
let locals = PyDict::new(py);
let call = PyCFunction::new_closure(py, None, None, move |args, _| {
let (retained, pause) = args.extract::<(Py<PyAny>, Py<PyAny>)>()?;
paused_call(args.py(), retained, pause, holder)
})
.unwrap();
locals.set_item("paused_call", call).unwrap();
py.run(
pyo3::ffi::c_str!(
r#"
import gc
import weakref
class Retained(Exception):
pass
class Pause:
def __await__(self):
yield
def cycle():
retained = Retained()
coroutine = paused_call(retained, Pause())
coroutine.send(None)
retained.coroutine = coroutine
return weakref.ref(retained)
reference = cycle()
gc.collect()
assert reference() is None
"#
),
Some(&locals),
Some(&locals),
)
.unwrap();
});
}

4
litellm-rust/crates/host-python/tests/hook_chain.rs
View file

@ -504,7 +504,9 @@ impl litellm_host_python::ExecutionBody for Preparing {
}
HookStep::Await(awaitable, resume) => {
self.resume = Some(resume);
Ok(litellm_host_python::ExecutionStep::Await(awaitable))
Ok(litellm_host_python::ExecutionStep::Suspended(
litellm_host_python::Suspension::Await(awaitable),
))
}
}
})