feat(rust-interop): count GIL acquisitions alongside releases

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<Py<PyAny>> {
     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::<PyDict>()
+            .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);
+        });
+    }
+}

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())?,
         }
     }
 }

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<T, F>(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)
 }

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};

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);
+}