[iris] Fix scheduling loop: filter reservation jobs at SQL level (#4179)

Add has_reservation column to jobs table so
_claim_workers_for_reservations
can filter at SQL level instead of deserializing all active job
protobufs.
On production DB (1,340 active jobs, 218MB of request_proto): p50 drops
from 156ms to 0.0ms, p95 from 2,669ms to 0.1ms. Adds before/after
comparison to the benchmark script.

Author: rjpower

lib/iris/scripts/benchmark_db_queries.py

@@ -30,6 +30,7 @@
 from iris.cluster.controller.controller import (
     _building_counts,
     _jobs_by_id,
+    _jobs_with_reservations,
     _schedulable_tasks,
 )
 from iris.cluster.controller.db import (
@@ -112,6 +113,30 @@ def benchmark_scheduling(db: ControllerDB, iterations: int) -> list[tuple[str, f
     else:
         print("  _jobs_by_id                                  (skipped, no pending jobs)")
 
+    # Reservation queries: compare old (fetch all + filter) vs new (SQL filter)
+    reservable_states = (
+        cluster_pb2.JOB_STATE_PENDING,
+        cluster_pb2.JOB_STATE_BUILDING,
+        cluster_pb2.JOB_STATE_RUNNING,
+    )
+
+    def _reservation_jobs_old():
+        with db.snapshot() as snapshot:
+            all_jobs = snapshot.select(JOBS, where=JOBS.c.state.in_(list(reservable_states)))
+        return [j for j in all_jobs if j.request.HasField("reservation")]
+
+    p50, p95 = bench("reservation_jobs (old: full scan)", _reservation_jobs_old, iterations=iterations)
+    results.append(("reservation_jobs (old: full scan)", p50, p95))
+    print_result("reservation_jobs (old: full scan)", p50, p95)
+
+    p50, p95 = bench(
+        "reservation_jobs (new: has_reservation)",
+        lambda: _jobs_with_reservations(db, reservable_states),
+        iterations=iterations,
+    )
+    results.append(("reservation_jobs (new: has_reservation)", p50, p95))
+    print_result("reservation_jobs (new: has_reservation)", p50, p95)
+
     return results
 
 

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

@@ -42,6 +42,7 @@
     Job,
     Task,
     Worker,
+    _decode_row,
     _tasks_with_attempts,
     healthy_active_workers_with_attributes,
     insert_task_profile,
@@ -284,6 +285,21 @@ def _jobs_by_id(queries: ControllerDB, job_ids: set[JobName]) -> dict[JobName, J
     return {job.job_id: job for job in jobs}
 
 
+def _jobs_with_reservations(queries: ControllerDB, states: tuple[int, ...]) -> list[Job]:
+    """Fetch only jobs that have reservations, filtering at the SQL level.
+
+    Uses the denormalized has_reservation column to avoid deserializing
+    request_proto for all active jobs.
+    """
+    placeholders = ",".join("?" for _ in states)
+    with queries.snapshot() as snapshot:
+        rows = snapshot._fetchall(
+            f"SELECT * FROM jobs WHERE state IN ({placeholders}) AND has_reservation = 1",
+            list(states),
+        )
+    return [_decode_row(Job, row) for row in rows]
+
+
 def _schedulable_tasks(queries: ControllerDB) -> list[Task]:
     # Only PENDING tasks can pass can_be_scheduled(); no need to fetch ASSIGNED/BUILDING/RUNNING.
     SCHEDULABLE_STATES = (cluster_pb2.TASK_STATE_PENDING,)
@@ -1181,11 +1197,8 @@ def _claim_workers_for_reservations(self, claims: dict[WorkerId, ReservationClai
             cluster_pb2.JOB_STATE_BUILDING,
             cluster_pb2.JOB_STATE_RUNNING,
         )
-        with self._db.snapshot() as snapshot:
-            reservable_jobs = snapshot.select(JOBS, where=JOBS.c.state.in_(list(reservable_states)))
-        for job in reservable_jobs:
-            if not job.request.HasField("reservation"):
-                continue
+        reservation_jobs = _jobs_with_reservations(self._db, reservable_states)
+        for job in reservation_jobs:
 
             job_wire = job.job_id.to_wire()
             for idx, res_entry in enumerate(job.request.reservation.entries):

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

@@ -579,6 +579,7 @@ class Job:
     exit_code: int | None = db_field("exit_code", _nullable(_decode_int))
     num_tasks: int = db_field("num_tasks", _decode_int)
     is_reservation_holder: bool = db_field("is_reservation_holder", _decode_bool_int)
+    has_reservation: bool = db_field("has_reservation", _decode_bool_int, default=False)
     name: str = db_field("name", _decode_str, default="")
     depth: int = db_field("depth", _decode_int, default=0)
 

lib/iris/src/iris/cluster/controller/migrations/0013_has_reservation.py

@@ -0,0 +1,33 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+import sqlite3
+
+
+def _has_column(conn: sqlite3.Connection, table: str, column: str) -> bool:
+    columns = {row[1] for row in conn.execute(f"PRAGMA table_info({table})").fetchall()}
+    return column in columns
+
+
+def migrate(conn: sqlite3.Connection) -> None:
+    if not _has_column(conn, "jobs", "has_reservation"):
+        conn.execute("ALTER TABLE jobs ADD COLUMN has_reservation INTEGER NOT NULL DEFAULT 0")
+    conn.execute(
+        "CREATE INDEX IF NOT EXISTS idx_jobs_has_reservation "
+        "ON jobs(has_reservation, state) WHERE has_reservation = 1"
+    )
+
+    # Backfill: scan all rows once and mark only those with reservations.
+    # Skip if any rows already have has_reservation=1 (migration already ran).
+    already_backfilled = conn.execute("SELECT 1 FROM jobs WHERE has_reservation = 1 LIMIT 1").fetchone()
+    if already_backfilled:
+        return
+
+    from iris.rpc import cluster_pb2
+
+    rows = conn.execute("SELECT job_id, request_proto FROM jobs WHERE request_proto IS NOT NULL").fetchall()
+    for row in rows:
+        proto = cluster_pb2.Controller.LaunchJobRequest()
+        proto.ParseFromString(row[1])
+        if proto.HasField("reservation") and proto.reservation.entries:
+            conn.execute("UPDATE jobs SET has_reservation = 1 WHERE job_id = ?", (row[0],))

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

@@ -596,12 +596,13 @@ def submit_job(
 
             state = cluster_pb2.JOB_STATE_PENDING if validation_error is None else cluster_pb2.JOB_STATE_FAILED
             finished_ms = None if validation_error is None else effective_submission_ms
+            has_reservation = 1 if request.HasField("reservation") and request.reservation.entries else 0
             cur.execute(
                 "INSERT INTO jobs("
                 "job_id, user_id, parent_job_id, root_job_id, depth, request_proto, state, submitted_at_ms, "
                 "root_submitted_at_ms, started_at_ms, finished_at_ms, scheduling_deadline_epoch_ms, "
-                "error, exit_code, num_tasks, is_reservation_holder, name"
-                ") VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, ?, NULL, ?, 0, ?)",
+                "error, exit_code, num_tasks, is_reservation_holder, has_reservation, name"
+                ") VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, ?, NULL, ?, 0, ?, ?)",
                 (
                     job_id.to_wire(),
                     job_id.user,
@@ -616,6 +617,7 @@ def submit_job(
                     deadline_epoch_ms,
                     validation_error,
                     replicas,
+                    has_reservation,
                     request.name,
                 ),
             )