Repository: seung-lab/cloud-volume
Pull request #593 (2023-07-26 02:43:33)

feat+refactor: add shard filename function for images
Refactors the image shard generation code to enable also extracting filenames.


Patch:
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):