afws_client: update

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,7 +101,15 @@ 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
 
@@ -84,6 +117,13 @@ class Client:
         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():
     parser = argparse.ArgumentParser()
@@ -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:

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 = ''