Merge pull request #41987 from BerriAI/litellm_rust_fork_safety

fix(rust): refuse native routes in processes forked after the runtime started
This commit is contained in:
yujonglee 2026-09-19 13:47:54 -07:00 • committed by GitHub
commit d675c1285b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 656 additions and 29 deletions

View file

@ -3046,7 +3046,6 @@ checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
dependencies = [
"base64 0.22.1",
"bytes",
"futures-channel",
"futures-core",
"futures-util",
"h2 0.4.15",

View file

@ -34,7 +34,7 @@ pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
pythonize = "0.29.0"
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] }
reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] }
rstest = "0.26.1"
rstest_reuse = "0.7.0"
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }

10
litellm-rust/clippy.toml Normal file
View file

@ -0,0 +1,10 @@
# The Tokio runtime is reached only through `host-python/src/execution.rs`, whose fork gate
# must see every entry. Going around it makes a fork-after-use hang instead of raising.
disallowed-methods = [
{ path = "pyo3_async_runtimes::tokio::get_runtime", reason = "use litellm_host_python::run_sync / run_sync_value" },
{ path = "pyo3_async_runtimes::tokio::future_into_py", reason = "use litellm_host_python::run_async / run_async_value" },
{ path = "pyo3_async_runtimes::tokio::future_into_py_with_locals", reason = "use litellm_host_python::run_async / run_async_value" },
{ path = "pyo3_async_runtimes::tokio::local_future_into_py", reason = "use litellm_host_python::run_async / run_async_value" },
{ path = "pyo3_async_runtimes::tokio::run", reason = "use litellm_host_python::run_sync / run_sync_value" },
{ path = "pyo3_async_runtimes::tokio::run_until_complete", reason = "use litellm_host_python::run_sync / run_sync_value" },
]

View file

