Merge pull request #13 from Flegyas/feature/stefan

Integrate feedback from user(s)

Author: lucmos

src/nn_core/callbacks.py

@@ -1,15 +1,31 @@
 import logging
-from typing import Any, Dict
+from pathlib import Path
+from typing import Any, Dict, Optional
 
 import pytorch_lightning as pl
+from omegaconf import DictConfig
 from pytorch_lightning import Callback, Trainer
 
 from nn_core.model_logging import NNLogger
+from nn_core.resume import parse_restore
+from nn_core.serialization import METADATA_KEY, NNCheckpointIO
 
 pylogger = logging.getLogger(__name__)
 
 
 class NNTemplateCore(Callback):
+    def __init__(self, restore_cfg: Optional[DictConfig]):
+        self.resume_ckpt_path, self.resume_run_version = parse_restore(restore_cfg)
+        self.restore_mode: Optional[str] = restore_cfg.get("mode", None) if restore_cfg is not None else None
+
+    @property
+    def resume_id(self) -> Optional[str]:
+        return self.resume_run_version
+
+    @property
+    def trainer_ckpt_path(self) -> Optional[str]:
+        return self.resume_ckpt_path if self.restore_mode != "finetune" else None
+
     @staticmethod
     def _is_nnlogger(trainer: Trainer) -> bool:
         return isinstance(trainer.logger, NNLogger)
@@ -21,6 +37,12 @@ def on_train_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) ->
             trainer.logger.log_configuration(model=pl_module)
             trainer.logger.watch_model(pl_module=pl_module)
 
