forked from M-Labs/artiq
1
0
Fork 0

afws_client: update

This commit is contained in:
Sebastien Bourdeauducq 2022-05-18 19:04:13 +08:00
parent 70503bee6f
commit 9896d78e07
2 changed files with 57 additions and 7 deletions

View File

@ -7,7 +7,10 @@ import socket
import ssl import ssl
import io import io
import zipfile import zipfile
import json
from prettytable import PrettyTable
from getpass import getpass from getpass import getpass
from tqdm import tqdm
def get_artiq_cert(): def get_artiq_cert():
@ -54,19 +57,41 @@ class Client:
self.fsocket.write((" ".join(command) + "\n").encode()) self.fsocket.write((" ".join(command) + "\n").encode())
self.fsocket.flush() self.fsocket.flush()
def read_line(self):
return self.fsocket.readline().decode("ascii")
def read_reply(self): def read_reply(self):
return self.fsocket.readline().decode("ascii").split() return self.fsocket.readline().decode("ascii").split()
def read_json(self):
return json.loads(self.fsocket.readline().decode("ascii"))
def login(self, username, password): def login(self, username, password):
self.send_command("LOGIN", username, password) self.send_command("LOGIN", username, password)
return self.read_reply() == ["HELLO"] return self.read_reply() == ["HELLO"]
def build(self, 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) self.send_command("BUILD", rev, variant)
reply = self.read_reply()[0] reply = self.read_reply()[0]
if reply != "BUILDING": if reply != "BUILDING":
return reply, None return reply, None
print("Build in progress. This may take 10-15 minutes.") print("Build in progress. This may take 10-15 minutes.")
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() reply, status = self.read_reply()
if reply != "DONE": if reply != "DONE":
raise ValueError("Unexpected server reply: expected 'DONE', got '{}'".format(reply)) raise ValueError("Unexpected server reply: expected 'DONE', got '{}'".format(reply))
@ -76,7 +101,15 @@ class Client:
reply, length = self.read_reply() reply, length = self.read_reply()
if reply != "PRODUCT": if reply != "PRODUCT":
raise ValueError("Unexpected server reply: expected 'PRODUCT', got '{}'".format(reply)) 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.") print("Download completed.")
return "OK", contents return "OK", contents
@ -84,6 +117,13 @@ class Client:
self.send_command("PASSWD", password) self.send_command("PASSWD", password)
return self.read_reply() == ["OK"] 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(): def main():
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
@ -95,9 +135,11 @@ def main():
action.required = True action.required = True
act_build = action.add_parser("build", help="build and download firmware") 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("--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("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_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() args = parser.parse_args()
cert = args.cert cert = args.cert
@ -146,14 +188,22 @@ def main():
if rev is None: if rev is None:
print("Unable to determine currently installed ARTIQ revision. Specify manually using --rev.") print("Unable to determine currently installed ARTIQ revision. Specify manually using --rev.")
sys.exit(1) sys.exit(1)
result, contents = client.build(rev, args.variant) result, contents = client.build(rev, args.variant, args.log)
if result != "OK": if result != "OK":
if result == "UNAUTHORIZED": if result == "UNAUTHORIZED":
print("You are not authorized to build this variant. Your firmware subscription may have expired. Contact helpdesk\x40m-labs.hk.") 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: else:
print("Build failed: {}".format(result)) print("Build failed: {}".format(result))
sys.exit(1) sys.exit(1)
zip_unarchive(contents, args.directory) 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: else:
raise ValueError raise ValueError
finally: finally:

View File

@ -147,7 +147,7 @@
nativeBuildInputs = [ pkgs.qt5.wrapQtAppsHook ]; nativeBuildInputs = [ pkgs.qt5.wrapQtAppsHook ];
# keep llvm_x and lld_x in sync with llvmlite # 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 ] 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; dontWrapQtApps = true;
postFixup = '' postFixup = ''