From acf40e77685eae750ef18f6941990eb39fdb6595 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 13 Oct 2026 19:58:21 -0700 Subject: [PATCH] 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 --- litellm-rust/crates/host-python/src/caller.rs | 58 +++++++++++++++++++ litellm-rust/crates/host-python/src/driver.rs | 58 +++++++------------ litellm-rust/crates/host-python/src/lib.rs | 2 +- litellm-rust/crates/host-python/src/task.rs | 32 ---------- 4 files changed, 81 insertions(+), 69 deletions(-) create mode 100644 litellm-rust/crates/host-python/src/caller.rs delete mode 100644 litellm-rust/crates/host-python/src/task.rs diff --git a/litellm-rust/crates/host-python/src/caller.rs b/litellm-rust/crates/host-python/src/caller.rs new file mode 100644 index 00000000000..3bae948c936 --- /dev/null +++ b/litellm-rust/crates/host-python/src/caller.rs @@ -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>; + +pub(crate) struct Yielded { + pub(crate) suspension: Suspension, + pub(crate) reply: Reply, +} + +/// Hands suspensions to the Python caller and receives its answers. +#[derive(Clone)] +pub(crate) struct PythonCaller { + co: Co, +} + +impl PythonCaller { + /// Connects the call body to its coroutine's yield handle. + pub(crate) fn new(co: Co) -> 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) -> 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) -> 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) -> 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()), + } + } +} diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index 2371597c738..da38588088e 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -1,4 +1,3 @@ -use std::ops::ControlFlow; use std::pin::pin; use std::sync::Arc; use std::task::{Context, Poll, Waker}; @@ -19,10 +18,10 @@ use litellm_host::{ }; use crate::PythonHostCalls; +use crate::caller::{Answer, PythonCaller, Yielded}; use crate::handle::{Execution, ExecutionBody, ExecutionStep, PythonLifecycle, Suspension}; use crate::hooks::{HookStep, PythonCallEvent, PythonCallHooks}; use crate::native::{NativeMachine, NativePoll}; -use crate::task::{Answer, PythonTask, Yielded}; use crate::{Error, InvokeError, PythonBinding}; type ProtocolOf = ::Protocol; @@ -124,19 +123,19 @@ fn classified(py: Python<'_>, binding: &H, error: ErrorOf) /// The call as one async body. Every hook step and host call is awaited in place on the /// caller's task; the parts lock is never held across an await. -struct Running +struct CallBody where H: PythonBinding, M: Machine, { - task: PythonTask, + caller: PythonCaller, parts: SharedParts, native: NativeMachine, start: Option, M>>, streaming: bool, } -impl Running +impl CallBody where L: PythonCallHooks + 'static, H: PythonBinding + PythonHostCalls + 'static, @@ -223,7 +222,7 @@ where HookStep::Ready(value) => return Ok(value), HookStep::Await(awaitable, resume) => (awaitable, resume), }; - let answer = self.task.suspend(Suspension::Await(awaitable)).await; + let answer = self.caller.await_python(awaitable).await; step = { let mut parts = self.parts(); Python::attach(|py| resume(&mut parts.hooks, py, answer))? @@ -237,16 +236,14 @@ where let poll = Python::attach(|py| self.native.resume(py, interruption.take()))?; let result = match poll { NativePoll::Ready(result) => result, - NativePoll::Suspend(awaitable) => { - match self.task.suspend(Suspension::Await(awaitable)).await { - Ok(_) => self.native.take_result()?, - Err(error) if self.parts().interrupted.is_none() => { - interruption = Some(self.interrupt(error)); - continue; - } - Err(error) => return Err(error), + NativePoll::Suspend(awaitable) => match self.caller.await_python(awaitable).await { + Ok(_) => self.native.take_result()?, + Err(error) if self.parts().interrupted.is_none() => { + interruption = Some(self.interrupt(error)); + continue; } - } + Err(error) => return Err(error), + }, }; let op = match result { Ok(MachineStep::Suspended(op)) => op, @@ -324,7 +321,7 @@ where Ok::<_, PyErr>(head) })? }; - self.consumed(Suspension::Open(head), reply).await; + reply.send(self.caller.open_stream(head).await); Ok(()) } HostRequest::Stream(StreamDelivery::Chunk(chunk, reply)) => { @@ -337,7 +334,7 @@ where Ok::<_, PyErr>(chunk) })? }; - self.consumed(Suspension::Yield(chunk), reply).await; + reply.send(self.caller.yield_chunk(chunk).await); Ok(()) } } @@ -352,7 +349,7 @@ where Python::attach(|py| parts.binding.begin_host_call(py, call))? }; while let Some(awaitable) = pending { - let result = self.task.suspend(Suspension::Await(awaitable)).await; + let result = self.caller.await_python(awaitable).await; pending = { let mut parts = self.parts(); Python::attach(|py| parts.binding.resume_host_call(py, result))? @@ -370,17 +367,6 @@ where self.settle(step).await } - /// Hands a stream boundary to the consumer. Reading on is demand for the next chunk; - /// anything else, closing the stream included, detaches the call from it. - async fn consumed(&self, suspension: Suspension, reply: Reply>) { - let read = self.task.suspend(suspension).await; - reply.send(if read.is_ok() { - ControlFlow::Continue(()) - } else { - ControlFlow::Break(()) - }); - } - async fn completed( &mut self, completion: HostedCompletion>, @@ -487,7 +473,7 @@ where } } -struct Driver +struct CallDriver where H: PythonBinding, L: PythonCallHooks, @@ -497,7 +483,7 @@ where parts: SharedParts, } -impl Driver +impl CallDriver where H: PythonBinding + 'static, L: PythonCallHooks + 'static, @@ -527,7 +513,7 @@ where } } -impl ExecutionBody for Driver +impl ExecutionBody for CallDriver where H: PythonBinding + 'static, L: PythonCallHooks + 'static, @@ -575,7 +561,7 @@ where } } -impl Drop for Driver +impl Drop for CallDriver where H: PythonBinding, L: PythonCallHooks, @@ -637,8 +623,8 @@ where let parts = Arc::clone(&parts); let asynchronous = options.asynchronous; move |co| { - Running { - task: PythonTask::new(co), + CallBody { + caller: PythonCaller::new(co), parts, native: NativeMachine::new(asynchronous), start: Some(Box::new(start)), @@ -647,7 +633,7 @@ where .run() } }); - let mut driver = Driver { + let mut driver = CallDriver { body: Mutex::new(body), reply: None, parts, diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 3deb95ffac7..580e6674b95 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -5,6 +5,7 @@ mod argument; mod binding; +mod caller; mod conversion_cache; mod driver; mod error; @@ -18,7 +19,6 @@ mod native; mod owned; mod runtime; mod services; -mod task; pub use argument::{effective_py_args, lookup, present}; pub use binding::PythonBinding; diff --git a/litellm-rust/crates/host-python/src/task.rs b/litellm-rust/crates/host-python/src/task.rs deleted file mode 100644 index a48a3f0bad9..00000000000 --- a/litellm-rust/crates/host-python/src/task.rs +++ /dev/null @@ -1,32 +0,0 @@ -use litellm_coroutine::{Abandoned, Co, Reply}; -use pyo3::prelude::*; - -use crate::Error; -use crate::handle::Suspension; - -pub(crate) type Answer = PyResult>; - -pub(crate) struct Yielded { - pub(crate) suspension: Suspension, - pub(crate) reply: Reply, -} - -/// The caller's Python task as the running call sees it. Every suspension is handed to -/// that task and the call continues with what the task resumes it with. -#[derive(Clone)] -pub(crate) struct PythonTask { - co: Co, -} - -impl PythonTask { - pub(crate) fn new(co: Co) -> Self { - Self { co } - } - - pub(crate) 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()), - } - } -}