+    def setup(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", stage: Optional[str] = None) -> None:
+        if self.restore_mode == "finetune":
+            checkpoint = NNCheckpointIO.load(path=Path(self.resume_ckpt_path))
+
+            pl_module.load_state_dict(checkpoint["state_dict"])
+
     def on_train_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
         if self._is_nnlogger(trainer):
             trainer.logger: NNLogger
@@ -31,3 +53,4 @@ def on_save_checkpoint(
     ) -> None:
         if self._is_nnlogger(trainer):
             trainer.logger.on_save_checkpoint(trainer=trainer, pl_module=pl_module, checkpoint=checkpoint)
+            checkpoint[METADATA_KEY] = trainer.datamodule.metadata

src/nn_core/common/utils.py

@@ -1,5 +1,6 @@
 import logging
 import os
+from contextlib import contextmanager
 from typing import List, Optional
 
 import dotenv
@@ -53,6 +54,31 @@ def load_envs(env_file: Optional[str] = None) -> None:
     dotenv.load_dotenv(dotenv_path=env_file, override=True)
 
 
+@contextmanager
+def environ(**kwargs):
+    """Temporarily set the process environment variables.
+
+    https://stackoverflow.com/a/34333710
+
+    >>> with environ(PLUGINS_DIR=u'test/plugins'):
+    ...   "PLUGINS_DIR" in os.environ
+    True
+
+    >>> "PLUGINS_DIR" in os.environ
+    False
+
+    :type kwargs: dict[str, unicode]
+    :param kwargs: Environment variables to set
+    """
+    old_environ = dict(os.environ)
+    os.environ.update(kwargs)
+    try:
+        yield
+    finally:
+        os.environ.clear()
+        os.environ.update(old_environ)
+
+
 def enforce_tags(tags: Optional[List[str]]) -> List[str]:
     if tags is None:
         if "id" in HydraConfig().cfg.hydra.job:
@@ -69,10 +95,10 @@ def enforce_tags(tags: Optional[List[str]]) -> List[str]:
     return tags
 
 
-def seed_index_everything(train_cfg: DictConfig) -> None:
+def seed_index_everything(train_cfg: DictConfig, sampling_seed: int = 42) -> None:
     if "seed_index" in train_cfg and train_cfg.seed_index is not None:
         seed_index = train_cfg.seed_index
-        seed_everything(42)
+        np.random.seed(sampling_seed)
         seeds = np.random.randint(np.iinfo(np.int32).max, size=max(42, seed_index + 1))
         seed = seeds[seed_index]
         seed_everything(seed)

src/nn_core/hooks.py

@@ -1,30 +1,42 @@
 import logging
 import re
+import tempfile
 from operator import xor
 from pathlib import Path
 from typing import Optional, Tuple
 
-import torch
 import wandb
 from omegaconf import DictConfig
 from wandb.apis.public import Run
 
+from nn_core.common import PROJECT_ROOT
+from nn_core.common.utils import environ
+from nn_core.serialization import NNCheckpointIO
+
 pylogger = logging.getLogger(__name__)
 
 RUN_PATH_PATTERN = re.compile(r"^([^/]+)/([^/]+)/([^/]+)$")
 
 RESUME_MODES = {
-    "continue": {
-        "restore_model": True,
-        "restore_run": True,
+    None: {
+        "logging": False,
+        "trainer": False,
+        "weights": False,
+    },
+    "finetune": {
+        "logging": False,
+        "trainer": False,
+        "weights": True,
     },
     "hotstart": {
-        "restore_model": True,
-        "restore_run": False,
+        "logging": False,
+        "trainer": True,
+        "weights": True,
     },
-    None: {
-        "restore_model": False,
-        "restore_run": False,
+    "continue": {
+        "logging": True,
+        "trainer": True,
+        "weights": True,
     },
 }
 
@@ -38,16 +50,26 @@ def resolve_ckpt(ckpt_or_run_path: str) -> str:
     Returns:
         an existing path towards the best checkpoint
     """
-    if Path(ckpt_or_run_path).exists():
-        return ckpt_or_run_path
-
-    try:
-        api = wandb.Api()
-        run: Run = api.run(path=ckpt_or_run_path)
-        ckpt_or_run_path = run.config["paths/checkpoints/best"]
+    if RUN_PATH_PATTERN.match(ckpt_or_run_path):
+        # If WANDB_DIR is set (as it is the case with our hydra configuration), the run dir is created by wandb in the
+        # project's root folder instead of in a temp directory.
+        with tempfile.TemporaryDirectory() as tmp_dir, environ(WANDB_DIR=tmp_dir):
+            # We are resolving the path from a wandb run id
+            try:
+                api = wandb.Api()
+                run: Run = api.run(path=ckpt_or_run_path)
+                ckpt_or_run_path = run.config["paths/checkpoints/best"]
+                return ckpt_or_run_path
+            except wandb.errors.CommError:
+                raise ValueError(f"Checkpoint or run not found: {ckpt_or_run_path}")
+
+    _ckpt_or_run_path: Path = Path(ckpt_or_run_path)
+    # If the path is relative, it is wrt the PROJECT_ROOT, so it is prepended.
+    if not _ckpt_or_run_path.is_absolute():
+        _ckpt_or_run_path = PROJECT_ROOT / _ckpt_or_run_path
+
+    if _ckpt_or_run_path.exists():
         return ckpt_or_run_path
-    except wandb.errors.CommError:
-        raise ValueError(f"Checkpoint or run not found: {ckpt_or_run_path}")
 
 
 def resolve_run_path(ckpt_or_run_path: str) -> str:
@@ -63,7 +85,7 @@ def resolve_run_path(ckpt_or_run_path: str) -> str:
         return ckpt_or_run_path
 
     try:
-        return torch.load(ckpt_or_run_path)["run_path"]
+        return NNCheckpointIO.load(path=Path(ckpt_or_run_path))["run_path"]
     except FileNotFoundError:
         raise ValueError(f"Checkpoint or run not found: {ckpt_or_run_path}")
 
@@ -83,7 +105,11 @@ def resolve_run_version(ckpt_or_run_path: Optional[str] = None, run_path: Option
     return RUN_PATH_PATTERN.match(run_path).group(3)
 
 
+# TODO: Refactor returning type to include restore mode too.
 def parse_restore(restore_cfg: DictConfig) -> Tuple[Optional[str], Optional[str]]:
+    if restore_cfg is None:
+        return None, None
+
     ckpt_or_run_path = restore_cfg.ckpt_or_run_path
     resume_mode = restore_cfg.mode
 
@@ -92,22 +118,21 @@ def parse_restore(restore_cfg: DictConfig) -> Tuple[Optional[str], Optional[str]
 
     if xor(bool(ckpt_or_run_path), bool(resume_mode)):
         pylogger.warning(f"Inconsistent resume modality {resume_mode} and checkpoint path '{ckpt_or_run_path}'")
+    else:
+        pylogger.info(f"Restoring with mode: <{resume_mode}>")
 
     if resume_mode not in RESUME_MODES:
         message = f"Unsupported resume mode {resume_mode}. Available resume modes are: {RESUME_MODES}"
         pylogger.error(message)
         raise ValueError(message)
 
     flags = RESUME_MODES[resume_mode]
-    restore_model = flags["restore_model"]
-    restore_run = flags["restore_run"]
 
     if ckpt_or_run_path is not None:
-        if restore_model:
-            resume_ckpt_path = resolve_ckpt(ckpt_or_run_path)
-            pylogger.info(f"Resume training from: '{resume_ckpt_path}'")
+        resume_ckpt_path = resolve_ckpt(ckpt_or_run_path)
+        pylogger.info(f"Resolved checkpoint path: '{resume_ckpt_path}'")
 
-        if restore_run:
+        if flags["logging"]:
             run_path = resolve_run_path(ckpt_or_run_path)
             resume_run_version = resolve_run_version(run_path=run_path)
             pylogger.info(f"Resume logging to: '{run_path}'")

src/nn_core/types.py renamed to src/nn_core/nn_types.py

@@ -0,0 +1,123 @@
+import importlib
+import inspect
+import logging
+import os
+import shutil
+import tempfile
+import zipfile
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, Callable, Dict, Optional, Type, Union
+
+import pytorch_lightning as pl
+from pytorch_lightning.plugins import TorchCheckpointIO
+
+METADATA_KEY: str = "metadata"
+
+pylogger = logging.getLogger(__name__)
+
+
+_METADATA_MODULE_KEY = f"{METADATA_KEY}_module"
+_METADATA_CLASS_KEY = f"{METADATA_KEY}_class"
+
+
+def _normalize_path(path: Union[Path, str]) -> Path:
+    if isinstance(path, str):
+        path = Path(path)
+    return (path.parent / path.stem.split(".")[0]).with_suffix(".ckpt.zip")
+
+
+class NNCheckpointIO(TorchCheckpointIO):
+    def __init__(self, jailing_dir: Optional[str] = None):
+        self.jailing_dir = jailing_dir
+
+    @classmethod
+    def load(cls, path: Path, map_location: Optional[Callable] = lambda storage, loc: storage):
+        return cls().load_checkpoint(path=str(path), map_location=map_location)
+
+    def save_checkpoint(self, checkpoint: Dict[str, Any], path, storage_options: Optional[Any] = None) -> None:
+        checkpoint_dir: Path = _normalize_path(path=path)
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            tmp_dir = Path(tmp_dir)
+
+            if METADATA_KEY in checkpoint:
+                metadata = checkpoint[METADATA_KEY]
+
+                metadata_path: Path = tmp_dir / METADATA_KEY
+                metadata_path.mkdir(exist_ok=True, parents=True)
+                metadata.save(dst_path=metadata_path)
+
+                checkpoint[_METADATA_MODULE_KEY] = inspect.getmodule(metadata).__name__
+                checkpoint[_METADATA_CLASS_KEY] = type(metadata).__name__
+
+                del checkpoint[METADATA_KEY]
+
+            super().save_checkpoint(
+                checkpoint=checkpoint, path=tmp_dir / "checkpoint.ckpt", storage_options=storage_options
+            )
+
+            compress_checkpoint(src_dir=tmp_dir, dst_file=checkpoint_dir)
+
+    def load_checkpoint(self, path, map_location: Optional[Callable] = lambda storage, loc: storage) -> Dict[str, Any]:
+        # load_checkpoint called from Trainer/Callbacks
+        with extract_checkpoint(ckpt_file=Path(path)) as ckpt_dir:
+            checkpoint = super().load_checkpoint(path=ckpt_dir / "checkpoint.ckpt", map_location=map_location)
+
+            if _METADATA_MODULE_KEY in checkpoint:
+                metadata_path: Path = ckpt_dir / METADATA_KEY
+                if metadata_path.exists():
+                    metadata_module = importlib.import_module(checkpoint[_METADATA_MODULE_KEY])
+                    metadata = getattr(metadata_module, checkpoint[_METADATA_CLASS_KEY]).load(src_path=metadata_path)
+                    checkpoint[METADATA_KEY] = metadata
+                    del checkpoint[_METADATA_MODULE_KEY]
+                    del checkpoint[_METADATA_CLASS_KEY]
+                else:
+                    raise FileNotFoundError(
+                        "This checkpoint is corrupted. It appears data info is required but missing."
+                    )
+
+            return checkpoint
+
+    def remove_checkpoint(self, path) -> None:
+        if self.jailing_dir is None or path.startswith(self.jailing_dir):
+            _normalize_path(path).unlink()
+            pylogger.debug(f"Removing checkpoint from {path}")
+        else:
+            pylogger.debug(
+                "Ignoring checkpoint deletion since it pertains to another run: "
+                "https://github.com/PyTorchLightning/pytorch-lightning/issues/11379"
+            )
+
+
+def compress_checkpoint(src_dir: Path, dst_file: Path, delete_dir: bool = 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)
+            for file in files:
+                zip_file.write(
+                    folder / file,
+                    os.path.relpath(os.path.join(folder, file), src_dir),
+                    compress_type=zipfile.ZIP_DEFLATED,
+                )
+
+    if delete_dir:
+        pylogger.debug(f"Deleting the checkpoint folder: '{src_dir}'")
+        shutil.rmtree(path=src_dir, ignore_errors=True)
+
+
+@contextmanager
+def extract_checkpoint(ckpt_file: Path) -> Path:
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        pylogger.debug(f"Extracting archive file '{ckpt_file}' to temp dir '{tmp_dir}'")
+        with zipfile.ZipFile(_normalize_path(ckpt_file), "r") as compressed_ckpt:
+            compressed_ckpt.extractall(tmp_dir)
+        yield Path(tmp_dir)
+
+
+def load_model(module_class: Type[pl.LightningModule], checkpoint_path: Path):
+    checkpoint = NNCheckpointIO.load(path=checkpoint_path)
+
+    model = module_class._load_model_state(checkpoint=checkpoint, metadata=checkpoint["metadata"])
+    return model