Merge pull request #26 from grok-ai/feature/bump-dependencies

Update to Lightning 1.7

setup.cfg

@@ -17,7 +17,7 @@ packages=find:
 install_requires =
     # Add project specific dependencies
     # Stuff easy to break with updates
-    pytorch-lightning>=1.5.8,<1.6
+    pytorch-lightning==1.7.*
     hydra-core
     wandb
 

src/nn_core/model_logging.py

@@ -39,7 +39,11 @@ def __init__(self, logging_cfg: DictConfig, cfg: DictConfig, resume_id: Optional
             self.logging_cfg.logger.mode = "offline"
 
         pylogger.info(f"Instantiating <{self.logging_cfg.logger['_target_'].split('.')[-1]}>")
-        self.wrapped: LightningLoggerBase = hydra.utils.instantiate(self.logging_cfg.logger, version=self.resume_id)
+        self.wrapped: LightningLoggerBase = hydra.utils.instantiate(
+            self.logging_cfg.logger,
+            version=self.resume_id,
+            dir=os.getenv("WANDB_DIR", "."),
+        )
 
         # force experiment lazy initialization
         _ = self.wrapped.experiment

src/nn_core/serialization.py

@@ -12,6 +12,7 @@
 
 import pytorch_lightning as pl
 import torch
+from pytorch_lightning.core.saving import _load_state
 from pytorch_lightning.plugins import TorchCheckpointIO
 
 METADATA_KEY: str = "metadata"
@@ -94,7 +95,7 @@ def remove_checkpoint(self, path) -> None:
 
 
 def compress_checkpoint(src_dir: Path, dst_file: Path, delete_dir: bool = True):
-
+    dst_file.parent.mkdir(parents=True, exist_ok=True)
     with zipfile.ZipFile(_normalize_path(dst_file), "w") as zip_file:
         for folder, subfolders, files in os.walk(src_dir):
             folder: Path = Path(folder)
@@ -161,7 +162,7 @@ def load_model(
         if substitute_values is not None:
             checkpoint = _substistute(checkpoint, substitute_values=substitute_values, substitute_keys=substitute_keys)
 
-        return module_class._load_model_state(checkpoint=checkpoint, metadata=checkpoint.get("metadata", None))
+        return _load_state(cls=module_class, checkpoint=checkpoint, metadata=checkpoint.get("metadata", None))
     else:
         pylogger.warning(f"Loading a legacy checkpoint (from vanilla PyTorch Lightning): '{checkpoint_path}'")
         module_class.load_from_checkpoint(checkpoint_path=str(checkpoint_path), map_location=map_location)