diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs index 2256a6888e2..a11162f1c55 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs @@ -538,11 +538,9 @@ struct PendingLogging { #[pymethods] impl PendingLogging { - fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> { + fn __call__(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<()> { let pending = slf.borrow_mut().pending.take(); - if let Some(pending) = pending - && success - { + if let Some(pending) = pending { match pending.asynchronous(py) { Err(error) if error.is_instance_of::(py) => { error.write_unraisable(py, Some(pending.logger.object(py))); @@ -563,10 +561,14 @@ impl PendingLogging { Ok(()) } - fn __clear__(slf: &Bound<'_, Self>) { + fn close(slf: &Bound<'_, Self>) { let pending = slf.borrow_mut().pending.take(); drop(pending); } + + fn __clear__(slf: &Bound<'_, Self>) { + Self::close(slf); + } } #[cfg(test)] @@ -989,7 +991,7 @@ sys.unraisablehook = old_hook } #[test] - fn deferred_release_uses_release_context_and_allows_reentry_once() { + fn deferred_logging_uses_call_context_and_allows_reentry_once() { let _guard = PYTHON_GLOBALS .lock() .unwrap_or_else(|error| error.into_inner()); @@ -1022,7 +1024,7 @@ class Coroutine: class Worker: def ensure_initialized_and_enqueue(self, coroutine): observed.append(marker.get()) - pending.release(True) + pending() coroutine.close() class Logger: @@ -1060,10 +1062,72 @@ logger = Logger() py.run( pyo3::ffi::c_str!( r#" -marker.set('release') -pending.release(True) -pending.release(True) -assert observed == ['created', 'release', 'closed'] +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'] "# ), Some(&locals), diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs index 6d5f40c5aa9..a703e9e736a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs @@ -1,3 +1,4 @@ +use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; use pyo3::types::{PyDict, PyTuple}; @@ -46,8 +47,19 @@ struct ProjectedOcrHost { secret_fields: Vec<&'static str>, azure_ad_token_provider: Option, pre_call: Option, - body: Option>, - headers: Option>, + payload: Option, +} + +struct CapturedOcrPayload { + body: Py, + headers: Py, +} + +impl CapturedOcrPayload { + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.body)?; + visit.call(&self.headers) + } } impl PythonOcrHost { @@ -74,8 +86,7 @@ impl PythonOcrHost { secret_fields: projected.secret_fields, azure_ad_token_provider: projected.azure_ad_token_provider, pre_call: None, - body: None, - headers: None, + payload: None, }); Ok(OcrHostResult::Request(Ok(( Box::new(projected.request), @@ -113,15 +124,23 @@ 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.body = Some(body.unbind()); - projected.headers = Some(headers.unbind()); + projected.payload = Some(CapturedOcrPayload { + body: body.unbind(), + headers: headers.unbind(), + }); Ok(request) } @@ -131,11 +150,12 @@ impl PythonOcrHost { request: OcrPostCallRequest, ) -> PyResult { let projected = self.projected()?; + let payload = projected.payload.as_ref(); self.state.logger()?.post_ocr( py, &request.original_response, - projected.body.as_ref(), - projected.headers.as_ref(), + payload.map(|payload| &payload.body), + payload.map(|payload| &payload.headers), )?; Ok(request) } @@ -223,15 +243,17 @@ impl PythonRoute for PythonOcrHost { self.projected = None; } - fn traverse(&self, visit: &pyo3::gc::PyVisit<'_>) -> Result<(), pyo3::gc::PyTraverseError> { + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { let Some(projected) = &self.projected else { return Ok(()); }; if let Some(provider) = &projected.azure_ad_token_provider { provider.traverse(visit)?; } - visit.call(&projected.body)?; - visit.call(&projected.headers) + if let Some(payload) = &projected.payload { + payload.traverse(visit)?; + } + Ok(()) } } diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 46b222a4fc9..285c731c140 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3212,7 +3212,10 @@ 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 - pending.release(not exception_raised) + if exception_raised: + pending.close() + else: + pending() _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 6295469c066..2c407f6c3a0 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -461,10 +461,16 @@ def test_native_pending_logging_is_released_only_for_ocr(call_type: str, excepti ) if call_type in ("ocr", "aocr"): - pending.release.assert_called_once_with(not exception_raised) + if exception_raised: + pending.close.assert_called_once_with() + pending.assert_not_called() + else: + pending.assert_called_once_with() + pending.close.assert_not_called() assert logger._native_pending_logging is None else: - pending.release.assert_not_called() + pending.assert_not_called() + pending.close.assert_not_called() assert logger._native_pending_logging is pending if exception_raised: enqueue.assert_not_called()