| import os |
| import sys |
| import time |
| import __main__ |
| import traceback |
| |
|
|
| import colbert.utils.distributed as distributed |
|
|
| from contextlib import contextmanager |
| from colbert.utils.logging import Logger |
| from colbert.utils.utils import timestamp, create_directory, print_message |
|
|
|
|
| class _RunManager(): |
| def __init__(self): |
| self.experiments_root = None |
| self.experiment = None |
| self.path = None |
| self.script = self._get_script_name() |
| self.name = self._generate_default_run_name() |
| self.original_name = self.name |
| self.exit_status = 'FINISHED' |
|
|
| self._logger = None |
| self.start_time = time.time() |
|
|
| def init(self, rank, root, experiment, name): |
| assert '/' not in experiment, experiment |
| assert '/' not in name, name |
|
|
| self.experiments_root = os.path.abspath(root) |
| self.experiment = experiment |
| self.name = name |
| self.path = os.path.join(self.experiments_root, self.experiment, self.script, self.name) |
|
|
| if rank < 1: |
| if os.path.exists(self.path): |
| print('\n\n') |
| print_message("It seems that ", self.path, " already exists.") |
| print_message("Do you want to overwrite it? \t yes/no \n") |
|
|
| |
|
|
| response = input() |
| if response.strip() != 'yes': |
| assert not os.path.exists(self.path), self.path |
| else: |
| create_directory(self.path) |
|
|
| distributed.barrier(rank) |
|
|
| self._logger = Logger(rank, self) |
| self._log_args = self._logger._log_args |
| self.warn = self._logger.warn |
| self.info = self._logger.info |
| self.info_all = self._logger.info_all |
| self.log_metric = self._logger.log_metric |
| self.log_new_artifact = self._logger.log_new_artifact |
|
|
| def _generate_default_run_name(self): |
| return timestamp() |
|
|
| def _get_script_name(self): |
| return os.path.basename(__main__.__file__) if '__file__' in dir(__main__) else 'none' |
|
|
| @contextmanager |
| def context(self, consider_failed_if_interrupted=True): |
| try: |
| yield |
|
|
| except KeyboardInterrupt as ex: |
| print('\n\nInterrupted\n\n') |
| self._logger._log_exception(ex.__class__, ex, ex.__traceback__) |
| self._logger._log_all_artifacts() |
|
|
| if consider_failed_if_interrupted: |
| self.exit_status = 'KILLED' |
|
|
| sys.exit(128 + 2) |
|
|
| except Exception as ex: |
| self._logger._log_exception(ex.__class__, ex, ex.__traceback__) |
| self._logger._log_all_artifacts() |
|
|
| self.exit_status = 'FAILED' |
|
|
| raise ex |
|
|
| finally: |
| total_seconds = str(time.time() - self.start_time) + '\n' |
| original_name = str(self.original_name) |
| name = str(self.name) |
|
|
| self.log_new_artifact(os.path.join(self._logger.logs_path, 'elapsed.txt'), total_seconds) |
| self.log_new_artifact(os.path.join(self._logger.logs_path, 'name.original.txt'), original_name) |
| self.log_new_artifact(os.path.join(self._logger.logs_path, 'name.txt'), name) |
|
|
| self._logger._log_all_artifacts() |
|
|
| |
|
|
|
|
| Run = _RunManager() |
|
|
