From 4aa6acc04580986831197ea10c5c0f5689383592 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 19:08:13 -0700 Subject: [PATCH] Revert "refactor(python-bridge): make deferred logging callable" This reverts commit 5a7ea7ef710b89cca377da5bc60bdc6603257764. --- .../crates/python-bridge/src/lifecycle/mod.rs | 86 +++---------------- .../python-bridge/src/routes/ocr/host.rs | 46 +++------- litellm/proxy/common_request_processing.py | 5 +- .../test_deferred_guardrail_logging.py | 10 +-- 4 files changed, 26 insertions(+), 121 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs index a11162f1c55..2256a6888e2 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs @@ -538,9 +538,11 @@ struct PendingLogging { #[pymethods] impl PendingLogging { - fn __call__(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<()> { + fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> { let pending = slf.borrow_mut().pending.take(); - if let Some(pending) = pending { + if let Some(pending) = pending + && success + { match pending.asynchronous(py) { Err(error) if error.is_instance_of::(py) => { error.write_unraisable(py, Some(pending.logger.object(py))); @@ -561,14 +563,10 @@ impl PendingLogging { Ok(()) } - fn close(slf: &Bound<'_, Self>) { + fn __clear__(slf: &Bound<'_, Self>) { let pending = slf.borrow_mut().pending.take(); drop(pending); } - - fn __clear__(slf: &Bound<'_, Self>) { - Self::close(slf); - } } #[cfg(test)] @@ -991,7 +989,7 @@ sys.unraisablehook = old_hook } #[test] - fn deferred_logging_uses_call_context_and_allows_reentry_once() { + fn deferred_release_uses_release_context_and_allows_reentry_once() { let _guard = PYTHON_GLOBALS .lock() .unwrap_or_else(|error| error.into_inner()); @@ -1024,7 +1022,7 @@ class Coroutine: class Worker: def ensure_initialized_and_enqueue(self, coroutine): observed.append(marker.get()) - pending() + pending.release(True) coroutine.close() class Logger: @@ -1062,72 +1060,10 @@ logger = Logger() py.run( pyo3::ffi::c_str!( r#" -marker.set('call') -pending() -pending() -assert observed == ['created', 'call', 'closed'] -"# - ), - Some(&locals), - Some(&locals), - ) - .unwrap(); - }); - } - - #[test] - fn deferred_logging_close_is_reentry_safe_and_invalidates_aliases() { - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - py.run( - pyo3::ffi::c_str!( - r#" -observed = [] - -class Retained: - def __del__(self): - observed.append('finalized') - alias() - -class Logger: - def async_success_handler(self, *args): - observed.append('enqueued') - -logger = Logger() -retained = Retained() -"# - ), - Some(&locals), - Some(&locals), - ) - .unwrap(); - let pending = Py::new( - py, - PendingLogging { - pending: Some(PendingSuccess { - logger: locals - .get_item("logger") - .unwrap() - .unwrap() - .extract() - .unwrap(), - response: Some(locals.get_item("retained").unwrap().unwrap().unbind()), - start: py.None(), - end: None, - }), - }, - ) - .unwrap(); - locals.set_item("pending", &pending).unwrap(); - locals.set_item("alias", &pending).unwrap(); - locals.del_item("retained").unwrap(); - py.run( - pyo3::ffi::c_str!( - r#" -pending.close() -alias() -assert observed == ['finalized'] +marker.set('release') +pending.release(True) +pending.release(True) +assert observed == ['created', 'release', 'closed'] "# ), Some(&locals), diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 91645aa04c1..3cbe4e03a10 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -1,6 +1,5 @@ use std::sync::Arc; -use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; use pyo3::types::PyDict; @@ -29,21 +28,10 @@ struct ProjectedOcrHost { secret_fields: Vec<&'static str>, azure_ad_token_provider: Option, pre_call: Option, - payload: Option, reader: Option, reader_failed: bool, -} - -struct CapturedOcrPayload { - body: Py, - headers: Py, -} - -impl CapturedOcrPayload { - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.body)?; - visit.call(&self.headers) - } + body: Option>, + headers: Option>, } impl PythonOcrHost { @@ -77,9 +65,10 @@ impl PythonOcrHost { secret_fields: projected.secret_fields, azure_ad_token_provider: projected.azure_ad_token_provider, pre_call: None, - payload: None, reader: projected.reader, reader_failed: false, + body: None, + headers: None, }); Ok(OcrHostResult::Request(Ok(( Box::new( @@ -121,23 +110,15 @@ impl PythonOcrHost { for (name, value) in &request.headers { headers.set_item(name, value)?; } - logger.pre_ocr( - py, - request.api_key.as_deref(), - &body, - &headers, - &request.url, - )?; + logger.pre_ocr(py, request.api_key.as_deref(), &body, &headers, &request.url)?; request.body = from_py(&body)?; request.headers = headers .iter() .map(|(name, value)| Ok((name.extract::()?, value.extract::()?))) .collect::>>()?; let projected = self.projected_mut()?; - projected.payload = Some(CapturedOcrPayload { - body: body.unbind(), - headers: headers.unbind(), - }); + projected.body = Some(body.unbind()); + projected.headers = Some(headers.unbind()); Ok(request) } @@ -147,12 +128,11 @@ impl PythonOcrHost { request: OcrPostCallRequest, ) -> PyResult { let projected = self.projected()?; - let payload = projected.payload.as_ref(); self.state.logger()?.post_ocr( py, &request.original_response, - payload.map(|payload| &payload.body), - payload.map(|payload| &payload.headers), + projected.body.as_ref(), + projected.headers.as_ref(), )?; Ok(request) } @@ -247,7 +227,7 @@ impl PythonRoute for PythonOcrHost { self.projected = None; } - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + fn traverse(&self, visit: &pyo3::gc::PyVisit<'_>) -> Result<(), pyo3::gc::PyTraverseError> { let Some(projected) = &self.projected else { return Ok(()); }; @@ -257,10 +237,8 @@ impl PythonRoute for PythonOcrHost { if let Some(reader) = &projected.reader { reader.traverse(visit)?; } - if let Some(payload) = &projected.payload { - payload.traverse(visit)?; - } - Ok(()) + visit.call(&projected.body)?; + visit.call(&projected.headers) } } diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 729b3bc5c82..2f39e6c71bc 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3307,10 +3307,7 @@ class ProxyBaseLLMRequestProcessing: pending: Final = getattr(logging_obj, "_native_pending_logging", None) if pending is not None: logging_obj._native_pending_logging = None # rebind-ok: consume the native OCR release signal once - if exception_raised: - pending.close() - else: - pending() + pending.release(not exception_raised) _enqueue_fn: Final = getattr(logging_obj, "_enqueue_deferred_logging", None) if _enqueue_fn is None: return diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index 2c407f6c3a0..6295469c066 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -461,16 +461,10 @@ def test_native_pending_logging_is_released_only_for_ocr(call_type: str, excepti ) if call_type in ("ocr", "aocr"): - if exception_raised: - pending.close.assert_called_once_with() - pending.assert_not_called() - else: - pending.assert_called_once_with() - pending.close.assert_not_called() + pending.release.assert_called_once_with(not exception_raised) assert logger._native_pending_logging is None else: - pending.assert_not_called() - pending.close.assert_not_called() + pending.release.assert_not_called() assert logger._native_pending_logging is pending if exception_raised: enqueue.assert_not_called()