Merge pull request #13 from manifoldai/ms/logging_decorator

Ms/logging decorator

{{ cookiecutter.repo_name }}/{{ cookiecutter.repo_name }}/scripts/evaluate.py

@@ -13,10 +13,9 @@
 import click
 import mlflow
 import numpy as np
-from dotenv import find_dotenv, load_dotenv
 
 from {{cookiecutter.repo_name}}.util.config import parse_config
-from {{cookiecutter.repo_name}}.util.logging import setup_logging
+from {{cookiecutter.repo_name}}.util.logging import setup_logging_env
 
 logger = logging.getLogger(__name__)
 
@@ -25,6 +24,7 @@
 @click.argument(
     "config_file", type=click.Path(exists=True), default="/mnt/configs/config.yml"
 )
+@setup_logging_env
 def main(config_file="config.yml"):
     """
     Main function that loads config, sets up logging, and runs evaluation
@@ -106,6 +106,4 @@ def log_experiment(
 
 
 if __name__ == "__main__":
-    setup_logging()
-    load_dotenv(find_dotenv(), verbose=True, override=True)
     main()

{{ cookiecutter.repo_name }}/{{ cookiecutter.repo_name }}/scripts/predict.py

@@ -11,12 +11,13 @@
 
 import click
 
-from {{cookiecutter.repo_name}}.util.logging import setup_logging
+from {{cookiecutter.repo_name}}.util.logging import setup_logging_env
 
 logger = logging.getLogger(__name__)
 
 
 @click.command()
+@setup_logging_env
 def main():
     """
     Main function that loads config, sets up logging, and runs predictions
@@ -31,6 +32,4 @@ def main():
 
 
 if __name__ == "__main__":
-    # set-up logging only once
-    setup_logging()
     main()

{{ cookiecutter.repo_name }}/{{ cookiecutter.repo_name }}/scripts/train.py

@@ -12,7 +12,7 @@
 import click
 
 from {{cookiecutter.repo_name}}.util.config import parse_config
-from {{cookiecutter.repo_name}}.util.logging import setup_logging
+from {{cookiecutter.repo_name}}.util.logging import setup_logging_env
 
 logger = logging.getLogger(__name__)
 
@@ -21,6 +21,7 @@
 @click.argument(
     "config_file", type=click.Path(exists=True), default="/mnt/configs/config.yml"
 )
+@setup_logging_env
 def main(config_file="config.yml"):
     """
     Main function that loads config, sets up logging, and runs training
@@ -38,5 +39,4 @@ def main(config_file="config.yml"):
 
 
 if __name__ == "__main__":
-    setup_logging()
     main()

{{ cookiecutter.repo_name }}/{{ cookiecutter.repo_name }}/util/logging.py

@@ -4,9 +4,15 @@
 """
 import logging.config
 import os
+import sys
+from datetime import timedelta
+from typing import Callable
 
 import coloredlogs
+import numpy as np
 import yaml
+from dotenv import find_dotenv, load_dotenv
+from pytictoc import TicToc
 
 
 logger = logging.getLogger(__name__)
@@ -56,3 +62,29 @@ def setup_logging(logging_config="logging.yml", default_level=logging.INFO):
         config_method = "default_level"
         coloredlogs.install(level="DEBUG")
     logger.info(f"Logging set from {config_method}")
+
+
+def setup_logging_env(main: Callable) -> Callable:
+    """Decorator to set up loggging and load env variables
+
+    Args:
+        main: top level function (typically main triggered by CLI)
+
+    Return:
+        function after setting up logging and loading env variables
+    """
+
+    def wrapper(*args, **kwargs):
+        setup_logging()
+        load_dotenv(find_dotenv())
+        logger.info(f"Loaded environment variables")
+        logger.info(f"Starting {main.__name__}() in {sys.argv[0]}")
+        t = TicToc()
+        t.tic()
+        main(*args, **kwargs)
+        logger.info(
+            f"Finished {main.__name__}() in "
+            f"{timedelta(seconds=np.ceil(t.tocvalue()))}"
+        )
+
+    return wrapper