tangoREST.py

# tangoREST.py
#
# Implements open, upload, addJob, and poll to be used for the RESTful
# interface of Tango.
#

import sys
import os
import inspect
import hashlib
import json
import logging
import docker

currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
parentdir = os.path.dirname(currentdir)
sys.path.insert(0, parentdir)

from config import Config
from tangoObjects import TangoJob, TangoMachine, InputFile
from tango import TangoServer


class Status(object):
    def __init__(self):
        self.found_dir = self.create(0, "Found directory")
        self.made_dir = self.create(0, "Created directory")
        self.file_uploaded = self.create(0, "Uploaded file")
        self.file_exists = self.create(0, "File exists")
        self.job_added = self.create(0, "Job added")
        self.obtained_info = self.create(0, "Found info successfully")
        self.obtained_jobs = self.create(0, "Found list of jobs")
        self.preallocated = self.create(0, "VMs preallocated")
        self.obtained_pool = self.create(0, "Found pool")
        self.obtained_all_pools = self.create(0, "Found all pools")
        self.partial_output_obtained = self.create(0, "Partial output obtained")
        self.image_built = self.create(0, "Docker image built")

        self.wrong_key = self.create(-1, "Key not recognized")
        self.wrong_courselab = self.create(-1, "Courselab not found")
        self.out_not_found = self.create(-1, "Output file not found")
        self.invalid_image = self.create(-1, "Invalid image name")
        self.invalid_prealloc_size = self.create(-1, "Invalid prealloc size")
        self.pool_not_found = self.create(-1, "Pool not found")
        self.prealloc_failed = self.create(-1, "Preallocate VM failed")
        self.image_build_failed = self.create(-1, "Image build failed")

    def create(self, id, msg):
        """create - Constructs a dict with the given ID and message"""
        result = {}
        result["statusId"] = id
        result["statusMsg"] = msg
        return result


