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 index d6939db2e95..cf39bd6435e 100644 --- a/litellm/proxy/common_utils/config_includes.py +++ b/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 diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index 927deb68826..e72272b1aca 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/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: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a49df166dca..7ac5e697b61 100644 --- a/litellm/proxy/proxy_server.py +++ b/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 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 c5d84460d39..b654320569c 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,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): 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 9b9c4af9e42..3d9138ad098 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 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 # ---------------------------------------------------------------------------