feat(download): async, status; download.py: use download_and_extract

Author: gadicc

api/download.py

@@ -70,27 +70,13 @@ async def download_model(
         if not filename:
             filename = normalized_model_id + ".tar.zst"
         model_file = os.path.join(MODELS_DIR, filename)
-        storage = Storage(model_url, default_path=normalized_model_id + ".tar.zst")
+        storage = Storage(
+            model_url, default_path=normalized_model_id + ".tar.zst", status=status
+        )
         exists = storage.file_exists()
         if exists:
-            storage.download_file(model_file)
-            # os.mkdir(id)
-            # Path(id).mkdir(parents=True, exist_ok=False)
             model_dir = os.path.join(MODELS_DIR, normalized_model_id)
-            os.mkdir(model_dir)
-            subprocess.run(
-                [
-                    "tar",
-                    "--use-compress-program=unzstd",
-                    "-C",
-                    model_dir,
-                    "-xvf",
-                    model_file,
-                ],
-                check=True,
-            )
-            subprocess.run(["ls", "-l"])
-            os.remove(model_file)
+            await asyncio.to_thread(storage.download_and_extract, model_file, model_dir)
         else:
             if checkpoint_url:
                 download_checkpoint(checkpoint_url)

api/utils/storage/BaseStorage.py

@@ -5,8 +5,13 @@
 
 
 class BaseArchive(ABC):
-    def __init__(self, path):
+    def __init__(self, path, status=None):
         self.path = path
+        self.status = status
+
+    def updateStatus(self, type, progress):
+        if hasattr(self, "status"):
+            self.status.update(type, progress)
 
     def extract(self):
         print("TODO")
@@ -23,6 +28,7 @@ def test(path):
         return re.search(r"\.tar\.zstd?$", path)
 
     def extract(self, dir, dry_run=False):
+        self.updateStatus("extract", 0)
         if not dir:
             dir = os.path.dirname(self.path)
         base, ext, subext = self.splitext()
@@ -41,8 +47,10 @@ def extract(self, dir, dry_run=False):
                 ],
                 check=True,
             )
+            subprocess.run(["ls", "-l"])
             os.remove(self.path)
 
+        self.updateStatus("extract", 1)
         return dir  # , base, ext, subext
 
 
@@ -63,6 +71,11 @@ def test(url):
 
     def __init__(self, url, **kwargs):
         self.url = url
+        self.status = kwargs.get("status", None)
+
+    def updateStatus(self, type, progress):
+        if hasattr(self, "status"):
+            self.status.update(type, progress)
 
     def splitext(self):
         base, ext = os.path.splitext(self.url)
@@ -77,7 +90,7 @@ def download_file(self, dest):
         """Download the file to `dest`"""
         pass
 
-    def download_and_extract(self, fname, dry_run=False):
+    def download_and_extract(self, fname, dir=None, dry_run=False):
         """
         Downloads the file, and if it's an archive, extract it too.  Returns
         the filename if not, or directory name (fname without extension) if
@@ -86,12 +99,11 @@ def download_and_extract(self, fname, dry_run=False):
         if not fname:
             fname = self.get_filename()
 
-        dir = None
-        archive = Archive(fname, dry_run=dry_run)
+        archive = Archive(fname, status=self.status)
         if archive:
             # TODO, streaming pipeline
             self.download_file(fname)
-            return archive.extract()
+            return archive.extract(dir)
         else:
             self.download_file(fname)
             return fname

api/utils/storage/HTTPStorage.py

@@ -16,7 +16,7 @@ def test(url):
         return re.search(r"^https?://", url)
 
     def __init__(self, url, **kwargs):
-        self.url = url
+        super().__init__(url, **kwargs)
 
     def upload_file(self, source, dest):
         raise RuntimeError("HTTP PUT not implemented yet")
@@ -37,6 +37,9 @@ def download_file(self, fname):
             unit_scale=True,
             unit_divisor=1024,
         ) as bar:
+            total_written = 0
             for data in resp.iter_content(chunk_size=1024):
                 size = file.write(data)
                 bar.update(size)
+                total_written += size
+                self.updateStatus("download", total_written / total)

api/utils/storage/S3Storage.py

@@ -24,7 +24,7 @@ def test(url):
         return re.search(r"^(https?\+)?s3://", url)
 
     def __init__(self, url, **kwargs):
-        self.url = url
+        super().__init__(url, **kwargs)
 
         if url.startswith("s3://"):
             url = "https://" + url[5:]
@@ -90,10 +90,16 @@ def upload_file(self, source, dest):
         upload_start = get_now()
         file_size = os.stat(source).st_size
         with tqdm(total=file_size, unit="B", unit_scale=True, desc="Uploading") as bar:
+            total_transferred = 0
+
+            def callback(bytes_transferred):
+                nonlocal total_transferred
+                bar.update(bytes_transferred),
+                total_transferred += bytes_transferred
+                self.updateStatus("upload", total_transferred / file_size)
+
             result = self.bucket().upload_file(
-                Filename=source,
-                Key=dest,
-                Callback=lambda bytes_transferred: bar.update(bytes_transferred),
+                Filename=source, Key=dest, Callback=callback
             )
         print(result)
         upload_total = get_now() - upload_start
@@ -110,10 +116,15 @@ def download_file(self, dest):
         with tqdm(
             total=object.content_length, unit="B", unit_scale=True, desc="Downloading"
         ) as bar:
-            object.download_file(
-                Filename=dest,
-                Callback=lambda bytes_transffered: bar.update(bytes_transffered),
-            )
+            total_transferred = 0
+
+            def callback(bytes_transferred):
+                nonlocal total_transferred
+                bar.update(bytes_transferred),
+                total_transferred += bytes_transferred
+                self.updateStatus("download", total_transferred / object.content_length)
+
+            object.download_file(Filename=dest, Callback=callback)
 
     def file_exists(self):
         # res = self.s3client().list_objects_v2(