[iris] Duplicate user budget tiers into cluster config

Migration 0037 is UPDATE-only and only acts on pre-existing user_budgets
rows, so a DB wipe or a listed user who hasn't submitted a job yet
lands on UserBudgetDefaults (budget_limit=0, max_band=BATCH) instead
of their intended tier. Move the tier membership into cluster config
(IrisClusterConfig.user_budgets) so it is:

  - Version-controlled next to the rest of the cluster config
  - Reconciled into the DB on every controller start via an upsert
  - Durable across DB wipes (migration is the one-shot fixup for the
    already-deployed prod DB; config is the source of truth going forward)

New proto message UserBudgetTier groups users sharing a budget_limit
and max_band; config entries accept real PriorityBand names
(PRIORITY_BAND_PRODUCTION etc.) in YAML.

Co-authored-by: William Held <Helw150@users.noreply.github.com>

Author: github-actions[bot]

lib/iris/examples/marin.yaml

@@ -176,3 +176,39 @@ scale_groups:
         service_account: iris-worker@hai-gcp-models.iam.gserviceaccount.com
         mode: GCP_SLICE_MODE_VM
         machine_type: n2-highmem-2
+
+# ---------------------------------------------------------------------------
+# User budget tiers
+#
+# Reconciled into user_budgets on every controller start. Mirrors the one-shot
+# fixup in migrations/0037_user_budget_default.py so fresh DBs (and listed
+# users who haven't yet submitted a job) land on the intended tier instead of
+# UserBudgetDefaults (budget_limit=0, max_band=BATCH).
+#
+# Unlisted submitters stay on the default BATCH cap; see
+# docs/priority-bands.md → "Max-band caps" for the recovery path when an
+# unlisted user hits PERMISSION_DENIED on an INTERACTIVE submission.
+# ---------------------------------------------------------------------------
+user_budgets:
+  # Admins: standard budget, may submit PRIORITY_BAND_PRODUCTION for work that
+  # must not be downgraded. INTERACTIVE work still counts against the budget.
+  - user_ids: [runner, power, dlwh, rav, romain, held, larry]
+    budget_limit: 75000
+    max_band: PRIORITY_BAND_PRODUCTION
+  # Researchers: standard budget, capped at PRIORITY_BAND_INTERACTIVE.
+  - user_ids:
+      - ruili
+      - quevedo
+      - pc0618
+      - konwoo
+      - ahmedah
+      - ahmed
+      - rohith
+      - tim
+      - eczech
+      - tonyhlee
+      - kevin
+      - calvinxu
+      - moojink
+    budget_limit: 75000
+    max_band: PRIORITY_BAND_INTERACTIVE

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

@@ -5,15 +5,20 @@
 
 from __future__ import annotations
 
+import json
+import logging
 from collections import defaultdict
+from collections.abc import Iterable
 from dataclasses import dataclass
 from typing import Generic, TypeVar
 
-import json
+from rigging.timing import Timestamp
 
-from iris.cluster.controller.db import ACTIVE_TASK_STATES, QuerySnapshot
+from iris.cluster.controller.db import ACTIVE_TASK_STATES, ControllerDB, QuerySnapshot
 from iris.cluster.types import JobName
-from iris.rpc import job_pb2
+from iris.rpc import config_pb2, job_pb2
+
+logger = logging.getLogger(__name__)
 
 T = TypeVar("T")
 
