checkpointing.cache.base

 1from __future__ import annotations
 2from abc import ABC, abstractmethod
 3from typing import Generic, TypeVar, Union
 4from checkpointing._typing import ContextId, ReturnValue
 5from checkpointing.exceptions import CheckpointNotExist
 6
 7from threading import Lock as ThreadLock
 8from multiprocessing import Lock as ProcessLock
 9
10
11class CacheBase(ABC, Generic[ContextId, ReturnValue]):
12    """
13    The base class for Cache.
14
15    To implement a concrete Cache class, you need to implement the following abstract methods:
16    - `save(self, context_id: ContextId, result: ReturnValue) -> None`
17    - `retrieve(self, context: Context) -> ReturnValue`
18    """
19
20    @abstractmethod
21    def save(self, context_id: ContextId, result: ReturnValue) -> None:
22        """
23        Save the result with the given context id.
24
25        Args:
26            context_id: identifier of the function call context
27            result: return value of the function call
28        """
29        pass
30
31    @abstractmethod
32    def retrieve(self, context_id: ContextId) -> ReturnValue:
33        """
34        Retrieve the function return value with the given context id.
35        If there is no cached results for the context_id, throws a checkpointing.exceptions.CheckpointNotExist
36
37        Args:
38            context_id: identifier of the function call context
39
40        Returns:
41            The return value of the function that corresponds to this context id
42        """
43        pass
44
45    def synchronize_with(self, lock) -> SynchronizedCache:
46        """
47        Args:
48            lock: the Lock object that supports the similar interface as threading/multiprocessing.Lock()
49
50        Returns:
51            A process-safe version of this cache using the given lock.
52        """
53        return SynchronizedCache(self, lock)
54
55
56CacheSubclass = TypeVar("CacheSubclass", bound=CacheBase)
57
58
59class SynchronizedCache(CacheBase, Generic[CacheSubclass]):
60    """
61    Base class for synchronized cache. In the constructor of a concrete subclass,
62    a `self.lock` attribute should be assigned as a thread/process lock, 
63    or any analog of such synchronization locks.
64    """
65
66    def __init__(self, cache, lock) -> None:
67        self.__cache = cache
68        self.__lock = lock
69
70    def save(self, context_id: ContextId, result: ReturnValue) -> None:
71        with self.__lock:
72            self.__cache.save(context_id, result)
73
74    def retrieve(self, context_id: ContextId) -> ReturnValue:
75        with self.__lock:
76            return self.__cache.retrieve(context_id)