Repository: seung-lab/cloud-volume
Pull request #592 (2023-07-20 17:15:45)

feat: add cache multi-processing locking feature
Adds `CloudVolume(..., cache_locking=True)` flag (true by default). 

This starts using the new CloudFiles inter-process locking mechanism which uses `fasteners` to create locking files for local cache access but not for `file://` access generally.

This may have some performance implications, but can be switched off. It's only enabled for cache because otherwise the number of lock files will grow without bound. Caches get cleared periodically.


Patch:
diff --git a/README.md b/README.md
index 41077ac..20ecb12 100644
--- a/README.md
+++ b/README.md
@@ -416,7 +416,8 @@ CloudVolume(
     background_color:int=0, green_threads:bool=False, use_https:bool=False,
     max_redirects:int=10, mesh_dir:Optional[str]=None, skel_dir:Optional[str]=None, 
     agglomerate:bool=False, secrets:SecretsType=None, 
-    spatial_index_db:Optional[str]=None, lru_bytes:int = 0
+    spatial_index_db:Optional[str]=None, lru_bytes:int = 0,
+    cache_locking:bool = True
 )
 ```
 
@@ -448,7 +449,7 @@ CloudVolume(
 
           Note: This cache is totally separate from the LRU controlled by 
           lru_bytes.
-
+*      cache_locking: (bool) The local cache will use file locks via fasteners to prevent issues with multi-process cache access. If this is not a concern, performance can be slightly improved by setting this to False. This uses CloudFiles' locking mechanism.
 *      cdn_cache: (int, bool, or str) Sets Cache-Control HTTP header on uploaded 
         image files. Most cloud providers perform some kind of caching. As of 
         this writing, Google defaults to 3600 seconds. Most of the time you'll 
diff --git a/cloudvolume/cacheservice.py b/cloudvolume/cacheservice.py
index 802fec2..ff5e7f0 100644
--- a/cloudvolume/cacheservice.py
+++ b/cloudvolume/cacheservice.py
@@ -94,8 +94,14 @@ class CacheService(object):
     else:
       return mip_size(self.config.mip)
 
+  def cloudfiles(self):
+    return CloudFiles(
+      'file://' + self.path, 
+      locking=self.config.cache_locking
+    )
+
   def has(self, filename):
-    return CloudFiles('file://' + self.path).exists(filename)
+    return self.cloudfiles().exists(filename)
     
   def list(self, mip=None):
     mip = self.config.mip if mip is None else mip
@@ -161,6 +167,8 @@ class CacheService(object):
     
     Return: void
     """
+    self.cloudfiles().clear_locks()
+
     if not os.path.exists(self.path):
       return
 
@@ -299,16 +307,15 @@ class CacheService(object):
       """.format(cached_prov.serialize(), fresh_prov.serialize()))
 
   def get_json(self, filename):
-    return CloudFiles('file://' + self.path).get_json(filename)
+    return self.cloudfiles().get_json(filename)
 
   def maybe_cache_info(self):
     if self.enabled:
-      cf = CloudFiles('file://' + self.path)
-      cf.put_json('info', self.meta.info)
+      self.cloudfiles().put_json('info', self.meta.info)
 
   def maybe_cache_provenance(self):
     if self.enabled and self.meta.provenance:
-      cf = CloudFiles('file://' + self.path)
+      cf = self.cloudfiles()
       cf.put('provenance', self.meta.provenance.serialize().encode('utf8'))
 
   def upload_single(self, filename, content, *args, **kwargs):
@@ -320,7 +327,12 @@ class CacheService(object):
 
     progress = progress if progress is not None else self.config.progress
 
-    cf = CloudFiles(self.meta.cloudpath, progress=progress, secrets=self.config.secrets)
+    cf = CloudFiles(
+      self.meta.cloudpath, 
+      progress=progress, 
+      secrets=self.config.secrets,
+      locking=self.config.cache_locking,
+    )
     files = list(compression.transcode(files, encoding=compress, level=compress_level))
     cf.puts( 
       files, 
@@ -333,7 +345,11 @@ class CacheService(object):
 
     if self.enabled:
       self.put(files, compress=compress)
-      cf_cache = CloudFiles('file://' + self.path, progress=('to Cache' if progress else None))
+      cf_cache = CloudFiles(
+        'file://' + self.path, 
+        progress=('to Cache' if progress else None),
+        locking=self.config.cache_locking,
+      )
       cf_cache.puts(
         files,
         compress=compress,
@@ -370,7 +386,11 @@ class CacheService(object):
       if locs['local']:
         return self.get_single(local_alias)
 
-    filedata = CloudFiles(self.meta.cloudpath, secrets=self.config.secrets)[path, start:end]
+    filedata = CloudFiles(
+      self.meta.cloudpath, 
+      secrets=self.config.secrets,
+      locking=self.config.cache_locking,
+    )[path, start:end]
 
     if self.enabled:
       self.put([ (local_alias, filedata) ], compress=compress)
@@ -422,6 +442,7 @@ class CacheService(object):
       progress=progress, 
       secrets=self.config.secrets,
       parallel=self.config.parallel,
+      locking=self.config.cache_locking,
     )
     remote_fragments = cf.get(
       ( { 'path': p[0], 'start': p[1], 'end': p[2] } for p in remote_path_tuples ),
@@ -483,6 +504,7 @@ class CacheService(object):
       progress=progress, 
       secrets=self.config.secrets,
       parallel=self.config.parallel,
+      locking=self.config.cache_locking,
     )
     remote_fragments = cf.get(locs['remote'], raw=True)
 
@@ -491,7 +513,11 @@ class CacheService(object):
         raise frag['error']
 
     if self.enabled:
-      cf_cache = CloudFiles('file://' + self.path, progress=('to Cache' if progress else None))
+      cf_cache = CloudFiles(
+        'file://' + self.path, 
+        progress=('to Cache' if progress else None),
+        locking=self.config.cache_locking,
+      )
       cf_cache.puts(
         compression.transcode(
           ( frag for frag in remote_fragments if frag['content'] is not None ),
@@ -516,7 +542,11 @@ class CacheService(object):
   def get(self, cloudpaths, progress=None):
     progress = self.config.progress if progress is None else progress
     
-    cf = CloudFiles('file://' + self.path, progress=progress)
+    cf = CloudFiles(
+      'file://' + self.path, 
+      progress=progress, 
+      locking=self.config.cache_locking,
+    )
     return cf.get(cloudpaths, return_dict=True)
 
   def put_single(self, path, content, *args, **kwargs):
@@ -540,7 +570,11 @@ class CacheService(object):
     
     save_location = 'file://' + self.path
     progress = 'to Cache' if progress else None
-    cf = CloudFiles(save_location, progress=progress)
+    cf = CloudFiles(
+      save_location, 
+      progress=progress, 
+      locking=self.config.cache_locking
+    )
     cf.puts(
       files, 
       compress=compress, 
diff --git a/cloudvolume/chunks.py b/cloudvolume/chunks.py
index eb00e23..637d2cd 100644
--- a/cloudvolume/chunks.py
+++ b/cloudvolume/chunks.py
@@ -17,6 +17,7 @@ import zlib
 import io
 import numpy as np
 
+import fpzip 
 import pyspng
 import simplejpeg
 import compresso
@@ -28,6 +29,7 @@ from PIL import Image
 
 from .lib import yellow, nvl
 
+# This is left in place mainly for testing
 try:
   import compressed_segmentation as cseg
   ACCELERATED_CSEG = True # C extension version
@@ -36,18 +38,6 @@ except ImportError:
 
 from . import py_compressed_segmentation as csegpy
 
-try:
-  import fpzip 
-except ImportError:
-  fpziperrormsg = yellow("CloudVolume: fpzip codec is not available. Was it installed? pip install fpzip")
-  class fpzip(): # type: ignore
-    @classmethod
-    def compress(cls, content):
-      raise NotImplementedError(fpziperrormsg)
-    @classmethod
-    def decompress(cls, content):
-      raise NotImplementedError(fpziperrormsg)
-
 SUPPORTED_ENCODINGS = (
   "raw", "kempressed", "fpzip",
   "compressed_segmentation", "compresso",
diff --git a/cloudvolume/cloudvolume.py b/cloudvolume/cloudvolume.py
index 219b89d..ebbe2dd 100644
--- a/cloudvolume/cloudvolume.py
+++ b/cloudvolume/cloudvolume.py
@@ -35,7 +35,7 @@ class SharedConfiguration(object):
     self, cdn_cache:bool, compress:CompressType, 
     compress_level:Optional[int], green:bool,
     mip:int, parallel:ParallelType, progress:bool, secrets:SecretsType,
-    spatial_index_db:Optional[str],
+    spatial_index_db:Optional[str], cache_locking:bool,
     *args, **kwargs
   ):
     if type(parallel) == bool:
@@ -46,6 +46,7 @@ class SharedConfiguration(object):
       parallel = int(parallel)
 
     self.cdn_cache = cdn_cache
+    self.cache_locking = bool(cache_locking)
     self.compress = compress
     self.compress_level = compress_level
     self.green = bool(green)
@@ -68,7 +69,8 @@ class CloudVolume:
     background_color:int=0, green_threads:bool=False, use_https:bool=False,
     max_redirects:int=10, mesh_dir:Optional[str]=None, skel_dir:Optional[str]=None, 
     agglomerate:bool=False, secrets:SecretsType=None, 
-    spatial_index_db:Optional[str]=None, lru_bytes:int = 0
+    spatial_index_db:Optional[str]=None, lru_bytes:int = 0,
+    cache_locking:bool = True
   ):
     """
     A "serverless" Python client for reading and writing arbitrarily large 
@@ -131,7 +133,8 @@ class CloudVolume:
 
           Note: This cache is totally separate from the LRU controlled by 
           lru_bytes.
-
+      cache_locking: (bool) Use CloudFiles inter-process file locking to make cache 
+        access safe for multiple simultaneous processes.
       cdn_cache: (int, bool, or str) Sets Cache-Control HTTP header on uploaded 
         image files. Most cloud providers perform some kind of caching. As of 
         this writing, Google defaults to 3600 seconds. Most of the time you'll 
diff --git a/cloudvolume/datasource/boss/__init__.py b/cloudvolume/datasource/boss/__init__.py
index b76d7ca..43e6a19 100644
--- a/cloudvolume/datasource/boss/__init__.py
+++ b/cloudvolume/datasource/boss/__init__.py
@@ -1,3 +1,5 @@
+from typing import Optional, Union
+
 from .image import BossImageSource
 from .metadata import BossMetadata
 
@@ -5,15 +7,20 @@ from ...frontends.precomputed import CloudVolumePrecomputed
 
 from .. import get_cache_path
 from ...cacheservice import CacheService
-from ...cloudvolume import SharedConfiguration, register_plugin
+from ...cloudvolume import (
+  register_plugin, SharedConfiguration,
+  CompressType, ParallelType, CacheType,
+  SecretsType
+)
 from ...paths import strict_extract
 
 def create_boss(
-    cloudpath, mip=0, bounded=True, autocrop=False,
-    fill_missing=False, cache=False, compress_cache=None,
-    cdn_cache=True, progress=False, info=None, provenance=None,
-    compress=None, non_aligned_writes=False, parallel=1,
-    delete_black_uploads=False, green_threads=False
+    cloudpath, mip:int = 0, bounded:bool = True, autocrop:bool = False,
+    fill_missing:bool = False, cache:CacheType = False, compress_cache:CompressType = None,
+    cdn_cache:bool = True, progress:bool = False, info:Optional[dict] = None, 
+    provenance:Optional[dict] = None, compress:CompressType = None, 
+    non_aligned_writes:bool = False, parallel:int = 1, delete_black_uploads:bool = False, 
+    green_threads:bool = False, cache_locking:bool = True,
   ):
     path = strict_extract(cloudpath)
     config = SharedConfiguration(
@@ -24,6 +31,9 @@ def create_boss(
       mip=mip,
       parallel=parallel,
       progress=progress,
+      secrets=secrets,
+      spatial_index_db=None,
+      cache_locking=cache_locking,
     )
     cache = CacheService(
       cloudpath=get_cache_path(cache, cloudpath),
diff --git a/cloudvolume/datasource/graphene/__init__.py b/cloudvolume/datasource/graphene/__init__.py
index ec95314..1afc3d5 100644
--- a/cloudvolume/datasource/graphene/__init__.py
+++ b/cloudvolume/datasource/graphene/__init__.py
@@ -26,7 +26,7 @@ def create_graphene(
     mesh_dir:Optional[str]=None, skel_dir:Optional[str]=None, 
     agglomerate:bool=False, secrets:SecretsType=None, 
     spatial_index_db:Optional[str]=None, 
-    lru_bytes:int = 0,
+    lru_bytes:int = 0, cache_locking:bool = True,
     **kwargs
   ):
     from ...frontends import CloudVolumeGraphene
@@ -42,6 +42,7 @@ def create_graphene(
       progress=progress,
       secrets=secrets,
       spatial_index_db=spatial_index_db,
+      cache_locking=cache_locking,
     )
 
     def mkcache(cloudpath):
diff --git a/cloudvolume/datasource/n5/__init__.py b/cloudvolume/datasource/n5/__init__.py
index 7972e67..342670c 100644
--- a/cloudvolume/datasource/n5/__init__.py
+++ b/cloudvolume/datasource/n5/__init__.py
@@ -20,7 +20,8 @@ def create_n5(
   cdn_cache:bool=True, progress:bool=False, 
   compress:CompressType=None, compress_level:Optional[int]=None,
   non_aligned_writes:bool=False, 
-  parallel:ParallelType=1,green_threads:bool=False, secrets:SecretsType=None, 
+  parallel:ParallelType=1,green_threads:bool=False, 
+  secrets:SecretsType=None, cache_locking:bool = True,
   **kwargs # absorb graphene arguments
 ):
     path = strict_extract(cloudpath)
diff --git a/cloudvolume/datasource/precomputed/__init__.py b/cloudvolume/datasource/precomputed/__init__.py
index 6085f0a..35502a2 100644
--- a/cloudvolume/datasource/precomputed/__init__.py
+++ b/cloudvolume/datasource/precomputed/__init__.py
@@ -27,7 +27,7 @@ def create_precomputed(
     green_threads:bool=False, use_https:bool=False,
     max_redirects:int=10, mesh_dir:Optional[str]=None, skel_dir:Optional[str]=None,
     secrets:SecretsType=None, spatial_index_db:Optional[str]=None, 
-    lru_bytes:int = 0,
+    lru_bytes:int = 0, cache_locking:bool = True,
     **kwargs # absorb graphene arguments
   ):
     path = strict_extract(cloudpath)
@@ -41,6 +41,7 @@ def create_precomputed(
       progress=progress,
       secrets=secrets,
       spatial_index_db=spatial_index_db,
+      cache_locking=cache_locking,
     )
 
     cache_service = CacheService(
diff --git a/cloudvolume/datasource/precomputed/image/__init__.py b/cloudvolume/datasource/precomputed/image/__init__.py
index 806dcba..faec1ef 100644
--- a/cloudvolume/datasource/precomputed/image/__init__.py
+++ b/cloudvolume/datasource/precomputed/image/__init__.py
@@ -549,28 +549,19 @@ class PrecomputedImageSource(ImageSourceInterface):
     spec = sharding.ShardingSpecification.from_dict(scale['sharding'])
     return sharding.ShardReader(self.meta, self.cache, spec)
 
-  def make_shard(self, img, bbox, mip=None, spec=None, progress=False):
-    """
-    Convert an image that represents a single complete shard 
-    into a shard file.
-  
-    img: a volumetric numpy array image
-    bbox: the bbox it represents in voxel coordinates
-    mip: if specified, use the sharding specification from 
-      this mip level, otherwise use the sharding spec from
-      the current implicit mip level in config.
-    spec: use the provided specification (overrides mip parameter)
-
-    Returns: (filename, shard_file)
-    """
-    mip = mip if mip is not None else self.config.mip
-    scale = self.meta.scale(mip)
-
+  def shard_spec(self, mip, spec=None):
     if spec is None:
+      scale = self.meta.scale(mip)
       if 'sharding' in scale:
         spec = sharding.ShardingSpecification.from_dict(scale['sharding'])
       else:
         raise ValueError("mip {} does not have a sharding specification.".format(mip))
+    return spec
+
+  def morton_codes(self, bbox, mip=None, spec=None):
+    mip = mip if mip is not None else self.config.mip
+    scale = self.meta.scale(mip)
+    spec = self.shard_spec(mip, spec)
 
     bbox = Bbox.create(bbox)
     if bbox.subvoxel():
@@ -610,6 +601,38 @@ class PrecomputedImageSource(ImageSourceInterface):
         "The gridpoints for this image did not all correspond to the same shard. Got: {}".format(bbox)
       )
 
+    return gpts, morton_codes
+
+  def shard_filename(self, bbox:BboxLikeType, mip:int) -> str:
+    mip = mip if mip is not None else self.config.mip
+    if not self.is_sharded(mip):
+      raise ValueError("Unable to compute filename for unsharded image.")
+
+    spec = self.shard_spec(mip)
+    gpts, morton_codes = self.morton_codes(bbox, mip=mip, spec=spec)
+    reader = self.shard_reader()
+    return reader.get_filename(first(morton_codes))
+
+  def make_shard(self, img, bbox, mip=None, spec=None, progress=False):
+    """
+    Convert an image that represents a single complete shard 
+    into a shard file.
+  
+    img: a volumetric numpy array image
+    bbox: the bbox it represents in voxel coordinates
+    mip: if specified, use the sharding specification from 
+      this mip level, otherwise use the sharding spec from
+      the current implicit mip level in config.
+    spec: use the provided specification (overrides mip parameter)
+
+    Returns: (filename, shard_file)
+    """
+    mip = mip if mip is not None else self.config.mip
+    
+    spec = self.shard_spec(mip, spec)
+    gpts, morton_codes = self.morton_codes(bbox, mip=mip, spec=spec)
+    chunk_size = self.meta.chunk_size(mip)
+
     labels = {}
     pt_anchor = gpts[0] * chunk_size
     for pt_abs, morton_code in zip(gpts, morton_codes):
@@ -626,6 +649,7 @@ class PrecomputedImageSource(ImageSourceInterface):
         compression_params=self.meta.compression_params(mip),
       )
 
+    reader = self.shard_reader(mip=mip)
     shard_filename = reader.get_filename(first(labels.keys()))
 
     return (shard_filename, spec.synthesize_shard(labels, progress=progress))
diff --git a/cloudvolume/datasource/precomputed/image/rx.py b/cloudvolume/datasource/precomputed/image/rx.py
index cc0d4e6..90503be 100644
--- a/cloudvolume/datasource/precomputed/image/rx.py
+++ b/cloudvolume/datasource/precomputed/image/rx.py
@@ -332,9 +332,11 @@ def download_single_voxel_unsharded(
   if locations["local"]:
     cloudpath = cachedir
     cache_enabled = False
+    locking = cache.config.cache_locking
   else:
     cloudpath = meta.cloudpath
     cache_enabled = cache.enabled
+    locking = False
 
   if filename is None:
     if fill_missing:
@@ -349,7 +351,8 @@ def download_single_voxel_unsharded(
       filename, fill_missing,
       cache_enabled, compress_cache,
       secrets, background_color,
-      partial(decode_single_voxel, requested_bbox.minpt - chunk_bbx.minpt)
+      partial(decode_single_voxel, requested_bbox.minpt - chunk_bbx.minpt),
+      decompress=True, locking=locking
     )
 
   if renumber:
@@ -497,12 +500,14 @@ def download_chunk(
   filename, fill_missing,
   enable_cache, compress_cache,
   secrets, background_color,
-  decode_fn, decompress=True
+  decode_fn, decompress=True, locking=False
 ):
   if lru is not None and filename in lru:
     content = lru[filename]
   else:
-    (file,) = CloudFiles(cloudpath, secrets=secrets).get([ filename ], raw=True)
+    (file,) = CloudFiles(
+      cloudpath, secrets=secrets, locking=locking
+    ).get([ filename ], raw=True)
     content = file['content']
 
     if enable_cache:
@@ -512,7 +517,7 @@ def download_chunk(
           in_place=(compress_cache == False)
         )
       )['content'] 
-      CloudFiles('file://' + cache.path).put(
+      cache.cloudfiles().put(
         path=filename, 
         content=(cache_content or b''), 
         content_type=content_type(meta.encoding(mip)), 
@@ -545,13 +550,13 @@ def download_chunks_threaded(
   locations = cache.compute_data_locations(cloudpaths)
   cachedir = 'file://' + cache.path
 
-  def process(cloudpath, filename, enable_cache):
+  def process(cloudpath, filename, enable_cache, locking):
     labels, bbox = download_chunk(
       meta, cache, lru, cloudpath, mip,
       filename, fill_missing,
       enable_cache, compress_cache,
       secrets, background_color,
-      decode_fn, decompress
+      decode_fn, decompress, locking
     )
     fn(labels, bbox)
 
@@ -567,10 +572,11 @@ def download_chunks_threaded(
   qualify = lambda fname: os.path.join(meta.key(mip), os.path.basename(fname))
 
   local_downloads = ( 
-    partial(process, cachedir, qualify(filename), False) for filename in locations['local'] 
+    partial(process, cachedir, qualify(filename), False, cache.config.cache_locking) 
+    for filename in locations['local'] 
   )
   remote_downloads = ( 
-    partial(process, meta.cloudpath, filename, cache.enabled) for filename in locations['remote'] 
+    partial(process, meta.cloudpath, filename, cache.enabled, False) for filename in locations['remote'] 
   )
 
   if progress and not isinstance(progress, str):
diff --git a/cloudvolume/datasource/precomputed/image/tx.py b/cloudvolume/datasource/precomputed/image/tx.py
index 55ba770..d39f23a 100644
--- a/cloudvolume/datasource/precomputed/image/tx.py
+++ b/cloudvolume/datasource/precomputed/image/tx.py
@@ -283,7 +283,7 @@ def threaded_upload_chunks(
     img = img[ ..., np.newaxis ]
 
   remote = CloudFiles(meta.cloudpath, secrets=secrets)
-  local = CloudFiles('file://' + cache.path, secrets=secrets)
+  local = cache.cloudfiles()
 
   remote_compress = should_compress(meta.encoding(mip), compress, cache)
   cache_compress = should_compress(meta.encoding(mip), compress, cache, iscache=True)
diff --git a/cloudvolume/mesh.py b/cloudvolume/mesh.py
index 51a9d16..f944443 100644
--- a/cloudvolume/mesh.py
+++ b/cloudvolume/mesh.py
@@ -80,6 +80,9 @@ class Mesh(object):
 
     if no_self_normals != no_other_normals:
       return False
+
+    if self.vertices.shape[0] != other.vertices.shape[0]:
+      return False
        
     equality = np.all(self.vertices == other.vertices) \
       and np.all(self.faces == other.faces)
diff --git a/requirements.txt b/requirements.txt
index 818c916..2c29af3 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,6 @@
 boto3>=1.4.7
 chardet>=3.0.4
-cloud-files>=2.0.0
+cloud-files>=4.18.1
 compressed-segmentation>=2.1.1
 compresso>=3.0.0
 crackle-codec


Test patch:
diff --git a/test/test_meshing.py b/test/test_meshing.py
index 1400aeb..256d1d3 100644
--- a/test/test_meshing.py
+++ b/test/test_meshing.py
@@ -128,23 +128,23 @@ def test_duplicate_vertices():
 
   # test that 4 is now affected
   mesh2 = deduplicate(mesh, x=4)
-  assert not np.all(mesh.vertices == mesh2.vertices)
+  assert mesh != mesh2
   assert mesh2.vertices.shape[0] == mesh.vertices.shape[0] - 1
 
   mesh2 = deduplicate(mesh, x=3)
-  assert not np.all(mesh.vertices == mesh2.vertices)
+  assert mesh != mesh2
   assert mesh2.vertices.shape[0] == mesh.vertices.shape[0] - 1
 
   mesh2 = deduplicate(mesh, x=4, offset_x=-1)
-  assert not np.all(mesh.vertices == mesh2.vertices)
+  assert mesh != mesh2
   assert mesh2.vertices.shape[0] == mesh.vertices.shape[0] - 1
 
   mesh2 = deduplicate(mesh, x=5)
-  assert not np.all(mesh.vertices == mesh2.vertices)
+  assert mesh != mesh2
   assert mesh2.vertices.shape[0] == mesh.vertices.shape[0] - 1
 
   mesh2 = deduplicate(mesh, x=1)
-  assert not np.all(mesh.vertices == mesh2.vertices)
+  assert mesh != mesh2
   assert mesh2.vertices.shape[0] == mesh.vertices.shape[0] - 3
 
 def test_get_mesh_caching(unsharded_vol):