class TangoREST(object):

    COURSELABS = Config.COURSELABS
    OUTPUT_FOLDER = Config.OUTPUT_FOLDER
    LOGFILE = Config.LOGFILE

    # Replace with choice of key store and override validateKey.
    # This key is just for testing.
    KEYS = Config.KEYS

    def __init__(self):

        logging.basicConfig(
            filename=self.LOGFILE,
            format="%(levelname)s|%(asctime)s|%(name)s|%(message)s",
            level=Config.LOGLEVEL,
        )
        self.log = logging.getLogger("TangoREST")
        self.log.info("Starting RESTful Tango server")

        self.tango = TangoServer()
        self.status = Status()

    def validateKey(self, key):
        """validateKey - Validates key provided by client"""
        result = False
        for el in self.KEYS:
            if el == key:
                result = True
        return result

    def getDirName(self, key, courselab):
        """getDirName - Computes directory name"""
        return "%s-%s" % (key, courselab)

    def getDirPath(self, key, courselab):
        """getDirPath - Computes directory path"""
        labName = self.getDirName(key, courselab)
        return "%s/%s" % (self.COURSELABS, labName)

    def getOutPath(self, key, courselab):
        """getOutPath - Computes output directory path"""
        labPath = self.getDirPath(key, courselab)
        return "%s/%s" % (labPath, self.OUTPUT_FOLDER)

    def checkFileExists(self, directory, filename, fileMD5):
        """checkFileExists - Checks if a file exists in a
        directory
        """
        for elem in os.listdir(directory):
            if elem == filename:
                try:
                    body = open("%s/%s" % (directory, elem), "rb").read()
                    md5hash = hashlib.md5(body).hexdigest()
                    return md5hash == fileMD5
                except IOError:
                    continue

    def createTangoMachine(self, image, vmms=Config.VMMS_NAME, vmObj=None):
        """createTangoMachine - Creates a tango machine object from image"""
        cores = getattr(Config, "DOCKER_CORES_LIMIT", None)
        memory = getattr(Config, "DOCKER_MEMORY_LIMIT", None)
        if vmObj and "cores" in vmObj and "memory" in vmObj:
            cores = vmObj["cores"]
            memory = vmObj["memory"]
        return TangoMachine(
            name=image,
            vmms=vmms,
            image="%s" % (image),
            cores=cores,
            memory=memory,
            disk=None,
            network=None,
        )

    def convertJobObj(self, dirName, jobObj):
        """convertJobObj - Converts a dictionary into a TangoJob object"""

        name = jobObj["jobName"]
        outputFile = "%s/%s/%s/%s" % (
            self.COURSELABS,
            dirName,
            self.OUTPUT_FOLDER,
            jobObj["output_file"],
        )
        timeout = jobObj["timeout"]
        notifyURL = None
        maxOutputFileSize = Config.MAX_OUTPUT_FILE_SIZE
        if "callback_url" in jobObj:
            notifyURL = jobObj["callback_url"]

        # List of input files
        input = []
        for file in jobObj["files"]:
            inFile = file["localFile"]
            vmFile = file["destFile"]
            handinfile = InputFile(
                localFile="%s/%s/%s" % (self.COURSELABS, dirName, inFile),
                destFile=vmFile,
            )
            input.append(handinfile)

        # VM object
        vm = self.createTangoMachine(jobObj["image"])

        # for backward compatibility
        accessKeyId = None
        accessKey = None
        if "accessKey" in jobObj and len(jobObj["accessKey"]) > 0:
            accessKeyId = jobObj["accessKeyId"]
            accessKey = jobObj["accessKey"]

        disableNetwork = False
        if "disable_network" in jobObj and isinstance(jobObj["disable_network"], bool):
            disableNetwork = jobObj["disable_network"]

        allowedOutgoingIPs = None
        if "allowed_outgoing_ips" in jobObj and isinstance(
            jobObj["allowed_outgoing_ips"], list
        ):
            allowedOutgoingIPs = jobObj["allowed_outgoing_ips"]

        job = TangoJob(
            name=name,
            vm=vm,
            outputFile=outputFile,
            input=input,
            timeout=timeout,
            notifyURL=notifyURL,
            maxOutputFileSize=maxOutputFileSize,
            accessKey=accessKey,
            accessKeyId=accessKeyId,
            disableNetwork=disableNetwork,
            allowedOutgoingIPs=allowedOutgoingIPs,
        )

        self.log.debug("inputFiles: %s" % [file.localFile for file in input])
        self.log.debug("outputFile: %s" % outputFile)
        return job

    def convertTangoMachineObj(self, tangoMachine):
        """convertVMObj - Converts a TangoMachine object into a dictionary"""
        # May need to convert instance_id
        vm = dict()
        vm["network"] = tangoMachine.network
        vm["resume"] = tangoMachine.resume
        vm["image"] = tangoMachine.image
        vm["memory"] = tangoMachine.memory
        vm["vmms"] = tangoMachine.vmms
        vm["cores"] = tangoMachine.cores
        vm["disk"] = tangoMachine.disk
        vm["id"] = tangoMachine.id
        vm["name"] = tangoMachine.name
        return vm

    def convertInputFileObj(self, inputFile):
        """convertInputFileObj - Converts an InputFile object into a dictionary"""
        input = dict()
        input["destFile"] = inputFile.destFile
        input["localFile"] = inputFile.localFile
        return input

    def convertTangoJobObj(self, tangoJobObj):
        """convertTangoJobObj - Converts a TangoJob object into a dictionary"""
        job = dict()
        # Convert scalar attribtues first
        job["retries"] = tangoJobObj.retries
        job["outputFile"] = tangoJobObj.outputFile
        job["name"] = tangoJobObj.name
        job["notifyURL"] = tangoJobObj.notifyURL
        job["maxOutputFileSize"] = tangoJobObj.maxOutputFileSize
        job["assigned"] = tangoJobObj.assigned
        job["timeout"] = tangoJobObj.timeout
        job["id"] = tangoJobObj.id
        job["trace"] = tangoJobObj.trace

        # Convert VM object
        job["vm"] = self.convertTangoMachineObj(tangoJobObj.vm)

        # Convert InputFile objects
        inputFiles = list()
        for inputFile in tangoJobObj.input:
            inputFiles.append(self.convertInputFileObj(inputFile))
        job["input"] = inputFiles

        return job

    ##
    # Tango RESTful API
    ##

    def open(self, key, courselab):
        """open - Return a dict of md5 hashes for each input file in the
        key-courselab directory and make one if the directory doesn't exist
        """
        self.log.debug("Received open request(%s, %s)" % (key, courselab))
        if self.validateKey(key):
            labPath = self.getDirPath(key, courselab)
            try:
                if os.path.exists(labPath):
                    self.log.info("Found directory for (%s, %s)" % (key, courselab))
                    statusObj = self.status.found_dir
                    statusObj["files"] = {}
                    return statusObj
                else:
                    outputPath = self.getOutPath(key, courselab)
                    os.makedirs(outputPath)
                    self.log.info("Created directory for (%s, %s)" % (key, courselab))
                    statusObj = self.status.made_dir
                    statusObj["files"] = {}
                    return statusObj
            except Exception as e:
                self.log.error("open request failed: %s" % str(e))
                return self.status.create(-1, str(e))
        else:
            self.log.info("Key not recognized: %s" % key)
            return self.status.wrong_key

    def upload(self, key, courselab, file, tempfile, fileMD5):
        """upload - Upload file as an input file in key-courselab if the
        same file doesn't exist already
        """
        self.log.debug("Received upload request(%s, %s, %s)" % (key, courselab, file))
        if self.validateKey(key):
            labPath = self.getDirPath(key, courselab)
            try:
                if os.path.exists(labPath):
                    if self.checkFileExists(labPath, file, fileMD5):
                        self.log.info(
                            "File (%s, %s, %s) exists" % (key, courselab, file)
                        )
                        os.unlink(tempfile)
                        return self.status.file_exists
                    absPath = "%s/%s" % (labPath, file)
                    os.rename(tempfile, absPath)
                    self.log.info(
                        "Uploaded file to (%s, %s, %s)" % (key, courselab, file)
                    )
                    return self.status.file_uploaded
                else:
                    self.log.info("Courselab for (%s, %s) not found" % (key, courselab))
                    os.unlink(tempfile)
                    return self.status.wrong_courselab
            except Exception as e:
                exc_type, exc_obj, exc_tb = sys.exc_info()
                fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1]
                print(exc_type, fname, exc_tb.tb_lineno)
                self.log.error("upload request failed: %s" % str(e))
                os.unlink(tempfile)
                return self.status.create(-1, str(e))
        else:
            self.log.info("Key not recognized: %s" % key)
            os.unlink(tempfile)
            return self.status.wrong_key

    def addJob(self, key, courselab, jobStr):
        """addJob - Add the job to be processed by Tango"""
        self.log.debug("Received addJob request(%s, %s, %s)" % (key, courselab, jobStr))
        if self.validateKey(key):
            labName = self.getDirName(key, courselab)
            try:
                jobObj = json.loads(jobStr)
                job = self.convertJobObj(labName, jobObj)
                jobId = self.tango.addJob(job)
                self.log.debug("Done adding job")
                if jobId == -1:
                    self.log.info("Failed to add job to tango")
                    return self.status.create(-1, job.trace)
                self.log.info("Successfully added job ID: %s to tango" % str(jobId))
                result = self.status.job_added
                result["jobId"] = jobId
                return result
            except Exception as e:
                exc_type, exc_obj, exc_tb = sys.exc_info()
                fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1]
                print(exc_type, fname, exc_tb.tb_lineno)
                self.log.error("addJob request failed: %s" % str(e))
                return self.status.create(-1, str(e))
        else:
            self.log.info("Key not recognized: %s" % key)
            return self.status.wrong_key

    def getPartialOutput(self, key, jobId):
        """getPartialOutput - Return the partial output of the job"""
        self.log.debug("Received getPartialOutput request(%s, %s)" % (key, jobId))
        if self.validateKey(key):
            try:
                output = self.tango.getPartialOutput(jobId)
                result = self.status.partial_output_obtained
                result["output"] = output
                return result
            except Exception as e:
                self.log.error("getPartialOutput request failed: %s" % str(e))
                return self.status.create(-1, str(e))
        else:
            self.log.info("Key not recognized: %s" % key)
            return self.status.wrong_key

    def poll(self, key, courselab, outputFile):
        """poll - Poll for the output file in key-courselab"""
        self.log.debug(
            "Received poll request(%s, %s, %s)" % (key, courselab, outputFile)
        )
        if self.validateKey(key):
            outputPath = self.getOutPath(key, courselab)
            outfilePath = "%s/%s" % (outputPath, outputFile)
            if os.path.exists(outfilePath):
                self.log.info(
                    "Output file (%s, %s, %s) found" % (key, courselab, outputFile)
                )
                output = open(outfilePath)
                result = output.read()
                output.close()
                return result
            self.log.info(
                "Output file (%s, %s, %s) not found" % (key, courselab, outputFile)
            )
            return self.status.out_not_found
        else:
            self.log.info("Key not recognized: %s" % key)
            return self.status.wrong_key

    def info(self, key):
        """info - Returns basic status for the Tango service such as uptime, number of jobs etc"""
        self.log.debug("Received info request (%s)" % (key))
        if self.validateKey(key):
            info = self.tango.getInfo()
            result = self.status.obtained_info
            result["info"] = info
            return result
        else:
            self.log.info("Key not recognized: %s" % key)
            return self.status.wrong_key

    def jobs(self, key, deadJobs):
        """jobs - Returns the list of live jobs (deadJobs == 0) or the list of dead jobs (deadJobs == 1)"""
        self.log.debug("Received jobs request (%s, %s)" % (key, deadJobs))
        if self.validateKey(key):
            jobs = list()
            result = self.status.obtained_jobs
            if int(deadJobs) == 0:
                jobs = self.tango.getJobs(0)
                self.log.debug("Retrieved live jobs (deadJobs = %s)" % deadJobs)
            elif int(deadJobs) == 1:
                jobs = self.tango.getJobs(-1)
                self.log.debug("Retrieved dead jobs (deadJobs = %s)" % deadJobs)
            result["jobs"] = list()
            for job in jobs:
                result["jobs"].append(self.convertTangoJobObj(job))

            return result
        else:
            self.log.info("Key not recognized: %s" % key)
            return self.status.wrong_key

    def pool(self, key, image):
        """pool - Get information about pool(s) of VMs"""
        self.log.debug("Received pool request(%s, %s)" % (key, image))
        if self.validateKey(key):
            pools = self.tango.preallocator.getAllPools()
            self.log.info("All pools found")
            if image == "":
                result = self.status.obtained_all_pools
            else:
                if image in pools:
                    pools = {image: pools[image]}
                    self.log.info("Pool image found: %s" % image)
                    result = self.status.obtained_pool
                else:
                    self.log.info("Invalid image name: %s" % image)
                    result = self.status.pool_not_found

            result["pools"] = pools
            return result
        else:
            self.log.info("Key not recognized: %s" % key)
            return self.status.wrong_key

    async def prealloc(self, key, image, num, vmStr):
        """prealloc - Create a pool of num instances spawned from image"""
        self.log.debug("Received prealloc request(%s, %s, %s)" % (key, image, num))
        if self.validateKey(key):
            if vmStr != "":
                vmObj = json.loads(vmStr)
                vm = self.createTangoMachine(image, vmObj=vmObj)
            else:
                vm = self.createTangoMachine(image)

            ret = self.tango.preallocVM(vm, int(num))

            if ret == -1:
                self.log.error("Prealloc failed")
                return self.status.prealloc_failed
            if ret == -2:
                self.log.error("Invalid prealloc size")
                return self.status.invalid_prealloc_size
            if ret == -3:
                self.log.error("Invalid image name")
                return self.status.invalid_image
            self.log.info("Successfully preallocated VMs")
            return self.status.preallocated
        else:
            self.log.info("Key not recognized: %s" % key)
            return self.status.wrong_key

    def build(self, key, tempfile, imageName):
        self.log.debug("Received build request(%s)" % (key))
        if self.validateKey(key):
            if Config.VMMS_NAME != "localDocker":
                self.log.error("Not using Docker backend, so cannot build image")
                os.unlink(tempfile)
                return self.status.image_build_failed
            try:
                client = docker.from_env()
                imageTarStr = open(tempfile, "rb").read()
                images = client.images.load(imageTarStr)
                if len(images) != 1:
                    for image in images:
                        client.images.remove(image.id, force=True)
                    raise Exception(
                        "Wrong number of images built:  %s" % str(len(images))
                    )
                image = images[0]
                imageName = imageName + ":latest"
                image.tag(imageName)
                for tag in image.tags:
                    if tag != imageName:
                        client.images.remove(tag)
            except Exception as e:
                self.log.error("Image build failed: " + str(e))
                os.unlink(tempfile)
                return self.status.image_build_failed

            self.log.info("Successfully loaded image: %s" % (imageName))
            os.unlink(tempfile)
            return self.status.image_built
        else:
            self.log.info("Key not recognized: %s" % key)
            os.unlink(tempfile)
            return self.status.wrong_key