diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 1bb84184e4c..3930584e3b5 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -251,25 +251,27 @@ impl LegacyLogging { if !self.asynchronous { return pending().sync(py); } - if !self.internal - && self - .call - .kwargs() - .bind(py) - .get_item("fallbacks")? - .is_none_or(|value| value.is_none()) - { - if logger.defers_async_logging(py) { - let pending = Py::new( - py, - PendingLogging { - pending: Some(pending()), - }, - )?; - logger.defer_success(py, pending.bind(py).as_any())?; - } else { - pending().asynchronous(py)?; - } + if self.internal { + return logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end); + } + let enqueues = self + .call + .kwargs() + .bind(py) + .get_item("fallbacks")? + .is_none_or(|value| value.is_none()); + if logger.defers_async_logging(py) { + let pending = Py::new( + py, + PendingLogging { + pending: Some(pending()), + enqueues, + }, + )?; + return logger.defer_success(py, pending.bind(py).as_any()); + } + if enqueues { + pending().asynchronous(py)?; } logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end) } @@ -1875,8 +1877,13 @@ mod terminal_tests { c"", &["async_success_handler", "enqueued", "sync_success_for_async_call"] )] - #[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])] + #[case::async_deferred(true, c"logger._defer_async_logging = True", &[])] #[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])] + #[case::async_deferred_with_fallbacks( + true, + c"logger._defer_async_logging = True\nkwargs = {'fallbacks': ['other']}", + &[] + )] fn success_reaches_the_logging_handlers( #[case] asynchronous: bool, #[case] script: &CStr, @@ -1908,6 +1915,43 @@ assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_asy }); } + #[rstest] + #[case::accepted(c"", true, &["sync_success_for_async_call", "async_success_handler", "enqueued"])] + #[case::accepted_with_fallbacks(c"kwargs = {'fallbacks': ['other']}", true, &["sync_success_for_async_call"])] + #[case::blocked(c"", false, &[])] + fn a_deferred_success_runs_sync_callbacks_only_once_released( + #[case] script: &CStr, + #[case] accepted: bool, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c"response = object()\nlogger._defer_async_logging = True", + ); + run(py, &locals, script); + let mut logging = logged(py, &locals, true); + succeed(py, &locals, &mut logging); + locals.set_item("accepted", accepted).unwrap(); + run( + py, + &locals, + c" +assert logger.names() == [], logger.calls +logger._native_pending_logging.release(accepted) +logger._native_pending_logging.release(True) +", + ); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + }); + } + #[rstest] fn dropped_observations_preserve_deferred_success_and_response_identity() { Python::initialize(); @@ -1933,7 +1977,7 @@ assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_asy )); assert_eq!(sender.dropped_events(), 1); run(py, &locals, c" -assert logger.names() == ['sync_success_for_async_call'], logger.calls +assert logger.names() == [], logger.calls logger._native_pending_logging.release(True) logger._native_pending_logging.release(True) assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'enqueued'], logger.calls diff --git a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs index 292648ba69c..5ee0c5b4f65 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs @@ -22,21 +22,37 @@ impl PendingSuccess { self.logger .enqueue_success(py, &self.response, &self.start, &self.end) } + + pub(crate) fn sync_for_async_call(&self, py: Python<'_>) -> PyResult<()> { + self.logger + .sync_success_for_async_call(py, &self.response, &self.start, &self.end) + } } #[pyclass] pub(crate) struct PendingLogging { pub(crate) pending: Option, + pub(crate) enqueues: bool, } #[pymethods] impl PendingLogging { fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> { - let pending = slf.borrow_mut().pending.take(); + let (pending, enqueues) = { + let mut slot = slf.borrow_mut(); + (slot.pending.take(), slot.enqueues) + }; if let Some(pending) = pending && success { - match pending.asynchronous(py) { + let released = pending.sync_for_async_call(py).and_then(|()| { + if enqueues { + pending.asynchronous(py) + } else { + Ok(()) + } + }); + match released { Err(error) if error.is_instance_of::(py) => { error.write_unraisable(py, Some(pending.logger.object(py))); } @@ -87,6 +103,7 @@ mod tests { start: py.None(), end: Some(py.None()), }), + enqueues: true, }, ) .unwrap(); @@ -121,8 +138,9 @@ marker.set('release') pending.release(True) pending.release(True) assert observed == ['release'], observed -assert logger.names() == ['async_success_handler', 'enqueued'], logger.calls +assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'enqueued'], logger.calls assert logger.calls[0][1] is response +assert logger.calls[1][1] is response ", ); }); @@ -182,7 +200,7 @@ logger.on_enqueue = on_enqueue &locals, c" pending.release(True) -assert logger.names() == ['async_success_handler', 'enqueued', 'closed'], logger.calls +assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'enqueued', 'closed'], logger.calls assert unraisable_from(logger) == ([] if propagates else [failure]) ", ); diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 611316fb52c..ec602b47f43 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -220,6 +220,39 @@ async def test_deferred_logging_requires_release_and_runs_at_most_once(server: R assert events[0].response is response +@pytest.mark.asyncio +@pytest.mark.parametrize("blocked", [False, True]) +async def test_deferred_sync_success_callbacks_wait_for_release(server: RecordingServer, blocked: bool) -> None: + finished: Final = threading.Event() + observations: Final = [] + + class Observe(CustomLogger): + def log_success_event(self, kwargs, response_obj, start_time, end_time): + observations.append(response_obj) + finished.set() + + logger: Final = Logging( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="deferred-sync", + function_id="deferred-sync", + dynamic_success_callbacks=[Observe()], + ) + logger._defer_async_logging = True + response: Final = await call_aocr(server, litellm_logging_obj=logger) + await drain_logging() + assert not logger._sync_success_scheduled + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked) + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked) + assert logger._sync_success_scheduled is not blocked + if not blocked: + assert await asyncio.to_thread(finished.wait, 5) + assert observations == [response] + + def test_sync_success_uses_executor_and_copied_caller_context(server: RecordingServer) -> None: context: Final = ContextVar("sync-lifecycle", default="missing") context.set("caller")