malsched

#!/usr/bin/env python3

import configparser
import glob
import json
import os
import pickle
import sys
import string
import subprocess
import tempfile
import threading
import time
import urllib

import pymonetdb
import lz4.frame

from malcom.mal_dict import MalDictionary
from malcom.mal_instr import MalInstruction
from malcom.stats import ColumnStatsD
from malcom.utils import Utils


class QueryError(Exception):
    def __init__(self, msg, exit_code = None):
        self.message = msg
        self.exit_code = exit_code


class Connection:
    def __init__(self, name, mapi, size):
        self.name = name
        self.mapi = mapi
        self.size = size
        self.conn = self.connect()
        self.conn.cursor().execute("SELECT 1;")

    def connect(self):
        if not self.mapi.startswith('mapi:'):
            raise RuntimeError("MAPI URL must start with 'mapi:': " + self.mapi)
        p = urllib.parse.urlparse(self.mapi[5:])
        username = p.username or 'monetdb'
        password = p.password or 'monetdb'
        hostname = p.hostname or 'localhost'
        port = p.port or 50000
        database = p.path[1:]
        if '/' in database:
            raise RuntimeError("Invalid database name " + repr(database))
        c = pymonetdb.Connection(
            database=database,
            hostname=hostname,
            port=port,
            username=username,
            password=password,
            connect_timeout=4.0
        )
        return c

    def cursor(self):
        return self.conn.cursor()

    def tracer(self):
        return Tracer(self.mapi)


class Tracer:
    def __init__(self, mapi_url):
        self.url = mapi_url
        self.tmp = tempfile.TemporaryFile()
        self.proc = subprocess.Popen(['stethoscope', '-j', '-d', self.url], stdout=self.tmp, stderr=subprocess.DEVNULL)
        self.data = None
        time.sleep(0.2)  # random sleeps in the name of reliability
        
    def finish(self, timeout=3):
        if not self.tmp:
            return self.data
        time.sleep(0.2)  # random sleeps in the name of reliability
        self.proc.terminate()
        self.proc.wait(timeout=timeout)
        self.tmp.seek(0)
        self.data = self.tmp.read()
        self.tmp.close()
        self.tmp = None
        return self.data


class Interaction:
    out = sys.stdout
    at_beginning_of_line = True
    max_delay = 0.8
    animation = [chr(0x1f550 + i) for i in range(12)]
    animation_interval = 0.05

    def __call__(self, fmt, *args, **opts):
        txt = fmt.format(*args, **opts)
        self.out.write(txt)
        self.out.flush()
        if txt:
            self.at_beginning_of_line = txt.endswith('\n')

    def close(self):
        self.go_to_start()

    def go_to_start(self):
        if not self.at_beginning_of_line:
            self('\n')

    def delay(self, seconds):
        if seconds < self.max_delay:
            time.sleep(seconds)
        else:
            self.play_animation()

    def play_animation(self):
        reset = '\b\b'
        clear = '  '
        for s in self.animation:
            self.out.write(s)
            self.out.flush()
            time.sleep(self.animation_interval)
            self.out.write(reset)
            self.out.flush()
        self.out.write(clear)
        self.out.flush()
        self.out.write(reset)
        self.out.flush()

    def prompt(self, prompt="> "):
        self.go_to_start()
        answer = input(prompt)
        self.at_beginning_of_line = True
        return answer

interact = Interaction()

def main(conf_file):
    if not os.path.isfile(conf_file):
        print("No such config file: {}".format(conf_file))
        return 1
    conf = default_config()
    conf.read([conf_file])

    interact("Welcome to the Malcom scheduler.\n")

    connections = connect_to_databases(conf)
    model = load_model(conf)

    while True:
        try:
            qname = ask_which_query()
            if not qname:
                return 0

            query = read_query(conf, qname)
            interact("Query:\n{}\n", query)

            expected = predict_memory_use(conf, model, qname, query)
            interact("Estimated intermediate data size: {}\n", format_size(expected))

            (db_name, size_gb) = chose_best_database(conf, expected)
            interact("Advise to send this to {} which has {}\n", db_name, format_size(size_gb))

            others = [ c for c in connections.values() if c.name != db_name ]
            for con in [connections[db_name]] + sorted(others, key=lambda c: -c.size):
                measure_query(conf, con, qname, query)

        except QueryError as e:
            if e.exit_code != None:
                interact("Error [{}]: {}\n\n", e.exit_code, e.message)
                return e.exit_code
            else:
                interact("Error: {}\n\n", e.message)


def default_config():
    conf = configparser.ConfigParser(None)
    return conf


