fix(dynamic_rate_limiter.py): add rpm allocation, priority + quota reservation to docs

2026-08-10 16:22:17 +00:00 · 2024-07-01 23:35:42 -07:00
parent 6b529d4e0e
commit 196b94455e
3 changed files with 181 additions and 82 deletions
@@ -152,11 +152,11 @@ litellm_remaining_team_budget_metric{team_alias="QA Prod Bot",team_id="de35b29e-
 ```


-### Dynamic TPM Allocation 
+### Dynamic TPM/RPM Allocation 

-Prevent projects from gobbling too much quota. 
+Prevent projects from gobbling too much tpm/rpm.

-Dynamically allocate TPM quota to api keys, based on active keys in that minute. [**See Code**](https://github.com/BerriAI/litellm/blob/9bffa9a48e610cc6886fc2dce5c1815aeae2ad46/litellm/proxy/hooks/dynamic_rate_limiter.py#L125)
+Dynamically allocate TPM/RPM quota to api keys, based on active keys in that minute. [**See Code**](https://github.com/BerriAI/litellm/blob/9bffa9a48e610cc6886fc2dce5c1815aeae2ad46/litellm/proxy/hooks/dynamic_rate_limiter.py#L125)

 1. Setup config.yaml 

@@ -247,4 +247,90 @@ except RateLimitError as e:

 ```
 This was rate limited b/c - Error code: 429 - {'error': {'message': {'error': 'Key=<hashed_token> over available TPM=0. Model TPM=0, Active keys=2'}, 'type': 'None', 'param': 'None', 'code': 429}}
+```
+
+
+#### ✨ [BETA] Set Priority / Reserve Quota
+
+Reserve tpm/rpm capacity for projects in prod.
+
+:::tip
+
+Reserving tpm/rpm on keys based on priority is a premium feature. Please [get an enterprise license](./enterprise.md) for it. 
+:::
+
+
+1. Setup config.yaml
+
+```yaml 
+model_list:
+  - model_name: gpt-3.5-turbo             
+    litellm_params:
+      model: "gpt-3.5-turbo"       
+      api_key: os.environ/OPENAI_API_KEY 
+      rpm: 100   
+
+litellm_settings:
+  callbacks: ["dynamic_rate_limiter"]
+  priority_reservation: {"dev": 0, "prod": 1}
+
+general_settings:
+  master_key: sk-1234 # OR set `LITELLM_MASTER_KEY=".."` in your .env
+  database_url: postgres://.. # OR set `DATABASE_URL=".."` in your .env
+```
+
+
+priority_reservation: 
+- Dict[str, float]
+  - str: can be any string
+  - float: from 0 to 1. Specify the % of tpm/rpm to reserve for keys of this priority.
+
+**Start Proxy**
+
+```
+litellm --config /path/to/config.yaml
+```
+
+2. Create a key with that priority
+
+```bash
+curl -X POST 'http://0.0.0.0:4000/key/generate' \
+-H 'Authorization: Bearer <your-master-key>' \
+-H 'Content-Type: application/json' \
+-D '{
+	"metadata": {"priority": "dev"} # 👈 KEY CHANGE
+}'
+```
+
+**Expected Response**
+
+```
+{
+  ...
+  "key": "sk-.."
+}
+```
+
+
+3. Test it!
+
+```bash
+curl -X POST 'http://0.0.0.0:4000/chat/completions' \
+  -H 'Content-Type: application/json' \
+  -H 'Authorization: sk-...' \ # 👈 key from step 2.
+  -D '{
+  "model": "gpt-3.5-turbo",
+  "messages": [
+      {
+      "role": "user",
+      "content": "what llm are you"
+      }
+  ],
+}'
+```
+
+**Expected Response**
+
+```
+Key=... over available RPM=0. Model RPM=100, Active keys=None
 ```
@@ -1,14 +1,16 @@
 model_list:
-  - model_name: "*"             # all requests where model not in your config go to this deployment
+  - model_name: gpt-3.5-turbo             # all requests where model not in your config go to this deployment
    litellm_params:
-      model: "openai/*"           
+      model: "gpt-3.5-turbo"        
+      rpm: 100   

 litellm_settings:
-  success_callback: ["s3"]
-  s3_callback_params:
-    s3_bucket_name: my-test-bucket-22-litellm   # AWS Bucket Name for S3
-    s3_region_name: us-west-2              # AWS Region Name for S3
-    s3_aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID  # us os.environ/<variable name> to pass environment variables. This is AWS Access Key ID for S3
-    s3_aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY  # AWS Secret Access Key for S3
-    s3_path: my-test-path
-
+  callbacks: ["dynamic_rate_limiter"]
+  priority_reservation: {"dev": 0, "prod": 1}
+#   success_callback: ["s3"]
+#   s3_callback_params:
+#     s3_bucket_name: my-test-bucket-22-litellm   # AWS Bucket Name for S3
+#     s3_region_name: us-west-2              # AWS Region Name for S3
+#     s3_aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID  # us os.environ/<variable name> to pass environment variables. This is AWS Access Key ID for S3
+#     s3_aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY  # AWS Secret Access Key for S3
+#     s3_path: my-test-path
@@ -102,78 +102,89 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
            - remaining_model_rpm: int or null. If available rpm is int, then this will be too.
            - active_projects: int or null
        """
-        weight: float = 1
-        if (
-            litellm.priority_reservation is None
-            or priority not in litellm.priority_reservation
-        ):
+        try:
+            weight: float = 1
+            if (
+                litellm.priority_reservation is None
+                or priority not in litellm.priority_reservation
+            ):
+                verbose_proxy_logger.error(
+                    "Priority Reservation not set. priority={}, but litellm.priority_reservation is {}.".format(
+                        priority, litellm.priority_reservation
+                    )
+                )
+            elif priority is not None and litellm.priority_reservation is not None:
+                if os.getenv("LITELLM_LICENSE", None) is None:
+                    verbose_proxy_logger.error(
+                        "PREMIUM FEATURE: Reserving tpm/rpm by priority is a premium feature. Please add a 'LITELLM_LICENSE' to your .env to enable this.\nGet a license: https://docs.litellm.ai/docs/proxy/enterprise."
+                    )
+                else:
+                    weight = litellm.priority_reservation[priority]
+
+            active_projects = await self.internal_usage_cache.async_get_cache(
+                model=model
+            )
+            current_model_tpm, current_model_rpm = (
+                await self.llm_router.get_model_group_usage(model_group=model)
+            )
+            model_group_info: Optional[ModelGroupInfo] = (
+                self.llm_router.get_model_group_info(model_group=model)
+            )
+            total_model_tpm: Optional[int] = None
+            total_model_rpm: Optional[int] = None
+            if model_group_info is not None:
+                if model_group_info.tpm is not None:
+                    total_model_tpm = model_group_info.tpm
+                if model_group_info.rpm is not None:
+                    total_model_rpm = model_group_info.rpm
+
+            remaining_model_tpm: Optional[int] = None
+            if total_model_tpm is not None and current_model_tpm is not None:
+                remaining_model_tpm = total_model_tpm - current_model_tpm
+            elif total_model_tpm is not None:
+                remaining_model_tpm = total_model_tpm
+
+            remaining_model_rpm: Optional[int] = None
+            if total_model_rpm is not None and current_model_rpm is not None:
+                remaining_model_rpm = total_model_rpm - current_model_rpm
+            elif total_model_rpm is not None:
+                remaining_model_rpm = total_model_rpm
+
+            available_tpm: Optional[int] = None
+
+            if remaining_model_tpm is not None:
+                if active_projects is not None:
+                    available_tpm = int(remaining_model_tpm * weight / active_projects)
+                else:
+                    available_tpm = int(remaining_model_tpm * weight)
+
+            if available_tpm is not None and available_tpm < 0:
+                available_tpm = 0
+
+            available_rpm: Optional[int] = None
+
+            if remaining_model_rpm is not None:
+                if active_projects is not None:
+                    available_rpm = int(remaining_model_rpm * weight / active_projects)
+                else:
+                    available_rpm = int(remaining_model_rpm * weight)
+
+            if available_rpm is not None and available_rpm < 0:
+                available_rpm = 0
+            return (
+                available_tpm,
+                available_rpm,
+                remaining_model_tpm,
+                remaining_model_rpm,
+                active_projects,
+            )
+        except Exception as e:
            verbose_proxy_logger.error(
-                "Priority Reservation not set. priority={}, but litellm.priority_reservation is {}.".format(
-                    priority, litellm.priority_reservation
+                "litellm.proxy.hooks.dynamic_rate_limiter.py::check_available_usage: Exception occurred - {}\n{}".format(
+                    str(e), traceback.format_exc()
                )
            )
-        elif priority is not None and litellm.priority_reservation is not None:
-            if os.getenv("LITELLM_LICENSE", None) is None:
-                verbose_proxy_logger.error(
-                    "PREMIUM FEATURE: Reserving tpm/rpm by priority is a premium feature. Please add a 'LITELLM_LICENSE' to your .env to enable this.\nGet a license: https://docs.litellm.ai/docs/proxy/enterprise."
-                )
-            else:
-                weight = litellm.priority_reservation[priority]
-        active_projects = await self.internal_usage_cache.async_get_cache(model=model)
-        current_model_tpm, current_model_rpm = (
-            await self.llm_router.get_model_group_usage(model_group=model)
-        )
-        model_group_info: Optional[ModelGroupInfo] = (
-            self.llm_router.get_model_group_info(model_group=model)
-        )
-        total_model_tpm: Optional[int] = None
-        total_model_rpm: Optional[int] = None
-        if model_group_info is not None:
-            if model_group_info.tpm is not None:
-                total_model_tpm = model_group_info.tpm
-            if model_group_info.rpm is not None:
-                total_model_rpm = model_group_info.rpm
-
-        remaining_model_tpm: Optional[int] = None
-        if total_model_tpm is not None and current_model_tpm is not None:
-            remaining_model_tpm = total_model_tpm - current_model_tpm
-        elif total_model_tpm is not None:
-            remaining_model_tpm = total_model_tpm
-
-        remaining_model_rpm: Optional[int] = None
-        if total_model_rpm is not None and current_model_rpm is not None:
-            remaining_model_rpm = total_model_rpm - current_model_rpm
-        elif total_model_rpm is not None:
-            remaining_model_rpm = total_model_rpm
-
-        available_tpm: Optional[int] = None
-
-        if remaining_model_tpm is not None:
-            if active_projects is not None:
-                available_tpm = int(remaining_model_tpm * weight / active_projects)
-            else:
-                available_tpm = int(remaining_model_tpm * weight)
-
-        if available_tpm is not None and available_tpm < 0:
-            available_tpm = 0
-
-        available_rpm: Optional[int] = None
-
-        if remaining_model_rpm is not None:
-            if active_projects is not None:
-                available_rpm = int(remaining_model_rpm * weight / active_projects)
-            else:
-                available_rpm = int(remaining_model_rpm * weight)
-
-        if available_rpm is not None and available_rpm < 0:
-            available_rpm = 0
-        return (
-            available_tpm,
-            available_rpm,
-            remaining_model_tpm,
-            remaining_model_rpm,
-            active_projects,
-        )
+            return None, None, None, None, None

    async def async_pre_call_hook(
        self,