Schedule budget resets at expectable times (#10331) (#10333)

* Schedule budget resets at expectable times (#10331)

* Enhance budget reset functionality with timezone support and standardized reset times

- Added `get_next_standardized_reset_time` function to calculate budget reset times based on specified durations and timezones.
- Introduced `timezone_utils.py` to manage timezone retrieval and budget reset time calculations.
- Updated budget reset logic in `reset_budget_job.py`, `internal_user_endpoints.py`, `key_management_endpoints.py`, and `team_endpoints.py` to utilize the new timezone-aware reset time calculations.
- Added unit tests for the new reset time functionality in `test_duration_parser.py`.
- Updated `.gitignore` to include `test.py` and made minor formatting adjustments in `docker-compose.yml` for consistency.

* Fixed linting

* Fix for mypy

* Fixed testcase for reset

* fix(duration_parser.py): move off zoneinfo - doesn't work with python 3.8

* test: update test

* refactor: improve budget reset time calculation and update related tests for accuracy

* clean up imports in team_endpoints.py

* test: update budget remaining hours assertions to reflect new reset time logic

* build(model_prices_and_context_window.json): update model

---------

Co-authored-by: Prathamesh Saraf <pratamesh1867@gmail.com>

.gitignore

@@ -89,3 +89,4 @@ litellm/proxy/migrations/*
 config.yaml
 tests/litellm/litellm_core_utils/llm_cost_calc/log.txt
 tests/test_custom_dir/*
+test.py

docker-compose.yml

@@ -16,23 +16,28 @@ services:
     ports:
       - "4000:4000" # Map the container port to the host, change the host port if necessary
     environment:
-        DATABASE_URL: "postgresql://llmproxy:dbpassword9090@db:5432/litellm"
-        STORE_MODEL_IN_DB: "True" # allows adding models to proxy via UI
+      DATABASE_URL: "postgresql://llmproxy:dbpassword9090@db:5432/litellm"
+      STORE_MODEL_IN_DB: "True" # allows adding models to proxy via UI
     env_file:
       - .env # Load local .env file
     depends_on:
-      - db  # Indicates that this service depends on the 'db' service, ensuring 'db' starts first
-    healthcheck:  # Defines the health check configuration for the container
-      test: [ "CMD", "curl", "-f", "http://localhost:4000/health/liveliness || exit 1" ]  # Command to execute for health check
-      interval: 30s  # Perform health check every 30 seconds
-      timeout: 10s   # Health check command times out after 10 seconds
-      retries: 3     # Retry up to 3 times if health check fails
-      start_period: 40s  # Wait 40 seconds after container start before beginning health checks
+      - db # Indicates that this service depends on the 'db' service, ensuring 'db' starts first
+    healthcheck: # Defines the health check configuration for the container
+      test: [
+          "CMD",
+          "curl",
+          "-f",
+          "http://localhost:4000/health/liveliness || exit 1",
+        ] # Command to execute for health check
+      interval: 30s # Perform health check every 30 seconds
+      timeout: 10s # Health check command times out after 10 seconds
+      retries: 3 # Retry up to 3 times if health check fails
+      start_period: 40s # Wait 40 seconds after container start before beginning health checks
 
- 
   db:
     image: postgres:16
     restart: always
+    container_name: litellm_db
     environment:
       POSTGRES_DB: litellm
       POSTGRES_USER: llmproxy
@@ -40,13 +45,13 @@ services:
     ports:
       - "5432:5432"
     volumes:
-      - postgres_data:/var/lib/postgresql/data  # Persists Postgres data across container restarts
+      - postgres_data:/var/lib/postgresql/data # Persists Postgres data across container restarts
     healthcheck:
       test: ["CMD-SHELL", "pg_isready -d litellm -U llmproxy"]
       interval: 1s
       timeout: 5s
       retries: 10
-  
+
   prometheus:
     image: prom/prometheus
     volumes:
@@ -55,14 +60,14 @@ services:
     ports:
       - "9090:9090"
     command:
-      - '--config.file=/etc/prometheus/prometheus.yml'
-      - '--storage.tsdb.path=/prometheus'
-      - '--storage.tsdb.retention.time=15d'
+      - "--config.file=/etc/prometheus/prometheus.yml"
+      - "--storage.tsdb.path=/prometheus"
+      - "--storage.tsdb.retention.time=15d"
     restart: always
 
 volumes:
   prometheus_data:
     driver: local
   postgres_data:
-    name: litellm_postgres_data  # Named volume for Postgres data persistence
+    name: litellm_postgres_data # Named volume for Postgres data persistence
 

docs/my-website/docs/proxy/budget_reset_and_tz.md

@@ -0,0 +1,33 @@
+## Budget Reset Times and Timezones
+
+LiteLLM now supports predictable budget reset times that align with natural calendar boundaries:
+
+- All budgets reset at midnight (00:00:00) in the configured timezone
+- Special handling for common durations:
+  - Daily (24h/1d): Reset at midnight every day
+  - Weekly (7d): Reset on Monday at midnight
+  - Monthly (30d): Reset on the 1st of each month at midnight
+
+### Configuring the Timezone
+
+You can specify the timezone for all budget resets in your configuration file:
+
+```yaml
+litellm_settings:
+  max_budget: 100 # (float) sets max budget as $100 USD
+  budget_duration: 30d # (number)(s/m/h/d)
+  timezone: "US/Eastern" # Any valid timezone string
+```
+
+This ensures that all budget resets happen at midnight in your specified timezone rather than in UTC.
+If no timezone is specified, UTC will be used by default.
+
+Common timezone values:
+
+- `UTC` - Coordinated Universal Time
+- `US/Eastern` - Eastern Time
+- `US/Pacific` - Pacific Time
+- `Europe/London` - UK Time
+- `Asia/Kolkata` - Indian Standard Time (IST)
+- `Asia/Tokyo` - Japan Standard Time
+- `Australia/Sydney` - Australian Eastern Time

litellm/litellm_core_utils/duration_parser.py

@@ -8,8 +8,8 @@ duration_in_seconds is used in diff parts of the code base, example
 
 import re
 import time
-from datetime import datetime, timedelta
-from typing import Tuple
+from datetime import datetime, timedelta, timezone
+from typing import Optional, Tuple
 
 
 def _extract_from_regex(duration: str) -> Tuple[int, str]:
@@ -93,3 +93,253 @@ def duration_in_seconds(duration: str) -> int:
 
     else:
         raise ValueError(f"Unsupported duration unit, passed duration: {duration}")
+
+
+def get_next_standardized_reset_time(
+    duration: str, current_time: datetime, timezone_str: str = "UTC"
+) -> datetime:
+    """
+    Get the next standardized reset time based on the duration.
+
+    All durations will reset at predictable intervals, aligned from the current time:
+    - Nd: If N=1, reset at next midnight; if N>1, reset every N days from now
+    - Nh: Every N hours, aligned to hour boundaries (e.g., 1:00, 2:00)
+    - Nm: Every N minutes, aligned to minute boundaries (e.g., 1:05, 1:10)
+    - Ns: Every N seconds, aligned to second boundaries
+
+    Parameters:
+    - duration: Duration string (e.g. "30s", "30m", "30h", "30d")
+    - current_time: Current datetime
+    - timezone_str: Timezone string (e.g. "UTC", "US/Eastern", "Asia/Kolkata")
+
+    Returns:
+    - Next reset time at a standardized interval in the specified timezone
+    """
+    # Set up timezone and normalize current time
+    current_time, timezone = _setup_timezone(current_time, timezone_str)
+
+    # Parse duration
+    value, unit = _parse_duration(duration)
+    if value is None:
+        # Fall back to default if format is invalid
+        return current_time.replace(
+            hour=0, minute=0, second=0, microsecond=0
+        ) + timedelta(days=1)
+
+    # Midnight of the current day in the specified timezone
+    base_midnight = current_time.replace(hour=0, minute=0, second=0, microsecond=0)
+
+    # Handle different time units
+    if unit == "d":
+        return _handle_day_reset(current_time, base_midnight, value, timezone)
+    elif unit == "h":
+        return _handle_hour_reset(current_time, base_midnight, value)
+    elif unit == "m":
+        return _handle_minute_reset(current_time, base_midnight, value)
+    elif unit == "s":
+        return _handle_second_reset(current_time, base_midnight, value)
+    else:
+        # Unrecognized unit, default to next midnight
+        return base_midnight + timedelta(days=1)
+
+
+def _setup_timezone(
+    current_time: datetime, timezone_str: str = "UTC"
+) -> Tuple[datetime, timezone]:
+    """Set up timezone and normalize current time to that timezone."""
+    try:
+        if timezone_str is None:
+            tz = timezone.utc
+        else:
+            # Map common timezone strings to their UTC offsets
+            timezone_map = {
+                "US/Eastern": timezone(timedelta(hours=-4)),  # EDT
+                "US/Pacific": timezone(timedelta(hours=-7)),  # PDT
+                "Asia/Kolkata": timezone(timedelta(hours=5, minutes=30)),  # IST
+                "Europe/London": timezone(timedelta(hours=1)),  # BST
+                "UTC": timezone.utc,
+            }
+            tz = timezone_map.get(timezone_str, timezone.utc)
+    except Exception:
+        # If timezone is invalid, fall back to UTC
+        tz = timezone.utc
+
+    # Convert current_time to the target timezone
+    if current_time.tzinfo is None:
+        # Naive datetime - assume it's UTC
+        utc_time = current_time.replace(tzinfo=timezone.utc)
+        current_time = utc_time.astimezone(tz)
+    else:
+        # Already has timezone - convert to target timezone
+        current_time = current_time.astimezone(tz)
+
+    return current_time, tz
+
+
+def _parse_duration(duration: str) -> Tuple[Optional[int], Optional[str]]:
+    """Parse the duration string into value and unit."""
+    match = re.match(r"(\d+)([a-z]+)", duration)
+    if not match:
+        return None, None
+
+    value, unit = match.groups()
+    return int(value), unit
+
+
+def _handle_day_reset(
+    current_time: datetime, base_midnight: datetime, value: int, timezone: timezone
+) -> datetime:
+    """Handle day-based reset times."""
+    if value == 1:  # Daily reset at midnight
+        return base_midnight + timedelta(days=1)
+    elif value == 7:  # Weekly reset on Monday at midnight
+        days_until_monday = (7 - current_time.weekday()) % 7
+        if days_until_monday == 0:  # If today is Monday
+            days_until_monday = 7
+        return base_midnight + timedelta(days=days_until_monday)
+    elif value == 30:  # Monthly reset on 1st at midnight
+        # Get 1st of next month at midnight
+        if current_time.month == 12:
+            next_reset = datetime(
+                year=current_time.year + 1,
+                month=1,
+                day=1,
+                hour=0,
+                minute=0,
+                second=0,
+                microsecond=0,
+                tzinfo=timezone,
+            )
+        else:
+            next_reset = datetime(
+                year=current_time.year,
+                month=current_time.month + 1,
+                day=1,
+                hour=0,
+                minute=0,
+                second=0,
+                microsecond=0,
+                tzinfo=timezone,
+            )
+        return next_reset
+    else:  # Custom day value - next interval is value days from current
+        return current_time.replace(
+            hour=0, minute=0, second=0, microsecond=0
+        ) + timedelta(days=value)
+
+
+def _handle_hour_reset(
+    current_time: datetime, base_midnight: datetime, value: int
+) -> datetime:
+    """Handle hour-based reset times."""
+    current_hour = current_time.hour
+    current_minute = current_time.minute
+    current_second = current_time.second
+    current_microsecond = current_time.microsecond
+
+    # Calculate next hour aligned with the value
+    if current_minute == 0 and current_second == 0 and current_microsecond == 0:
+        next_hour = (
+            current_hour + value - (current_hour % value)
+            if current_hour % value != 0
+            else current_hour + value
+        )
+    else:
+        next_hour = (
+            current_hour + value - (current_hour % value)
+            if current_hour % value != 0
+            else current_hour + value
+        )
+
+    # Handle overnight case
+    if next_hour >= 24:
+        next_hour = next_hour % 24
+        next_day = base_midnight + timedelta(days=1)
+        return next_day.replace(hour=next_hour)
+
+    return current_time.replace(hour=next_hour, minute=0, second=0, microsecond=0)
+
+
+def _handle_minute_reset(
+    current_time: datetime, base_midnight: datetime, value: int
+) -> datetime:
+    """Handle minute-based reset times."""
+    current_hour = current_time.hour
+    current_minute = current_time.minute
+    current_second = current_time.second
+    current_microsecond = current_time.microsecond
+
+    # Calculate next minute aligned with the value
+    if current_second == 0 and current_microsecond == 0:
+        next_minute = (
+            current_minute + value - (current_minute % value)
+            if current_minute % value != 0
+            else current_minute + value
+        )
+    else:
+        next_minute = (
+            current_minute + value - (current_minute % value)
+            if current_minute % value != 0
+            else current_minute + value
+        )
+
+    # Handle hour rollover
+    next_hour = current_hour + (next_minute // 60)
+    next_minute = next_minute % 60
+
+    # Handle overnight case
+    if next_hour >= 24:
+        next_hour = next_hour % 24
+        next_day = base_midnight + timedelta(days=1)
+        return next_day.replace(
+            hour=next_hour, minute=next_minute, second=0, microsecond=0
+        )
+
+    return current_time.replace(
+        hour=next_hour, minute=next_minute, second=0, microsecond=0
+    )
+
+
+def _handle_second_reset(
+    current_time: datetime, base_midnight: datetime, value: int
+) -> datetime:
+    """Handle second-based reset times."""
+    current_hour = current_time.hour
+    current_minute = current_time.minute
+    current_second = current_time.second
+    current_microsecond = current_time.microsecond
+
+    # Calculate next second aligned with the value
+    if current_microsecond == 0:
+        next_second = (
+            current_second + value - (current_second % value)
+            if current_second % value != 0
+            else current_second + value
+        )
+    else:
+        next_second = (
+            current_second + value - (current_second % value)
+            if current_second % value != 0
+            else current_second + value
+        )
+
+    # Handle minute rollover
+    additional_minutes = next_second // 60
+    next_second = next_second % 60
+    next_minute = current_minute + additional_minutes
+
+    # Handle hour rollover
+    next_hour = current_hour + (next_minute // 60)
+    next_minute = next_minute % 60
+
+    # Handle overnight case
+    if next_hour >= 24:
+        next_hour = next_hour % 24
+        next_day = base_midnight + timedelta(days=1)
+        return next_day.replace(
+            hour=next_hour, minute=next_minute, second=next_second, microsecond=0
+        )
+
+    return current_time.replace(
+        hour=next_hour, minute=next_minute, second=next_second, microsecond=0
+    )

litellm/proxy/common_utils/reset_budget_job.py

@@ -1,11 +1,10 @@
 import asyncio
 import json
 import time
-from datetime import datetime, timedelta
+from datetime import datetime
 from typing import List, Literal, Optional, Union
 
 from litellm._logging import verbose_proxy_logger
-from litellm.litellm_core_utils.duration_parser import duration_in_seconds
 from litellm.proxy._types import (
     LiteLLM_TeamTable,
     LiteLLM_UserTable,
@@ -328,8 +327,11 @@ class ResetBudgetJob:
         try:
             item.spend = 0.0
             if hasattr(item, "budget_duration") and item.budget_duration is not None:
-                duration_s = duration_in_seconds(duration=item.budget_duration)
-                item.budget_reset_at = current_time + timedelta(seconds=duration_s)
+                # Get standardized reset time based on budget duration
+                from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
+                item.budget_reset_at = get_budget_reset_time(
+                    budget_duration=item.budget_duration
+                )
             return item
         except Exception as e:
             verbose_proxy_logger.exception(

litellm/proxy/common_utils/timezone_utils.py

@@ -0,0 +1,29 @@
+from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time
+
+from datetime import datetime, timezone
+def get_budget_reset_timezone():
+    """
+    Get the budget reset timezone from general_settings.
+    Falls back to UTC if not specified.
+    """
+    # Import at function level to avoid circular imports
+    from litellm.proxy.proxy_server import general_settings
+    if general_settings:
+        litellm_settings = general_settings.get("litellm_settings", {})
+        if litellm_settings and "timezone" in litellm_settings:
+            return litellm_settings["timezone"]
+    
+    return "UTC"
+
+
+def get_budget_reset_time(budget_duration: str):
+    """
+    Get the budget reset time from general_settings.
+    Falls back to UTC if not specified.
+    """
+    reset_at = get_next_standardized_reset_time(
+        duration=budget_duration, 
+        current_time=datetime.now(timezone.utc),
+        timezone_str=get_budget_reset_timezone()
+    )
+    return reset_at

litellm/proxy/management_endpoints/internal_user_endpoints.py

@@ -14,7 +14,7 @@ These are members of a Team on LiteLLM
 import asyncio
 import traceback
 import uuid
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timezone
 from typing import Any, Dict, List, Optional, Union, cast
 
 import fastapi
@@ -22,7 +22,6 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
 
 import litellm
 from litellm._logging import verbose_proxy_logger
-from litellm.litellm_core_utils.duration_parser import duration_in_seconds
 from litellm.proxy._types import *
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity
@@ -651,9 +650,10 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest) -> di
         is_internal_user = True
 
     if "budget_duration" in non_default_values:
-        duration_s = duration_in_seconds(duration=non_default_values["budget_duration"])
-        user_reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s)
-        non_default_values["budget_reset_at"] = user_reset_at
+        from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
+        non_default_values["budget_reset_at"] = get_budget_reset_time(
+            budget_duration=non_default_values["budget_duration"]
+        )
 
     if "max_budget" not in non_default_values:
         if (
@@ -668,11 +668,10 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest) -> di
             non_default_values[
                 "budget_duration"
             ] = litellm.internal_user_budget_duration
-            duration_s = duration_in_seconds(
-                duration=non_default_values["budget_duration"]
+            from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
+            non_default_values["budget_reset_at"] = get_budget_reset_time(
+                budget_duration=non_default_values["budget_duration"]
             )
-            user_reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s)
-            non_default_values["budget_reset_at"] = user_reset_at
 
     return non_default_values
 

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -663,8 +663,8 @@ def prepare_key_update_data(
             and (isinstance(budget_duration, str))
             and len(budget_duration) > 0
         ):
-            duration_s = duration_in_seconds(duration=budget_duration)
-            key_reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s)
+            from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
+            key_reset_at = get_budget_reset_time(budget_duration=budget_duration)
             non_default_values["budget_reset_at"] = key_reset_at
             non_default_values["budget_duration"] = budget_duration
 

litellm/proxy/management_endpoints/team_endpoints.py

@@ -13,7 +13,7 @@ import asyncio
 import json
 import traceback
 import uuid
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timezone
 from typing import Any, Dict, List, Optional, Tuple, Union, cast
 
 import fastapi
@@ -198,7 +198,6 @@ async def new_team(  # noqa: PLR0915
     try:
         from litellm.proxy.proxy_server import (
             create_audit_log_for_update,
-            duration_in_seconds,
             litellm_proxy_admin_name,
             prisma_client,
         )
@@ -313,11 +312,10 @@ async def new_team(  # noqa: PLR0915
 
         # If budget_duration is set, set `budget_reset_at`
         if complete_team_data.budget_duration is not None:
-            duration_s = duration_in_seconds(
-                duration=complete_team_data.budget_duration
+            from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
+            complete_team_data.budget_reset_at = get_budget_reset_time(
+                budget_duration=complete_team_data.budget_duration, 
             )
-            reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s)
-            complete_team_data.budget_reset_at = reset_at
 
         complete_team_data_dict = complete_team_data.model_dump(exclude_none=True)
         complete_team_data_dict = prisma_client.jsonify_team_object(
@@ -468,7 +466,6 @@ async def update_team(
     from litellm.proxy.auth.auth_checks import _cache_team_object
     from litellm.proxy.proxy_server import (
         create_audit_log_for_update,
-        duration_in_seconds,
         litellm_proxy_admin_name,
         prisma_client,
         proxy_logging_obj,
@@ -496,9 +493,9 @@ async def update_team(
 
     # Check budget_duration and budget_reset_at
     if data.budget_duration is not None:
-        duration_s = duration_in_seconds(duration=data.budget_duration)
-        reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s)
-
+        from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
+        reset_at = get_budget_reset_time(budget_duration=data.budget_duration)
+        
         # set the budget_reset_at in DB
         updated_kv["budget_reset_at"] = reset_at
 

tests/litellm/litellm_core_utils/test_duration_parser.py

@@ -0,0 +1,113 @@
+import unittest
+from datetime import datetime, timezone
+from zoneinfo import ZoneInfo
+from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time
+
+class TestStandardizedResetTime(unittest.TestCase):
+    def test_day_based_resets(self):
+        """Test day-based reset durations (1d, 7d, 30d)"""
+        # Base time: 2023-05-15 10:30:00 UTC
+        base_time = datetime(2023, 5, 15, 10, 30, 0, tzinfo=timezone.utc)
+        
+        # Daily reset (1d) - should reset at next midnight
+        daily_expected = datetime(2023, 5, 16, 0, 0, 0, tzinfo=timezone.utc)
+        daily_result = get_next_standardized_reset_time("1d", base_time, "UTC")
+        self.assertEqual(daily_result, daily_expected)
+        
+        # Weekly reset (7d) - should reset on next Monday
+        wednesday = datetime(2023, 5, 17, 15, 45, 0, tzinfo=timezone.utc)  # A Wednesday
+        weekly_expected = datetime(2023, 5, 22, 0, 0, 0, tzinfo=timezone.utc)  # Next Monday
+        weekly_result = get_next_standardized_reset_time("7d", wednesday, "UTC")
+        self.assertEqual(weekly_result, weekly_expected)
+        
+        # Monthly reset (30d) - should reset on 1st of next month
+        monthly_expected = datetime(2023, 6, 1, 0, 0, 0, tzinfo=timezone.utc)
+        monthly_result = get_next_standardized_reset_time("30d", base_time, "UTC")
+        self.assertEqual(monthly_result, monthly_expected)
+        
+        # Custom day reset (3d) - should reset after 3 days
+        custom_day_expected = datetime(2023, 5, 18, 0, 0, 0, tzinfo=timezone.utc)
+        custom_day_result = get_next_standardized_reset_time("3d", base_time, "UTC")
+        self.assertEqual(custom_day_result, custom_day_expected)
+
+    def test_hour_minute_second_resets(self):
+        """Test hour, minute, and second based reset durations"""
+        # Base time: 2023-05-15 15:20:30 UTC (3:20:30 PM)
+        base_time = datetime(2023, 5, 15, 15, 20, 30, tzinfo=timezone.utc)
+        
+        # 2-hour reset - should reset at next even hour (16:00)
+        hour_expected = datetime(2023, 5, 15, 16, 0, 0, tzinfo=timezone.utc)
+        hour_result = get_next_standardized_reset_time("2h", base_time, "UTC")
+        self.assertEqual(hour_result, hour_expected)
+        
+        # 30-minute reset - should reset at next 30-minute mark (15:30)
+        minute_expected = datetime(2023, 5, 15, 15, 30, 0, tzinfo=timezone.utc)
+        minute_result = get_next_standardized_reset_time("30m", base_time, "UTC")
+        self.assertEqual(minute_result, minute_expected)
+        
+        # 15-second reset - should reset at next 15-second mark (15:20:45)
+        second_expected = datetime(2023, 5, 15, 15, 20, 45, tzinfo=timezone.utc)
+        second_result = get_next_standardized_reset_time("15s", base_time, "UTC")
+        self.assertEqual(second_result, second_expected)
+
+    def test_timezone_handling(self):
+        """Test timezone handling with different regions"""
+        # Base time: 2023-05-15 22:30:00 UTC (late in UTC day)
+        base_time = datetime(2023, 5, 15, 22, 30, 0, tzinfo=timezone.utc)
+        
+        # Test daily reset in different timezones
+        # US/Eastern (UTC-4): 6:30 PM, so next reset is midnight same day
+        eastern = ZoneInfo("US/Eastern")
+        eastern_expected = datetime(2023, 5, 16, 0, 0, 0, tzinfo=eastern)
+        eastern_result = get_next_standardized_reset_time("1d", base_time, "US/Eastern")
+        self.assertEqual(eastern_result, eastern_expected)
+        
+        # Asia/Kolkata (UTC+5:30): 4:00 AM next day, so next reset is midnight the day after
+        ist = ZoneInfo("Asia/Kolkata")
+        ist_expected = datetime(2023, 5, 17, 0, 0, 0, tzinfo=ist)
+        ist_result = get_next_standardized_reset_time("1d", base_time, "Asia/Kolkata")
+        self.assertEqual(ist_result, ist_expected)
+        
+        # Test hourly reset in different timezones
+        # US/Pacific (UTC-7): 3:30 PM, so next 2h reset is 4:00 PM
+        pacific = ZoneInfo("US/Pacific")
+        pacific_expected = datetime(2023, 5, 15, 16, 0, 0, tzinfo=pacific)
+        pacific_result = get_next_standardized_reset_time("2h", base_time, "US/Pacific")
+        self.assertEqual(pacific_result, pacific_expected)
+        
+        # Test minute reset in different timezones
+        # Europe/London (UTC+1): 11:30 PM, so next 15m reset is 11:45 PM
+        london = ZoneInfo("Europe/London")
+        london_expected = datetime(2023, 5, 15, 23, 45, 0, tzinfo=london)
+        london_result = get_next_standardized_reset_time("15m", base_time, "Europe/London")
+        self.assertEqual(london_result, london_expected)
+
+    def test_edge_cases(self):
+        """Test edge cases and boundary conditions"""
+        # Exactly on hour boundary
+        on_hour = datetime(2023, 5, 15, 14, 0, 0, tzinfo=timezone.utc)
+        hour_expected = datetime(2023, 5, 15, 16, 0, 0, tzinfo=timezone.utc)
+        hour_result = get_next_standardized_reset_time("2h", on_hour, "UTC")
+        self.assertEqual(hour_result, hour_expected)
+        
+        # Exactly on minute boundary
+        on_minute = datetime(2023, 5, 15, 14, 30, 0, tzinfo=timezone.utc)
+        minute_expected = datetime(2023, 5, 15, 15, 0, 0, tzinfo=timezone.utc)
+        minute_result = get_next_standardized_reset_time("30m", on_minute, "UTC")
+        self.assertEqual(minute_result, minute_expected)
+        
+        # Near day boundary
+        near_midnight = datetime(2023, 5, 15, 23, 50, 0, tzinfo=timezone.utc)
+        
+        # 30m near midnight - should roll over to next day
+        midnight_minute_expected = datetime(2023, 5, 16, 0, 0, 0, tzinfo=timezone.utc)
+        midnight_minute_result = get_next_standardized_reset_time("30m", near_midnight, "UTC")
+        self.assertEqual(midnight_minute_result, midnight_minute_expected)
+        
+        # Invalid timezone - should fall back to UTC
+        invalid_tz_expected = datetime(2023, 5, 16, 0, 0, 0, tzinfo=timezone.utc)
+        invalid_tz_result = get_next_standardized_reset_time("1d", on_hour, "NonExistentTimeZone")
+        self.assertEqual(invalid_tz_result, invalid_tz_expected)
+
+if __name__ == "__main__":
+    unittest.main()

tests/litellm/proxy/common_utils/test_reset_budget_job.py

@@ -3,7 +3,7 @@ import json
 import os
 import sys
 import time
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 
 import pytest
 from fastapi.testclient import TestClient
@@ -68,8 +68,8 @@ async def run_async_test(coro):
 
 # Tests
 def test_reset_budget_for_key(reset_budget_job, mock_prisma_client):
-    # Setup test data
-    now = datetime.utcnow()
+    # Setup test data with timezone-aware datetime
+    now = datetime.now(timezone.utc)
     test_key = type(
         "LiteLLM_VerificationToken",
         (),
@@ -94,8 +94,8 @@ def test_reset_budget_for_key(reset_budget_job, mock_prisma_client):
 
 
 def test_reset_budget_for_user(reset_budget_job, mock_prisma_client):
-    # Setup test data
-    now = datetime.utcnow()
+    # Setup test data with timezone-aware datetime
+    now = datetime.now(timezone.utc)
     test_user = type(
         "LiteLLM_UserTable",
         (),
@@ -120,8 +120,8 @@ def test_reset_budget_for_user(reset_budget_job, mock_prisma_client):
 
 
 def test_reset_budget_for_team(reset_budget_job, mock_prisma_client):
-    # Setup test data
-    now = datetime.utcnow()
+    # Setup test data with timezone-aware datetime
+    now = datetime.now(timezone.utc)
     test_team = type(
         "LiteLLM_TeamTable",
         (),
@@ -146,8 +146,8 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client):
 
 
 def test_reset_budget_all(reset_budget_job, mock_prisma_client):
-    # Setup test data
-    now = datetime.utcnow()
+    # Setup test data with timezone-aware datetime
+    now = datetime.now(timezone.utc)
 
     # Create test objects for all three types
     test_key = type(
@@ -198,4 +198,4 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client):
     # Check that all spends were reset to 0
     assert mock_prisma_client.updated_data["key"][0].spend == 0.0
     assert mock_prisma_client.updated_data["user"][0].spend == 0.0
-    assert mock_prisma_client.updated_data["team"][0].spend == 0.0
+    assert mock_prisma_client.updated_data["team"][0].spend == 0.0

tests/otel_tests/test_prometheus.py

@@ -9,7 +9,6 @@ import uuid
 import os
 import sys
 from openai import AsyncOpenAI
-import time
 from typing import Dict, Any
 
 sys.path.insert(
@@ -385,10 +384,9 @@ async def test_team_budget_metrics():
         ), "remaining budget should be less than 10.0 after first request"
         assert first_budget["total"] == 10.0, "Total budget metric is incorrect"
         print("first_budget['remaining_hours']", first_budget["remaining_hours"])
-        # Verify remaining hours matches 7 days (with small delta for processing time)
-        assert (
-            abs(first_budget["remaining_hours"] - (7 * 24)) <= 0.1
-        ), "Budget remaining hours should be approximately 7 days (168 hours)"
+        # The budget reset time is now midnight, not exactly 7 days (168 hours) from creation
+        # So we'll check if it's within a reasonable range (5-7 days)
+        assert 120 <= first_budget["remaining_hours"] <= 168, "Budget remaining hours should be within a reasonable range (5-7 days)"
 
         # Get team info and verify spend matches prometheus metrics
         team_info = await get_team_info(session, team_id)
@@ -518,10 +516,9 @@ async def test_key_budget_metrics():
         ), "remaining budget should be less than 10.0 after first request"
         assert first_budget["total"] == 10.0, "Total budget metric is incorrect"
         print("first_budget['remaining_hours']", first_budget["remaining_hours"])
-        # Verify remaining hours matches 7 days (with small delta for processing time)
-        assert (
-            abs(first_budget["remaining_hours"] - (7 * 24)) <= 0.1
-        ), "Budget remaining hours should be approximately 7 days (168 hours)"
+        # The budget reset time is now midnight, not exactly 7 days (168 hours) from creation
+        # So we'll check if it's within a reasonable range (5-7 days)
+        assert 120 <= first_budget["remaining_hours"] <= 168, "Budget remaining hours should be within a reasonable range (5-7 days)"
 
         # Get key info and verify spend matches prometheus metrics
         key_info = await get_key_info(session, key)

tests/proxy_unit_tests/test_key_generate_prisma.py

@@ -1361,8 +1361,19 @@ def test_generate_and_update_key(prisma_client):
             )
             current_time = datetime.now(timezone.utc)
 
-            # assert budget_reset_at is 30 days from now
-            assert 31 >= (budget_reset_at - current_time).days >= 27
+            # Calculate days until end of current month
+            if current_time.month == 12:
+                end_of_month = datetime(current_time.year + 1, 1, 1, tzinfo=timezone.utc)
+            else:
+                end_of_month = datetime(current_time.year, current_time.month + 1, 1, tzinfo=timezone.utc)
+            days_until_end_of_month = (end_of_month - current_time).days - 1
+
+            # assert budget_reset_at is at the end of the current month
+            # assert days_until_end_of_month >= (budget_reset_at - current_time).days >= days_until_end_of_month - 2 # handle time zone differences
+            # Check that budget_reset_at is on the first day of next month
+            next_month_first_day = end_of_month 
+            # Assert that the reset date is the 1st of next month (0 or 1 day difference)
+            assert abs((budget_reset_at - next_month_first_day).days) <= 1
 
             # cleanup - delete key
             delete_key_request = KeyRequest(keys=[generated_key])
@@ -1924,7 +1935,7 @@ async def test_call_with_key_never_over_budget(prisma_client):
         pytest.fail(f"This should have not failed!. They key uses max_budget=None. {e}")
 
 
-@pytest.mark.asyncio()
+@pytest.mark.asyncio
 async def test_call_with_key_over_budget_stream(prisma_client):
     # 14. Make a call with a key over budget, expect to fail
     setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
@@ -2729,10 +2740,14 @@ async def test_create_update_team(prisma_client):
     budget_reset_at = _updated_info["budget_reset_at"].replace(tzinfo=timezone.utc)
     current_time = datetime.datetime.now(timezone.utc)
 
-    # assert budget_reset_at is 2 days from now
-    assert (
-        abs((budget_reset_at - current_time).total_seconds() - 2 * 24 * 60 * 60) <= 10
-    )
+    # Verify that budget_reset_at is at midnight (hour, minute, second are all 0)
+    assert budget_reset_at.hour == 0
+    assert budget_reset_at.minute == 0
+    assert budget_reset_at.second == 0
+
+    # Calculate days difference - should be close to 2 days (within 1 day to account for time of test execution)
+    days_diff = (budget_reset_at.date() - current_time.date()).days
+    assert 1 <= days_diff <= 2
 
     # now hit team_info
     try:
@@ -2859,12 +2874,18 @@ async def test_update_user_unit_test(prisma_client):
     assert _user_info["rpm_limit"] == 100
     assert _user_info["metadata"] == {"very-new-metadata": "something"}
 
-    # budget reset at should be 10 days from now
+    # budget_reset_at should be at midnight 10 days from now
     budget_reset_at = _user_info["budget_reset_at"].replace(tzinfo=timezone.utc)
     current_time = datetime.now(timezone.utc)
-    assert (
-        abs((budget_reset_at - current_time).total_seconds() - 10 * 24 * 60 * 60) <= 10
-    )
+    
+    # Verify that budget_reset_at is at midnight (hour, minute, second are all 0)
+    assert budget_reset_at.hour == 0
+    assert budget_reset_at.minute == 0
+    assert budget_reset_at.second == 0
+    
+    # Calculate days difference - should be close to 10 days (within 1 day to account for time of test execution)
+    days_diff = (budget_reset_at.date() - current_time.date()).days
+    assert 9 <= days_diff <= 10
 
 
 @pytest.mark.asyncio()