@ -4,6 +4,7 @@ use std::pin::Pin;
use std::task::{Context, Poll, Waker};
use std::time::Duration;
use crate::fork_gate::{ForkGate, Refused, RuntimeAlreadyStarted};
use crate::{Pythonized, panic_to_pyerr, release_gil};
use futures_util::FutureExt;
use pyo3::exceptions::PyRuntimeError;
@ -12,6 +13,67 @@ use serde::Serialize;
use tokio::runtime::{Handle, Runtime};
use tokio::time::{self, MissedTickBehavior};
pyo3::create_exception!(
_native,
ForkedAfterNativeRuntimeStarted,
PyRuntimeError,
"This process was forked after the native runtime started. Runtime threads do not survive fork(), so native routes cannot run here."
);
pyo3::create_exception!(
_native,
ProcessReservedForForking,
PyRuntimeError,
"This process was reserved for forking workers, so native routes cannot run here."
);
static FORK_GATE: ForkGate = ForkGate::new();
/// Whether this process has started the Tokio runtime.
pub fn runtime_started() -> bool {
FORK_GATE.started(std::process::id())
}
/// Declares that this process exists to fork workers, so it must never start the runtime.
/// Fails if it already has. Workers are unaffected: the reservation is keyed by pid.
pub fn reserve_process_for_forking() -> Result<(), RuntimeAlreadyStarted> {
FORK_GATE.reserve(std::process::id())
}
/// The only door to the Tokio runtime: every route reaches it through this module, which is
/// what lets the gate speak for the whole extension. `clippy.toml` disallows going around it.
fn enter_runtime() -> PyResult<()> {
FORK_GATE
.enter(std::process::id())
.map_err(|refused| match refused {
Refused::ReservedForForking => ProcessReservedForForking::new_err(
"this process is reserved for forking workers and cannot run native routes; \
move the call into a worker, after the fork",
),
Refused::ForkedAfterStart => ForkedAfterNativeRuntimeStarted::new_err(
"this process was forked after the native runtime started, and runtime threads \
do not survive fork(); start workers with spawn or forkserver, or fork before \
the first native call",
),
})
}
#[expect(clippy::disallowed_methods, reason = "this is the gated door")]
fn runtime() -> PyResult<&'static Runtime> {
enter_runtime()?;
Ok(pyo3_async_runtimes::tokio::get_runtime())
}
#[expect(clippy::disallowed_methods, reason = "this is the gated door")]
fn future_into_py<F, T>(py: Python<'_>, future: F) -> PyResult<Bound<'_, PyAny>>
where
F: Future<Output = PyResult<T>> + Send + 'static,
T: for<'py> IntoPyObject<'py> + Send + 'static,
{
enter_runtime()?;
pyo3_async_runtimes::tokio::future_into_py(py, future)
}
pub fn run_sync<T, E, F>(
py: Python<'_>,
future: F,
@ -22,12 +84,7 @@ where
E: Send + 'static,
F: Future<Output = Result<T, E>> + Send + 'static,
{
run_sync_on(
py,
pyo3_async_runtimes::tokio::get_runtime(),
future,
map_error,
)
run_sync_on(py, runtime()?, future, map_error)
}
pub fn run_sync_value<T, F>(py: Python<'_>, future: F) -> PyResult<T>
@ -35,7 +92,7 @@ where
T: Send + 'static,
F: Future<Output = PyResult<T>> + Send + 'static,
{
run_sync_value_on(py, pyo3_async_runtimes::tokio::get_runtime(), future)
run_sync_value_on(py, runtime()?, future)
}
fn run_sync_value_on<T, F>(py: Python<'_>, runtime: &Runtime, future: F) -> PyResult<T>
@ -83,7 +140,7 @@ where
E: Send + 'static,
F: Future<Output = Result<T, E>> + Send + 'static,
{
pyo3_async_runtimes::tokio::future_into_py(py, async move {
future_into_py(py, async move {
let result = catch_future_panic(future).await?;
let result = map_core_result(result, map_error)?;
Ok(Pythonized(result))
@ -95,7 +152,7 @@ where
T: for<'py> IntoPyObject<'py> + Send + 'static,
F: Future<Output = PyResult<T>> + Send + 'static,
{
pyo3_async_runtimes::tokio::future_into_py(py, async move { catch_future_panic(future).await? })
future_into_py(py, async move { catch_future_panic(future).await? })
}
pub fn poll_async_value<T, F>(py: Python<'_>, future: Pin<&mut F>) -> PyResult<Poll<T>>
@ -103,8 +160,9 @@ where
T: Send,
F: Future<Output = PyResult<T>> + Send,
{
let runtime = runtime()?;
let result = release_gil(py, || {
let _runtime = pyo3_async_runtimes::tokio::get_runtime().enter();
let _runtime = runtime.enter();
std::panic::catch_unwind(AssertUnwindSafe(|| {
future.poll(&mut Context::from_waker(Waker::noop()))
}))
@ -286,27 +344,25 @@ mod tests {
}
#[pyfunction]
fn runtime_worker_count() -> usize {
pyo3_async_runtimes::tokio::get_runtime()
.metrics()
.num_workers()
fn runtime_worker_count() -> PyResult<usize> {
Ok(runtime()?.metrics().num_workers())
}
#[pyfunction]
fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool {
fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> PyResult<bool> {
let completion_deadline = Instant::now() + Duration::from_secs(2);
while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions {
if Instant::now() >= completion_deadline {
return false;
return Ok(false);
}
thread::sleep(Duration::from_millis(1));
}
let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1);
pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
runtime()?.spawn(async move {
let _ = heartbeat_tx.send(());
});
heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok()
Ok(heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok())
}
fn extract_bool(py: Python<'_>, result: PyResult<Py<PyAny>>) -> bool {
@ -317,6 +373,16 @@ mod tests {
.expect("result should convert")
}
#[rstest]
fn reaching_the_runtime_marks_the_process_as_started(
#[from(initialized_python)] python: &InitializedPython,
) {
python.attach(|py| {
run_sync_value(py, async { Ok(()) }).unwrap();
assert!(runtime_started());
});
}
#[rstest]
fn inline_poll_releases_gil_and_enters_runtime(
#[from(initialized_python)] python: &InitializedPython,

View file

@ -0,0 +1,139 @@
use std::sync::atomic::{AtomicU32, Ordering};
const UNSET: u32 = 0;
/// Decides which process may use the Tokio runtime. Its worker threads do not survive
/// `fork()`: a child forked after they started hangs on its first native call. The gate turns
/// both halves of that hazard into errors, keyed by pid so a fork needs no hook to be seen:
/// a process reserved for forking can never start the runtime, and a child of a process that
/// did start it is refused instead of hanging.
pub(crate) struct ForkGate {
runtime_pid: AtomicU32,
fork_only_pid: AtomicU32,
}
#[derive(Debug, PartialEq, Eq)]
pub(crate) enum Refused {
ReservedForForking,
ForkedAfterStart,
}
#[derive(Debug, PartialEq, Eq)]
pub struct RuntimeAlreadyStarted;
impl ForkGate {
pub(crate) const fn new() -> Self {
Self {
runtime_pid: AtomicU32::new(UNSET),
fork_only_pid: AtomicU32::new(UNSET),
}
}
/// Claims the runtime for `pid`. Claim first, then look for a reservation: `reserve` does
/// the mirror image, so when the two race at least one of them sees the other.
pub(crate) fn enter(&self, pid: u32) -> Result<(), Refused> {
match self
.runtime_pid
.compare_exchange(UNSET, pid, Ordering::SeqCst, Ordering::SeqCst)
{
Err(owner) if owner != pid => return Err(Refused::ForkedAfterStart),
_ => {}
}
if self.fork_only_pid.load(Ordering::SeqCst) == pid {
// Nothing was started, so the workers forked from here must still find it unclaimed.
let _ =
self.runtime_pid
.compare_exchange(pid, UNSET, Ordering::SeqCst, Ordering::SeqCst);
return Err(Refused::ReservedForForking);
}
Ok(())
}
/// Reserves `pid` for forking. Reserve first, then look for a started runtime: `enter` does
/// the mirror image, so when the two race at least one of them sees the other. A refused
/// reservation leaves the gate exactly as it was, so a process already running the runtime
/// keeps refusing the children it forks.
pub(crate) fn reserve(&self, pid: u32) -> Result<(), RuntimeAlreadyStarted> {
self.fork_only_pid.store(pid, Ordering::SeqCst);
if self.runtime_pid.load(Ordering::SeqCst) == pid {
let _ =
self.fork_only_pid
.compare_exchange(pid, UNSET, Ordering::SeqCst, Ordering::SeqCst);
return Err(RuntimeAlreadyStarted);
}
Ok(())
}
pub(crate) fn started(&self, pid: u32) -> bool {
self.runtime_pid.load(Ordering::SeqCst) == pid
}
}
#[cfg(test)]
mod tests {
use super::*;
const MASTER: u32 = 100;
const WORKER: u32 = 101;
#[test]
fn unreserved_process_starts_the_runtime_and_stays_started() {
let gate = ForkGate::new();
assert!(!gate.started(MASTER));
assert_eq!(gate.enter(MASTER), Ok(()));
assert_eq!(gate.enter(MASTER), Ok(()));
assert!(gate.started(MASTER));
}
#[test]
fn reserved_process_can_never_start_the_runtime() {
let gate = ForkGate::new();
assert_eq!(gate.reserve(MASTER), Ok(()));
assert_eq!(gate.enter(MASTER), Err(Refused::ReservedForForking));
assert_eq!(gate.enter(MASTER), Err(Refused::ReservedForForking));
assert!(!gate.started(MASTER));
}
#[test]
fn workers_forked_from_a_reserved_process_start_their_own_runtime() {
let gate = ForkGate::new();
gate.reserve(MASTER).unwrap();
gate.enter(MASTER).unwrap_err();
assert_eq!(gate.enter(WORKER), Ok(()));
assert!(gate.started(WORKER));
}
#[test]
fn reserving_after_the_runtime_started_is_refused() {
let gate = ForkGate::new();
gate.enter(MASTER).unwrap();
assert_eq!(gate.reserve(MASTER), Err(RuntimeAlreadyStarted));
}
#[test]
fn a_refused_reservation_leaves_the_runtime_claimed_and_its_children_refused() {
let gate = ForkGate::new();
gate.enter(MASTER).unwrap();
assert_eq!(gate.reserve(MASTER), Err(RuntimeAlreadyStarted));
assert_eq!(gate.enter(MASTER), Ok(()));
assert!(gate.started(MASTER));
assert_eq!(gate.enter(WORKER), Err(Refused::ForkedAfterStart));
}
#[test]
fn child_forked_after_the_runtime_started_is_refused_instead_of_hanging() {
let gate = ForkGate::new();
gate.enter(MASTER).unwrap();
assert_eq!(gate.enter(WORKER), Err(Refused::ForkedAfterStart));
assert!(!gate.started(WORKER));
assert_eq!(gate.enter(MASTER), Ok(()));
}
}

View file

@ -8,6 +8,7 @@ mod argument;
mod callable;
mod driver;
mod execution;
mod fork_gate;
mod gil;
mod handle;
mod marshal;
@ -18,7 +19,12 @@ pub use adapter::{
pub use argument::lookup;
pub use callable::wrap_failure;
pub use driver::run_call;
pub use execution::{poll_async_value, run_async, run_async_value, run_sync, run_sync_value};
pub use execution::{
ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, poll_async_value,
reserve_process_for_forking, run_async, run_async_value, run_sync, run_sync_value,
runtime_started,
};
pub use fork_gate::RuntimeAlreadyStarted;
pub use gil::{release_count, release_gil};
pub use handle::{Execution, ExecutionBody, ExecutionStep};
pub use marshal::{Pythonized, from_py, from_py_argument, panic_to_pyerr, to_py};

View file

@ -1,5 +1,5 @@
use litellm_host_python::release_count;
use pyo3::{prelude::*, types::PyDict};
use litellm_host_python::{release_count, runtime_started};
use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyDict};
#[pyfunction]
pub(crate) fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
@ -8,6 +8,20 @@ pub(crate) fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
Ok(stats.into_any().unbind())
}
/// True once this process has started the native runtime, which does not survive `fork()`.
#[pyfunction]
pub(crate) fn process_state_started() -> bool {
runtime_started()
}
/// Declares that this process only forks workers: from now on every native route raises here,
/// so the runtime can never start. Raises if it already has. Forked workers are unaffected.
#[pyfunction]
pub(crate) fn reserve_process_for_forking() -> PyResult<()> {
litellm_host_python::reserve_process_for_forking()
.map_err(|_| PyRuntimeError::new_err("the native runtime already started in this process"))
}
#[cfg(feature = "panic-test")]
#[pyfunction]
pub(crate) fn _panic_for_test() {

View file

@ -13,7 +13,7 @@ mod _native {
#[pymodule_export]
use crate::diagnostics::_panic_for_test;
#[pymodule_export]
use crate::diagnostics::gil_stats;
use crate::diagnostics::{gil_stats, process_state_started, reserve_process_for_forking};
#[pymodule_export]
use crate::errors::{RustBridgeDeclined, RustUpstreamError};
#[pymodule_export]
@ -30,6 +30,8 @@ mod _native {
use crate::routes::responses::ResponsesWebSocketConnection;
#[pymodule_export]
use crate::token_counter::TokenCounter;
#[pymodule_export]
use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking};
}
use pyo3::prelude::*;
@ -50,6 +52,8 @@ mod tests {
let mut expected = vec![
"RustBridgeDeclined",
"RustUpstreamError",
"ForkedAfterNativeRuntimeStarted",
"ProcessReservedForForking",
"ocr",
"aocr",
"transcription",
@ -62,6 +66,8 @@ mod tests {
"ResponsesWebSocketConnection",
"TokenCounter",
"gil_stats",
"process_state_started",
"reserve_process_for_forking",
];
expected.sort_unstable();

View file

@ -25,7 +25,7 @@ impl ResponsesWebSocketConnection {
) -> PyResult<Bound<'py, PyAny>> {
let headers = marshal_headers(headers)?;
let timeout = optional_timeout(timeout_seconds);
pyo3_async_runtimes::tokio::future_into_py(py, async move {
litellm_host_python::run_async_value(py, async move {
let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout)
.await
.map_err(responses_error_to_pyerr)?;
@ -35,7 +35,7 @@ impl ResponsesWebSocketConnection {
fn send_text<'py>(&self, py: Python<'py>, text: String) -> PyResult<Bound<'py, PyAny>> {
let inner = self.inner.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move {
litellm_host_python::run_async_value(py, async move {
inner
.send_text(text)
.await
@ -45,14 +45,14 @@ impl ResponsesWebSocketConnection {
fn recv_text<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let inner = self.inner.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move {
litellm_host_python::run_async_value(py, async move {
inner.recv_text().await.map_err(responses_error_to_pyerr)
})
}
fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let inner = self.inner.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move {
litellm_host_python::run_async_value(py, async move {
inner.close().await.map_err(responses_error_to_pyerr)
})
}
@ -68,6 +68,10 @@ mod tests {
use tokio_tungstenite::{accept_async, tungstenite::Message};
#[test]
#[expect(
clippy::disallowed_methods,
reason = "the test server shares the routes' runtime"
)]
fn responses_websocket_connection_round_trips_through_python() {
Python::initialize();
let runtime = pyo3_async_runtimes::tokio::get_runtime();

View file

@ -589,6 +589,11 @@ class ProxyInitializationHelpers:
gunicorn_options["certfile"] = ssl_certfile_path
gunicorn_options["keyfile"] = ssl_keyfile_path
# The master preloads the app and then forks every worker, so native routes are
# forbidden in it: their runtime threads would not survive the fork.
from litellm.rust_bridge.fork_guard import reserve_process_for_forking
reserve_process_for_forking("the gunicorn master")
start_query_engine_reaper()
StandaloneApplication(app=app, options=gunicorn_options).run() # Run gunicorn

View file

@ -9,6 +9,8 @@ from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMe
class RustBridgeDeclined(Exception): ...
class RustUpstreamError(Exception): ...
class ForkedAfterNativeRuntimeStarted(RuntimeError): ...
class ProcessReservedForForking(RuntimeError): ...
def ocr(
request: LiteLLMOcrRequest,
@ -101,8 +103,12 @@ class TokenCounter:
def acount_request(self, body: bytes) -> Future[dict[str, object]]: ...
def gil_stats() -> dict[str, int]: ...
def process_state_started() -> bool: ...
def reserve_process_for_forking() -> None: ...
__all__ = [
"ForkedAfterNativeRuntimeStarted",
"ProcessReservedForForking",
"ResponsesWebSocketConnection",
"RustBridgeDeclined",
"RustUpstreamError",
@ -116,5 +122,7 @@ __all__ = [
"gil_stats",
"messages",
"ocr",
"process_state_started",
"reserve_process_for_forking",
"transcription",
]

View file

@ -0,0 +1,47 @@
"""Fork safety of the Rust extension.
Its runtime threads do not survive ``fork``, so a child forked after the first native call
cannot run native routes: it raises ``ForkedAfterNativeRuntimeStarted`` instead of hanging.
Fork before the first native call, or start workers with ``spawn`` / ``forkserver``.
A process whose job is to fork workers (the gunicorn master under ``preload``) reserves itself:
from then on any native route called in it raises ``ProcessReservedForForking`` at the call
site, so the runtime can never start there. Workers forked from it are unaffected.
"""
from __future__ import annotations
from typing import Final
from litellm.rust_bridge.loader import get_native_bridge
class NativeStateStartedBeforeFork(RuntimeError):
pass
class _NeverRaised(RuntimeError):
"""Stands in for a native exception when the extension is unavailable or predates it."""
_native: Final = get_native_bridge()
ForkedAfterNativeRuntimeStarted: Final[type[RuntimeError]] = getattr(
_native, "ForkedAfterNativeRuntimeStarted", _NeverRaised
)
ProcessReservedForForking: Final[type[RuntimeError]] = getattr(_native, "ProcessReservedForForking", _NeverRaised)
def reserve_process_for_forking(where: str) -> None:
"""Forbid native routes in this process. Raises if one already ran here."""
native: Final = get_native_bridge()
reserve: Final = getattr(native, "reserve_process_for_forking", None)
if not callable(reserve):
return
try:
reserve()
except RuntimeError as error:
raise NativeStateStartedBeforeFork(
f"The LiteLLM Rust extension already ran a native route in {where}, and its runtime "
"threads do not survive fork(). Move the native call (warm-up, health check, "
"import-time initialization) into the worker, after the fork."
) from error

View file

@ -60,6 +60,13 @@ TQ007 A module global that a conftest saves before every test and restores aft
names are read from the keys the conftest assigns directly and from whatever the
save loop iterates, including a module-level tuple or dict it names rather than
spells out.
TQ009 A child interpreter spawned as `subprocess.run([sys.executable, ...])` without
`-I`/`-P` as its first flag. Without isolation the child's sys.path leads with
the working directory, so a source checkout shadows the installed package and
the child tests a different `litellm` than the parent imported -- TQ003 is the
same working-directory hazard seen from the child's side. Use
tests.test_litellm_rust.support.child_interpreter.run_child_interpreter, which
also asserts the child resolved the same `litellm.__file__` as the parent.
Every rule is suppressible with `# test-quality-ok: <reason>` on the reported
line, following the repo's `*-ok: <reason>` convention. A suppression without a
@ -140,6 +147,9 @@ SKIP_CALLS: Final = frozenset(("pytest.skip", "skip"))
CONFTEST_NAME: Final = "conftest.py"
SDK_MODULE: Final = "litellm"
SUBPROCESS_SPAWNS: Final = frozenset(("run", "Popen", "check_output", "check_call", "call"))
INTERPRETER_ISOLATION_FLAGS: Final = frozenset(("-I", "-P"))
CREDENTIAL_NAME_RE: Final = re.compile(
r"(?:API_KEY|_KEY|TOKEN|SECRET|PASSWORD|CREDENTIAL|DATABASE_URL|ACCESS_KEY_ID)$"
)
@ -709,6 +719,35 @@ def _snapshotted_names(tree: ast.Module) -> Iterator[tuple[str, int]]:
yield from _string_members(iterable)
def iter_child_interpreter_violations(path: Path, tree: ast.Module) -> Iterator[Violation]:
for node in ast.walk(tree):
if not (isinstance(node, ast.Call) and node.args):
continue
if _dotted_name(node.func).rsplit(".", 1)[-1] not in SUBPROCESS_SPAWNS:
continue
argv: Final = node.args[0]
if not isinstance(argv, (ast.List, ast.Tuple)) or not argv.elts:
continue
if _dotted_name(argv.elts[0]) != "sys.executable":
continue
isolated: Final = (
len(argv.elts) > 1
and isinstance(argv.elts[1], ast.Constant)
and argv.elts[1].value in INTERPRETER_ISOLATION_FLAGS
)
if isolated:
continue
yield Violation(
path,
node.lineno,
"TQ009",
"child interpreter spawned without -I/-P; the working directory lands on sys.path "
"and a source checkout can shadow the installed package, use "
"tests.test_litellm_rust.support.child_interpreter.run_child_interpreter or pass -I "
f"(suppress: `# {SUPPRESSION_TOKEN}: <reason>`)",
)
def iter_conftest_inventory_violations(path: Path, tree: ast.Module) -> Iterator[Violation]:
if path.name != CONFTEST_NAME:
return
@ -746,6 +785,7 @@ def check_file(path: Path) -> tuple[Violation, ...]:
*iter_credential_skip_violations(path, tree),
*iter_conftest_inventory_violations(path, tree),
*iter_internal_patch_violations(path, tree),
*iter_child_interpreter_violations(path, tree),
)
if violation.line not in skip
)

View file

@ -22,5 +22,8 @@
},
"TQ008": {
"limit": 10993
},
"TQ009": {
"limit": 59
}
}

View file

@ -21,6 +21,15 @@ from uvicorn.importer import import_from_string
from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server
@pytest.fixture(autouse=True)
def fork_reservation():
"""Reserving is irreversible: it would forbid native routes in this pytest worker for good"""
with patch( # test-quality-ok: process-global native state, a real reservation would poison every later test in the worker
"litellm.rust_bridge.fork_guard.reserve_process_for_forking"
) as reserve:
yield reserve
@pytest.mark.xdist_group("proxy_cli")
class TestProxyInitializationHelpers:
@patch("importlib.metadata.version")
@ -1574,6 +1583,32 @@ class TestProxyInitializationHelpers:
assert captured["options"]["max_requests"] == 1000
assert captured["options"]["max_requests_jitter"] == 50
@pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows")
def test_gunicorn_master_is_reserved_for_forking_before_it_runs(self, fork_reservation):
"""preload forks workers from the master, so native routes are forbidden there first"""
pytest.importorskip("gunicorn")
reserved_before_run: list = []
def capture_run(self):
reserved_before_run.append(fork_reservation.call_args)
with (
patch("gunicorn.app.base.BaseApplication.run", capture_run),
patch( # test-quality-ok: option tests must not start a thread or change the pytest worker's child ownership
"litellm.proxy.proxy_cli.start_query_engine_reaper"
),
):
ProxyInitializationHelpers._run_gunicorn_server(
host="127.0.0.1",
port=4012,
app=MagicMock(),
num_workers=2,
ssl_certfile_path=None,
ssl_keyfile_path=None,
)
assert [call.args for call in reserved_before_run] == [("the gunicorn master",)]
@pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows")
def test_gunicorn_jitter_without_base_warns(self):
"""gunicorn path warns when jitter is set without --max_requests_before_restart"""

View file

@ -0,0 +1,36 @@
from types import SimpleNamespace
import pytest
from litellm.rust_bridge import fork_guard
def _reserve_with(monkeypatch: pytest.MonkeyPatch, native: object) -> None:
monkeypatch.setattr(fork_guard, "get_native_bridge", lambda: native)
fork_guard.reserve_process_for_forking("the gunicorn master")
def test_missing_extension_has_nothing_to_reserve(monkeypatch: pytest.MonkeyPatch) -> None:
assert _reserve_with(monkeypatch, None) is None
def test_extension_built_before_reservation_existed_passes(monkeypatch: pytest.MonkeyPatch) -> None:
assert _reserve_with(monkeypatch, SimpleNamespace()) is None
def test_unused_extension_is_reserved(monkeypatch: pytest.MonkeyPatch) -> None:
calls: list[None] = []
_reserve_with(monkeypatch, SimpleNamespace(reserve_process_for_forking=lambda: calls.append(None)))
assert calls == [None]
def test_used_extension_refuses_and_names_the_place(monkeypatch: pytest.MonkeyPatch) -> None:
def reserve() -> None:
raise RuntimeError("the native runtime already started in this process")
with pytest.raises(fork_guard.NativeStateStartedBeforeFork, match="the gunicorn master") as raised:
_reserve_with(monkeypatch, SimpleNamespace(reserve_process_for_forking=reserve))
assert isinstance(raised.value.__cause__, RuntimeError)

View file

@ -737,3 +737,28 @@ def test_a_fanned_out_run_reports_each_generated_file_exactly_once(tmp_path):
assert len(reported) == len(paths)
assert len({line.split(":")[0] for line in reported}) == len(paths)
assert all(" TQ001 " in line for line in reported)
def test_sys_executable_child_without_isolation_flag_is_flagged(tmp_path):
source = 'import subprocess, sys\nsubprocess.run([sys.executable, "-c", "pass"])\n'
assert _codes(tmp_path, source) == ["TQ009"]
def test_sys_executable_child_with_dash_i_is_clean(tmp_path):
source = 'import subprocess, sys\nsubprocess.run([sys.executable, "-I", "-c", "pass"])\n'
assert _codes(tmp_path, source) == []
def test_sys_executable_child_with_dash_p_is_clean(tmp_path):
source = 'import subprocess, sys\nsubprocess.run([sys.executable, "-P", "-c", "pass"])\n'
assert _codes(tmp_path, source) == []
def test_non_interpreter_subprocess_call_is_untouched(tmp_path):
source = 'import subprocess\nsubprocess.run(["python", "-c", "pass"])\n'
assert _codes(tmp_path, source) == []
def test_popen_sys_executable_tuple_is_flagged(tmp_path):
source = 'import subprocess, sys\nsubprocess.Popen((sys.executable, "script.py"))\n'
assert _codes(tmp_path, source) == ["TQ009"]

View file

@ -136,7 +136,9 @@ def test_the_shipped_budget_covers_every_rule_the_checker_can_emit():
import json
budget = json.loads((_REPO_ROOT / "test-quality-budget.json").read_text())
assert set(budget) == {"TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ008"}
assert set(budget) == {
"TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ008", "TQ009"
}
assert all(spec["limit"] >= 0 for spec in budget.values())

View file

@ -0,0 +1,36 @@
from __future__ import annotations
import os
import subprocess
import sys
from collections.abc import Mapping
from typing import Final
import litellm
PARENT_LITELLM_FILE: Final = "LITELLM_TEST_PARENT_LITELLM_FILE"
_PROLOGUE: Final = (
"import os as _os, litellm as _litellm; _parent = _os.environ.pop({key!r}); "
'assert _litellm.__file__ == _parent, f"child imported litellm from {{_litellm.__file__}}, parent from {{_parent}}"; '
"del _os, _litellm, _parent\n"
)
def run_child_interpreter(
source: str, *, env: Mapping[str, str] | None = None, timeout: float
) -> subprocess.CompletedProcess[str]:
"""Run `source` in a fresh interpreter that imports the same `litellm` as this process.
`-I` keeps the working directory off sys.path so a source checkout cannot shadow an
installed wheel, and the prologue fails fast with both paths if the child still
resolves a different package.
"""
environment: Final = {**(os.environ if env is None else env), PARENT_LITELLM_FILE: litellm.__file__}
return subprocess.run(
[sys.executable, "-I", "-c", _PROLOGUE.format(key=PARENT_LITELLM_FILE) + source],
capture_output=True,
text=True,
timeout=timeout,
env=environment,
)

View file

@ -0,0 +1,146 @@
import os
import textwrap
import pytest
from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter
pytestmark = pytest.mark.requires_rust_extension
_NATIVE_CONTRACT = textwrap.dedent(
"""
import os
from litellm.rust_bridge import _native
from litellm.rust_bridge.fork_guard import reserve_process_for_forking
def native_route_error():
import asyncio
async def call():
await _native.ResponsesWebSocketConnection.connect("ws://127.0.0.1:1", {}, 0.2)
try:
asyncio.run(call())
except Exception as error:
return f"{type(error).__name__}: {error}"
return ""
assert _native.process_state_started() is False
reserve_process_for_forking("the test master")
assert native_route_error().startswith("ProcessReservedForForking: ")
assert _native.process_state_started() is False
pid = os.fork()
if pid == 0:
error = native_route_error()
started = _native.process_state_started()
os._exit(0 if started and "reserved" not in error and "forked" not in error else 1)
assert os.waitpid(pid, 0)[1] == 0
pid = os.fork()
if pid == 0:
native_route_error()
grandchild = os.fork()
if grandchild == 0:
os._exit(0 if native_route_error().startswith("ForkedAfterNativeRuntimeStarted: ") else 1)
os._exit(os.waitpid(grandchild, 0)[1])
assert os.waitpid(pid, 0)[1] == 0
"""
)
@pytest.mark.skipif(not hasattr(os, "fork"), reason="fork only")
def test_compiled_extension_forbids_the_master_and_frees_its_workers() -> None:
env = {**os.environ, "OBJC_DISABLE_INITIALIZE_FORK_SAFETY": "YES"}
result = run_child_interpreter(_NATIVE_CONTRACT, env=env, timeout=60)
assert result.returncode == 0, result.stderr
_SDK_CONTRACT = textwrap.dedent(
"""
import asyncio, json, multiprocessing, os, threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import litellm
from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted
class Handler(BaseHTTPRequestHandler):
def do_POST(self):
self.rfile.read(int(self.headers["Content-Length"]))
if self.headers.get("User-Agent", "").startswith("python-httpx"):
self.send_response(418)
self.end_headers()
return
body = json.dumps({
"pages": [{"index": 0, "markdown": "native", "images": [], "dimensions": None}],
"model": "mistral-ocr-latest",
"usage_info": {"pages_processed": 1, "doc_size_bytes": 3},
}).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def log_message(self, *args):
pass
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
threading.Thread(target=server.serve_forever, daemon=True).start()
arguments = {
"model": "mistral/mistral-ocr-latest",
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
"api_key": "test-key",
"api_base": f"http://127.0.0.1:{server.server_port}",
"num_retries": 0,
}
litellm.rust(True)
SERVED, REFUSED, OTHER = 0, 3, 4
def outcome(asynchronous):
try:
response = asyncio.run(litellm.aocr(**arguments)) if asynchronous else litellm.ocr(**arguments)
except ForkedAfterNativeRuntimeStarted:
return REFUSED
except Exception:
return OTHER
return SERVED if response.pages[0].markdown == "native" else OTHER
def forked(asynchronous):
pid = os.fork()
if pid == 0:
os._exit(outcome(asynchronous))
return os.waitstatus_to_exitcode(os.waitpid(pid, 0)[1])
def pooled(asynchronous):
with multiprocessing.get_context("fork").Pool(1) as pool:
return pool.apply(outcome, (asynchronous,))
# Forking before the first native call is fine: the child starts its own runtime.
assert [forked(False), forked(True)] == [SERVED, SERVED]
assert outcome(False) == SERVED
# After it, a forked child is told so instead of hanging on threads that do not exist.
assert [forked(False), forked(True)] == [REFUSED, REFUSED]
assert [pooled(False), pooled(True)] == [REFUSED, REFUSED]
# The parent is not poisoned by any of it.
assert [outcome(False), outcome(True)] == [SERVED, SERVED]
"""
)
@pytest.mark.skipif(not hasattr(os, "fork"), reason="fork only")
def test_sdk_call_in_a_child_forked_after_native_use_raises_instead_of_hanging() -> None:
env = {
**os.environ,
"OBJC_DISABLE_INITIALIZE_FORK_SAFETY": "YES",
"LITELLM_RUST": "1",
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
}
result = run_child_interpreter(_SDK_CONTRACT, env=env, timeout=120)
assert result.returncode == 0, result.stderr