diff --git a/litellm/containers/main.py b/litellm/containers/main.py index 625a291fb5..105e999ffe 100644 --- a/litellm/containers/main.py +++ b/litellm/containers/main.py @@ -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( diff --git a/litellm/llms/openai/containers/transformation.py b/litellm/llms/openai/containers/transformation.py index 46718816f3..e67bfbe0c6 100644 --- a/litellm/llms/openai/containers/transformation.py +++ b/litellm/llms/openai/containers/transformation.py @@ -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" diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 2059246674..ddc79a2865 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -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( diff --git a/tests/test_litellm/containers/test_container_regional_api_base.py b/tests/test_litellm/containers/test_container_regional_api_base.py new file mode 100644 index 0000000000..7c6154867f --- /dev/null +++ b/tests/test_litellm/containers/test_container_regional_api_base.py @@ -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 + diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 5f03ef1817..99b4ebba06 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -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"""