diff --git a/litellm/proxy/db/connection_budget.py b/litellm/proxy/db/connection_budget.py index b55a22d4383..f8c59df0a29 100644 --- a/litellm/proxy/db/connection_budget.py +++ b/litellm/proxy/db/connection_budget.py @@ -7,8 +7,9 @@ queries next to ported ones. These invariants hold until Prisma is gone: `database_extra_connection_params`). The Rust share is carved out of it, never added on top, so a Rust pool does not grow the connection budget. - The Rust share defaults to `RUST_RESERVED_CONNECTIONS`. An operator moves connections from Prisma to the - Rust pool by setting `LITELLM_DATABASE_RUST_WRITER_CONNECTION_LIMIT` before startup, and `proxy_cli` - clamps it so Prisma keeps at least one. + Rust pool by setting `LITELLM_DATABASE_RUST_WRITER_CONNECTION_LIMIT` or + `LITELLM_DATABASE_RUST_READER_CONNECTION_LIMIT` before startup, and `proxy_cli` clamps it so Prisma keeps at + least one. - The split is Python's. It is about Prisma and Prisma lives only in the Python gateway, so this module is deleted with Prisma. The process that starts the workers (`proxy_cli`) decides it once and hands each worker plain numbers. diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 1a1e443748d..7d1fb14155d 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1621,7 +1621,8 @@ def run_server( ConnectionBudget.split( TypeAdapter(PositiveInt).validate_python( _url_query_value(read_replica_url, "connection_limit") or writer_budget.total - ) + ), + rust_reserved=rust_reserved.reader_connection_limit, ) if read_replica_url else None diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 0b5099fc800..8c99876c132 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -3665,6 +3665,7 @@ def _run_server_and_capture_urls( read_replica_url: str | None = None, extra_args: tuple[str, ...] = (), rust_writer_connection_limit: str | None = None, + rust_reader_connection_limit: str | None = None, ) -> dict: loaded_config = yaml.safe_load(Path(config_path).read_text()) mock_proxy_config = MagicMock() @@ -3683,6 +3684,8 @@ def _run_server_and_capture_urls( clean_env["DATABASE_URL_READ_REPLICA"] = read_replica_url if rust_writer_connection_limit is not None: clean_env["LITELLM_DATABASE_RUST_WRITER_CONNECTION_LIMIT"] = rust_writer_connection_limit + if rust_reader_connection_limit is not None: + clean_env["LITELLM_DATABASE_RUST_READER_CONNECTION_LIMIT"] = rust_reader_connection_limit with ( patch.dict(os.environ, clean_env, clear=True), @@ -3860,6 +3863,23 @@ class TestReadReplicaConnectionParams: assert int(query["connection_limit"][0]) + int(captured["LITELLM_DATABASE_RUST_READER_CONNECTION_LIMIT"]) == 50 assert query["pool_timeout"] == ["20"] + def test_operator_set_rust_share_is_carved_out_of_the_reader_limit(self, tmp_path: Path) -> None: + config_path: Final = tmp_path / "config.yaml" + config_path.write_text(yaml.dump({"model_list": [], "general_settings": {}})) + + captured: Final = _CAPTURED_URLS.validate_python( + _run_server_and_capture_urls( + str(config_path), + read_replica_url="postgresql://t:t@reader:5432/t?connection_limit=12", + rust_reader_connection_limit="5", + ) + ) + + reader: Final = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL_READ_REPLICA"]).query) + writer: Final = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query) + assert (reader["connection_limit"], captured["LITELLM_DATABASE_RUST_READER_CONNECTION_LIMIT"]) == (["7"], "5") + assert (writer["connection_limit"], captured["LITELLM_DATABASE_RUST_WRITER_CONNECTION_LIMIT"]) == (["8"], "2") + def test_connection_budget_line_counts_the_limits_the_final_urls_carry( self, tmp_path: Path,