|
import os
|
|
import sys
|
|
import time
|
|
import __main__
|
|
import traceback
|
|
import mlflow
|
|
|
|
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()
|
|
|
|
mlflow.end_run(status=self.exit_status)
|
|
|
|
|
|
Run = _RunManager()
|
|
|
