artiq/artiq/frontend/artiq_client.py

#!/usr/bin/env python3

import argparse
import time
import asyncio
import sys
from operator import itemgetter
from dateutil.parser import parse as parse_date

from prettytable import PrettyTable

from artiq.protocols.pc_rpc import Client
from artiq.protocols.sync_struct import Subscriber
from artiq.protocols import pyon


def clear_screen():
    sys.stdout.write("\x1b[2J\x1b[H")


def get_argparser():
    parser = argparse.ArgumentParser(description="ARTIQ CLI client")
    parser.add_argument(
        "-s", "--server", default="::1",
        help="hostname or IP of the master to connect to")
    parser.add_argument(
        "--port", default=None, type=int,
        help="TCP port to use to connect to the master")

    subparsers = parser.add_subparsers(dest="action")
    subparsers.required = True

    parser_add = subparsers.add_parser("submit", help="submit an experiment")
    parser_add.add_argument("-p", "--pipeline", default="main", type=str,
                            help="pipeline to run the experiment in "
                                 "(default: %(default)s)")
    parser_add.add_argument("-P", "--priority", default=0, type=int,
                            help="priority (higher value means sooner "
                                 "scheduling, default: %(default)s)")
    parser_add.add_argument("-t", "--timed", default=None, type=str,
                            help="set a due date for the experiment")
    parser_add.add_argument("-f", "--flush", default=False, action="store_true",
                            help="flush the pipeline before preparing "
                            "the experiment")
    parser_add.add_argument("-R", "--repository", default=False,
                            action="store_true",
                            help="use the experiment repository")
    parser_add.add_argument("-r", "--revision", default=None,
                            help="use a specific repository revision "
                                 "(defaults to head, ignored without -R)")
    parser_add.add_argument("-c", "--class-name", default=None,
                            help="name of the class to run")
    parser_add.add_argument("file",
                            help="file containing the experiment to run")
    parser_add.add_argument("arguments", nargs="*",
                            help="run arguments")

    parser_delete = subparsers.add_parser("delete",
                                          help="delete an experiment "
                                               "from the schedule")
    parser_delete.add_argument("rid", type=int,
                               help="run identifier (RID)")

    parser_set_parameter = subparsers.add_parser(
        "set-parameter", help="add or modify a parameter")
    parser_set_parameter.add_argument("name", help="name of the parameter")
    parser_set_parameter.add_argument("value",
                                      help="value in PYON format")

    parser_del_parameter = subparsers.add_parser(
        "del-parameter", help="delete a parameter")
    parser_del_parameter.add_argument("name", help="name of the parameter")

    parser_show = subparsers.add_parser(
        "show", help="show schedule, log, devices or parameters")
    parser_show.add_argument(
        "what",
        help="select object to show: schedule/log/devices/parameters")

    subparsers.add_parser(
        "scan-ddb", help="trigger a device database (re)scan")

    parser_scan_repos = subparsers.add_parser(
        "scan-repository", help="trigger a repository (re)scan")
    parser_scan_repos.add_argument("revision", default=None, nargs="?",
                                   help="use a specific repository revision "
                                        "(defaults to head)")

    return parser


def _parse_arguments(arguments):
    d = {}
    for argument in arguments:
        name, value = argument.split("=")
        d[name] = pyon.decode(value)
    return d


def _action_submit(remote, args):
    try:
        arguments = _parse_arguments(args.arguments)
    except:
        print("Failed to parse run arguments")
        sys.exit(1)

    expid = {
        "file": args.file,
        "class_name": args.class_name,
        "arguments": arguments,
    }
    if args.repository:
        expid["repo_rev"] = args.revision
    if args.timed is None:
        due_date = None
    else:
        due_date = time.mktime(parse_date(args.timed).timetuple())
    rid = remote.submit(args.pipeline, expid,
                        args.priority, due_date, args.flush)
    print("RID: {}".format(rid))


def _action_delete(remote, args):
    remote.delete(args.rid)


