Allow AMI and Security Group selection (in combination with PR on Autolab)

jobManager.py

@@ -81,7 +81,7 @@ def __manage(self):
                     newVM = copy.deepcopy(job.vm)
                     newVM.id = self._getNextID()
                     try:
-                        preVM = vmms.initializeVM(newVM)
+                        preVM = vmms.initializeVM(newVM, ami = job.ami, security_group = job.security_group)
                     except Exception as e:
                         self.log.error("ERROR initialization VM: %s", e)
                         self.log.error(traceback.format_exc())

restful_tango/tangoREST.py

@@ -165,8 +165,10 @@ def convertJobObj(self, dirName, jobObj):
             input.append(handinfile)
 
         ec2_vmms = False
-        if "ec2Vmms" in jobObj:
+        if Config.VMMS_NAME == "ec2SSH":
             ec2_vmms = True
+        if "ec2Vmms" in jobObj:
+            ec2_vmms = jobObj["ec2Vmms"]
 
         instance_type = None
         if "instanceType" in jobObj and len(jobObj["instanceType"]) > 0:
@@ -182,6 +184,14 @@ def convertJobObj(self, dirName, jobObj):
         if "accessKey" in jobObj and len(jobObj["accessKey"]) > 0:
             accessKeyId = jobObj["accessKeyId"]
             accessKey = jobObj["accessKey"]
+
+        ami = None
+        if "ami" in jobObj and len(jobObj["ami"]) > 0:
+            ami = jobObj["ami"]
+
+        security_group = None
+        if "security_group" in jobObj and len(jobObj["security_group"]) > 0:
+            security_group = jobObj["security_group"]
 
         disableNetwork = False
         if "disable_network" in jobObj and isinstance(jobObj["disable_network"], bool):
@@ -197,6 +207,8 @@ def convertJobObj(self, dirName, jobObj):
             maxOutputFileSize=maxOutputFileSize,
             accessKey=accessKey,
             accessKeyId=accessKeyId,
+            ami = ami,
+            security_group = security_group,
             disableNetwork=disableNetwork,
         )
 

tango.py

@@ -48,6 +48,7 @@
 from jobQueue import JobQueue
 from tangoObjects import TangoJob
 from config import Config
+from vmms.ec2SSH import Ec2SSH
 
 
 class TangoServer(object):
@@ -204,7 +205,19 @@ def getInfo(self):
         stats["runjob_errors"] = Config.runjob_errors
         stats["copyout_errors"] = Config.copyout_errors
         stats["num_threads"] = threading.activeCount()
-
+
+        # Fetch tagged AMIs and security groups from Ec2SSH
+        try:
+            ec2_ssh = Ec2SSH()
+            stats["tagged_amis"] = [
+                {"name": key, "id": ec2_ssh.img2ami[key].id} for key in ec2_ssh.img2ami
+            ]
+            stats["security_groups"] = ec2_ssh.security_groups
+        except Exception as e:
+            stats["tagged_amis"] = []
+            stats["security_groups"] = []
+            logging.error(f"Failed to fetch tagged AMIs or security groups: {e}")
+
         return stats
 
     def getPartialOutput(self, jobid):

tangoObjects.py

@@ -99,6 +99,8 @@ def __init__(
         maxOutputFileSize=Config.MAX_OUTPUT_FILE_SIZE,
         accessKeyId=None,
         accessKey=None,
+        ami = None,
+        security_group = None,
         disableNetwork=None,
     ):
         self.assigned = False
