Merge pull request #88 from neuro-ml/dev

Support for key labels in cache to disk

Author: maxme1

connectome/__version__.py

@@ -1 +1 @@
-__version__ = '0.6.0'
+__version__ = '0.6.1'

connectome/cache/disk.py

@@ -1,16 +1,18 @@
 import logging
 
 from tarn import PickleKeyStorage
+from tarn.interface import MaybeLabels
 
 from .base import Cache
 
 logger = logging.getLogger(__name__)
 
 
 class DiskCache(Cache):
-    def __init__(self, pool: PickleKeyStorage):
+    def __init__(self, pool: PickleKeyStorage, labels: MaybeLabels = None):
         super().__init__()
         self.cache = pool
+        self.labels = labels
 
     def prepare(self, param):
         raw = param.value
@@ -21,4 +23,4 @@ def get(self, key, context):
         return self.cache.read(context, error=False)
 
     def set(self, key, value, context):
-        self.cache.write(context, value, error=False)
+        self.cache.write(context, value, error=False, labels=self.labels)

connectome/layers/cache.py

@@ -6,6 +6,7 @@
 import numpy as np
 from tarn import DiskDict, HashKeyStorage, PickleKeyStorage
 from tarn.config import StorageConfig, init_storage
+from tarn.interface import MaybeLabels
 
 from ..cache import Cache, DiskCache, MemoryCache
 from ..containers import EdgesBag, IdentityContext
@@ -114,16 +115,17 @@ class CacheToDisk(CacheToStorage):
     """
 
     def __init__(self, index: PathLikes, storage: HashKeyStorage, serializer: SerializersLike, names: StringsLike, *,
-                 impure: bool = False):
+                 impure: bool = False, labels: MaybeLabels = None):
         super().__init__(names=names, impure=impure)
         names, serializer = _normalize_disk_arguments(names, serializer)
-        self.storage = DiskCache(PickleKeyStorage(index, storage, serializer))
+        self.storage = DiskCache(PickleKeyStorage(index, storage, serializer), labels=labels)
 
     def _get_storage(self) -> Cache:
         return self.storage
 
     @classmethod
-    def simple(cls, *names, root: PathLike, serializer: Union[Serializer, Sequence[Serializer]] = None):
+    def simple(cls, *names, root: PathLike, serializer: Union[Serializer, Sequence[Serializer]] = None, 
+               labels: MaybeLabels = None):
         """
         A simple version of caching to disk with adequate default settings.
 
