Fix saving large pipelines (#1335)

tinkoff-ai · Jul 31, 2023 · cceb500 · cceb500 · github-actions · Jul 31, 2023
1 parent e4c121c
commit cceb500
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 3 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -27,6 +27,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `mrmr` feature selection working with categoricals ([#1311](https://github.com/tinkoff-ai/etna/pull/1311))
 - Fix version of `statsforecast` to 1.4 to avoid dependency conflicts during installation ([#1313](https://github.com/tinkoff-ai/etna/pull/1313))
 - Add inverse transformation into `predict` method of pipelines ([#1314](https://github.com/tinkoff-ai/etna/pull/1314))
+- Allow saving large pipelines ([#1335](https://github.com/tinkoff-ai/etna/pull/1335))
 
 ### Removed
 - Building docker images with cuda 10.2 ([#1306](https://github.com/tinkoff-ai/etna/pull/1306))

diff --git a/etna/core/mixins.py b/etna/core/mixins.py
@@ -230,7 +230,7 @@ def _save_metadata(self, archive: zipfile.ZipFile):
             output_file.write(metadata_bytes)
 
     def _save_state(self, archive: zipfile.ZipFile):
-        with archive.open("object.pkl", "w") as output_file:
+        with archive.open("object.pkl", "w", force_zip64=True) as output_file:
             dill.dump(self, output_file)
 
     def save(self, path: pathlib.Path):

diff --git a/etna/models/mixins.py b/etna/models/mixins.py
@@ -630,7 +630,7 @@ def get_model(self) -> Any:
 
 
 class SaveNNMixin(SaveMixin):
-    """Implementation of ``AbstractSaveable``  torch related classes.
+    """Implementation of ``AbstractSaveable`` torch related classes.
 
     It saves object to the zip archive with 2 files:
 
@@ -642,7 +642,7 @@ class SaveNNMixin(SaveMixin):
     def _save_state(self, archive: zipfile.ZipFile):
         import torch
 
-        with archive.open("object.pt", "w") as output_file:
+        with archive.open("object.pt", "w", force_zip64=True) as output_file:
             torch.save(self, output_file, pickle_module=dill)
 
     @classmethod