run_exp

#!/usr/bin/env python3


import os
import subprocess
import uuid
import datetime as dt
import re
import argparse
import socket
import json
import config
import sys
import signal
from utils import bcolors


class ExpInfoAPI:
    def __init__(self, ip, port):
        self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.s.connect((ip, port))
        self.f = self.s.makefile()

    def _write(self, j):
        self.s.send(json.dumps(j).encode())
        self.s.send(b'\n')

    def _read(self):
        msg = self.f.readline()
        return json.loads(msg)

    def get_running(self):
        self._write({'method': 'read'})
        return self._read()

    def register_new(self, exp_id, **info):
        self._write({'method': 'insert', 'key': exp_id, 'value': info})
        assert(self._read()['error'] is None)

    def update_running(self, exp_id, **fields):
        self._write({'method': 'update', 'key': exp_id, 'value': fields})
        assert(self._read()['error'] is None)


# echo '{"method": "insert", "key": 1234, "value": {"exclusive": false, "user": "Anton", "start": "23:22", "end":"", "cmd":"asdf", "msg":"trolol", "pid":0, "numa":"0"}}' | nc localhost 9090


def print_prompt():
    with open(os.path.join(config.LOCATION, config.PROMPT_FILENAME), 'r') as f:
        sys.stdout.write(f.read().replace('\\n', '\n'))
        sys.stdout.flush()


def print_motd():
    with open(os.path.join(config.LOCATION, config.MOTD_FILENAME), 'r') as f:
        sys.stdout.write(f.read())
        sys.stdout.flush()


def get_args(print_help=False):

    parser = argparse.ArgumentParser()
    parser.add_argument('--prompt', default=False,
                        action='store_true', help='print small info for terminal prompt')
    parser.add_argument(
        '-t', '--time', help='expected runtime [hours:minutes]')

    parser.add_argument('-ex', '--exclusive', default=False, action='store_true',
                        help='request exclusive experiment access (shared by default)')
    parser.add_argument('-p', '--pin', default=False,
                        action='store_true', help='pin executeable to specified numa nodes (does invoke numactl)')

    required_group = parser.add_argument_group('required arguments')
    required_group.add_argument('-m', '--message',
                                help='descriptive message for the job')
    required_group.add_argument(
        '-n', '--numa', help='Pin to numa node (does *not* invoke numactl)')

    parser.add_argument(
        'cmd', nargs='*', help='run command as experiment job. pro-tip: write -- before your command.')

    args = parser.parse_args()

    if args.prompt:
        print_prompt()
        sys.exit(0)

    if not args.cmd:
        print_motd()
        sys.exit(0)

    if not args.message:
        parser.print_help()
        print('error: the following arguments are required: -m/--message')
        sys.exit(1)

    if not args.numa:
        parser.print_help()
        print('error: the following arguments are required: -n/--numa')
        sys.exit(1)

    return args


def get_user():
    if 'SUDO_USER' in os.environ:
        return os.environ['SUDO_USER']
    else:
        return os.environ['USER']


if __name__ == '__main__':
    args = get_args()

    cmd = ' '.join(args.cmd)
    if args.pin:
        cmd = 'numactl --membind={} --cpunodebind={} -- {}'.format(
            args.numa, args.numa, cmd)

    exp_id = uuid.uuid4().hex
    
    start_time = dt.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    end_time = None
    if args.time:
        m = re.search(r'([0-9]+)\:([0-9]{2})', args.time)
        try:
            hours, minutes = map(int, m.groups())
        except AttributeError:
            raise Exception('Could not parse time format')

        end_time = dt.datetime.now() + dt.timedelta(hours=hours, minutes=minutes)
        end_time = end_time.strftime('%Y-%m-%d %H:%M:%S')

    api = ExpInfoAPI(config.DAEMON_IP, config.DAEMON_PORT)

    jobs = api.get_running()
    for job in jobs.values():
        if job['exclusive']:
            print_motd()
            print('{}{}An exclusive job is running, please wait until it has finished.{}'.format(
                bcolors.BOLD, bcolors.FAIL, bcolors.ENDC))
            sys.exit(1)

    if args.exclusive and len(jobs) > 0:
        print_motd()
        print('{}{}Cannot start job exclusively, please wait until all other jobs finished.{}'.format(
            bcolors.BOLD, bcolors.FAIL, bcolors.ENDC))
        sys.exit(1)

    api.register_new(
        exp_id=exp_id,
        user=get_user(),
        start=start_time,
        end=end_time,
        cmd=cmd,
        msg=args.message,
        exclusive=args.exclusive,
        pid=-1,
        numa=args.numa
    )

    print('Starting cmd={}'.format(cmd), file=sys.stderr)

    try:
        p = subprocess.Popen(cmd, shell=True, preexec_fn=os.setsid)
        api.update_running(exp_id=exp_id, pid=p.pid)
        p.wait()
    except KeyboardInterrupt:
        pass

    try:
        os.killpg(os.getpgid(p.pid), signal.SIGTERM)
    except ProcessLookupError:
        pass