Add sharding support for Neuroglancer Precomputed tile source

The Rust/WASM implementation fo tile source 14 supports sharded datasets now. The changes of this commit make the client side in CATMAID work with these changes. Some additional improvements are included, too: parallel loading of non-sharded blocks and prefetching data into the block cache with batch support.
catmaid · Jan 16, 2025 · 6ddd131 · 6ddd131
1 parent 98652dd
commit 6ddd131
Show file tree

Hide file tree

Showing 7 changed files with 2,591 additions and 1,937 deletions.
diff --git a/.travis.jshintexpected b/.travis.jshintexpected
@@ -150,8 +150,34 @@ django/applications/catmaid/static/libs/catmaid/time-series.js: line 19, col 12,
 django/applications/catmaid/static/libs/catmaid/volumes.js: line 613, col 9, 'abs' is defined but never used.
 django/applications/catmaid/static/libs/catmaid/volumes.js: line 805, col 15, 'dSimplicesT' is defined but never used.
 
-django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 1817, col 1, 'async functions' is only available in ES8 (use 'esversion: 8').
-django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 1850, col 1, 'async functions' is only available in ES8 (use 'esversion: 8').
-django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 2009, col 29, The object literal notation {} is preferable.
-
-115 errors
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 8, col 9, It's not necessary to initialize 'wasm' to 'undefined'.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 35, col 188, Missing semicolon.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 37, col 76, Unnecessary semicolon.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 55, col 154, Missing semicolon.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 137, col 67, Missing semicolon.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 1942, col 5, 'async functions' is only available in ES8 (use 'esversion: 8').
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 1979, col 22, Missing semicolon.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 1987, col 22, Missing semicolon.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 1991, col 22, Missing semicolon.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 2007, col 22, Missing semicolon.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 2026, col 22, Missing semicolon.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 2030, col 22, Missing semicolon.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 2118, col 22, Missing semicolon.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 2131, col 34, The array literal notation [] is preferable.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 2157, col 35, The object literal notation {} is preferable.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 2171, col 22, Missing semicolon.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 2175, col 22, Missing semicolon.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 2226, col 22, Missing semicolon.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 2333, col 35, Invalid typeof value 'bigint'
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 2371, col 49, Invalid typeof value 'bigint'
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 2469, col 36, Missing semicolon.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 2471, col 107, Missing semicolon.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 2488, col 5, 'async functions' is only available in ES8 (use 'esversion: 8').
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 2494, col 52, Missing semicolon.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm.js: line 2496, col 122, Missing semicolon.
+
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm_worker.js: line 28, col 43, 'use_cache' is defined but never used.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm_worker.js: line 57, col 43, 'use_cache' is defined but never used.
+django/applications/catmaid/static/libs/ngpre-wasm/ngpre_wasm_worker.js: line 57, col 54, '_' is defined but never used.
+
+140 errors
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -634,8 +634,8 @@ Data sources:
 
 - A new tile source has been added: Neuroglancer Precomputed data, with ID 14.
   This image block source works very similar to the N5 tile source. At the
