artiq/artiq/frontend/artiq_devtool.py

#!/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("-H", "--host", metavar="HOSTNAME",
                        type=str, default="lab.m-labs.hk",
                        help="SSH host where the development board is located")
    parser.add_argument("-D", "--device", metavar="HOSTNAME",
                        type=str, default="kc705.lab.m-labs.hk",
                        help="address or domain corresponding to the development board")
    parser.add_argument("-s", "--serial", metavar="PATH",
                        type=str, default="/dev/ttyUSB_kc705",
                        help="TTY device corresponding to the development board")
    parser.add_argument("-l", "--lockfile", metavar="PATH",
                        type=str, default="/run/boards/kc705",
                        help="The lockfile to be acquired for the duration of the script")
    parser.add_argument("-t", "--target", metavar="TARGET",
                        type=str, default="kc705_dds",
                        help="Target to build, one of: "
                             "kc705_dds kc705_drtio_master kc705_drtio_satellite")
    parser.add_argument("-c", "--config", metavar="PATH",
                        type=str, default="openocd-kc705.cfg",
                        help="OpenOCD configuration file corresponding to the development board")

    parser.add_argument("actions", metavar="ACTION",
                        type=str, default=[], nargs="+",
                        help="actions to perform, sequence of: "
                             "build reset boot boot+log connect hotswap clean")

    return parser


def main():
    args = get_argparser().parse_args()
    init_logger(args)
    if args.verbose == args.quiet == 0:
        logging.getLogger().setLevel(logging.INFO)

    if args.target == "kc705_dds" or args.target == "kc705_drtio_master":
        firmware = "runtime"
    elif args.target == "kc705_drtio_satellite":
        firmware = "satman"
    else:
        raise NotImplementedError("unknown target {}".format(args.target))

    client = SSHClient(args.host)
    substs = {
        "env":      "bash -c 'export PATH=$HOME/miniconda/bin:$PATH; exec $0 $*' ",
        "serial":   args.serial,
        "firmware": firmware,
    }

    lock_acquired = False
    def lock():
        nonlocal lock_acquired

        if not lock_acquired:
            logger.info("Acquiring device lock")
            flock = client.spawn_command("flock --verbose --nonblock {} /bin/sleep 86400"
                                            .format(args.lockfile),
                                         get_pty=True)
            flock_file = flock.makefile('r')
            while not lock_acquired:
                line = flock_file.readline()
                if not line:
                    break
                logger.debug(line.rstrip())
                if line.startswith("flock: executing"):
                    lock_acquired = True
                elif line.startswith("flock: failed"):
                    logger.error("Failed to get lock")
                    sys.exit(1)

    for action in args.actions:
        if action == "build":
            logger.info("Building firmware")
            try:
                subprocess.check_call(["python3",
                                        "-m", "artiq.gateware.targets." + args.target,
                                        "--no-compile-gateware",
                                        "--output-dir",
                                        "/tmp/{target}".format(target=args.target)])
            except subprocess.CalledProcessError:
                logger.error("Build failed")
                sys.exit(1)

        elif action == "clean":
            logger.info("Cleaning build directory")
            target_dir = "/tmp/{target}".format(target=args.target)
            if os.path.isdir(target_dir):
                shutil.rmtree(target_dir)

        elif action == "reset":
            lock()

            logger.info("Resetting device")
            client.run_command(
                "{env} artiq_flash start" +
                (" --target-file " + args.config if args.config else ""),
                **substs)

        elif action == "boot" or action == "boot+log":
            lock()

            logger.info("Uploading firmware")
            client.get_sftp().put("/tmp/{target}/software/{firmware}/{firmware}.bin"
                                      .format(target=args.target, firmware=firmware),
                                  "{tmp}/{firmware}.bin"
                                      .format(tmp=client.tmp, firmware=firmware))

            logger.info("Booting firmware")
            flterm = client.spawn_command(
                "{env} python3 flterm.py {serial} " +
                "--kernel {tmp}/{firmware}.bin " +
                ("--upload-only" if action == "boot" else "--output-only"),
                **substs)
            artiq_flash = client.spawn_command(
                "{env} artiq_flash start" +
                (" --target-file " + args.config if args.config else ""),
                **substs)
            client.drain(flterm)

        elif action == "connect":
            lock()

            transport = client.get_transport()

            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(
                "{env} python3 flterm.py {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()