@@ -119,6 +121,8 @@ def __init__(
         self._remoteLocation = None
         self.accessKeyId = accessKeyId
         self.accessKey = accessKey
+        self.ami = ami
+        self.security_group = security_group
         self.disableNetwork = disableNetwork
 
     def __repr__(self):

vmms/ec2SSH.py

@@ -192,16 +192,30 @@ def __init__(self, accessKeyId=None, accessKey=None):
 
         self.img2ami = {}
         self.images = []
+        self.security_groups = []
         try:
             self.boto3resource = boto3.resource("ec2", config.Config.EC2_REGION)
             self.boto3client = boto3.client("ec2", config.Config.EC2_REGION)
 
             # Get images from ec2
             images = self.boto3resource.images.filter(Owners=["self"])
+
+            # Get all existing security groups
+            response = self.boto3client.describe_security_groups()
+            for sg in response["SecurityGroups"]:
+                self.security_groups.append({
+                    "name": sg["GroupName"],
+                    "id": sg["GroupId"],
+                    "description": sg.get("Description", "")
+                })
         except Exception as e:
             self.log.error("EC2SSH failed initialization: %s" % (e))
             raise
-
+
+        self.log.info("Existing Security Groups:")
+        for sg in self.security_groups:
+            self.log.info(f"Group Name: {sg['name']}, Group ID: {sg['id']}, Description: {sg['description']}")
+
         for image in images:
             if image.tags:
                 for tag in image.tags:
@@ -256,10 +270,10 @@ def domainName(self, vm):
     # VMMS helper methods
     #
 
-    def tangoMachineToEC2Instance(self, vm: TangoMachine):
+    def tangoMachineToEC2Instance(self, vm: TangoMachine, ami = None):
         """tangoMachineToEC2Instance - returns an object with EC2 instance
-        type and AMI. Only general-purpose instances are used. Defalt AMI
-        is currently used.
+        type and AMI. Only general-purpose instances are used.
+        Allow user to select AMI on Autolab or default is used if nothing selected.
         """
         ec2instance = dict()
 
@@ -274,8 +288,12 @@ def tangoMachineToEC2Instance(self, vm: TangoMachine):
         else:
             ec2instance["instance_type"] = config.Config.DEFAULT_INST_TYPE
 
-        # for now, ami is config default
-        ec2instance["ami"] = self.img2ami[vm.image].id
+        if((ami is None) or ami == ""):
+            # use ami associated with image if user did not specify preference
+            ec2instance["ami"] = self.img2ami[vm.image].id
+        else:
+            # use ami user specified
+            ec2instance["ami"] = ami
 
         self.log.info("tangoMachineToEC2Instance: %s" % str(ec2instance))
         return ec2instance
@@ -304,23 +322,44 @@ def deleteKeyPair(self):
         except OSError:
             pass
 
-    def createSecurityGroup(self):
+    def createSecurityGroup(self, security_group = None):
+        # if security_group entered, try it; otherwise, try default
+        first_group = security_group or config.Config.DEFAULT_SECURITY_GROUP
         try:
             # Check if the security group already exists
             response = self.boto3client.describe_security_groups(
                 Filters=[
                     {
                         "Name": "group-name",
-                        "Values": [config.Config.DEFAULT_SECURITY_GROUP],
+                        "Values": [first_group],
                     }
                 ]
             )
             if response["SecurityGroups"]:
                 security_group_id = response["SecurityGroups"][0]["GroupId"]
                 return
         except Exception as e:
-            self.log.debug("ERROR checking for existing security group: %s", e)
+            self.log.debug("ERROR checking for existing security group '%s': %s", first_group, e)
 
+        # if security_group passed in but fails, try Config file's DEFAULT_SECURITY_GROUP
+        if security_group and security_group != config.Config.DEFAULT_SECURITY_GROUP:
+            try:
+                # Check if the security group already exists
+                response = self.boto3client.describe_security_groups(
+                    Filters=[
+                        {
+                            "Name": "group-name",
+                            "Values": [config.Config.DEFAULT_SECURITY_GROUP],
+                        }
+                    ]
+                )
+                if response["SecurityGroups"]:
+                    security_group_id = response["SecurityGroups"][0]["GroupId"]
+                    return
+            except Exception as e:
+                self.log.debug("ERROR checking for existing security group '%s': %s", config.Config.DEFAULT_SECURITY_GROUP, e)
+
+        # if neither exists, credit with default's name and allow all traffic
         try:
             response = self.boto3resource.create_security_group(
                 GroupName=config.Config.DEFAULT_SECURITY_GROUP,
@@ -336,7 +375,7 @@ def createSecurityGroup(self):
     #
     # VMMS API functions
     #
-    def initializeVM(self, vm):
+    def initializeVM(self, vm, ami=None, security_group=None):
         """initializeVM - Tell EC2 to create a new VM instance.
 
         Returns a boto.ec2.instance.Instance object.
@@ -345,10 +384,10 @@ def initializeVM(self, vm):
         # Create the instance and obtain the reservation
         try:
             instanceName = self.instanceName(vm.id, vm.name)
-            ec2instance = self.tangoMachineToEC2Instance(vm)
+            ec2instance = self.tangoMachineToEC2Instance(vm, ami)
             self.log.debug("instanceName: %s" % instanceName)
             # ensure that security group exists
-            self.createSecurityGroup()
+            self.createSecurityGroup(security_group)
             if self.useDefaultKeyPair:
                 self.key_pair_name = config.Config.SECURITY_KEY_NAME
                 self.key_pair_path = config.Config.SECURITY_KEY_PATH