-  moment only non-compressed and non-sharded image volumes are supported. In
-  order for the voxel space coordinates to match between CATMAID and
+  moment only GZip and JPEG compression is supported. Sharded datasets can be
+  loaded. In order for the voxel space coordinates to match between CATMAID and
   Neuroglancer, if the Neuroglancer dataset defines a voxel offset, the
   respective CATMAID stack needs to have its zoom-level zero voxel offset
   defined in the stack meta data in the admin view, e.g. `{"voxelOffset":

diff --git a/django/applications/catmaid/static/js/image-block.js b/django/applications/catmaid/static/js/image-block.js
@@ -35,12 +35,59 @@
       CATMAID.asEventSource(this);
 
       this.queue = new CATMAID.MultiQueueDispatcher(
-        coord => this.readBlock(...coord),
+        this.source.prefersCombinedRequests() ?
+            // Assume same zoom level
+            coords => this.readBlocks(coords[0][0], coords) :
+            coord => this.readBlock(...coord),
         4,
-        () => this._deduper.pending()
+        () => this._deduper.pending(),
+        this.source.prefersCombinedRequests(),
+        // Filter function that will only allow set requests of same zoom level
+        // as last coord in queue.
+        (coord, i, coords) => coord[0] === coords[coords.length - 1][0]
       );
     }
 
+    readBlocks(zoomLevel, blockCoords) {
+      if (!CATMAID.tools.isFn(this.source.readBlocks)) {
+        return Promise.resolve([]);
+      }
+
+      // Drop all block coords that are available from the cache already
+      const blockKeys = new Map();
+      const uncachedBlocks = blockCoords.map(coord => {
+        const blockKey = [coord[0], coord[1], coord[2], coord[3]].join('/');
+        blockKeys.set(blockKey, coord);
+        return blockKey;
+      }).filter(blockKey => {
+        return !this._cache.has(blockKey);
+      });
+
+      // Drop all block coords that are currently requested
+      let blockPromise = this._deduper.dedup_many(
+        uncachedBlocks,
+        (keysToRequest) => {
+          const coordsToRequest = keysToRequest.map(k => blockKeys.get(k));
+          return this.source.readBlocks(zoomLevel, coordsToRequest);
+        },
+        (blockData, key) => {
+          let [z, c1, c2, c3] = key.split('/');
+          return blockData.find(b => z == zoomLevel && b.gridPosition[0] == c1
+              && b.gridPosition[1] == c2 && b.gridPosition[2] == c3);
+        });
+
+      return blockPromise
+        .then(block_data => {
+          for (let {block, etag, gridPosition} of block_data) {
+            if (!block || !etag || !gridPosition) continue;
+            let blockKey = [zoomLevel, ...gridPosition].join('/');
+            this._stateIDs.set(blockKey, etag);
+            this._cache.set(blockKey, block);
+          }
+          return block_data;
+        });
+    }
+
     readBlock(...zoomBlockCoord) {
       let blockKey = zoomBlockCoord.join('/');
       let block = this._cache.get(blockKey);

diff --git a/django/applications/catmaid/static/js/layers/pixi-image-block-layer.js b/django/applications/catmaid/static/js/layers/pixi-image-block-layer.js
@@ -327,20 +327,29 @@
         // immediately, so that the buffer will be cleared.
         window.clearTimeout(this._swapBuffersTimeout);
         this._swapBuffersTimeout = window.setTimeout(this._swapBuffers.bind(this, true), 3000);
-        Promise.all(toLoad.map(([[i, j], coord]) => this
-            ._readBlock(...coord.slice(0, 4))
-            .then(block => {
-              if (this._context.renderer === null) return; // Layer was destroyed.
-              if (!this._tilesBuffer || !this._tilesBuffer[i] || !this._tilesBuffer[i][j] ||
-                  !CATMAID.tools.arraysEqual(this._tilesBuffer[i][j].coord, coord)) return;
 
-              let slice = this._sliceBlock(block, blockZ);
-
-              this._sliceToTexture(slice, this._tilesBuffer[i][j].texture);
-              this._tilesBuffer[i][j].coord = coord;
-              this._tilesBuffer[i][j].loaded = true;
-            })
-        )).then(this._swapBuffers.bind(this, false, this._swapBuffersTimeout));
+        // Attempt to get entire data for view at the same time and populate
+        // cache with it. This is usually much faster with sharded data. For
+        // non-sharded data, it is faster to skip this step at the moment.
+        let prepare = this.tileSource.prefersCombinedRequests() ?
+            this._populateCache(tileInfo.zoom, toLoad.map(x => x[1])) :
+            Promise.resolve([]);
+        prepare.then(() => {
+              Promise.all(toLoad.map(([[i, j], coord]) => this
+                  ._readBlock(...coord.slice(0, 4))
+                  .then(block => {
+                    if (this._context.renderer === null) return; // Layer was destroyed.
+                    if (!this._tilesBuffer || !this._tilesBuffer[i] || !this._tilesBuffer[i][j] ||
+                        !CATMAID.tools.arraysEqual(this._tilesBuffer[i][j].coord, coord)) return;
+
+                    let slice = this._sliceBlock(block, blockZ);
+
+                    this._sliceToTexture(slice, this._tilesBuffer[i][j].texture);
+                    this._tilesBuffer[i][j].coord = coord;
+                    this._tilesBuffer[i][j].loaded = true;
+                  })
+              )).then(this._swapBuffers.bind(this, false, this._swapBuffersTimeout));
+            });
         loading = true;
       } else if (!loading) {
         this._oldZoom = this._swapZoom;
@@ -361,6 +370,10 @@
       }
     }
 
+    _populateCache(zoomLevel, locations) {
+      return this._blockCache.readBlocks(zoomLevel, locations);
+    }
+
     _readBlock(zoomLevel, x, y, z) {
       let blockCoord = CATMAID.tools.permute([x, y, z], this.recipDimPerm);
 

diff --git a/django/applications/catmaid/static/js/tile-source.js b/django/applications/catmaid/static/js/tile-source.js
@@ -579,6 +579,14 @@
     readBlock(zoomLevel, xi, yi, zi) {
       throw new CATMAID.NotImplementedError();
     }
+
+    /**
+     * Whether or not it can be benefitial to combine multiple block requests
+     * into a single one. This can be the case e.g. for sharded data.
+     */
+    prefersCombinedRequests() {
+      return false;
+    }
   };
 
 
@@ -1026,7 +1034,9 @@
         this.promiseNeuroglancerPrecomputedwasm = (new Function(`return import('${jsPath}')`))()
             .then(module =>
                 // The global ngpre_wasm variable is created in ngpre_wasm.js
-                ngpre_wasm(CATMAID.makeStaticURL('libs/ngpre-wasm/ngpre_wasm_bg.wasm'))
+                ngpre_wasm({
+                  module_or_path: CATMAID.makeStaticURL('libs/ngpre-wasm/ngpre_wasm_bg.wasm'),
+                })
                 .then(() => ngpre_wasm));
       }
 
@@ -1142,6 +1152,41 @@
       });
     }
 
+    readBlocks(zoomLevel, blockCoords) {
+      return this.promiseReady.then(() => {
+          let path = this.datasetPath(zoomLevel);
+          let dataAttrs = this.datasetAttributes;
+
+          if (blockCoords.length === 0) {
+            return [];
+          }
+
+          const gridCoords = blockCoords.map(coord => coord.slice(1,4).map(BigInt)).flat();
+          return this.reader.read_blocks_with_etag(path, dataAttrs, gridCoords);
+        }).then(block_data => {
+          let desBlocks = new Array(block_data.length);
+          for (let i=0; i<block_data.length; ++i) {
+            block = block_data[i];
+            if (block) {
+              let etag = block.get_etag();
+              let size = block.get_size();
+              let gridPosition = block.get_grid_position();
+              let n = 1;
+              let stride = size.map(s => { let rn = n; n *= s; return rn; });
+              desBlocks[i] = {
+                etag,
+                block: new nj.NdArray(nj.ndarray(block.into_data(), size, stride))
+                    .transpose(...this.sliceDims),
+                gridPosition,
+              };
+            } else {
+              desBlocks[i] = {block, etag: undefined, gridPosition: undefined};
+            }
+          }
+          return desBlocks;
+        });
+    }
+
     datasetPath(zoomLevel) {
       return this.datasetPathFormat
           .replace('%SCALE_DATASET%', this.scaleLevelPath(zoomLevel));
@@ -1184,6 +1229,14 @@
           corsTime:   result[1][1]
       })));
     }
+
+    /**
+     * The Neuroglancer WASM implementation supports downloading multiple blocks
+     * in parallel.
+     */
+    prefersCombinedRequests() {
+      return true;
+    }
   };
 
   /**
@@ -1239,6 +1292,56 @@
             });
       });
     }
+
+    readBlocks(zoomLevel, blockCoords) {
+      return this.promiseReady.then(() => {
+          let path = this.datasetPath(zoomLevel);
+          let dataAttrs = this.datasetAttributes;
+
+          if (blockCoords.length === 0) {
+            return [];
+          }
+
+          // Flatten input grid coords (needed for wasm-bindgen)
+          const gridCoords = blockCoords.map(coord => coord.slice(1,4).map(BigInt)).flat();
+          // FIXME: Optionally, run bundle downloads in separate web workers. This
+          // seems currently slower and might not be worthwile afterall, because
+          // all block requests in a single call to WASM will be executed in
+          // parallel already.
+          let bundleRequests = false;
+          let request_block_data = bundleRequests ?
+            this.workers
+              .postMessage([path, dataAttrs.to_json(), gridCoords, true, true])
+              .then(bundles => {
+                return Promise.all(bundles.map(bundle => {
+                  return bundle && bundle.length > 0 ?
+                    this.workers.postMessage([path, dataAttrs.to_json(), bundle.flat(), true]) :
+                    [];
+                }));
+              }) :
+            this.workers.postMessage([path, dataAttrs.to_json(), gridCoords, true]).then(block_data => [block_data]);
+
+            return request_block_data
+              .then(bundled_block_data => {
+                return bundled_block_data.map(block_data => {
+                  return block_data.map(block => {
+                    if (block) {
+                      let n = 1;
+                      let stride = block.size.map(s => { let rn = n; n *= s; return rn; });
+                      return {
+                        etag: block.etag,
+                        block: new nj.NdArray(nj.ndarray(block.data, block.size, stride))
+                            .transpose(...this.sliceDims),
+                        gridPosition: block.gridPosition,
+                      };
+                    } else {
+                      return {block, etag: undefined, gridPosition: undefined};
+                    }
+                  });
+                }).flat();
+              });
+      });
+    }
   };
 
   /**