Make keepalive_timeout parameter work for Gunicorn (#19087)

* [Fix] Containers API - Allow routing to regional endpoints (#19118)

* fix get_complete_url

* fix url resolution containers API

* TestContainerRegionalApiBase

* feat(proxy): add keepalive_timeout support for Gunicorn server

Add configurable keepalive timeout parameter for Gunicorn workers to
match existing Uvicorn functionality. This allows users to tune the
keep-alive connection timeout based on their deployment requirements.

Changes:
- Add keepalive_timeout parameter to _run_gunicorn_server method
- Configure Gunicorn's keepalive setting (defaults to 90s if not specified)
- Update --keepalive_timeout CLI help text to document both Uvicorn and Gunicorn behavior
- Pass keepalive_timeout from run_server to _run_gunicorn_server

Tests:
- Add test to verify keepalive_timeout flag is properly passed to Gunicorn
- Add test to verify default 90s timeout when flag is not specified

Co-Authored-By: lizhen921 <294474470@qq.com>
Signed-off-by: Kris Xia <xiajiayi0506@gmail.com>

---------

Signed-off-by: Kris Xia <xiajiayi0506@gmail.com>
Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Co-authored-by: lizhen921 <294474470@qq.com>
This commit is contained in:
Kris Xia 2026-01-16 06:02:59 +08:00 committed by GitHub
parent 92827ead65
commit ccc0e342f2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 288 additions and 9 deletions

View File

@ -199,7 +199,13 @@ def create_container(
return response
# get llm provider logic
litellm_params = GenericLiteLLMParams(**kwargs)
# Pass credential params explicitly since they're named args, not in kwargs
litellm_params = GenericLiteLLMParams(
api_key=api_key,
api_base=api_base,
api_version=api_version,
**kwargs,
)
# get provider config
container_provider_config: Optional[BaseContainerConfig] = (
ProviderConfigManager.get_provider_container_config(
@ -406,7 +412,13 @@ def list_containers(
return response
# get llm provider logic
litellm_params = GenericLiteLLMParams(**kwargs)
# Pass credential params explicitly since they're named args, not in kwargs
litellm_params = GenericLiteLLMParams(
api_key=api_key,
api_base=api_base,
api_version=api_version,
**kwargs,
)
# get provider config
container_provider_config: Optional[BaseContainerConfig] = (
ProviderConfigManager.get_provider_container_config(
@ -594,7 +606,13 @@ def retrieve_container(
return response
# get llm provider logic
litellm_params = GenericLiteLLMParams(**kwargs)
# Pass credential params explicitly since they're named args, not in kwargs
litellm_params = GenericLiteLLMParams(
api_key=api_key,
api_base=api_base,
api_version=api_version,
**kwargs,
)
# get provider config
container_provider_config: Optional[BaseContainerConfig] = (
ProviderConfigManager.get_provider_container_config(
@ -774,7 +792,13 @@ def delete_container(
return response
# get llm provider logic
litellm_params = GenericLiteLLMParams(**kwargs)
# Pass credential params explicitly since they're named args, not in kwargs
litellm_params = GenericLiteLLMParams(
api_key=api_key,
api_base=api_base,
api_version=api_version,
**kwargs,
)
# get provider config
container_provider_config: Optional[BaseContainerConfig] = (
ProviderConfigManager.get_provider_container_config(
@ -968,7 +992,13 @@ def list_container_files(
return response
# get llm provider logic
litellm_params = GenericLiteLLMParams(**kwargs)
# Pass credential params explicitly since they're named args, not in kwargs
litellm_params = GenericLiteLLMParams(
api_key=api_key,
api_base=api_base,
api_version=api_version,
**kwargs,
)
# get provider config
container_provider_config: Optional[BaseContainerConfig] = (
ProviderConfigManager.get_provider_container_config(
@ -1203,7 +1233,13 @@ def upload_container_file(
return response
# get llm provider logic
litellm_params = GenericLiteLLMParams(**kwargs)
# Pass credential params explicitly since they're named args, not in kwargs
litellm_params = GenericLiteLLMParams(
api_key=api_key,
api_base=api_base,
api_version=api_version,
**kwargs,
)
# get provider config
container_provider_config: Optional[BaseContainerConfig] = (
ProviderConfigManager.get_provider_container_config(

View File

@ -83,8 +83,13 @@ class OpenAIContainerConfig(BaseContainerConfig):
) -> str:
"""Get the complete URL for OpenAI container API.
"""
if api_base is None:
api_base = "https://api.openai.com/v1"
api_base = (
api_base
or litellm.api_base
or get_secret_str("OPENAI_BASE_URL")
or get_secret_str("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
return f"{api_base.rstrip('/')}/containers"

View File

@ -187,6 +187,7 @@ class ProxyInitializationHelpers:
ssl_certfile_path: str,
ssl_keyfile_path: str,
max_requests_before_restart: Optional[int] = None,
keepalive_timeout: Optional[int] = None,
):
"""
Run litellm with `gunicorn`
@ -267,6 +268,10 @@ class ProxyInitializationHelpers:
"access_log_format": '%(h)s %(l)s %(u)s %(t)s "%(r)s" %(s)s %(b)s',
}
# Optional: set keepalive timeout if specified by user
if keepalive_timeout is not None:
gunicorn_options["keepalive"] = keepalive_timeout
# Optional: recycle workers after N requests to mitigate memory growth
if max_requests_before_restart is not None:
gunicorn_options["max_requests"] = max_requests_before_restart
@ -489,7 +494,7 @@ class ProxyInitializationHelpers:
"--keepalive_timeout",
default=None,
type=int,
help="Set the uvicorn keepalive timeout in seconds (uvicorn timeout_keep_alive parameter)",
help="Set the keepalive timeout in seconds. For Uvicorn: timeout_keep_alive parameter. For Gunicorn: keepalive parameter. Default: Uvicorn uses ~75s, Gunicorn uses 90s",
envvar="KEEPALIVE_TIMEOUT",
)
@click.option(
@ -859,6 +864,7 @@ def run_server( # noqa: PLR0915
ssl_certfile_path=ssl_certfile_path,
ssl_keyfile_path=ssl_keyfile_path,
max_requests_before_restart=max_requests_before_restart,
keepalive_timeout=keepalive_timeout,
)
elif run_hypercorn is True:
ProxyInitializationHelpers._init_hypercorn_server(

View File

@ -0,0 +1,163 @@
"""
Tests for OpenAI Containers API regional api_base support.
Validates that litellm.create_container and litellm.upload_container_file
correctly use regional endpoints like https://us.api.openai.com/v1 for
US Data Residency instead of defaulting to https://api.openai.com/v1.
"""
import os
import sys
from unittest.mock import MagicMock, patch
import httpx
import pytest
sys.path.insert(0, os.path.abspath("../../.."))
import litellm
class TestContainerRegionalApiBase:
"""Test suite for container API regional api_base support."""
def setup_method(self):
"""Set up test fixtures."""
os.environ["OPENAI_API_KEY"] = "sk-test123"
def teardown_method(self):
"""Clean up after tests."""
if "OPENAI_API_KEY" in os.environ:
del os.environ["OPENAI_API_KEY"]
if "OPENAI_BASE_URL" in os.environ:
del os.environ["OPENAI_BASE_URL"]
if "OPENAI_API_BASE" in os.environ:
del os.environ["OPENAI_API_BASE"]
litellm.api_base = None
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_create_container_uses_regional_api_base(self, mock_post):
"""
Test that litellm.create_container uses the regional api_base when provided.
This validates the fix for US Data Residency support where requests should
go to https://us.api.openai.com/v1 instead of https://api.openai.com/v1.
"""
mock_response = MagicMock(spec=httpx.Response)
mock_response.status_code = 200
mock_response.json.return_value = {
"id": "cntr_123456",
"object": "container",
"created_at": 1747857508,
"status": "running",
"expires_after": {"anchor": "last_active_at", "minutes": 20},
"last_active_at": 1747857508,
"name": "Test Container"
}
mock_post.return_value = mock_response
litellm.create_container(
name="Test Container",
custom_llm_provider="openai",
api_base="https://us.api.openai.com/v1",
)
mock_post.assert_called_once()
call_args = mock_post.call_args
called_url = call_args[1]["url"]
assert "us.api.openai.com" in called_url, f"Expected US regional URL, got: {called_url}"
assert called_url == "https://us.api.openai.com/v1/containers"
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_create_container_uses_env_var_openai_base_url(self, mock_post):
"""
Test that litellm.create_container uses OPENAI_BASE_URL env var.
"""
os.environ["OPENAI_BASE_URL"] = "https://us.api.openai.com/v1"
mock_response = MagicMock(spec=httpx.Response)
mock_response.status_code = 200
mock_response.json.return_value = {
"id": "cntr_123456",
"object": "container",
"created_at": 1747857508,
"status": "running",
"expires_after": {"anchor": "last_active_at", "minutes": 20},
"last_active_at": 1747857508,
"name": "Test Container"
}
mock_post.return_value = mock_response
litellm.create_container(
name="Test Container",
custom_llm_provider="openai",
)
mock_post.assert_called_once()
call_args = mock_post.call_args
called_url = call_args[1]["url"]
assert "us.api.openai.com" in called_url, f"Expected US regional URL, got: {called_url}"
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_create_container_defaults_to_standard_openai(self, mock_post):
"""
Test that litellm.create_container defaults to standard OpenAI URL
when no regional api_base is configured.
"""
mock_response = MagicMock(spec=httpx.Response)
mock_response.status_code = 200
mock_response.json.return_value = {
"id": "cntr_123456",
"object": "container",
"created_at": 1747857508,
"status": "running",
"expires_after": {"anchor": "last_active_at", "minutes": 20},
"last_active_at": 1747857508,
"name": "Test Container"
}
mock_post.return_value = mock_response
litellm.create_container(
name="Test Container",
custom_llm_provider="openai",
)
mock_post.assert_called_once()
call_args = mock_post.call_args
called_url = call_args[1]["url"]
assert called_url == "https://api.openai.com/v1/containers"
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_upload_container_file_uses_regional_api_base(self, mock_post):
"""
Test that litellm.upload_container_file uses the regional api_base when provided.
"""
mock_response = MagicMock(spec=httpx.Response)
mock_response.status_code = 200
mock_response.json.return_value = {
"id": "file_123456",
"object": "container.file",
"created_at": 1747857508,
"container_id": "cntr_123456",
"path": "/mnt/user/data.csv",
"source": "user",
}
mock_post.return_value = mock_response
litellm.upload_container_file(
container_id="cntr_123456",
file=("data.csv", b"col1,col2\n1,2", "text/csv"),
custom_llm_provider="openai",
api_base="https://us.api.openai.com/v1",
)
mock_post.assert_called_once()
call_args = mock_post.call_args
called_url = call_args[1]["url"]
assert "us.api.openai.com" in called_url, f"Expected US regional URL, got: {called_url}"
assert "cntr_123456/files" in called_url

View File

@ -483,6 +483,75 @@ class TestProxyInitializationHelpers:
# Verify that uvicorn.run was called again
mock_uvicorn_run.assert_called_once()
@patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers._run_gunicorn_server")
@patch("builtins.print")
def test_gunicorn_keepalive_timeout_flag(self, mock_print, mock_gunicorn):
"""Test that the keepalive_timeout flag is properly passed to Gunicorn"""
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_app = MagicMock()
mock_proxy_config = MagicMock()
mock_key_mgmt = MagicMock()
mock_save_worker_config = MagicMock()
with patch.dict(
"sys.modules",
{
"proxy_server": MagicMock(
app=mock_app,
ProxyConfig=mock_proxy_config,
KeyManagementSettings=mock_key_mgmt,
save_worker_config=mock_save_worker_config,
)
},
):
result = runner.invoke(
run_server, ["--local", "--run_gunicorn", "--keepalive_timeout", "120"]
)
assert result.exit_code == 0
# Verify _run_gunicorn_server was called with keepalive_timeout
mock_gunicorn.assert_called_once()
call_kwargs = mock_gunicorn.call_args.kwargs
assert call_kwargs["keepalive_timeout"] == 120
@patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers._run_gunicorn_server")
@patch("builtins.print")
def test_gunicorn_keepalive_default(self, mock_print, mock_gunicorn):
"""Test that Gunicorn uses default 90s when keepalive_timeout not specified"""
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_app = MagicMock()
mock_proxy_config = MagicMock()
mock_key_mgmt = MagicMock()
mock_save_worker_config = MagicMock()
with patch.dict(
"sys.modules",
{
"proxy_server": MagicMock(
app=mock_app,
ProxyConfig=mock_proxy_config,
KeyManagementSettings=mock_key_mgmt,
save_worker_config=mock_save_worker_config,
)
},
):
result = runner.invoke(run_server, ["--local", "--run_gunicorn"])
assert result.exit_code == 0
# Verify default behavior (keepalive_timeout is None, Gunicorn will use 90)
call_kwargs = mock_gunicorn.call_args.kwargs
assert call_kwargs.get("keepalive_timeout") is None
class TestHealthAppFactory:
"""Test cases for the health app factory module"""