DI-engine/ding/utils/file_helper.py

import io
from ditk import logging
import os
import pickle
import time
from functools import lru_cache
from typing import Union

import torch

from .import_helper import try_import_ceph, try_import_redis, try_import_rediscluster, try_import_mc
from .lock_helper import get_file_lock

_memcached = None
_redis_cluster = None

if os.environ.get('DI_STORE', 'off').lower() == 'on':
    print('Enable DI-store')
    from di_store import Client

    di_store_config_path = os.environ.get("DI_STORE_CONFIG_PATH", './di_store.yaml')
    di_store_client = Client(di_store_config_path)

    def save_to_di_store(data):
        return di_store_client.put(data)

    def read_from_di_store(object_ref):
        data = di_store_client.get(object_ref)
        di_store_client.delete(object_ref)
        return data
else:
    save_to_di_store = read_from_di_store = None


@lru_cache()
def get_ceph_package():
    return try_import_ceph()


@lru_cache()
def get_redis_package():
    return try_import_redis()


@lru_cache()
def get_rediscluster_package():
    return try_import_rediscluster()


@lru_cache()
def get_mc_package():
    return try_import_mc()


def read_from_ceph(path: str) -> object:
    """
    Overview:
        Read file from ceph
    Arguments:
        - path (:obj:`str`): File path in ceph, start with ``"s3://"``
    Returns:
        - (:obj:`data`): Deserialized data
    """
    value = get_ceph_package().Get(path)
    if not value:
        raise FileNotFoundError("File({}) doesn't exist in ceph".format(path))

    return pickle.loads(value)


@lru_cache()
def _get_redis(host='localhost', port=6379):
    """
    Overview:
        Ensures redis usage
    Arguments:
        - host (:obj:`str`): Host string
        - port (:obj:`int`): Port number
    Returns:
        - (:obj:`Redis(object)`): Redis object with given ``host``, ``port``, and ``db=0``
    """
    return get_redis_package().StrictRedis(host=host, port=port, db=0)


def read_from_redis(path: str) -> object:
    """
    Overview:
        Read file from redis
    Arguments:
        - path (:obj:`str`): Dile path in redis, could be a string key
    Returns:
        - (:obj:`data`): Deserialized data
    """
    return pickle.loads(_get_redis().get(path))


def _ensure_rediscluster(startup_nodes=[{"host": "127.0.0.1", "port": "7000"}]):
    """
    Overview:
        Ensures redis usage
    Arguments:
        - List of startup nodes (:obj:`dict`) of
            - host (:obj:`str`): Host string
            - port (:obj:`int`): Port number
    Returns:
        - (:obj:`RedisCluster(object)`): RedisCluster object with given ``host``, ``port``, \
            and ``False`` for ``decode_responses`` in default.
    """
    global _redis_cluster
    if _redis_cluster is None:
        _redis_cluster = get_rediscluster_package().RedisCluster(startup_nodes=startup_nodes, decode_responses=False)
    return


def read_from_rediscluster(path: str) -> object:
    """
    Overview:
        Read file from rediscluster
    Arguments:
        - path (:obj:`str`): Dile path in rediscluster, could be a string key
    Returns:
        - (:obj:`data`): Deserialized data
    """
    _ensure_rediscluster()
    value_bytes = _redis_cluster.get(path)
    value = pickle.loads(value_bytes)
    return value


def read_from_file(path: str) -> object:
    """
    Overview:
        Read file from local file system
    Arguments:
        - path (:obj:`str`): File path in local file system
    Returns:
        - (:obj:`data`): Deserialized data
    """
    with open(path, "rb") as f:
        value = pickle.load(f)

    return value


def _ensure_memcached():
    """
    Overview:
        Ensures memcache usage
    Returns:
        - (:obj:`MemcachedClient instance`): MemcachedClient's class instance built with current \
            memcached_client's ``server_list.conf`` and ``client.conf`` files
    """
    global _memcached
    if _memcached is None:
        server_list_config_file = "/mnt/lustre/share/memcached_client/server_list.conf"
        client_config_file = "/mnt/lustre/share/memcached_client/client.conf"
        _memcached = get_mc_package().MemcachedClient.GetInstance(server_list_config_file, client_config_file)
    return


def read_from_mc(path: str, flush=False) -> object:
    """
    Overview:
        Read file from memcache, file must be saved by `torch.save()`
    Arguments:
        - path (:obj:`str`): File path in local system
    Returns:
        - (:obj:`data`): Deserialized data
    """
    _ensure_memcached()
    while True:
        try:
            value = get_mc_package().pyvector()
            if flush:
                _memcached.Get(path, value, get_mc_package().MC_READ_THROUGH)
                return
            else:
                _memcached.Get(path, value)
            value_buf = get_mc_package().ConvertBuffer(value)
            value_str = io.BytesIO(value_buf)
            value_str = torch.load(value_str, map_location='cpu')
            return value_str
        except Exception:
            print('read mc failed, retry...')
            time.sleep(0.01)


