mirror of
https://github.com/m-labs/artiq.git
synced 2024-12-28 20:53:35 +08:00
artiq_flash: implement network transparency.
This commit is contained in:
whitequark
parent
80cbef0031
commit
b553804e5a
@ -18,7 +18,9 @@ from artiq.tools import verbosity_args, init_logger, logger, SSHClient
|
||||
|
||||
|
||||
def get_argparser():
|
||||
parser = argparse.ArgumentParser(description="ARTIQ core device development tool")
|
||||
parser = argparse.ArgumentParser(
|
||||
description="ARTIQ core device development tool",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
|
||||
verbosity_args(parser)
|
||||
|
||||
@ -27,19 +29,19 @@ def get_argparser():
|
||||
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",
|
||||
parser.add_argument("-H", "--host",
|
||||
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,
|
||||
parser.add_argument('-b', "--board",
|
||||
type=str, default="{boardtype}-1",
|
||||
help="Board to connect to on the development SSH host")
|
||||
parser.add_argument("-d", "--device", metavar="DEVICENAME",
|
||||
type=str, default="{board}.{hostname}",
|
||||
parser.add_argument("-d", "--device",
|
||||
type=str, default="{board}.{host}",
|
||||
help="Address or domain corresponding to the development board")
|
||||
parser.add_argument("-s", "--serial", metavar="SERIAL",
|
||||
parser.add_argument("-s", "--serial",
|
||||
type=str, default="/dev/ttyUSB_{board}",
|
||||
help="TTY device corresponding to the development board")
|
||||
parser.add_argument("-l", "--lockfile", metavar="LOCKFILE",
|
||||
parser.add_argument("-l", "--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",
|
||||
@ -59,47 +61,30 @@ def main():
|
||||
if args.verbose == args.quiet == 0:
|
||||
logging.getLogger().setLevel(logging.INFO)
|
||||
|
||||
def build_dir(*path, target=args.target):
|
||||
return os.path.join("/tmp", target, *path)
|
||||
|
||||
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"]
|
||||
boardtype, firmware = "sayma_amc", "runtime"
|
||||
build_args += ["--rtm-csr-csv", build_dir("sayma_rtm_csr.csv", target="sayma_rtm")]
|
||||
elif args.target == "sayma_amc_drtio_master":
|
||||
boardtype, firmware = "sayma", "runtime"
|
||||
boardtype, firmware = "sayma_amc", "runtime"
|
||||
elif args.target == "sayma_amc_drtio_satellite":
|
||||
boardtype, firmware = "sayma", "satman"
|
||||
boardtype, firmware = "sayma_amc", "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))
|
||||
board = args.board.format(boardtype=boardtype)
|
||||
device = args.device.format(board=board, host=args.host)
|
||||
lockfile = args.lockfile.format(board=board)
|
||||
serial = args.serial.format(board=board)
|
||||
|
||||
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
|
||||
@ -109,10 +94,13 @@ def main():
|
||||
|
||||
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_args = ["flock"]
|
||||
if not args.wait:
|
||||
flock_args.append("--nonblock")
|
||||
flock_args += ["--verbose", lockfile]
|
||||
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()
|
||||
@ -125,65 +113,52 @@ def main():
|
||||
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
|
||||
def flash(*steps):
|
||||
flash_args = ["artiq_flash"]
|
||||
for _ in range(args.verbose):
|
||||
flash_args.append("-v")
|
||||
flash_args += ["-H", args.host, "-t", boardtype]
|
||||
flash_args += ["--srcbuild", build_dir()]
|
||||
flash_args += ["--preinit-command", "source /var/boards/{}".format(board)]
|
||||
flash_args += steps
|
||||
subprocess.check_call(flash_args)
|
||||
|
||||
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)])
|
||||
subprocess.check_call([
|
||||
"python3", "-m", "artiq.gateware.targets." + args.target,
|
||||
"--no-compile-gateware",
|
||||
*build_args,
|
||||
"--output-dir", build_dir()])
|
||||
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)
|
||||
shutil.rmtree(build_dir, ignore_errors=True)
|
||||
|
||||
elif action == "reset":
|
||||
lock()
|
||||
|
||||
logger.info("Resetting device")
|
||||
artiq_flash(["reset"])
|
||||
flash("start")
|
||||
|
||||
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")
|
||||
lock()
|
||||
|
||||
logger.info("Flashing firmware")
|
||||
artiq_flash(["-d", "{tmp}", "proxy", "bootloader", "firmware",
|
||||
"start" if action == "flash" else ""])
|
||||
flash("proxy", "bootloader", "firmware")
|
||||
|
||||
logger.info("Booting firmware")
|
||||
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)
|
||||
flterm = client.spawn_command(["flterm", serial, "--output-only"])
|
||||
flash("start")
|
||||
client.drain(flterm)
|
||||
else:
|
||||
flash("start")
|
||||
|
||||
elif action == "connect":
|
||||
lock()
|
||||
@ -218,10 +193,10 @@ def main():
|
||||
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)
|
||||
*peer_addr, device, port)
|
||||
try:
|
||||
remote_stream = \
|
||||
transport.open_channel('direct-tcpip', (args.device, port), peer_addr)
|
||||
transport.open_channel('direct-tcpip', (device, port), peer_addr)
|
||||
except Exception:
|
||||
logger.exception("Cannot open channel on port %s", port)
|
||||
continue
|
||||
@ -238,17 +213,13 @@ def main():
|
||||
|
||||
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)
|
||||
client.run_command(["flterm", serial, "--output-only"])
|
||||
|
||||
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)])
|
||||
subprocess.check_call(["artiq_coreboot", "hotswap",
|
||||
build_dir("software", firmware, firmware + ".bin")])
|
||||
except subprocess.CalledProcessError:
|
||||
logger.error("Build failed")
|
||||
sys.exit(1)
|
||||
|
||||
@ -5,12 +5,13 @@ import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
import shutil
|
||||
import re
|
||||
from functools import partial
|
||||
|
||||
from artiq import __artiq_dir__ as artiq_dir
|
||||
from artiq.tools import verbosity_args, init_logger, logger, SSHClient, LocalClient
|
||||
from artiq.frontend.bit2bin import bit2bin
|
||||
|
||||
|
||||
def get_argparser():
|
||||
parser = argparse.ArgumentParser(
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
@ -35,6 +36,12 @@ Prerequisites:
|
||||
and replug the device. Ensure you are member of the
|
||||
plugdev group: 'sudo adduser $USER plugdev' and re-login.
|
||||
""")
|
||||
|
||||
verbosity_args(parser)
|
||||
|
||||
parser.add_argument("-H", "--host", metavar="HOSTNAME",
|
||||
type=str, default=None,
|
||||
help="SSH host where the development board is located")
|
||||
parser.add_argument("-t", "--target", default="kc705",
|
||||
help="target board, default: %(default)s, one of: "
|
||||
"kc705 kasli sayma_amc sayma_rtm")
|
||||
@ -73,14 +80,34 @@ def proxy_path():
|
||||
|
||||
|
||||
class Programmer:
|
||||
def __init__(self, target_file, preinit_commands):
|
||||
self.target_file = target_file
|
||||
def __init__(self, client, target_file, preinit_commands):
|
||||
self.client = client
|
||||
if target_file:
|
||||
self.target_file = self._transfer_script(target_file)
|
||||
else:
|
||||
self.target_file = None
|
||||
self.preinit_commands = preinit_commands
|
||||
self.prog = []
|
||||
|
||||
def _transfer_script(self, script):
|
||||
if isinstance(self.client, LocalClient):
|
||||
return script
|
||||
|
||||
def rewriter(content):
|
||||
def repl(match):
|
||||
return "[find {}]".format(self._transfer_script(match.group(1)))
|
||||
return re.sub(r"\[find (.+?)\]", repl, content, re.DOTALL)
|
||||
|
||||
script = os.path.join(scripts_path(), script)
|
||||
return self.client.transfer_file(script, rewriter)
|
||||
|
||||
def _command(self, cmd):
|
||||
self.prog.append(cmd.replace("{", "{{").replace("}", "}}"))
|
||||
|
||||
def init(self):
|
||||
self.prog.extend(self.preinit_commands)
|
||||
self.prog.append("init")
|
||||
for command in self.preinit_commands:
|
||||
self._command(command)
|
||||
self._command("init")
|
||||
|
||||
def load(self, bitfile):
|
||||
raise NotImplementedError
|
||||
@ -95,46 +122,49 @@ class Programmer:
|
||||
raise NotImplementedError
|
||||
|
||||
def do(self):
|
||||
self.prog.append("exit")
|
||||
cmdline = [
|
||||
"openocd",
|
||||
"-s", scripts_path()
|
||||
]
|
||||
self._command("exit")
|
||||
|
||||
cmdline = ["openocd"]
|
||||
if isinstance(self.client, LocalClient):
|
||||
cmdline += ["-s", scripts_path()]
|
||||
if self.target_file is not None:
|
||||
cmdline += ["-f", self.target_file]
|
||||
cmdline += ["-c", "; ".join(self.prog)]
|
||||
subprocess.check_call(cmdline)
|
||||
|
||||
self.client.run_command(cmdline)
|
||||
|
||||
|
||||
class ProgrammerJtagSpi7(Programmer):
|
||||
def __init__(self, target, preinit_commands):
|
||||
Programmer.__init__(self, os.path.join("board", target + ".cfg"),
|
||||
def __init__(self, client, target, preinit_commands):
|
||||
Programmer.__init__(self, client, os.path.join("board", target + ".cfg"),
|
||||
preinit_commands)
|
||||
self.init()
|
||||
|
||||
def load(self, bitfile, pld=0):
|
||||
self.prog.append("pld load {} {{{}}}".format(pld, bitfile))
|
||||
bitfile = self.client.transfer_file(bitfile)
|
||||
self._command("pld load {} {{{}}}".format(pld, bitfile))
|
||||
|
||||
def proxy(self, proxy_bitfile, pld=0):
|
||||
self.prog.append("jtagspi_init {} {{{}}}".format(pld, proxy_bitfile))
|
||||
proxy_bitfile = self.client.transfer_file(proxy_bitfile)
|
||||
self._command("jtagspi_init {} {{{}}}".format(pld, proxy_bitfile))
|
||||
|
||||
def flash_binary(self, flashno, address, filename):
|
||||
# jtagspi_program supports only one flash
|
||||
assert flashno == 0
|
||||
self.prog.append("jtagspi_program {{{}}} 0x{:x}".format(
|
||||
filename = self.client.transfer_file(filename)
|
||||
self._command("jtagspi_program {{{}}} 0x{:x}".format(
|
||||
filename, address))
|
||||
|
||||
def start(self):
|
||||
self.prog.append("xc7_program xc7.tap")
|
||||
self._command("xc7_program xc7.tap")
|
||||
|
||||
|
||||
class ProgrammerSayma(Programmer):
|
||||
sector_size = 0x10000
|
||||
|
||||
def __init__(self, preinit_commands):
|
||||
# TODO: support Sayma RTM
|
||||
Programmer.__init__(self, None, preinit_commands)
|
||||
self.proxy_loaded = False
|
||||
def __init__(self, client, preinit_commands):
|
||||
Programmer.__init__(self, client, None, preinit_commands)
|
||||
|
||||
self.prog += [
|
||||
"interface ftdi",
|
||||
"ftdi_device_desc \"Quad RS232-HS\"",
|
||||
@ -161,11 +191,12 @@ class ProgrammerSayma(Programmer):
|
||||
self.init()
|
||||
|
||||
def load(self, bitfile, pld=1):
|
||||
self.prog.append("pld load {} {{{}}}".format(pld, bitfile))
|
||||
bitfile = self.client.transfer_file(bitfile)
|
||||
self._command("pld load {} {{{}}}".format(pld, bitfile))
|
||||
|
||||
def proxy(self, proxy_bitfile, pld=1):
|
||||
self.load(proxy_bitfile, pld)
|
||||
self.prog.append("reset halt")
|
||||
self._command("reset halt")
|
||||
|
||||
def flash_binary(self, flashno, address, filename):
|
||||
sector_first = address // self.sector_size
|
||||
@ -173,25 +204,23 @@ class ProgrammerSayma(Programmer):
|
||||
assert size
|
||||
sector_last = sector_first + (size - 1) // self.sector_size
|
||||
assert sector_last >= sector_first
|
||||
self.prog += [
|
||||
"flash probe xcu.spi{}".format(flashno),
|
||||
"flash erase_sector {} {} {}".format(flashno, sector_first, sector_last),
|
||||
"flash write_bank {} {{{}}} 0x{:x}".format(flashno, filename, address),
|
||||
"flash verify_bank {} {{{}}} 0x{:x}".format(flashno, filename, address),
|
||||
]
|
||||
filename = self.client.transfer_file(filename)
|
||||
self._command("flash probe xcu.spi{}".format(flashno))
|
||||
self._command("flash erase_sector {} {} {}".format(flashno, sector_first, sector_last))
|
||||
self._command("flash write_bank {} {{{}}} 0x{:x}".format(flashno, filename, address))
|
||||
self._command("flash verify_bank {} {{{}}} 0x{:x}".format(flashno, filename, address))
|
||||
|
||||
def start(self):
|
||||
self.proxy_loaded = False
|
||||
self.prog.append("xcu_program xcu.tap")
|
||||
self._command("xcu_program xcu.tap")
|
||||
|
||||
|
||||
def main():
|
||||
parser = get_argparser()
|
||||
opts = parser.parse_args()
|
||||
args = get_argparser().parse_args()
|
||||
init_logger(args)
|
||||
|
||||
config = {
|
||||
"kc705": {
|
||||
"programmer_factory": partial(ProgrammerJtagSpi7, "kc705"),
|
||||
"programmer_factory": partial(ProgrammerJtagSpi7, target="kc705"),
|
||||
"proxy_bitfile": "bscan_spi_xc7k325t.bit",
|
||||
"variants": ["nist_clock", "nist_qc2"],
|
||||
"gateware": (0, 0x000000),
|
||||
@ -200,7 +229,7 @@ def main():
|
||||
"firmware": (0, 0xb40000),
|
||||
},
|
||||
"kasli": {
|
||||
"programmer_factory": partial(ProgrammerJtagSpi7, "kasli"),
|
||||
"programmer_factory": partial(ProgrammerJtagSpi7, target="kasli"),
|
||||
"proxy_bitfile": "bscan_spi_xc7a100t.bit",
|
||||
"variants": ["opticlock"],
|
||||
"gateware": (0, 0x000000),
|
||||
@ -222,29 +251,34 @@ def main():
|
||||
"proxy_bitfile": "bscan_spi_xcku040-sayma.bit",
|
||||
"gateware": (1, 0x150000),
|
||||
},
|
||||
}[opts.target]
|
||||
}[args.target]
|
||||
|
||||
variant = opts.variant
|
||||
variant = args.variant
|
||||
if "variants" in config:
|
||||
if variant is not None and variant not in config["variants"]:
|
||||
raise SystemExit("Invalid variant for this board")
|
||||
if variant is None:
|
||||
variant = config["variants"][0]
|
||||
bin_dir = opts.dir
|
||||
bin_dir = args.dir
|
||||
if bin_dir is None:
|
||||
if variant:
|
||||
bin_name = "{}-{}".format(opts.target, variant)
|
||||
bin_name = "{}-{}".format(args.target, variant)
|
||||
else:
|
||||
bin_name = opts.target
|
||||
bin_name = args.target
|
||||
bin_dir = os.path.join(artiq_dir, "binaries", bin_name)
|
||||
if opts.srcbuild is None and not os.path.exists(bin_dir) and opts.action != ["start"]:
|
||||
if args.srcbuild is None and not os.path.exists(bin_dir) and args.action != ["start"]:
|
||||
raise SystemExit("Binaries directory '{}' does not exist"
|
||||
.format(bin_dir))
|
||||
|
||||
programmer = config["programmer_factory"](opts.preinit_command)
|
||||
if args.host is None:
|
||||
client = LocalClient()
|
||||
else:
|
||||
client = SSHClient(args.host)
|
||||
|
||||
programmer = config["programmer_factory"](client, preinit_commands=args.preinit_command)
|
||||
|
||||
conv = False
|
||||
for action in opts.action:
|
||||
for action in args.action:
|
||||
if action == "proxy":
|
||||
proxy_found = False
|
||||
for p in [bin_dir, proxy_path(), os.path.expanduser("~/.migen"),
|
||||
@ -258,10 +292,10 @@ def main():
|
||||
raise SystemExit(
|
||||
"proxy gateware bitstream {} not found".format(config["proxy_bitfile"]))
|
||||
elif action == "gateware":
|
||||
if opts.srcbuild is None:
|
||||
if args.srcbuild is None:
|
||||
path = bin_dir
|
||||
else:
|
||||
path = os.path.join(opts.srcbuild, "gateware")
|
||||
path = os.path.join(args.srcbuild, "gateware")
|
||||
bin = os.path.join(path, "top.bin")
|
||||
if not os.access(bin, os.R_OK):
|
||||
bin_handle, bin = tempfile.mkstemp()
|
||||
@ -271,29 +305,29 @@ def main():
|
||||
conv = True
|
||||
programmer.flash_binary(*config["gateware"], bin)
|
||||
elif action == "bootloader":
|
||||
if opts.srcbuild is None:
|
||||
if args.srcbuild is None:
|
||||
path = bin_dir
|
||||
else:
|
||||
path = os.path.join(opts.srcbuild, "software", "bootloader")
|
||||
path = os.path.join(args.srcbuild, "software", "bootloader")
|
||||
programmer.flash_binary(*config["bootloader"], os.path.join(path, "bootloader.bin"))
|
||||
elif action == "storage":
|
||||
programmer.flash_binary(*config["storage"], opts.storage)
|
||||
programmer.flash_binary(*config["storage"], args.storage)
|
||||
elif action == "firmware":
|
||||
if variant == "satellite":
|
||||
firmware_name = "satman"
|
||||
else:
|
||||
firmware_name = "runtime"
|
||||
if opts.srcbuild is None:
|
||||
if args.srcbuild is None:
|
||||
path = bin_dir
|
||||
else:
|
||||
path = os.path.join(opts.srcbuild, "software", firmware_name)
|
||||
path = os.path.join(args.srcbuild, "software", firmware_name)
|
||||
programmer.flash_binary(*config["firmware"],
|
||||
os.path.join(path, firmware_name + ".fbi"))
|
||||
elif action == "load":
|
||||
if opts.srcbuild is None:
|
||||
if args.srcbuild is None:
|
||||
path = bin_dir
|
||||
else:
|
||||
path = os.path.join(opts.srcbuild, "gateware")
|
||||
path = os.path.join(args.srcbuild, "gateware")
|
||||
programmer.load(os.path.join(path, "top.bit"))
|
||||
elif action == "start":
|
||||
programmer.start()
|
||||
|
||||
@ -3,10 +3,9 @@ import logging
|
||||
import sys
|
||||
import asyncio
|
||||
import collections
|
||||
import os
|
||||
import atexit
|
||||
import string
|
||||
import random
|
||||
import os, random, tempfile, shutil, shlex, subprocess
|
||||
|
||||
import numpy as np
|
||||
|
||||
@ -256,7 +255,45 @@ def get_user_config_dir():
|
||||
return dir
|
||||
|
||||
|
||||
class SSHClient:
|
||||
class Client:
|
||||
def transfer_file(self, filename, rewriter=None):
|
||||
raise NotImplementedError
|
||||
|
||||
def run_command(self, cmd, **kws):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class LocalClient(Client):
|
||||
def __init__(self):
|
||||
tmpname = "".join([random.Random().choice("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
|
||||
for _ in range(6)])
|
||||
self.tmp = os.path.join(tempfile.gettempdir(), "artiq" + tmpname)
|
||||
self._has_tmp = False
|
||||
|
||||
def _prepare_tmp(self):
|
||||
if not self._has_tmp:
|
||||
os.mkdir(self.tmp)
|
||||
atexit.register(lambda: shutil.rmtree(self.tmp, ignore_errors=True))
|
||||
self._has_tmp = True
|
||||
|
||||
def transfer_file(self, filename, rewriter=None):
|
||||
logger.debug("Transferring {}".format(filename))
|
||||
if rewriter is None:
|
||||
return filename
|
||||
else:
|
||||
tmp_filename = os.path.join(self.tmp, filename.replace(os.sep, "_"))
|
||||
with open(filename) as local:
|
||||
self._prepare_tmp()
|
||||
with open(tmp_filename, 'w') as tmp:
|
||||
tmp.write(rewriter(local.read()))
|
||||
return tmp_filename
|
||||
|
||||
def run_command(self, cmd, **kws):
|
||||
logger.debug("Executing {}".format(cmd))
|
||||
subprocess.check_call([arg.format(tmp=self.tmp, **kws) for arg in cmd])
|
||||
|
||||
|
||||
class SSHClient(Client):
|
||||
def __init__(self, host):
|
||||
self.host = host
|
||||
self.ssh = None
|
||||
@ -284,16 +321,28 @@ class SSHClient:
|
||||
if self.sftp is None:
|
||||
self.sftp = self.get_ssh().open_sftp()
|
||||
self.sftp.mkdir(self.tmp)
|
||||
atexit.register(lambda: self.run_command("rm -rf {tmp}"))
|
||||
atexit.register(lambda: self.run_command(["rm", "-rf", "{tmp}"]))
|
||||
return self.sftp
|
||||
|
||||
def transfer_file(self, filename, rewriter=None):
|
||||
remote_filename = "{}/{}".format(self.tmp, filename.replace("/", "_"))
|
||||
logger.debug("Transferring {}".format(filename))
|
||||
if rewriter is None:
|
||||
self.get_sftp().put(filename, remote_filename)
|
||||
else:
|
||||
with open(filename) as local:
|
||||
with self.get_sftp().open(remote_filename, 'w') as remote:
|
||||
remote.write(rewriter(local.read()))
|
||||
return remote_filename
|
||||
|
||||
def spawn_command(self, cmd, get_pty=False, **kws):
|
||||
chan = self.get_transport().open_session()
|
||||
chan.set_combine_stderr(True)
|
||||
if get_pty:
|
||||
chan.get_pty()
|
||||
cmd = " ".join([shlex.quote(arg.format(tmp=self.tmp, **kws)) for arg in cmd])
|
||||
logger.debug("Executing {}".format(cmd))
|
||||
chan.exec_command(cmd.format(tmp=self.tmp, **kws))
|
||||
chan.exec_command(cmd)
|
||||
return chan
|
||||
|
||||
def drain(self, chan):
|
||||
|
||||
Loading…
Reference in New Issue
Block a user