@@ -158,7 +160,7 @@ def simple(cls, *names, root: PathLike, serializer: Union[Serializer, Sequence[S
                 PickleSerializer(),
             )
 
-        return cls(index, HashKeyStorage(DiskDict(storage)), serializer, names)
+        return cls(index, HashKeyStorage(DiskDict(storage)), serializer, names, labels=labels)
 
 
 def _normalize_disk_arguments(names, serializer):

connectome/layers/columns.py

@@ -3,6 +3,7 @@
 from typing import Any, Generator, Union
 
 from tarn import HashKeyStorage, PickleKeyStorage
+from tarn.interface import MaybeLabels
 from tqdm.auto import tqdm
 
 from ..cache import DiskCache, MemoryCache
@@ -41,7 +42,7 @@ class CacheColumns(DynamicConnectLayer, CacheLayer):
     """
 
     def __init__(self, index: PathLikes, storage: HashKeyStorage, serializer: SerializersLike, names: StringsLike, *,
-                 verbose: bool = False, shard_size: Union[int, float, None] = None):
+                 verbose: bool = False, shard_size: Union[int, float, None] = None, labels: MaybeLabels = None):
         if shard_size == 1:
             raise ValueError(f'Shard size of 1 is ambiguous. Use None if you want to have a single shard')
         names, serializer = _normalize_disk_arguments(names, serializer)
@@ -50,7 +51,7 @@ def __init__(self, index: PathLikes, storage: HashKeyStorage, serializer: Serial
         self.names = names
         self.shard_size = shard_size
         self.verbose = verbose
-        self.disk = DiskCache(PickleKeyStorage(index, storage, serializer))
+        self.disk = DiskCache(PickleKeyStorage(index, storage, serializer), labels=labels)
         self.ram = MemoryCache(None)
 
     def _prepare_container(self, previous: EdgesBag) -> EdgesBag:

connectome/layers/debug.py

@@ -1,17 +1,21 @@
 import hashlib
-from typing import Any, Generator
+from typing import Any, Generator, Type, Union
 
-from tarn.cache.storage import key_to_digest
+from tarn.compat import HashAlgorithm
+from tarn.pickler import dumps
 
 from ..containers import EdgesBag
-from ..engine import Command, Details, Node, NodeHash, NodeHashes, Request, Response, StaticGraph, StaticHash
+from ..engine import (
+    Command, Details, Node, NodeHash, NodeHashes, Request, Response, StaticGraph, StaticHash, CustomHash, LeafHash
+)
 from ..utils import StringsLike
 from .base import CallableLayer
 from .cache import to_seq
 
 
 class HashDigest(CallableLayer):
-    def __init__(self, names: StringsLike, algorithm):
+    def __init__(self, names: StringsLike, algorithm: Union[Type[HashAlgorithm], str, None] = None,
+                 return_value: bool = False):
         if isinstance(algorithm, str):
             algorithm = getattr(hashlib, algorithm)
 
@@ -22,7 +26,7 @@ def __init__(self, names: StringsLike, algorithm):
             inp, out = Node(name, details), Node(name, details)
             inputs.append(inp)
             outputs.append(out)
-            edges.append(HashDigestEdge(algorithm).bind(inp, out))
+            edges.append(HashDigestEdge(algorithm, return_value).bind(inp, out))
 
         super().__init__(EdgesBag(
             inputs, outputs, edges,
@@ -31,16 +35,27 @@ def __init__(self, names: StringsLike, algorithm):
 
 
 class HashDigestEdge(StaticGraph, StaticHash):
-    def __init__(self, algorithm):
+    def __init__(self, algorithm, return_value):
         super().__init__(arity=1)
         self.algorithm = algorithm
+        self.return_value = return_value
 
     def _make_hash(self, inputs: NodeHashes) -> NodeHash:
-        return inputs[0]
+        return CustomHash('connectome.HashDigest', LeafHash(self.algorithm), LeafHash(self.return_value), *inputs)
 
     def evaluate(self) -> Generator[Request, Response, Any]:
-        value = yield Command.ParentValue, 0
-        output = yield Command.CurrentHash,
+        result = []
+        if self.return_value:
+            value = yield Command.ParentValue, 0
+            result.append(value)
 
-        pickled, digest = key_to_digest(self.algorithm, output.value)
-        return value, output.value, digest, pickled
+        node_hash = yield Command.CurrentHash,
+        node_hash = node_hash.value
+        result.append(node_hash)
+
+        pickled = dumps(node_hash)
+        result.append(pickled)
+        if self.algorithm is not None:
+            result.append(self.algorithm(pickled).digest())
+
+        return tuple(result)

requirements.txt

@@ -1,5 +1,5 @@
 numpy
 tqdm
 pylru
-tarn>=0.1.0,<1.0.0
+tarn>=0.5.0,<1.0.0
 jboc<1.0.0

tests/disk_fixtures.py

@@ -5,9 +5,8 @@
 from typing import Iterator
 
 import pytest
-
 from tarn import HashKeyStorage
-from tarn.config import init_storage, StorageConfig
+from tarn.config import StorageConfig, init_storage
 
 from connectome import CacheToDisk
 

tests/graph_fixtures.py

@@ -1,7 +1,7 @@
 import pytest
 
 from connectome.containers import ReversibleContainer
-from connectome.engine import Node, BoundEdge, FunctionEdge
+from connectome.engine import BoundEdge, FunctionEdge, Node
 from connectome.utils import extract_signature
 
 

tests/interface_fixtures.py

@@ -1,7 +1,8 @@
 import re
+
 import pytest
 
-from connectome import Source, Transform, inverse, optional, positional, meta, Function
+from connectome import Function, Source, Transform, inverse, meta, optional, positional
 from connectome.engine.compiler import identity
 
 

tests/test_cache/test_cache.py

@@ -6,17 +6,14 @@
 from threading import Thread
 
 import pytest
+from tarn.config import StorageConfig, init_storage
+from utils import Counter
 
-from tarn import DiskDict
-from tarn.config import init_storage, StorageConfig
-
-from connectome import CacheToRam, Apply, CacheToDisk, CacheColumns, Transform, optional
+from connectome import Apply, CacheColumns, CacheToDisk, CacheToRam, Transform, optional
 from connectome.engine.edges import CacheEdge
 from connectome.interface.nodes import Silent
 from connectome.serializers import JsonSerializer
 
-from utils import Counter
-
 
 def sleeper(s):
     def f(x):
@@ -142,8 +139,6 @@ def visit(storage, root):
         with tempfile.TemporaryDirectory() as temp, storage_factory(locker=locker) as temp_storage:
             temp = Path(temp) / 'cache'
             init_storage(StorageConfig(hash='blake2b', levels=[1, 63], locker=locker), temp)
-            # TODO: remove after tarn is updated
-            DiskDict(temp)
 
             th = Process(target=visit, args=(temp_storage, temp))
             th.start()

tests/test_cache/test_disk_cache.py

@@ -1,11 +1,11 @@
 import os
 
 import pytest
+from tarn.cache.pickler import LATEST_VERSION
 
-from connectome import Transform, CacheToDisk
+from connectome import CacheToDisk, Transform
 from connectome.exceptions import StorageCorruption
 from connectome.serializers import JsonSerializer
-from tarn.cache.pickler import LATEST_VERSION
 
 
 def setup_cache(temp_dir):

tests/test_cache/test_hash.py

@@ -1,14 +1,14 @@
 from pathlib import Path
 from tempfile import TemporaryDirectory
 
-from tarn.config import init_storage, StorageConfig
+from tarn.config import StorageConfig, init_storage
 
-from connectome import Chain, CacheToRam, CacheToDisk, CacheColumns, HashDigest
+from connectome import CacheColumns, CacheToDisk, CacheToRam, Chain, HashDigest
 from connectome.serializers import JsonSerializer
 
 
 def test_hash(block_maker, storage_factory):
-    hash_layer = HashDigest(['image'], 'blake2b')
+    hash_layer = HashDigest(['image'], 'blake2b', return_value=True)
     pipeline = Chain(
         block_maker.first_ds(first_constant=2, ids_arg=15),
         block_maker.crop(),

tests/test_cache/test_layer.py

@@ -1,4 +1,4 @@
-from connectome import Transform, CacheToRam
+from connectome import CacheToRam, Transform
 
 
 def test_nested_virtual():

tests/test_container.py

@@ -3,7 +3,7 @@
 import pytest
 
 from connectome import Transform
-from connectome.exceptions import GraphError, DependencyError
+from connectome.exceptions import DependencyError, GraphError
 
 
 def test_normalization():

tests/test_interface/test_factory.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 
-from connectome import Transform, Function, Source, meta
+from connectome import Function, Source, Transform, meta
 from connectome.exceptions import FieldError, GraphError
 
 

tests/test_interface/test_interface.py

@@ -2,8 +2,8 @@
 
 import pytest
 
-from connectome import positional, Source, Chain, Transform, Merge, meta, Output
-from connectome.interface.metaclasses import TransformBase, SourceBase
+from connectome import Chain, Merge, Output, Source, Transform, meta, positional
+from connectome.interface.metaclasses import SourceBase, TransformBase
 
 
 def test_single_with_params():

tests/test_interface/test_inverse.py

@@ -1,6 +1,6 @@
 import pytest
 
-from connectome import Transform, Apply, inverse
+from connectome import Apply, Transform, inverse
 from connectome.engine.compiler import identity
 from connectome.exceptions import FieldError
 

tests/test_interface/test_metaclasses.py

@@ -1,6 +1,6 @@
 import pytest
 
-from connectome import Source, Transform, Mixin
+from connectome import Mixin, Source, Transform
 
 
 def test_subclasses():

tests/test_interface/test_mixins.py

@@ -1,6 +1,6 @@
 import pytest
 
-from connectome import Mixin, Source, meta, Transform, Output, inverse
+from connectome import Mixin, Output, Source, Transform, inverse, meta
 
 
 class SizeMixin(Mixin):

tests/test_interface/test_prepared.py

@@ -1,6 +1,6 @@
 import pytest
 
-from connectome import Source, meta, impure, Transform, HashDigest, CacheToRam, Merge
+from connectome import CacheToRam, HashDigest, Merge, Source, Transform, impure, meta
 from connectome.interface.complex_edges import hash_by_value
 
 
@@ -23,8 +23,8 @@ class Stripped(Source):
 
 
 def test_hash():
-    one = Computable(length=4) >> HashDigest(['field'], 'blake2b')
-    two = Stripped() >> HashDigest(['field'], 'blake2b')
+    one = Computable(length=4) >> HashDigest(['field'], 'blake2b', return_value=True)
+    two = Stripped() >> HashDigest(['field'], 'blake2b', return_value=True)
     assert one.field('12345678')[0] == 'received 1234'
     assert one.field('12345678') == one.field('1234----') == two.field('1234')
 
@@ -35,7 +35,7 @@ class A(Transform):
         def x(a, b, c):
             return a + b + c
 
-    ds = A() >> HashDigest(['x'], 'blake2b')
+    ds = A() >> HashDigest(['x'], 'blake2b', return_value=True)
     assert ds.x(1, 2, 3) == ds.x(3, 2, 1)
 
 

tests/test_interface/test_validation.py

@@ -1,5 +1,6 @@
 import pytest
-from connectome import Source, Filter
+
+from connectome import Filter, Source
 
 
 def test_default_args():

tests/test_layers/test_chain.py

@@ -2,7 +2,8 @@
 
 import pytest
 from tarn.config import StorageConfig, init_storage
-from connectome import Source, Transform, Chain, CacheToRam, meta, LazyChain, HashDigest, CacheColumns, Merge, optional
+
+from connectome import CacheColumns, CacheToRam, Chain, HashDigest, LazyChain, Merge, Source, Transform, meta, optional
 from connectome.exceptions import DependencyError, FieldError
 
 
@@ -33,7 +34,7 @@ def test_chain():
 def test_nested(block_maker):
     one = block_maker.first_ds(first_constant=2, ids_arg=15)
     two = block_maker.crop()
-    hash_layer = HashDigest('image', 'blake2b')
+    hash_layer = HashDigest('image', 'blake2b', return_value=True)
 
     base, *variants = [
         Chain(one, two, hash_layer),

tests/test_layers/test_filter.py

@@ -1,5 +1,6 @@
 import pytest
-from connectome import Chain, Filter, Source, meta, impure
+
+from connectome import Chain, Filter, Source, impure, meta
 from connectome.engine.base import HashError
 from connectome.exceptions import DependencyError
 from connectome.interface.blocks import HashDigest
@@ -18,7 +19,7 @@ def test_filter(block_maker):
     ids = pipeline.ids
     assert ids == ('4', '14')
 
-    hash_layer = HashDigest(['image', 'lungs', 'spacing'], 'blake2b')
+    hash_layer = HashDigest(['image', 'lungs', 'spacing'], 'blake2b', return_value=True)
     hashed = Chain(block, hash_layer)
     pipeline = Chain(
         block, Filter(lambda image: image.endswith('4')), hash_layer,