Remove compression from .nemo files (#3626)

Signed-off-by: Oleksii Kuchaiev <[email protected]> Co-authored-by: Somshubra Majumdar <[email protected]> Co-authored-by: Eric Harper <[email protected]>
NVIDIA · Mar 2, 2022 · 45355e9 · 45355e9
1 parent 7290a93
commit 45355e9
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/nemo/core/connectors/save_restore_connector.py b/nemo/core/connectors/save_restore_connector.py
@@ -384,14 +384,23 @@ def _inject_model_parallel_rank_for_ckpt(self, dirname, basename):
     def _make_nemo_file_from_folder(filename, source_dir):
         dirname = os.path.dirname(filename)
         os.makedirs(dirname, exist_ok=True)
-        with tarfile.open(filename, "w:gz") as tar:
+        with tarfile.open(filename, "w:") as tar:
             tar.add(source_dir, arcname=".")
 
     @staticmethod
     def _unpack_nemo_file(path2file: str, out_folder: str) -> str:
         if not os.path.exists(path2file):
             raise FileNotFoundError(f"{path2file} does not exist")
-        tar = tarfile.open(path2file, "r:gz")
+        # we start with an assumption of uncompressed tar,
+        # which should be true for versions 1.7.0 and above
+        tar_header = "r:"
+        try:
+            tar_test = tarfile.open(path2file, tar_header)
+            tar_test.close()
+        except tarfile.ReadError:
+            # can be older checkpoint => try compressed tar
+            tar_header = "r:gz"
+        tar = tarfile.open(path2file, tar_header)
         tar.extractall(path=out_folder)
         tar.close()
         return out_folder