# Copyright 2021 MosaicML. All Rights Reserved.

import datetime
import logging
import os
import pathlib
import time

from composer.utils import dist

log = logging.getLogger(__name__)

_RUN_DIRECTORY_KEY = "COMPOSER_RUN_DIRECTORY"

_start_time_str = datetime.datetime.now().isoformat()


def get_node_run_directory() -> str:
    """Returns the run directory for the node. This folder is shared by all ranks on the node.
    
    Returns:
        str: The node run directory.
    """
    node_run_directory = os.environ.get(_RUN_DIRECTORY_KEY, os.path.join("runs", _start_time_str))
    if node_run_directory.endswith(os.path.sep):
        # chop off the training slash so os.path.basename would work as expected
        node_run_directory = node_run_directory[:-1]
    os.makedirs(node_run_directory, exist_ok=True)
    return os.path.abspath(node_run_directory)


def get_run_directory() -> str:
    """Returns the run directory for the current rank.
    
    Returns:
        str: The run directory.
    """
    run_dir = os.path.join(get_node_run_directory(), f"rank_{dist.get_global_rank()}")
    os.makedirs(run_dir, exist_ok=True)
    return run_dir


def get_modified_files(modified_since_timestamp: float, *, ignore_hidden: bool = True):
    """Returns a list of files (recursively) in the run directory that have been modified since
    ``modified_since_timestamp``.

    Args:
        modified_since_timestamp (float): Minimum last modified timestamp(in seconds since EPOCH)
            of files to include.
        ignore_hidden (bool, optional): Whether to ignore hidden files and folders (default: ``True``)
    Returns:
        List[str]: List of filepaths that have been modified since ``modified_since_timestamp``
    """
    modified_files = []
    run_directory = get_run_directory()
    if run_directory is None:
        raise RuntimeError("Run directory is not defined")
    for root, dirs, files in os.walk(run_directory):
        del dirs  # unused
        for file in files:
            if ignore_hidden and any(x.startswith(".") for x in file.split(os.path.sep)):
                # skip hidden files and folders
                continue
            filepath = os.path.join(root, file)
            modified_time = os.path.getmtime(filepath)
            if modified_time >= modified_since_timestamp:
                modified_files.append(filepath)
    return modified_files


def get_run_directory_timestamp() -> float:
    """Returns the current timestamp on the run directory filesystem.
    Note that the disk time can differ from system time (e.g. when using
    network filesystems).

    Returns:
        float: the current timestamp on the run directory filesystem.
    """
    run_directory = get_run_directory()
    if run_directory is None:
        raise RuntimeError("Run directory is not defined")
    python_time = time.time()
    touch_file = (pathlib.Path(run_directory) / f".{python_time}")
    touch_file.touch()
    new_last_uploaded_timestamp = os.path.getmtime(str(touch_file))
    os.remove(str(touch_file))
    return new_last_uploaded_timestamp

[((11, 6, 11, 33), 'logging.getLogger', 'logging.getLogger', ({(11, 24, 11, 32): '__name__'}, {}), '(__name__)', False, 'import logging\n'), ((28, 4, 28, 50), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((29, 11, 29, 46), 'os.path.abspath', 'os.path.abspath', ({(29, 27, 29, 45): 'node_run_directory'}, {}), '(node_run_directory)', False, 'import os\n'), ((39, 4, 39, 39), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((58, 29, 58, 51), 'os.walk', 'os.walk', ({(58, 37, 58, 50): 'run_directory'}, {}), '(run_directory)', False, 'import os\n'), ((82, 18, 82, 29), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((15, 18, 15, 41), 'datetime.datetime.now', 'datetime.datetime.now', ({}, {}), '()', False, 'import datetime\n'), ((24, 60, 24, 97), 'os.path.join', 'os.path.join', ({(24, 73, 24, 79): '"""runs"""', (24, 81, 24, 96): '_start_time_str'}, {}), "('runs', _start_time_str)", False, 'import os\n'), ((83, 18, 83, 45), 'pathlib.Path', 'pathlib.Path', ({(83, 31, 83, 44): 'run_directory'}, {}), '(run_directory)', False, 'import pathlib\n'), ((64, 23, 64, 47), 'os.path.join', 'os.path.join', ({(64, 36, 64, 40): 'root', (64, 42, 64, 46): 'file'}, {}), '(root, file)', False, 'import os\n'), ((65, 28, 65, 54), 'os.path.getmtime', 'os.path.getmtime', ({(65, 45, 65, 53): 'filepath'}, {}), '(filepath)', False, 'import os\n'), ((38, 61, 38, 83), 'composer.utils.dist.get_global_rank', 'dist.get_global_rank', ({}, {}), '()', False, 'from composer.utils import dist\n')]