From 9896d78e07b93e04214d75029ae94202ecdda1fb Mon Sep 17 00:00:00 2001 From: Sebastien Bourdeauducq Date: Wed, 18 May 2022 19:04:13 +0800 Subject: [PATCH] afws_client: update --- artiq/frontend/afws_client.py | 62 +++++++++++++++++++++++++++++++---- flake.nix | 2 +- 2 files changed, 57 insertions(+), 7 deletions(-) diff --git a/artiq/frontend/afws_client.py b/artiq/frontend/afws_client.py index 7d9a78265..25ddfc1f9 100755 --- a/artiq/frontend/afws_client.py +++ b/artiq/frontend/afws_client.py @@ -7,7 +7,10 @@ import socket import ssl import io import zipfile +import json +from prettytable import PrettyTable from getpass import getpass +from tqdm import tqdm def get_artiq_cert(): @@ -54,20 +57,42 @@ class Client: self.fsocket.write((" ".join(command) + "\n").encode()) self.fsocket.flush() + def read_line(self): + return self.fsocket.readline().decode("ascii") + def read_reply(self): return self.fsocket.readline().decode("ascii").split() + def read_json(self): + return json.loads(self.fsocket.readline().decode("ascii")) + def login(self, username, password): self.send_command("LOGIN", username, password) return self.read_reply() == ["HELLO"] - def build(self, rev, variant): - self.send_command("BUILD", rev, variant) + def build(self, rev, variant, log): + if not variant: + variants = self.get_variants() + if len(variants) != 1: + raise ValueError("User can build more than 1 variant - need to specify") + variant = variants[0][0] + print("Building variant: {}".format(variant)) + if log: + self.send_command("BUILD", rev, variant, "LOG_ENABLE") + else: + self.send_command("BUILD", rev, variant) reply = self.read_reply()[0] if reply != "BUILDING": return reply, None print("Build in progress. This may take 10-15 minutes.") - reply, status = self.read_reply() + if log: + line = self.read_line() + while line != "" and line.startswith("LOG"): + print(line[4:], end="") + line = self.read_line() + reply, status = line.split() + else: + reply, status = self.read_reply() if reply != "DONE": raise ValueError("Unexpected server reply: expected 'DONE', got '{}'".format(reply)) if status != "done": @@ -76,13 +101,28 @@ class Client: reply, length = self.read_reply() if reply != "PRODUCT": raise ValueError("Unexpected server reply: expected 'PRODUCT', got '{}'".format(reply)) - contents = self.fsocket.read(int(length)) + length = int(length) + contents = bytearray() + with tqdm(total=length, unit="iB", unit_scale=True, unit_divisor=1024) as progress_bar: + total = 0 + while total != length: + chunk_len = min(4096, length-total) + contents += self.fsocket.read(chunk_len) + total += chunk_len + progress_bar.update(chunk_len) print("Download completed.") return "OK", contents def passwd(self, password): self.send_command("PASSWD", password) return self.read_reply() == ["OK"] + + def get_variants(self): + self.send_command("GET_VARIANTS") + reply = self.read_reply()[0] + if reply != "OK": + raise ValueError("Unexpected server reply: expected 'OK', got '{}'".format(reply)) + return self.read_json() def main(): @@ -95,9 +135,11 @@ def main(): action.required = True act_build = action.add_parser("build", help="build and download firmware") act_build.add_argument("--rev", default=None, help="revision to build (default: currently installed ARTIQ revision)") - act_build.add_argument("variant", help="variant to build") + act_build.add_argument("--log", action="store_true", help="Display the build log") act_build.add_argument("directory", help="output directory") + act_build.add_argument("variant", nargs="?", default=None, help="variant to build (can be omitted if user is authorised to build only one)") act_passwd = action.add_parser("passwd", help="change password") + act_get_variants = action.add_parser("get_variants", help="get available variants and expiry dates") args = parser.parse_args() cert = args.cert @@ -146,14 +188,22 @@ def main(): if rev is None: print("Unable to determine currently installed ARTIQ revision. Specify manually using --rev.") sys.exit(1) - result, contents = client.build(rev, args.variant) + result, contents = client.build(rev, args.variant, args.log) if result != "OK": if result == "UNAUTHORIZED": print("You are not authorized to build this variant. Your firmware subscription may have expired. Contact helpdesk\x40m-labs.hk.") + elif result == "TOOMANY": + print("Too many builds in a queue. Please wait for others to finish.") else: print("Build failed: {}".format(result)) sys.exit(1) zip_unarchive(contents, args.directory) + elif args.action == "get_variants": + data = client.get_variants() + table = PrettyTable() + table.field_names = ["Variant", "Expiry date"] + table.add_rows(data) + print(table) else: raise ValueError finally: diff --git a/flake.nix b/flake.nix index ed0cad534..c025214a9 100644 --- a/flake.nix +++ b/flake.nix @@ -147,7 +147,7 @@ nativeBuildInputs = [ pkgs.qt5.wrapQtAppsHook ]; # keep llvm_x and lld_x in sync with llvmlite propagatedBuildInputs = [ pkgs.llvm_11 pkgs.lld_11 llvmlite-new sipyco.packages.x86_64-linux.sipyco pythonparser artiq-comtools.packages.x86_64-linux.artiq-comtools ] - ++ (with pkgs.python3Packages; [ pyqtgraph pygit2 numpy dateutil scipy prettytable pyserial python-Levenshtein h5py pyqt5 qasync ]); + ++ (with pkgs.python3Packages; [ pyqtgraph pygit2 numpy dateutil scipy prettytable pyserial python-Levenshtein h5py pyqt5 qasync tqdm ]); dontWrapQtApps = true; postFixup = ''