From 15c3276f983a7b54fc581a9fd1b4c8f697668dbf Mon Sep 17 00:00:00 2001 From: Jannik Lindemann Date: Sat, 1 Aug 2026 12:33:00 +0200 Subject: [PATCH] [OOC] Add Materialized Planning Support --- .../instructions/ooc/CachingStream.java | 33 ++++- .../instructions/ooc/OOCStreamable.java | 10 ++ .../instructions/ooc/PlaybackStream.java | 7 +- .../runtime/ooc/planning/OOCPlanner.java | 29 ++++ .../runtime/ooc/planning/OOCStoreLayout.java | 36 +++++ .../ooc/primitives/MappingOOCPrimitive.java | 6 +- .../primitives/MaterializeOOCPrimitive.java | 117 +++++++++++++++ .../runtime/ooc/primitives/OOCPrimitive.java | 140 +++++++++++++++--- .../ooc/primitives/TransposeOOCPrimitive.java | 6 +- .../store/IndexedMaterializedStoreReader.java | 9 +- .../runtime/ooc/store/MaterializedStore.java | 53 ++++++- .../ooc/store/OOCStreamMaterializer.java | 13 +- .../store/OrderedMaterializedStoreReader.java | 9 +- .../sysds/runtime/ooc/store/StateTable.java | 33 +++-- .../sysds/runtime/ooc/store/StoreLease.java | 87 ++++++++--- .../ooc/stream/SourceOOCStreamable.java | 11 +- .../runtime/ooc/util/StateTableUtils.java | 11 +- .../ooc/OOCInstructionUtilsTest.java | 2 +- .../test/component/ooc/OOCPrimitiveTest.java | 82 +++++++++- .../component/ooc/StateTableUtilsTest.java | 10 +- 20 files changed, 613 insertions(+), 91 deletions(-) create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/planning/OOCStoreLayout.java create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/primitives/MaterializeOOCPrimitive.java diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/CachingStream.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/CachingStream.java index dab3eca3a0f..3b80175b73a 100644 --- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/CachingStream.java +++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/CachingStream.java @@ -79,6 +79,7 @@ public class CachingStream implements OOCStreamable { private boolean _deletable = false; private int _maxConsumptionCount = 0; + private int _lazyHandleReservations = 0; private String _watchdogId = null; public CachingStream(OOCStream source) { @@ -356,7 +357,7 @@ public synchronized void scheduleDeletion() { if (_deletable) return; // Deletion already scheduled - if (_cacheInProgress && _maxConsumptionCount == 0) + if(_cacheInProgress && _maxConsumptionCount == 0 && _lazyHandleReservations == 0) System.out.println("[WARN] Scheduling deletion for caching stream with no listeners: " + this); _deletable = true; @@ -370,6 +371,8 @@ public String toString() { } private synchronized void tryDeleteBlock(int i) { + if(_lazyHandleReservations > 0) + return; int cnt = _consumptionCounts.getInt(i); if (cnt > _maxConsumptionCount) throw new DMLRuntimeException("Cannot have more than " + _maxConsumptionCount + " consumptions."); @@ -582,6 +585,11 @@ public OOCStream getReadStream() { return new PlaybackStream(this); } + @Override + public OOCStream getReservedReadStream() { + return new PlaybackStream(this, consumeLazyHandleReservation()); + } + @Override public OOCStream getWriteStream() { return _source.getWriteStream(); @@ -721,6 +729,29 @@ public synchronized void incrSubscriberCount(int count) { _maxConsumptionCount += count; } + @Override + public synchronized void reserveLazyHandle() { + _lazyHandleReservations++; + } + + @Override + public synchronized void discardHandle() { + if(_lazyHandleReservations <= 0) + return; + _lazyHandleReservations--; + if(_deletable) + for(int i = 0; i < _consumptionCounts.size(); i++) + tryDeleteBlock(i); + } + + private synchronized boolean consumeLazyHandleReservation() { + if(_lazyHandleReservations <= 0) + return false; + _lazyHandleReservations--; + _maxConsumptionCount++; + return true; + } + /** * Artificially increase the processing count of a block. */ diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/OOCStreamable.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/OOCStreamable.java index 0f087a7f55b..10a4dc88174 100644 --- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/OOCStreamable.java +++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/OOCStreamable.java @@ -47,4 +47,14 @@ default OOCPrimitive getPrimitive() { default void assignPrimitive(OOCPrimitive primitive) { throw new UnsupportedOperationException("Stream does not support primitive assignment"); } + + default OOCStream getReservedReadStream() { + return getReadStream(); + } + + default void reserveLazyHandle() { + } + + default void discardHandle() { + } } diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/PlaybackStream.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/PlaybackStream.java index 167d7ba7aed..9b62405b460 100644 --- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/PlaybackStream.java +++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/PlaybackStream.java @@ -37,10 +37,15 @@ public class PlaybackStream implements OOCStream { private QueueCallback _lastDequeue; public PlaybackStream(CachingStream streamCache) { + this(streamCache, false); + } + + public PlaybackStream(CachingStream streamCache, boolean reserved) { this._streamCache = streamCache; this._streamIdx = new AtomicInteger(0); this._subscriberSet = new AtomicBoolean(false); - streamCache.incrSubscriberCount(1); + if(!reserved) + streamCache.incrSubscriberCount(1); } @Override diff --git a/src/main/java/org/apache/sysds/runtime/ooc/planning/OOCPlanner.java b/src/main/java/org/apache/sysds/runtime/ooc/planning/OOCPlanner.java index 3db13c2691a..38a5c609525 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/planning/OOCPlanner.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/planning/OOCPlanner.java @@ -25,10 +25,14 @@ import java.util.List; import java.util.Set; +import org.apache.sysds.runtime.instructions.ooc.OOCStreamable; +import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue; +import org.apache.sysds.runtime.ooc.primitives.MaterializeOOCPrimitive; import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive; public final class OOCPlanner { public static void compile(OOCPrimitive root) { + injectMaterializations(root, Collections.newSetFromMap(new IdentityHashMap<>()), new IdentityHashMap<>()); List primitives = new ArrayList<>(); collect(root, Collections.newSetFromMap(new IdentityHashMap<>()), primitives); if(primitives.isEmpty()) @@ -44,6 +48,31 @@ public static void compile(OOCPrimitive root) { primitive.tryStartExecution(); } + @SuppressWarnings("unchecked") + private static void injectMaterializations(OOCPrimitive primitive, Set visited, + IdentityHashMap, MaterializeOOCPrimitive> boundaries) { + if(primitive.hasStartedExecution() || !visited.add(primitive)) + return; + for(OOCPrimitive.OOCMaterializedInputRequest request : primitive.requiredMaterializedInputs()) { + OOCStreamable input = (OOCStreamable) primitive + .getInput(request.inputIndex()); + MaterializeOOCPrimitive boundary = boundaries.compute(input, (k, v) -> { + if(v == null) { + MaterializeOOCPrimitive p = new MaterializeOOCPrimitive(input, request.layout(), + primitive.getContext()); + primitive.transferInputHandle(request.inputIndex()); + return p; + } + primitive.discardInputHandle(request.inputIndex()); + return v; + }); + boundary.registerRequest(request.expectedReaders()); + primitive.installMaterializedInput(request.inputIndex(), boundary); + } + for(OOCPrimitive child : primitive.getChildren()) + injectMaterializations(child, visited, boundaries); + } + private static void collect(OOCPrimitive primitive, Set visited, List result) { if(primitive.hasStartedExecution() || !visited.add(primitive)) return; diff --git a/src/main/java/org/apache/sysds/runtime/ooc/planning/OOCStoreLayout.java b/src/main/java/org/apache/sysds/runtime/ooc/planning/OOCStoreLayout.java new file mode 100644 index 00000000000..bafd5ec79f7 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/planning/OOCStoreLayout.java @@ -0,0 +1,36 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.ooc.planning; + +import org.apache.sysds.runtime.matrix.data.MatrixIndexes; +import org.apache.sysds.runtime.meta.DataCharacteristics; + +public enum OOCStoreLayout { + ROW_MAJOR; + + public int linearize(MatrixIndexes indexes, DataCharacteristics characteristics) { + if(characteristics == null || !characteristics.dimsKnown() || characteristics.getBlocksize() <= 0) + throw new IllegalArgumentException("Materialized store layout requires known dimensions and block size."); + long columns = characteristics.getNumColBlocks(); + long index = Math.addExact(Math.multiplyExact(indexes.getRowIndex() - 1, columns), + indexes.getColumnIndex() - 1); + return Math.toIntExact(index); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/MappingOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/MappingOOCPrimitive.java index 345925cc83c..c4a96bc01dc 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/MappingOOCPrimitive.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/MappingOOCPrimitive.java @@ -50,8 +50,8 @@ private MappingOOCPrimitive(OOCStream input, OOCStreamable c.requestPattern(accessPattern)); } @Override diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/MaterializeOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/MaterializeOOCPrimitive.java new file mode 100644 index 00000000000..eaa9e8232c7 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/MaterializeOOCPrimitive.java @@ -0,0 +1,117 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.ooc.primitives; + +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; + +import org.apache.sysds.runtime.DMLRuntimeException; +import org.apache.sysds.runtime.instructions.ooc.CachingStream; +import org.apache.sysds.runtime.instructions.ooc.OOCStream; +import org.apache.sysds.runtime.instructions.ooc.OOCStreamable; +import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue; +import org.apache.sysds.runtime.ooc.cache.OOCCacheManager; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; +import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern; +import org.apache.sysds.runtime.ooc.planning.OOCStoreLayout; +import org.apache.sysds.runtime.ooc.store.MaterializedStore; +import org.apache.sysds.runtime.ooc.store.OOCStreamMaterializer; +import org.apache.sysds.runtime.ooc.stream.StreamContext; + +public final class MaterializeOOCPrimitive extends OOCPrimitive { + private final OOCStreamable _source; + private final OOCStoreLayout _layout; + private final OOCFuture> _store; + private final AtomicBoolean _finished; + private int _expectedReaders; + private int _consumers; + + public MaterializeOOCPrimitive(OOCStreamable source, OOCStoreLayout layout, + StreamContext context) { + super(context, source.getPrimitive() == null ? List.of() : List.of(source.getPrimitive())); + _source = source; + _layout = layout; + _store = new OOCFuture<>(); + _finished = new AtomicBoolean(); + } + + public synchronized void registerRequest(int expectedReaders) { + if(expectedReaders <= 0) + throw new IllegalArgumentException("Materialization request requires at least one reader."); + if(hasStartedExecution()) + throw new IllegalStateException("Cannot register a consumer after materialization started."); + _expectedReaders = Math.addExact(_expectedReaders, expectedReaders); + _consumers = Math.addExact(_consumers, 1); + } + + public OOCFuture> store() { + return _store; + } + + @Override + protected void inferPatternsInternal() { + _pattern = OOCAccessPattern.ROW_MAJOR; + for(OOCPrimitive child : getChildren()) + child.requestPattern(OOCAccessPattern.ROW_MAJOR); + inferParentPatterns(); + } + + @Override + protected void requestPatternInternal(OOCAccessPattern accessPattern) { + _pattern = OOCAccessPattern.ROW_MAJOR; + for(OOCPrimitive child : getChildren()) + child.requestPattern(OOCAccessPattern.ROW_MAJOR); + } + + @Override + protected void startExecution() { + try { + OOCStream source = _source.getReservedReadStream(); + MaterializedStore store = new MaterializedStore<>(OOCCacheManager.getGlobalCache(), + CachingStream._streamSeq.getNextID(), _expectedReaders, _consumers); + OOCStreamMaterializer materializer = new OOCStreamMaterializer(store, + indexes -> _layout.linearize(indexes, _source.getDataCharacteristics()), _allowance); + materializer.completion().whenComplete((ignored, error) -> { + if(error != null) + fail(error); + finish(); + }); + if(getContext() != null) + getContext().addInStream(source); + _store.complete(store); + materializer.attach(source); + } + catch(Throwable failure) { + _store.completeExceptionally(failure); + fail(failure); + finish(); + } + } + + private void fail(Throwable error) { + if(getContext() != null) + getContext().failAll(DMLRuntimeException.of(error)); + } + + private void finish() { + if(_finished.compareAndSet(false, true)) + onComplete(); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java index aacb35acdad..0bc20d54bd4 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java @@ -20,36 +20,55 @@ package org.apache.sysds.runtime.ooc.primitives; import java.util.ArrayList; +import java.util.HashSet; import java.util.List; +import java.util.Objects; +import java.util.Set; import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.sysds.runtime.instructions.ooc.OOCStream; +import org.apache.sysds.runtime.instructions.ooc.OOCStreamable; +import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; import org.apache.sysds.runtime.ooc.memory.GlobalMemoryBroker; import org.apache.sysds.runtime.ooc.memory.MemoryAllowance; import org.apache.sysds.runtime.ooc.memory.SyncMemoryAllowance; import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern; import org.apache.sysds.runtime.ooc.planning.OOCPlanner; +import org.apache.sysds.runtime.ooc.planning.OOCStoreLayout; +import org.apache.sysds.runtime.ooc.store.MaterializedStore; import org.apache.sysds.runtime.ooc.stream.StreamContext; public abstract class OOCPrimitive { private final StreamContext _context; - private final List _children; - private final List _parents; + private final Set _children; + private final Set _parents; + private final List _inputs; private final AtomicBoolean _started; private final AtomicBoolean _executionStarted; protected OOCAccessPattern _pattern; protected MemoryAllowance _allowance; protected OOCPrimitive(StreamContext context, List children) { + this(context); + children.stream().filter(Objects::nonNull).forEach(child -> { + _children.add(child); + child._parents.add(this); + }); + } + + protected OOCPrimitive(StreamContext context, OOCStreamable... inputs) { + this(context); + for(OOCStreamable input : inputs) + _inputs.add(new InputSlot(input)); + rebuildInputChildren(); + } + + private OOCPrimitive(StreamContext context) { _context = context; - _parents = new ArrayList<>(); - List uniqueChildren = new ArrayList<>(children.size()); - for(OOCPrimitive child : children) { - if(containsIdentity(uniqueChildren, child)) - continue; - uniqueChildren.add(child); - child.addParent(this); - } - _children = List.copyOf(uniqueChildren); + _children = new HashSet<>(); + _parents = new HashSet<>(); + _inputs = new ArrayList<>(); _started = new AtomicBoolean(); _executionStarted = new AtomicBoolean(); _pattern = OOCAccessPattern.UNSET; @@ -59,17 +78,12 @@ public final StreamContext getContext() { return _context; } - public final List getChildren() { + public final Set getChildren() { return _children; } - public final List getParents() { - return List.copyOf(_parents); - } - - private void addParent(OOCPrimitive parent) { - if(!containsIdentity(_parents, parent)) - _parents.add(parent); + public final Set getParents() { + return _parents; } protected final void inferParentPatterns() { @@ -86,6 +100,59 @@ public final boolean hasStartedExecution() { return _executionStarted.get(); } + public List requiredMaterializedInputs() { + return List.of(); + } + + public final OOCStreamable getInput(int index) { + return _inputs.get(index)._source; + } + + public final OOCPrimitive getChildPrimitiveAt(int index) { + return _inputs.get(index)._primitive; + } + + public final void installMaterializedInput(int index, MaterializeOOCPrimitive boundary) { + if(hasStartedExecution()) + throw new IllegalStateException("Cannot replace an input after primitive execution started."); + InputSlot input = _inputs.get(index); + input._primitive = boundary; + rebuildInputChildren(); + } + + public final synchronized void transferInputHandle(int index) { + InputSlot input = _inputs.get(index); + if(!input._handleReserved) + throw new IllegalStateException("Input " + index + " no longer owns a lazy handle."); + input._handleReserved = false; + } + + public final void discardInputHandle(int index) { + OOCStreamable source; + synchronized(this) { + InputSlot input = _inputs.get(index); + if(!input._handleReserved) + return; + input._handleReserved = false; + source = input._source; + } + source.discardHandle(); + } + + @SuppressWarnings("unchecked") + protected final OOCStream getInputReadStream(int index) { + transferInputHandle(index); + return (OOCStream) _inputs.get(index)._source.getReservedReadStream(); + } + + protected final OOCFuture> getMaterializedInput(int index) { + OOCFuture> materialized = ((MaterializeOOCPrimitive) _inputs + .get(index)._primitive).store(); + if(materialized == null) + throw new IllegalStateException("Input " + index + " was not materialized by the planner."); + return materialized; + } + public final void start() { if(_started.compareAndSet(false, true)) OOCPlanner.compile(this); @@ -99,6 +166,8 @@ public final void tryStartExecution() { } public final void onComplete() { + for(int i = 0; i < _inputs.size(); i++) + discardInputHandle(i); _allowance.shutdown(); } @@ -112,11 +181,18 @@ public final void requestPattern(OOCAccessPattern accessPattern) { requestPatternInternal(accessPattern); } - private static boolean containsIdentity(List primitives, OOCPrimitive primitive) { - for(OOCPrimitive current : primitives) - if(current == primitive) - return true; - return false; + private void rebuildInputChildren() { + List next = new ArrayList<>(); + for(InputSlot input : _inputs) + if(input._primitive != null) + next.add(input._primitive); + for(OOCPrimitive child : _children) + if(!next.contains(child)) + child._parents.remove(this); + for(OOCPrimitive child : next) + child._parents.add(this); + _children.clear(); + _children.addAll(next); } protected abstract void startExecution(); @@ -124,4 +200,20 @@ private static boolean containsIdentity(List primitives, OOCPrimit protected abstract void inferPatternsInternal(); protected abstract void requestPatternInternal(OOCAccessPattern accessPattern); + + private static final class InputSlot { + private final OOCStreamable _source; + private OOCPrimitive _primitive; + private boolean _handleReserved; + + private InputSlot(OOCStreamable source) { + _source = source; + _primitive = source.getPrimitive(); + _handleReserved = true; + source.reserveLazyHandle(); + } + } + + public record OOCMaterializedInputRequest(int inputIndex, OOCStoreLayout layout, int expectedReaders) { + } } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/TransposeOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/TransposeOOCPrimitive.java index 1cbe504546f..18b127bfd09 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/TransposeOOCPrimitive.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/TransposeOOCPrimitive.java @@ -51,7 +51,7 @@ private TransposeOOCPrimitive(OOCStream input, OOCStreamable @Override protected void inferPatternsInternal() { - _pattern = (getChildren().isEmpty() ? OOCAccessPattern.ANY : getChildren().get(0).getAccessPattern()) + _pattern = (getChildren().isEmpty() ? OOCAccessPattern.ANY : getChildren().iterator().next().getAccessPattern()) .transposed(); inferParentPatterns(); } @@ -59,8 +59,8 @@ protected void inferPatternsInternal() { @Override protected void requestPatternInternal(OOCAccessPattern accessPattern) { _pattern = accessPattern; - if(!getChildren().isEmpty()) - getChildren().get(0).requestPattern(accessPattern.transposed()); + for(OOCPrimitive child : getChildren()) + child.requestPattern(accessPattern.transposed()); } @Override diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/IndexedMaterializedStoreReader.java b/src/main/java/org/apache/sysds/runtime/ooc/store/IndexedMaterializedStoreReader.java index 3106705a34a..1c3a4c9005f 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/store/IndexedMaterializedStoreReader.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/IndexedMaterializedStoreReader.java @@ -81,7 +81,7 @@ else if(entry == null) { result.complete(null); } else - result.complete(new StoreLease<>(entry, () -> release(index, entry, requestAllowance))); + result.complete(StoreLease.createAsync(entry, () -> release(index, entry, requestAllowance))); }); return result; } @@ -94,13 +94,14 @@ public StoreLease requestIfLive(int index, MemoryAllowance requestAllowance) _liveness.unreserve(index); return null; } - return new StoreLease<>(entry, () -> release(index, entry, requestAllowance)); + return StoreLease.createAsync(entry, () -> release(index, entry, requestAllowance)); } - private void release(int index, BlockEntry entry, MemoryAllowance requestAllowance) { - _cache.unpin(entry, requestAllowance); + private OOCFuture release(int index, BlockEntry entry, MemoryAllowance requestAllowance) { + OOCCache.UnpinHandle unpin = _cache.unpin(entry, requestAllowance); _liveness.consumed(index); _afterRelease.accept(index); + return unpin.getCompletionFuture(); } private void reserve(int index) { diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStore.java b/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStore.java index d1ed01ffa3e..aab137d8298 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStore.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStore.java @@ -41,20 +41,38 @@ public final class MaterializedStore { private final BitSet _forgotten; private final AtomicInteger _published; private final AtomicInteger _publishedCount; + private final OOCFuture _completion; + private final OOCFuture _readersSealedFuture; + private final boolean _autoSealReaders; private volatile List _readers; private volatile int _completedSize; private volatile boolean _complete; private volatile boolean _readersSealed; private volatile boolean _closed; + private int _pendingReaders; + private int _consumers; public MaterializedStore(OOCCache cache, long streamId) { + this(cache, streamId, -1, 1); + } + + public MaterializedStore(OOCCache cache, long streamId, int expectedReaders, int consumers) { + if(expectedReaders == 0 || expectedReaders < -1) + throw new IllegalArgumentException("Expected reader count must be positive or disabled."); + if(consumers <= 0) + throw new IllegalArgumentException("Materialized store requires at least one consumer."); _cache = cache; _streamId = streamId; _registeredReaders = new ArrayList<>(); _forgotten = new BitSet(); _published = new AtomicInteger(); _publishedCount = new AtomicInteger(); + _completion = new OOCFuture<>(); + _readersSealedFuture = new OOCFuture<>(); + _autoSealReaders = expectedReaders > 0; + _pendingReaders = expectedReaders; + _consumers = consumers; _readers = Collections.emptyList(); } @@ -74,9 +92,10 @@ StoreLease publishPinnedLive(int index, T value, long bytes, MemoryAllowance } _publishedCount.incrementAndGet(); updatePublished(index + 1); - return new StoreLease<>(entry, () -> { - _cache.unpin(entry, allowance); + return StoreLease.createAsync(entry, () -> { + OOCCache.UnpinHandle unpin = _cache.unpin(entry, allowance); tryForget(index); + return unpin.getCompletionFuture(); }); } @@ -93,6 +112,19 @@ public synchronized void complete() { throw new IllegalStateException("Incomplete publication: " + _publishedCount.get() + " published items for logical range [0, " + _completedSize + ")"); _complete = true; + _completion.complete(null); + } + + void failMaterialization(Throwable error) { + _completion.completeExceptionally(error); + } + + public OOCFuture completion() { + return _completion; + } + + public OOCFuture readersSealed() { + return _readersSealedFuture; } public synchronized OrderedMaterializedStoreReader openReader(AccessPattern pattern, MemoryAllowance allowance, @@ -109,6 +141,7 @@ public synchronized OrderedMaterializedStoreReader openReader(AccessPattern p OrderedMaterializedStoreReader reader = new OrderedMaterializedStoreReader<>(_cache, _streamId, pattern, allowance, Math.max(1, maxPrefetch), softOrdering, this::forgetAfterReaderClose, this::tryForget); _registeredReaders.add(reader); + readerRegistered(); return reader; } @@ -120,6 +153,7 @@ public synchronized IndexedMaterializedStoreReader openIndexedReader(Liveness IndexedMaterializedStoreReader reader = new IndexedMaterializedStoreReader<>(_cache, _streamId, () -> _completedSize, liveness, this::forgetAfterReaderClose, this::tryForget); _registeredReaders.add(reader); + readerRegistered(); return reader; } @@ -136,7 +170,8 @@ public OOCFuture> requestPublished(int index, MemoryAllowance allo else if(entry == null) result.complete(null); else - result.complete(new StoreLease<>(entry, () -> _cache.unpin(entry, allowance))); + result.complete( + StoreLease.createAsync(entry, () -> _cache.unpin(entry, allowance).getCompletionFuture())); }); return result; } @@ -151,6 +186,16 @@ public synchronized void sealReaders() { int publishedSize = _complete ? _completedSize : _published.get(); for(int i = 0; i < publishedSize; i++) tryForget(i); + _readersSealedFuture.complete(null); + } + + private void readerRegistered() { + if(!_autoSealReaders) + return; + if(_pendingReaders <= 0) + throw new IllegalStateException("More materialized readers opened than declared."); + if(--_pendingReaders == 0) + sealReaders(); } public int size() { @@ -162,6 +207,8 @@ public void close() { synchronized(this) { if(_closed) return; + if(--_consumers > 0) + return; _closed = true; localReaders = _readersSealed ? _readers : new ArrayList<>(_registeredReaders); } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/OOCStreamMaterializer.java b/src/main/java/org/apache/sysds/runtime/ooc/store/OOCStreamMaterializer.java index 2cede71f804..3462c2cd3af 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/store/OOCStreamMaterializer.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/OOCStreamMaterializer.java @@ -57,7 +57,14 @@ public OOCStreamMaterializer(MaterializedStore store, ToIntF } public void attach(OOCStream source) { - source.setSubscriber(this); + try { + source.setSubscriber(this); + } + catch(Throwable failure) { + DMLRuntimeException wrapped = DMLRuntimeException.of(failure); + fail(wrapped); + throw wrapped; + } } public OOCFuture completion() { @@ -106,7 +113,7 @@ private void publish(OOCStream.QueueCallback callback) { } try(lease) { for(Consumer> liveConsumer : _liveConsumers) { - try(OOCStream.QueueCallback alias = new MaterializedCallback(lease.retain())) { + try(OOCStream.QueueCallback alias = new MaterializedCallback<>(lease.retain())) { liveConsumer.accept(alias); } } @@ -120,6 +127,7 @@ private void finish() { _store.complete(); } catch(RuntimeException ex) { + _store.failMaterialization(ex); deliverEos(DMLRuntimeException.of(ex)); _completion.completeExceptionally(ex); return; @@ -131,6 +139,7 @@ private void finish() { private void fail(DMLRuntimeException failure) { if(!_done.compareAndSet(false, true)) return; + _store.failMaterialization(failure); deliverEos(failure); _completion.completeExceptionally(failure); } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/OrderedMaterializedStoreReader.java b/src/main/java/org/apache/sysds/runtime/ooc/store/OrderedMaterializedStoreReader.java index 2dd048750c1..dae4ebaadd2 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/store/OrderedMaterializedStoreReader.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/OrderedMaterializedStoreReader.java @@ -140,13 +140,14 @@ public StoreLease next() throws InterruptedException { throw ex; } } - return new StoreLease<>(entry, () -> release(request._index, entry)); + return StoreLease.createAsync(entry, () -> release(request._index, entry)); } - public void release(int index, BlockEntry entry) { - _cache.unpin(entry, _allowance); + public OOCFuture release(int index, BlockEntry entry) { + OOCCache.UnpinHandle unpin = _cache.unpin(entry, _allowance); _pattern.consumed(index); _afterRelease.accept(index); + return unpin.getCompletionFuture(); } private void checkReady() { @@ -172,7 +173,7 @@ private StoreLease nextSoft() throws InterruptedException { if(entry == null) throw new IllegalStateException("Reader is closed"); fillSoft(); - return new StoreLease<>(entry, () -> release(request._index, entry)); + return StoreLease.createAsync(entry, () -> release(request._index, entry)); } private void fillStrict() { diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/StateTable.java b/src/main/java/org/apache/sysds/runtime/ooc/store/StateTable.java index 80eb57dcfe4..b70a7f4988a 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/store/StateTable.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/StateTable.java @@ -63,8 +63,16 @@ public StateTable(OOCCache cache, long streamId, int numSlots) { public void addEvictionPolicy(IntToLongFunction slotPolicy) { _evictionPolicies.add(slotPolicy); - if(_evictionPolicyInstalled.compareAndSet(false, true)) + if(_evictionPolicyInstalled.compareAndSet(false, true)) { + synchronized(this) { + for(int index = 0; index < _slots.length; index++) { + Slot slot = _slots[index]; + if(slot != null && slot._putFuture == null && slot._tableOwnedKey) + registerGeneration(index, slot._key); + } + } _cache.addEvictionPolicy(_streamId, this::scoreTableEntry); + } } public void put(int index, ManagedPayload payload) { @@ -169,8 +177,8 @@ public OOCFuture> acquire(int index, MemoryAllowance leaseAllowanc if(error != null) result.completeExceptionally(error); else - result.complete( - entry == null ? null : new StoreLease<>(entry, () -> _cache.unpin(entry, leaseAllowance))); + result.complete(entry == null ? null : StoreLease.createAsync(entry, + () -> _cache.unpin(entry, leaseAllowance).getCompletionFuture())); }); return result; } @@ -187,7 +195,8 @@ public StoreLease peek(int index, MemoryAllowance leaseAllowance) { key = slot._key; } BlockEntry entry = _cache.pinIfLive(key.getStreamId(), key.getSequenceNumber(), leaseAllowance); - return entry == null ? null : new StoreLease<>(entry, () -> _cache.unpin(entry, leaseAllowance)); + return entry == null ? null : StoreLease.createAsync(entry, + () -> _cache.unpin(entry, leaseAllowance).getCompletionFuture()); } public void clear(int index) { @@ -255,9 +264,8 @@ private void finalizeOwnedPut(int index, Slot slot, ManagedPayload payload) { synchronized(this) { slot._key = key; slot._tableOwnedKey = true; - int generation = blockIndex(key.getSequenceNumber()); - ensureGenerationCapacity(generation); - _generationSlots.set(generation, index + 1); + if(_evictionPolicyInstalled.get()) + registerGeneration(index, key); cleared = slot._cleared; putFuture = slot._putFuture; slot._putFuture = null; @@ -329,13 +337,14 @@ private OOCFuture> pinTaken(Slot slot, MemoryAllowance leaseAllowa result.completeExceptionally(completionError); return; } - result.complete(new StoreLease<>(entry, () -> _cache.unpin(entry, leaseAllowance))); + result.complete( + StoreLease.createAsync(entry, () -> _cache.unpin(entry, leaseAllowance).getCompletionFuture())); }); return result; } private void releaseSlot(Slot slot) { - if(slot._tableOwnedKey) { + if(slot._tableOwnedKey && _evictionPolicyInstalled.get()) { int generation = blockIndex(slot._key.getSequenceNumber()); AtomicIntegerArray slots = _generationSlots; if(generation < slots.length()) @@ -344,6 +353,12 @@ private void releaseSlot(Slot slot) { _cache.dereference(slot._key); } + private void registerGeneration(int index, BlockKey key) { + int generation = blockIndex(key.getSequenceNumber()); + ensureGenerationCapacity(generation); + _generationSlots.set(generation, index + 1); + } + private long scoreTableEntry(long generation) { int index = blockIndex(generation); AtomicIntegerArray slots = _generationSlots; diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/StoreLease.java b/src/main/java/org/apache/sysds/runtime/ooc/store/StoreLease.java index c9d4ca6f8cd..abc6d20079f 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/store/StoreLease.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/StoreLease.java @@ -19,32 +19,43 @@ package org.apache.sysds.runtime.ooc.store; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Supplier; + import org.apache.sysds.runtime.ooc.cache.BlockEntry; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; import org.apache.sysds.runtime.ooc.cache.io.SpillableObject; -import java.util.concurrent.atomic.AtomicInteger; - public final class StoreLease implements AutoCloseable { - private final Runnable _releaser; private final T _value; private final BlockEntry _entry; - private final AtomicInteger _shared; + private final SharedStoreLease _sharedLease; private boolean _open; - public StoreLease(BlockEntry entry, Runnable releaser) { - this(null, entry, releaser, new AtomicInteger(1)); + private StoreLease(T value, BlockEntry entry, SharedStoreLease release) { + _value = value; + _entry = entry; + _sharedLease = release; + _open = true; } - public StoreLease(T value, Runnable releaser) { - this(value, null, releaser, new AtomicInteger(1)); + public static StoreLease create(T value, Runnable releaser) { + return new StoreLease<>(value, null, new SharedStoreLease(() -> { + releaser.run(); + return OOCFuture.completed(null); + }, new AtomicInteger(1), new OOCFuture<>())); } - private StoreLease(T value, BlockEntry entry, Runnable releaser, AtomicInteger shared) { - _releaser = releaser; - _value = value; - _entry = entry; - _shared = shared; - _open = true; + public static StoreLease create(BlockEntry entry, Runnable releaser) { + return new StoreLease<>(null, entry, new SharedStoreLease(() -> { + releaser.run(); + return OOCFuture.completed(null); + }, new AtomicInteger(1), new OOCFuture<>())); + } + + public static StoreLease createAsync(BlockEntry entry, + Supplier> releaser) { + return new StoreLease<>(null, entry, new SharedStoreLease(releaser, new AtomicInteger(1), new OOCFuture<>())); } @SuppressWarnings("unchecked") @@ -63,16 +74,48 @@ synchronized BlockEntry entry() { public synchronized StoreLease retain() { if(!_open) throw new IllegalStateException("Lease is closed"); - _shared.incrementAndGet(); - return new StoreLease<>(_value, _entry, _releaser, _shared); + _sharedLease.references.incrementAndGet(); + return new StoreLease<>(_value, _entry, _sharedLease); + } + + public OOCFuture closeAsync() { + boolean release; + synchronized(this) { + if(!_open) + return _sharedLease.future; + _open = false; + release = _sharedLease.references.decrementAndGet() == 0; + } + if(release) { + OOCFuture released; + try { + released = _sharedLease.releaser.get(); + } + catch(Throwable error) { + _sharedLease.future.completeExceptionally(error); + return _sharedLease.future; + } + if(released == null) { + _sharedLease.future + .completeExceptionally(new NullPointerException("Asynchronous lease releaser returned null")); + return _sharedLease.future; + } + released.whenComplete((ignored, error) -> { + if(error == null) + _sharedLease.future.complete(null); + else + _sharedLease.future.completeExceptionally(error); + }); + } + return _sharedLease.future; } @Override - public synchronized void close() { - if(!_open) - return; - _open = false; - if(_shared.decrementAndGet() == 0) - _releaser.run(); + public void close() { + closeAsync(); + } + + private record SharedStoreLease(Supplier> releaser, AtomicInteger references, + OOCFuture future) { } } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/stream/SourceOOCStreamable.java b/src/main/java/org/apache/sysds/runtime/ooc/stream/SourceOOCStreamable.java index 4c9ca3683af..f8672701880 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/stream/SourceOOCStreamable.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/stream/SourceOOCStreamable.java @@ -28,6 +28,7 @@ public class SourceOOCStreamable implements OOCStreamable { private final CacheableData _data; + private OOCStream _reservedReadStream; public SourceOOCStreamable(CacheableData data) { _data = data; @@ -35,12 +36,12 @@ public SourceOOCStreamable(CacheableData data) { @Override public OOCStream getReadStream() { - return _data.getStreamHandle(); + return _reservedReadStream != null ? _reservedReadStream : _data.getStreamHandle(); } @Override public OOCStream getWriteStream() { - return _data.getStreamHandle(); + return _reservedReadStream != null ? _reservedReadStream : _data.getStreamHandle(); } @Override @@ -72,4 +73,10 @@ public CacheableData getData() { public void setData(CacheableData data) { throw new UnsupportedOperationException(); } + + @Override + public synchronized void reserveLazyHandle() { + if(_reservedReadStream == null) + _reservedReadStream = _data.getStreamHandle(); + } } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/util/StateTableUtils.java b/src/main/java/org/apache/sysds/runtime/ooc/util/StateTableUtils.java index a6aecd40c71..641192c16a6 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/util/StateTableUtils.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/util/StateTableUtils.java @@ -33,7 +33,7 @@ public final class StateTableUtils { public static OOCFuture putOrTake(StateTable table, int slot, OOCStream.QueueCallback tile, MemoryAllowance allowance) { - if(tile instanceof MaterializedCallback pinned && pinned.pinnedEntry() != null) + if(tile instanceof MaterializedCallback pinned && pinned.pinnedEntry() != null) return putReferenceOrTake(table, slot, pinned, allowance); ManagedPayload payload; if(tile instanceof InMemoryQueueCallback managed && managed.getManagedBytes() > 0) { @@ -64,14 +64,15 @@ public static OOCFuture putOrTake(StateTable table, i else if(lease == null) result.complete(null); else - result.complete(new Match(new MaterializedCallback(new StoreLease<>(payload.value(), payload::release)), - new MaterializedCallback(lease))); + result.complete( + new Match(new MaterializedCallback<>(StoreLease.create(payload.value(), payload::release)), + new MaterializedCallback<>(lease))); }); return result; } private static OOCFuture putReferenceOrTake(StateTable table, int slot, - MaterializedCallback pinned, MemoryAllowance allowance) { + MaterializedCallback pinned, MemoryAllowance allowance) { OOCFuture result = new OOCFuture<>(); OOCFuture> matched; try { @@ -91,7 +92,7 @@ else if(lease == null) { result.complete(null); } else - result.complete(new Match(pinned, new MaterializedCallback(lease))); + result.complete(new Match(pinned, new MaterializedCallback<>(lease))); }); return result; } diff --git a/src/test/java/org/apache/sysds/test/component/ooc/OOCInstructionUtilsTest.java b/src/test/java/org/apache/sysds/test/component/ooc/OOCInstructionUtilsTest.java index 42ca04cbbb6..a889cc73e83 100644 --- a/src/test/java/org/apache/sysds/test/component/ooc/OOCInstructionUtilsTest.java +++ b/src/test/java/org/apache/sysds/test/component/ooc/OOCInstructionUtilsTest.java @@ -56,7 +56,7 @@ public void testSubmitTasksClosesCallbacksAfterCompletion() throws Exception { }, new StreamContext().addOutStream()); IndexedMatrixValue value = new IndexedMatrixValue(new MatrixIndexes(1, 1), new MatrixBlock(1, 1, 1.0)); - source.enqueue(new MaterializedCallback<>(new StoreLease<>(value, released::incrementAndGet))); + source.enqueue(new MaterializedCallback<>(StoreLease.create(value, released::incrementAndGet))); source.closeInput(); completion.get(10, TimeUnit.SECONDS); diff --git a/src/test/java/org/apache/sysds/test/component/ooc/OOCPrimitiveTest.java b/src/test/java/org/apache/sysds/test/component/ooc/OOCPrimitiveTest.java index be6886ce1a6..7a9c6ba1e51 100644 --- a/src/test/java/org/apache/sysds/test/component/ooc/OOCPrimitiveTest.java +++ b/src/test/java/org/apache/sysds/test/component/ooc/OOCPrimitiveTest.java @@ -22,12 +22,14 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Set; import org.apache.sysds.common.Types.FileFormat; import org.apache.sysds.common.Types.ValueType; import org.apache.sysds.runtime.controlprogram.caching.MatrixObject; import org.apache.sysds.runtime.instructions.ooc.CachingStream; import org.apache.sysds.runtime.instructions.ooc.OOCStream; +import org.apache.sysds.runtime.instructions.ooc.OOCStreamable; import org.apache.sysds.runtime.instructions.ooc.SubscribableTaskQueue; import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue; import org.apache.sysds.runtime.matrix.data.MatrixBlock; @@ -36,7 +38,11 @@ import org.apache.sysds.runtime.meta.MetaDataFormat; import org.apache.sysds.runtime.ooc.cache.OOCCacheManager; import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern; +import org.apache.sysds.runtime.ooc.planning.OOCStoreLayout; +import org.apache.sysds.runtime.ooc.primitives.MaterializeOOCPrimitive; import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive; +import org.apache.sysds.runtime.ooc.store.CountingLiveness; +import org.apache.sysds.runtime.ooc.store.IndexedMaterializedStoreReader; import org.apache.sysds.runtime.ooc.stream.FilteredOOCStream; import org.apache.sysds.runtime.ooc.stream.StreamContext; import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; @@ -64,8 +70,8 @@ public void testGraphPatternsAndExecution() { TestPrimitive source = new TestPrimitive(List.of()); TestPrimitive sink = new TestPrimitive(List.of(source, source)); - Assert.assertEquals(List.of(source), sink.getChildren()); - Assert.assertEquals(List.of(sink), source.getParents()); + Assert.assertEquals(Set.of(source), sink.getChildren()); + Assert.assertEquals(Set.of(sink), source.getParents()); source.inferPatterns(); Assert.assertEquals(OOCAccessPattern.ANY, source.getAccessPattern()); Assert.assertEquals(OOCAccessPattern.ANY, sink.getAccessPattern()); @@ -87,6 +93,30 @@ public void testGraphPatternsAndExecution() { Assert.assertEquals(OOCAccessPattern.ROW_MAJOR, sink.getAccessPattern()); } + @Test + public void testPlannerDoubleMaterialize() { + OOCCacheManager.reset(); + try { + SubscribableTaskQueue source = new SubscribableTaskQueue<>(); + source.setData(new MatrixObject(ValueType.FP64, "/dev/null", + new MetaDataFormat(new MatrixCharacteristics(0, 0, 1), FileFormat.BINARY))); + CachingStream cached = new CachingStream(source); + source.closeInput(); + MaterializingTestPrimitive sink = new MaterializingTestPrimitive(cached); + cached.getReadStream().setSubscriber(OOCStream.QueueCallback::close); + cached.scheduleDeletion(); + + sink.start(); + + Assert.assertEquals(1, sink.getChildren().size()); + Assert.assertTrue(sink.getChildPrimitiveAt(0) instanceof MaterializeOOCPrimitive); + Assert.assertEquals(1, sink._executions); + } + finally { + OOCCacheManager.reset(); + } + } + @Test public void testDataGenMapTransposePipeline() { SubscribableTaskQueue generated = new SubscribableTaskQueue<>(); @@ -156,6 +186,54 @@ public void testJoinOutOfOrder() { cachedLeft.scheduleDeletion(); } + private static final class MaterializingTestPrimitive extends OOCPrimitive { + private int _executions; + + private MaterializingTestPrimitive(OOCStreamable source) { + super(new StreamContext(), source, source); + } + + @Override + public List requiredMaterializedInputs() { + return List.of(new OOCMaterializedInputRequest(0, OOCStoreLayout.ROW_MAJOR, 1), + new OOCMaterializedInputRequest(1, OOCStoreLayout.ROW_MAJOR, 1)); + } + + @Override + protected void startExecution() { + getMaterializedInput(0).whenComplete((store, error) -> { + if(error != null) + return; + store.completion().whenComplete((ignored, completionError) -> { + if(completionError != null) + return; + IndexedMaterializedStoreReader first = store + .openIndexedReader(new CountingLiveness(0, 0)); + IndexedMaterializedStoreReader second = store + .openIndexedReader(new CountingLiveness(0, 0)); + Assert.assertTrue(store.readersSealed().isDone()); + first.close(); + second.close(); + store.close(); + store.close(); + _executions++; + onComplete(); + }); + }); + } + + @Override + protected void inferPatternsInternal() { + _pattern = OOCAccessPattern.ANY; + inferParentPatterns(); + } + + @Override + protected void requestPatternInternal(OOCAccessPattern accessPattern) { + _pattern = accessPattern; + } + } + private static final class TestPrimitive extends OOCPrimitive { private int _executions; diff --git a/src/test/java/org/apache/sysds/test/component/ooc/StateTableUtilsTest.java b/src/test/java/org/apache/sysds/test/component/ooc/StateTableUtilsTest.java index 7debcd04e3a..5ecbb070b48 100644 --- a/src/test/java/org/apache/sysds/test/component/ooc/StateTableUtilsTest.java +++ b/src/test/java/org/apache/sysds/test/component/ooc/StateTableUtilsTest.java @@ -78,7 +78,7 @@ public void testCallbackPutOrTake() throws Exception { _source.put(0, new ManagedPayload<>(tile(1.0), TILE_BYTES, _producer)); StoreLease pinned = _source.peek(0, _reader); Assert.assertNotNull(pinned); - Assert.assertNull(StateTableUtils.putOrTake(_table, 0, new MaterializedCallback(pinned), _reader) + Assert.assertNull(StateTableUtils.putOrTake(_table, 0, new MaterializedCallback<>(pinned), _reader) .get(WAIT_SECONDS, TimeUnit.SECONDS)); Assert.assertEquals(0, _reader.getUsedMemory()); @@ -122,10 +122,10 @@ public void testStateTableLifecycle() throws Exception { Assert.assertNotNull(lease); Assert.assertEquals(5.0, lease.value().getValue().get(0, 0), 0.0); } - try(StoreLease lease = _table.take(0, _reader).get(WAIT_SECONDS, TimeUnit.SECONDS)) { - Assert.assertNotNull(lease); - Assert.assertEquals(5.0, lease.value().getValue().get(0, 0), 0.0); - } + StoreLease taken = _table.take(0, _reader).get(WAIT_SECONDS, TimeUnit.SECONDS); + Assert.assertNotNull(taken); + Assert.assertEquals(5.0, taken.value().getValue().get(0, 0), 0.0); + taken.closeAsync().get(WAIT_SECONDS, TimeUnit.SECONDS); Assert.assertNull(_table.take(0, _reader).get(WAIT_SECONDS, TimeUnit.SECONDS)); _producer.reserveBlocking(TILE_BYTES);