Move functions from template to core (#9)

* Move parse_restore to nn-core

* Move enforce_tags to nn-core

* Move seed_index logic to nn-core

Author: lucmos

src/nn_core/common/utils.py

@@ -1,10 +1,15 @@
 import logging
 import os
-from typing import Optional
+from typing import List, Optional
 
 import dotenv
+import numpy as np
+from hydra.core.hydra_config import HydraConfig
+from omegaconf import DictConfig
+from pytorch_lightning import seed_everything
+from rich.prompt import Prompt
 
-logger = logging.getLogger(__name__)
+pylogger = logging.getLogger(__name__)
 
 
 def get_env(env_name: str, default: Optional[str] = None) -> str:
@@ -19,13 +24,17 @@ def get_env(env_name: str, default: Optional[str] = None) -> str:
     """
     if env_name not in os.environ:
         if default is None:
-            raise KeyError(f"{env_name} not defined and no default value is present!")
+            message = f"{env_name} not defined and no default value is present!"
+            pylogger.error(message)
+            raise KeyError(message)
         return default
 
     env_value: str = os.environ[env_name]
     if not env_value:
         if default is None:
-            raise ValueError(f"{env_name} has yet to be configured and no default value is present!")
+            message = f"{env_name} has yet to be configured and no default value is present!"
+            pylogger.error(message)
+            raise ValueError(message)
         return default
 
     return env_value
@@ -42,3 +51,31 @@ def load_envs(env_file: Optional[str] = None) -> None:
                      it searches for a `.env` file in the project.
     """
     dotenv.load_dotenv(dotenv_path=env_file, override=True)
+
+
+def enforce_tags(tags: Optional[List[str]]) -> List[str]:
+    if tags is None:
+        if "id" in HydraConfig().cfg.hydra.job:
+            # We are in multi-run setting (either via a sweep or a scheduler)
+            message: str = "You need to specify 'core.tags' in a multi-run setting!"
+            pylogger.error(message)
+            raise ValueError(message)
+
+        pylogger.warning("No tags provided, asking for tags...")
+        tags = Prompt.ask("Enter a list of comma separated tags", default="develop")
+        tags = [x.strip() for x in tags.split(",")]
+
+    pylogger.info(f"Tags: {tags if tags is not None else []}")
+    return tags
+
+
+def seed_index_everything(train_cfg: DictConfig) -> None:
+    if "seed_index" in train_cfg and train_cfg.seed_index is not None:
+        seed_index = train_cfg.seed_index
+        seed_everything(42)
+        seeds = np.random.randint(np.iinfo(np.int32).max, size=max(42, seed_index + 1))
+        seed = seeds[seed_index]
+        seed_everything(seed)
+        pylogger.info(f"Setting seed {seed} from seeds[{seed_index}]")
+    else:
+        pylogger.warning("The seed has not been set! The reproducibility is not guaranteed.")

src/nn_core/resume.py

@@ -1,13 +1,33 @@
+import logging
 import re
+from operator import xor
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Tuple
 
 import torch
 import wandb
+from omegaconf import DictConfig
 from wandb.apis.public import Run
 
+pylogger = logging.getLogger(__name__)
+
 RUN_PATH_PATTERN = re.compile(r"^([^/]+)/([^/]+)/([^/]+)$")
 
+RESUME_MODES = {
+    "continue": {
+        "restore_model": True,
+        "restore_run": True,
+    },
+    "hotstart": {
+        "restore_model": True,
+        "restore_run": False,
+    },
+    None: {
+        "restore_model": False,
+        "restore_run": False,
+    },
+}
+
 
 def resolve_ckpt(ckpt_or_run_path: str) -> str:
     """Resolve the run path or ckpt to a checkpoint.
@@ -61,3 +81,35 @@ def resolve_run_version(ckpt_or_run_path: Optional[str] = None, run_path: Option
     if run_path is None:
         run_path = resolve_run_path(ckpt_or_run_path)
     return RUN_PATH_PATTERN.match(run_path).group(3)
+
+
+def parse_restore(restore_cfg: DictConfig) -> Tuple[Optional[str], Optional[str]]:
+    ckpt_or_run_path = restore_cfg.ckpt_or_run_path
+    resume_mode = restore_cfg.mode
+
+    resume_ckpt_path = None
+    resume_run_version = None
+
+    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}'")
+
+    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}'")
+
+        if restore_run:
+            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}'")
+
+    return resume_ckpt_path, resume_run_version