mirror of
https://github.com/clearml/clearml-server
synced 2025-01-31 10:56:48 +00:00
80 lines
2.6 KiB
Python
80 lines
2.6 KiB
Python
from contextlib import contextmanager
|
|
from typing import Optional, TypeVar, Generic, Type, Callable
|
|
|
|
from redis import StrictRedis
|
|
|
|
import database
|
|
from timing_context import TimingContext
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
def _do_nothing(_: T):
|
|
return
|
|
|
|
|
|
class RedisCacheManager(Generic[T]):
|
|
"""
|
|
Class for store/retrieve of state objects from redis
|
|
|
|
self.state_class - class of the state
|
|
self.redis - instance of redis
|
|
self.expiration_interval - expiration interval in seconds
|
|
"""
|
|
|
|
def __init__(
|
|
self, state_class: Type[T], redis: StrictRedis, expiration_interval: int
|
|
):
|
|
self.state_class = state_class
|
|
self.redis = redis
|
|
self.expiration_interval = expiration_interval
|
|
|
|
def set_state(self, state: T) -> None:
|
|
redis_key = self._get_redis_key(state.id)
|
|
with TimingContext("redis", "cache_set_state"):
|
|
self.redis.set(redis_key, state.to_json())
|
|
self.redis.expire(redis_key, self.expiration_interval)
|
|
|
|
def get_state(self, state_id) -> Optional[T]:
|
|
redis_key = self._get_redis_key(state_id)
|
|
with TimingContext("redis", "cache_get_state"):
|
|
response = self.redis.get(redis_key)
|
|
if response:
|
|
return self.state_class.from_json(response)
|
|
|
|
def delete_state(self, state_id) -> None:
|
|
with TimingContext("redis", "cache_delete_state"):
|
|
self.redis.delete(self._get_redis_key(state_id))
|
|
|
|
def _get_redis_key(self, state_id):
|
|
return f"{self.state_class}/{state_id}"
|
|
|
|
@contextmanager
|
|
def get_or_create_state(
|
|
self,
|
|
state_id=None,
|
|
init_state: Callable[[T], None] = _do_nothing,
|
|
validate_state: Callable[[T], None] = _do_nothing,
|
|
):
|
|
"""
|
|
Try to retrieve state with the given id from the Redis cache if yes then validates it
|
|
If no then create a new one with randomly generated id
|
|
Yield the state and write it back to redis once the user code block exits
|
|
:param state_id: id of the state to retrieve
|
|
:param init_state: user callback to init the newly created state
|
|
If not passed then no init except for the id generation is done
|
|
:param validate_state: user callback to validate the state if retrieved from cache
|
|
Should throw an exception if the state is not valid. If not passed then no validation is done
|
|
"""
|
|
state = self.get_state(state_id) if state_id else None
|
|
if state:
|
|
validate_state(state)
|
|
else:
|
|
state = self.state_class(id=database.utils.id())
|
|
init_state(state)
|
|
|
|
try:
|
|
yield state
|
|
finally:
|
|
self.set_state(state)
|