feat(storage): support misc tar compression; progress

Author: gadicc

api/utils/storage/BaseStorage.py

@@ -2,6 +2,7 @@
 import re
 import subprocess
 from abc import ABC, abstractmethod
+import xtarfile as tarfile
 
 
 class BaseArchive(ABC):
@@ -22,10 +23,10 @@ def splitext(self):
         return base, ext, subext
 
 
-class TarZstdArchive(BaseArchive):
+class TarArchive(BaseArchive):
     @staticmethod
     def test(path):
-        return re.search(r"\.tar\.zstd?$", path)
+        return re.search(r"\.tar", path)
 
     def extract(self, dir, dry_run=False):
         self.updateStatus("extract", 0)
@@ -36,25 +37,26 @@ def extract(self, dir, dry_run=False):
 
         if not dry_run:
             os.mkdir(dir)
-            subprocess.run(
-                [
-                    "tar",
-                    "--use-compress-program=unzstd",
-                    "-C",
-                    dir,
-                    "-xvf",
-                    self.path,
-                ],
-                check=True,
-            )
+
+            def track_progress(tar):
+                i = 0
+                members = tar.getmembers()
+                for member in members:
+                    i += 1
+                    self.updateStatus("extract", i / len(members))
+                    yield member
+
+            with tarfile.open(self.path, "r") as tar:
+                tar.extractall(path=dir, members=track_progress(tar))
+                tar.close()
             subprocess.run(["ls", "-l"])
             os.remove(self.path)
 
         self.updateStatus("extract", 1)
         return dir  # , base, ext, subext
 
 
-archiveClasses = [TarZstdArchive]
+archiveClasses = [TarArchive]
 
 
 def Archive(path, **kwargs):

requirements.txt

@@ -55,3 +55,5 @@ datasets==2.8.0
 omegaconf==2.3.0
 pytorch_lightning==1.9.2
 tensorboard==2.12.0
+
+xtarfile[zstd]==0.1.0