scripts: Add an abstraction for how to restart pods.

Author: guyf-starkware

scripts/prod/restart_all_nodes_together.py

@@ -10,6 +10,7 @@
     NamespaceAndInstructionArgs,
     RestartStrategy,
     Service,
+    ServiceRestarter,
     get_configmap,
     get_current_block_number,
     get_logs_explorer_url,
@@ -124,15 +125,22 @@ def main():
             f"Please check logs and verify that the node has proposed a block that was accepted. Logs URL: {url}"
         )
 
+    namespace_and_instruction_args = NamespaceAndInstructionArgs(
+        namespace_list,
+        context_list,
+        post_restart_instructions,
+    )
+    restarter = ServiceRestarter.from_restart_strategy(
+        RestartStrategy.ALL_AT_ONCE,
+        namespace_and_instruction_args,
+        Service.Core,
+    )
+
     update_config_and_restart_nodes(
         config_overrides,
-        NamespaceAndInstructionArgs(
-            namespace_list,
-            context_list,
-            post_restart_instructions,
-        ),
+        namespace_and_instruction_args,
         Service.Core,
-        RestartStrategy.ALL_AT_ONCE,
+        restarter,
     )
 
 

scripts/prod/set_node_revert_mode.py

@@ -9,6 +9,7 @@
     NamespaceAndInstructionArgs,
     RestartStrategy,
     Service,
+    ServiceRestarter,
     get_current_block_number,
     get_logs_explorer_url,
     print_colored,
@@ -51,16 +52,23 @@ def set_revert_mode(
         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,
     )
+    restarter = ServiceRestarter.from_restart_strategy(
+        restart_strategy,
+        namespace_and_instruction_args,
+        Service.Core,
+    )
+
     update_config_and_restart_nodes(
         config_overrides,
         namespace_and_instruction_args,
         Service.Core,
-        restart_strategy,
+        restarter,
     )
 
 

scripts/prod/update_config_and_restart_nodes.py

@@ -8,8 +8,10 @@
 from update_config_and_restart_nodes_lib import (
     ApolloArgsParserBuilder,
     Colors,
+    RestartStrategy,
     NamespaceAndInstructionArgs,
     Service,
+    ServiceRestarter,
     print_colored,
     print_error,
     update_config_and_restart_nodes,
@@ -69,7 +71,7 @@ def service_type_converter(service_name: str) -> Service:
     if service_name.startswith("Service."):
         service_name = service_name[8:]
 
-    service_name = service_name.lower()
+    # service_name = service_name.lower()
 
     try:
         return Service[service_name]
@@ -150,15 +152,20 @@ def main():
         print_error("No config overrides provided")
         sys.exit(1)
 
-    update_config_and_restart_nodes(
-        config_overrides,
-        NamespaceAndInstructionArgs(
-            NamespaceAndInstructionArgs.get_namespace_list_from_args(args),
-            NamespaceAndInstructionArgs.get_context_list_from_args(args),
-            None,
-        ),
-        args.service,
+    namespace_and_instruction_args = NamespaceAndInstructionArgs(
+        NamespaceAndInstructionArgs.get_namespace_list_from_args(args),
+        NamespaceAndInstructionArgs.get_context_list_from_args(args),
+        None,
+    )
+
+    restarter = ServiceRestarter.from_restart_strategy(
         args.restart_strategy,
+        namespace_and_instruction_args,
+        args.service,
+    )
+
+    update_config_and_restart_nodes(
+        config_overrides, namespace_and_instruction_args, args.service, restarter
     )