workspace.py

import logging
import os
import time

import git
import omegaconf

logger = logging.getLogger(__name__)


def load_config(config_path, command_line_args=None):
    """Loads the config file using OmegaConf, performing merges with base configs and the command line arguments"""

    logger.info(f"Loading from: {config_path}")

    # Load the config using OmegaConf
    config = omegaconf.OmegaConf.load(config_path)

    # Load all the base configs to include, and merge them with the current config, giving precedence to the current config
    if hasattr(config, "include"):
        base_config_paths = [os.path.join(os.path.dirname(config_path), include_path) for include_path in config.include]
        base_configs = [load_config(base_config_path) for base_config_path in base_config_paths]
        config = omegaconf.OmegaConf.merge(*base_configs, config)

    # Load the command line arguments, and merge them with the current config, giving precedence to the command line
    if command_line_args is not None:
        command_line_config = omegaconf.OmegaConf.from_dotlist(command_line_args)
        config = omegaconf.OmegaConf.merge(config, command_line_config)

    return config


def save_git_commit_info(save_path):
    """Use gitpython to save info about the current git commit to a file"""

    repo = git.Repo(search_parent_directories=True)
    head_commit = repo.head.commit
    git_commit_info = {
        "hexsha": head_commit.hexsha,
        "authored": {
            "author": head_commit.author.name,
            "authored_time": head_commit.authored_date,
        },
        "committed": {
            "commit": head_commit.committer.name,
            "committed_time": head_commit.committed_date,
        },
        "message": head_commit.message.strip(),
    }

    git_commit_info = omegaconf.OmegaConf.create(git_commit_info)
    omegaconf.OmegaConf.save(git_commit_info, save_path)
    return git_commit_info


def create_workspace(config):
    """Create a results folder in the target directory"""

    # Treat the name as a time.strftime format string (so that every experiment is named after when it was run)
    config.name = time.strftime(config.name, time.localtime())

    # Create the results directory
    os.makedirs(config.save_dir)

    # Save the config to the results directory
    omegaconf.OmegaConf.save(config, os.path.join(config.save_dir, "config.yaml"))
    save_git_commit_info(os.path.join(config.save_dir, "git.yaml"))

    # Set up the print loggers by removing all handlers associated with the root logger object,
    # then setting up the logger to print messages *and* save them to a file
    for handler in logging.root.handlers:
        logging.root.removeHandler(handler)
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logging.INFO,
        handlers=[
            logging.FileHandler(os.path.join(config.save_dir, "output.log")),
            logging.StreamHandler(),
        ],
    )

    return config