From 490bf065a49e330dd024f21681d6d33c69ffb33d Mon Sep 17 00:00:00 2001 From: Chris Allan Date: Thu, 14 Mar 2024 11:34:11 +0000 Subject: [PATCH 1/2] Switch from PixelBuffer to Zarr metadata and array cache --- .../omero/ms/core/PixelsService.java | 46 ++++++++++++------- .../omero/ms/core/ZarrPixelBuffer.java | 29 ++++++++---- .../resources/blitz/blitz-ZarrPixelBuffer.xml | 2 +- 3 files changed, 51 insertions(+), 26 deletions(-) diff --git a/src/main/java/com/glencoesoftware/omero/ms/core/PixelsService.java b/src/main/java/com/glencoesoftware/omero/ms/core/PixelsService.java index 31e8406..0e3750a 100644 --- a/src/main/java/com/glencoesoftware/omero/ms/core/PixelsService.java +++ b/src/main/java/com/glencoesoftware/omero/ms/core/PixelsService.java @@ -35,7 +35,9 @@ import com.google.common.base.Splitter; import com.upplication.s3fs.OmeroS3FilesystemProvider; -import com.github.benmanes.caffeine.cache.Cache; +import com.bc.zarr.ZarrArray; +import com.bc.zarr.ZarrGroup; +import com.github.benmanes.caffeine.cache.AsyncLoadingCache; import com.github.benmanes.caffeine.cache.Caffeine; import ome.api.IQuery; @@ -70,34 +72,45 @@ public class PixelsService extends ome.io.nio.PixelsService { /** Max Plane Height */ protected final Integer maxPlaneHeight; - /** OME NGFF LRU cache size */ - private final long omeNgffPixelBufferCacheSize; + /** Zarr metadata and array cache size */ + private final long zarrCacheSize; /** Copy of private IQuery also provided to ome.io.nio.PixelsService */ private final IQuery iQuery; - /** LRU cache of pixels ID vs OME NGFF pixel buffers */ - private Cache omeNgffPixelBufferCache; + /** Root path vs. metadata cache */ + private final + AsyncLoadingCache> zarrMetadataCache; + + /** Array path vs. ZarrArray cache */ + private final AsyncLoadingCache zarrArrayCache; public PixelsService( String path, boolean isReadOnlyRepo, File memoizerDirectory, long memoizerWait, FilePathResolver resolver, BackOff backOff, TileSizes sizes, IQuery iQuery, - long omeNgffPixelBufferCacheSize, + long zarrCacheSize, int maxPlaneWidth, int maxPlaneHeight) { super( path, isReadOnlyRepo, memoizerDirectory, memoizerWait, resolver, backOff, sizes, iQuery ); - this.omeNgffPixelBufferCacheSize = omeNgffPixelBufferCacheSize; - log.info("OME NGFF pixel buffer cache size: {}", - omeNgffPixelBufferCacheSize); + this.zarrCacheSize = zarrCacheSize; + log.info("Zarr metadata and array cache size: {}", zarrCacheSize); this.maxPlaneWidth = maxPlaneWidth; this.maxPlaneHeight = maxPlaneHeight; this.iQuery = iQuery; - omeNgffPixelBufferCache = Caffeine.newBuilder() - .maximumSize(this.omeNgffPixelBufferCacheSize) - .build(); + zarrMetadataCache = Caffeine.newBuilder() + .maximumSize(this.zarrCacheSize) + .buildAsync(key -> { + ZarrGroup rootGroup = ZarrGroup.open(key); + return rootGroup.getAttributes(); + }); + zarrArrayCache = Caffeine.newBuilder() + .maximumSize(this.zarrCacheSize) + .buildAsync(key -> { + return ZarrArray.open(key); + }); } /** @@ -241,7 +254,8 @@ public ZarrPixelBuffer getLabelImagePixelBuffer(Mask mask) "No root for Mask:" + mask.getId()); } return new ZarrPixelBuffer( - pixels, asPath(root), maxPlaneWidth, maxPlaneHeight); + pixels, asPath(root), maxPlaneWidth, maxPlaneHeight, + zarrMetadataCache, zarrArrayCache); } /** @@ -269,7 +283,8 @@ protected ZarrPixelBuffer createOmeNgffPixelBuffer(Pixels pixels) { log.info("OME-NGFF root is: " + uri); try { ZarrPixelBuffer v = new ZarrPixelBuffer( - pixels, root, maxPlaneWidth, maxPlaneHeight); + pixels, root, maxPlaneWidth, maxPlaneHeight, + zarrMetadataCache, zarrArrayCache); log.info("Using OME-NGFF pixel buffer"); return v; } catch (Exception e) { @@ -299,8 +314,7 @@ protected ZarrPixelBuffer createOmeNgffPixelBuffer(Pixels pixels) { */ @Override public PixelBuffer getPixelBuffer(Pixels pixels, boolean write) { - PixelBuffer pixelBuffer = omeNgffPixelBufferCache.get( - pixels.getId(), key -> createOmeNgffPixelBuffer(pixels)); + PixelBuffer pixelBuffer = createOmeNgffPixelBuffer(pixels); if (pixelBuffer != null) { return pixelBuffer; } diff --git a/src/main/java/com/glencoesoftware/omero/ms/core/ZarrPixelBuffer.java b/src/main/java/com/glencoesoftware/omero/ms/core/ZarrPixelBuffer.java index b645491..ed2f9d9 100644 --- a/src/main/java/com/glencoesoftware/omero/ms/core/ZarrPixelBuffer.java +++ b/src/main/java/com/glencoesoftware/omero/ms/core/ZarrPixelBuffer.java @@ -29,13 +29,13 @@ import java.util.List; import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; import java.util.stream.IntStream; import org.slf4j.LoggerFactory; import com.bc.zarr.DataType; import com.bc.zarr.ZarrArray; -import com.bc.zarr.ZarrGroup; import com.github.benmanes.caffeine.cache.AsyncLoadingCache; import com.github.benmanes.caffeine.cache.Caffeine; @@ -70,9 +70,6 @@ public class ZarrPixelBuffer implements PixelBuffer { /** Max Plane Height */ private final Integer maxPlaneHeight; - /** Root group of the OME-NGFF multiscale we are operating on */ - private final ZarrGroup rootGroup; - /** Zarr attributes present on the root group */ private final Map rootGroupAttributes; @@ -85,6 +82,13 @@ public class ZarrPixelBuffer implements PixelBuffer { /** Whether or not the Zarr is on S3 or similar */ private final boolean isRemote; + /** Root path vs. metadata cache */ + private final + AsyncLoadingCache> zarrMetadataCache; + + /** Array path vs. ZarrArray cache */ + private final AsyncLoadingCache zarrArrayCache; + /** * Default constructor * @param pixels Pixels metadata for the pixel buffer @@ -94,14 +98,21 @@ public class ZarrPixelBuffer implements PixelBuffer { * @throws IOException */ public ZarrPixelBuffer(Pixels pixels, Path root, Integer maxPlaneWidth, - Integer maxPlaneHeight) + Integer maxPlaneHeight, + AsyncLoadingCache> zarrMetadataCache, + AsyncLoadingCache zarrArrayCache) throws IOException { log.info("Creating ZarrPixelBuffer"); this.pixels = pixels; this.root = root; + this.zarrMetadataCache = zarrMetadataCache; + this.zarrArrayCache = zarrArrayCache; this.isRemote = root.toString().startsWith("s3://")? true : false; - rootGroup = ZarrGroup.open(this.root); - rootGroupAttributes = rootGroup.getAttributes(); + try { + rootGroupAttributes = this.zarrMetadataCache.get(this.root).get(); + } catch (ExecutionException|InterruptedException e) { + throw new IOException(e); + } if (!rootGroupAttributes.containsKey("multiscales")) { throw new IllegalArgumentException("Missing multiscales metadata!"); } @@ -766,8 +777,8 @@ public void setResolutionLevel(int resolutionLevel) { "This Zarr file has no pixel data"); } try { - array = ZarrArray.open( - root.resolve(Integer.toString(this.resolutionLevel))); + array = zarrArrayCache.get( + root.resolve(Integer.toString(this.resolutionLevel))).get(); } catch (Exception e) { // FIXME: Throw the right exception throw new RuntimeException(e); diff --git a/src/main/resources/blitz/blitz-ZarrPixelBuffer.xml b/src/main/resources/blitz/blitz-ZarrPixelBuffer.xml index 94013e3..ae818c0 100644 --- a/src/main/resources/blitz/blitz-ZarrPixelBuffer.xml +++ b/src/main/resources/blitz/blitz-ZarrPixelBuffer.xml @@ -15,7 +15,7 @@ - + From 78bdb2d47342d969fb9f9345da38c21a6ac57158 Mon Sep 17 00:00:00 2001 From: Chris Allan Date: Thu, 14 Mar 2024 11:57:07 +0000 Subject: [PATCH 2/2] Fix unit tests --- .../omero/ms/core/PixelsService.java | 33 ++++++++++--- .../omero/ms/core/ZarrPixelBufferTest.java | 47 ++++++++++++------- 2 files changed, 57 insertions(+), 23 deletions(-) diff --git a/src/main/java/com/glencoesoftware/omero/ms/core/PixelsService.java b/src/main/java/com/glencoesoftware/omero/ms/core/PixelsService.java index 0e3750a..aa8917e 100644 --- a/src/main/java/com/glencoesoftware/omero/ms/core/PixelsService.java +++ b/src/main/java/com/glencoesoftware/omero/ms/core/PixelsService.java @@ -102,15 +102,34 @@ public PixelsService( this.iQuery = iQuery; zarrMetadataCache = Caffeine.newBuilder() .maximumSize(this.zarrCacheSize) - .buildAsync(key -> { - ZarrGroup rootGroup = ZarrGroup.open(key); - return rootGroup.getAttributes(); - }); + .buildAsync(PixelsService::getZarrMetadata); zarrArrayCache = Caffeine.newBuilder() .maximumSize(this.zarrCacheSize) - .buildAsync(key -> { - return ZarrArray.open(key); - }); + .buildAsync(PixelsService::getZarrArray); + } + + /** + * Retrieves Zarr metadata from a given path. + * @param root path to get Zarr metadata from + * @return See above. + * @throws IOException + */ + public static Map getZarrMetadata(Path path) + throws IOException { + // FIXME: Really should be ZarrUtils.readAttributes() to allow for + // attribute retrieval from either a ZarrArray or ZarrGroup but ZarrPath + // is package private at the moment. + return ZarrGroup.open(path).getAttributes(); + } + + /** + * Opens a Zarr array at a given path. + * @param root path to open a Zarr array from + * @return See above. + * @throws IOException + */ + public static ZarrArray getZarrArray(Path path) throws IOException { + return ZarrArray.open(path); } /** diff --git a/src/test/java/com/glencoesoftware/omero/ms/core/ZarrPixelBufferTest.java b/src/test/java/com/glencoesoftware/omero/ms/core/ZarrPixelBufferTest.java index 25679be..b4dc861 100644 --- a/src/test/java/com/glencoesoftware/omero/ms/core/ZarrPixelBufferTest.java +++ b/src/test/java/com/glencoesoftware/omero/ms/core/ZarrPixelBufferTest.java @@ -34,6 +34,7 @@ import org.junit.Test; import com.bc.zarr.ZarrArray; +import com.github.benmanes.caffeine.cache.Caffeine; import loci.formats.FormatTools; import loci.formats.in.FakeReader; @@ -47,6 +48,20 @@ public class ZarrPixelBufferTest extends AbstractZarrPixelBufferTest { + public ZarrPixelBuffer createPixelBuffer( + Pixels pixels, Path path, + Integer maxPlaneWidth, Integer maxPlaneHeight) throws IOException { + return new ZarrPixelBuffer( + pixels, path, maxPlaneWidth, maxPlaneHeight, + Caffeine.newBuilder() + .maximumSize(0) + .buildAsync(PixelsService::getZarrMetadata), + Caffeine.newBuilder() + .maximumSize(0) + .buildAsync(PixelsService::getZarrArray) + ); + } + @Test public void testGetChunks() throws IOException { int sizeT = 1; @@ -60,7 +75,7 @@ public void testGetChunks() throws IOException { Path output = writeTestZarr( sizeT, sizeC, sizeZ, sizeY, sizeX, "uint16", resolutions); try (ZarrPixelBuffer zpbuf = - new ZarrPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { + createPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { int[][] chunks = zpbuf.getChunks(); int[][] expectedChunks = new int[][] { new int[] {1, 1, 1, 512, 1024}, @@ -88,7 +103,7 @@ public void testGetDatasets() throws IOException { sizeT, sizeC, sizeZ, sizeY, sizeX, "uint16", resolutions); try (ZarrPixelBuffer zpbuf = - new ZarrPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { + createPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { List> datasets = zpbuf.getDatasets(); List> expectedDatasets = getDatasets(3); for (int i = 0; i < datasets.size(); i++) { @@ -110,7 +125,7 @@ public void testGetResolutionDescriptions() throws IOException { Path output = writeTestZarr( sizeT, sizeC, sizeZ, sizeY, sizeX, "uint16", resolutions); try (ZarrPixelBuffer zpbuf = - new ZarrPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { + createPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { List> expected = new ArrayList>(); expected.add(Arrays.asList(new Integer[] {2048, 512})); expected.add(Arrays.asList(new Integer[] {1024, 256})); @@ -158,7 +173,7 @@ public void testGetTile() throws IOException, InvalidRangeException { } test.write(data, new int[] {2,3,4,5,6}, new int[] {0,0,0,0,0}); try (ZarrPixelBuffer zpbuf = - new ZarrPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { + createPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { PixelData pixelData = zpbuf.getTile(0, 0, 0, 0, 0, 2, 2); ByteBuffer bb = pixelData.getData(); bb.order(ByteOrder.BIG_ENDIAN); @@ -242,7 +257,7 @@ public void testGetTimepointStackPlaneRowCol() Path output = writeTestZarr( sizeT, sizeC, sizeZ, sizeY, sizeX, "int32", resolutions); try (ZarrPixelBuffer zpbuf = - new ZarrPixelBuffer(pixels, output.resolve("0"), 2048, 2048)) { + createPixelBuffer(pixels, output.resolve("0"), 2048, 2048)) { for (int t = 0; t < sizeT; t++) { // Assert timepoint byte[] timepoint = zpbuf.getTimepoint(t).getData().array(); @@ -309,7 +324,7 @@ public void testGetTileLargerThanImage() } test.write(data, new int[] {2,3,4,5,6}, new int[] {0,0,0,0,0}); try (ZarrPixelBuffer zpbuf = - new ZarrPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { + createPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { zpbuf.setResolutionLevel(0); PixelData pixelData = zpbuf.getTile(0, 0, 0, 0, 0, 10, 10); ByteBuffer bb = pixelData.getData(); @@ -337,7 +352,7 @@ public void testTileExceedsMax() throws IOException, InvalidRangeException { resolutions); try (ZarrPixelBuffer zpbuf = - new ZarrPixelBuffer(pixels, output.resolve("0"), 32, 32)) { + createPixelBuffer(pixels, output.resolve("0"), 32, 32)) { PixelData pixelData = zpbuf.getTile(0, 0, 0, 0, 0, 32, 33); Assert.assertNull(pixelData); } @@ -356,7 +371,7 @@ public void testCheckBoundsValidZeros() throws IOException { Path output = writeTestZarr( sizeT, sizeC, sizeZ, sizeY, sizeX, "uint16", resolutions); try (ZarrPixelBuffer zpbuf = - new ZarrPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { + createPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { zpbuf.checkBounds(0, 0, 0, 0, 0); } } @@ -374,7 +389,7 @@ public void testCheckBoundsValidEnd() throws IOException { Path output = writeTestZarr( sizeT, sizeC, sizeZ, sizeY, sizeX, "uint16", resolutions); try (ZarrPixelBuffer zpbuf = - new ZarrPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { + createPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { zpbuf.checkBounds(2047, 511, 2, 1, 0); } } @@ -392,7 +407,7 @@ public void testCheckBoundsOutOfRange() throws IOException { Path output = writeTestZarr( sizeT, sizeC, sizeZ, sizeY, sizeX, "uint16", resolutions); try (ZarrPixelBuffer zpbuf = - new ZarrPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { + createPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { zpbuf.checkBounds(2048, 511, 2, 1, 0); } } @@ -410,7 +425,7 @@ public void testCheckBounds() throws IOException { Path output = writeTestZarr( sizeT, sizeC, sizeZ, sizeY, sizeX, "uint16", resolutions); try (ZarrPixelBuffer zpbuf = - new ZarrPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { + createPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { zpbuf.checkBounds(-1, 0, 0, 0, 0); } } @@ -428,7 +443,7 @@ public void testGetTileSize() throws IOException { Path output = writeTestZarr( sizeT, sizeC, sizeZ, sizeY, sizeX, "uint16", resolutions); try (ZarrPixelBuffer zpbuf = - new ZarrPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { + createPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { Dimension tileSize = zpbuf.getTileSize(); Assert.assertEquals(1024, tileSize.getWidth(), 0.1); Assert.assertEquals(1024, tileSize.getHeight(), 0.1); @@ -449,7 +464,7 @@ public void testUint16() throws IOException { Path output = writeTestZarr( sizeT, sizeC, sizeZ, sizeY, sizeX, "uint16", resolutions); try (ZarrPixelBuffer zpbuf = - new ZarrPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { + createPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { Assert.assertEquals(FormatTools.UINT16, zpbuf.getPixelsType()); Assert.assertEquals(false, zpbuf.isSigned()); Assert.assertEquals(false, zpbuf.isFloat()); @@ -471,7 +486,7 @@ public void testFloat() throws IOException { Path output = writeTestZarr( sizeT, sizeC, sizeZ, sizeY, sizeX, "float", resolutions); try (ZarrPixelBuffer zpbuf = - new ZarrPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { + createPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { Assert.assertEquals(FormatTools.FLOAT, zpbuf.getPixelsType()); Assert.assertEquals(true, zpbuf.isSigned()); Assert.assertEquals(true, zpbuf.isFloat()); @@ -493,7 +508,7 @@ public void testSizes() throws IOException { Path output = writeTestZarr( sizeT, sizeC, sizeZ, sizeY, sizeX, "uint16", resolutions); try (ZarrPixelBuffer zpbuf = - new ZarrPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { + createPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { // Plane size Assert.assertEquals( sizeX * sizeY * bytesPerPixel, @@ -534,7 +549,7 @@ public void testSetResolutionLevelOutOfBounds() throws IOException { Path output = writeTestZarr( sizeT, sizeC, sizeZ, sizeY, sizeX, "uint16", resolutions); try (ZarrPixelBuffer zpbuf = - new ZarrPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { + createPixelBuffer(pixels, output.resolve("0"), 1024, 1024)) { zpbuf.setResolutionLevel(3); } }