fix(proxy): read a bucket config's include tree off the event loop

Reading a config from a bucket ran a blocking boto3 GET straight from the
event loop for every object in the include tree, and on GCS it built a new
bucket client per object, each one starting a flush task that never ends.

S3 reads now go through a worker thread, and one bucket client serves the
whole include tree.

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,

litellm/proxy/common_utils/config_includes.py

@@ -6,7 +6,7 @@ INCLUDE_KEY: Final = "include"
 
 
 class ConfigLoader(Protocol):
-    def __call__(self, include_entry: str, /) -> Awaitable[Mapping[str, object]]: ...
+    def __call__(self, include_entry: str, declared_in: str, /) -> Awaitable[tuple[str, Mapping[str, object]]]: ...
 
 
 def _merged_value(base_value: object, included_value: object) -> object:
@@ -44,25 +44,40 @@ def include_entries(config: Mapping[str, object]) -> tuple[str, ...]:
     return paths
 
 
-async def _resolve(config: Mapping[str, object], pending: tuple[str, ...], load: ConfigLoader) -> Mapping[str, object]:
+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],
+    load: ConfigLoader,
+) -> Mapping[str, object]:
     if not pending:
         return _without_include(config)
 
-    included: Final = await load(pending[0])
+    entry, declared_in = pending[0]
+    location, included = await load(entry, declared_in)
+    if location in loaded:
+        return await _resolve(config, pending[1:], loaded, load)
+
     return await _resolve(
         _merged(config, _without_include(included)),
-        (*pending[1:], *include_entries(included)),
+        (*pending[1:], *_pending_from(included, location)),
+        loaded | frozenset((location,)),
         load,
     )
 
 
-async def resolve_includes(config: Mapping[str, object], load: ConfigLoader) -> dict[str, object]:
+async def resolve_includes(config: Mapping[str, object], location: str, load: ConfigLoader) -> 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, an included config may declare
-    further includes, and `load` decides where an entry is read from, so the same merge applies to
-    configs on disk and to configs hosted in a bucket.
+    List values are extended and every other value is overridden, each entry is resolved relative to
+    the config that declares it, a config already pulled in is not merged a second time, and `load`
+    decides where an entry is read from, so the same merge applies to configs on disk and to configs
+    hosted in a bucket.
     """
-    merged: Final = await _resolve(config, include_entries(config), load)
+    merged: Final = await _resolve(config, _pending_from(config, location), frozenset((location,)), load)
     return dict(merged)  # mutable-ok: the proxy mutates the config it loads

litellm/proxy/common_utils/load_config_utils.py

@@ -1,7 +1,8 @@
+import asyncio
 import os
 import posixpath
 from collections.abc import Awaitable, Mapping
-from typing import Final, Protocol
+from typing import TYPE_CHECKING, Final, Protocol
 
 import yaml
 from pydantic import TypeAdapter, ValidationError
@@ -9,6 +10,9 @@ 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])
 
 
@@ -16,6 +20,10 @@ 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]: ...
+
+
 def get_file_contents_from_s3(bucket_name, object_key):
     try:
         # v0 rely on boto3 for authentication - allowing boto3 to handle IAM credentials etc
@@ -51,14 +59,22 @@ def get_file_contents_from_s3(bucket_name, object_key):
         return None
 
 
-async def get_config_file_contents_from_gcs(bucket_name, object_key):
+def gcs_config_bucket(bucket_name: str) -> "GCSBucketBase | 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)
+        return GCSBucketLogger(bucket_name=bucket_name)
+    except Exception as e:
+        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):
+    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)
         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
@@ -90,14 +106,35 @@ async def resolve_bucket_includes(
     object_key: str,
     fetch: BucketObjectFetcher,
 ) -> dict[str, object]:
-    async def load(include_entry: str) -> Mapping[str, object]:
-        include_key: Final = resolve_include_object_key(object_key, include_entry)
+    async def load(include_entry: str, declared_in: str) -> tuple[str, Mapping[str, object]]:
+        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}")
-        return included
+        return include_key, included
 
-    return await resolve_includes(config=config, load=load)
+    return await resolve_includes(config=config, location=object_key, load=load)
+
+
+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":
+
+        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 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(
@@ -106,12 +143,10 @@ 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)
+
     async def fetch(key: str) -> Mapping[str, object] | None:
-        raw: Final = (
-            await get_config_file_contents_from_gcs(bucket_name=bucket_name, object_key=key)
-            if bucket_type == "gcs"
-            else get_file_contents_from_s3(bucket_name=bucket_name, object_key=key)
-        )
+        raw: Final = await read(key)
         if raw is None:
             return None
         try:

litellm/proxy/proxy_server.py

@@ -4536,12 +4536,12 @@ class ProxyConfig:
         if config is None:
             raise Exception("Config cannot be None or Empty.")
         # Process includes
-        config = await 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
 
-    async 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
 
@@ -4557,13 +4557,19 @@ class ProxyConfig:
         ```
         """
 
