scripts: add an abstraction for the args that depend on each other and validate them in one place

Author: guyf-starkware

scripts/prod/restart_all_nodes_together.py

@@ -7,13 +7,12 @@
 import urllib.request
 from update_config_and_restart_nodes_lib import (
     ApolloArgsParserBuilder,
+    NamespaceAndInstructionArgs,
     RestartStrategy,
     Service,
     get_configmap,
-    get_context_list_from_args,
     get_current_block_number,
     get_logs_explorer_url,
-    get_namespace_list_from_args,
     parse_config_from_yaml,
     print_colored,
     update_config_and_restart_nodes,
@@ -107,12 +106,8 @@ def main():
         "consensus_manager_config.immediate_active_height": next_block_number,
     }
 
-    namespace_list = get_namespace_list_from_args(args)
-    context_list = get_context_list_from_args(args)
-    if context_list is not None:
-        assert len(namespace_list) == len(
-            context_list
-        ), "namespace_list and context_list must have the same length"
+    namespace_list = NamespaceAndInstructionArgs.get_namespace_list_from_args(args)
+    context_list = NamespaceAndInstructionArgs.get_context_list_from_args(args)
 
     # Generate logs explorer URLs if needed
     post_restart_instructions = []
@@ -131,11 +126,13 @@ def main():
 
     update_config_and_restart_nodes(
         config_overrides,
-        namespace_list,
+        NamespaceAndInstructionArgs(
+            namespace_list,
+            context_list,
+            post_restart_instructions,
+        ),
         Service.Core,
-        context_list,
         RestartStrategy.ALL_AT_ONCE,
-        post_restart_instructions,
     )
 
 

scripts/prod/set_node_revert_mode.py

