forked from M-Labs/artiq
1
0
Fork 0
artiq/artiq/frontend/artiq_devtool.py

284 lines
11 KiB
Python
Executable File

#!/usr/bin/env python3
# This script makes the following assumptions:
# * miniconda is installed remotely at ~/miniconda
# * misoc and artiq are installed remotely via conda
import sys
import argparse
import logging
import subprocess
import socket
import select
import threading
import os
import shutil
import re
import shlex
from artiq.tools import verbosity_args, init_logger
from artiq.remoting import SSHClient
from artiq.coredevice.comm_mgmt import CommMgmt
logger = logging.getLogger(__name__)
def get_argparser():
parser = argparse.ArgumentParser(
description="ARTIQ core device development tool",
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
verbosity_args(parser)
parser.add_argument("-t", "--target", metavar="TARGET",
type=str, default="kasli",
help="target to build, one of: "
"kc705 kasli sayma")
parser.add_argument("-V", "--variant", metavar="VARIANT",
type=str, default=None,
help="variant to build, dependent on the target")
parser.add_argument("-g", "--gateware",
default=False, action="store_true",
help="build/flash gateware, not just software")
parser.add_argument("--args", metavar="ARGS",
type=shlex.split, default=[],
help="extra arguments for gateware/firmware build")
parser.add_argument("-H", "--host",
type=str, default="lab.m-labs.hk",
help="SSH host where the development board is located")
parser.add_argument("-J", "--jump",
type=str, default=None,
help="SSH host to jump through")
parser.add_argument("-b", "--board",
type=str, default="{board_type}-1",
help="board to connect to on the development SSH host")
parser.add_argument("-B", "--board-file",
type=str, default="/var/lib/artiq/boards/{board}",
help="the board file containing the openocd initialization commands; "
"it is also used as the lock file")
parser.add_argument("-s", "--serial",
type=str, default="/dev/ttyUSB_{board}",
help="TTY device corresponding to the development board")
parser.add_argument("-d", "--device",
type=str, default="{board}",
help="address or domain corresponding to the development board")
parser.add_argument("-w", "--wait", action="store_true",
help="wait for the board to unlock instead of aborting the actions")
parser.add_argument("actions", metavar="ACTION",
type=str, default=[], nargs="+",
help="actions to perform, sequence of: "
"build clean reset flash flash+log load connect hotswap")
return parser
def main():
args = get_argparser().parse_args()
init_logger(args)
if args.verbose == args.quiet == 0:
logging.getLogger().setLevel(logging.INFO)
def build_dir(*path, target=args.target):
return os.path.join("/tmp/", "artiq_" + target, *path)
if args.target == "kc705":
board_type, firmware = "kc705", "runtime"
variant = "nist_clock" if args.variant is None else args.variant
elif args.target == "sayma":
board_type, firmware = "sayma", "runtime"
variant = "standalone" if args.variant is None else args.variant
elif args.target == "kasli":
board_type, firmware = "kasli", "runtime"
variant = "tester" if args.variant is None else args.variant
else:
raise NotImplementedError("unknown target {}".format(args.target))
board = args.board.format(board_type=board_type)
board_file = args.board_file.format(board=board)
device = args.device.format(board=board, host=args.host)
serial = args.serial.format(board=board)
client = SSHClient(args.host, args.jump)
flock_acquired = False
flock_file = None # GC root
def lock():
nonlocal flock_acquired
nonlocal flock_file
if not flock_acquired:
fuser_args = ["fuser", "-u", board_file]
fuser = client.spawn_command(fuser_args)
fuser_file = fuser.makefile('r')
fuser_match = re.search(r"\((.+?)\)", fuser_file.readline())
if fuser_match and fuser_match.group(1) == os.getenv("USER"):
logger.info("Lock already acquired by {}".format(os.getenv("USER")))
flock_acquired = True
return
logger.info("Acquiring device lock")
flock_args = ["flock"]
if not args.wait:
flock_args.append("--nonblock")
flock_args += ["--verbose", board_file]
flock_args += ["sleep", "86400"]
flock = client.spawn_command(flock_args, get_pty=True)
flock_file = flock.makefile('r')
while not flock_acquired:
line = flock_file.readline()
if not line:
break
logger.debug(line.rstrip())
if line.startswith("flock: executing"):
flock_acquired = True
elif line.startswith("flock: failed"):
logger.error("Failed to get lock")
sys.exit(1)
def command(*args, on_failure="Command failed"):
logger.debug("Running {}".format(" ".join([shlex.quote(arg) for arg in args])))
try:
subprocess.check_call(args)
except subprocess.CalledProcessError:
logger.error(on_failure)
sys.exit(1)
def build(target, *extra_args, output_dir=build_dir(), variant=variant):
build_args = ["python3", "-m", "artiq.gateware.targets." + target, *extra_args]
if not args.gateware:
build_args.append("--no-compile-gateware")
if variant:
build_args += ["--variant", variant]
build_args += ["--output-dir", output_dir]
command(*build_args, on_failure="Build failed")
def flash(*steps):
lock()
flash_args = ["artiq_flash"]
for _ in range(args.verbose):
flash_args.append("-v")
flash_args += ["-H", args.host]
if args.jump:
flash_args += ["-J", args.jump]
flash_args += ["-t", board_type]
flash_args += ["-V", variant]
flash_args += ["-I", "source {}".format(board_file)]
flash_args += ["--srcbuild", build_dir()]
flash_args += steps
command(*flash_args, on_failure="Flashing failed")
for action in args.actions:
if action == "build":
logger.info("Building target")
if args.target == "sayma":
build("sayma_rtm", output_dir=build_dir("rtm_gateware"), variant=None)
build("sayma_amc", "--rtm-csr-csv", build_dir("rtm_gateware", "rtm_csr.csv"))
else:
build(args.target)
elif action == "clean":
logger.info("Cleaning build directory")
shutil.rmtree(build_dir(), ignore_errors=True)
elif action == "reset":
logger.info("Resetting device")
flash("start")
elif action == "flash":
gateware = ["gateware"] if args.gateware else []
logger.info("Flashing and booting")
flash(*gateware, "bootloader", "firmware", "start")
elif action == "flash+log":
gateware = ["gateware"] if args.gateware else []
logger.info("Flashing")
flash(*gateware, "bootloader", "firmware")
flterm = client.spawn_command(["flterm", serial, "--output-only"])
logger.info("Booting")
flash("start")
client.drain(flterm)
elif action == "load":
logger.info("Loading gateware")
flash("load")
elif action == "connect":
lock()
transport = client.get_transport()
transport.set_keepalive(30)
def forwarder(local_stream, remote_stream):
try:
while True:
r, _, _ = select.select([local_stream, remote_stream], [], [])
if local_stream in r:
data = local_stream.recv(65535)
if data == b"":
break
remote_stream.sendall(data)
if remote_stream in r:
data = remote_stream.recv(65535)
if data == b"":
break
local_stream.sendall(data)
except Exception as err:
logger.error("Cannot forward on port %s: %s", port, repr(err))
local_stream.close()
remote_stream.close()
def listener(port):
listener = socket.socket()
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
listener.bind(('localhost', port))
listener.listen(8)
while True:
local_stream, peer_addr = listener.accept()
logger.info("Accepting %s:%s and opening SSH channel to %s:%s",
*peer_addr, device, port)
try:
remote_stream = \
transport.open_channel('direct-tcpip', (device, port), peer_addr)
except Exception:
logger.exception("Cannot open channel on port %s", port)
continue
thread = threading.Thread(target=forwarder, args=(local_stream, remote_stream),
name="forward-{}".format(port), daemon=True)
thread.start()
ports = [1380, 1381, 1382, 1383]
for port in ports:
thread = threading.Thread(target=listener, args=(port,),
name="listen-{}".format(port), daemon=True)
thread.start()
logger.info("Forwarding ports {} to core device and logs from core device"
.format(", ".join(map(str, ports))))
client.run_command(["flterm", serial, "--output-only"])
elif action == "hotswap":
lock()
logger.info("Hotswapping firmware")
firmware = build_dir(variant, "software", firmware, firmware + ".bin")
mgmt = CommMgmt(device)
mgmt.open(ssh_transport=client.get_transport())
with open(firmware, "rb") as f:
mgmt.hotswap(f.read())
else:
logger.error("Unknown action {}".format(action))
sys.exit(1)
if __name__ == "__main__":
main()