[iris] Grant OpenAthena principals full CoreWeave budgets

lib/iris/config/cw-rno2a.yaml

@@ -163,7 +163,7 @@ defaults:
     default_task_image: ghcr.io/marin-community/iris-task:latest
 
 user_budget_defaults:
-  budget_limit: 400000
+  budget_limit: 10688
   max_band: interactive
 
 scale_groups:
@@ -212,23 +212,30 @@ scale_groups:
 # 42752 is 4 h100-8x nodes (32 GPUs) at full node resources.
 user_budgets:
   - user_ids:
-      - benjaminfeuer
       - betsy
-      - dlwh
       - eczech
-      - held
       - larry
       - marin
-      - muchanem
-      - mwittmann
-      - power
-      - rav
-      - romain
       - wmoss
       - zack
     budget_limit: 42752
     max_band: PRIORITY_BAND_INTERACTIVE
-  # 342016 is 32 h100-8x nodes: half of this cluster's 64-node fleet.
-  - user_ids: [benjaminfeuer]
-    budget_limit: 342016
+  - user_ids:
+      - ben.feuer@openathena.ai
+      - benjaminfeuer
+      - david.hall@openathena.ai
+      - dlwh
+      - mark.muchane@openathena.ai
+      - muchanem
+      - matt.wittmann@openathena.ai
+      - mwittmann
+      - rafal.wojdyla@openathena.ai
+      - rav
+      - romain.yon@openathena.ai
+      - romain
+      - russell.power@openathena.ai
+      - power
+      - will.held@openathena.ai
+      - held
+    budget_limit: 0
     max_band: PRIORITY_BAND_INTERACTIVE

lib/iris/config/cw-us-east-02a.yaml

@@ -163,7 +163,7 @@ defaults:
     default_task_image: ghcr.io/marin-community/iris-task:latest
 
 user_budget_defaults:
-  budget_limit: 400000
+  budget_limit: 10688
   max_band: interactive
 
 scale_groups:
@@ -212,23 +212,30 @@ scale_groups:
 # 42752 is 4 h100-8x nodes (32 GPUs) at full node resources.
 user_budgets:
   - user_ids:
-      - benjaminfeuer
       - betsy
-      - dlwh
       - eczech
-      - held
       - larry
       - marin
-      - muchanem
-      - mwittmann
-      - power
-      - rav
-      - romain
       - wmoss
       - zack
     budget_limit: 42752
     max_band: PRIORITY_BAND_INTERACTIVE
-  # 171008 is 16 h100-8x nodes: half of this cluster's 32-node fleet.
-  - user_ids: [benjaminfeuer]
-    budget_limit: 171008
+  - user_ids:
+      - ben.feuer@openathena.ai
+      - benjaminfeuer
+      - david.hall@openathena.ai
+      - dlwh
+      - mark.muchane@openathena.ai
+      - muchanem
+      - matt.wittmann@openathena.ai
+      - mwittmann
+      - rafal.wojdyla@openathena.ai
+      - rav
+      - romain.yon@openathena.ai
+      - romain
+      - russell.power@openathena.ai
+      - power
+      - will.held@openathena.ai
+      - held
+    budget_limit: 0
     max_band: PRIORITY_BAND_INTERACTIVE

lib/iris/config/cw-us-east-08a.yaml

@@ -186,7 +186,7 @@ defaults:
     default_task_image: ghcr.io/marin-community/iris-task:latest
 
 user_budget_defaults:
-  budget_limit: 400000
+  budget_limit: 5680
   max_band: interactive
 
 scale_groups:
@@ -237,19 +237,30 @@ scale_groups:
 # 128000 is about one NVL72 rack: 18 gb200-4x nodes cost ~102,000.
 user_budgets:
   - user_ids:
-      - benjaminfeuer
       - betsy
-      - dlwh
       - eczech
-      - held
       - larry
       - marin
+      - wmoss
+      - zack
+    budget_limit: 128000
+    max_band: PRIORITY_BAND_INTERACTIVE
+  - user_ids:
+      - ben.feuer@openathena.ai
+      - benjaminfeuer
+      - david.hall@openathena.ai
+      - dlwh
+      - mark.muchane@openathena.ai
       - muchanem
+      - matt.wittmann@openathena.ai
       - mwittmann
-      - power
+      - rafal.wojdyla@openathena.ai
       - rav
+      - romain.yon@openathena.ai
       - romain
-      - wmoss
-      - zack
-    budget_limit: 128000
+      - russell.power@openathena.ai
+      - power
+      - will.held@openathena.ai
+      - held
+    budget_limit: 0
     max_band: PRIORITY_BAND_INTERACTIVE

lib/iris/config/cw-us-west-04a.yaml

@@ -115,7 +115,7 @@ defaults:
     default_task_image: ghcr.io/marin-community/iris-task:latest
 
 user_budget_defaults:
