diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/ParForProgramBlock.java b/src/main/java/org/apache/sysds/runtime/controlprogram/ParForProgramBlock.java index 793cf39ce69..26716507b51 100644 --- a/src/main/java/org/apache/sysds/runtime/controlprogram/ParForProgramBlock.java +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/ParForProgramBlock.java @@ -91,6 +91,7 @@ import org.apache.sysds.runtime.instructions.cp.ScalarObject; import org.apache.sysds.runtime.instructions.cp.StringObject; import org.apache.sysds.runtime.instructions.cp.VariableCPInstruction; +import org.apache.sysds.runtime.instructions.ooc.TeeOOCInstruction; import org.apache.sysds.runtime.lineage.Lineage; import org.apache.sysds.runtime.lineage.LineageCacheConfig; import org.apache.sysds.runtime.lineage.LineageItem; @@ -730,6 +731,7 @@ private void executeLocalParFor( ExecutionContext ec, IntObject from, IntObject final LocalTaskQueue queue = new LocalTaskQueue<>(); final Thread[] threads = new Thread[_numThreads]; final LocalParWorker[] workers = new LocalParWorker[_numThreads]; + final Set resultVarNames = _resultVars.stream().map(v -> v._name).collect(Collectors.toSet()); @SuppressWarnings("unchecked") final HashMap[] workerBaselines = DMLScript.USE_OOC ? new HashMap[_numThreads] : null; try @@ -740,8 +742,11 @@ private void executeLocalParFor( ExecutionContext ec, IntObject from, IntObject workers[i] = createParallelWorker( _pwIDs[i], queue, ec, i); if(DMLScript.USE_OOC) { workerBaselines[i] = new HashMap<>(); - for(Map.Entry e : workers[i].getVariables().entrySet()) + for(Map.Entry e : workers[i].getVariables().entrySet()) { workerBaselines[i].put(e.getKey(), e.getValue()); + if(!resultVarNames.contains(e.getKey()) && e.getValue() instanceof MatrixObject matrix) + TeeOOCInstruction.incrRef(matrix.getStreamable(), 1); + } } threads[i] = new Thread( workers[i] , "PARFOR"); threads[i].setPriority(Thread.MAX_PRIORITY); @@ -785,8 +790,6 @@ private void executeLocalParFor( ExecutionContext ec, IntObject from, IntObject // Step 4) collecting results from each parallel worker //obtain results and cleanup other intermediates before result merge - Set resultVarNames = _resultVars.stream() - .map(v -> v._name).collect(Collectors.toSet()); LocalVariableMap [] localVariables = new LocalVariableMap [_numThreads]; for( int i=0; i<_numThreads; i++ ) { localVariables[i] = workers[i].getVariables(); @@ -796,6 +799,8 @@ private void executeLocalParFor( ExecutionContext ec, IntObject from, IntObject Data current = localVariables[i].get(var); if(current != null && current != workerBaselines[i].get(var)) VariableCPInstruction.processRmvarInstruction(workers[i].getExecutionContext(), var); + else if(current instanceof MatrixObject matrix) + TeeOOCInstruction.incrRef(matrix.getStreamable(), -1); } } } diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/caching/CacheableData.java b/src/main/java/org/apache/sysds/runtime/controlprogram/caching/CacheableData.java index ab27540fce7..622f6fd00a2 100644 --- a/src/main/java/org/apache/sysds/runtime/controlprogram/caching/CacheableData.java +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/caching/CacheableData.java @@ -491,7 +491,7 @@ public synchronized OOCStream getStreamHandle() { } OOCStream stream = _streamHandle.getReadStream(); - if(!stream.hasStreamCache()) + if(!_streamHandle.hasStreamCache() && !_streamHandle.hasMaterializedStore()) _streamHandle = null; // To ensure read once return stream; } 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 10a4dc88174..9f3c92dc749 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 @@ -32,6 +32,13 @@ public interface OOCStreamable { CachingStream getStreamCache(); + default boolean hasMaterializedStore() { + return false; + } + + default void scheduleMaterializedStoreDeletion() { + } + boolean isProcessed(); DataCharacteristics getDataCharacteristics(); diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/TeeOOCInstruction.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/TeeOOCInstruction.java index 548e80df942..1a46185fd72 100644 --- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/TeeOOCInstruction.java +++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/TeeOOCInstruction.java @@ -24,22 +24,24 @@ import org.apache.sysds.runtime.instructions.InstructionUtils; import org.apache.sysds.runtime.instructions.cp.CPOperand; import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue; +import org.apache.sysds.runtime.ooc.store.MaterializedStoreStreamable; import java.util.concurrent.ConcurrentHashMap; public class TeeOOCInstruction extends ComputationOOCInstruction { - private static final ConcurrentHashMap refCtr = new ConcurrentHashMap<>(); + private static final ConcurrentHashMap, Integer> refCtr = new ConcurrentHashMap<>(); public static void reset() { if (!refCtr.isEmpty()) { System.err.println("There are some dangling streams still in the cache: " + refCtr); - for(CachingStream cache : refCtr.keySet()) { + for(OOCStreamable stream : refCtr.keySet()) { try { - cache.scheduleDeletion(); + scheduleDeletion(stream); } catch(Exception ex) { - System.err.println("Failed to schedule deletion for dangling stream " + cache + ": " + ex.getMessage()); + System.err + .println("Failed to schedule deletion for dangling stream " + stream + ": " + ex.getMessage()); } } refCtr.clear(); @@ -50,19 +52,26 @@ public static void reset() { * Increments the reference counter of a stream by the set amount. */ public static void incrRef(OOCStreamable stream, int incr) { - if (!stream.hasStreamCache()) + if(!stream.hasStreamCache() && !stream.hasMaterializedStore()) return; - CachingStream cache = stream.getStreamCache(); + OOCStreamable handle = stream.hasStreamCache() ? stream.getStreamCache() : stream; - Integer ref = refCtr.compute(cache, (k, v) -> { + Integer ref = refCtr.compute(handle, (k, v) -> { if (v == null) v = 0; v += incr; return v <= 0 ? null : v; }); - if (ref == null) - cache.scheduleDeletion(); + if(ref == null) + scheduleDeletion(handle); + } + + private static void scheduleDeletion(OOCStreamable stream) { + if(stream.hasMaterializedStore()) + stream.scheduleMaterializedStoreDeletion(); + else + stream.getStreamCache().scheduleDeletion(); } protected TeeOOCInstruction(OOCType type, CPOperand in1, CPOperand out, String opcode, String istr) { @@ -82,15 +91,15 @@ public void processInstruction(ExecutionContext ec) { //get input stream MatrixObject min = ec.getMatrixObject(input1); OOCStreamable streamable = min.getStreamable(); - CachingStream handle; + OOCStreamable handle; - if(streamable.hasStreamCache()) { - handle = streamable.getStreamCache(); + if(streamable.hasStreamCache() || streamable.hasMaterializedStore()) { + handle = streamable.hasStreamCache() ? streamable.getStreamCache() : streamable; incrRef(handle, 1); } else { - // We also set the input stream handle - handle = new CachingStream(min.getStreamHandle()); + // The input and output matrix objects both retain the new reusable handle. + handle = new MaterializedStoreStreamable(min.getStreamHandle(), min); min.setStreamHandle(handle); incrRef(handle, 2); } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java index ad08c637bcf..bdc4d69f263 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java @@ -43,6 +43,7 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.locks.LockSupport; import java.util.function.LongUnaryOperator; +import java.util.function.Supplier; public final class OOCPackedCache implements OOCCache { private static final long PACKED_STREAM_ID = CachingStream._streamSeq.getNextID(); @@ -231,11 +232,8 @@ public OOCFuture pin(long sId, long tId, MemoryAllowance allowance) if(!(location instanceof SealedPackLocation packed)) return _physical.pin(sId, tId, allowance); - return packed.state().pin(_physical, allowance, false).map(physicalEntry -> { - if(physicalEntry == null) - return null; - return createLogicalPin(new BlockKey(sId, tId), packed); - }); + BlockKey key = new BlockKey(sId, tId); + return pinLogical(key, packed, () -> packed.state().pin(_physical, allowance, false)); } @Override @@ -248,11 +246,8 @@ public OOCFuture pinAdmitted(long sId, long tId, MemoryAllowance all if(!(location instanceof SealedPackLocation packed)) return _physical.pinAdmitted(sId, tId, allowance); - return packed.state().pinAdmitted(_physical, allowance).map(physicalEntry -> { - if(physicalEntry == null) - return null; - return createLogicalPin(new BlockKey(sId, tId), packed); - }); + BlockKey key = new BlockKey(sId, tId); + return pinLogical(key, packed, () -> packed.state().pinAdmitted(_physical, allowance)); } @Override @@ -265,9 +260,19 @@ public BlockEntry pinIfLive(long sId, long tId, MemoryAllowance allowance) { if(!(location instanceof SealedPackLocation packed)) return _physical.pinIfLive(sId, tId, allowance); - if(packed.state().pinIfLive(_physical, allowance) == null) - return null; - return createLogicalPin(new BlockKey(sId, tId), packed); + BlockKey key = new BlockKey(sId, tId); + packed.retain(); + try { + if(packed.state().pinIfLive(_physical, allowance) == null) { + releaseLocation(key, packed); + return null; + } + return createLogicalPin(key, packed); + } + catch(RuntimeException | Error error) { + releaseLocation(key, packed); + throw error; + } } @Override @@ -410,6 +415,7 @@ private UnpinHandle unpinPending(BlockEntry entry, PendingLogicalPin pin, Memory entry.unpin(); entry.setCacheMeta(null); PackedUnpinHandle handle = pin.builder().unpinProducer(entry, pin.slot(), allowance); + releasePendingPin(entry.getKey(), pin.builder(), pin.slot()); if(pin.builder().sealed && pin.builder().activePins == 0) pin.builder().transferProducerOwnership(_physical); scheduleSeal(pin.builder()); @@ -426,7 +432,9 @@ private UnpinHandle unpinPacked(BlockEntry entry, PackedLogicalPin pin, MemoryAl } entry.unpin(); entry.setCacheMeta(null); - return pin.location().state().unpin(this, _packReleaseDelayMs, allowance); + UnpinHandle handle = pin.location().state().unpin(this, _packReleaseDelayMs, allowance); + releaseLocation(entry.getKey(), pin.location()); + return handle; } void enqueueRelease(PackedPinState state) { @@ -492,6 +500,34 @@ private SealedPackLocation forceSeal(PendingPackLocation pending) { } } + private OOCFuture pinLogical(BlockKey key, SealedPackLocation location, + Supplier> pin) { + location.retain(); + OOCFuture physical; + try { + physical = pin.get(); + } + catch(RuntimeException | Error error) { + releaseLocation(key, location); + throw error; + } + physical.whenComplete((entry, error) -> { + if(entry == null || error != null) + releaseLocation(key, location); + }); + return physical.map(entry -> entry == null ? null : createLogicalPin(key, location)); + } + + private void releasePendingPin(BlockKey key, PackBuilder builder, int slot) { + if(builder.sealed) { + PackedCacheLocation location = getLocation(key.getStreamId(), key.getSequenceNumber()); + if(location instanceof SealedPackLocation packed) + releaseLocation(key, packed); + } + else if(builder.releaseSlot(slot) == 0) + clearLocation(key); + } + private static BlockEntry createLogicalPin(BlockKey logicalKey, SealedPackLocation location) { PackedBlock block = (PackedBlock) location.state().physicalEntry.getDataUnsafe(); Object data = block.values[location.slot()]; diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackBuilder.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackBuilder.java index 9a5f998e17d..82f074394d4 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackBuilder.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackBuilder.java @@ -59,7 +59,7 @@ int append(long streamId, long tileId, Object value, long size) { tileIds[slot] = tileId; values[slot] = value; sizes[slot] = size; - refCounts[slot] = 1; + refCounts[slot] = 2; bytes += size; activePins++; return slot; 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 index eaa9e8232c7..e1002e98005 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/MaterializeOOCPrimitive.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/MaterializeOOCPrimitive.java @@ -21,12 +21,16 @@ import java.util.List; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.ToIntFunction; 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.matrix.data.MatrixIndexes; +import org.apache.sysds.runtime.meta.DataCharacteristics; import org.apache.sysds.runtime.ooc.cache.OOCCacheManager; import org.apache.sysds.runtime.ooc.cache.OOCFuture; import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern; @@ -40,19 +44,32 @@ public final class MaterializeOOCPrimitive extends OOCPrimitive { private final OOCStoreLayout _layout; private final OOCFuture> _store; private final AtomicBoolean _finished; + private final boolean _reusable; private int _expectedReaders; private int _consumers; public MaterializeOOCPrimitive(OOCStreamable source, OOCStoreLayout layout, StreamContext context) { + this(source, layout, context, false); + } + + private MaterializeOOCPrimitive(OOCStreamable source, OOCStoreLayout layout, + StreamContext context, boolean reusable) { super(context, source.getPrimitive() == null ? List.of() : List.of(source.getPrimitive())); _source = source; _layout = layout; _store = new OOCFuture<>(); _finished = new AtomicBoolean(); + _reusable = reusable; + } + + public static MaterializeOOCPrimitive reusable(OOCStreamable source) { + return new MaterializeOOCPrimitive(source, OOCStoreLayout.ROW_MAJOR, null, true); } public synchronized void registerRequest(int expectedReaders) { + if(_reusable) + throw new IllegalStateException("Reusable materialization registers readers dynamically."); if(expectedReaders <= 0) throw new IllegalArgumentException("Materialization request requires at least one reader."); if(hasStartedExecution()) @@ -84,10 +101,19 @@ protected void requestPatternInternal(OOCAccessPattern accessPattern) { 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); + MaterializedStore store = _reusable ? new MaterializedStore<>( + OOCCacheManager.getGlobalCache(), + CachingStream._streamSeq.getNextID()) : new MaterializedStore<>(OOCCacheManager.getGlobalCache(), + CachingStream._streamSeq.getNextID(), _expectedReaders, _consumers); + DataCharacteristics characteristics = _source.getDataCharacteristics(); + AtomicInteger nextIndex = new AtomicInteger(); + ToIntFunction linearize; + if(_reusable && + (characteristics == null || !characteristics.dimsKnown() || characteristics.getBlocksize() <= 0)) + linearize = ignored -> nextIndex.getAndIncrement(); + else + linearize = indexes -> _layout.linearize(indexes, characteristics); + OOCStreamMaterializer materializer = new OOCStreamMaterializer(store, linearize, _allowance); materializer.completion().whenComplete((ignored, error) -> { if(error != null) fail(error); diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStoreStreamable.java b/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStoreStreamable.java new file mode 100644 index 00000000000..5d3879de72e --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStoreStreamable.java @@ -0,0 +1,344 @@ +/* + * 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.store; + +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Consumer; + +import org.apache.sysds.runtime.DMLRuntimeException; +import org.apache.sysds.runtime.controlprogram.caching.CacheableData; +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.meta.DataCharacteristics; +import org.apache.sysds.runtime.ooc.memory.GlobalMemoryBroker; +import org.apache.sysds.runtime.ooc.memory.SyncMemoryAllowance; +import org.apache.sysds.runtime.ooc.primitives.MaterializeOOCPrimitive; +import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive; + +public final class MaterializedStoreStreamable implements OOCStreamable { + private static final int REPLAY_PREFETCH = 8; + + private final MaterializeOOCPrimitive _primitive; + private MaterializedStore _store; + private SyncMemoryAllowance _readerAllowance; + private CacheableData _data; + private boolean _deleteScheduled; + private boolean _materializationDone; + private boolean _readersSealed; + private boolean _closed; + private int _reservedReaders; + private int _pendingReaders; + private int _activeReaders; + + public MaterializedStoreStreamable(OOCStream source, CacheableData data) { + if(source == null) + throw new IllegalArgumentException("Materialized stream requires a source."); + _data = data; + _primitive = MaterializeOOCPrimitive.reusable(source); + _primitive.store().whenComplete((store, error) -> { + if(error != null) { + markMaterializationDone(); + return; + } + synchronized(this) { + _store = store; + } + store.completion().whenComplete((ignored, completionError) -> markMaterializationDone()); + tryFinalize(); + }); + } + + @Override + public OOCStream getReadStream() { + return createReader(false); + } + + @Override + public OOCStream getReservedReadStream() { + return createReader(true); + } + + private synchronized OOCStream createReader(boolean reserved) { + if(reserved && _reservedReaders > 0) + _reservedReaders--; + else if(_deleteScheduled) + throw new DMLRuntimeException("Cannot open a reader on a materialized stream scheduled for deletion."); + _pendingReaders++; + DeferredReader stream = new DeferredReader(this); + stream.setData(_data); + stream.assignPrimitive(_primitive); + return stream; + } + + private void openReader(DeferredReader output) { + _primitive.store().whenComplete((store, storeError) -> { + if(storeError != null) { + failPendingReader(output, storeError); + return; + } + store.completion().whenComplete((ignored, completionError) -> { + if(completionError != null) { + failPendingReader(output, completionError); + return; + } + OrderedMaterializedStoreReader reader = null; + try { + reader = store.openReader(new SequentialAccessPattern(store.size()), readerAllowance(), + REPLAY_PREFETCH); + synchronized(this) { + _pendingReaders--; + _activeReaders++; + } + tryFinalize(); + drive(output, reader); + } + catch(Throwable failure) { + if(reader == null) + failPendingReader(output, failure); + else { + reader.close(); + try { + output.propagateFailure(DMLRuntimeException.of(failure)); + } + finally { + finishReader(output); + } + } + } + }); + }); + } + + private void drive(DeferredReader output, OrderedMaterializedStoreReader reader) { + StoreBackedStream replay = new StoreBackedStream<>(reader); + replay.setData(_data); + replay.setSubscriber(callback -> { + if(callback.isFailure()) { + DMLRuntimeException failure; + try { + callback.get(); + failure = new DMLRuntimeException("Materialized replay failed."); + } + catch(Throwable error) { + failure = DMLRuntimeException.of(error); + } + try { + output.propagateFailure(failure); + } + finally { + finishReader(output); + } + } + else if(callback.isEos()) { + try { + output.closeInput(); + } + finally { + finishReader(output); + } + } + else { + OOCStream.QueueCallback retained = callback.keepOpen(); + try { + output.enqueue(retained); + } + catch(Throwable failure) { + retained.close(); + throw DMLRuntimeException.of(failure); + } + } + }); + } + + private void failPendingReader(DeferredReader output, Throwable error) { + if(!output.finish()) + return; + synchronized(this) { + _pendingReaders--; + } + try { + output.propagateFailure(DMLRuntimeException.of(error)); + } + finally { + tryFinalize(); + } + } + + private void finishReader(DeferredReader output) { + if(!output.finish()) + return; + synchronized(this) { + _activeReaders--; + } + tryFinalize(); + } + + private synchronized SyncMemoryAllowance readerAllowance() { + if(_readerAllowance == null) + _readerAllowance = new SyncMemoryAllowance(GlobalMemoryBroker.get()); + return _readerAllowance; + } + + private void markMaterializationDone() { + synchronized(this) { + _materializationDone = true; + } + tryFinalize(); + } + + @Override + public synchronized void reserveLazyHandle() { + if(_deleteScheduled) + throw new DMLRuntimeException("Cannot reserve a reader on a materialized stream scheduled for deletion."); + _reservedReaders++; + } + + @Override + public void discardHandle() { + synchronized(this) { + if(_reservedReaders <= 0) + return; + _reservedReaders--; + } + tryFinalize(); + } + + @Override + public void scheduleMaterializedStoreDeletion() { + synchronized(this) { + _deleteScheduled = true; + } + tryFinalize(); + } + + private void tryFinalize() { + MaterializedStore store; + SyncMemoryAllowance allowance = null; + boolean seal = false; + boolean close = false; + synchronized(this) { + store = _store; + if(!_deleteScheduled || _reservedReaders != 0 || _pendingReaders != 0) + return; + if(store != null && !_readersSealed) { + _readersSealed = true; + seal = true; + } + if(_materializationDone && _activeReaders == 0 && !_closed) { + _closed = true; + close = store != null; + allowance = _readerAllowance; + } + } + if(seal) + store.sealReaders(); + if(close) + store.close(); + if(allowance != null) + allowance.shutdown(); + } + + @Override + public boolean hasMaterializedStore() { + return true; + } + + @Override + public OOCStream getWriteStream() { + throw new UnsupportedOperationException("Materialized streams are read-only."); + } + + @Override + public boolean hasStreamCache() { + return false; + } + + @Override + public CachingStream getStreamCache() { + return null; + } + + @Override + public boolean isProcessed() { + return false; + } + + @Override + public synchronized DataCharacteristics getDataCharacteristics() { + return _data == null ? null : _data.getDataCharacteristics(); + } + + @Override + public synchronized CacheableData getData() { + return _data; + } + + @Override + public synchronized void setData(CacheableData data) { + _data = data; + } + + @Override + public OOCPrimitive getPrimitive() { + return _primitive; + } + + private static final class DeferredReader extends SubscribableTaskQueue { + private final MaterializedStoreStreamable _owner; + private final AtomicBoolean _activated; + private final AtomicBoolean _finished; + + private DeferredReader(MaterializedStoreStreamable owner) { + _owner = owner; + _activated = new AtomicBoolean(); + _finished = new AtomicBoolean(); + } + + @Override + public void setSubscriber(Consumer> subscriber) { + super.setSubscriber(subscriber); + activate(); + } + + @Override + public IndexedMatrixValue dequeue() { + activate(); + return super.dequeue(); + } + + @Override + public QueueCallback dequeueCB() { + activate(); + return super.dequeueCB(); + } + + private void activate() { + if(_activated.compareAndSet(false, true)) + _owner.openReader(this); + } + + private boolean finish() { + return _finished.compareAndSet(false, true); + } + } +} 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 3462c2cd3af..1c85ea4f4b0 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 @@ -91,7 +91,17 @@ public void accept(OOCStream.QueueCallback callback) { finish(); return; } - publish(callback); + if(callback instanceof OOCStream.GroupQueueCallback grouped) { + @SuppressWarnings("unchecked") + OOCStream.GroupQueueCallback group = (OOCStream.GroupQueueCallback) grouped; + for(int i = 0; i < group.size(); i++) { + try(OOCStream.QueueCallback item = group.getCallback(i)) { + publish(item); + } + } + } + else + publish(callback); } catch(RuntimeException ex) { fail(DMLRuntimeException.of(ex)); 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 7a9c6ba1e51..0fe0f2e01c3 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 @@ -43,6 +43,7 @@ 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.store.MaterializedStoreStreamable; import org.apache.sysds.runtime.ooc.stream.FilteredOOCStream; import org.apache.sysds.runtime.ooc.stream.StreamContext; import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; @@ -117,6 +118,41 @@ public void testPlannerDoubleMaterialize() { } } + @Test + public void testReusableMaterializedStream() { + OOCCacheManager.reset(); + try { + MatrixObject data = new MatrixObject(ValueType.FP64, "/dev/null", + new MetaDataFormat(new MatrixCharacteristics(1, 2, 1), FileFormat.BINARY)); + SubscribableTaskQueue source = new SubscribableTaskQueue<>(); + source.setData(data); + source.enqueue(new IndexedMatrixValue(new MatrixIndexes(1, 1), new MatrixBlock(1, 1, 3d))); + source.enqueue(new IndexedMatrixValue(new MatrixIndexes(1, 2), new MatrixBlock(1, 1, 4d))); + source.closeInput(); + + MaterializedStoreStreamable handle = new MaterializedStoreStreamable(source, data); + handle.reserveLazyHandle(); + handle.reserveLazyHandle(); + handle.scheduleMaterializedStoreDeletion(); + OOCStream first = handle.getReservedReadStream(); + OOCStream second = handle.getReservedReadStream(); + first.start(); + + for(OOCStream replay : List.of(first, second)) { + double sum = 0; + OOCStream.QueueCallback callback; + while((callback = replay.dequeueCB()) != null) + try(OOCStream.QueueCallback current = callback) { + sum += current.get().getValue().get(0, 0); + } + Assert.assertEquals(7, sum, 0); + } + } + finally { + OOCCacheManager.reset(); + } + } + @Test public void testDataGenMapTransposePipeline() { SubscribableTaskQueue generated = new SubscribableTaskQueue<>(); diff --git a/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCPackedCacheTest.java b/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCPackedCacheTest.java index 7a29f8781b4..3dbceb88d0a 100644 --- a/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCPackedCacheTest.java +++ b/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCPackedCacheTest.java @@ -179,8 +179,8 @@ public void testReferenceAndDereferencePackedLocations() throws Exception { try { producer.reserveBlocking(BYTES); BlockEntry pending = cache.putPinned(STREAM_ID, 0, value(13.0), BYTES, producer); - Assert.assertEquals(2, cache.reference(pending)); - Assert.assertEquals(1, cache.dereference(pending)); + Assert.assertEquals(3, cache.reference(pending)); + Assert.assertEquals(2, cache.dereference(pending)); unpinAndFlush(cache, producer, new BlockEntry[] {pending}); awaitUsedMemory(producer, 0, WAIT_TIMEOUT_SEC); @@ -188,9 +188,11 @@ public void testReferenceAndDereferencePackedLocations() throws Exception { BlockEntry pinned = cache.pin(STREAM_ID, 0, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); Assert.assertNotNull(pinned); Assert.assertEquals(13.0, scalar(pinned), 0.0); + Assert.assertEquals(3, cache.reference(pinned)); + Assert.assertEquals(2, cache.dereference(pinned)); + Assert.assertEquals(1, cache.dereference(new BlockKey(STREAM_ID, 0))); Assert.assertEquals(2, cache.reference(pinned)); Assert.assertEquals(1, cache.dereference(pinned)); - Assert.assertEquals(0, cache.dereference(new BlockKey(STREAM_ID, 0))); await(cache.unpin(pinned, reader), WAIT_TIMEOUT_SEC); awaitUsedMemory(reader, 0, WAIT_TIMEOUT_SEC);