From 331e490ca963c50a787ef51cd0a56a0221854b48 Mon Sep 17 00:00:00 2001 From: Jannik Lindemann Date: Thu, 16 Jul 2026 18:18:11 +0200 Subject: [PATCH 1/2] [OOC] Wire First Reworked OOC Primitives --- .../controlprogram/caching/CacheableData.java | 4 +- .../controlprogram/caching/MatrixObject.java | 6 +- .../instructions/ooc/CachingStream.java | 13 +- .../ooc/CtableOOCInstruction.java | 6 + .../ooc/DataGenOOCInstruction.java | 38 +-- .../ooc/MatrixIndexingOOCInstruction.java | 5 + .../instructions/ooc/OOCInstruction.java | 297 +++++++++--------- .../instructions/ooc/PlaybackStream.java | 6 + .../instructions/ooc/ReorgOOCInstruction.java | 15 +- .../instructions/ooc/UnaryOOCInstruction.java | 13 +- .../runtime/ooc/cache/OOCCacheManager.java | 11 +- .../ooc/cache/io/OOCMatrixIOHandler.java | 82 +---- .../ooc/primitives/MappingOOCPrimitive.java | 14 +- .../ooc/primitives/TransposeOOCPrimitive.java | 13 +- .../runtime/ooc/util/OOCInstructionUtils.java | 8 +- .../test/component/ooc/OOCPrimitiveTest.java | 17 + 16 files changed, 259 insertions(+), 289 deletions(-) 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 e12d2caa5e3..ab27540fce7 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 @@ -635,7 +635,9 @@ && getRDDHandle() == null) ) { _requiresLocalWrite = false; } else if( hasStreamHandle() ) { - _data = readBlobFromStream( getStreamHandle() ); + OOCStream stream = getStreamHandle(); + stream.start(); + _data = readBlobFromStream(stream); } else if( getRDDHandle()==null || getRDDHandle().allowsShortCircuitRead() ) { if( DMLScript.STATISTICS ) diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/caching/MatrixObject.java b/src/main/java/org/apache/sysds/runtime/controlprogram/caching/MatrixObject.java index 5430a7a6c32..28fa70f7741 100644 --- a/src/main/java/org/apache/sysds/runtime/controlprogram/caching/MatrixObject.java +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/caching/MatrixObject.java @@ -637,8 +637,10 @@ protected long writeStreamToHDFS(String fname, String ofmt, int rep, FileFormatP MetaDataFormat iimd = (MetaDataFormat) _metaData; FileFormat fmt = (ofmt != null ? FileFormat.safeValueOf(ofmt) : iimd.getFileFormat()); MatrixWriter writer = MatrixWriterFactory.createMatrixWriter(fmt, rep, fprop); - return writer.writeMatrixFromStream(fname, getStreamHandle(), - getNumRows(), getNumColumns(), ConfigurationManager.getBlocksize()); + OOCStream stream = getStreamHandle(); + stream.start(); + return writer.writeMatrixFromStream(fname, stream, getNumRows(), getNumColumns(), + ConfigurationManager.getBlocksize()); } @Override 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 ff4362fe285..dab3eca3a0f 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 @@ -30,6 +30,7 @@ import org.apache.sysds.runtime.ooc.cache.GroupedBlockKey; import org.apache.sysds.runtime.ooc.cache.io.OOCIOHandler; import org.apache.sysds.runtime.ooc.cache.OOCCacheManager; +import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive; import org.apache.sysds.runtime.ooc.stream.SourceOOCStream; import org.apache.sysds.runtime.ooc.util.OOCUtils; import shaded.parquet.it.unimi.dsi.fastutil.ints.IntArrayList; @@ -272,9 +273,10 @@ else if(tmp instanceof OOCCacheManager.CachedSubCallback cachedSub) { } } } - } catch (DMLRuntimeException e) { + } + catch(RuntimeException e) { // Propagate failure to subscribers - _failure = e; + _failure = DMLRuntimeException.of(e); synchronized (this) { notifyAll(); } @@ -595,6 +597,11 @@ public CachingStream getStreamCache() { return this; } + @Override + public OOCPrimitive getPrimitive() { + return _source.getPrimitive(); + } + @Override public boolean isProcessed() { return false; @@ -617,7 +624,7 @@ public void setData(CacheableData data) { @SuppressWarnings("unchecked") public void setSubscriber(Consumer> subscriber, boolean incrConsumers) { - if(_deletable) + if(_deletable && incrConsumers) throw new DMLRuntimeException("Cannot register a new subscriber on " + this + " because has been flagged for deletion"); if(_failure != null) throw _failure; diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/CtableOOCInstruction.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/CtableOOCInstruction.java index 01fd348d101..f633cca82d4 100644 --- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/CtableOOCInstruction.java +++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/CtableOOCInstruction.java @@ -106,6 +106,12 @@ public void processInstruction( ExecutionContext ec ) { } else cst3 = ec.getScalarInput(input3).getDoubleValue(); + qIn1.start(); + if(qIn2 != null) + qIn2.start(); + if(qIn3 != null) + qIn3.start(); + HashMap blocksIn2 = new HashMap<>(), blocksIn3 = new HashMap<>(); MatrixBlock block2, block3; diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/DataGenOOCInstruction.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/DataGenOOCInstruction.java index 348d6f6930c..c5c2e299cbb 100644 --- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/DataGenOOCInstruction.java +++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/DataGenOOCInstruction.java @@ -192,40 +192,32 @@ public void processInstruction(ExecutionContext ec) { long lcols = ec.getScalarInput(cols).getLongValue(); checkValidDimensions(lrows, lcols); - OOCStream qIn = createWritableStream(); - int nrb = (int)((lrows-1) / blen)+1; - int ncb = (int)((lcols-1) / blen)+1; - - for (int row = 0; row < nrb; row++) - for (int col = 0; col < ncb; col++) - qIn.enqueue(new MatrixIndexes(row+1, col+1)); - - qIn.closeInput(); - if(sparsity == 0.0 && lrows < Integer.MAX_VALUE && lcols < Integer.MAX_VALUE) { - mapOOC(qIn, qOut, idx -> { - long rlen = Math.min(blen, lrows - (idx.getRowIndex()-1) * blen); - long clen = Math.min(blen, lcols - (idx.getColumnIndex()-1) * blen); - return new IndexedMatrixValue(idx, new MatrixBlock((int)rlen, (int)clen, 0.0)); - }); + OOCInstructionUtils.dataGen(qOut, idx -> { + long rlen = Math.min(blen, lrows - (idx.getRowIndex() - 1) * blen); + long clen = Math.min(blen, lcols - (idx.getColumnIndex() - 1) * blen); + return new MatrixBlock((int) rlen, (int) clen, 0.0); + }, getContext()); return; } if(sparsity == 1.0 && minValue == maxValue) { - mapOOC(qIn, qOut, idx -> { - long rlen = Math.min(blen, lrows - (idx.getRowIndex()-1) * blen); - long clen = Math.min(blen, lcols - (idx.getColumnIndex()-1) * blen); - return new IndexedMatrixValue(idx, new MatrixBlock((int)rlen, (int)clen, minValue)); - }); + OOCInstructionUtils.dataGen(qOut, idx -> { + long rlen = Math.min(blen, lrows - (idx.getRowIndex() - 1) * blen); + long clen = Math.min(blen, lcols - (idx.getColumnIndex() - 1) * blen); + return new MatrixBlock((int) rlen, (int) clen, minValue); + }, getContext()); return; } Well1024a bigrand = LibMatrixDatagen.setupSeedsForRand(lSeed); + int nrb = (int) ((lrows - 1) / blen) + 1; + int ncb = (int) ((lcols - 1) / blen) + 1; int nb = nrb * ncb; long[] seeds = new long[nb]; for(int i = 0; i < nb; i++) seeds[i] = bigrand.nextLong(); - mapOOC(qIn, qOut, idx -> { + OOCInstructionUtils.dataGen(qOut, idx -> { long rlen = Math.min(blen, lrows - (idx.getRowIndex()-1) * blen); long clen = Math.min(blen, lcols - (idx.getColumnIndex()-1) * blen); @@ -242,8 +234,8 @@ public void processInstruction(ExecutionContext ec) { LibMatrixDatagen.genRandomNumbers(false, 0, 1, 0, 1, mout, getGenerator(rlen, clen), bSeed, null); mout.recomputeNonZeros(); - return new IndexedMatrixValue(idx, mout); - }); + return mout; + }, getContext()); } else if(method == Types.OpOpDG.SEQ) { double lfrom = ec.getScalarInput(seq_from).getDoubleValue(); diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/MatrixIndexingOOCInstruction.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/MatrixIndexingOOCInstruction.java index 90e79d2373f..33c9b680813 100644 --- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/MatrixIndexingOOCInstruction.java +++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/MatrixIndexingOOCInstruction.java @@ -88,6 +88,7 @@ public void processInstruction(ExecutionContext ec) { Double scalarOut = null; IndexedMatrixValue tmp; + qIn.start(); while((tmp = qIn.dequeue()) != LocalTaskQueue.NO_MORE_TASKS) { if(tmp.getIndexes().getRowIndex() == firstBlockRow + 1 && tmp.getIndexes().getColumnIndex() == firstBlockCol + 1) { @@ -301,6 +302,7 @@ public void processInstruction(ExecutionContext ec) { qOut.propagateFailure(DMLRuntimeException.of(err)); return null; }); + qIn.start(); if(hasIntermediateStream) cachedStream.scheduleDeletion(); // We can immediately delete blocks after consumption @@ -356,6 +358,7 @@ else if(opcode.equalsIgnoreCase(Opcodes.LEFT_INDEX.toString())) { qOutRaw.closeInput(); return null; }); + qLhs.start(); return; } @@ -452,6 +455,8 @@ else if(opcode.equalsIgnoreCase(Opcodes.LEFT_INDEX.toString())) { qOutRaw.closeInput(); return null; }); + qLhs.start(); + qRhs.start(); } else throw new DMLRuntimeException( diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/OOCInstruction.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/OOCInstruction.java index 805f8723486..b263b307e1d 100644 --- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/OOCInstruction.java +++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/OOCInstruction.java @@ -153,6 +153,12 @@ public void postprocessInstruction(ExecutionContext ec) { OOCEventLog.onComputeEvent(_callerId, nanoTime, System.nanoTime()); } + protected StreamContext getContext() { + if(_streamContext == null) + _streamContext = new StreamContext(_callerId, getExtendedOpcode()); + return _streamContext; + } + protected void addInStream(OOCStream... queue) { if(_streamContext == null) _streamContext = new StreamContext(_callerId, getExtendedOpcode()); @@ -185,7 +191,11 @@ protected CompletableFuture filterOOC(OOCStream qIn, Consumer pr if (!inStreamsDefined() || !outStreamsDefined()) throw new NotImplementedException("filterOOC requires manual specification of all input and output streams for error propagation"); - return submitOOCTasks(qIn, c -> processor.accept(c.get()), p -> predicate.apply(p.get()), onNotProcessed != null ? (i, tmp) -> onNotProcessed.accept(tmp.get()) : null); + CompletableFuture future = submitOOCTasks(qIn, c -> processor.accept(c.get()), + p -> predicate.apply(p.get()), + onNotProcessed != null ? (i, tmp) -> onNotProcessed.accept(tmp.get()) : null); + qIn.start(); + return future; } protected OOCStream filteredOOCStream(OOCStream qIn, Function predicate) { @@ -222,27 +232,23 @@ protected CompletableFuture expandOOC(OOCStream qIn, OOCStream future = new CompletableFuture<>(); submitOOCTasks(qIn, tmp -> { - Collection out; - try(tmp) { - out = op.apply(tmp.get()); - } - if(!out.isEmpty()) { - deferredCtr.getAndIncrement(); - TaskContext.defer(() -> { - out.forEach(qOut::enqueue); - if(deferredCtr.decrementAndGet() == 0) - future.complete(null); - }); - } - }) - .thenRun(() -> { - if(deferredCtr.decrementAndGet() == 0) - future.complete(null); - }) - .exceptionally(err -> { - future.completeExceptionally(err); - return null; - }); + Collection out = op.apply(tmp.get()); + if(!out.isEmpty()) { + deferredCtr.getAndIncrement(); + TaskContext.defer(() -> { + out.forEach(qOut::enqueue); + if(deferredCtr.decrementAndGet() == 0) + future.complete(null); + }); + } + }).thenRun(() -> { + if(deferredCtr.decrementAndGet() == 0) + future.complete(null); + }).exceptionally(err -> { + future.completeExceptionally(err); + return null; + }); + qIn.start(); return future.thenRun(qOut::closeInput).exceptionally(err -> { DMLRuntimeException dmlErr = DMLRuntimeException.of(err); @@ -260,7 +266,7 @@ protected CompletableFuture mapOptionalOOC(OOCStream qIn, OOCStr Consumer> exec = tmp -> { Optional r; - try(tmp) { + try { r = optionalMapper.apply(tmp.get()); } catch(Exception e) { @@ -293,6 +299,7 @@ protected CompletableFuture mapOptionalOOC(OOCStream qIn, OOCStr future.completeExceptionally(err); return null; }); + qIn.start(); return future.thenRun(qOut::closeInput).exceptionally(err -> { DMLRuntimeException dmlErr = DMLRuntimeException.of(err); @@ -329,60 +336,61 @@ protected CompletableFuture broadcastJoinOOC(OOCStream rightReadStream = rightCached ? broadcast : rightCache.getReadStream(); CompletableFuture fut1 = submitOOCTasks(List.of(leftReadStream, rightReadStream), (i, tmp) -> { - try(tmp) { - P key = i == 0 ? onLeft.apply(tmp.get()) : onRight.apply(tmp.get()); + P key = i == 0 ? onLeft.apply(tmp.get()) : onRight.apply(tmp.get()); - if(i == 0) { // qIn stream - BroadcastedElement b; + if(i == 0) { // qIn stream + BroadcastedElement b; - synchronized(lock) { - b = availableBroadcastInput.get(key); + synchronized(lock) { + b = availableBroadcastInput.get(key); - if(b == null) { - availableLeftInput.compute(key, (k, v) -> { - if(v == null) - v = new ArrayList<>(); - v.add(tmp.get().getIndexes()); - return v; - }); - return; - } - } - - // Then items are present in cache - waitCtr.incrementAndGet(); - OOCCacheManager.requestManyBlocks( - List.of(leftCache.peekCachedBlockKey(tmp.get().getIndexes()), rightCache.peekCachedBlockKey(b.idx))) - .whenComplete((items, err) -> { - try { - broadcastingQueue.enqueue(new Tuple4<>(key, items.get(0).keepOpen(), items.get(1).keepOpen(), b)); - } finally { - items.forEach(OOCStream.QueueCallback::close); - } + if(b == null) { + availableLeftInput.compute(key, (k, v) -> { + if(v == null) + v = new ArrayList<>(); + v.add(tmp.get().getIndexes()); + return v; }); - } - else { // broadcast stream - BroadcastedElement b = new BroadcastedElement(tmp.get().getIndexes()); - List queued; - synchronized(lock) { - availableBroadcastInput.put(key, b); - queued = availableLeftInput.remove(key); + return; } + } - if(queued != null) { - for(MatrixIndexes idx : queued) { - waitCtr.incrementAndGet(); - - OOCCacheManager.requestManyBlocks( - List.of(leftCache.peekCachedBlockKey(idx), rightCache.peekCachedBlockKey(tmp.get().getIndexes()))) - .whenComplete((items, err) -> { - try{ - broadcastingQueue.enqueue(new Tuple4<>(key, items.get(0).keepOpen(), items.get(1).keepOpen(), b)); - } finally { - items.forEach(OOCStream.QueueCallback::close); - } - }); + // Then items are present in cache + waitCtr.incrementAndGet(); + OOCCacheManager.requestManyBlocks( + List.of(leftCache.peekCachedBlockKey(tmp.get().getIndexes()), rightCache.peekCachedBlockKey(b.idx))) + .whenComplete((items, err) -> { + try { + broadcastingQueue + .enqueue(new Tuple4<>(key, items.get(0).keepOpen(), items.get(1).keepOpen(), b)); + } + finally { + items.forEach(OOCStream.QueueCallback::close); } + }); + } + else { // broadcast stream + BroadcastedElement b = new BroadcastedElement(tmp.get().getIndexes()); + List queued; + synchronized(lock) { + availableBroadcastInput.put(key, b); + queued = availableLeftInput.remove(key); + } + + if(queued != null) { + for(MatrixIndexes idx : queued) { + waitCtr.incrementAndGet(); + + OOCCacheManager.requestManyBlocks(List.of(leftCache.peekCachedBlockKey(idx), + rightCache.peekCachedBlockKey(tmp.get().getIndexes()))).whenComplete((items, err) -> { + try { + broadcastingQueue.enqueue( + new Tuple4<>(key, items.get(0).keepOpen(), items.get(1).keepOpen(), b)); + } + finally { + items.forEach(OOCStream.QueueCallback::close); + } + }); } } } @@ -394,25 +402,24 @@ protected CompletableFuture broadcastJoinOOC(OOCStream fut2 = submitOOCTasks(List.of(broadcastingQueue), (i, tpl) -> { - try(tpl) { - final BroadcastedElement b = tpl.get()._4(); - final OOCStream.QueueCallback lValue = tpl.get()._2(); - final OOCStream.QueueCallback bValue = tpl.get()._3(); - - try(lValue; bValue) { - b.value = bValue.get(); - leftCache.incrProcessingCount(lValue.get().getIndexes(), 1); - qOut.enqueue(mapper.apply(lValue.get(), b)); - - if(b.tryRelease()) { - availableBroadcastInput.remove(tpl.get()._1()); - rightCache.incrProcessingCount(b.idx, 1); // Correct for incremented subscriber count to allow block deletion - } + final BroadcastedElement b = tpl.get()._4(); + final OOCStream.QueueCallback lValue = tpl.get()._2(); + final OOCStream.QueueCallback bValue = tpl.get()._3(); + + try(lValue; bValue) { + b.value = bValue.get(); + leftCache.incrProcessingCount(lValue.get().getIndexes(), 1); + qOut.enqueue(mapper.apply(lValue.get(), b)); + + if(b.tryRelease()) { + availableBroadcastInput.remove(tpl.get()._1()); + rightCache.incrProcessingCount(b.idx, 1); // Correct for incremented subscriber count to allow block + // deletion } - - if(waitCtr.decrementAndGet() == 0) - broadcastingQueue.closeInput(); } + + if(waitCtr.decrementAndGet() == 0) + broadcastingQueue.closeInput(); }); if(!qIn.hasStreamCache()) @@ -422,6 +429,8 @@ protected CompletableFuture broadcastJoinOOC(OOCStream fut = CompletableFuture.allOf(fut1, fut2); final StreamContext context = _streamContext.copy(); + qIn.start(); + broadcast.start(); return fut.thenRun(() -> { availableBroadcastInput.forEach((k, v) -> { rightCache.incrProcessingCount(v.idx, 1); @@ -458,50 +467,45 @@ protected CompletableFuture joinManyOOC(OOCStream leftReadStream = leftCached ? left : leftCache.getReadStream(); OOCStream rightReadStream = rightCached ? right : rightCache.getReadStream(); - CompletableFuture fut1 = submitOOCTasks(List.of(leftReadStream, rightReadStream), - (i, tmp) -> { - try(tmp) { - boolean leftItem = i == 0; - P key = (leftItem ? leftOn : rightOn).apply(tmp.get()); - Tuple2, List> tuple = joinMap.computeIfAbsent(key, - k -> new Tuple2<>(new ArrayList<>(releaseRightCount), new ArrayList<>(releaseLeftCount))); - BroadcastedElement b = new BroadcastedElement(tmp.get().getIndexes()); - List matches = leftItem ? tuple._2 : tuple._1; - List toInsert = leftItem ? tuple._1 : tuple._2; - int matchesSize; - boolean remove; - synchronized(tuple) { - toInsert.add(b); - matchesSize = matches.size(); - waitCtr.addAndGet(matchesSize); - remove = tuple._1.size() == releaseRightCount && tuple._2.size() == releaseLeftCount; - } - - // We have the guarantee that matches is append only so we don't need to synchronize for this - for(int mIdx = 0; mIdx < matchesSize; mIdx++) { - BroadcastedElement e = matches.get(mIdx); - OOCCacheManager.requestManyBlocks( - List.of(leftCache.peekCachedBlockKey(leftItem ? b.idx : e.idx), - rightCache.peekCachedBlockKey(leftItem ? e.idx : b.idx))).thenApply(joined -> { - try { - joinQueue.enqueue( - new Tuple5<>(key, joined.get(0).keepOpen(), joined.get(1).keepOpen(), - leftItem ? b : e, leftItem ? e : b)); - } - finally { - joined.forEach(OOCStream.QueueCallback::close); - } - return null; - }).exceptionally(t -> { - joinQueue.propagateFailure(DMLRuntimeException.of(t)); - return null; - }); - } + CompletableFuture fut1 = submitOOCTasks(List.of(leftReadStream, rightReadStream), (i, tmp) -> { + boolean leftItem = i == 0; + P key = (leftItem ? leftOn : rightOn).apply(tmp.get()); + Tuple2, List> tuple = joinMap.computeIfAbsent(key, + k -> new Tuple2<>(new ArrayList<>(releaseRightCount), new ArrayList<>(releaseLeftCount))); + BroadcastedElement b = new BroadcastedElement(tmp.get().getIndexes()); + List matches = leftItem ? tuple._2 : tuple._1; + List toInsert = leftItem ? tuple._1 : tuple._2; + int matchesSize; + boolean remove; + synchronized(tuple) { + toInsert.add(b); + matchesSize = matches.size(); + waitCtr.addAndGet(matchesSize); + remove = tuple._1.size() == releaseRightCount && tuple._2.size() == releaseLeftCount; + } - if(remove) - joinMap.remove(key); + // We have the guarantee that matches is append only so we don't need to synchronize for this + for(int mIdx = 0; mIdx < matchesSize; mIdx++) { + BroadcastedElement e = matches.get(mIdx); + OOCCacheManager.requestManyBlocks(List.of(leftCache.peekCachedBlockKey(leftItem ? b.idx : e.idx), + rightCache.peekCachedBlockKey(leftItem ? e.idx : b.idx))).thenApply(joined -> { + try { + joinQueue.enqueue(new Tuple5<>(key, joined.get(0).keepOpen(), joined.get(1).keepOpen(), + leftItem ? b : e, leftItem ? e : b)); + } + finally { + joined.forEach(OOCStream.QueueCallback::close); + } + return null; + }).exceptionally(t -> { + joinQueue.propagateFailure(DMLRuntimeException.of(t)); + return null; + }); } - }); + + if(remove) + joinMap.remove(key); + }); fut1 = fut1.thenApply(v -> { if(waitCtr.decrementAndGet() == 0) joinQueue.closeInput(); @@ -537,6 +541,8 @@ protected CompletableFuture joinManyOOC(OOCStream groupedReduceOOC(OOCStream CompletableFuture outFuture = new CompletableFuture<>(); CompletableFuture pipeFuture = pipeOOC(qIn, cb -> { - try(cb) { - Aggregator agg = aggregators.compute(cb.get().getIndexes(), (k, v) -> { - if(v == null) { - v = new Aggregator(reduce, emitCount); - busyCtr.incrementAndGet(); - v.getFuture().thenApply(imv -> { - qOut.enqueue(imv); - if(busyCtr.decrementAndGet() == 0) - outFuture.complete(null); - return null; - }) - .exceptionally(outFuture::completeExceptionally); - } - return v; - }); - agg.insert(cb.get()); - } + Aggregator agg = aggregators.compute(cb.get().getIndexes(), (k, v) -> { + if(v == null) { + v = new Aggregator(reduce, emitCount); + busyCtr.incrementAndGet(); + v.getFuture().thenApply(imv -> { + qOut.enqueue(imv); + if(busyCtr.decrementAndGet() == 0) + outFuture.complete(null); + return null; + }).exceptionally(outFuture::completeExceptionally); + } + return v; + }); + agg.insert(cb.get()); }); pipeFuture.thenRun(() -> { @@ -742,6 +746,7 @@ protected CompletableFuture groupedReduceOOC(OOCStream outFuture.complete(null); }); + qIn.start(); return outFuture.thenRun(qOut::closeInput); } 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 7526b09f592..167d7ba7aed 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 @@ -23,6 +23,7 @@ import org.apache.sysds.runtime.controlprogram.caching.CacheableData; import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue; import org.apache.sysds.runtime.meta.DataCharacteristics; +import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive; import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicBoolean; @@ -140,4 +141,9 @@ public boolean hasStreamCache() { public CachingStream getStreamCache() { return _streamCache; } + + @Override + public OOCPrimitive getPrimitive() { + return _streamCache.getPrimitive(); + } } diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/ReorgOOCInstruction.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/ReorgOOCInstruction.java index 94d896a3546..40a677e5d71 100644 --- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/ReorgOOCInstruction.java +++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/ReorgOOCInstruction.java @@ -30,9 +30,9 @@ import org.apache.sysds.runtime.instructions.cp.CPOperand; import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue; import org.apache.sysds.runtime.matrix.data.MatrixBlock; -import org.apache.sysds.runtime.matrix.data.MatrixIndexes; import org.apache.sysds.runtime.matrix.operators.Operator; import org.apache.sysds.runtime.matrix.operators.ReorgOperator; +import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; import org.apache.sysds.runtime.util.DataConverter; public class ReorgOOCInstruction extends ComputationOOCInstruction { @@ -105,19 +105,12 @@ public void processInstruction( ExecutionContext ec ) { ec.releaseMatrixInput(input1.getName()); ec.setMatrixOutput(output.getName(), soresBlock); } else if(r_op.fn instanceof SwapIndex) { - OOCStream qIn = min.getStreamHandle(); + OOCStreamable qIn = min.getStreamable(); OOCStream qOut = createWritableStream(); ec.getMatrixObject(output).setStreamHandle(qOut); - // Transpose operation - mapOOC(qIn, qOut, tmp -> { - MatrixBlock inBlock = (MatrixBlock) tmp.getValue(); - long oldRowIdx = tmp.getIndexes().getRowIndex(); - long oldColIdx = tmp.getIndexes().getColumnIndex(); - - MatrixBlock outBlock = inBlock.reorgOperations((ReorgOperator) _optr, new MatrixBlock(), -1, -1, -1); - return new IndexedMatrixValue(new MatrixIndexes(oldColIdx, oldRowIdx), outBlock); - }); + OOCInstructionUtils.transposedMap(qIn, qOut, + block -> block.reorgOperations((ReorgOperator) _optr, new MatrixBlock(), -1, -1, -1), getContext()); } } } diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/UnaryOOCInstruction.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/UnaryOOCInstruction.java index df2f50bf573..8f7946ef701 100644 --- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/UnaryOOCInstruction.java +++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/UnaryOOCInstruction.java @@ -35,6 +35,7 @@ import org.apache.sysds.runtime.matrix.data.MatrixIndexes; import org.apache.sysds.runtime.matrix.operators.UnaryOperator; import org.apache.sysds.runtime.meta.DataCharacteristics; +import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; public class UnaryOOCInstruction extends ComputationOOCInstruction { private UnaryOperator _uop = null; @@ -69,20 +70,17 @@ public void processInstruction( ExecutionContext ec ) { boolean cumSumProd = Builtin.isBuiltinCode(uop.fn, BuiltinCode.CUMSUMPROD); ec.getDataCharacteristics(output.getName()).set(min.getNumRows(), cumSumProd ? 1 : min.getNumColumns(), min.getBlocksize(), -1); - OOCStream qIn = min.getStreamHandle(); OOCStream qOut; boolean cumulative = isCumulativeUnary(uop); if(cumulative) { - qOut = processCumulativeUnaryInstruction(ec, uop, qIn); + qOut = processCumulativeUnaryInstruction(ec, uop, min.getStreamHandle()); } else { qOut = createWritableStream(); - mapOOC(qIn, qOut, tmp -> { - IndexedMatrixValue tmpOut = new IndexedMatrixValue(); - tmpOut.set(tmp.getIndexes(), tmp.getValue().unaryOperations(uop, new MatrixBlock())); - return tmpOut; - }); + OOCStreamable input = min.getStreamable(); + OOCInstructionUtils.equiMapBlock(input, qOut, block -> block.unaryOperations(uop, new MatrixBlock()), + getContext()); } ec.getMatrixObject(output).setStreamHandle(qOut); @@ -140,6 +138,7 @@ private OOCStream processCumulativeUnaryInstruction(Executio }); } + qIn.start(); return mergeOOCStreams(splitOutputs); } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheManager.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheManager.java index 78d12348fe5..07e1206a293 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheManager.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheManager.java @@ -286,8 +286,12 @@ public static class CachedQueueCallback implements OOCStream.QueueCallback @SuppressWarnings("unchecked") CachedQueueCallback(BlockEntry result, DMLRuntimeException failure) { + this(result, (T) result.getData(), failure); + } + + private CachedQueueCallback(BlockEntry result, T data, DMLRuntimeException failure) { this._result = result; - this._data = (T)result.getData(); + this._data = data; this._failure = failure; this._pinned = new AtomicBoolean(true); } @@ -305,8 +309,11 @@ public T get() { public OOCStream.QueueCallback keepOpen() { if(!_pinned.get()) throw new IllegalStateException("Cannot keep open an already closed callback"); + T data = _data; + if(data == null) + throw new IllegalStateException("Cannot keep open an empty callback"); pin(_result); - return new CachedQueueCallback<>(_result, _failure); + return new CachedQueueCallback<>(_result, data, _failure); } @Override diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCMatrixIOHandler.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCMatrixIOHandler.java index bfe1565ab3a..e9487f7bd45 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCMatrixIOHandler.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/io/OOCMatrixIOHandler.java @@ -70,9 +70,6 @@ public class OOCMatrixIOHandler implements OOCIOHandler { private static final int READER_SIZE = 16; private static final long OVERFLOW = 8192 * 1024; private static final long MAX_PARTITION_SIZE = 8192 * 8192; - private static final long GROUP_TARGET_BYTES = 8L * 1024 * 1024; - private static final long GROUP_MAX_BYTES = 16L * 1024 * 1024; - private static final int GROUP_MAX_COUNT = 64; private static final long IDLE_FLUSH_MS = 1; private final String _spillDir; @@ -371,12 +368,6 @@ private void readSequenceFile(JobConf job, Path path, SourceReadRequest request, AtomicLong bytesRead, long byteLimit, Object budgetLock, ConcurrentLinkedDeque descriptors) throws IOException { MatrixIndexes key = new MatrixIndexes(); - List groupValues = new ArrayList<>(); - List groupDescs = new ArrayList<>(); - long groupBytes = 0; - long groupSerialized = 0; - long groupStart = -1; - long groupEnd = -1; try(SequenceFile.Reader reader = new SequenceFile.Reader(job, SequenceFile.Reader.file(path))) { long pos = filePositions.get(fileIdx); @@ -409,50 +400,11 @@ else if (bytesRead.get() + blockSize > byteLimit) { SourceBlockDescriptor descriptor = new SourceBlockDescriptor(path.toString(), request.format, outIdx, recordStart, (int)(recordEnd - recordStart), blockSize); - boolean small = blockSize <= GROUP_TARGET_BYTES; - boolean contiguous = groupValues.isEmpty() || recordStart == groupEnd; - boolean canAdd = small - && contiguous - && groupValues.size() < GROUP_MAX_COUNT - && (groupBytes + (recordEnd - recordStart)) <= GROUP_MAX_BYTES; - - if (!canAdd && !groupValues.isEmpty()) { - flushSourceGroup(request, groupValues, groupDescs, groupStart, groupEnd, groupSerialized, - descriptors); - groupValues.clear(); - groupDescs.clear(); - groupBytes = 0; - groupSerialized = 0; - groupStart = -1; - groupEnd = -1; - } - - if (small) { - if (groupValues.isEmpty()) - groupStart = recordStart; - groupEnd = recordEnd; - groupValues.add(imv); - groupDescs.add(descriptor); - groupBytes += (recordEnd - recordStart); - groupSerialized += blockSize; - if (groupSerialized >= GROUP_TARGET_BYTES || groupBytes >= GROUP_MAX_BYTES || groupValues.size() >= GROUP_MAX_COUNT) { - flushSourceGroup(request, groupValues, groupDescs, groupStart, groupEnd, groupSerialized, - descriptors); - groupValues.clear(); - groupDescs.clear(); - groupBytes = 0; - groupSerialized = 0; - groupStart = -1; - groupEnd = -1; - } - } - else { - if (request.target instanceof SourceOOCStream src) - src.enqueue(imv, descriptor); - else - request.target.enqueue(imv); - descriptors.add(descriptor); - } + if(request.target instanceof SourceOOCStream src) + src.enqueue(imv, descriptor); + else + request.target.enqueue(imv); + descriptors.add(descriptor); filePositions.set(fileIdx, reader.getPosition()); if (DMLScript.OOC_LOG_EVENTS) { @@ -465,34 +417,11 @@ else if (bytesRead.get() + blockSize > byteLimit) { break; // Note that we knowingly go over limit, which could result in READER_SIZE*8MB overshoot } - if (!groupValues.isEmpty()) { - flushSourceGroup(request, groupValues, groupDescs, groupStart, groupEnd, groupSerialized, - descriptors); - } - if (!stop.get()) completed.set(fileIdx, 1); } } - private void flushSourceGroup(SourceReadRequest request, List values, - List descs, long start, long end, long totalSerialized, - ConcurrentLinkedDeque descriptors) { - if (values.isEmpty()) - return; - SourceBlockDescriptor first = descs.get(0); - OOCIOHandler.GroupSourceBlockDescriptor group = - new OOCIOHandler.GroupSourceBlockDescriptor(first.path, first.format, first.indexes, start, - (int)(end - start), totalSerialized, new ArrayList<>(descs)); - if (request.target instanceof SourceOOCStream src) - src.enqueueGroup(new ArrayList<>(values), group); - else { - for (IndexedMatrixValue v : values) - request.target.enqueue(v); - } - descriptors.addAll(group.blocks); - } - private void closeTarget(org.apache.sysds.runtime.instructions.ooc.OOCStream target, boolean close) { if(close) { try { @@ -504,7 +433,6 @@ private void closeTarget(org.apache.sysds.runtime.instructions.ooc.OOCStream _input; + private final OOCStream _input; private final OOCStreamable _output; private final Function _operation; public MappingOOCPrimitive(OOCStreamable input, OOCStreamable output, Function operation, StreamContext context) { - this(input.getPrimitive(), input, output, operation, context); + this(input.getReadStream(), output, operation, context); } - private MappingOOCPrimitive(OOCPrimitive inputPrimitive, OOCStreamable input, - OOCStreamable output, Function operation, - StreamContext context) { - super(context, inputPrimitive == null ? List.of() : List.of(inputPrimitive)); + private MappingOOCPrimitive(OOCStream input, OOCStreamable output, + Function operation, StreamContext context) { + super(context, input.getPrimitive() == null ? List.of() : List.of(input.getPrimitive())); _input = input; _output = output; _operation = operation; @@ -66,10 +65,9 @@ protected void requestPatternInternal(OOCAccessPattern accessPattern) { @Override protected void startExecution() { - OOCStream input = _input.getReadStream(); OOCStream output = _output.getWriteStream(); OOCInstructionUtils - .submitAdmittedOOCTasks(input, output, + .submitAdmittedOOCTasks(_input, output, value -> new IndexedMatrixValue(value.getIndexes(), _operation.apply(value)), _allowance, getContext()) .thenRun(this::onComplete); } 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 938c753a197..1cbe504546f 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 @@ -32,18 +32,18 @@ import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; public class TransposeOOCPrimitive extends OOCPrimitive { - private final OOCStreamable _input; + private final OOCStream _input; private final OOCStreamable _output; private final Function _operation; public TransposeOOCPrimitive(OOCStreamable input, OOCStreamable output, Function operation, StreamContext context) { - this(input.getPrimitive(), input, output, operation, context); + this(input.getReadStream(), output, operation, context); } - private TransposeOOCPrimitive(OOCPrimitive inputPrimitive, OOCStreamable input, - OOCStreamable output, Function operation, StreamContext context) { - super(context, inputPrimitive == null ? List.of() : List.of(inputPrimitive)); + private TransposeOOCPrimitive(OOCStream input, OOCStreamable output, + Function operation, StreamContext context) { + super(context, input.getPrimitive() == null ? List.of() : List.of(input.getPrimitive())); _input = input; _output = output; _operation = operation; @@ -65,9 +65,8 @@ protected void requestPatternInternal(OOCAccessPattern accessPattern) { @Override protected void startExecution() { - OOCStream input = _input.getReadStream(); OOCStream output = _output.getWriteStream(); - OOCInstructionUtils.submitAdmittedOOCTasks(input, output, value -> { + OOCInstructionUtils.submitAdmittedOOCTasks(_input, output, value -> { MatrixIndexes indexes = value.getIndexes(); return new IndexedMatrixValue(new MatrixIndexes(indexes.getColumnIndex(), indexes.getRowIndex()), _operation.apply((MatrixBlock) value.getValue())); diff --git a/src/main/java/org/apache/sysds/runtime/ooc/util/OOCInstructionUtils.java b/src/main/java/org/apache/sysds/runtime/ooc/util/OOCInstructionUtils.java index 98bab492d53..781cdd04ea2 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/util/OOCInstructionUtils.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/util/OOCInstructionUtils.java @@ -71,12 +71,12 @@ public static void equiMap(OOCStreamable input, OOCStream input, OOCStream output, + public static void transposedMap(OOCStreamable input, OOCStream output, Function operation, StreamContext context) { output.assignPrimitive(new TransposeOOCPrimitive(input, output, operation, context)); } - public static void transpose(OOCStream input, OOCStream output, + public static void transpose(OOCStreamable input, OOCStream output, StreamContext context) { transposedMap(input, output, MatrixBlock::transpose, context); } @@ -191,8 +191,12 @@ private static void subscribe(OOCStream queue, int streamIndex, future.complete(null); return; } + if(callback.isFailure() && !(callback instanceof OOCStream.GroupQueueCallback)) + callback.get(); Consumer> process = item -> { + if(item.isFailure()) + item.get(); if(predicate != null && !predicate.apply(streamIndex, item)) { if(onNotProcessed != null) onNotProcessed.accept(streamIndex, item); 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 40abcf26348..630cbbd12ae 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 @@ -30,8 +30,10 @@ 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; +import org.apache.sysds.runtime.matrix.data.MatrixIndexes; import org.apache.sysds.runtime.meta.MatrixCharacteristics; 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.primitives.OOCPrimitive; import org.apache.sysds.runtime.ooc.stream.FilteredOOCStream; @@ -41,6 +43,21 @@ import org.junit.Test; public class OOCPrimitiveTest { + @Test + public void testRetainForgottenCacheCallback() { + OOCCacheManager.reset(); + try(OOCStream.QueueCallback callback = OOCCacheManager.putAndPin(1, 1, + new IndexedMatrixValue(new MatrixIndexes(1, 1), new MatrixBlock(1, 1, 7)))) { + OOCCacheManager.forget(1, 1); + try(OOCStream.QueueCallback retained = callback.keepOpen()) { + Assert.assertEquals(7, retained.get().getValue().get(0, 0), 0); + } + } + finally { + OOCCacheManager.reset(); + } + } + @Test public void testGraphPatternsAndExecution() { TestPrimitive source = new TestPrimitive(List.of()); From 41dd93ea653ed814653a015a908c32e64270bd03 Mon Sep 17 00:00:00 2001 From: Jannik Lindemann Date: Tue, 21 Jul 2026 18:16:43 +0200 Subject: [PATCH 2/2] [OOC] Bugfix --- .../org/apache/sysds/test/component/ooc/OOCPrimitiveTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 630cbbd12ae..0ce1e1ad38c 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 @@ -47,7 +47,7 @@ public class OOCPrimitiveTest { public void testRetainForgottenCacheCallback() { OOCCacheManager.reset(); try(OOCStream.QueueCallback callback = OOCCacheManager.putAndPin(1, 1, - new IndexedMatrixValue(new MatrixIndexes(1, 1), new MatrixBlock(1, 1, 7)))) { + new IndexedMatrixValue(new MatrixIndexes(1, 1), new MatrixBlock(1, 1, 7d)))) { OOCCacheManager.forget(1, 1); try(OOCStream.QueueCallback retained = callback.keepOpen()) { Assert.assertEquals(7, retained.get().getValue().get(0, 0), 0);