-  budget_limit: 400000
+  budget_limit: 10688
   max_band: interactive
 
 scale_groups:
@@ -160,3 +160,24 @@ scale_groups:
       coreweave:
         region: US-WEST-04A
         instance_type: gd-8xh100ib-i128
+
+user_budgets:
+  - user_ids:
+      - ben.feuer@openathena.ai
+      - benjaminfeuer
+      - david.hall@openathena.ai
+      - dlwh
+      - mark.muchane@openathena.ai
+      - muchanem
+      - matt.wittmann@openathena.ai
+      - mwittmann
+      - rafal.wojdyla@openathena.ai
+      - rav
+      - romain.yon@openathena.ai
+      - romain
+      - russell.power@openathena.ai
+      - power
+      - will.held@openathena.ai
+      - held
+    budget_limit: 0
+    max_band: PRIORITY_BAND_INTERACTIVE

lib/iris/docs/priority-bands.md

@@ -99,15 +99,23 @@ config at startup are:
   default budget; jobs run INTERACTIVE while within budget and degrade to
   BATCH once exceeded. SYSTEM and PRODUCTION submissions are rejected.
 
+Budget rows are keyed by the authenticated principal stored in
+`jobs.submitting_user`. IAP and JWT submissions use the verified email or
+service-account identity. The friendly owner in a job path remains independent:
+an admin may submit `/power/train` while Iris accounts the job to
+`russell.power@openathena.ai`. Child jobs inherit the root submitter, and a
+federated handoff carries the same principal to the receiving cluster. Trusted
+local submissions (`local_admin`) and rows with an empty `submitting_user` use the
+job-path owner as a fallback budget key.
+
 If a higher-band submission is rejected:
 
 1. **Use the appropriate lower band.** Ordinary research should run at
    INTERACTIVE or BATCH.
