DAY = 86400
HOUR = 3600
MINUTE = 60
def remove_backspaces(line):
def report_fatal(msg: str, logger):
def sha256_checksum(filename, block_size=65536):
def filehash(filename, block_size=65536, hashobj=None):
def rand_string(length):
def _looks_like_url(name):
def parse_s3_path(qualified: str):
def retry(f,
          no_retries=5, sleep_time=1,
          exception_class=BaseException, logger=None):
def compression_to_extension(compression):
def compression_to_taropt(compression):
def _compression_to_extension_taropt(compression):
def timeit(method):
    def timed(*args, **kw):
def sixdecode(sval):
def shquote(sval):
def parse_duration(duration_str):
def parse_verbosity(verbosity=None):
def str2duration(sval):
def get_temp_filename() -> str:
def delete_local_path(local_path: str, root: str, shallow: bool):
def delete_local_folders(local_folder_path: str, root: str):
def rm_rf(path):
def check_for_kb_interrupt():
def add_packages(list1, list2): DEBUG = logging.DEBUG
INFO = logging.INFO
ERROR = logging.ERROR
def get_logger(name):
def debug(line):
def error(line):
def info(line): class StorageHandler:
    def __init__(self, storage_type: StorageType,
                 logger,
                 measure_timestamp_diff=False,
                 compression=None):
        self.type = storage_type
        self.logger = logger
        self.compression = compression
        self._timestamp_shift = 0
        if measure_timestamp_diff:
            try:
                self._timestamp_shift = self._measure_timestamp_diff()
            except BaseException:
                check_for_kb_interrupt()
                self._timestamp_shift = 0

    @classmethod
    def get_id(cls, config: Dict) -> str:
        raise NotImplementedError("Not implemented: upload_file")

    def upload_file(self, key, local_path):
        raise NotImplementedError("Not implemented: upload_file")

    def download_file(self, key, local_path):
        raise NotImplementedError("Not implemented: download_file")

    def download_remote_path(self, remote_path, local_path):
        raise NotImplementedError("Not implemented: download_remote_path")

    def delete_file(self, key, shallow=True):
        raise NotImplementedError("Not implemented: delete_file")

    def get_file_url(self, key, method='GET'):
        raise NotImplementedError("Not implemented: get_file_url")

    def get_file_timestamp(self, key):
        raise NotImplementedError("Not implemented: get_file_timestamp")

    def get_qualified_location(self, key):
        raise NotImplementedError("Not implemented: get_qualified_location")

    def get_local_destination(self, remote_path: str):
        raise NotImplementedError("Not implemented: get_local_destination")

    def get_timestamp_shift(self):
        return self._timestamp_shift

    def get_compression(self):
        return self.compression

    def cleanup(self):
        pass

    def _measure_timestamp_diff(self):
        max_diff = 60
        tmpfile = get_temp_filename() + '.txt'
        with open(tmpfile, 'w') as f:
            f.write('timestamp_diff_test')
        key = 'tests/' + str(uuid.uuid4())
        self.upload_file(key, tmpfile)
        remote_timestamp = self.get_file_timestamp(key)

        if remote_timestamp is not None:
            now_remote_diff = time.time() - remote_timestamp
            self.storage_handler.delete_file(key)
            os.remove(tmpfile)

            assert -max_diff < now_remote_diff and \
                now_remote_diff < max_diff, \
                "Timestamp difference is more than 60 seconds. " + \
                "You'll need to adjust local clock for caching " + \
                "to work correctly"

            return -now_remote_diff if now_remote_diff < 0 else 0 class Artifact:

    def __init__(self, art_name, art_dict, logger=None):
        self.name = art_name
        self.key: str = None
        self.local_path: str = None
        self.remote_path: str = None
        self.credentials = None
        self.hash = None

        self.logger = logger
        if self.logger is None:
            self.logger = logs.get_logger(self.__class__.__name__)
            self.logger.setLevel(storage_setup.get_storage_verbose_level())

        self.storage_handler: StorageHandler = None

        self.unpack: bool = art_dict.get('unpack')
        self.is_mutable: bool = art_dict.get('mutable')
        if 'key' in art_dict.keys():
            self.key = art_dict['key']
        if 'local' in art_dict.keys():
            self.local_path = art_dict['local']
        if 'qualified' in art_dict.keys():
            self.remote_path = art_dict['qualified']
        if 'url' in art_dict.keys():
            self.remote_path = art_dict['url']
        if 'hash' in art_dict.keys():
            self.hash = art_dict['hash']
        self.credentials = credentials.Credentials.get_credentials(art_dict)

        self._setup_storage_handler(art_dict)

    def upload(self, local_path=None):
        if self.storage_handler is None:
            msg: str = "No storage handler is set for artifact {0}"\
                .format(self.key)
            util.report_fatal(msg, self.logger)

        if local_path is None:
            local_path = self.local_path

        if self.in_blobstore:
            msg: str = ('Artifact with key {0} exists in blobstore, ' +
                        'skipping the upload').format(self.key)
            self.logger.debug(msg)
            return self.key

        if os.path.exists(local_path):
            tar_filename =\
                tar_artifact(local_path, self.key,
                             self.get_compression(), self.logger)
            if self.key is None:
                self.key = 'blobstore/' + util.sha256_checksum(tar_filename) \
                      + '.tar' + util.compression_to_extension(self.get_compression())
                time_stamp = self.storage_handler.get_file_timestamp(self.key)
                if time_stamp is not None:
                    self.logger.debug(
                        'Artifact with key %s exists in blobstore, skipping the upload',
                        self.key)
                    os.remove(tar_filename)
                    return self.key

            self.storage_handler.upload_file(self.key, tar_filename)
            os.remove(tar_filename)
            return self.key
        self.logger.debug(
            "Local path %s does not exist. Not uploading anything.",
            local_path)
        return None

    def get_compression(self):
        if self.storage_handler is not None:
            return self.storage_handler.get_compression()
        return None

    def _download_no_key_artifact(self):
        if self.is_mutable:
            self.logger.info("Downloading mutable artifact: %s",
                             self.name)
        if self.remote_path is None:
            msg: str =\
                "CANNOT download artifact without remote path: {0}"\
                    .format(self.name)
            util.report_fatal(msg, self.logger)

        key = self._generate_key()
        local_path = artifacts_tracker.get_blob_cache(key)
        local_path =\
            self._get_target_local_path(local_path, self.remote_path)
        if os.path.exists(local_path):
            msg: str = ('Immutable artifact exists at local_path {0},' +
                        ' skipping the download').format(local_path)
            self.logger.debug(msg)
            self.local_path = local_path
            return local_path

        if self.storage_handler.type == StorageType.storageDockerHub or \
           self.storage_handler.type == StorageType.storageSHub:
            msg: str = ('Qualified {0} points to a shub or dockerhub,' +
                        ' skipping the download').format(self.remote_path)
            self.logger.debug(msg)
            return self.remote_path

        self.storage_handler.download_remote_path(
            self.remote_path, local_path)

        self.logger.debug('Downloaded file %s from external source %s',
                          local_path, self.remote_path)
        self.local_path = local_path
        #self.key = key
        return self.local_path

    def _has_newer_artifact(self, local_path) -> bool:
        self.logger.debug(
            'Comparing date of the artifact %s in storage with local %s',
            self.key, local_path)
        storage_time = self.storage_handler.get_file_timestamp(self.key)
        local_time = os.path.getmtime(local_path)
        if storage_time is None:
            msg: str = \
                ("Unable to get storage timestamp for {0}, storage is either " + \
                 "corrupted or has not finished uploading").format(self.key)
            self.logger.info(msg)
            return False

        timestamp_shift = self.storage_handler.get_timestamp_shift()
        if local_time > storage_time - timestamp_shift:
            self.logger.debug(
                "Local path %s is younger than stored %s, skipping the download",
                local_path, self.key)
            return False

        return True

    def _download_and_untar_artifact(self, local_path):
        tar_filename: str = util.get_temp_filename()
        self.logger.debug("tar_filename = %s", tar_filename)

        # Now download our artifact from studio.storage and untar it:
        try:
            result: bool = \
                self.storage_handler.download_file(self.key, tar_filename)
            if not result:
                msg: str = \
                    "FAILED to download {0}.".format(self.key)
                self.logger.info(msg)
                return None
        except BaseException as exc:
            util.check_for_kb_interrupt()
            msg: str = \
                "FAILED to download {0}: {1}.".format(self.key, exc)
            self.logger.info(msg)
            return None

        if os.path.exists(tar_filename):
            untar_artifact(local_path, tar_filename, self.logger)
            os.remove(tar_filename)
            self.local_path = local_path
            return local_path
        self.logger.info('file %s download failed', tar_filename)
        return None

    def download(self, local_path=None, only_newer=True):
        if self.storage_handler is None:
            msg: str = "No storage handler is set for artifact {0}" \
                .format(self.key)
            util.report_fatal(msg, self.logger)

        if self.key is None:
            return self._download_no_key_artifact()

        if local_path is None:
            if self.local_path is not None and \
                    os.path.exists(self.local_path):
                local_path = self.local_path
            else:
                if self.is_mutable:
                    local_path = artifacts_tracker.get_artifact_cache(self.key)
                else:
                    local_path = artifacts_tracker.get_blob_cache(self.key)
                    if os.path.exists(local_path):
                        msg: str = ('Immutable artifact exists at local_path {0},' +
                                    ' skipping the download').format(local_path)
                        self.logger.debug(msg)
                        self.local_path = local_path
                        return local_path

        local_path = re.sub(r'\/\Z', '', local_path)
        self.logger.debug("Downloading dir %s to local path %s from studio.storage...",
                          self.key, local_path)

        if only_newer and os.path.exists(local_path):
            if not self._has_newer_artifact(local_path):
                return local_path

        # Now download our artifact from studio.storage and untar it:
        return self._download_and_untar_artifact(local_path)

    def _get_target_local_path(self, local_path: str, remote_path: str):
        result: str = local_path
        dir_name, file_name = \
            self.storage_handler.get_local_destination(remote_path)
        if dir_name is not None:
            result = os.path.join(result, dir_name)
        if file_name is not None:
            result = os.path.join(result, file_name)
        return result

    def delete(self):
        if self.key is not None:
            self.logger.debug('Deleting artifact: %s', self.key)
            self.storage_handler.delete_file(self.key, shallow=False)


    def get_url(self, method='GET', get_timestamp=False):
        if self.key is not None:
            url = self.storage_handler.get_file_url(self.key, method=method)
        elif self.storage_handler.type == StorageType.storageHTTP:
            url = self.remote_path
        else:
            url = None

        if get_timestamp:
            timestamp = self.storage_handler.get_file_timestamp(self.key)
            return url, timestamp
        return url

    def _looks_like_local_file(self, url: str) -> bool:
        if url is None:
            return False

        local_prefix: str = "/"
        result = url.startswith(local_prefix)
        return result

    def stream(self):
        url = self.get_url()
        if url is None:
            return None

        # pylint: disable=consider-using-with

        # if our url is actually a local file reference
        # (can happen in local execution mode)
        # then we just open a local file:
        if self._looks_like_local_file(url):
            fileobj = open(url, 'rb')
        else:
            fileobj = urlopen(url)

        if fileobj:
            try:
                retval = tarfile.open(fileobj=fileobj, mode='r|*')
                return retval
            except BaseException as exc:
                util.check_for_kb_interrupt()
                fileobj.close()
                msg: str = 'FAILED to stream artifact {0}: {1}'.format(url, exc)
                util.report_fatal(msg, self.logger)
        return None


    def get_hash(self, local_path=None):

        if local_path is None:
            local_path = self.local_path

        if local_path is None or not os.path.exists(local_path):
            return self._generate_key()

        tar_filename =\
            tar_artifact(local_path, self.key,
                         self.get_compression(), self.logger)

        try:
            retval = util.sha256_checksum(tar_filename)
            os.remove(tar_filename)
            self.logger.debug('deleted local artifact file %s', tar_filename)
            return retval
        except BaseException as exc:
            util.check_for_kb_interrupt()
            self.logger.error(
                'error generating a hash for %s: %s',
                    tar_filename, repr(exc))
        return None

    def _is_s3_endpoint(self) -> bool:
        if self.remote_path is None:
            return False
        if self.remote_path.startswith('s3://'):
            return True
        if self.credentials is not None and\
            self.credentials.get_type() == credentials.AWS_TYPE:
            return True
        return False

    def _build_s3_config(self, art_dict):
        """
        For art_dict representing external S3-based artifact,
        build configuration suitable for constructing
        S3-based storage handler for this artifact.
        Returns: (configuration dictionary, artifact's S3 key)
        """
        url, bucket, key = util.parse_s3_path(self.remote_path)
        config = dict()
        config['endpoint'] = "http://{0}".format(url)
        config['bucket'] = bucket
        config[credentials.KEY_CREDENTIALS] =\
            self.credentials.to_dict() if self.credentials else dict()
        if 'region' in art_dict.keys():
            config['region'] = art_dict['region']

        return config, key

    def _build_http_config(self):
        """
        For external Http-based artifact,
        build configuration suitable for constructing
        Http-based storage handler for this artifact.
        """
        config = dict()
        config['endpoint'] = self.remote_path
        config[credentials.KEY_CREDENTIALS] =\
            self.credentials.to_dict() if self.credentials else dict()
        return config

    def _setup_storage_handler(self, art_dict):
        if self.key is not None:
            # Artifact is already stored in our shared blob-cache:
            self.storage_handler = storage_setup.get_storage_artifact_store()
            return

        if self.remote_path is not None:
            if self._is_s3_endpoint():
                s3_config_dict, _ = self._build_s3_config(art_dict)
                factory = StorageHandlerFactory.get_factory()
                self.storage_handler =\
                    factory.get_handler(StorageType.storageS3, s3_config_dict)
                return

            if self.remote_path.startswith('http://') or \
               self.remote_path.startswith('https://'):
                http_config_dict: Dict = self._build_http_config()
                factory = StorageHandlerFactory.get_factory()
                self.storage_handler =\
                    factory.get_handler(StorageType.storageHTTP, http_config_dict)
                return

        if self.local_path is not None:
            self.storage_handler =\
                storage_setup.get_storage_artifact_store()
            return

        raise NotImplementedError(
            "FAILED to setup storage handler for artifact: {0} {1}"
                .format(self.name, repr(art_dict)))

    def to_dict(self):
        result = dict()
        result['unpack'] = self.unpack
        result['mutable'] = self.is_mutable
        if self.key is not None:
            result['key'] = self.key
        if self.local_path is not None:
            result['local'] = self.local_path
        if self.remote_path is not None:
            if self.storage_handler.type == StorageType.storageHTTP:
                result['url'] = self.remote_path
            else:
                result['qualified'] = self.remote_path

        if self.storage_handler.type == StorageType.storageS3:
            # Get artifact bucket directly from remote_path:
            _, bucket, _ = util.parse_s3_path(self.remote_path)
            result['bucket'] = bucket

        if self.credentials is not None:
            result[credentials.KEY_CREDENTIALS] = self.credentials.to_dict()

        return result


    @property
    def in_blobstore(self) -> bool:
        if self.key is not None and self.key.startswith('blobstore/') and \
           self.storage_handler.get_file_timestamp(self.key) is not None:
            return True
        return False

    def _generate_key(self):
        return hashlib.sha256(self.remote_path.encode()).hexdigest() class Experiment:
    """Experiment information."""

    def __init__(self, key, filename, args, pythonenv,
                 project=None,
                 artifacts=None,
                 status='waiting',
                 resources_needed=None,
                 time_added=None,
                 time_started=None,
                 time_last_checkpoint=None,
                 time_finished=None,
                 info=None,
                 git=None,
                 metric=None,
                 pythonver=None,
                 max_duration=None,
                 owner=None,
                 from_compl_service: bool = False):
        if info is None:
            info = {}

        self.key = key
        self._build_args(filename, args)
        self.owner = owner if owner else None
        self.from_compl_service = from_compl_service

        self.pythonenv = pythonenv
        self.project = project
        self.pythonver = pythonver
        if self.pythonver is None:
            self.pythonver = "{0}.{1}"\
                .format(sys.version_info[0], sys.version_info[1])

        self._build_artifacts(key, artifacts)

        self.resources_needed = resources_needed
        self.status = status
        self.time_added = time_added
        self.time_started = time_started
        self.time_last_checkpoint = time_last_checkpoint
        self.time_finished = time_finished
        self.info = info
        self.git = git
        self.metric = metric
        self.max_duration = max_duration

    def _build_args(self, filename, args):
        self.args = []
        self.filename = filename

        if filename and '::' in filename:
            self.filename = '-m'
            module_name = filename.replace('::', '.')
            if module_name.startswith('.'):
                module_name = module_name[1:]

            self.args.append(module_name)

        if args:
            self.args += args

        self.args = [shquote(a) for a in self.args]

    def _build_artifacts(self, key, artifacts):
        if artifacts is not None:
            for tag, art_dict in artifacts.items():
                if art_dict.get('mutable', False):
                    art_dict['local'] = artifacts_tracker.get_artifact_cache(tag, key)

        # Build table of experiment artifacts:
        self.artifacts = dict()
        for tag, art in artifacts.items():
            self.artifacts[tag] = Artifact(tag, art)

    def to_dict(self):
        result = dict()
        result['key'] = self.key
        result['args'] = self.args
        result['filename'] = self.filename
        result['pythonenv'] = self.pythonenv
        result['project'] = self.project
        result['pythonver'] = self.pythonver
        result['resources_needed'] = self.resources_needed
        result['status'] = self.status
        result['time_added'] = self.time_added
        result['time_started'] = self.time_started
        result['time_last_checkpoint'] = self.time_last_checkpoint
        result['time_finished'] = self.time_finished
        result['info'] = self.info
        result['git'] = self.git
        result['metric'] = self.metric
        result['max_duration'] = self.max_duration
        result['owner'] = self.owner
        result['from_compl_service'] = self.from_compl_service

        # Process artifacts:
        artifacts_dict = dict()
        for art_name, art in self.artifacts.items():
            artifacts_dict[art_name] = art.to_dict()
        result['artifacts'] = artifacts_dict
        return result def experiment_from_dict(data, info=None):
    if info is None:
        info = {}

    try:
        return Experiment(
            key=data['key'],
            filename=data['filename'],
            args=data.get('args'),
            pythonenv=data['pythonenv'],
            project=data.get('project'),
            status=data['status'],
            artifacts=data.get('artifacts'),
            resources_needed=data.get('resources_needed'),
            time_added=data['time_added'],
            time_started=data.get('time_started'),
            time_last_checkpoint=data.get('time_last_checkpoint'),
            time_finished=data.get('time_finished'),
            info=info if any(info) else data.get('info'),
            git=data.get('git'),
            metric=data.get('metric'),
            pythonver=data.get('pythonver'),
            max_duration=data.get('max_duration'),
            owner=data.get('owner'),
            from_compl_service=data.get('from_compl_service', False)
        )
    except KeyError as exc:
        raise exc def get_storage_verbose_level():
    global _storage_verbose_level
    return _storage_verbose_level def retry(f,
          no_retries=5, sleep_time=1,
          exception_class=BaseException, logger=None):
    for i in range(no_retries):
        try:
            return f()
        except exception_class as exc:
            check_for_kb_interrupt()
            if i == no_retries - 1:
                raise exc

            if logger:
                logger.info(
                    ('Exception {0} is caught, ' +
                     'sleeping {1}s and retrying (attempt {2} of {3})')
                        .format(exc, sleep_time, i, no_retries))
            time.sleep(sleep_time)
    return None def report_fatal(msg: str, logger):
    logger.error(msg)
    raise ValueError(msg) def compression_to_extension(compression):
    return _compression_to_extension_taropt(compression)[0] def check_for_kb_interrupt():
    current_exc = sys.exc_info()[1]
    if current_exc is None:
        return
    if isinstance(current_exc, KeyboardInterrupt):
        raise current_exc
    return
import time
from concurrent.futures import ThreadPoolExecutor, wait
from studio.util import util, logs
from studio.storage.storage_handler import StorageHandler
from studio.artifacts.artifact import Artifact
from studio.experiments.experiment import Experiment, experiment_from_dict
from studio.storage.storage_setup import get_storage_verbose_level
from studio.util.util import retry, report_fatal,\
    compression_to_extension, check_for_kb_interrupt



class KeyValueProvider:
    """Data provider for managing experiment lifecycle."""

    def __init__(
            self,
            db_config,

            handler: StorageHandler,