#!/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 from artiq.tools import verbosity_args, init_logger, logger, SSHClient def get_argparser(): parser = argparse.ArgumentParser(description="ARTIQ core device development tool") verbosity_args(parser) parser.add_argument("-t", "--target", metavar="TARGET", type=str, default="kc705_dds", help="Target to build, one of: " "kc705_dds kasli sayma_rtm sayma_amc_standalone " "sayma_amc_drtio_master sayma_amc_drtio_satellite") parser.add_argument("-H", "--host", metavar="HOSTNAME", type=str, default="lab.m-labs.hk", help="SSH host where the development board is located") parser.add_argument('-b', "--board", metavar="BOARD", type=str, default=None, help="Board to connect to on the development SSH host") parser.add_argument("-d", "--device", metavar="DEVICENAME", type=str, default="{board}.{hostname}", help="Address or domain corresponding to the development board") parser.add_argument("-s", "--serial", metavar="SERIAL", type=str, default="/dev/ttyUSB_{board}", help="TTY device corresponding to the development board") parser.add_argument("-l", "--lockfile", metavar="LOCKFILE", type=str, default="/run/boards/{board}", help="The lockfile to be acquired for the duration of the actions") 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 hotswap") return parser def main(): args = get_argparser().parse_args() init_logger(args) if args.verbose == args.quiet == 0: logging.getLogger().setLevel(logging.INFO) build_args = [] if args.target == "kc705_dds": boardtype, firmware = "kc705", "runtime" elif args.target == "sayma_amc_standalone": boardtype, firmware = "sayma", "runtime" build_args += ["--rtm-csr-csv", "/tmp/sayma_rtm/sayma_rtm_csr.csv"] elif args.target == "sayma_amc_drtio_master": boardtype, firmware = "sayma", "runtime" elif args.target == "sayma_amc_drtio_satellite": boardtype, firmware = "sayma", "satman" elif args.target == "sayma_rtm": boardtype, firmware = "sayma_rtm", None else: raise NotImplementedError("unknown target {}".format(args.target)) flash_args = ["-t", boardtype] if boardtype == "sayma": if args.board is None: args.board = "sayma_1" if args.board == "sayma_1": flash_args += ["--preinit-command", "ftdi_location 5:2"] elif args.board == "sayma_2": flash_args += ["--preinit-command", "ftdi_location 3:10"] elif args.board == "sayma_3": flash_args += ["--preinit-command", "ftdi_location 5:1"] else: raise NotImplementedError("unknown --preinit-command for {}".format(boardtype)) client = SSHClient(args.host) substs = { "target": args.target, "hostname": args.host, "boardtype": boardtype, "board": args.board if args.board else boardtype + "_1", "firmware": firmware, } substs.update({ "devicename": args.device.format(**substs), "lockfile": args.lockfile.format(**substs), "serial": args.serial.format(**substs), }) flock_acquired = False flock_file = None # GC root def lock(): nonlocal flock_acquired nonlocal flock_file if not flock_acquired: logger.info("Acquiring device lock") flock = client.spawn_command("flock --verbose {block} {lockfile} sleep 86400" .format(block="" if args.wait else "--nonblock", **substs), 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 artiq_flash(args, synchronous=True): args = flash_args + args args = ["'{}'".format(arg) if " " in arg else arg for arg in args] cmd = client.spawn_command( "artiq_flash " + " ".join(args), **substs) if synchronous: client.drain(cmd) else: return cmd for action in args.actions: if action == "build": logger.info("Building target") try: subprocess.check_call(["python3", "-m", "artiq.gateware.targets." + args.target, "--no-compile-gateware", *build_args, "--output-dir", "/tmp/{target}".format(**substs)]) except subprocess.CalledProcessError: logger.error("Build failed") sys.exit(1) elif action == "clean": logger.info("Cleaning build directory") target_dir = "/tmp/{target}".format(**substs) if os.path.isdir(target_dir): shutil.rmtree(target_dir) elif action == "reset": logger.info("Resetting device") artiq_flash(["reset"]) elif action == "flash" or action == "flash+log": def upload_product(product, ext): logger.info("Uploading {}".format(product)) client.get_sftp().put("/tmp/{target}/software/{product}/{product}.{ext}" .format(target=args.target, product=product, ext=ext), "{tmp}/{product}.{ext}" .format(tmp=client.tmp, product=product, ext=ext)) upload_product("bootloader", "bin") upload_product(firmware, "fbi") logger.info("Flashing firmware") artiq_flash(["-d", "{tmp}", "proxy", "bootloader", "firmware", "start" if action == "flash" else ""]) if action == "flash+log": logger.info("Booting firmware") flterm = client.spawn_command( "flterm {serial} " + "--kernel {tmp}/{firmware}.bin " + ("--upload-only" if action == "boot" else "--output-only"), **substs) artiq_flash(["start"], synchronous=False) client.drain(flterm) 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, args.device, port) try: remote_stream = \ transport.open_channel('direct-tcpip', (args.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", **substs) elif action == "hotswap": logger.info("Hotswapping firmware") try: subprocess.check_call(["python3", "-m", "artiq.frontend.artiq_coreboot", "hotswap", "/tmp/{target}/software/{firmware}/{firmware}.bin" .format(target=args.target, firmware=firmware)]) except subprocess.CalledProcessError: logger.error("Build failed") sys.exit(1) else: logger.error("Unknown action {}".format(action)) sys.exit(1) if __name__ == "__main__": main()