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:
parent
92827ead65
commit
ccc0e342f2
@ -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(
|
||||
|
||||
@ -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"
|
||||
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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"""
|
||||
|
||||
Loading…
Reference in New Issue
Block a user