launch: Add logic to set and discover Amazon Linux 2022 AMI

kislyuk · Jul 1, 2022 · a4caace · a4caace
1 parent dc9c8e6
commit a4caace
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 6 deletions.
diff --git a/aegea/base_config.yml b/aegea/base_config.yml
@@ -732,6 +732,7 @@ launch:
     - AmazonElasticFileSystemClientFullAccess
     - AWSOpsWorksCloudWatchLogs
     - AmazonSSMManagedInstanceCore
+  amazon_linux_release: 2
 
 ssh:
   server_alive_interval: 15

diff --git a/aegea/launch.py b/aegea/launch.py
@@ -22,7 +22,7 @@
     aegea launch my-instance --spot-price 1 --instance-type r5d.xlarge --efs-home
 """
 
-import os, sys, time, datetime, base64, json
+import os, sys, time, datetime, base64, json, argparse
 from typing import Dict, List
 
 import yaml
@@ -111,7 +111,7 @@ def launch(args):
     if args.ubuntu_linux_ami:
         args.ami = locate_ami("Ubuntu", release="20.04", architecture=arch)
     elif args.amazon_linux_ami:
-        args.ami = locate_ami("Amazon Linux", release="2", architecture=arch)
+        args.ami = locate_ami("Amazon Linux", release=str(args.amazon_linux_release), architecture=arch)
     else:
         try:
             if not (args.ami or args.ami_tags or args.ami_tag_keys):
@@ -293,6 +293,7 @@ def instance_types(self, **kwargs):
                     help="Use the most recent AMI with these tag names")
 parser.add_argument("--ubuntu-linux-ami", action="store_true", help="Use the most recent Ubuntu Linux LTS AMI")
 parser.add_argument("--amazon-linux-ami", action="store_true", help="Use the most recent Amazon Linux 2 AMI")
+parser.add_argument("--amazon-linux-release", help=argparse.SUPPRESS)
 parser.add_argument("--spot", action="store_true",
                     help="Launch a preemptible spot instance, which is cheaper but could be forced to shut down")
 parser.add_argument("--duration-hours", type=float, help="Terminate the spot instance after this number of hours")

diff --git a/aegea/util/__init__.py b/aegea/util/__init__.py
@@ -109,7 +109,7 @@ def describe_cidr(cidr):
         whois_names += [o.get("contact", {}).get("name", "") for o in whois.get("objects", {}).values()]
     except Exception:
         try:
-            whois_names = [socket.gethostbyaddr(address)]
+            whois_names = [socket.gethostbyaddr(address)]  # type: ignore
         except Exception:
             whois_names = [cidr]
     return ", ".join(str(n) for n in whois_names)

diff --git a/aegea/util/aws/__init__.py b/aegea/util/aws/__init__.py
@@ -23,11 +23,14 @@ def locate_ami(distribution, release, architecture):
         locate_ami(distribution="Ubuntu", release="20.04", architecture="amd64")
         locate_ami(distribution="Amazon Linux", release="2", architecture="arm64")
     """
-    if distribution == "Amazon Linux" and release == "2":
+    if distribution == "Amazon Linux" and release in {"2", "2022"}:
         if architecture == "amd64":
             architecture = "x86_64"
-        ssm_param_name = "/aws/service/ami-amazon-linux-latest/amzn2-ami-hvm-{}-gp2".format(architecture)
-        ami_id = clients.ssm.get_parameters(Names=[ssm_param_name])["Parameters"][0]["Value"]
+        ssm_param_names = {
+            "2": f"/aws/service/ami-amazon-linux-latest/amzn2-ami-hvm-{architecture}-gp2",
+            "2022": f"/aws/service/ami-amazon-linux-latest/al2022-ami-kernel-default-{architecture}"
+        }
+        ami_id = clients.ssm.get_parameters(Names=[ssm_param_names[release]])["Parameters"][0]["Value"]
         logger.info("Found %s for %s %s %s", ami_id, distribution, release, architecture)
         return resources.ec2.Image(ami_id)
     elif distribution == "Ubuntu":