add commander prepare

aws-experiment-launch/new-pipeline/commander_prepare.py

@@ -0,0 +1,62 @@
+import argparse
+import logging
+import os
+import stat
+import sys
+
+from utils import utils
+
+logging.basicConfig(level=logging.INFO, format='%(threadName)s %(asctime)s - %(name)s - %(levelname)s - %(message)s')
+LOGGER = logging.getLogger(__file__)
+LOGGER.setLevel(logging.INFO)
+
+PEMS = [
+    "virginia-key-benchmark.pem",
+    "ohio-key-benchmark.pem",
+    "california-key-benchmark.pem",
+    "oregon-key-benchmark.pem",
+    "tokyo-key-benchmark.pem",
+    "singapore-key-benchmark.pem",
+    "frankfurt-key-benchmark.pem",
+    "ireland-key-benchmark.pem",
+]
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='This script helps you to genereate distribution config')
+    parser.add_argument('--distribution_config', type=str,
+                        dest='distribution_config', default='distribution_config.txt')
+    parser.add_argument('--commander_logging', type=str,
+                        dest='commander_logging', default='commander_logging.sh')
+
+    args = parser.parse_args()
+
+    if not os.path.exists(args.distribution_config):
+        sys.exit(1)
+    with open(args.distribution_config, "r") as fin:
+        lines = fin.readlines()
+
+    commander_address = None
+    commander_region = None
+    commander_output = None
+    with open(args.distribution_config, "w") as fout:
+        for line in lines:
+            if "commander" in line:
+                items = [item.strip() for item in line.split(" ")]
+                commander_address = items[0]
+                commander_region = int(items[4][0])
+                commander_output = "\n".join(items)
+            else:
+                fout.write(line.strip() + "\n")
+    LOGGER.info("Generated %s" % args.distribution_config)
+
+    if not commander_address or not commander_region:
+        LOGGER.info("Failed to extract commander address and commander region.")
+        sys.exit(1)
+
+    with open(args.commander_logging, "w") as fout:
+        fout.write("ssh -i ../keys/%s ec2-user@%s\n" % (PEMS[commander_region - 1], commander_address))
+    st = os.stat(args.commander_logging)
+    os.chmod(args.commander_logging, st.st_mode | stat.S_IEXEC)
+    LOGGER.info("Generated %s" % args.commander_logging)
+    LOGGER.info("DONE.")
+