diff --git a/litellm/proxy/common_utils/config_includes.py b/litellm/proxy/common_utils/config_includes.py index cf39bd6435e..23e6c615a09 100644 --- a/litellm/proxy/common_utils/config_includes.py +++ b/litellm/proxy/common_utils/config_includes.py @@ -1,10 +1,38 @@ +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. + """ + declared_relative: Final = os.path.abspath(os.path.join(os.path.dirname(declared_in), include_file)) + if os.path.exists(declared_relative): + return declared_relative + + 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 + + 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 + + class ConfigLoader(Protocol): def __call__(self, include_entry: str, declared_in: str, /) -> Awaitable[tuple[str, Mapping[str, object]]]: ... diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index e72272b1aca..d62286a50f6 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -2,6 +2,7 @@ import asyncio import os import posixpath from collections.abc import Awaitable, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol import yaml @@ -24,7 +25,19 @@ class BucketObjectReader(Protocol): def __call__(self, object_key: str, /) -> Awaitable[object | None]: ... -def get_file_contents_from_s3(bucket_name, object_key): +class SyncBucketObjectReader(Protocol): + def __call__(self, object_key: str, /) -> object | None: ... + + +def _parsed_config(file_contents: str) -> object: + parsed: Final = yaml.safe_load(file_contents) + 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 @@ -39,24 +52,28 @@ 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) - return None + 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) + return _parsed_config(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 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": @@ -64,27 +81,27 @@ def gcs_config_bucket(bucket_name: str) -> "GCSBucketBase | None": from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger return GCSBucketLogger(bucket_name=bucket_name) - except Exception as e: + 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, gcs_bucket=None): +async def get_config_file_contents_from_gcs( + bucket_name: str, + object_key: str, + gcs_bucket: "GCSBucketBase | None" = None, +) -> object | None: try: bucket: Final = gcs_config_bucket(bucket_name) if gcs_bucket is None else gcs_bucket if bucket is None: return None - file_contents = await bucket.download_gcs_object(object_key) + 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 + return _parsed_config(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 @@ -110,20 +127,24 @@ async def resolve_bucket_includes( include_key: Final = resolve_include_object_key(declared_in, include_entry) included: Final = await fetch(include_key) if included is None: - raise FileNotFoundError(f"Included config could not be read from bucket: {include_key}") + raise FileNotFoundError( + f"Included config could not be read from bucket: {include_key}. " + "The underlying bucket error is logged above." + ) return include_key, included return await resolve_includes(config=config, location=object_key, load=load) -def bucket_object_reader(bucket_type: str | None, bucket_name: str) -> BucketObjectReader: +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(get_file_contents_from_s3, bucket_name, object_key) + return await asyncio.to_thread(read_object, object_key) return read_from_s3 @@ -143,7 +164,7 @@ async def get_config_from_bucket( bucket_name: str, object_key: str, ) -> dict[str, object] | None: - read: Final = bucket_object_reader(bucket_type, bucket_name) + read: Final = await bucket_object_reader(bucket_type, bucket_name) async def fetch(key: str) -> Mapping[str, object] | None: raw: Final = await read(key) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7ac5e697b61..7fada09febf 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -341,7 +341,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_includes +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 @@ -4560,7 +4560,7 @@ class ProxyConfig: included_config_adapter: Final = TypeAdapter(dict[str, object]) async def load_included(include_file: str, declared_in: str) -> tuple[str, Mapping[str, object]]: - file_path: Final = os.path.abspath(os.path.join(os.path.dirname(declared_in), include_file)) + file_path: Final = resolve_include_file_path(include_file, declared_in, config_file_path) if not os.path.exists(file_path): raise FileNotFoundError(f"Included file not found: {file_path}") try: 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 b654320569c..75aa5a3f0f8 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 @@ -93,11 +93,6 @@ class TestGetFileContentsFromS3: class TestBucketConfigIncludes: - """`include:` directives in a bucket-hosted config.yaml (LIT-6982). - - They used to be dropped silently: the proxy booted with the root config applied and everything - the included objects declared missing, with nothing logged. - """ @staticmethod def _bucket(objects): @@ -161,7 +156,6 @@ class TestBucketConfigIncludes: @pytest.mark.asyncio async def test_a_nested_include_resolves_against_the_object_that_declares_it(self): - """A nested `include` names a neighbour of the object declaring it, not of the root config.""" merged = await resolve_bucket_includes( config={"include": ["shared/models.yaml"]}, object_key="configs/config.yaml", @@ -265,8 +259,8 @@ class TestBucketConfigIncludes: "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, } monkeypatch.setattr( - "litellm.proxy.common_utils.load_config_utils.get_file_contents_from_s3", - lambda bucket_name, object_key: objects.get(object_key), + "litellm.proxy.common_utils.load_config_utils.s3_object_reader", + lambda bucket_name: objects.get, ) config = await get_config_from_bucket( @@ -279,21 +273,72 @@ class TestBucketConfigIncludes: } @pytest.mark.asyncio - async def test_the_blocking_s3_read_runs_off_the_event_loop_thread(self, monkeypatch): + async def test_the_blocking_s3_work_runs_off_the_event_loop_thread(self, monkeypatch): loop_thread = threading.current_thread() - read_threads = [] + threads = [] - def record_thread(bucket_name, object_key): - read_threads.append(threading.current_thread()) - return {"model_list": [{"model_name": "a-model"}]} + def build_reader(bucket_name): + threads.append(threading.current_thread()) - monkeypatch.setattr( - "litellm.proxy.common_utils.load_config_utils.get_file_contents_from_s3", record_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 read_threads and loop_thread not in read_threads + 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): @@ -336,8 +381,8 @@ class TestBucketConfigIncludes: @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.get_file_contents_from_s3", - lambda bucket_name, object_key: None, + "litellm.proxy.common_utils.load_config_utils.s3_object_reader", + lambda bucket_name: (lambda object_key: None), ) assert ( 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 3d9138ad098..e90f101a555 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -743,7 +743,6 @@ async def test_ProxyConfig__process_includes_follows_nested_includes(tmp_path): @pytest.mark.asyncio async def test_ProxyConfig__process_includes_resolves_a_nested_include_next_to_its_own_file(tmp_path): - """A nested `include` names a sibling of the file that declares it, not of the root config.""" (tmp_path / "shared").mkdir() (tmp_path / "shared" / "models.yaml").write_text( "include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n" @@ -758,6 +757,21 @@ async def test_ProxyConfig__process_includes_resolves_a_nested_include_next_to_i 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_merges_a_shared_file_once(tmp_path): (tmp_path / "shared.yaml").write_text("model_list:\n - model_name: shared\n") @@ -1110,7 +1124,6 @@ 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): - """A bucket-hosted config.yaml used to drop its `include:` entries silently (LIT-6982).""" objects = { "lit6982/config.yaml": { "include": ["model_config.yaml"], @@ -1121,8 +1134,8 @@ async def test_ProxyConfig_get_config_from_a_bucket_merges_includes(monkeypatch) 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.get_file_contents_from_s3", - lambda bucket_name, object_key: objects.get(object_key), + "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")