Skip to content

Commit 32378c3

Browse files
PhiRhoPhiRho
committed
Enhance Makefile and worker logic for budget management
- Updated Makefile to enforce regex parameter requirement for backend tests. - Refactored budget management logic in worker.py to improve handling of budget duration and amount updates based on reset time alignment. - Added detailed logging for budget updates and conditions triggering updates. - Enhanced test coverage for budget management scenarios, including handling of None and "0d" budget durations.
1 parent c843d9c commit 32378c3

File tree

3 files changed

+409
-306
lines changed

3 files changed

+409
-306
lines changed

‎Makefile‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,12 @@ test-postgres: test-clean test-network
2424
sleep 5
2525

2626
# Run backend tests for a specific regex
27+
# Usage: make backend-test-regex regex="test_pattern"
2728
backend-test-regex: test-clean backend-test-build test-postgres
28-
@read -p "Enter regex: " regex; \
29+
@if [ -z "$(regex)" ]; then \
30+
echo "Error: regex parameter is required. Usage: make backend-test-regex regex=\"test_pattern\""; \
31+
exit 1; \
32+
fi
2933
docker run --rm \
3034
--network amazeeai_default \
3135
-e DATABASE_URL="postgresql://postgres:postgres@amazee-test-postgres/postgres_service" \
@@ -40,7 +44,7 @@ backend-test-regex: test-clean backend-test-build test-postgres
4044
-e ENV_SUFFIX="test" \
4145
-v $(PWD)/app:/app/app \
4246
-v $(PWD)/tests:/app/tests \
43-
amazee-backend-test pytest -vv -k "$$regex"
47+
amazee-backend-test pytest -vv -k "$(regex)"
4448

4549
# Run backend tests in a new container
4650
backend-test: test-clean backend-test-build test-postgres

‎app/core/worker.py‎

