diff --git a/dags/hotfix_ic_os_to_subnets.py b/dags/hotfix_ic_os_to_subnets.py new file mode 100644 index 00000000..38e9c7e8 --- /dev/null +++ b/dags/hotfix_ic_os_to_subnets.py @@ -0,0 +1,217 @@ +""" +Hotfix rollout of IC OS to all subnets. + +Unlike the regular rollout DAG, this DAG does not follow a time-based +schedule. It rolls out to all subnets as fast as possible, with manual +approval gates for the Monday subnets (canary) and the NNS subnet +(tdb26, always last). After the canary, subnets are grouped into +batches of 3, matching the regular rollout's batch size. + +The subnet list is derived from DEFAULT_GUESTOS_ROLLOUT_PLANS — the same +subnets as the weekly rollout. +""" + +import datetime +import os +import sys +from typing import Any, cast + +import operators.ic_os_rollout as ic_os_rollout +import pendulum +import sensors.ic_os_rollout as ic_os_sensor +from airflow.decorators import task +from airflow.models.baseoperator import chain +from airflow.models.param import Param +from airflow.operators.empty import EmptyOperator +from airflow.utils.task_group import TaskGroup +from dfinity.ic_os_rollout import ( + HOTFIX_MAX_BATCHES, + SubnetIdWithRevision, + SubnetRolloutPlanWithRevision, + extract_subnets_from_plan, +) +from dfinity.ic_types import IC_NETWORKS + +from airflow import DAG + +# Temporarily add the DAGs folder to import defaults.py. +sys.path.append(os.path.dirname(__file__)) +try: + from defaults import DEFAULT_GUESTOS_ROLLOUT_PLANS +finally: + sys.path.pop() + + +DAGS: dict[str, DAG] = {} +for network_name, network in IC_NETWORKS.items(): + _canary_subnets, _batch_subnets = extract_subnets_from_plan( + DEFAULT_GUESTOS_ROLLOUT_PLANS[network_name] + ) + with DAG( + dag_id=f"rollout_ic_os_to_{network_name}_subnets_hotfix", + schedule=None, + start_date=pendulum.datetime(2020, 1, 1, tz="UTC"), + catchup=False, + dagrun_timeout=datetime.timedelta(days=3), + tags=["hotfix", "rollout", "DRE", "IC OS", "GuestOS"], + render_template_as_native_obj=True, + params={ + "git_revision": Param( + "0000000000000000000000000000000000000000", + type="string", + pattern="^[a-f0-9]{40}$", + title="Git revision", + description="Git revision of the IC OS GuestOS release to roll out;" + " the version must have been elected before but the rollout will" + " check for that.", + ), + "simulate": Param( + True, + type="boolean", + title="Simulate", + description="If enabled (the default), the update proposal will be" + " simulated but not created, and its acceptance will be simulated too.", + ), + }, + ) as dag: + DAGS[network_name] = dag + retries = int(86400 / 60 / 5) # one day worth of retries + + @task + def revisions(schedule, **context): # type: ignore + revs = set() + for batch in schedule.values(): + for instance in batch[1]: + revs.add(instance.git_revision) + return list(revs) + + def make_me_a_batch(group_name: str, batch: int) -> None: + @task + def collect_batch_subnets( + current_batch_index: int, **kwargs: Any + ) -> list[SubnetIdWithRevision]: + batch = cast( + SubnetRolloutPlanWithRevision, + kwargs["ti"].xcom_pull("hotfix_schedule"), + ).get(str(current_batch_index)) + if not batch: + print("This batch is empty.") + return [] + subnets = batch[1] + return [ + {"subnet_id": s.subnet_id, "git_revision": s.git_revision} + for s in subnets + ] + + proceed = collect_batch_subnets(batch) + + join = EmptyOperator( + task_id="join", + trigger_rule="none_failed_min_one_success", + ) + + # When proceed returns empty, all other tasks downstream + # from it, which use the expand() function, skip. + # But the join task must run unconditionally, else the + # downstream task (next batch) will be skipped, so we have + # to add an explicit linkage between proceed and join, + # such that join will always succeed instead of being skipped + # and therefore the next batch will run. + proceed >> join + + ( + ic_os_sensor.WaitForManualApproval.partial( + task_id="wait_for_manual_approval", + batch_index=batch, + slow_start_batches=len(_canary_subnets), + total_batches="""{{ + ti.xcom_pull(task_ids='hotfix_schedule') | length + }}""", + ).expand(_ignored=proceed) + >> ic_os_rollout.CreateSubnetUpdateProposalIdempotently.partial( + task_id="create_proposal_if_none_exists", + git_revision="{{ params.git_revision }}", + simulate_proposal=cast(bool, "{{ params.simulate }}"), + retries=retries, + network=network, + ).expand(subnet_id=proceed) + >> ( + ic_os_rollout.RequestProposalVote.partial( + task_id="request_proposal_vote", + source_task_id=f"{group_name}.create_proposal_if_none_exists", + retries=retries, + ).expand(_ignored=proceed), + ( + ic_os_sensor.WaitForProposalAcceptance.partial( + task_id="wait_until_proposal_is_accepted", + git_revision="{{ params.git_revision }}", + simulate_proposal_acceptance=cast( + bool, """{{ params.simulate }}""" + ), + retries=retries, + network=network, + ).expand(subnet_id=proceed) + ), + ) + >> ic_os_sensor.WaitForReplicaRevisionUpdated.partial( + task_id="wait_for_replica_revision", + git_revision="{{ params.git_revision }}", + retries=retries, + network=network, + expected_replica_count="""{{ + ti.xcom_pull( + task_ids='%(group_name)s.""" + % locals() + + """create_proposal_if_none_exists', + key='replica_count' + ) | int + }}""", + ).expand(subnet_id=proceed) + >> ic_os_sensor.WaitUntilNoAlertsOnSubnet.partial( + task_id="wait_until_no_alerts", + git_revision="{{ params.git_revision }}", + retries=retries, + network=network, + ).expand(subnet_id=proceed) + >> join + ) + + sched = ic_os_rollout.hotfix_schedule( + network, _canary_subnets, _batch_subnets + ) + revs = revisions(sched) + wait_for_election = ic_os_sensor.WaitForRevisionToBeElected.partial( + task_id="wait_for_revision_to_be_elected", + simulate_elected=cast(bool, "{{ params.simulate }}"), + network=network, + retries=retries, + ).expand(git_revision=revs) + + wait_for_other_rollouts = ic_os_sensor.WaitForOtherDAGs( + task_id="wait_for_other_rollouts" + ) + + upgrade_unassigned_nodes = ic_os_rollout.UpgradeUnassignedNodes( + task_id="upgrade_unassigned_nodes", + simulate=cast(bool, "{{ params.simulate }}"), + network=network, + retries=retries, + ) + + task_groups = [] + for batch in range(HOTFIX_MAX_BATCHES): + if batch < len(_canary_subnets): + group_name = f"canary_{batch + 1}" + else: + group_name = f"batch_{batch - len(_canary_subnets) + 1}" + with TaskGroup(group_id=group_name) as group: + make_me_a_batch(group_name, batch) + task_groups.append(group) + chain( + ( + wait_for_election, + wait_for_other_rollouts, + ), + *task_groups, + upgrade_unassigned_nodes, + ) diff --git a/plugins/operators/ic_os_rollout.py b/plugins/operators/ic_os_rollout.py index 97fbf6eb..63d7ff2f 100644 --- a/plugins/operators/ic_os_rollout.py +++ b/plugins/operators/ic_os_rollout.py @@ -25,6 +25,7 @@ SubnetRolloutPlanWithRevision, assign_default_revision, check_plan, + hotfix_planner, rollout_planner, subnet_id_and_git_revision_from_args, ) @@ -274,6 +275,31 @@ def __init__( ) +class NotifyManualApprovalNeeded(slack.SlackAPIPostOperator): + template_fields = ("text",) + + def __init__( + self, + batch_name: str, + _ignored: Any = None, + **kwargs: Any, + ) -> None: + dr_dre_slack_id = DR_DRE_SLACK_ID + text = ( + f"""Hotfix rollout batch {batch_name} needs""" + f""" manual approval. """ + """ please approve in the Airflow UI to proceed.""" + ) + slack.SlackAPIPostOperator.__init__( + self, + channel=SLACK_CHANNEL, + username="Airflow", + text=text, + slack_conn_id=SLACK_CONNECTION_ID, + **kwargs, + ) + + class UpgradeUnassignedNodes(BaseOperator): template_fields = ("simulate",) network: ic_types.ICNetwork @@ -350,6 +376,48 @@ def schedule( return plan +@task +def hotfix_schedule( + network: ic_types.ICNetwork, + canary_abbreviations: list[str], + batch_abbreviations: list[str], + **context: dict[str, Any], +) -> SubnetRolloutPlanWithRevision: + git_revision = "{:040}".format( + context["task"].render_template( # type: ignore + "{{ params.git_revision }}", + context, + ) + ) + subnet_list_source = dre.DRE( + network=network, + subprocess_hook=SubprocessHook(), + ).get_subnet_list + + plan = hotfix_planner( + canary_abbreviations, + batch_abbreviations, + subnet_list_source, + git_revision, + ) + + canary_count = len(canary_abbreviations) + for nstr, (_, members) in plan.items(): + batch_idx = int(nstr) + if batch_idx < canary_count: + label = f"Canary {batch_idx + 1}" + else: + label = f"Batch {batch_idx - canary_count + 1}" + print(f"{label}:") + for item in members: + print( + f" Subnet {item.subnet_id} ({item.subnet_num}) will be" + f" rolled out to git revision {item.git_revision}." + ) + + return plan + + def create_api_boundary_nodes_proposal_if_none_exists( api_boundary_node_ids: list[str], git_revision: str, diff --git a/plugins/sensors/ic_os_rollout.py b/plugins/sensors/ic_os_rollout.py index 1f0e29ec..870997e5 100644 --- a/plugins/sensors/ic_os_rollout.py +++ b/plugins/sensors/ic_os_rollout.py @@ -23,6 +23,7 @@ from airflow.utils.context import Context from airflow.utils.state import DagRunState from dfinity.ic_os_rollout import ( + DR_DRE_SLACK_ID, SLACK_CHANNEL, SLACK_CONNECTION_ID, subnet_id_and_git_revision_from_args, @@ -547,6 +548,91 @@ def remember_messaging() -> None: post(f"Alerts have subsided. Rollout of {subnet_id} can proceed.") +class WaitForManualApproval(BaseSensorOperator): + """Conditionally waits for manual approval in the Airflow UI. + + If the batch requires approval (first N batches or last batch), + sends a Slack notification and defers until manually marked as Success. + Otherwise, passes through immediately. + """ + + template_fields: Sequence[str] = ("total_batches",) + + def __init__( + self, + *, + task_id: str, + batch_index: int, + slow_start_batches: int, + total_batches: str | int, + _ignored: Any = None, + **kwargs: Any, + ) -> None: + BaseSensorOperator.__init__(self, task_id=task_id, **kwargs) + self.batch_index = batch_index + self.slow_start_batches = slow_start_batches + self.total_batches = total_batches + + def _needs_approval(self) -> bool: + total = int(self.total_batches) + return ( + self.batch_index < self.slow_start_batches + or self.batch_index == total - 1 + ) + + def execute(self, context: Context, event: Any = None) -> None: + if not self._needs_approval(): + self.log.info( + "Batch %d is not a canary batch. No manual approval needed.", + self.batch_index + 1, + ) + return + + if event is None: + # First execution — send Slack notification. + self.log.info( + "Sending Slack notification for batch %d.", + self.batch_index + 1, + ) + try: + from airflow.configuration import conf + + base_url = conf.get("webserver", "base_url") + dag_id = context["dag_run"].dag_id + run_id = context["dag_run"].run_id + dag_url = f"{base_url}/dags/{dag_id}/grid?dag_run_id={run_id}" + slack.SlackAPIPostOperator( + task_id="notify_manual_approval_needed", + channel=SLACK_CHANNEL, + username="Airflow", + text=( + f"Hotfix rollout batch {self.batch_index + 1} needs" + f" manual approval. " + f" please <{dag_url}|approve in the Airflow UI>" + " to proceed." + ), + slack_conn_id=SLACK_CONNECTION_ID, + ).execute(context=context) + except Exception: + self.log.warning( + "Failed to send Slack notification." + " Continuing without notification.", + exc_info=True, + ) + + self.log.info( + "MANUAL APPROVAL REQUIRED for batch %d.", + self.batch_index + 1, + ) + self.log.info( + "Mark this task as Success in the Airflow UI to proceed." + ) + self.defer( + trigger=TimeDeltaTrigger(datetime.timedelta(minutes=5)), + method_name="execute", + ) + + class WaitForOtherDAGs(BaseSensorOperator): source_dag_id: str | None diff --git a/shared/dfinity/ic_os_rollout.py b/shared/dfinity/ic_os_rollout.py index c7553a83..0413089f 100644 --- a/shared/dfinity/ic_os_rollout.py +++ b/shared/dfinity/ic_os_rollout.py @@ -4,7 +4,7 @@ import datetime import re -from typing import Callable, NotRequired, TypeAlias, TypedDict, cast +from typing import Any, Callable, NotRequired, TypeAlias, TypedDict, cast from dfinity.ic_types import ( SubnetRolloutInstance, @@ -26,6 +26,8 @@ SLACK_CONNECTION_ID = "slack.ic_os_rollout" DR_DRE_SLACK_ID = "S05GPUNS7EX" MAX_BATCHES: int = 30 +HOTFIX_MAX_BATCHES: int = 30 +NNS_SUBNET_PREFIX = "tdb26" # To be deleted when we upgrade to Airflow 2.11. PLAN_FORM = """ @@ -351,6 +353,134 @@ def check_plan(plan: SubnetRolloutPlanWithRevision) -> None: ) +def extract_subnets_from_plan( + plan_yaml: str, +) -> tuple[list[str], list[str]]: + """Extract canary and batch subnets from a rollout plan. + + Parses the weekly rollout plan YAML and returns Monday subnets + (canary) and all remaining subnets separately. + + Returns: + (canary_subnets, batch_subnets) where canary_subnets are the + Monday subnets and batch_subnets are all others. + """ + import yaml + + def _extract(spec: "list[Any] | dict[str, Any]") -> list[str]: + result: list[str] = [] + if isinstance(spec, list): + result.extend(str(s) for s in spec) + elif isinstance(spec, dict): + for s in spec.get("subnets", []): + result.append(str(s["subnet"]) if isinstance(s, dict) else str(s)) + return result + + plan = yaml.safe_load(plan_yaml) + canary_subnets: list[str] = [] + batch_subnets: list[str] = [] + + for day, hours in plan.items(): + for _time, spec in hours.items(): + if day == "Monday": + canary_subnets.extend(_extract(spec)) + else: + batch_subnets.extend(_extract(spec)) + + return canary_subnets, batch_subnets + + +def hotfix_planner( + canary_abbreviations: list[str], + batch_abbreviations: list[str], + subnet_list_source: Callable[[], list[str]], + git_revision: str, + subnets_per_batch: int = 3, +) -> SubnetRolloutPlanWithRevision: + """Plan a hotfix rollout: canary subnets then batches of N, NNS last. + + Canary subnets each get their own batch (for manual approval). + The remaining non-NNS subnets are grouped into batches of + `subnets_per_batch`. The NNS subnet (tdb26) is always in the + final batch alone. + + Args: + canary_abbreviations: Subnet ID prefixes for canary (one per batch). + batch_abbreviations: Subnet ID prefixes for remaining subnets. + subnet_list_source: Callable returning full subnet principal IDs. + git_revision: The git revision to deploy to all subnets. + subnets_per_batch: Subnets per batch after the canary. + + Raises: + ValueError: If the resulting batch count exceeds HOTFIX_MAX_BATCHES. + """ + subnet_list = subnet_list_source() + + def resolve(abbreviations: list[str]) -> list[tuple[int, str]]: + resolved: list[tuple[int, str]] = [] + for abbrev in abbreviations: + matches = [ + (i, s) for i, s in enumerate(subnet_list) if s.startswith(abbrev) + ] + if not matches: + raise ValueError( + f"subnet abbreviation {abbrev!r} not found in subnet list" + ) + if len(matches) > 1: + raise ValueError( + f"subnet abbreviation {abbrev!r} is ambiguous" + ) + resolved.append(matches[0]) + return resolved + + canary_resolved = resolve(canary_abbreviations) + batch_resolved = resolve(batch_abbreviations) + + # Separate NNS subnet from batch subnets. + nns = [(n, s) for n, s in batch_resolved if s.startswith(NNS_SUBNET_PREFIX)] + non_nns = [(n, s) for n, s in batch_resolved if not s.startswith(NNS_SUBNET_PREFIX)] + + # Build batches: canary (one per batch), then groups, then NNS. + batches: list[list[tuple[int, str]]] = [] + + for subnet in canary_resolved: + batches.append([subnet]) + + for i in range(0, len(non_nns), subnets_per_batch): + batches.append(non_nns[i : i + subnets_per_batch]) + + # NNS last. + if nns: + batches.append(nns) + + if len(batches) > HOTFIX_MAX_BATCHES: + raise ValueError( + f"Too many batches ({len(batches)}) for hotfix rollout." + f" Maximum supported is {HOTFIX_MAX_BATCHES}." + ) + + dummy_time = datetime.datetime( + 2000, 1, 1, tzinfo=datetime.timezone.utc + ) + + plan: SubnetRolloutPlanWithRevision = {} + for batch_index, members in enumerate(batches): + plan[str(batch_index)] = ( + dummy_time, + [ + SubnetRolloutInstanceWithRevision( + start_at=dummy_time, + subnet_num=subnet_num, + subnet_id=subnet_id, + git_revision=git_revision, + ) + for subnet_num, subnet_id in members + ], + ) + + return plan + + class TimetableSpec(TypedDict): resume_at: str suspend_at: str diff --git a/tests/test_hotfix_planner.py b/tests/test_hotfix_planner.py new file mode 100644 index 00000000..db4f2ae0 --- /dev/null +++ b/tests/test_hotfix_planner.py @@ -0,0 +1,190 @@ +import unittest + +from dfinity.ic_os_rollout import ( + NNS_SUBNET_PREFIX, + SubnetRolloutPlanWithRevision, + extract_subnets_from_plan, + hotfix_planner, +) + + +class TestExtractSubnetsFromPlan(unittest.TestCase): + def test_monday_subnets_are_canary(self) -> None: + plan = """ +Monday: + 9:00: [aaa, bbb] + 11:00: [ccc] +Tuesday: + 7:00: [ddd, eee, fff] +""" + canary, batch = extract_subnets_from_plan(plan) + assert canary == ["aaa", "bbb", "ccc"] + assert batch == ["ddd", "eee", "fff"] + + def test_monday_next_week_is_not_canary(self) -> None: + plan = """ +Monday: + 9:00: [aaa] +Tuesday: + 7:00: [bbb] +Monday next week: + 7:00: + subnets: [nns] + batch: 30 +""" + canary, batch = extract_subnets_from_plan(plan) + assert canary == ["aaa"] + assert batch == ["bbb", "nns"] + + def test_dict_subnet_spec(self) -> None: + plan = """ +Monday: + 9:00: + subnets: + - subnet: aaa + git_revision: 0123456789012345678901234567890123456789 +Tuesday: + 7:00: [bbb] +""" + canary, batch = extract_subnets_from_plan(plan) + assert canary == ["aaa"] + assert batch == ["bbb"] + + def test_empty_plan(self) -> None: + plan = """ +Monday: + 9:00: [] +""" + canary, batch = extract_subnets_from_plan(plan) + assert canary == [] + assert batch == [] + + def test_no_monday(self) -> None: + plan = """ +Tuesday: + 9:00: [aaa, bbb] +""" + canary, batch = extract_subnets_from_plan(plan) + assert canary == [] + assert batch == ["aaa", "bbb"] + + +class TestHotfixPlanner(unittest.TestCase): + NNS = f"{NNS_SUBNET_PREFIX}-full-principal-id" + + def fake_subnet_list(self) -> list[str]: + return [ + "aaa-full-principal-id", + "bbb-full-principal-id", + "ccc-full-principal-id", + "ddd-full-principal-id", + "eee-full-principal-id", + "fff-full-principal-id", + "ggg-full-principal-id", + "hhh-full-principal-id", + "iii-full-principal-id", + self.NNS, + ] + + def subnet_ids_per_batch(self, plan: SubnetRolloutPlanWithRevision) -> list[list[str]]: + """Extract subnet IDs grouped by batch in order.""" + result = [] + for nstr in sorted(plan.keys(), key=int): + _, members = plan[nstr] + result.append([m.subnet_id for m in members]) + return result + + def test_canary_are_single_batches(self) -> None: + plan = hotfix_planner( + canary_abbreviations=["aaa", "bbb"], + batch_abbreviations=["ccc", "ddd", "eee"], + subnet_list_source=self.fake_subnet_list, + git_revision="a" * 40, + ) + batches = self.subnet_ids_per_batch(plan) + # 2 canary (one each) + 1 batch of 3 + assert len(batches) == 3 + assert len(batches[0]) == 1 # canary 1 + assert len(batches[1]) == 1 # canary 2 + assert len(batches[2]) == 3 # batch of 3 + + def test_nns_is_always_last(self) -> None: + plan = hotfix_planner( + canary_abbreviations=["aaa"], + batch_abbreviations=["bbb", "ccc", NNS_SUBNET_PREFIX], + subnet_list_source=self.fake_subnet_list, + git_revision="a" * 40, + ) + batches = self.subnet_ids_per_batch(plan) + last_batch = batches[-1] + assert len(last_batch) == 1 + assert last_batch[0].startswith(NNS_SUBNET_PREFIX) + + def test_nns_not_grouped_with_others(self) -> None: + plan = hotfix_planner( + canary_abbreviations=["aaa"], + batch_abbreviations=["bbb", "ccc", NNS_SUBNET_PREFIX], + subnet_list_source=self.fake_subnet_list, + git_revision="a" * 40, + ) + batches = self.subnet_ids_per_batch(plan) + # canary(aaa) + batch(bbb,ccc) + nns + assert len(batches) == 3 + assert batches[0] == ["aaa-full-principal-id"] + assert batches[1] == ["bbb-full-principal-id", "ccc-full-principal-id"] + assert batches[2] == [self.NNS] + + def test_batches_of_three(self) -> None: + plan = hotfix_planner( + canary_abbreviations=["aaa"], + batch_abbreviations=[ + "bbb", "ccc", "ddd", "eee", "fff", "ggg", "hhh", + ], + subnet_list_source=self.fake_subnet_list, + git_revision="a" * 40, + subnets_per_batch=3, + ) + batches = self.subnet_ids_per_batch(plan) + # canary(1) + 3 batches(3,3,1) + assert len(batches) == 4 + assert len(batches[0]) == 1 # canary + assert len(batches[1]) == 3 + assert len(batches[2]) == 3 + assert len(batches[3]) == 1 # remainder + + def test_git_revision_applied_to_all(self) -> None: + rev = "b" * 40 + plan = hotfix_planner( + canary_abbreviations=["aaa"], + batch_abbreviations=["bbb"], + subnet_list_source=self.fake_subnet_list, + git_revision=rev, + ) + for _, (_, members) in plan.items(): + for m in members: + assert m.git_revision == rev + + def test_unknown_abbreviation_raises(self) -> None: + self.assertRaises( + ValueError, + lambda: hotfix_planner( + canary_abbreviations=["zzz"], + batch_abbreviations=[], + subnet_list_source=self.fake_subnet_list, + git_revision="a" * 40, + ), + ) + + def test_ambiguous_abbreviation_raises(self) -> None: + def ambiguous_list() -> list[str]: + return ["aaa-001", "aaa-002"] + + self.assertRaises( + ValueError, + lambda: hotfix_planner( + canary_abbreviations=["aaa"], + batch_abbreviations=[], + subnet_list_source=ambiguous_list, + git_revision="a" * 40, + ), + )