@@ -6,12 +6,11 @@
 import urllib.parse
 from update_config_and_restart_nodes_lib import (
     ApolloArgsParserBuilder,
+    NamespaceAndInstructionArgs,
     RestartStrategy,
     Service,
-    get_context_list_from_args,
     get_current_block_number,
     get_logs_explorer_url,
-    get_namespace_list_from_args,
     print_colored,
     print_error,
     update_config_and_restart_nodes,
@@ -46,20 +45,22 @@ def set_revert_mode(
     }
 
     post_restart_instructions = []
-    for namespace, context in zip(namespace_list, context_list):
+    for namespace, context in zip(namespace_list, context_list or [None] * len(namespace_list)):
         url = get_logs_explorer_url_for_revert(namespace, revert_up_to_block, project_name)
 
         post_restart_instructions.append(
             f"Please check logs and verify that revert has completed (both in the batcher and for sync). Logs URL: {url}"
         )
-
+    namespace_and_instruction_args = NamespaceAndInstructionArgs(
+        namespace_list,
+        context_list,
+        post_restart_instructions,
+    )
     update_config_and_restart_nodes(
         config_overrides,
-        namespace_list,
+        namespace_and_instruction_args,
         Service.Core,
-        context_list,
         restart_strategy,
-        post_restart_instructions,
     )
 
 
@@ -155,8 +156,8 @@ def main():
             print_error("Error: --revert-up-to-block (-b) cannot be set when disabling revert.")
             sys.exit(1)
 
-    namespace_list = get_namespace_list_from_args(args)
-    context_list = get_context_list_from_args(args)
+    namespace_list = NamespaceAndInstructionArgs.get_namespace_list_from_args(args)
+    context_list = NamespaceAndInstructionArgs.get_context_list_from_args(args)
 
     should_disable_revert = not args.revert_only
     if should_revert:

scripts/prod/update_config_and_restart_nodes.py

@@ -8,9 +8,8 @@
 from update_config_and_restart_nodes_lib import (
     ApolloArgsParserBuilder,
     Colors,
+    NamespaceAndInstructionArgs,
     Service,
-    get_context_list_from_args,
-    get_namespace_list_from_args,
     print_colored,
     print_error,
     update_config_and_restart_nodes,
@@ -153,11 +152,13 @@ def main():
 
     update_config_and_restart_nodes(
         config_overrides,
-        get_namespace_list_from_args(args),
+        NamespaceAndInstructionArgs(
+            NamespaceAndInstructionArgs.get_namespace_list_from_args(args),
+            NamespaceAndInstructionArgs.get_context_list_from_args(args),
+            None,
+        ),
         args.service,
-        get_context_list_from_args(args),
         args.restart_strategy,
-        None,
     )
 
 

scripts/prod/update_config_and_restart_nodes_lib.py

@@ -211,35 +211,6 @@ def validate_arguments(args: argparse.Namespace) -> None:
             sys.exit(1)
 
 
-def get_namespace_list_from_args(
-    args: argparse.Namespace,
-) -> list[str]:
-    """Get a list of namespaces based on the arguments"""
-    if args.namespace_list:
-        return args.namespace_list
-
-    return [
-        f"{args.namespace_prefix}-{i}"
-        for i in range(args.start_index, args.start_index + args.num_nodes)
-    ]
-
-
-def get_context_list_from_args(
-    args: argparse.Namespace,
-) -> list[str]:
-    """Get a list of contexts based on the arguments"""
-    if args.cluster_list:
-        return args.cluster_list
-
-    if args.cluster_prefix is None:
-        return None
-
-    return [
-        f"{args.cluster_prefix}-{i}"
-        for i in range(args.start_index, args.start_index + args.num_nodes)
-    ]
-
-
 def get_logs_explorer_url(
     query: str,
     project_name: Optional[str] = None,
@@ -511,45 +482,103 @@ def restart_pod(
             sys.exit(1)
 
 
+class NamespaceAndInstructionArgs:
+    def __init__(
+        self,
+        namespace_list: list[str],
+        cluster_list: Optional[list[str]],
+        instruction_list: Optional[list[str]] = None,
+    ):
+        assert (
+            namespace_list is not None and len(namespace_list) > 0
+        ), "Namespace list cannot be None or empty."
+        self.namespace_list = namespace_list
+        assert cluster_list is None or len(cluster_list) == len(
+            namespace_list
+        ), "cluster_list must have the same length as namespace_list"
+        self.cluster_list = cluster_list
+        assert instruction_list is None or len(namespace_list) == len(
+            instruction_list
+        ), "instruction_list must have the same length as namespace_list"
+        self.instruction_list = instruction_list
+
+    def size(self) -> int:
+        return len(self.namespace_list)
+
+    def get_namespace(self, index: int) -> str:
+        return self.namespace_list[index]
+
+    def get_cluster(self, index: int) -> Optional[str]:
+        return self.cluster_list[index] if self.cluster_list else None
+
+    def get_instruction(self, index: int) -> Optional[str]:
+        return self.instruction_list[index] if self.instruction_list else None
+
+    @staticmethod
+    def get_namespace_list_from_args(
+        args: argparse.Namespace,
+    ) -> list[str]:
+        """Get a list of namespaces based on the arguments"""
+        if args.namespace_list:
+            return args.namespace_list
+
+        return [
+            f"{args.namespace_prefix}-{i}"
+            for i in range(args.start_index, args.start_index + args.num_nodes)
+        ]
+
+    @staticmethod
+    def get_context_list_from_args(
+        args: argparse.Namespace,
+    ) -> list[str]:
+        """Get a list of contexts based on the arguments"""
+        if args.cluster_list:
+            return args.cluster_list
+
+        if args.cluster_prefix is None:
+            return None
+
+        return [
+            f"{args.cluster_prefix}-{i}"
+            for i in range(args.start_index, args.start_index + args.num_nodes)
+        ]
+
+
 def update_config_and_restart_nodes(
     config_overrides: dict[str, Any],
-    namespace_list: list[str],
+    namespace_and_instruction_args: NamespaceAndInstructionArgs,
     service: Service,
-    cluster_list: Optional[list[str]],
     restart_strategy: RestartStrategy,
-    # TODO(guy.f): Add more abstraction to restart strategy so that the instructions are abstracted as part of it.
-    post_restart_instructions: Optional[list[str]] = None,
 ) -> None:
     assert config_overrides is not None, "config_overrides must be provided"
-    assert namespace_list is not None and len(namespace_list) > 0, "namespaces must be provided"
-
-    if post_restart_instructions is not None:
-        assert len(post_restart_instructions) == len(
-            namespace_list
-        ), f"post_restart_instructions must have the same length as namespace_list. logs_explorer_urls: {len(post_restart_instructions)}, namespace_list: {len(namespace_list)}"
+    assert namespace_and_instruction_args.namespace_list is not None, "namespaces must be provided"
 
-    if not cluster_list:
+    if not namespace_and_instruction_args.cluster_list:
         print_colored(
             "cluster-prefix/cluster-list not provided. Assuming all nodes are on the current cluster",
             Colors.RED,
         )
-    else:
-        assert len(cluster_list) == len(
-            namespace_list
-        ), f"cluster_list must have the same number of values as namespace_list. cluster_list: {cluster_list}, namespace_list: {namespace_list}"
 
     # Store original and updated configs for all nodes
     configs = []
 
     # Process each node's configuration
-    for index, namespace in enumerate(namespace_list):
-        cluster = cluster_list[index] if cluster_list else None
+    for index in range(namespace_and_instruction_args.size()):
+        namespace = namespace_and_instruction_args.get_namespace(index)
+        cluster = namespace_and_instruction_args.get_cluster(index)
+
         print_colored(
             f"\nProcessing node for namespace {namespace} (cluster: {cluster if cluster else 'current cluster'})..."
         )
 
         # Get current config and normalize it (e.g. " vs ') to ensure not showing bogus diffs.
-        original_config = normalize_config(get_configmap(namespace, cluster, service))
+        original_config = normalize_config(
+            get_configmap(
+                namespace,
+                cluster,
+                service,
+            )
+        )
 
         # Update config
         updated_config = update_config_values(original_config, config_overrides)
@@ -568,19 +597,22 @@ def update_config_and_restart_nodes(
     print_colored("\nApplying configurations...")
     for index, config in enumerate(configs):
         print(f"Applying config {index}...")
-        apply_configmap(
-            config["updated"],
-            namespace_list[index],
-            index,
-            cluster_list[index] if cluster_list else None,
-        )
+        # apply_configmap(
+        #     config["updated"],
+        #     namespace_and_instruction_args.get_namespace(index),
+        #     index,
+        #     namespace_and_instruction_args.get_cluster(index),
+        # )
 
     if restart_strategy != RestartStrategy.NO_RESTART:
         for index, config in enumerate(configs):
-            restart_pod(
-                namespace_list[index], service, index, cluster_list[index] if cluster_list else None
-            )
-            instructions = post_restart_instructions[index] if post_restart_instructions else None
+            # restart_pod(
+            #     namespace_and_instruction_args.get_namespace(index),
+            #     service,
+            #     index,
+            #     namespace_and_instruction_args.get_cluster(index),
+            # )
+            instructions = namespace_and_instruction_args.get_instruction(index)
             print_colored(f"Restarted pod.\n{instructions if instructions else ''} ", Colors.YELLOW)
             if restart_strategy == RestartStrategy.ONE_BY_ONE:
                 # Don't ask in the case of the last job.