fix: add missing steps for standard engine replica version

.gitignore

@@ -10,3 +10,4 @@ plugins/dags/*
 tests/logs
 tests/airflow.cfg
 tests/airflow.db
+.serena/*

dags/rollout_ic_os_to_subnets.py

@@ -16,13 +16,16 @@
 from airflow.models.baseoperator import chain
 from airflow.models.param import Param
 from airflow.operators.empty import EmptyOperator
+from airflow.sensors.base import PokeReturnValue
 from airflow.utils.task_group import TaskGroup
+from dfinity.dre import FAKE_PROPOSAL_NUMBER
 from dfinity.ic_os_rollout import (
     MAX_BATCHES,
     SubnetIdWithRevision,
     SubnetRolloutPlanWithRevision,
 )
 from dfinity.ic_types import IC_NETWORKS
+from dfinity.rollout_types import ProposalInfo
 
 from airflow import DAG, __version__
 
@@ -298,13 +301,47 @@ def collect_step(step_index: int, **kwargs: Any) -> list[float]:
             )
             proceed >> join
 
+            @task.sensor(
+                task_id="wait_until_proposal_is_accepted",
+                retries=retries,
+                poke_interval=120,
+                timeout=86400 * 7,
+                mode="reschedule",
+            )
+            def wait_until_proposal_is_accepted(
+                proposal_info: ProposalInfo, **kwargs: Any
+            ) -> PokeReturnValue:
+                # An idempotent no-op (the engine is already at or above the
+                # target deployment_progress) yields a fake proposal that never
+                # needs a vote, so there is nothing to wait for.
+                if (
+                    proposal_info["proposal_id"] == FAKE_PROPOSAL_NUMBER
+                    or not proposal_info["needs_vote"]
+                ):
+                    print(
+                        "No real proposal to wait for; the standard engine is"
+                        " already at or above the target deployment progress."
+                    )
+                    return PokeReturnValue(is_done=True)
+                simulate = cast(bool, kwargs["params"]["simulate"])
+                return PokeReturnValue(
+                    is_done=ic_os_sensor.has_proposal_executed(
+                        proposal_info, network, simulate
+                    )
+                )
+
+            @task(trigger_rule="none_failed_min_one_success")
+            def engine_subnet_ids(engines: list[dict[str, Any]]) -> list[str]:
+                return [e["subnet_id"] for e in engines]
+
             # Collect the engines upgraded in this step so we can monitor their
-            # alerts.  This runs after the proposal so the version has been
-            # elected/voted and the engines have begun upgrading.  It is not
-            # mapped: it reads its own and the previous step's target progress
-            # from the schedule XCom, and returns a flat list of subnet IDs.  It
-            # runs unconditionally (like join) so an empty step still lets the
-            # alert wait (and the chain) proceed.
+            # alerts.  This runs after the proposal has been accepted so the
+            # version has been elected/voted and the engines have begun
+            # upgrading.  It is not mapped: it reads its own and the previous
+            # step's target progress from the schedule XCom, and returns, per
+            # engine subnet, its subnet ID and current replica count.  It runs
+            # unconditionally (like join) so an empty step still lets the
+            # downstream waits (and the chain) proceed.
             upgraded_engines = ic_os_rollout.CollectStandardEngineUpgradedSubnets(
                 task_id="collect_upgraded_engines",
                 git_revision="{{ params.git_revision }}",
@@ -314,6 +351,16 @@ def collect_step(step_index: int, **kwargs: Any) -> list[float]:
                 trigger_rule="none_failed_min_one_success",
             )
 
+            create_proposal = (
+                ic_os_rollout.CreateStandardEngineProposalIdempotently.partial(
+                    task_id="create_proposal_if_none_exists",
+                    git_revision="{{ params.git_revision }}",
+                    simulate_proposal=cast(bool, "{{ params.simulate }}"),
+                    retries=retries,
+                    network=network,
+                ).expand(deployment_progress=proceed)
+            )
+
             (
                 ic_os_sensor.CustomDateTimeSensorAsync.partial(
                     task_id="wait_until_start_time",
@@ -325,26 +372,31 @@ def collect_step(step_index: int, **kwargs: Any) -> list[float]:
                     % step_index,
                     simulate="{{ params.simulate }}",
                 ).expand(_ignored=proceed)
-                >> ic_os_rollout.CreateStandardEngineProposalIdempotently.partial(
-                    task_id="create_proposal_if_none_exists",
-                    git_revision="{{ params.git_revision }}",
-                    simulate_proposal=cast(bool, "{{ params.simulate }}"),
-                    retries=retries,
-                    network=network,
-                ).expand(deployment_progress=proceed)
+                >> create_proposal
                 >> ic_os_rollout.RequestProposalVote.partial(
                     task_id="request_proposal_vote",
                     source_task_id="standard_engine.step_%d.create_proposal_if_none_exists"
                     % step_index,
                     retries=retries,
                 ).expand(_ignored=proceed)
+                >> wait_until_proposal_is_accepted.expand(
+                    proposal_info=create_proposal.output
+                )
                 >> upgraded_engines
+                >> ic_os_sensor.WaitForReplicaRevisionUpdated.partial(
+                    task_id="wait_for_replica_revision",
+                    git_revision="{{ params.git_revision }}",
+                    retries=retries,
+                    network=network,
+                ).expand_kwargs(upgraded_engines.output)
                 >> ic_os_sensor.WaitUntilNoAlertsOnSubnet.partial(
                     task_id="wait_until_no_alerts",
                     git_revision="{{ params.git_revision }}",
                     retries=retries,
                     network=network,
-                ).expand(subnet_id=upgraded_engines.output)
+                ).expand(
+                    subnet_id=engine_subnet_ids(upgraded_engines.output)  # type: ignore
+                )
                 >> join
             )
 

plugins/operators/ic_os_rollout.py

@@ -627,6 +627,12 @@ class CollectStandardEngineUpgradedSubnets(BaseOperator):
     The returned (flat) list is intended to be used to expand alert-monitoring
     tasks so we only watch the engines that actually changed version in this step.
 
+    Each returned item is a mapping with the engine's `subnet_id` and its
+    `expected_replica_count` (the current number of replicas on that subnet,
+    queried from Prometheus).  The count lets a downstream
+    `WaitForReplicaRevisionUpdated` task (mapped via `expand_kwargs`) know how
+    many nodes must reach the new revision.
+
     The step's own target and the previous step's target are read from the
     `standard_engine_schedule` XCom using this step's index, so this task does
     not need to be dynamically mapped.  Returns an empty list when the step is
@@ -657,7 +663,7 @@ def __init__(
         self.network = network
         BaseOperator.__init__(self, task_id=task_id, **kwargs)
 
-    def execute(self, context: Context) -> list[str]:
+    def execute(self, context: Context) -> list[dict[str, Any]]:
         _, git_revision = subnet_id_and_git_revision_from_args("", self.git_revision)
 
         plan = cast(
@@ -688,7 +694,36 @@ def execute(self, context: Context) -> list[str]:
             git_revision,
             subnets,
         )
-        return subnets
+
+        # For each engine subnet, remember the current number of replicas so a
+        # downstream WaitForReplicaRevisionUpdated (mapped via expand_kwargs)
+        # knows how many nodes must reach the new revision.
+        result: list[dict[str, Any]] = []
+        for subnet_id in subnets:
+            try:
+                replica_count = int(
+                    prom.query_prometheus_servers(
+                        self.network.prometheus_urls,
+                        f'sum(ic_replica_info{{ic_subnet="{subnet_id}"}})'
+                        " by (ic_subnet)",
+                    )[0]["value"]
+                )
+            except IndexError:
+                raise RuntimeError(
+                    f"No replicas have been found with subnet {subnet_id}"
+                )
+            self.log.info(
+                "Engine subnet %s currently has %s replicas.",
+                subnet_id,
+                replica_count,
+            )
+            result.append(
+                {
+                    "subnet_id": subnet_id,
+                    "expected_replica_count": replica_count,
+                }
+            )
+        return result
 
 
 def create_api_boundary_nodes_proposal_if_none_exists(

rollout-dashboard/frontend/src/lib/types.ts

@@ -97,6 +97,14 @@ const GuestOsStandardEngineStepState = {
     pending: { icon: "🕐", name: "pending" },
     waiting: { icon: "⌛", name: "waiting" },
     proposing: { icon: "📝", name: "proposing deployment progress increase" },
+    waiting_for_acceptance: {
+        icon: "🗳️",
+        name: "waiting for the proposal to be accepted",
+    },
+    waiting_for_adoption: {
+        icon: "⏳",
+        name: "waiting for the engines to adopt the new revision",
+    },
     waiting_for_alerts_gone: {
         icon: "📢",
         name: "monitoring upgraded engines until no more alerts",

rollout-dashboard/server/src/live_state/guestos_rollout.rs

@@ -719,7 +719,13 @@ impl Parser {
                                 "create_proposal_if_none_exists" => {
                                     se_trans_min!(StandardEngineStepState::Proposing)
                                 }
-                                "collect_upgraded_engines" | "wait_until_no_alerts" => {
+                                "wait_until_proposal_is_accepted" => {
+                                    se_trans_min!(StandardEngineStepState::WaitingForAcceptance)
+                                }
+                                "collect_upgraded_engines" | "wait_for_replica_revision" => {
+                                    se_trans_min!(StandardEngineStepState::WaitingForAdoption)
+                                }
+                                "wait_until_no_alerts" => {
                                     se_trans_min!(StandardEngineStepState::WaitingForAlertsGone)
                                 }
                                 "request_proposal_vote" | "join" => {}
@@ -737,9 +743,15 @@ impl Parser {
                                     se_trans_exact!(StandardEngineStepState::Proposing);
                                 }
                                 "create_proposal_if_none_exists" => {
-                                    se_trans_exact!(StandardEngineStepState::WaitingForAlertsGone)
+                                    se_trans_exact!(StandardEngineStepState::WaitingForAcceptance)
+                                }
+                                "wait_until_proposal_is_accepted" => {
+                                    se_trans_exact!(StandardEngineStepState::WaitingForAdoption)
                                 }
                                 "collect_upgraded_engines" => {
+                                    se_trans_exact!(StandardEngineStepState::WaitingForAdoption)
+                                }
+                                "wait_for_replica_revision" => {
                                     se_trans_exact!(StandardEngineStepState::WaitingForAlertsGone)
                                 }
                                 "wait_until_no_alerts" => {

rollout-dashboard/server/src/types.rs

@@ -329,8 +329,14 @@ pub mod v2 {
             Waiting,
             /// The step is submitting the proposal that raises deployment_progress.
             Proposing,
-            /// The proposal has been submitted; the step is monitoring the
-            /// upgraded engines until their alerts subside.
+            /// The proposal has been submitted; the step is waiting for it to be
+            /// voted in and executed.
+            WaitingForAcceptance,
+            /// The proposal has executed; the step is waiting for the upgraded
+            /// engines to adopt the new replica revision.
+            WaitingForAdoption,
+            /// The upgraded engines have adopted the new revision; the step is
+            /// monitoring them until their alerts subside.
             WaitingForAlertsGone,
             /// The proposal has been submitted and the upgraded engines are
             /// healthy; the step is complete.