def read_from_path(path: str):
    """
    Overview:
        Read file from ceph
    Arguments:
        - path (:obj:`str`): File path in ceph, start with ``"s3://"``, or use local file system
    Returns:
        - (:obj:`data`): Deserialized data
    """
    if get_ceph_package() is None:
        logging.info(
            "You do not have ceph installed! Loading local file!"
            " If you are not testing locally, something is wrong!"
        )
        return read_from_file(path)
    else:
        return read_from_ceph(path)


def save_file_ceph(path, data):
    """
    Overview:
        Save pickle dumped data file to ceph
    Arguments:
        - path (:obj:`str`): File path in ceph, start with ``"s3://"``, use file system when not
        - data (:obj:`Any`): Could be dict, list or tensor etc.
    """
    data = pickle.dumps(data)
    save_path = os.path.dirname(path)
    file_name = os.path.basename(path)
    ceph = get_ceph_package()
    if ceph is not None:
        if hasattr(ceph, 'save_from_string'):
            ceph.save_from_string(save_path, file_name, data)
        elif hasattr(ceph, 'put'):
            ceph.put(os.path.join(save_path, file_name), data)
        else:
            raise RuntimeError('ceph can not save file, check your ceph installation')
    else:
        size = len(data)
        if save_path == 'do_not_save':
            logging.info(
                "You do not have ceph installed! ignored file {} of size {}!".format(file_name, size) +
                " If you are not testing locally, something is wrong!"
            )
            return
        p = os.path.join(save_path, file_name)
        with open(p, 'wb') as f:
            logging.info(
                "You do not have ceph installed! Saving as local file at {} of size {}!".format(p, size) +
                " If you are not testing locally, something is wrong!"
            )
            f.write(data)


def save_file_redis(path, data):
    """
    Overview:
        Save pickle dumped data file to redis
    Arguments:
        - path (:obj:`str`): File path (could be a string key) in redis
        - data (:obj:`Any`): Could be dict, list or tensor etc.
    """
    _get_redis().set(path, pickle.dumps(data))


def save_file_rediscluster(path, data):
    """
    Overview:
        Save pickle dumped data file to rediscluster
    Arguments:
        - path (:obj:`str`): File path (could be a string key) in redis
        - data (:obj:`Any`): Could be dict, list or tensor etc.
    """
    _ensure_rediscluster()
    data = pickle.dumps(data)
    _redis_cluster.set(path, data)
    return


def read_file(path: str, fs_type: Union[None, str] = None, use_lock: bool = False) -> object:
    """
    Overview:
        Read file from path
    Arguments:
        - path (:obj:`str`): The path of file to read
        - fs_type (:obj:`str` or :obj:`None`): The file system type, support ``{'normal', 'ceph'}``
        - use_lock (:obj:`bool`): Whether ``use_lock`` is in local normal file system
    """
    if fs_type is None:
        if path.lower().startswith('s3'):
            fs_type = 'ceph'
        elif get_mc_package() is not None:
            fs_type = 'mc'
        else:
            fs_type = 'normal'
    assert fs_type in ['normal', 'ceph', 'mc']
    if fs_type == 'ceph':
        data = read_from_path(path)
    elif fs_type == 'normal':
        if use_lock:
            with get_file_lock(path, 'read'):
                data = torch.load(path, map_location='cpu')
        else:
            data = torch.load(path, map_location='cpu')
    elif fs_type == 'mc':
        data = read_from_mc(path)
    return data


def save_file(path: str, data: object, fs_type: Union[None, str] = None, use_lock: bool = False) -> None:
    """
    Overview:
        Save data to file of path
    Arguments:
        - path (:obj:`str`): The path of file to save to
        - data (:obj:`object`): The data to save
        - fs_type (:obj:`str` or :obj:`None`): The file system type, support ``{'normal', 'ceph'}``
        - use_lock (:obj:`bool`): Whether ``use_lock`` is in local normal file system
    """
    if fs_type is None:
        if path.lower().startswith('s3'):
            fs_type = 'ceph'
        elif get_mc_package() is not None:
            fs_type = 'mc'
        else:
            fs_type = 'normal'
    assert fs_type in ['normal', 'ceph', 'mc']
    if fs_type == 'ceph':
        save_file_ceph(path, data)
    elif fs_type == 'normal':
        if use_lock:
            with get_file_lock(path, 'write'):
                torch.save(data, path)
        else:
            torch.save(data, path)
    elif fs_type == 'mc':
        torch.save(data, path)
        read_from_mc(path, flush=True)


def remove_file(path: str, fs_type: Union[None, str] = None) -> None:
    """
    Overview:
        Remove file
    Arguments:
        - path (:obj:`str`): The path of file you want to remove
        - fs_type (:obj:`str` or :obj:`None`): The file system type, support ``{'normal', 'ceph'}``
    """
    if fs_type is None:
        fs_type = 'ceph' if path.lower().startswith('s3') else 'normal'
    assert fs_type in ['normal', 'ceph']
    if fs_type == 'ceph':
        os.popen("aws s3 rm --recursive {}".format(path))
    elif fs_type == 'normal':
        os.popen("rm -rf {}".format(path))