-        async def load_included(include_file: str) -> Mapping[str, object]:
-            file_path: Final = os.path.join(base_dir, include_file)
+        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))
             if not os.path.exists(file_path):
                 raise FileNotFoundError(f"Included file not found: {file_path}")
-            return self._load_yaml_file(file_path)
+            try:
+                included: Final = 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
+            return file_path, included
 
-        return await resolve_includes(config=config, load=load_included)
+        return await resolve_includes(config=config, location=config_file_path, load=load_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

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"

tests/test_litellm/proxy/common_utils/test_load_config_utils.py

@@ -1,4 +1,6 @@
+import asyncio
 import re
+import threading
 from unittest.mock import MagicMock, mock_open, patch
 
 import pytest
@@ -157,6 +159,60 @@ class TestBucketConfigIncludes:
 
         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):
+        """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",
+            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(
@@ -222,6 +278,23 @@ class TestBucketConfigIncludes:
             "model_list": [{"model_name": "included-model"}],
         }
 
+    @pytest.mark.asyncio
+    async def test_the_blocking_s3_read_runs_off_the_event_loop_thread(self, monkeypatch):
+        loop_thread = threading.current_thread()
+        read_threads = []
+
+        def record_thread(bucket_name, object_key):
+            read_threads.append(threading.current_thread())
+            return {"model_list": [{"model_name": "a-model"}]}
+
+        monkeypatch.setattr(
+            "litellm.proxy.common_utils.load_config_utils.get_file_contents_from_s3", record_thread
+        )
+
+        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
+
     @pytest.mark.asyncio
     async def test_get_config_from_bucket_merges_includes_over_gcs(self, monkeypatch):
         objects = {
@@ -232,11 +305,20 @@ class TestBucketConfigIncludes:
             "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]},
         }
 
-        async def fake_gcs(bucket_name, object_key):
-            return objects.get(object_key)
+        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.get_config_file_contents_from_gcs", fake_gcs
+            "litellm.proxy.common_utils.load_config_utils.gcs_config_bucket",
+            lambda bucket_name: FakeGCSBucket(),
         )
 
         config = await get_config_from_bucket(
@@ -247,6 +329,9 @@ class TestBucketConfigIncludes:
             "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):

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 os
 import re
@@ -716,7 +717,7 @@ async def test_ProxyConfig__process_includes_merges_files(tmp_path):
     inc.write_text("model_list:\n  - model_name: gpt-4\n")
     pc = ProxyConfig()
     cfg = {"include": ["models.yaml"], "model_list": [], "litellm_settings": {}}
-    result = await 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": {},
@@ -727,17 +728,72 @@ async def test_ProxyConfig__process_includes_merges_files(tmp_path):
 async def test_ProxyConfig__process_includes_missing_file_raises(tmp_path):
     pc = ProxyConfig()
     with pytest.raises(FileNotFoundError):
-        await 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"]}, base_dir=str(tmp_path))
+    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):
+    """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"
+    )
+    (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_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"}]}
+
+
 # ---------------------------------------------------------------------------
 # ProxyConfig.save_config
 # ---------------------------------------------------------------------------