diff --git a/litellm-rust/crates/python-bridge/src/diagnostics.rs b/litellm-rust/crates/python-bridge/src/diagnostics.rs index cc153a89b8f..fb67268d857 100644 --- a/litellm-rust/crates/python-bridge/src/diagnostics.rs +++ b/litellm-rust/crates/python-bridge/src/diagnostics.rs @@ -1,4 +1,4 @@ -use litellm_python_interop::release_count; +use litellm_python_interop::{acquisition_count, release_count}; use pyo3::prelude::*; use pyo3::types::PyDict; @@ -6,6 +6,7 @@ use pyo3::types::PyDict; fn gil_stats(py: Python<'_>) -> PyResult> { let stats = PyDict::new(py); stats.set_item("releases", release_count())?; + stats.set_item("acquisitions", acquisition_count())?; Ok(stats.into_any().unbind()) } @@ -21,3 +22,54 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(_panic_for_test, module)?)?; Ok(()) } + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use super::*; + + fn stat(py: Python<'_>, key: &str) -> u64 { + gil_stats(py) + .expect("stats should build") + .into_bound(py) + .cast_into::() + .expect("stats should be a dict") + .get_item(key) + .expect("key should resolve") + .expect("key should be present") + .extract() + .expect("stat should be an integer") + } + + #[test] + fn gil_stats_reports_both_interop_counters() { + Python::initialize(); + Python::attach(|py| { + assert_eq!(stat(py, "releases"), release_count()); + assert_eq!(stat(py, "acquisitions"), acquisition_count()); + }); + } + + #[test] + fn gil_stats_counts_a_sync_route_run_on_both_counters() { + Python::initialize(); + Python::attach(|py| { + let releases_before = stat(py, "releases"); + let acquisitions_before = stat(py, "acquisitions"); + + crate::execution::run_sync( + py, + async { + tokio::time::sleep(Duration::from_millis(120)).await; + Ok(true) + }, + crate::errors::core_error_to_pyerr, + ) + .expect("sync route should complete"); + + assert!(stat(py, "releases") > releases_before); + assert!(stat(py, "acquisitions") > acquisitions_before); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/execution.rs b/litellm-rust/crates/python-bridge/src/execution.rs index f3648158cf6..3e19ef75956 100644 --- a/litellm-rust/crates/python-bridge/src/execution.rs +++ b/litellm-rust/crates/python-bridge/src/execution.rs @@ -4,7 +4,7 @@ use std::time::Duration; use futures_util::FutureExt; use litellm_core::error::Error; -use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil}; +use litellm_python_interop::{Pythonized, attach, panic_to_pyerr, release_gil}; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; use serde::Serialize; @@ -99,7 +99,7 @@ where loop { tokio::select! { result = &mut future => return result, - _ = signal_checks.tick() => Python::attach(|py| py.check_signals())?, + _ = signal_checks.tick() => attach(|py| py.check_signals())?, } } } diff --git a/litellm-rust/crates/python-interop/src/gil.rs b/litellm-rust/crates/python-interop/src/gil.rs index 04b966a6002..1f256c0cf04 100644 --- a/litellm-rust/crates/python-interop/src/gil.rs +++ b/litellm-rust/crates/python-interop/src/gil.rs @@ -2,8 +2,18 @@ use std::sync::atomic::{AtomicU64, Ordering}; use pyo3::prelude::*; +static GIL_ACQUISITIONS: AtomicU64 = AtomicU64::new(0); static GIL_RELEASES: AtomicU64 = AtomicU64::new(0); +/// Runs work attached to the interpreter and records the acquisition. +pub fn attach(f: F) -> T +where + F: for<'py> FnOnce(Python<'py>) -> T, +{ + GIL_ACQUISITIONS.fetch_add(1, Ordering::Relaxed); + Python::attach(f) +} + /// Runs work detached from the interpreter and records the release. /// /// `f` must not access Python state while the interpreter is detached. @@ -16,6 +26,10 @@ where py.detach(f) } +pub fn acquisition_count() -> u64 { + GIL_ACQUISITIONS.load(Ordering::Relaxed) +} + pub fn release_count() -> u64 { GIL_RELEASES.load(Ordering::Relaxed) } diff --git a/litellm-rust/crates/python-interop/src/lib.rs b/litellm-rust/crates/python-interop/src/lib.rs index 2e562bdae70..73d98d47a5b 100644 --- a/litellm-rust/crates/python-interop/src/lib.rs +++ b/litellm-rust/crates/python-interop/src/lib.rs @@ -1,5 +1,5 @@ mod gil; mod marshal; -pub use gil::{release_count, release_gil}; +pub use gil::{acquisition_count, attach, release_count, release_gil}; pub use marshal::{Pythonized, from_py, panic_to_pyerr, to_py}; diff --git a/litellm-rust/crates/python-interop/tests/interop.rs b/litellm-rust/crates/python-interop/tests/interop.rs index 9c456dcb938..b6532ff7aac 100644 --- a/litellm-rust/crates/python-interop/tests/interop.rs +++ b/litellm-rust/crates/python-interop/tests/interop.rs @@ -2,7 +2,9 @@ use pyo3::Python; use rstest::{fixture, rstest}; use serde_json::{Value, json}; -use litellm_python_interop::{from_py, release_count, release_gil, to_py}; +use litellm_python_interop::{ + acquisition_count, attach, from_py, release_count, release_gil, to_py, +}; struct InitializedPython; @@ -42,3 +44,12 @@ fn release_gil_runs_work_and_records_it(#[from(initialized_python)] python: &Ini assert_eq!(result, 42); assert_eq!(release_count(), before + 1); } + +#[rstest] +fn attach_runs_work_and_records_it(#[from(initialized_python)] python: &InitializedPython) { + let before = acquisition_count(); + let result = python.attach(|_| attach(|_| 42)); + + assert_eq!(result, 42); + assert_eq!(acquisition_count(), before + 1); +}