Cleanups.

Author: rjpower

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

@@ -75,6 +75,7 @@
 from iris.cluster.controller.provider import TaskProvider
 from iris.cluster.controller.reads import SchedulableWorker
 from iris.cluster.controller.scheduler import (
+    DEFAULT_MAX_ASSIGNMENTS_PER_WORKER,
     DEFAULT_MAX_BUILDING_TASKS_PER_WORKER,
     JobRequirements,
     Scheduler,
@@ -328,13 +329,24 @@ def compute_demand_entries(
         dry_run_workers = _inject_reservation_taints(snapshots, claims)
         dry_run_jobs = _inject_taint_constraints(jobs, has_reservation, has_direct_reservation)
 
-        context = scheduler.create_scheduling_context(
-            dry_run_workers,
+        # Dry-run scheduling context — only the per-(task, worker) matching loop
+        # consumes capacities/jobs/pending_tasks, so the raw-read fields stay
+        # empty. Building/assignment limits are disabled so big workers can
+        # absorb multiple tasks (prevents false demand on idle clusters).
+        context = SchedulingContext(
+            workers=dry_run_workers,
             building_counts=building_counts,
-            pending_tasks=task_ids,
-            jobs=dry_run_jobs,
             max_building_tasks=_UNLIMITED,
             max_assignments_per_worker=_UNLIMITED,
+            pending_tasks=task_ids,
+            jobs=dry_run_jobs,
+            pending_task_rows=[],
+            user_spend={},
+            user_budget_limits={},
+            requested_bands={},
+            reserved_job_ids=frozenset(),
+            reservation_entry_counts={},
+            user_budget_defaults=UserBudgetDefaults(),
         )
         result = scheduler.find_assignments(context)
         for task_id, _ in result.assignments:
@@ -1032,10 +1044,13 @@ def build_scheduling_context(
 
     snapshots = [worker_snapshot_from_row(w, usage_by_worker.get(w.worker_id)) for w in workers]
     sorted_pending = _sort_pending_tasks_by_resolved_band(pending, requested_bands)
-    return SchedulingContext.from_workers(
-        snapshots,
+    return SchedulingContext(
+        workers=snapshots,
         building_counts=building_counts,
         max_building_tasks=max_building_tasks,
+        max_assignments_per_worker=DEFAULT_MAX_ASSIGNMENTS_PER_WORKER,
+        pending_tasks=[],
+        jobs={},
         pending_task_rows=sorted_pending,
         user_spend=user_spend,
         user_budget_limits=user_budget_limits,
@@ -1468,6 +1483,12 @@ def __init__(
         self._scheduling_diagnostics: dict[str, str] = {}
         self._scheduling_round: int = 0
 
+        # Last completed scheduling context — None until the first tick runs.
+        # The dashboard diagnostics path reads this instead of rebuilding from
+        # the DB. This is the only ``| None`` attribute on Controller: it is
+        # genuinely None before the first scheduling tick has run.
+        self._last_scheduling_context: SchedulingContext | None = None
+
         # Set to True once start() is called. Used to gate operations that
         # are only valid before the controller loops begin (e.g. LoadCheckpoint).
         self._started = False
@@ -1974,6 +1995,7 @@ def _run_scheduling(self) -> SchedulingOutcome:
 
         if not ctx.pending_task_rows:
             self._scheduling_diagnostics = {}
+            self._last_scheduling_context = ctx
             return SchedulingOutcome.NO_PENDING_TASKS
 
         gated = apply_scheduling_gates(
@@ -1989,6 +2011,7 @@ def _run_scheduling(self) -> SchedulingOutcome:
 
         if not gated.schedulable_task_ids:
             self._scheduling_diagnostics = {}
+            self._last_scheduling_context = ctx
             return SchedulingOutcome.NO_PENDING_TASKS
 
         order = compute_scheduling_order(ctx, gated, trace=trace)
@@ -1998,6 +2021,9 @@ def _run_scheduling(self) -> SchedulingOutcome:
         preemptions = self._apply_preemptions(order, tainted_jobs, all_assignments, claims, context)
 
         self._cache_scheduling_diagnostics(context, tainted_jobs, all_assignments, order.ordered_task_ids)
+        # Post-taint context (or the un-tainted ctx when no claims were active)
+        # — exposed via ``last_scheduling_context`` for dashboard diagnostics.
+        self._last_scheduling_context = context
 
         if all_assignments or preemptions:
             log_event(
@@ -2048,11 +2074,12 @@ def _run_scheduler_pass(
         if claims:
             modified_workers = _inject_reservation_taints(list(ctx.workers), claims)
             building_counts = {wid: cap.building_task_count for wid, cap in ctx.capacities.items()}
-            context = self._scheduler.create_scheduling_context(
-                modified_workers,
-                building_counts=building_counts,
-                pending_tasks=order.ordered_task_ids,
+            ctx.pending_tasks = list(order.ordered_task_ids)
+            context = ctx.evolve_with_workers(
+                workers=modified_workers,
                 jobs=modified_jobs,
+                building_counts=building_counts,
+                max_building_tasks=self._scheduler.max_building_tasks_per_worker,
             )
         else:
             ctx.pending_tasks = list(order.ordered_task_ids)
@@ -2274,16 +2301,16 @@ def _mark_task_unschedulable(self, task: Any) -> None:
                 reason=f"Scheduling timeout exceeded ({timeout})",
             )
 
-    def create_scheduling_context(self, workers: list[SchedulableWorker]) -> SchedulingContext:
-        """Create a worker-only scheduling context for diagnostics and dashboard RPCs."""
-        with self._db.read_snapshot() as snap:
-            building_counts = reads.building_counts(snap, [w.worker_id for w in workers])
-            usage_by_worker = reads.resource_usage_by_worker(snap)
-        snapshots = [worker_snapshot_from_row(w, usage_by_worker.get(w.worker_id)) for w in workers]
-        return self._scheduler.create_scheduling_context(
-            snapshots,
-            building_counts=building_counts,
-        )
+    @property
+    def last_scheduling_context(self) -> "SchedulingContext | None":
+        """Return the most recent finalized scheduling context.
+
+        ``None`` before the first scheduling tick has run; otherwise the
+        post-taint context from the last completed ``_run_scheduling`` pass.
+        Consumed by dashboard diagnostics that need a snapshot of capacities
+        and pending tasks without rebuilding from the DB.
+        """
+        return self._last_scheduling_context
 
     # =========================================================================
     # Worker lifecycle RPC dispatch (Reconcile / Ping)

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

@@ -1074,8 +1074,8 @@ class SchedulableWorker:
     """Worker shape consumed by the scheduler.
 
     Field names mirror the :class:`scheduler.WorkerSnapshot` protocol so
-    instances flow into ``Scheduler.create_scheduling_context`` without
-    an adapter.
+    instances flow through ``worker_snapshot_from_row`` into
+    ``SchedulingContext`` without an adapter.
     """
 
     worker_id: WorkerId

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

@@ -11,8 +11,8 @@
 The scheduler operates exclusively on scheduler-owned types (JobRequirements,
 WorkerCapacity, SchedulingContext) and has ZERO runtime imports from controller
 state. Callers project worker rows into ``WorkerSnapshot`` (via
-``worker_snapshot_from_row``) at the boundary before invoking
-``create_scheduling_context``.
+``worker_snapshot_from_row``) at the boundary before constructing
+``SchedulingContext`` directly.
 
 """
 
@@ -328,100 +328,96 @@ class SchedulingContext:
     gates/order pipeline and derived structures (``index``, ``capacities``) for
     the per-(task, worker) matching hot loop.
 
+    Construction is direct: callers supply ``workers``, ``building_counts``,
+    ``max_building_tasks``, and the raw-read fields; ``__post_init__`` derives
+    ``capacities``, ``index``, and ``_str_to_wid`` once. To rebuild the index
+    after taint injection mid-tick, use :meth:`evolve_with_workers` which reuses
+    the raw-read fields and only redoes the per-worker derivation.
+
     Posting lists are read-only after construction; capacity deductions don't
     touch them. Workers are tracked via ``assignment_counts`` to bound tasks
     per worker per cycle.
     """
 
-    index: ConstraintIndex
-
-    # Worker capacities indexed by worker ID
-    capacities: dict[WorkerId, WorkerCapacity]
-
-    # Reverse map from string ID back to WorkerId
-    _str_to_wid: dict[str, WorkerId]
-
-    assignment_counts: dict[WorkerId, int] = field(default_factory=dict)
-    max_assignments_per_worker: int = DEFAULT_MAX_ASSIGNMENTS_PER_WORKER
-
-    # Task IDs in scheduling priority order; populated after gates+order resolve.
-    pending_tasks: list[JobName] = field(default_factory=list)
-    jobs: dict[JobName, JobRequirements] = field(default_factory=dict)
-
-    # Raw per-tick reads — consumed by gates/order helpers; empty for diagnostics/dry-run.
-    pending_task_rows: list[PendingTask] = field(default_factory=list)
-    workers: list[WorkerSnapshot] = field(default_factory=list)
-    user_spend: dict[str, int] = field(default_factory=dict)
-    user_budget_limits: dict[str, int] = field(default_factory=dict)
-    requested_bands: dict[JobName, int] = field(default_factory=dict)
-    reserved_job_ids: frozenset[JobName] = field(default_factory=frozenset)
-    reservation_entry_counts: dict[JobName, int] = field(default_factory=dict)
-    user_budget_defaults: UserBudgetDefaults | None = None
+    workers: list[WorkerSnapshot]
+    building_counts: dict[WorkerId, int]
+    max_building_tasks: int
+    max_assignments_per_worker: int
+    pending_tasks: list[JobName]
+    jobs: dict[JobName, JobRequirements]
+    pending_task_rows: list[PendingTask]
+    user_spend: dict[str, int]
+    user_budget_limits: dict[str, int]
+    requested_bands: dict[JobName, int]
+    reserved_job_ids: frozenset[JobName]
+    reservation_entry_counts: dict[JobName, int]
+    user_budget_defaults: UserBudgetDefaults
+
+    # Derived from ``workers`` in __post_init__.
+    capacities: dict[WorkerId, WorkerCapacity] = field(init=False)
+    index: ConstraintIndex = field(init=False)
+    _str_to_wid: dict[str, WorkerId] = field(init=False)
+
+    # Per-cycle mutable state; always starts empty.
+    assignment_counts: dict[WorkerId, int] = field(init=False)
 
     # Scores memoized per (worker, soft-constraints) tuple; worker attributes are
     # stable within a tick so the same pair always yields the same score.
-    _soft_score_cache: dict[tuple[WorkerId, tuple[Constraint, ...]], int] = field(default_factory=dict)
-
-    @property
-    def all_worker_ids(self) -> set[WorkerId]:
-        return {self._str_to_wid[s] for s in self.index._all_ids}
+    _soft_score_cache: dict[tuple[WorkerId, tuple[Constraint, ...]], int] = field(init=False)
 
-    @classmethod
-    def from_workers(
-        cls,
-        workers: list[WorkerSnapshot],
-        building_counts: dict[WorkerId, int] | None = None,
-        max_building_tasks: int = DEFAULT_MAX_BUILDING_TASKS_PER_WORKER,
-        pending_tasks: list[JobName] | None = None,
-        jobs: dict[JobName, JobRequirements] | None = None,
-        max_assignments_per_worker: int = DEFAULT_MAX_ASSIGNMENTS_PER_WORKER,
-        pending_task_rows: list[PendingTask] | None = None,
-        user_spend: dict[str, int] | None = None,
-        user_budget_limits: dict[str, int] | None = None,
-        requested_bands: dict[JobName, int] | None = None,
-        reserved_job_ids: frozenset[JobName] | None = None,
-        reservation_entry_counts: dict[JobName, int] | None = None,
-        user_budget_defaults: UserBudgetDefaults | None = None,
-    ) -> "SchedulingContext":
-        """Build scheduling context from worker list.
-
-        Creates capacity snapshots and a ConstraintIndex for fast attribute matching.
-        """
-        building_counts = building_counts or {}
-
-        capacities = {
+    def __post_init__(self) -> None:
+        self.capacities = {
             w.worker_id: WorkerCapacity.from_worker(
                 w,
-                building_count=building_counts.get(w.worker_id, 0),
-                max_building_tasks=max_building_tasks,
+                building_count=self.building_counts.get(w.worker_id, 0),
+                max_building_tasks=self.max_building_tasks,
             )
-            for w in workers
+            for w in self.workers
         }
-
         str_to_wid: dict[str, WorkerId] = {}
         entity_attrs: dict[str, dict[str, AttributeValue]] = {}
-        for wid, cap in capacities.items():
+        for wid, cap in self.capacities.items():
             key = str(wid)
             str_to_wid[key] = wid
             entity_attrs[key] = dict(cap.attributes)
+        self._str_to_wid = str_to_wid
+        self.index = ConstraintIndex.build(entity_attrs)
+        self.assignment_counts = {}
+        self._soft_score_cache = {}
 
-        index = ConstraintIndex.build(entity_attrs)
-
-        return cls(
-            index=index,
-            capacities=capacities,
-            _str_to_wid=str_to_wid,
-            pending_tasks=pending_tasks or [],
-            jobs=jobs or {},
-            max_assignments_per_worker=max_assignments_per_worker,
-            pending_task_rows=list(pending_task_rows or []),
-            workers=list(workers),
-            user_spend=dict(user_spend or {}),
-            user_budget_limits=dict(user_budget_limits or {}),
-            requested_bands=dict(requested_bands or {}),
-            reserved_job_ids=frozenset(reserved_job_ids or ()),
-            reservation_entry_counts=dict(reservation_entry_counts or {}),
-            user_budget_defaults=user_budget_defaults,
+    @property
+    def all_worker_ids(self) -> set[WorkerId]:
+        return {self._str_to_wid[s] for s in self.index._all_ids}
+
+    def evolve_with_workers(
+        self,
+        workers: list[WorkerSnapshot],
+        jobs: dict[JobName, JobRequirements],
+        building_counts: dict[WorkerId, int],
+        max_building_tasks: int,
+    ) -> "SchedulingContext":
+        """Rebuild capacities/index for taint-injected workers.
+
+        Reuses all raw-read fields (``pending_task_rows``, ``user_spend``, etc.)
+        verbatim. The caller supplies updated ``workers``/``jobs`` (e.g. after
+        reservation taint injection) and fresh ``building_counts``. The
+        returned context starts a fresh placement pass with empty
+        ``assignment_counts`` and an empty soft-score cache.
+        """
+        return SchedulingContext(
+            workers=workers,
+            building_counts=building_counts,
+            max_building_tasks=max_building_tasks,
+            max_assignments_per_worker=self.max_assignments_per_worker,
+            pending_tasks=self.pending_tasks,
+            jobs=jobs,
+            pending_task_rows=self.pending_task_rows,
+            user_spend=self.user_spend,
+            user_budget_limits=self.user_budget_limits,
+            requested_bands=self.requested_bands,
+            reserved_job_ids=self.reserved_job_ids,
+            reservation_entry_counts=self.reservation_entry_counts,
+            user_budget_defaults=self.user_budget_defaults,
         )
 
     def matching_workers(self, constraints: Sequence[Constraint]) -> set[WorkerId]:
@@ -618,6 +614,11 @@ def __init__(
     ):
         self._max_building_tasks_per_worker = max_building_tasks_per_worker
 
+    @property
+    def max_building_tasks_per_worker(self) -> int:
+        """Per-worker BUILDING-state limit applied to fresh scheduling contexts."""
+        return self._max_building_tasks_per_worker
+
     def find_assignments(
         self,
         context: SchedulingContext,
@@ -816,35 +817,6 @@ def _group_soft_score(group_worker_ids: list[WorkerId]) -> int:
         )
         return None
 
-    def create_scheduling_context(
-        self,
-        workers: list[WorkerSnapshot],
-        building_counts: dict[WorkerId, int] | None = None,
-        pending_tasks: list[JobName] | None = None,
-        jobs: dict[JobName, JobRequirements] | None = None,
-        max_building_tasks: int | None = None,
-        max_assignments_per_worker: int | None = None,
-    ) -> SchedulingContext:
-        """Create a scheduling context for the given workers.
-
-        Convenience wrapper for tests, diagnostics, and the autoscaler dry-run
-        path. Does not populate the raw read fields (``pending_task_rows``,
-        ``user_spend``, etc.); use ``build_scheduling_context`` for the full
-        scheduling-loop construction.
-        """
-        limit = max_building_tasks if max_building_tasks is not None else self._max_building_tasks_per_worker
-        assignments_limit = (
-            max_assignments_per_worker if max_assignments_per_worker is not None else DEFAULT_MAX_ASSIGNMENTS_PER_WORKER
-        )
-        return SchedulingContext.from_workers(
-            workers,
-            building_counts=building_counts,
-            max_building_tasks=limit,
-            pending_tasks=pending_tasks,
-            jobs=jobs,
-            max_assignments_per_worker=assignments_limit,
-        )
-
     def get_job_scheduling_diagnostics(
         self,
         req: JobRequirements,