def load_model(conf):
    prefix = conf['malsched']['trace_directory']
    globs = [word.strip() for word in conf['traces']['training'].split(',') if word]
    paths = [p for g in globs for p in glob.glob(os.path.join(prefix, g))]
    if not paths:
        raise RuntimeError("no training traces configured in [traces]training")

    interact("Loading {} traces\n", len(paths))
    maldicts = []
    for i, p in enumerate(paths):
        interact("{}/{}\r", i, len(paths))
        maldict = load_trace(p)
        maldicts.append(maldict)
    interact(' ' * 16 + '\r')

    interact("Processing...")
    [first, *rest] = maldicts
    union = first.union(*rest)
    return union

def load_trace(path):
    with open(path, 'rb') as raw_file:
        if path.endswith('.lz4'):
            with lz4.frame.LZ4FrameFile(raw_file, mode='rb') as lz_file:
                maldict = pickle.load(lz_file)
        else:
            maldict = pickle.load(raw_file)
    return maldict


def connect_to_databases(conf):
    databases = [ name for name in conf.sections() if 'mapi' in conf[name] ]
    if not databases:
        raise RuntimeError("No database configured")
    for db in databases:
        if not 'size_gb' in conf[db]:
            raise QueryError("DB {} does not have size_gb configured".format(db))

    clients = {}
    for db in databases:
        interact("Connecting to {}..", db)
        mapi = conf[db]['mapi']
        size = float(conf[db]['size_gb']) * 1024 ** 3
        clients[db] = Connection(db, mapi, size)
        interact("connected\n")

    return clients


def ask_which_query():
    """Interacts with the user via `interact`, returning either a query name
    or None if the user wants to quit
    """
    try:
        interact.go_to_start()
        interact("\n===============================\n")
        interact("Enter the name of a TPC-H query, for example Q22\n")
        name = interact.prompt()
        if name:
            return name.strip()
        interact("(Two more ENTERs mean QUIT)\n")
        for _ in range(2):
            name = interact.prompt()
            if name:
                return name.strip()
        return None
    except EOFError:
        return None


def read_query(conf, qname):
    prefix = conf['malsched']['sql_directory']
    path = os.path.join(prefix, qname.upper() + '.sql')
    if not os.path.exists(path):
        raise QueryError("File not found: {}".format(path))
    with open(path) as f:
        return f.read()


def predict_memory_use(conf, model, qname, query):
    trace_file_name = conf['traces'].get(qname.upper())
    if not trace_file_name:
        raise QueryError("No trace file configured for query {}".format(qname))
    prefix = conf['malsched']['trace_directory']
    path = os.path.join(prefix, trace_file_name)
    trace = load_trace(path)
    graph = trace.buildApproxGraph(model)
    return trace.predictMaxMem(graph)


def chose_best_database(conf, bytes_needed):
    databases = [
        (name, float(conf[name]['size_gb']))
        for name in conf.sections()
        if 'mapi' in conf[name]
    ]
    choice = None
    best = None
    for (name, gb_max) in databases:
        size = gb_max * 1024 * 1024 * 1024
        if not choice:
            # Anything is better than nothing
            better = True
        elif best < bytes_needed and size > best:
            # Still not enough, but better than what we had
            better = True
        elif bytes_needed <= size < best:
            # We had something suitable but this is better
            better = True
        else:
            # Sorry, we'll pass
            better = False

        if better:
            choice = name
            best = size

    return (choice, best)


def measure_query(conf, con, qname, query):
    interact("\nExecuting {} on {}.. ", qname, con.name)
    tracer = con.tracer()
    c = con.cursor()
    t0 = time.time()
    nrows = c.execute(query)
    t1 = time.time()
    c.close()
    trace = tracer.finish()
    interact("{} rows returned in {:.2f}s,", nrows, t1 - t0)
    #
    blacklist = Utils.init_blacklist(conf['malsched']['blacklist'])
    stats = ColumnStatsD.fromFile(conf['malsched']['stats'])
    d = MalDictionary.fromJsonFile(trace, blacklist, stats)
    max_mem = d.getMaxMem()
    #
    trace_lines = trace.splitlines()
    max_rss = 1024 * 1024 * max((json.loads(line).get('rss', 0) for line in trace_lines), default=0)
    #
    interact("\nmax intermediate data size {}, server rss {}\n",
        format_size(max_mem), format_size(max_rss))



# Copied from StackOverflow
# https://stackoverflow.com/questions/1094841/reusable-library-to-get-human-readable-version-of-file-size
def format_size(num, suffix='B'):
    for unit in ['','Ki','Mi','Gi','Ti','Pi','Ei','Zi']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'Yi', suffix)


if __name__ == "__main__":
    try:
        if len(sys.argv) == 2:
            ret = main(sys.argv[1]) or 0
            sys.exit(ret)
        else:
            print("Usage: {} CONFIG_FILE.ini".format(sys.argv[0]))
            sys.exit(1)
    finally:
        interact.close()