petrel-oss-python-sdk2 / tests /io_retry_test.py
Weiyun1025's picture
Upload folder using huggingface_hub
2abfccb verified
import logging
import time
import random
from os.path import expanduser, abspath
from petrel_client.client_base import ClientBase
from petrel_client.common.io_profile import profile
from petrel_client.common.io_retry import retry
from petrel_client.common import exception
from petrel_client.common.exception import ObjectNotFoundError, ResourceNotFoundError
from petrel_client.common.config import Config
from petrel_client.common.log import init_log
from petrel_client.common.io_profile import Profiler
LOG = logging.getLogger(__name__)
class Client(ClientBase):
def __init__(self, conf_path, name, count_disp):
conf_path = abspath(expanduser(conf_path))
config = Config(conf_path)
self._default_config = config.default()
init_log(self._default_config)
LOG.info('init io_retry_test.Client, conf_path %s', conf_path)
Profiler.set_default_conf(self._default_config)
super(Client, self).__init__(name=name, count_disp=count_disp)
@profile('get')
def get(self, key):
def not_found():
raise ObjectNotFoundError(key)
def resource_not_found():
raise ResourceNotFoundError()
def error():
raise Exception(key)
def found():
return 'content'
action = random.choice([found, not_found, resource_not_found, error])
time.sleep(0.001)
return action()
@profile('put')
def put(self, key, content):
def normal():
return len(content)
def error():
raise Exception(key)
action = random.choice([normal, error])
return action()
class FakeMixedClient(object):
def __init__(self, client):
self.client = client
def get(self, uri, **kwargs):
@retry('get', exceptions=(Exception,), raises=(exception.ResourceNotFoundError,), tries=3)
def do_get(self, uri, **kwargs):
try:
self.client.get(uri)
except exception.ObjectNotFoundError as err:
LOG.debug(err)
return None
except exception.ResourceNotFoundError as err:
raise
except Exception as err:
raise
return do_get(self, uri, **kwargs)
def put(self, uri, content, **kwargs):
@retry('put', exceptions=(Exception,), tries=3)
def do_put(self, uri, content, **kwargs):
try:
self.client.put(uri, content)
except Exception as err:
raise
return do_put(self, uri, content, **kwargs)
c = Client(conf_path='~/petreloss.conf', name='cluster1', count_disp=50)
mc = FakeMixedClient(c)
for _ in range(50):
try:
mc.get('key')
except Exception:
pass
for _ in range(50):
try:
mc.put('key', '!@#$%'*10)
except Exception:
pass