minerva/minerva/machine.py

import time
#from pexpect import pxssh
from fabric import Connection
import shlex
import threading

class Machine:
    def __init__(self,
                 pier,
                 ami           = "ami-0a538467cc9da9bb2", # ubuntu 22
                 instance_type = "t2.micro",
                 variables     = {},
                 username      = None,
                 key_pair      = None,
                 name          = "Minerva Instance",
                 public        = True,
                 disk_size     = 8):

        self.pier       = pier
        self.ami        = ami
        self.instance_type = instance_type
        self.username   = username
        self.key_pair   = key_pair
        self.variables  = variables
        self.name       = name
        self.ready      = False
        self.info       = None
        self.ssh        = None
        self.terminated = False
        self.public     = public
        self.disk_size  = disk_size

    def create(self):
        if self.info:
            return

        iam = {'Name': self.pier.iam} if self.pier.iam else {}
        res = self.pier.ec2.run_instances(
            ImageId      = self.ami,
            InstanceType = self.instance_type,
            KeyName      = self.key_pair or self.pier.key_pair_name,
            MinCount     = 1,
            MaxCount     = 1,
            TagSpecifications = [{'ResourceType': 'instance',
                                  'Tags': [{'Key': 'Name', 'Value': self.name}]}],
            NetworkInterfaces = [{'AssociatePublicIpAddress': self.public,
                                  'SubnetId':    self.pier.subnet_id,
                                  'Groups':      self.pier.groups,
                                  'DeviceIndex': 0}],
            BlockDeviceMappings = [{'DeviceName': '/dev/sda1',
                                    'Ebs': {'VolumeSize':          self.disk_size,
                                            'DeleteOnTermination': True}}],
            IamInstanceProfile = iam
        )

        self.info = res['Instances'][0]
        self.private_ip = self.info['NetworkInterfaces'][0]['PrivateIpAddress']

        # FIXME there should be a check here in case some instances fail to
        # start up in a timely manner
        # Start a countdown in the background
        # to give time for the instance to start up
        wait_time = 30
        self.thread = threading.Thread(target = self.wait,
                                       args   = (wait_time,),
                                       daemon = True)
        self.thread.start()

    def status(self):
        resp = self.pier.ec2.describe_instance_status(InstanceIds=[self.info['InstanceId']],
                                                      IncludeAllInstances=True)
        return resp['InstanceStatuses'][0]['InstanceState']['Name']

    def join(self):
        self.thread.join()

    def wait(self, n):
        i = 0
        # Time for the server to show as "running"
        # and time for the server to finish getting daemons running
        while self.status() != "running":
            time.sleep(10)
            i += 1

            if i > 18:
                reason = f"{self.info['InstanceId']} took too long to start ({i} attempts)"
                raise Exception(reason)

        # Final wait, now that the server is up and running -- need
        # some time for daemons to start
        time.sleep(25)
        self.ready = True

    # alternatively, could maybe implement this with SSM so that we can access
    # private subnets? TODO
    def login(self):
        if self.ssh:
            return True

        if not self.public:
            raise Exception("Can only log into server that has a public IP")

        # Machine must be running first, so we need to wait for the countdown to finish
        self.join()

        resp = self.pier.ec2.describe_instances(InstanceIds=[self.info['InstanceId']])
        self.description = resp['Reservations'][0]['Instances'][0]
        self.public_ip   = self.description['PublicIpAddress']

        print(f"\t{self.name} ({self.info['InstanceId']}) => {self.public_ip} ({self.private_ip})")

        self.ssh = Connection(self.public_ip,
                              self.username,
                              connect_kwargs = {
                                  "key_filename": self.pier.key_path
                              }
                             )
        return True

    def prep_variables(self):
        varz = [f"export {var}={repr(val)}" for var, val in self.variables.items()]
        base = ['source ~/.profile']
        return "; ".join([*base, *varz])

    # Unfortunately, under the hood, it's running /bin/bash -c '...'
    # You stand informed
    def cmd(self, command, hide=True, disown=False):
        res = self.ssh.run(f"{self.prep_variables()}; {command}",
                           warn=True,
                           hide=hide,
                           disown=disown)

        return res

    def write_env_file(self, variables, fname="~/env.list"):
        vals = "\n".join([f"{var}={val}" for var, val in variables.items()])
        self.cmd(f"echo {shlex.quote(vals)} > {fname}")
        return fname

    def aws_docker_login(self, ecr):
        return self.cmd(f"aws ecr get-login-password --region {self.pier.session.region_name} | " +
                        f"docker login --username AWS --password-stdin {ecr}"
)

    def docker_run(self, uri, cmd="", env={}):
        if env:
            fname   = self.write_env_file(env)
            environ = f"--env-file {fname}"
        else:
            environ = ""

        return self.cmd(f"docker run {environ} {uri} {cmd}")

    def docker_pull(self, uri):
        return self.cmd(f"docker pull {uri}")

    def terminate(self):
        if self.terminated:
            return

        self.pier.ec2.terminate_instances(
            InstanceIds=[self.info['InstanceId']],
            DryRun=False
        )
        print(f"terminated {self.name} ({self.info['InstanceId']})")
        self.terminated = True