Merge pull request #721 from tisnik/lcore-741-quota-limiters-in-configuration

LCORE-741: quota limiters in configuration

Author: tisnik

src/configuration.py

@@ -20,11 +20,14 @@
     InferenceConfiguration,
     DatabaseConfiguration,
     ConversationCacheConfiguration,
+    QuotaHandlersConfiguration,
 )
 
 from cache.cache import Cache
 from cache.cache_factory import CacheFactory
 
+from quota.quota_limiter import QuotaLimiter
+from quota.quota_limiter_factory import QuotaLimiterFactory
 
 logger = logging.getLogger(__name__)
 
@@ -48,6 +51,7 @@ def __init__(self) -> None:
         """Initialize the class instance."""
         self._configuration: Optional[Configuration] = None
         self._conversation_cache: Optional[Cache] = None
+        self._quota_limiters: list[QuotaLimiter] = []
 
     def load_configuration(self, filename: str) -> None:
         """Load configuration from YAML file."""
@@ -59,6 +63,10 @@ def load_configuration(self, filename: str) -> None:
 
     def init_from_dict(self, config_dict: dict[Any, Any]) -> None:
         """Initialize configuration from a dictionary."""
+        # clear cached values when configuration changes
+        self._conversation_cache = None
+        self._quota_limiters = []
+        # now it is possible to re-read configuration
         self._configuration = Configuration(**config_dict)
 
     @property
@@ -143,6 +151,13 @@ def database_configuration(self) -> DatabaseConfiguration:
             raise LogicError("logic error: configuration is not loaded")
         return self._configuration.database
 
+    @property
+    def quota_handlers_configuration(self) -> QuotaHandlersConfiguration:
+        """Return quota handlers configuration."""
+        if self._configuration is None:
+            raise LogicError("logic error: configuration is not loaded")
+        return self._configuration.quota_handlers
+
     @property
     def conversation_cache(self) -> Cache:
         """Return the conversation cache."""
@@ -154,5 +169,16 @@ def conversation_cache(self) -> Cache:
             )
         return self._conversation_cache
 
+    @property
+    def quota_limiters(self) -> list[QuotaLimiter]:
+        """Return list of all setup quota limiters."""
+        if self._configuration is None:
+            raise LogicError("logic error: configuration is not loaded")
+        if not self._quota_limiters:
+            self._quota_limiters = QuotaLimiterFactory.quota_limiters(
+                self._configuration.quota_handlers
+            )
+        return self._quota_limiters
+
 
 configuration: AppConfig = AppConfig()

tests/unit/test_configuration.py

@@ -16,12 +16,14 @@ def _reset_app_config_between_tests() -> Generator:
     # ensure clean state before each test
     try:
         AppConfig()._configuration = None  # type: ignore[attr-defined]
+        AppConfig()._quota_limiters = []  # type: ignore[attr-defined]
     except Exception:
         pass
     yield
     # ensure clean state after each test
     try:
         AppConfig()._configuration = None  # type: ignore[attr-defined]
+        AppConfig()._quota_limiters = []  # type: ignore[attr-defined]
     except Exception:
         pass
 
@@ -78,10 +80,18 @@ def test_default_configuration() -> None:
         # try to read property
         _ = cfg.conversation_cache_configuration  # pylint: disable=pointless-statement
 
+    with pytest.raises(Exception, match="logic error: configuration is not loaded"):
+        # try to read property
+        _ = cfg.quota_handlers_configuration  # pylint: disable=pointless-statement
+
     with pytest.raises(Exception, match="logic error: configuration is not loaded"):
         # try to read property
         _ = cfg.conversation_cache  # pylint: disable=pointless-statement
 
+    with pytest.raises(Exception, match="logic error: configuration is not loaded"):
+        # try to read property
+        _ = cfg.quota_limiters  # pylint: disable=pointless-statement
+
 
 def test_configuration_is_singleton() -> None:
     """Test that configuration is singleton."""