def _action_set_parameter(remote, args):
    remote.set(args.name, pyon.decode(args.value))


def _action_del_parameter(remote, args):
    remote.delete(args.name)


def _action_scan_ddb(remote, args):
    remote.scan()


def _action_scan_repository(remote, args):
    remote.scan_async(args.revision)


def _show_schedule(schedule):
    clear_screen()
    if schedule:
        l = sorted(schedule.items(),
                   key=lambda x: (-x[1]["priority"],
                                  x[1]["due_date"] or 0,
                                  x[0]))
        table = PrettyTable(["RID", "Pipeline", "    Status    ", "Prio",
                             "Due date", "Revision", "File", "Class name"])
        for rid, v in l:
            row = [rid, v["pipeline"], v["status"], v["priority"]]
            if v["due_date"] is None:
                row.append("")
            else:
                row.append(time.strftime("%m/%d %H:%M:%S",
                           time.localtime(v["due_date"])))
            expid = v["expid"]
            if "repo_rev" in expid:
                row.append(expid["repo_rev"])
            else:
                row.append("Outside repo.")
            row.append(expid["file"])
            if expid["class_name"] is None:
                row.append("")
            else:
                row.append(expid["class_name"])
            table.add_row(row)
        print(table)
    else:
        print("Schedule is empty")


def _show_devices(devices):
    clear_screen()
    table = PrettyTable(["Name", "Description"])
    table.align["Description"] = "l"
    for k, v in sorted(devices.items(), key=itemgetter(0)):
        table.add_row([k, pyon.encode(v, True)])
    print(table)


def _show_parameters(parameters):
    clear_screen()
    table = PrettyTable(["Parameter", "Value"])
    for k, v in sorted(parameters.items(), key=itemgetter(0)):
        table.add_row([k, str(v)])
    print(table)


def _run_subscriber(host, port, subscriber):
    if port is None:
        port = 3250
    loop = asyncio.get_event_loop()
    try:
        loop.run_until_complete(subscriber.connect(host, port))
        try:
            loop.run_until_complete(asyncio.wait_for(subscriber.receive_task,
                                                     None))
            print("Connection to master lost")
        finally:
            loop.run_until_complete(subscriber.close())
    finally:
        loop.close()


def _show_dict(args, notifier_name, display_fun):
    d = dict()
    def init_d(x):
        d.clear()
        d.update(x)
        return d
    subscriber = Subscriber(notifier_name, init_d,
                            lambda mod: display_fun(d))
    _run_subscriber(args.server, args.port, subscriber)


class _LogPrinter:
    def __init__(self, init):
        for rid, msg in init:
            print(rid, msg)

    def append(self, x):
        rid, msg = x
        print(rid, msg)

    def insert(self, i, x):
        rid, msg = x
        print(rid, msg)

    def pop(self, i=-1):
        pass

    def __delitem__(self, x):
        pass

    def __setitem__(self, k, v):
        pass


def _show_log(args):
    subscriber = Subscriber("log", _LogPrinter)
    _run_subscriber(args.server, args.port, subscriber)


def main():
    args = get_argparser().parse_args()
    action = args.action.replace("-", "_")
    if action == "show":
        if args.what == "schedule":
            _show_dict(args, "schedule", _show_schedule)
        elif args.what == "log":
            _show_log(args)
        elif args.what == "devices":
            _show_dict(args, "devices", _show_devices)
        elif args.what == "parameters":
            _show_dict(args, "parameters", _show_parameters)
        else:
            print("Unknown object to show, use -h to list valid names.")
            sys.exit(1)
    else:
        port = 3251 if args.port is None else args.port
        target_name = {
            "submit": "master_schedule",
            "delete": "master_schedule",
            "set_parameter": "master_pdb",
            "del_parameter": "master_pdb",
            "scan_ddb": "master_ddb",
            "scan_repository": "master_repository"
        }[action]
        remote = Client(args.server, port, target_name)
        try:
            globals()["_action_" + action](remote, args)
        finally:
            remote.close_rpc()

if __name__ == "__main__":
    main()