Lines changed: 64 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -298,83 +298,84 @@ async def monitor_team_keys(
298298
# Check current values and only update if they don't match the parameters
299299
current_budget_duration = info.get("budget_duration")
300300
current_max_budget = info.get("max_budget")
301+
budget_reset_at = info.get("budget_reset_at")
301302

302303
needs_update = False
303304
data = {"litellm_token": key.litellm_token}
304305

305-
# Check if budget_duration needs updating
306-
expected_budget_duration = f"{renewal_period_days}d"
307-
if current_budget_duration != expected_budget_duration:
308-
data["budget_duration"] = expected_budget_duration
309-
needs_update = True
310-
logger.info(f"Key {key.id} budget_duration will be updated from '{current_budget_duration}' to '{expected_budget_duration}'")
311-
312-
# Check if budget_amount needs updating
306+
# Rule 1: If the budget amount mis-matches, always update that field
313307
if max_budget_amount is not None and current_max_budget != max_budget_amount:
314308
data["budget_amount"] = max_budget_amount
315309
needs_update = True
316-
logger.info(f"Key {key.id} budget_amount will be updated from {current_max_budget} to {max_budget_amount}")
317310

318-
# Only check reset timestamp if we need to update
311+
# Rule 2: If budget_duration is None, always update it
312+
if current_budget_duration is None:
313+
data["budget_duration"] = f"{renewal_period_days}d"
314+
needs_update = True
315+
logger.info(f"Key {key.id} budget update triggered by None budget_duration")
316+
317+
# Rule 2.5: If budget_duration is "0d", always update it (fix for expired keys)
318+
elif current_budget_duration == "0d":
319+
data["budget_duration"] = f"{renewal_period_days}d"
320+
needs_update = True
321+
logger.info(f"Key {key.id} budget update triggered by 0d duration (expired key fix)")
322+
323+
# Rule 3: If the duration mismatches, and the next reset time is within an hour of (now + current duration) then reset the duration
324+
else:
325+
expected_budget_duration = f"{renewal_period_days}d"
326+
if current_budget_duration != expected_budget_duration:
327+
# Check if reset time alignment heuristic is met
328+
should_update_duration = False
329+
330+
if budget_reset_at:
331+
try:
332+
# Parse the current budget duration to get the time delta
333+
# Support formats like: 10m, 2h, 30d, etc.
334+
duration_match = re.match(r'(\d+)([mhd])', current_budget_duration)
335+
if duration_match:
336+
duration_value = int(duration_match.group(1))
337+
duration_unit = duration_match.group(2)
338+
339+
# Calculate the time delta based on the unit
340+
if duration_unit == 'm':
341+
current_period_end = current_time + timedelta(minutes=duration_value)
342+
elif duration_unit == 'h':
343+
current_period_end = current_time + timedelta(hours=duration_value)
344+
elif duration_unit == 'd':
345+
current_period_end = current_time + timedelta(days=duration_value)
346+
else:
347+
# Unknown unit, skip duration update
348+
logger.warning(f"Unknown duration unit '{duration_unit}' for key {key.id}, skipping duration update")
349+
continue
350+
351+
# Parse the budget reset time
352+
reset_time = datetime.fromisoformat(budget_reset_at.replace('Z', '+00:00'))
353+
354+
# Check if reset time is within 1 hour of current period end
355+
time_diff = abs((reset_time - current_period_end).total_seconds())
356+
if time_diff <= 3600: # 1 hour in seconds
357+
should_update_duration = True
358+
logger.info(f"Key {key.id} duration update triggered by reset time alignment: reset at {reset_time}, period ends at {current_period_end}")
359+
except (ValueError, AttributeError) as e:
360+
logger.warning(f"Error parsing budget reset time for key {key.id}: {str(e)}")
361+
362+
if should_update_duration:
363+
data["budget_duration"] = expected_budget_duration
364+
needs_update = True
365+
366+
# Update when needs_update is True
319367
if needs_update:
320-
budget_reset_at_str = info.get("budget_reset_at")
321-
if budget_reset_at_str:
322-
try:
323-
# Parse the budget_reset_at timestamp
324-
budget_reset_at = datetime.fromisoformat(budget_reset_at_str.replace('Z', '+00:00'))
325-
if budget_reset_at.tzinfo is None:
326-
budget_reset_at = budget_reset_at.replace(tzinfo=UTC)
327-
logger.info(f"Key {key.id} budget_reset_at_str: {budget_reset_at_str}, budget_reset_at: {budget_reset_at}")
328-
329-
# Check if budget was reset recently using heuristics
330-
# budget_reset_at represents when the next reset will occur
331-
current_spend = info.get("spend", 0) or 0.0
332-
current_budget_duration = info.get("budget_duration")
333-
334-
should_update = False
335-
update_reason = ""
336-
337-
# Heuristic 1: Check if (now + current_budget_duration) is within an hour of budget_reset_at
338-
if current_budget_duration is not None:
339-
try:
340-
# Parse current budget duration (e.g., "30d" -> 30 days)
341-
duration_match = re.match(r'(\d+)d', current_budget_duration)
342-
if duration_match:
343-
duration_days = int(duration_match.group(1))
344-
expected_reset_time = current_time + timedelta(days=duration_days)
345-
hours_diff = abs((expected_reset_time - budget_reset_at).total_seconds() / 3600)
346-
347-
if hours_diff <= 1.0:
348-
should_update = True
349-
update_reason = f"reset time alignment (within {hours_diff:.2f} hours)"
350-
except (ValueError, AttributeError):
351-
logger.warning(f"Key {key.id} has invalid budget_duration format: {current_budget_duration}")
352-
else:
353-
logger.debug(f"Key {key.id} has no budget_duration set, skipping reset time alignment heuristic")
354-
should_update = True
355-
update_reason = "no budget_duration set, forcing update"
356-
357-
# Heuristic 2: Update if amount spent is $0.00 (indicating fresh reset)
358-
if current_spend == 0.0:
359-
should_update = True
360-
update_reason = "zero spend (fresh reset)"
361-
362-
if should_update:
363-
logger.info(f"Key {key.id} budget update triggered: {update_reason}, updating budget settings")
364-
await litellm_service.update_budget(**data)
365-
logger.info(f"Updated key {key.id} budget settings")
366-
else:
367-
logger.debug(f"Key {key.id} budget update not triggered, skipping update")
368-
except ValueError:
369-
logger.warning(f"Key {key.id} has invalid budget_reset_at timestamp: {budget_reset_at_str}")
370-
else:
371-
logger.warning(f"Key {key.id} has no budget_reset_at timestamp, forcing update")
372-
await litellm_service.update_budget(**data)
368+
# Determine what the new budget_duration will be for logging
369+
new_budget_duration = data.get("budget_duration", current_budget_duration)
370+
logger.info(f"Key {key.id} budget update triggered: changing from {current_budget_duration}, {current_max_budget} to {new_budget_duration}, {max_budget_amount}")
371+
await litellm_service.update_budget(**data)
372+
logger.info(f"Updated key {key.id} budget settings")
373373
else:
374374
logger.info(f"Key {key.id} budget settings already match the expected values, no update needed")
375375

376376
# Set the key duration to 0 days to end its usability.
377377
if expire_keys:
378+
logger.info(f"Key {key.id} expiring, setting duration to 0 days")
378379
await litellm_service.update_key_duration(key.litellm_token, "0d")
379380

380381
# Add to team total

0 commit comments

Comments
 (0)