mirror of https://github.com/m-labs/artiq.git
155 lines
5.1 KiB
Python
Executable File
155 lines
5.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import time
|
|
import asyncio
|
|
|
|
from prettytable import PrettyTable
|
|
|
|
from artiq.management.pc_rpc import Client
|
|
from artiq.management.sync_struct import Subscriber
|
|
from artiq.management.tools import clear_screen
|
|
|
|
|
|
def _get_args():
|
|
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", "--periodic", default=None, type=float,
|
|
help="run the experiment periodically every given number of seconds")
|
|
parser_add.add_argument(
|
|
"-t", "--timeout", default=None, type=float,
|
|
help="specify a timeout for the experiment to complete")
|
|
parser_add.add_argument("-f", "--function", default="run",
|
|
help="function to run")
|
|
parser_add.add_argument("-u", "--unit", default=None,
|
|
help="unit to run")
|
|
parser_add.add_argument("file", help="file containing the unit to run")
|
|
|
|
parser_cancel = subparsers.add_parser("cancel",
|
|
help="cancel an experiment")
|
|
parser_cancel.add_argument("-p", "--periodic", default=False,
|
|
action="store_true",
|
|
help="cancel a periodic experiment")
|
|
parser_cancel.add_argument("rid", type=int,
|
|
help="run identifier (RID/PRID)")
|
|
|
|
parser_show_queue = subparsers.add_parser(
|
|
"show-queue", help="show the experiment queue")
|
|
|
|
parser_show_periodic = subparsers.add_parser(
|
|
"show-periodic", help="show the periodic experiment table")
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
def _action_submit(remote, args):
|
|
run_params = {
|
|
"file": args.file,
|
|
"unit": args.unit,
|
|
"function": args.function
|
|
}
|
|
if args.periodic is None:
|
|
rid = remote.run_once(run_params, args.timeout)
|
|
print("RID: {}".format(rid))
|
|
else:
|
|
prid = remote.run_periodic(run_params, args.timeout,
|
|
args.periodic)
|
|
print("PRID: {}".format(prid))
|
|
|
|
|
|
def _action_cancel(remote, args):
|
|
if args.periodic:
|
|
remote.cancel_periodic(args.rid)
|
|
else:
|
|
remote.cancel_once(args.rid)
|
|
|
|
|
|
def _show_queue(queue):
|
|
clear_screen()
|
|
if queue:
|
|
table = PrettyTable(["RID", "File", "Unit", "Function", "Timeout"])
|
|
for rid, run_params, timeout in queue:
|
|
row = [rid, run_params["file"]]
|
|
for x in run_params["unit"], run_params["function"], timeout:
|
|
row.append("-" if x is None else x)
|
|
table.add_row(row)
|
|
print(table)
|
|
else:
|
|
print("Queue is empty")
|
|
|
|
|
|
def _show_periodic(periodic):
|
|
clear_screen()
|
|
if periodic:
|
|
table = PrettyTable(["Next run", "PRID", "File", "Unit", "Function",
|
|
"Timeout", "Period"])
|
|
sp = sorted(periodic.items(), key=lambda x: (x[1][0], x[0]))
|
|
for prid, (next_run, run_params, timeout, period) in sp:
|
|
row = [time.strftime("%m/%d %H:%M:%S", time.localtime(next_run)),
|
|
prid, run_params["file"]]
|
|
for x in run_params["unit"], run_params["function"], timeout:
|
|
row.append("-" if x is None else x)
|
|
row.append(period)
|
|
table.add_row(row)
|
|
print(table)
|
|
else:
|
|
print("No periodic schedule")
|
|
|
|
|
|
def _run_subscriber(host, port, subscriber):
|
|
if port is None:
|
|
port = 8887
|
|
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 main():
|
|
args = _get_args()
|
|
if args.action == "show-queue":
|
|
queue = []
|
|
def init_queue(x):
|
|
queue[:] = x
|
|
return queue
|
|
subscriber = Subscriber("queue", init_queue,
|
|
lambda: _show_queue(queue))
|
|
_run_subscriber(args.server, args.port, subscriber)
|
|
elif args.action == "show-periodic":
|
|
periodic = dict()
|
|
def init_periodic(x):
|
|
periodic.clear()
|
|
periodic.update(x)
|
|
return periodic
|
|
subscriber = Subscriber("periodic", init_periodic,
|
|
lambda: _show_periodic(periodic))
|
|
_run_subscriber(args.server, args.port, subscriber)
|
|
else:
|
|
port = 8888 if args.port is None else args.port
|
|
remote = Client(args.server, port, "master_schedule")
|
|
try:
|
|
globals()["_action_" + args.action](remote, args)
|
|
finally:
|
|
remote.close_rpc()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|