Skip to content

Commit d708e77

Browse files
committed
[iris] Include default-budget users in scheduler state response
Users without an explicit user_budgets row inherit UserBudgetDefaults, but get_scheduler_state only emitted SchedulerUserBudget entries from the DB, so the dashboard rendered Spent/Limit/Utilization as '-' for anyone running under defaults. Synthesize default entries for users with active spend so the scheduler view matches the budget logic.
1 parent 95bb136 commit d708e77

2 files changed

Lines changed: 23 additions & 7 deletions

File tree

lib/iris/src/iris/cluster/controller/service.py

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2526,24 +2526,33 @@ def get_scheduler_state(
25262526
)
25272527

25282528
# --- User budgets for response ---
2529+
# Users without an explicit user_budgets row inherit UserBudgetDefaults;
2530+
# synthesize entries for any user with active spend so the dashboard
2531+
# renders their Spent/Limit/Utilization instead of '-'.
25292532
budget_protos: list[controller_pb2.Controller.SchedulerUserBudget] = []
2530-
for b in budgets:
2531-
spent = user_spend.get(b.user_id, 0)
2532-
utilization = (spent / b.budget_limit * 100.0) if b.budget_limit > 0 else 0.0
2533+
defaults = self._user_budget_defaults
2534+
seen_users = {b.user_id for b in budgets}
2535+
budget_rows: list[tuple[str, int, int]] = [(b.user_id, b.budget_limit, b.max_band) for b in budgets]
2536+
for uid in user_spend:
2537+
if uid not in seen_users:
2538+
budget_rows.append((uid, defaults.budget_limit, defaults.max_band))
2539+
for user_id, budget_limit, max_band in budget_rows:
2540+
spent = user_spend.get(user_id, 0)
2541+
utilization = (spent / budget_limit * 100.0) if budget_limit > 0 else 0.0
25332542
# Show effective band: use INTERACTIVE as the test band to see if user is downgraded
25342543
eff = compute_effective_band(
25352544
job_pb2.PRIORITY_BAND_INTERACTIVE,
2536-
b.user_id,
2545+
user_id,
25372546
user_spend,
25382547
budget_limits,
25392548
self._user_budget_defaults,
25402549
)
25412550
budget_protos.append(
25422551
controller_pb2.Controller.SchedulerUserBudget(
2543-
user_id=b.user_id,
2544-
budget_limit=b.budget_limit,
2552+
user_id=user_id,
2553+
budget_limit=budget_limit,
25452554
budget_spent=spent,
2546-
max_band=b.max_band,
2555+
max_band=max_band,
25472556
effective_band=eff,
25482557
utilization_percent=utilization,
25492558
)

lib/iris/tests/cluster/controller/test_service.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1264,5 +1264,12 @@ def test_get_scheduler_state_with_running_task(controller_service, state):
12641264
assert resp.total_running == 1
12651265
assert resp.running_tasks[0].job_id == job_id.to_wire()
12661266
assert resp.running_tasks[0].user_id == "alice"
1267+
# alice has no explicit user_budgets row but has an active task — the
1268+
# scheduler state must report her spend using UserBudgetDefaults so the
1269+
# dashboard renders Spent/Limit/Utilization instead of '-'.
1270+
alice_budget = next((b for b in resp.user_budgets if b.user_id == "alice"), None)
1271+
assert alice_budget is not None
1272+
assert alice_budget.budget_spent > 0
1273+
assert alice_budget.budget_limit == controller_service._user_budget_defaults.budget_limit
12671274
finally:
12681275
_verified_identity.reset(token)

0 commit comments

Comments
 (0)