diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 31ceb338dcd..e338f490496 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -29,8 +29,6 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): def __init__(self, bucket_name: str | None = None) -> None: from litellm.proxy.proxy_server import premium_user - super().__init__(bucket_name=bucket_name) - self.batch_size = int(os.getenv("GCS_BATCH_SIZE", GCS_DEFAULT_BATCH_SIZE)) self.flush_interval = int(os.getenv("GCS_FLUSH_INTERVAL", GCS_DEFAULT_FLUSH_INTERVAL_SECONDS)) self.use_batched_logging = ( @@ -38,6 +36,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): ) self.flush_lock = asyncio.Lock() super().__init__( + bucket_name=bucket_name, flush_lock=self.flush_lock, batch_size=self.batch_size, flush_interval=self.flush_interval, diff --git a/litellm/proxy/common_utils/config_includes.py b/litellm/proxy/common_utils/config_includes.py new file mode 100644 index 00000000000..c1bb5ae952f --- /dev/null +++ b/litellm/proxy/common_utils/config_includes.py @@ -0,0 +1,132 @@ +import os +from collections.abc import Awaitable, Mapping +from types import MappingProxyType +from typing import Final, Protocol + +from litellm._logging import verbose_proxy_logger + +INCLUDE_KEY: Final = "include" + + +def resolve_include_file_path(include_file: str, declared_in: str, root_config_path: str) -> str: + """ + Resolve one `include` entry to the file it names, next to the config that declares it. + + A config written before nested entries resolved this way can name a file sitting next to the root + config instead, so that file is still read, with a warning naming where it was found. When both + files exist the one next to the declaring config wins and the other is named in a warning. + """ + declared_relative: Final = os.path.abspath(os.path.join(os.path.dirname(declared_in), include_file)) + root_relative: Final = os.path.abspath(os.path.join(os.path.dirname(root_config_path), include_file)) + if root_relative == declared_relative or not os.path.exists(root_relative): + return declared_relative + + if not os.path.exists(declared_relative): + verbose_proxy_logger.warning( + "Config include '%s' declared in %s was not found next to it, so %s was read instead. " + "Move the included file next to the config that declares it.", + include_file, + declared_in, + root_relative, + ) + return root_relative + + verbose_proxy_logger.warning( + "Config include '%s' declared in %s matches two files. %s sits next to that config and was read, " + "so %s was skipped. Rename one of the two to say which one you meant.", + include_file, + declared_in, + declared_relative, + root_relative, + ) + return declared_relative + + +class IncludeResolver(Protocol): + def __call__(self, include_entry: str, declared_in: str, /) -> str: ... + + +class ConfigReader(Protocol): + def __call__(self, location: str, /) -> Awaitable[Mapping[str, object]]: ... + + +def _merged_value(base_value: object, included_value: object) -> object: + if isinstance(included_value, list) and isinstance(base_value, list): + return [*base_value, *included_value] # mutable-ok: a merged config value stays the plain list the proxy loads + return included_value + + +def _merged_entry(base: Mapping[str, object], included: Mapping[str, object], key: str) -> object: + if key not in included: + return base[key] + return _merged_value(base.get(key), included[key]) + + +def _merged(base: Mapping[str, object], included: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType({key: _merged_entry(base, included, key) for key in (*base, *included)}) + + +def _without_include(config: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType({key: value for key, value in config.items() if key != INCLUDE_KEY}) + + +def include_entries(config: Mapping[str, object]) -> tuple[str, ...]: + if INCLUDE_KEY not in config: + return () + + entries: Final = config[INCLUDE_KEY] + if not isinstance(entries, list): + raise ValueError("'include' must be a list of file paths") + + paths: Final = tuple(entry for entry in entries if isinstance(entry, str)) + if len(paths) != len(entries): + raise ValueError("'include' must be a list of file paths") + + return paths + + +def _pending_from(config: Mapping[str, object], location: str) -> tuple[tuple[str, str], ...]: + return tuple((entry, location) for entry in include_entries(config)) + + +async def _resolve( + config: Mapping[str, object], + pending: tuple[tuple[str, str], ...], + loaded: frozenset[str], + resolve: IncludeResolver, + read: ConfigReader, +) -> Mapping[str, object]: + if not pending: + return _without_include(config) + + entry, declared_in = pending[0] + location: Final = resolve(entry, declared_in) + if location in loaded: + return await _resolve(config, pending[1:], loaded, resolve, read) + + included: Final = await read(location) + return await _resolve( + _merged(config, _without_include(included)), + (*pending[1:], *_pending_from(included, location)), + loaded | frozenset((location,)), + resolve, + read, + ) + + +async def resolve_includes( + config: Mapping[str, object], + location: str, + resolve: IncludeResolver, + read: ConfigReader, +) -> dict[str, object]: + """ + Merge every config named by the `include` directive into the config that declares it. + + List values are extended and every other value is overridden, `resolve` turns each entry into the + location it names relative to the config that declares it, a config already pulled in is neither + read nor merged a second time, and `read` decides where a location is read from, so the same merge + applies to configs on disk and to configs hosted in a bucket. + """ + merged: Final = await _resolve(config, _pending_from(config, location), frozenset((location,)), resolve, read) + return dict(merged) # mutable-ok: the proxy mutates the config it loads diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index 62649ad6ca1..4a082eb307b 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -1,12 +1,47 @@ +import asyncio import os -from typing import Final +import posixpath +from collections.abc import Awaitable, Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Protocol import yaml +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger +from litellm.proxy.common_utils.config_includes import resolve_includes + +if TYPE_CHECKING: + from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase + +_BUCKET_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, object]) -def get_file_contents_from_s3(bucket_name, object_key): +class BucketObjectFetcher(Protocol): + def __call__(self, object_key: str, /) -> Awaitable[Mapping[str, object] | None]: ... + + +class BucketObjectReader(Protocol): + def __call__(self, object_key: str, /) -> Awaitable[object | None]: ... + + +class SyncBucketObjectReader(Protocol): + def __call__(self, object_key: str, /) -> object | None: ... + + +def _parsed_config(object_key: str, file_contents: str) -> object | None: + try: + parsed: Final = yaml.safe_load(file_contents) + except yaml.YAMLError as e: + verbose_proxy_logger.error("Config object %s is not valid YAML: %s", object_key, e) + return None + return MappingProxyType({}) if parsed is None else parsed + + +def s3_object_reader(bucket_name: str) -> SyncBucketObjectReader: + """ + Build one reader for a whole config, so an `include` tree costs one S3 client rather than one per object. + """ try: # v0 rely on boto3 for authentication - allowing boto3 to handle IAM credentials etc import boto3 @@ -21,46 +56,147 @@ def get_file_contents_from_s3(bucket_name, object_key): aws_secret_access_key=credentials.secret_key, aws_session_token=credentials.token, # Optional, if using temporary credentials ) - verbose_proxy_logger.debug("Retrieving %s from S3 bucket: %s", object_key, bucket_name) - response: Final = s3_client.get_object(Bucket=bucket_name, Key=object_key) - verbose_proxy_logger.debug("Response: %s", response) - - # Read the file contents and directly parse YAML - file_contents: Final = response["Body"].read().decode("utf-8") - verbose_proxy_logger.debug("File contents retrieved from S3") - - # Parse YAML directly from string - config: Final = yaml.safe_load(file_contents) - return config - except ImportError as e: # this is most likely if a user is not using the litellm docker container verbose_proxy_logger.error("ImportError: %s", e) + return lambda object_key: None except Exception as e: - verbose_proxy_logger.error("Error retrieving file contents: %s", e) + verbose_proxy_logger.error("Error creating the S3 client for bucket %s: %s", bucket_name, e) + return lambda object_key: None + + def read(object_key: str) -> object | None: + try: + verbose_proxy_logger.debug("Retrieving %s from S3 bucket: %s", object_key, bucket_name) + response: Final = s3_client.get_object(Bucket=bucket_name, Key=object_key) + file_contents: Final = response["Body"].read().decode("utf-8") + except Exception as e: # noqa: BLE001 # any boto3 error must read as a missing object + verbose_proxy_logger.error("Error retrieving %s from S3 bucket %s: %s", object_key, bucket_name, e) + return None + + return _parsed_config(object_key, file_contents) + + return read + + +def get_file_contents_from_s3(bucket_name: str, object_key: str) -> object | None: + return s3_object_reader(bucket_name)(object_key) + + +def gcs_config_bucket(bucket_name: str) -> "GCSBucketBase | None": + """ + Build a plain GCS client for reading config objects. + + Reading a config out of a bucket is not GCS logging, so it neither needs the enterprise license + that gate covers nor the batching task the logger starts and never stops. + """ + try: + from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase + + return GCSBucketBase(bucket_name=bucket_name) + except Exception as e: # noqa: BLE001 # an unbuildable client must read as an unreadable bucket + verbose_proxy_logger.error("Error creating the GCS client for bucket %s: %s", bucket_name, e) return None -async def get_config_file_contents_from_gcs(bucket_name, object_key): +async def get_config_file_contents_from_gcs( + bucket_name: str, + object_key: str, + gcs_bucket: "GCSBucketBase | None" = None, +) -> object | None: try: - from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger - - gcs_bucket: Final = GCSBucketLogger( - bucket_name=bucket_name, - ) - file_contents = await gcs_bucket.download_gcs_object(object_key) + bucket: Final = gcs_config_bucket(bucket_name) if gcs_bucket is None else gcs_bucket + if bucket is None: + return None + file_contents: Final = await bucket.download_gcs_object(object_key) if file_contents is None: raise Exception(f"File contents are None for {object_key}") - # file_contentis is a bytes object, so we need to convert it to yaml - file_contents = file_contents.decode("utf-8") - # convert to yaml - config: Final = yaml.safe_load(file_contents) - return config + decoded: Final = file_contents.decode("utf-8") except Exception as e: - verbose_proxy_logger.error("Error retrieving file contents: %s", e) + verbose_proxy_logger.error("Error retrieving %s from GCS bucket %s: %s", object_key, bucket_name, e) return None + return _parsed_config(object_key, decoded) + + +def resolve_include_object_key(config_object_key: str, include_entry: str) -> str: + """ + Resolve one `include` entry to the object key it names, relative to the config object's prefix. + + A leading "/" means the bucket root, mirroring how an absolute path on disk ignores the + directory the including config sits in. + """ + if include_entry.startswith("/"): + return posixpath.normpath(include_entry).lstrip("/") + return posixpath.normpath(posixpath.join(posixpath.dirname(config_object_key), include_entry)) + + +async def resolve_bucket_includes( + *, + config: Mapping[str, object], + object_key: str, + fetch: BucketObjectFetcher, +) -> dict[str, object]: + async def read(include_key: str) -> Mapping[str, object]: + included: Final = await fetch(include_key) + if included is None: + raise FileNotFoundError( + f"Included config could not be read from bucket: {include_key}. " + "The underlying bucket error is logged above." + ) + return included + + def resolve(include_entry: str, declared_in: str) -> str: + return resolve_include_object_key(declared_in, include_entry) + + return await resolve_includes(config=config, location=object_key, resolve=resolve, read=read) + + +async def bucket_object_reader(bucket_type: str | None, bucket_name: str) -> BucketObjectReader: + """ + Build one reader for a whole config, so an `include` tree costs one bucket client rather than one per object. + """ + if bucket_type != "gcs": + read_object: Final = await asyncio.to_thread(s3_object_reader, bucket_name) + + async def read_from_s3(object_key: str) -> object | None: + return await asyncio.to_thread(read_object, object_key) + + return read_from_s3 + + gcs_bucket: Final = gcs_config_bucket(bucket_name) + + async def read_from_gcs(object_key: str) -> object | None: + if gcs_bucket is None: + return None + return await get_config_file_contents_from_gcs(bucket_name, object_key, gcs_bucket) + + return read_from_gcs + + +async def get_config_from_bucket( + *, + bucket_type: str | None, + bucket_name: str, + object_key: str, +) -> dict[str, object] | None: + read: Final = await bucket_object_reader(bucket_type, bucket_name) + + async def fetch(key: str) -> Mapping[str, object] | None: + raw: Final = await read(key) + if raw is None: + return None + try: + return _BUCKET_CONFIG_ADAPTER.validate_python(raw) + except ValidationError as e: + raise ValueError(f"Config object in bucket is not a YAML mapping: {key}") from e + + config: Final = await fetch(object_key) + if not config: + return None + + return await resolve_bucket_includes(config=config, object_key=object_key, fetch=fetch) + def download_python_file_from_s3( bucket_name: str, @@ -136,11 +272,9 @@ async def download_python_file_from_gcs( bool: True if successful, False otherwise """ try: - from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger + from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase - gcs_bucket: Final = GCSBucketLogger( - bucket_name=bucket_name, - ) + gcs_bucket: Final = GCSBucketBase(bucket_name=bucket_name) file_contents = await gcs_bucket.download_gcs_object(object_key) if file_contents is None: raise Exception(f"File contents are None for {object_key}") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b36a2f71691..d9f8e04ebda 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -353,6 +353,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AuthCacheInvalidationSubscriber, ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy +from litellm.proxy.common_utils.config_includes import resolve_include_file_path, resolve_includes from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber from litellm.proxy.common_utils.debug_utils import init_verbose_loggers from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router @@ -371,10 +372,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( check_file_size_under_limit, get_form_data, ) -from litellm.proxy.common_utils.load_config_utils import ( - get_config_file_contents_from_gcs, - get_file_contents_from_s3, -) +from litellm.proxy.common_utils.load_config_utils import get_config_from_bucket from litellm.proxy.common_utils.model_deprecation import collect_model_deprecations from litellm.proxy.common_utils.model_listing_utils import ( ClaudeCodeRoutingNames, @@ -4806,12 +4804,12 @@ class ProxyConfig: if config is None: raise Exception("Config cannot be None or Empty.") # Process includes - config = self._process_includes(config=config, base_dir=os.path.dirname(os.path.abspath(file_path or ""))) + config = await self._process_includes(config=config, config_file_path=os.path.abspath(file_path or "")) # verbose_proxy_logger.debug(f"loaded config={json.dumps(config, indent=4)}") return config - def _process_includes(self, config: dict, base_dir: str) -> dict: + async def _process_includes(self, config: dict, config_file_path: str) -> dict: """ Process includes by appending their contents to the main config @@ -4826,29 +4824,21 @@ class ProxyConfig: callbacks: ["prometheus"] ``` """ - if "include" not in config: - return config - if not isinstance(config["include"], list): - raise ValueError("'include' must be a list of file paths") + included_config_adapter: Final = TypeAdapter(dict[str, object]) - # Load and append all included files - for include_file in config["include"]: - file_path = os.path.join(base_dir, include_file) + def resolve(include_file: str, declared_in: str) -> str: + return resolve_include_file_path(include_file, declared_in, config_file_path) + + async def read_included(file_path: str) -> Mapping[str, object]: if not os.path.exists(file_path): raise FileNotFoundError(f"Included file not found: {file_path}") + try: + return included_config_adapter.validate_python(self._load_yaml_file(file_path)) + except ValidationError as e: + raise ValueError(f"Included config file is not a YAML mapping: {file_path}") from e - included_config = self._load_yaml_file(file_path) - # Simply update/extend the main config with included config - for key, value in included_config.items(): - if isinstance(value, list) and key in config: - config[key].extend(value) - else: - config[key] = value - - # Remove the include directive - del config["include"] - return config + return await resolve_includes(config=config, location=config_file_path, resolve=resolve, read=read_included) async def save_config(self, new_config: dict, include_env_vars: bool = False): global prisma_client, general_settings, user_config_file_path, store_model_in_db @@ -5206,15 +5196,19 @@ class ProxyConfig: global prisma_client, store_model_in_db # Load existing config - if os.environ.get("LITELLM_CONFIG_BUCKET_NAME") is not None: - bucket_name: Final = os.environ.get("LITELLM_CONFIG_BUCKET_NAME") + bucket_name: Final = os.environ.get("LITELLM_CONFIG_BUCKET_NAME") + if bucket_name is not None: object_key: Final = os.environ.get("LITELLM_CONFIG_BUCKET_OBJECT_KEY") bucket_type: Final = os.environ.get("LITELLM_CONFIG_BUCKET_TYPE") verbose_proxy_logger.debug("bucket_name: %s, object_key: %s", bucket_name, object_key) - if bucket_type == "gcs": - config = await get_config_file_contents_from_gcs(bucket_name=bucket_name, object_key=object_key) - else: - config = get_file_contents_from_s3(bucket_name=bucket_name, object_key=object_key) + if object_key is None: + raise Exception("LITELLM_CONFIG_BUCKET_OBJECT_KEY must be set to load the config from a bucket.") + + config = await get_config_from_bucket( + bucket_type=bucket_type, + bucket_name=bucket_name, + object_key=object_key, + ) if config is None: raise Exception("Unable to load config from given source.") diff --git a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py index 8d662311da1..a458752bed0 100644 --- a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py +++ b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py @@ -128,3 +128,20 @@ class TestGCSBucketBase: assert object_name.endswith("-target_uploadType_media") assert ".." not in object_name assert "?" not in object_name + + +class TestGCSBucketLoggerBucketName: + @pytest.mark.asyncio + async def test_the_bucket_name_it_is_constructed_with_survives(self, monkeypatch): + """Reading config.yaml out of a GCS bucket asks for that bucket, not the logging one (LIT-6982).""" + monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket") + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + + assert GCSBucketLogger(bucket_name="config-bucket").BUCKET_NAME == "config-bucket" + + @pytest.mark.asyncio + async def test_no_bucket_name_still_falls_back_to_the_environment(self, monkeypatch): + monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket") + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + + assert GCSBucketLogger().BUCKET_NAME == "logging-bucket" diff --git a/tests/test_litellm/proxy/common_utils/test_load_config_utils.py b/tests/test_litellm/proxy/common_utils/test_load_config_utils.py index 524c260e94a..1042dbe9653 100644 --- a/tests/test_litellm/proxy/common_utils/test_load_config_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_load_config_utils.py @@ -1,9 +1,18 @@ +import asyncio +import logging +import re +import threading from unittest.mock import MagicMock, mock_open, patch import pytest import yaml -from litellm.proxy.common_utils.load_config_utils import get_file_contents_from_s3 +from litellm.proxy.common_utils.load_config_utils import ( + gcs_config_bucket, + get_config_from_bucket, + get_file_contents_from_s3, + resolve_bucket_includes, +) class TestGetFileContentsFromS3: @@ -83,3 +92,385 @@ class TestGetFileContentsFromS3: # Verify yaml.safe_load was called with the decoded content mock_yaml_load.assert_called_once_with(yaml_content) + + +class TestBucketConfigIncludes: + + @staticmethod + def _bucket(objects): + async def fetch(object_key): + return objects.get(object_key) + + return fetch + + @pytest.mark.asyncio + async def test_include_resolves_against_the_config_objects_prefix(self): + merged = await resolve_bucket_includes( + config={"include": ["model_config.yaml"], "general_settings": {"master_key": "sk-1234"}}, + object_key="configs/prod/config.yaml", + fetch=self._bucket( + {"configs/prod/model_config.yaml": {"model_list": [{"model_name": "gpt-4o-mini"}]}} + ), + ) + + assert merged == { + "general_settings": {"master_key": "sk-1234"}, + "model_list": [{"model_name": "gpt-4o-mini"}], + } + + @pytest.mark.asyncio + async def test_include_with_a_leading_slash_reads_from_the_bucket_root(self): + merged = await resolve_bucket_includes( + config={"include": ["/shared/models.yaml"]}, + object_key="configs/prod/config.yaml", + fetch=self._bucket({"shared/models.yaml": {"model_list": [{"model_name": "shared"}]}}), + ) + + assert merged == {"model_list": [{"model_name": "shared"}]} + + @pytest.mark.asyncio + async def test_include_walks_out_of_the_prefix_with_dot_dot(self): + merged = await resolve_bucket_includes( + config={"include": ["../shared/models.yaml"]}, + object_key="configs/prod/config.yaml", + fetch=self._bucket({"configs/shared/models.yaml": {"model_list": [{"model_name": "shared"}]}}), + ) + + assert merged == {"model_list": [{"model_name": "shared"}]} + + @pytest.mark.asyncio + async def test_included_configs_may_declare_further_includes(self): + merged = await resolve_bucket_includes( + config={"include": ["models.yaml"]}, + object_key="configs/config.yaml", + fetch=self._bucket( + { + "configs/models.yaml": { + "include": ["extra/more_models.yaml"], + "model_list": [{"model_name": "first"}], + }, + "configs/extra/more_models.yaml": {"model_list": [{"model_name": "second"}]}, + } + ), + ) + + assert merged == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} + + @pytest.mark.asyncio + async def test_a_nested_include_resolves_against_the_object_that_declares_it(self): + merged = await resolve_bucket_includes( + config={"include": ["shared/models.yaml"]}, + object_key="configs/config.yaml", + fetch=self._bucket( + { + "configs/shared/models.yaml": { + "include": ["more_models.yaml"], + "model_list": [{"model_name": "first"}], + }, + "configs/shared/more_models.yaml": {"model_list": [{"model_name": "second"}]}, + "configs/more_models.yaml": {"model_list": [{"model_name": "wrong-prefix"}]}, + } + ), + ) + + assert merged == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} + + @pytest.mark.asyncio + async def test_an_object_pulled_in_twice_is_merged_once(self): + merged = await resolve_bucket_includes( + config={"include": ["a.yaml", "b.yaml"]}, + object_key="configs/config.yaml", + fetch=self._bucket( + { + "configs/a.yaml": {"include": ["shared.yaml"]}, + "configs/b.yaml": {"include": ["./shared.yaml"]}, + "configs/shared.yaml": {"model_list": [{"model_name": "shared"}]}, + } + ), + ) + + assert merged == {"model_list": [{"model_name": "shared"}]} + + @pytest.mark.asyncio + async def test_a_cycle_between_included_objects_terminates(self): + merged = await asyncio.wait_for( + resolve_bucket_includes( + config={"include": ["a.yaml"]}, + object_key="configs/config.yaml", + fetch=self._bucket( + { + "configs/a.yaml": {"include": ["b.yaml"], "model_list": [{"model_name": "from-a"}]}, + "configs/b.yaml": {"include": ["a.yaml"], "model_list": [{"model_name": "from-b"}]}, + } + ), + ), + timeout=10, + ) + + assert merged == {"model_list": [{"model_name": "from-a"}, {"model_name": "from-b"}]} + + @pytest.mark.asyncio + async def test_list_values_are_extended_and_other_values_are_overridden(self): + merged = await resolve_bucket_includes( + config={ + "include": ["models.yaml"], + "model_list": [{"model_name": "from-root"}], + "litellm_settings": {"drop_params": True}, + }, + object_key="config.yaml", + fetch=self._bucket( + { + "models.yaml": { + "model_list": [{"model_name": "from-include"}], + "litellm_settings": {"num_retries": 3}, + } + } + ), + ) + + assert merged == { + "model_list": [{"model_name": "from-root"}, {"model_name": "from-include"}], + "litellm_settings": {"num_retries": 3}, + } + + @pytest.mark.asyncio + async def test_a_missing_included_object_fails_loudly_with_its_key(self): + with pytest.raises(FileNotFoundError, match=re.escape("configs/prod/model_config.yaml")): + await resolve_bucket_includes( + config={"include": ["model_config.yaml"]}, + object_key="configs/prod/config.yaml", + fetch=self._bucket({}), + ) + + @pytest.mark.asyncio + async def test_a_non_list_include_fails_loudly(self): + with pytest.raises(ValueError, match="'include' must be a list of file paths"): + await resolve_bucket_includes( + config={"include": "model_config.yaml"}, + object_key="config.yaml", + fetch=self._bucket({}), + ) + + @pytest.mark.asyncio + async def test_get_config_from_bucket_merges_includes_over_s3(self, monkeypatch): + objects = { + "lit6982/config.yaml": { + "include": ["model_config.yaml"], + "general_settings": {"master_key": "sk-1234"}, + }, + "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, + } + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.s3_object_reader", + lambda bucket_name: objects.get, + ) + + config = await get_config_from_bucket( + bucket_type="s3", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert config == { + "general_settings": {"master_key": "sk-1234"}, + "model_list": [{"model_name": "included-model"}], + } + + @pytest.mark.asyncio + async def test_the_blocking_s3_work_runs_off_the_event_loop_thread(self, monkeypatch): + loop_thread = threading.current_thread() + threads = [] + + def build_reader(bucket_name): + threads.append(threading.current_thread()) + + def read(object_key): + threads.append(threading.current_thread()) + return {"model_list": [{"model_name": "a-model"}]} + + return read + + monkeypatch.setattr("litellm.proxy.common_utils.load_config_utils.s3_object_reader", build_reader) + + await get_config_from_bucket(bucket_type="s3", bucket_name="litellm-configs", object_key="config.yaml") + + assert len(threads) == 2 and loop_thread not in threads + + @pytest.mark.asyncio + async def test_one_s3_client_serves_the_whole_include_tree(self, monkeypatch): + objects = { + "lit6982/config.yaml": {"include": ["model_config.yaml"]}, + "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, + } + readers = [] + + def build_reader(bucket_name): + requested = [] + readers.append(requested) + + def read(object_key): + requested.append(object_key) + return objects.get(object_key) + + return read + + monkeypatch.setattr("litellm.proxy.common_utils.load_config_utils.s3_object_reader", build_reader) + + await get_config_from_bucket( + bucket_type="s3", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert readers == [["lit6982/config.yaml", "lit6982/model_config.yaml"]] + + @pytest.mark.asyncio + async def test_an_empty_included_object_merges_as_an_empty_config(self, monkeypatch): + objects = { + "lit6982/config.yaml": "include:\n - empty.yaml\nmodel_list:\n - model_name: only-model\n", + "lit6982/empty.yaml": "", + } + + class FakeGCSBucket: + async def download_gcs_object(self, object_key): + return objects[object_key].encode("utf-8") + + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.gcs_config_bucket", + lambda bucket_name: FakeGCSBucket(), + ) + + config = await get_config_from_bucket( + bucket_type="gcs", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert config == {"model_list": [{"model_name": "only-model"}]} + + @pytest.mark.asyncio + async def test_get_config_from_bucket_merges_includes_over_gcs(self, monkeypatch): + objects = { + "lit6982/config.yaml": { + "include": ["model_config.yaml"], + "general_settings": {"master_key": "sk-1234"}, + }, + "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, + } + + buckets = [] + + class FakeGCSBucket: + def __init__(self): + self.requested = [] + buckets.append(self) + + async def download_gcs_object(self, object_key): + self.requested.append(object_key) + return yaml.dump(objects[object_key]).encode("utf-8") + + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.gcs_config_bucket", + lambda bucket_name: FakeGCSBucket(), + ) + + config = await get_config_from_bucket( + bucket_type="gcs", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert config == { + "general_settings": {"master_key": "sk-1234"}, + "model_list": [{"model_name": "included-model"}], + } + assert [bucket.requested for bucket in buckets] == [ + ["lit6982/config.yaml", "lit6982/model_config.yaml"] + ] + + @pytest.mark.asyncio + async def test_get_config_from_bucket_returns_none_when_the_root_object_is_missing(self, monkeypatch): + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.s3_object_reader", + lambda bucket_name: (lambda object_key: None), + ) + + assert ( + await get_config_from_bucket( + bucket_type="s3", bucket_name="litellm-configs", object_key="missing.yaml" + ) + is None + ) + + @pytest.mark.asyncio + async def test_an_object_pulled_in_twice_is_read_once(self): + objects = { + "configs/a.yaml": {"include": ["shared.yaml"]}, + "configs/b.yaml": {"include": ["./shared.yaml"]}, + "configs/shared.yaml": {"model_list": [{"model_name": "shared"}]}, + } + requested = [] + + async def fetch(object_key): + requested.append(object_key) + return objects.get(object_key) + + await resolve_bucket_includes( + config={"include": ["a.yaml", "b.yaml"]}, + object_key="configs/config.yaml", + fetch=fetch, + ) + + assert requested == ["configs/a.yaml", "configs/b.yaml", "configs/shared.yaml"] + + @pytest.mark.asyncio + async def test_an_empty_root_object_does_not_boot_an_empty_proxy(self, monkeypatch): + class FakeGCSBucket: + async def download_gcs_object(self, object_key): + return b"" + + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.gcs_config_bucket", + lambda bucket_name: FakeGCSBucket(), + ) + + config = await get_config_from_bucket( + bucket_type="gcs", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert config is None + + @pytest.mark.asyncio + async def test_an_object_that_is_not_valid_yaml_is_reported_as_a_yaml_error(self, monkeypatch, caplog): + class FakeGCSBucket: + async def download_gcs_object(self, object_key): + return b"model_list: [\n" + + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.gcs_config_bucket", + lambda bucket_name: FakeGCSBucket(), + ) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + config = await get_config_from_bucket( + bucket_type="gcs", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert config is None + assert [ + record + for record in caplog.records + if "not valid YAML" in record.getMessage() and "lit6982/config.yaml" in record.getMessage() + ] + + +class TestGCSConfigBucketClient: + @pytest.mark.asyncio + async def test_reading_a_config_from_gcs_does_not_need_an_enterprise_license(self, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) + + bucket = gcs_config_bucket("litellm-configs") + + assert bucket is not None + assert bucket.BUCKET_NAME == "litellm-configs" + + @pytest.mark.asyncio + async def test_reading_a_config_from_gcs_starts_no_background_task(self, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + running_before = asyncio.all_tasks() + + gcs_config_bucket("litellm-configs") + + assert asyncio.all_tasks() - running_before == set() diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index e79448d0620..e109b650da7 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -8,6 +8,7 @@ Pins covered: from __future__ import annotations +import asyncio import json import logging import os @@ -712,22 +713,124 @@ async def test_ProxyConfig__get_config_from_file_missing_path_raises(): # --------------------------------------------------------------------------- -def test_ProxyConfig__process_includes_merges_files(tmp_path): +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_merges_files(tmp_path): inc = tmp_path / "models.yaml" inc.write_text("model_list:\n - model_name: gpt-4\n") pc = ProxyConfig() cfg = {"include": ["models.yaml"], "model_list": [], "litellm_settings": {}} - result = pc._process_includes(cfg, base_dir=str(tmp_path)) + result = await pc._process_includes(cfg, config_file_path=str(tmp_path / "config.yaml")) assert result == { "model_list": [{"model_name": "gpt-4"}], "litellm_settings": {}, } -def test_ProxyConfig__process_includes_missing_file_raises(tmp_path): +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_missing_file_raises(tmp_path): pc = ProxyConfig() with pytest.raises(FileNotFoundError): - pc._process_includes({"include": ["nope.yaml"]}, base_dir=str(tmp_path)) + await pc._process_includes({"include": ["nope.yaml"]}, config_file_path=str(tmp_path / "config.yaml")) + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_follows_nested_includes(tmp_path): + (tmp_path / "models.yaml").write_text("include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n") + (tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: second\n") + result = await ProxyConfig()._process_includes( + {"include": ["models.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) + assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_resolves_a_nested_include_next_to_its_own_file(tmp_path): + (tmp_path / "shared").mkdir() + (tmp_path / "shared" / "models.yaml").write_text( + "include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n" + ) + (tmp_path / "shared" / "more_models.yaml").write_text("model_list:\n - model_name: second\n") + (tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: wrong-directory\n") + + result = await ProxyConfig()._process_includes( + {"include": ["shared/models.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) + + assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_still_reads_a_nested_include_left_beside_the_root_config(tmp_path): + (tmp_path / "shared").mkdir() + (tmp_path / "shared" / "models.yaml").write_text( + "include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n" + ) + (tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: second\n") + + result = await ProxyConfig()._process_includes( + {"include": ["shared/models.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) + + assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_names_both_files_when_a_nested_include_matches_two(tmp_path, caplog): + (tmp_path / "shared").mkdir() + (tmp_path / "shared" / "models.yaml").write_text( + "include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n" + ) + (tmp_path / "shared" / "more_models.yaml").write_text("model_list:\n - model_name: next-to-the-declaring-file\n") + (tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: next-to-the-root-config\n") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await ProxyConfig()._process_includes( + {"include": ["shared/models.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) + + assert result == {"model_list": [{"model_name": "first"}, {"model_name": "next-to-the-declaring-file"}]} + assert [ + record + for record in caplog.records + if str(tmp_path / "shared" / "more_models.yaml") in record.getMessage() + and str(tmp_path / "more_models.yaml") in record.getMessage() + ] + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_merges_a_shared_file_once(tmp_path): + (tmp_path / "shared.yaml").write_text("model_list:\n - model_name: shared\n") + (tmp_path / "a.yaml").write_text("include:\n - shared.yaml\n") + (tmp_path / "b.yaml").write_text("include:\n - ./shared.yaml\n") + + result = await ProxyConfig()._process_includes( + {"include": ["a.yaml", "b.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) + + assert result == {"model_list": [{"model_name": "shared"}]} + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_names_the_file_when_it_is_not_a_mapping(tmp_path): + (tmp_path / "models.yaml").write_text("- model_name: gpt-4\n") + + with pytest.raises(ValueError, match=re.escape(str(tmp_path / "models.yaml"))): + await ProxyConfig()._process_includes( + {"include": ["models.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_terminates_on_a_cycle(tmp_path): + (tmp_path / "a.yaml").write_text("include:\n - b.yaml\nmodel_list:\n - model_name: from-a\n") + (tmp_path / "b.yaml").write_text("include:\n - a.yaml\nmodel_list:\n - model_name: from-b\n") + + result = await asyncio.wait_for( + ProxyConfig()._process_includes({"include": ["a.yaml"]}, config_file_path=str(tmp_path / "config.yaml")), + timeout=10, + ) + + assert result == {"model_list": [{"model_name": "from-a"}, {"model_name": "from-b"}]} # --------------------------------------------------------------------------- @@ -1044,6 +1147,31 @@ async def test_ProxyConfig_get_config_loads_from_file(tmp_path, monkeypatch): } +@pytest.mark.asyncio +async def test_ProxyConfig_get_config_from_a_bucket_merges_includes(monkeypatch): + objects = { + "lit6982/config.yaml": { + "include": ["model_config.yaml"], + "general_settings": {"master_key": "sk-1234"}, + }, + "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, + } + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.s3_object_reader", + lambda bucket_name: objects.get, + ) + monkeypatch.setenv("LITELLM_CONFIG_BUCKET_NAME", "litellm-configs") + monkeypatch.setenv("LITELLM_CONFIG_BUCKET_OBJECT_KEY", "lit6982/config.yaml") + monkeypatch.setenv("LITELLM_CONFIG_BUCKET_TYPE", "s3") + + cfg = await ProxyConfig().get_config() + + assert cfg["model_list"] == [{"model_name": "included-model"}] + assert "include" not in cfg + + @pytest.mark.asyncio async def test_ProxyConfig_get_config_missing_file_raises(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)