Spaces:
Runtime error
Runtime error
File size: 3,467 Bytes
58627fa |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
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")
# TODO: This should timeout and exit (i.e., fail) given no response for 60 seconds.
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' # mlflow.entities.RunStatus.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' # mlflow.entities.RunStatus.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()
|