update script to launch timekeeper first (#346)

* update script to launch timekeeper first * fix POT for bootstrap node * add option to disable timekeeper
autonomys · Oct 7, 2024 · 8bbd1de · 8bbd1de
1 parent 8e30fd3
commit 8bbd1de
Showing 1 changed file with 33 additions and 46 deletions.
diff --git a/scripts/launch-nodes/manage_subspace.py b/scripts/launch-nodes/manage_subspace.py
@@ -44,7 +44,7 @@ def run_command(client, command):
 
         # Treat Docker status updates as INFO instead of ERROR
         if error:
-            if any(keyword in error for keyword in ["Stopping", "Stopped", "Creating", "Started", "Removing", "Removed"]):
+            if any (keyword in error for keyword in ["Stopping", "Stopped", "Creating", "Started", "Removing", "Removed"]):
                 logger.info(f"Command output: {error.strip()}")
             else:
                 logger.error(f"Error running command: {error.strip()}")
@@ -138,8 +138,9 @@ def main():
     parser.add_argument('--config', required=True, help='Path to the TOML config file')
     parser.add_argument('--release_version', required=True, help='Release version to update in the .env file')
     parser.add_argument('--subspace_dir', default='/home/ubuntu/subspace', help='Path to the Subspace directory (default: /home/ubuntu/subspace)')
-    parser.add_argument('--pot_external_entropy', help='POT_EXTERNAL_ENTROPY value for the timekeeper node')
+    parser.add_argument('--pot_external_entropy', help='POT_EXTERNAL_ENTROPY value for all nodes')
     parser.add_argument('--log_level', default='INFO', help='Set the logging level (DEBUG, INFO, WARNING, ERROR)')
+    parser.add_argument('--no-timekeeper', action='store_true', help='Disable launching of the timekeeper node')
     args = parser.parse_args()
 
     # Set logging level based on user input
@@ -159,35 +160,48 @@ def main():
     release_version = args.release_version
     subspace_dir = args.subspace_dir
 
-    # Step 1: sudo docker compose down -v on all farmer and RPC nodes
-    for node in farmer_rpc_nodes:
+    # Step 1: Handle the timekeeper node first, if present and --no-timekeeper is not set
+    if not args.no_timekeeper and timekeeper_node:
         client = None  # Initialize the client variable
         try:
-            logger.info(f"Connecting to {node['host']} for sudo docker compose down -v...")
-            client = ssh_connect(node['host'], node['user'], node['ssh_key'])
+            logger.info(f"Connecting to the timekeeper node {timekeeper_node['host']}...")
+            client = ssh_connect(timekeeper_node['host'], timekeeper_node['user'], timekeeper_node['ssh_key'])
 
-            # Run sudo docker compose down -v
+            # Run sudo docker compose down -v for the timekeeper node
             docker_compose_down(client, subspace_dir)
 
-            # Close connection after shutdown
-            client.close()
+            # Modify the .env file with the POT_EXTERNAL_ENTROPY value
+            logger.debug(f"Modifying .env file for timekeeper with POT_EXTERNAL_ENTROPY={args.pot_external_entropy}")
+            modify_env_file(client, subspace_dir, release_version, pot_external_entropy=args.pot_external_entropy)
+
+            # Start the timekeeper node
+            docker_compose_up(client, subspace_dir)
+
+            logger.info("Timekeeper node started with the updated POT_EXTERNAL_ENTROPY value.")
         except Exception as e:
-            logger.error(f"Error during sudo docker compose down -v on {node['host']}: {e}")
+            logger.error(f"Error during timekeeper node update: {e}")
         finally:
             if client:
                 client.close()
-                logger.debug(f"Closed connection for node {node['host']}")
+                logger.debug(f"Closed connection to timekeeper node {timekeeper_node['host']}")
+    elif args.no_timekeeper:
+        logger.info("Skipping timekeeper node as --no-timekeeper flag is set.")
+    else:
+        logger.warning("Timekeeper node not found, proceeding with other nodes.")
 
-    # Step 2: Update .env and start sudo docker compose for RPC and Farmer nodes
+    # Step 2: Start the other farmer and RPC nodes after the timekeeper node
     protocol_version_hash = None
     for node in farmer_rpc_nodes:
         client = None  # Initialize the client variable
         try:
-            logger.info(f"Connecting to {node['host']}...")
+            logger.info(f"Connecting to {node['host']} for sudo docker compose down -v...")
             client = ssh_connect(node['host'], node['user'], node['ssh_key'])
 
-            # Modify the .env file
-            modify_env_file(client, subspace_dir, release_version)
+            # Run sudo docker compose down -v
+            docker_compose_down(client, subspace_dir)
+
+            # Modify the .env file for farmer and RPC nodes
+            modify_env_file(client, subspace_dir, release_version, pot_external_entropy=args.pot_external_entropy)
 
             # Start sudo docker compose up -d
             docker_compose_up(client, subspace_dir)
@@ -212,34 +226,7 @@ def main():
                 client.close()
                 logger.debug(f"Closed connection for node {node['host']}")
 
-    if timekeeper_node:
-        client = None  # Initialize the client variable
-        if args.pot_external_entropy:
-            try:
-                logger.info(f"Connecting to the timekeeper node {timekeeper_node['host']}...")
-                client = ssh_connect(timekeeper_node['host'], timekeeper_node['user'], timekeeper_node['ssh_key'])
-
-                # Run sudo docker compose down -v
-                docker_compose_down(client, subspace_dir)
-
-                # Modify the .env file with the POT_EXTERNAL_ENTROPY value
-                logger.debug(f"Modifying .env file with POT_EXTERNAL_ENTROPY={args.pot_external_entropy}")
-                modify_env_file(client, subspace_dir, release_version, pot_external_entropy=args.pot_external_entropy)
-
-                # Start the timekeeper node
-                docker_compose_up(client, subspace_dir)
-
-                logger.info("Timekeeper node started with the updated POT_EXTERNAL_ENTROPY value.")
-            except Exception as e:
-                logger.error(f"Error during timekeeper node update: {e}")
-            finally:
-                if client:
-                    client.close()
-                    logger.debug(f"Closed connection to timekeeper node {timekeeper_node['host']}")
-        else:
-            logger.warning(f"POT_EXTERNAL_ENTROPY not provided for the timekeeper node, skipping update.")
-
-    # Step 3: SSH into the bootstrap node and update GENESIS_HASH, then start it
+    # Step 3: SSH into the bootstrap node and update GENESIS_HASH and POT_EXTERNAL_ENTROPY, then start it
     if protocol_version_hash:
         client = None  # Initialize the client variable
         try:
@@ -249,14 +236,14 @@ def main():
             # Run sudo docker compose down -v for the bootstrap node
             docker_compose_down(client, subspace_dir)
 
-            # Modify .env with the new GENESIS_HASH
-            modify_env_file(client, subspace_dir, release_version, genesis_hash=protocol_version_hash)
+            # Modify .env with the new GENESIS_HASH and POT_EXTERNAL_ENTROPY
+            modify_env_file(client, subspace_dir, release_version, genesis_hash=protocol_version_hash, pot_external_entropy=args.pot_external_entropy)
 
             # Start the bootstrap node
             docker_compose_up(client, subspace_dir)
 
             client.close()
-            logger.info("Bootstrap node started with the updated Genesis Hash.")
+            logger.info("Bootstrap node started with the updated Genesis Hash and POT_EXTERNAL_ENTROPY.")
         except Exception as e:
             logger.error(f"Error during bootstrap node update: {e}")
         finally: