modelverse/performance/utils.py

import unittest
import sys
import os

import sys
import time
import json
import urllib
import urllib2
import subprocess
import signal
import random

sys.path.append("interface/HUTN")
sys.path.append("scripts")
from hutn_compiler.compiler import main as do_compile
from check_objects import to_recompile

USERNAME = "test_user"
PARALLEL_PUSH = True

BOOTSTRAP_FOLDER_NAME = "bootstrap"
CURRENT_FOLDER_NAME = "performance"

PORTS = set()

class ModelverseTerminated(Exception):
    """An exception that tells the user that the Modelverse has terminated."""
    pass

def get_code_folder_name():
    """Gets the name of the code folder."""
    return '%s/code' % CURRENT_FOLDER_NAME

def get_free_port():
    """Gets a unique new port."""
    while 1:
        port = random.randint(10000, 20000)
        # Check if this port is in the set of ports.
        if port not in PORTS:
            # We have found a unique port. Add it to the set and return.
            PORTS.add(port)
            return port

def execute(scriptname, parameters=None, wait=False):
    """Runs a script."""
    if os.name not in ["nt", "posix"]:
        # Stop now, as we would have no clue on how to kill its subtree
        raise Exception("Unknown OS version: " + str(os.name))

    command = [sys.executable, "scripts/%s.py" % scriptname] + (
        [] if parameters is None else parameters)

    if wait:
        return subprocess.call(command, shell=False)
    else:
        return subprocess.Popen(command, shell=False)

def kill(process):
    """Kills the given process."""
    if os.name == "nt":
        subprocess.call(["taskkill", "/F", "/T", "/PID", "%i" % process.pid])
    elif os.name == "posix":
        subprocess.call(["pkill", "-P", "%i" % process.pid])

def set_input_data(address, data):
    """Sets the Modelverse program's input data."""
    if data is not None:
        urllib2.urlopen(
            urllib2.Request(
                address,
                urllib.urlencode(
                    {"op": "set_input", "data": json.dumps(data), "username": USERNAME})),
            timeout=10).read()
    else:
        return []

def compile_file(address, mod_filename, filename, mode, proc):
    """Compiles the given file."""
    # Load in the file required
    try:
        timeout_val = 240
        username = str(random.random())
        while 1:
            proc2 = execute(
                "compile", [address, mod_filename, username, filename, mode], wait=False)

            if proc.returncode is not None:
                # Modelverse has already terminated, which isn't a good sign!
                raise Exception("Modelverse died!")

            while proc2.returncode is None:
                time.sleep(0.01)
                proc2.poll()
                timeout_val -= 0.01
                if timeout_val < 0:
                    kill(proc2)
                    print("Compilation timeout expired!")
                    return False

            if proc2.returncode != 2:
                break

        # Make sure everything stopped correctly
        assert proc2.returncode == 0
        if proc2.returncode != 0:
            return False
    except:
        raise
    finally:
        try:
            kill(proc2)
        except UnboundLocalError:
            pass

def run_file(files, parameters, mode, handle_output):
    """Compiles the given sequence of files, feeds them the given input in the given mode,
       and handles their output."""
    # Resolve file
    import os.path

    time.sleep(0.01)
    port = get_free_port()
    address = "http://127.0.0.1:%i" % port
    try:
        # Run Modelverse server
        proc = execute("run_local_modelverse", [str(port)], wait=False)

        threads = []
        mod_files = []
        for filename in files:
            if os.path.isfile("%s/%s" % (get_code_folder_name(), filename)):
                mod_filename = "%s/%s" % (get_code_folder_name(), filename)
            elif os.path.isfile("%s/%s" % (BOOTSTRAP_FOLDER_NAME, filename)):
                mod_filename = "%s/%s" % (BOOTSTRAP_FOLDER_NAME, filename)
            else:
                raise Exception("File not found: %s" % filename)
            mod_files.append(mod_filename)

        to_compile = to_recompile(address, mod_files)

        for mod_filename in to_compile:
            if PARALLEL_PUSH:
                import threading
                threads.append(
                    threading.Thread(
                        target=compile_file,
                        args=[address, mod_filename, mod_filename, mode, proc]))
                threads[-1].start()
            else:
                compile_file(address, mod_filename, mod_filename, mode, proc)

        if PARALLEL_PUSH:
            for t in threads:
                t.join()

        if mode[-1] == "O":
            # Fire up the linker
            val = execute("link_and_load", [address, USERNAME] + mod_files, wait=True)
            if val != 0:
                raise Exception("Linking error")

        # Send the request ...
        set_input_data(address, parameters)

        # ... and wait for replies
        while 1:
            val = urllib2.urlopen(
                urllib2.Request(
                    address,
                    urllib.urlencode({"op": "get_output", "username": USERNAME})),
                timeout=240).read()
            val = json.loads(val)

            if proc.returncode is not None:
                # Modelverse has terminated. This may or may not be what we want.
                raise ModelverseTerminated()

            if not handle_output(val):
                return

        # All passed!
        return
    except:
        raise
    finally:
        try:
            kill(proc)
        except UnboundLocalError:
            pass

def run_file_to_completion(files, parameters, mode):
    """Compiles the given sequence of files, feeds them the given input in the given mode,
       then collects and returns their output."""
    results = []
    def handle_output(output):
        """Appends the given output to the list of results."""
        results.append(output)
        return True

    try:
        run_file(files, parameters, mode, handle_output)
    except ModelverseTerminated:
        return results