@@ -139,3 +144,59 @@ def interleave_by_user(
             break
         round_idx += 1
     return result
+
+
+# Bands accepted in user_budgets config entries. UNSPECIFIED is kept out of the
+# set so a missing/zeroed max_band field surfaces as a config error rather than
+# silently granting BATCH; callers must pick a real band.
+_VALID_TIER_BANDS = frozenset(
+    (
+        job_pb2.PRIORITY_BAND_PRODUCTION,
+        job_pb2.PRIORITY_BAND_INTERACTIVE,
+        job_pb2.PRIORITY_BAND_BATCH,
+    )
+)
+
+
+def reconcile_user_budget_tiers(
+    db: ControllerDB,
+    tiers: Iterable[config_pb2.UserBudgetTier],
+    now: Timestamp,
+) -> int:
+    """Upsert per-user budgets from cluster config into the user_budgets table.
+
+    Runs at controller startup after auth is resolved. Each tier entry lists
+    a set of user_ids that all receive the same budget_limit and max_band.
+    Tiers are applied in order, so later tiers override earlier ones for
+    users listed in both — lets ops promote a user by appending a later tier
+    without editing earlier ones.
+
+    This complements migration 0037 (which fixes prod DBs that already have
+    rows) by handling fresh DBs and listed users who haven't submitted yet:
+    those users would otherwise land on the :class:`UserBudgetDefaults` row
+    created via INSERT OR IGNORE at first submission time.
+
+    Returns the number of (user_id, tier) pairs applied; duplicate user_ids
+    across tiers are counted per-apply since the later tier overwrites.
+    """
+    count = 0
+    for tier in tiers:
+        if tier.max_band not in _VALID_TIER_BANDS:
+            raise ValueError(
+                f"UserBudgetTier.max_band must be one of PRODUCTION/INTERACTIVE/BATCH; "
+                f"got {tier.max_band} for users {list(tier.user_ids)}"
+            )
+        for user_id in tier.user_ids:
+            if not user_id:
+                raise ValueError("UserBudgetTier.user_ids contains an empty entry")
+            db.ensure_user(user_id, now)
+            db.set_user_budget(
+                user_id=user_id,
+                budget_limit=tier.budget_limit,
+                max_band=tier.max_band,
+                now=now,
+            )
+            count += 1
+    if count:
+        logger.info("Reconciled %d user budget assignment(s) from cluster config", count)
+    return count

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

@@ -23,14 +23,15 @@
 import click
 
 from iris.cluster.controller.auth import ControllerAuth, create_controller_auth
+from iris.cluster.controller.budget import reconcile_user_budget_tiers
 from iris.cluster.controller.controller import Controller, ControllerConfig
 from iris.cluster.providers.types import port_is_open, resolve_external_host
 from iris.log_server.main import (
     AUTH_STRICT_ENV_VAR as LOG_SERVER_AUTH_STRICT_ENV_VAR,
     JWT_KEY_ENV_VAR as LOG_SERVER_JWT_KEY_ENV_VAR,
 )
 from iris.rpc import config_pb2
-from rigging.timing import Duration, ExponentialBackoff
+from rigging.timing import Duration, ExponentialBackoff, Timestamp
 
 logger = logging.getLogger(__name__)
 
@@ -220,6 +221,14 @@ def run_controller_serve(
     if auth.worker_token and base_worker_config is not None:
         base_worker_config.auth_token = auth.worker_token
 
+    # Reconcile per-user budget tiers from the cluster config into the DB.
+    # Runs after auth so users listed in admin_users already have rows; the
+    # upsert here adds/updates budget_limit and max_band for everyone in
+    # cluster_config.user_budgets. Unlisted users still fall through to
+    # UserBudgetDefaults at first-submit time.
+    if cluster_config and cluster_config.user_budgets:
+        reconcile_user_budget_tiers(db, cluster_config.user_budgets, Timestamp.now())
+
     # --- Start log server subprocess ---
     log_port = port + LOG_SERVER_PORT_OFFSET
     log_dir = local_state_dir / "logs"

lib/iris/src/iris/rpc/config.proto

@@ -16,6 +16,7 @@ edition = "2023";
 
 package iris.config;
 
+import "job.proto";
 import "time.proto";
 
 // Enable explicit presence tracking for all fields to support HasField on scalars
@@ -445,6 +446,28 @@ message KubernetesProviderConfig {
   string controller_address = 8; // Controller URL injected into task pods
 }
 
+// ============================================================================
+// USER BUDGETS
+// ============================================================================
+
+// Pre-seeds a row in the user_budgets table for each listed user at controller
+// startup. This is the source of truth for tier membership — the controller
+// upserts these into the DB on every start, so a DB wipe or a listed user
+// who hasn't yet submitted a job still lands on the intended budget + max
+// band. Tiers are applied in order; later tiers override earlier tiers for
+// the same user_id.
+message UserBudgetTier {
+  // Users assigned this tier. A user_id must match the authenticated identity
+  // the submitter presents (see AuthConfig); a mismatch silently drops the
+  // user to the INSERT OR IGNORE default row created at first job submission.
+  repeated string user_ids = 1;
+  // Max resource-value spend before tasks are downgraded to BATCH. 0 = unlimited.
+  int64 budget_limit = 2;
+  // Highest priority band this user is allowed to submit to. Submissions above
+  // this band are rejected with PERMISSION_DENIED.
+  iris.job.PriorityBand max_band = 3;
+}
+
 // ============================================================================
 // ROOT CLUSTER CONFIGURATION
 // ============================================================================
@@ -473,4 +496,9 @@ message IrisClusterConfig {
     KubernetesProviderConfig kubernetes_provider = 70;
     WorkerProviderConfig worker_provider = 71;
   }
+
+  // Per-user budget/max-band policy reconciled into the DB at every controller
+  // start. Unlisted users fall through to the default in UserBudgetDefaults
+  // (set at first job submission via INSERT OR IGNORE).
+  repeated UserBudgetTier user_budgets = 80;
 }