Use cache for computation Examples -------- >>> cache = Cache(1e9) # doctest: +SKIP The cache can be used locally as a context manager around ``compute`` or ``get`` calls: >>> with cache: # doctest: +SKIP ... result = x.compute() You can also register a cac
| 10 | |
| 11 | |
| 12 | class Cache(Callback): |
| 13 | """Use cache for computation |
| 14 | |
| 15 | Examples |
| 16 | -------- |
| 17 | |
| 18 | >>> cache = Cache(1e9) # doctest: +SKIP |
| 19 | |
| 20 | The cache can be used locally as a context manager around ``compute`` or |
| 21 | ``get`` calls: |
| 22 | |
| 23 | >>> with cache: # doctest: +SKIP |
| 24 | ... result = x.compute() |
| 25 | |
| 26 | You can also register a cache globally, so that it works for all |
| 27 | computations: |
| 28 | |
| 29 | >>> cache.register() # doctest: +SKIP |
| 30 | >>> cache.unregister() # doctest: +SKIP |
| 31 | """ |
| 32 | |
| 33 | def __init__(self, cache, *args, **kwargs): |
| 34 | try: |
| 35 | import cachey |
| 36 | except ImportError as ex: |
| 37 | raise ImportError( |
| 38 | f'Cache requires cachey, "{str(ex)}" problem importing' |
| 39 | ) from ex |
| 40 | self._nbytes = cachey.nbytes |
| 41 | if isinstance(cache, Number): |
| 42 | cache = cachey.Cache(cache, *args, **kwargs) |
| 43 | else: |
| 44 | assert not args and not kwargs |
| 45 | self.cache = cache |
| 46 | self.starttimes = dict() |
| 47 | |
| 48 | def _start(self, dsk): |
| 49 | self.durations = dict() |
| 50 | overlap = set(dsk) & set(self.cache.data) |
| 51 | for key in overlap: |
| 52 | dsk[key] = self.cache.data[key] |
| 53 | |
| 54 | def _pretask(self, key, dsk, state): |
| 55 | self.starttimes[key] = default_timer() |
| 56 | |
| 57 | def _posttask(self, key, value, dsk, state, id): |
| 58 | duration = default_timer() - self.starttimes[key] |
| 59 | deps = state["dependencies"][key] |
| 60 | if deps: |
| 61 | duration += max(self.durations.get(k, 0) for k in deps) |
| 62 | self.durations[key] = duration |
| 63 | nb = self._nbytes(value) + overhead + sys.getsizeof(key) * 4 |
| 64 | self.cache.put(key, value, cost=duration / nb / 1e9, nbytes=nb) |
| 65 | |
| 66 | def _finish(self, dsk, state, errored): |
| 67 | self.starttimes.clear() |
| 68 | self.durations.clear() |
no outgoing calls