-2. **Check your username.** The `max_band` cap is keyed on the verified
-   identity the controller sees. If the username in the error message isn't
-   what you expect — e.g. it's an email local-part or an SSO id rather than
-   your GitHub handle — your identity probably doesn't match the `user_id`
-   listed in the cluster config, and you'll land on the default tier.
+2. **Check your budget identity.** The error reports the authenticated email or
+   service account used for budget lookup. Confirm that exact value appears in
+   the cluster config. Trusted local submissions instead use the nickname at
+   the start of the job path.
 3. **Request an uplift.** If your work needs INTERACTIVE budget headroom or an
    admin-only band, ping [@Helw150](https://github.com/Helw150).
 

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

@@ -15,7 +15,7 @@
 from iris.cluster.controller import reads, writes
 from iris.cluster.controller.codec import device_counts_from_json
 from iris.cluster.controller.db import ControllerDB, Tx
-from iris.cluster.types import UserBudgetDefaults
+from iris.cluster.types import LOCAL_ADMIN_SUBMITTER, JobName, UserBudgetDefaults
 from iris.rpc import job_pb2
 from iris.rpc.proto_display import ADMIN_PRIORITY_BAND_VALUES, PRIORITY_BAND_VALUES
 
@@ -30,6 +30,13 @@ class UserTask(Generic[T]):
     task: T
 
 
+def budget_user_id(job_id: JobName, submitting_user: str) -> str:
+    """Use the authenticated submitter, with job-owner fallback for local or empty identities."""
+    if submitting_user and submitting_user != LOCAL_ADMIN_SUBMITTER:
+        return submitting_user
+    return job_id.user
+
+
 def resource_value(cpu_millicores: int, memory_bytes: int, accelerator_count: int) -> int:
     """Compute a scalar resource value for budget tracking.
 
@@ -47,14 +54,13 @@ def compute_user_spend(tx: Tx) -> dict[str, int]:
     Sums ``resource_value * task_count`` per user over the active, non-BATCH
     task rows returned by :func:`reads.user_spend_rows`.
 
-    Returns ``{user_id: total_resource_value}`` for users with active tasks.
+    Returns ``{budget_user_id: total_resource_value}`` for principals with active tasks.
     """
     rows = reads.user_spend_rows(tx)
 
     spend: dict[str, int] = defaultdict(int)
     for row in rows:
-        # job_id is decoded by JobNameType to JobName
-        user_id = row.job_id.user
+        user_id = budget_user_id(row.job_id, str(row.submitting_user))
         cpu = row.res_cpu_millicores
         mem = row.res_memory_bytes
         counts = device_counts_from_json(row.res_device_json)

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

@@ -70,6 +70,7 @@
 from iris.cluster.federation.store import FederationDirection, HandoffState
 from iris.cluster.runtime.env import TASK_OUTPUT_FINALIZING_STATUS
 from iris.cluster.types import (
+    LOCAL_ADMIN_SUBMITTER,
     LOCAL_CLUSTER,
     TERMINAL_JOB_STATES,
     AttemptUid,
@@ -552,6 +553,15 @@ def get_job_state(tx: Tx, job_id: JobName) -> int | None:
     return int(row.state) if row is not None else None
 
 
+def get_job_submitting_user(tx: Tx, job_id: JobName) -> str | None:
+    """Return the authenticated submitter stored for ``job_id``, or None if absent."""
+    row = tx.execute(
+        select(jobs_table.c.submitting_user).where(jobs_table.c.job_id == bindparam("job_id")),
+        {"job_id": job_id},
+    ).first()
+    return str(row.submitting_user) if row is not None else None
+
+
 def find_prunable_job(tx: Tx, terminal_states: Iterable[int], before_ts: Timestamp) -> JobName | None:
     """Return one terminal *local* job finished before ``before_ts``, or None.
 
@@ -858,6 +868,7 @@ def _row_to_pending_task(row: Row) -> PendingTask:
     return PendingTask(
         task_id=row.task_id,
         job_id=row.job_id,
+        submitting_user=str(row.submitting_user),
         backend_id=str(row.backend_id),
         state=int(row.state),
         current_attempt_id=int(row.current_attempt_id),
@@ -886,6 +897,7 @@ def _row_to_pending_task(row: Row) -> PendingTask:
         *PENDING_TASK_COLS,
         # job columns (label job_state to avoid clash with tasks.state)
         jobs_table.c.state.label("job_state"),
+        jobs_table.c.submitting_user,
         jobs_table.c.scheduling_deadline_epoch_ms,
         # job_config columns
         job_config_table.c.scheduling_timeout_ms,
@@ -954,24 +966,29 @@ def running_task_band_rows(tx: Tx) -> Sequence[Row]:
 _USER_SPEND_STMT = (
     select(
         local_tasks.c.job_id,
+        jobs_table.c.submitting_user,
         job_config_table.c.res_cpu_millicores,
         job_config_table.c.res_memory_bytes,
         job_config_table.c.res_device_json,
         func.count().label("task_count"),
     )
-    .select_from(local_tasks.join(job_config_table, job_config_table.c.job_id == local_tasks.c.job_id))
+    .select_from(
+        local_tasks.join(jobs_table, jobs_table.c.job_id == local_tasks.c.job_id).join(
+            job_config_table, job_config_table.c.job_id == local_tasks.c.job_id
+        )
+    )
     .where(hint_rare_state(local_tasks.c.state.in_(bindparam("states", expanding=True))))
     .where(job_config_table.c.priority_band != job_pb2.PRIORITY_BAND_BATCH)
-    .group_by(local_tasks.c.job_id)
+    .group_by(local_tasks.c.job_id, jobs_table.c.submitting_user)
 )
 
 
 def user_spend_rows(tx: Tx) -> Sequence[Row]:
     """Return per-job resource rows for active, non-BATCH tasks (budget spend basis).
 
-    Each row carries ``(job_id, res_cpu_millicores, res_memory_bytes,
-    res_device_json, task_count)``. ``job_config.priority_band`` (the user's
-    requested band) drives the BATCH exclusion, not the stamped
+    Each row carries ``(job_id, submitting_user, res_cpu_millicores,
+    res_memory_bytes, res_device_json, task_count)``. ``job_config.priority_band``
+    (the user's requested band) drives the BATCH exclusion, not the stamped
     ``tasks.priority_band``, so scheduler-downgraded jobs still count.
     """
     return tx.execute(_USER_SPEND_STMT, {"states": list(ACTIVE_TASK_STATES)}).all()
@@ -1435,13 +1452,20 @@ def list_active_tasks_for_jobs(
     return result
 
 
-def count_active_tasks_for_user(tx: Tx, user_id: str) -> int:
-    """Return the number of non-terminal tasks across all jobs owned by ``user_id``."""
+def count_active_tasks_for_budget_user(tx: Tx, user_id: str) -> int:
+    """Return non-terminal local tasks attributed to the budget identity ``user_id``."""
+    budget_user = case(
+        (
+            jobs_table.c.submitting_user.in_(("", LOCAL_ADMIN_SUBMITTER)),
+            jobs_table.c.user_id,
+        ),
+        else_=jobs_table.c.submitting_user,
+    )
     return int(
         tx.execute(
             select(func.count())
             .select_from(local_tasks.join(jobs_table, jobs_table.c.job_id == local_tasks.c.job_id))
-            .where(jobs_table.c.user_id == bindparam("user_id"))
+            .where(budget_user == bindparam("user_id"))
             .where(local_tasks.c.state.in_(bindparam("states", expanding=True))),
             {"user_id": user_id, "states": list(NON_TERMINAL_TASK_STATES)},
         ).scalar()