@@ -675,3 +685,124 @@ def test_configuration_with_in_memory_conversation_cache(tmpdir: Path) -> None:
     assert cfg.conversation_cache_configuration.memory is not None
     assert cfg.conversation_cache is not None
     assert isinstance(cfg.conversation_cache, InMemoryCache)
+
+
+def test_configuration_with_quota_handlers_no_storage(tmpdir: Path) -> None:
+    """Test loading configuration from YAML file with quota handlers configuration."""
+    cfg_filename = tmpdir / "config.yaml"
+    with open(cfg_filename, "w", encoding="utf-8") as fout:
+        fout.write(
+            """
+name: test service
+service:
+  host: localhost
+  port: 8080
+  auth_enabled: false
+  workers: 1
+  color_log: true
+  access_log: true
+llama_stack:
+  use_as_library_client: false
+  url: http://localhost:8321
+  api_key: test-key
+user_data_collection:
+  feedback_enabled: false
+quota_handlers:
+  limiters:
+    - name: user_monthly_limits
+      type: user_limiter
+      initial_quota: 10
+      quota_increase: 10
+      period: "2 seconds"
+    - name: cluster_monthly_limits
+      type: cluster_limiter
+      initial_quota: 100
+      quota_increase: 10
+      period: "10 seconds"
+  scheduler:
+    # scheduler ticks in seconds
+    period: 1
+            """
+        )
+
+    cfg = AppConfig()
+    cfg.load_configuration(str(cfg_filename))
+
+    assert cfg.quota_handlers_configuration is not None
+    assert cfg.quota_handlers_configuration.sqlite is None
+    assert cfg.quota_handlers_configuration.postgres is None
+    assert cfg.quota_handlers_configuration.limiters is not None
+    assert cfg.quota_handlers_configuration.scheduler is not None
+
+    # check the quota limiters configuration
+    assert len(cfg.quota_limiters) == 0
+
+    # check the scheduler configuration
+    assert cfg.quota_handlers_configuration.scheduler.period == 1
+
+
+def test_configuration_with_quota_handlers(tmpdir: Path) -> None:
+    """Test loading configuration from YAML file with quota handlers configuration."""
+    cfg_filename = tmpdir / "config.yaml"
+    with open(cfg_filename, "w", encoding="utf-8") as fout:
+        fout.write(
+            """
+name: test service
+service:
+  host: localhost
+  port: 8080
+  auth_enabled: false
+  workers: 1
+  color_log: true
+  access_log: true
+llama_stack:
+  use_as_library_client: false
+  url: http://localhost:8321
+  api_key: test-key
+user_data_collection:
+  feedback_enabled: false
+quota_handlers:
+  sqlite:
+      db_path: ":memory:"
+  limiters:
+    - name: user_monthly_limits
+      type: user_limiter
+      initial_quota: 10
+      quota_increase: 10
+      period: "2 seconds"
+    - name: cluster_monthly_limits
+      type: cluster_limiter
+      initial_quota: 100
+      quota_increase: 10
+      period: "10 seconds"
+  scheduler:
+    # scheduler ticks in seconds
+    period: 1
+            """
+        )
+
+    cfg = AppConfig()
+    cfg.load_configuration(str(cfg_filename))
+
+    assert cfg.quota_handlers_configuration is not None
+    assert cfg.quota_handlers_configuration.sqlite is not None
+    assert cfg.quota_handlers_configuration.postgres is None
+    assert cfg.quota_handlers_configuration.limiters is not None
+    assert cfg.quota_handlers_configuration.scheduler is not None
+
+    # check the storage
+    assert cfg.quota_handlers_configuration.sqlite.db_path == ":memory:"
+
+    # check the quota limiters configuration
+    assert len(cfg.quota_limiters) == 2
+    assert (
+        str(cfg.quota_limiters[0])
+        == "UserQuotaLimiter: initial quota: 10 increase by: 10"
+    )
+    assert (
+        str(cfg.quota_limiters[1])
+        == "ClusterQuotaLimiter: initial quota: 100 increase by: 10"
+    )
+
+    # check the scheduler configuration
+    assert cfg.quota_handlers_configuration.scheduler.period == 1