Skip to content

Switch from PixelBuffer to Zarr metadata and array cache #3

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 14, 2024
Merged

Switch from PixelBuffer to Zarr metadata and array cache #3

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
490bf06
Switch from PixelBuffer to Zarr metadata and array cache
chris-allan Mar 14, 2024
78bdb2d
Fix unit tests
chris-allan Mar 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 49 additions & 16 deletions src/main/java/com/glencoesoftware/omero/ms/core/PixelsService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -70,34 +72,64 @@ 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<Long, ZarrPixelBuffer> omeNgffPixelBufferCache;
/** Root path vs. metadata cache */
private final
AsyncLoadingCache<Path, Map<String, Object>> zarrMetadataCache;

/** Array path vs. ZarrArray cache */
private final AsyncLoadingCache<Path, ZarrArray> 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(PixelsService::getZarrMetadata);
zarrArrayCache = Caffeine.newBuilder()
.maximumSize(this.zarrCacheSize)
.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<String, Object> 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.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the same spirit as the work we did with @melissalinkert on bioformats2raw/raw2ometiff, I was considering switching to the more active zarr-developers/jzarr fork on this repository as well
This might be an opportunity to discuss any API improvements we should consider as a preamble of this update.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Opened as zarr-developers/jzarr#19. I think it makes sense to switch to that fork for testing as soon as possible.

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);
}

/**
Expand Down Expand Up @@ -241,7 +273,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);
}

/**
Expand Down Expand Up @@ -269,7 +302,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) {
Expand Down Expand Up @@ -299,8 +333,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;
}
Expand Down
29 changes: 20 additions & 9 deletions src/main/java/com/glencoesoftware/omero/ms/core/ZarrPixelBuffer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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<String, Object> rootGroupAttributes;

Expand All @@ -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<Path, Map<String, Object>> zarrMetadataCache;

/** Array path vs. ZarrArray cache */
private final AsyncLoadingCache<Path, ZarrArray> zarrArrayCache;

/**
* Default constructor
* @param pixels Pixels metadata for the pixel buffer
Expand All @@ -94,14 +98,21 @@ public class ZarrPixelBuffer implements PixelBuffer {
* @throws IOException
*/
public ZarrPixelBuffer(Pixels pixels, Path root, Integer maxPlaneWidth,
Integer maxPlaneHeight)
Integer maxPlaneHeight,
AsyncLoadingCache<Path, Map<String, Object>> zarrMetadataCache,
AsyncLoadingCache<Path, ZarrArray> 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!");
}
Expand Down Expand Up @@ -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);
Expand Down
2 changes: 1 addition & 1 deletion src/main/resources/blitz/blitz-ZarrPixelBuffer.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
<constructor-arg ref="backOff"/>
<constructor-arg ref="tileSizes"/>
<constructor-arg ref="internal-ome.api.IQuery"/>
<constructor-arg type="long" value="${omero.pixeldata.ome_ngff_pixel_buffer_cache_size:100}"/>
<constructor-arg type="long" value="${omero.pixeldata.zarr_cache_size:500}"/>
<constructor-arg type="int" value="${omero.pixeldata.max_plane_height}"/>
<constructor-arg type="int" value="${omero.pixeldata.max_plane_width}"/>
<property name="metrics" ref="metrics"/>
Expand Down
47 changes: 31 additions & 16 deletions src/test/java/com/glencoesoftware/omero/ms/core/ZarrPixelBufferTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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)
);
}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍 I'll port similar changes to the omero-ms-image-region unit tests instantiating ZarrPixelBuffer .

This raises the question of whether AbstractZarrPixelBufferTest (which originates from the omero-ms-image-region tests) is still relevant and could be consumed by downstream components. In that case, there's probably a case for moving this utility method under this abstract class.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good. I'm not sure the utility now to be honest.


@Test
public void testGetChunks() throws IOException {
int sizeT = 1;
Expand All @@ -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},
Expand Down Expand Up @@ -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<Map<String,String>> datasets = zpbuf.getDatasets();
List<Map<String,String>> expectedDatasets = getDatasets(3);
for (int i = 0; i < datasets.size(); i++) {
Expand All @@ -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<List<Integer>> expected = new ArrayList<List<Integer>>();
expected.add(Arrays.asList(new Integer[] {2048, 512}));
expected.add(Arrays.asList(new Integer[] {1024, 256}));
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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);
}
Expand All @@ -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);
}
}
Expand All @@ -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);
}
}
Expand All @@ -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);
}
}
Expand All @@ -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);
}
}
Expand All @@ -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);
Expand All @@ -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());
Expand All @@ -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());
Expand All @@ -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,
Expand Down Expand Up @@ -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);
}
}
Expand Down