diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/BroadcastOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/BroadcastOOCPrimitive.java index e2a9be63c25..49980c28959 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/BroadcastOOCPrimitive.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/BroadcastOOCPrimitive.java @@ -51,7 +51,6 @@ public final class BroadcastOOCPrimitive extends OOCPrimitive { private final Supplier _liveness; private final BiFunction _operation; private final AtomicBoolean _cleaned; - private final AtomicBoolean _failed; private final AtomicBoolean _sourceComplete; private final AtomicInteger _active; private MaterializedStore _store; @@ -70,7 +69,6 @@ public BroadcastOOCPrimitive(OOCStreamable streamed, _liveness = liveness; _operation = operation; _cleaned = new AtomicBoolean(); - _failed = new AtomicBoolean(); _sourceComplete = new AtomicBoolean(); _active = new AtomicInteger(1); } @@ -100,7 +98,7 @@ protected void startExecution() { _outputStream = _output.getWriteStream(); _ready = new SubscribableTaskQueue<>(); getContext().addOutStream(_outputStream, _ready); - OOCInstructionUtils.submitOOCTasks(_ready, callback -> process(callback.get()), getContext()) + OOCInstructionUtils.submitCloseableOOCTasks(_ready, this::process, getContext()) .whenComplete((ignored, error) -> { try { if(error != null) @@ -232,7 +230,6 @@ private void process(BroadcastWork work) { fail(failure); } finally { - work.close(); if(budget != null) budget.close(); completeOne(); @@ -255,14 +252,6 @@ private void completeOne() { } } - private void fail(Throwable error) { - if(!_failed.compareAndSet(false, true)) - return; - DMLRuntimeException failure = DMLRuntimeException.of(error); - _outputStream.propagateFailure(failure); - getContext().failAll(failure); - } - private void cleanup() { if(!_cleaned.compareAndSet(false, true)) return; diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/GroupedReduceOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/GroupedReduceOOCPrimitive.java index aa67497206e..b380848aa7d 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/GroupedReduceOOCPrimitive.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/GroupedReduceOOCPrimitive.java @@ -49,7 +49,6 @@ public final class GroupedReduceOOCPrimitive extends OOCPrimitive { private final OOCStreamable _output; private final BiFunction _merge; private final AtomicBoolean _cleaned; - private final AtomicBoolean _failed; private final AtomicBoolean _sourceComplete; private final AtomicInteger _active; private final AtomicInteger _finalizedGroups; @@ -66,7 +65,6 @@ public GroupedReduceOOCPrimitive(OOCStreamable input, OOCStr _output = output; _merge = merge; _cleaned = new AtomicBoolean(); - _failed = new AtomicBoolean(); _sourceComplete = new AtomicBoolean(); _active = new AtomicInteger(1); _finalizedGroups = new AtomicInteger(); @@ -100,7 +98,7 @@ protected void startExecution() { getContext().addInStream(input).addOutStream(_outputStream, _ready); _table = new StateTable<>(OOCCacheManager.getGlobalCache(), CachingStream._streamSeq.getNextID()); - OOCInstructionUtils.submitOOCTasks(_ready, callback -> process(callback.get()), getContext()) + OOCInstructionUtils.submitCloseableOOCTasks(_ready, this::process, getContext()) .whenComplete((ignored, error) -> { try { _outputStream.closeInput(); @@ -222,7 +220,6 @@ private void process(MergeWork work) { catch(Throwable failure) { if(merged != null) merged.release(); - work.close(); budget.close(); fail(failure); completeOne(); @@ -276,7 +273,7 @@ private void completeOne() { int remaining = _active.decrementAndGet(); if(remaining != 0) return; - if(!_failed.get() && _finalizedGroups.get() != _numGroups) + if(!hasFailed() && _finalizedGroups.get() != _numGroups) fail(new DMLRuntimeException( "Grouped reduction completed " + _finalizedGroups.get() + " of " + _numGroups + " row groups.")); try { @@ -287,14 +284,6 @@ private void completeOne() { } } - private void fail(Throwable error) { - if(!_failed.compareAndSet(false, true)) - return; - DMLRuntimeException failure = DMLRuntimeException.of(error); - _outputStream.propagateFailure(failure); - getContext().failAll(failure); - } - private void cleanup() { if(!_cleaned.compareAndSet(false, true)) return; diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/JoinOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/JoinOOCPrimitive.java index a4ef2c30ced..0ec802b181c 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/JoinOOCPrimitive.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/JoinOOCPrimitive.java @@ -19,7 +19,8 @@ package org.apache.sysds.runtime.ooc.primitives; -import java.util.concurrent.ExecutionException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiFunction; import org.apache.sysds.runtime.DMLRuntimeException; @@ -44,6 +45,9 @@ public class JoinOOCPrimitive extends OOCPrimitive { private final OOCStreamable _right; private final OOCStreamable _output; private final BiFunction _operation; + private final AtomicInteger _pending = new AtomicInteger(1); + private final AtomicInteger _unmatched = new AtomicInteger(); + private final CompletableFuture _pendingCompletion = new CompletableFuture<>(); private StateTable _table; public JoinOOCPrimitive(OOCStreamable left, OOCStreamable right, @@ -87,16 +91,13 @@ protected void startExecution() { long taskBytes = outputBytes + 2 * inputBytes; getContext().addOutStream(output); - OOCInstructionUtils.submitOOCTasks(matches, callback -> { - try(JoinWork work = callback.get()) { - IndexedMatrixValue mleft = work._left.get(); - IndexedMatrixValue mright = work._right.get(); - OOCUtils.enqueueExact(output, - new IndexedMatrixValue(mleft.getIndexes(), - _operation.apply((MatrixBlock) mleft.getValue(), (MatrixBlock) mright.getValue())), - work._budget); - } - }, callback -> true, (index, callback) -> callback.get().close(), getContext()).thenRun(() -> { + CompletableFuture processing = OOCInstructionUtils.submitCloseableOOCTasks(matches, (JoinWork work) -> { + IndexedMatrixValue mleft = work._left.get(); + IndexedMatrixValue mright = work._right.get(); + OOCUtils.enqueueExact(output, new IndexedMatrixValue(mleft.getIndexes(), + _operation.apply((MatrixBlock) mleft.getValue(), (MatrixBlock) mright.getValue())), work._budget); + }, getContext()); + CompletableFuture.allOf(processing, _pendingCompletion).thenRun(() -> { try { _table.close(); onComplete(); @@ -113,7 +114,6 @@ protected void startExecution() { private void drive(OOCStream leftInput, OOCStream rightInput, OOCStream matches, long taskBytes) { long cols = _right.getDataCharacteristics().getNumColBlocks(); - int unmatched = 0; try { while(true) { OOCStream.QueueCallback left = leftInput.dequeueCB(); @@ -129,23 +129,26 @@ private void drive(OOCStream leftInput, OOCStream callback, boolean left, long cols, long taskBytes, + private void accept(OOCStream.QueueCallback callback, boolean left, long cols, long taskBytes, OOCStream matches) { if(callback == null) - return 0; + return; OOCStream.QueueCallback owned = null; ReservationBudget budget = null; + boolean pending = false; try { owned = callback.keepOpen(); callback.close(); @@ -155,25 +158,18 @@ private int accept(OOCStream.QueueCallback callback, boolean long row = value.getIndexes().getRowIndex() - 1; long col = value.getIndexes().getColumnIndex() - 1; int slot = Math.toIntExact(row * cols + col); + _pending.incrementAndGet(); + pending = true; OOCFuture future = StateTableUtils.putOrTake(_table, slot, owned, budget); owned = null; - StateTableUtils.Match match = await(future); - if(match == null) - return 1; - JoinWork work = left ? new JoinWork(match.left(), match.right(), budget) : new JoinWork(match.right(), - match.left(), budget); + ReservationBudget pendingBudget = budget; budget = null; - try { - matches.enqueue(work); - work = null; - } - finally { - if(work != null) - work.close(); - } - return -1; + future.whenComplete((match, error) -> matchReady(match, left, pendingBudget, error, matches)); + pending = false; } finally { + if(pending) + completePending(matches); if(callback != null) callback.close(); if(owned != null) @@ -183,16 +179,57 @@ private int accept(OOCStream.QueueCallback callback, boolean } } - private static StateTableUtils.Match await(OOCFuture future) { + private void matchReady(StateTableUtils.Match match, boolean left, ReservationBudget budget, Throwable error, + OOCStream matches) { + JoinWork work = null; try { - return future.get(); + if(error != null) + throw DMLRuntimeException.of(error); + if(match == null) { + _unmatched.incrementAndGet(); + return; + } + _unmatched.decrementAndGet(); + work = left ? new JoinWork(match.left(), match.right(), budget) : new JoinWork(match.right(), match.left(), + budget); + match = null; + budget = null; + matches.enqueue(work); + work = null; + } + catch(Throwable failure) { + fail(failure); + } + finally { + if(work != null) + work.close(); + if(match != null) { + match.left().close(); + match.right().close(); + } + if(budget != null) + budget.close(); + completePending(matches); } - catch(InterruptedException error) { - Thread.currentThread().interrupt(); - throw new DMLRuntimeException(error); + } + + private void completePending(OOCStream matches) { + if(_pending.decrementAndGet() != 0) + return; + try { + int unmatched = _unmatched.get(); + if(unmatched != 0) + fail(new DMLRuntimeException("Join inputs contain " + unmatched + " unmatched blocks")); + else { + try { + matches.closeInput(); + } + catch(Exception ignored) { + } + } } - catch(ExecutionException error) { - throw DMLRuntimeException.of(error.getCause()); + finally { + _pendingCompletion.complete(null); } } 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 92bc28f8537..4c5842fb607 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 @@ -23,7 +23,6 @@ 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; @@ -130,11 +129,6 @@ protected void startExecution() { } } - 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 12d2d3a6d1f..45b7a3ff031 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 @@ -26,6 +26,7 @@ import java.util.Set; import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.sysds.runtime.DMLRuntimeException; 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; @@ -46,6 +47,7 @@ public abstract class OOCPrimitive { private final List _inputs; private final AtomicBoolean _started; private final AtomicBoolean _executionStarted; + private final AtomicBoolean _failed; protected OOCAccessPattern _pattern; protected MemoryAllowance _allowance; @@ -71,6 +73,7 @@ protected OOCPrimitive(StreamContext context) { _inputs = new ArrayList<>(); _started = new AtomicBoolean(); _executionStarted = new AtomicBoolean(); + _failed = new AtomicBoolean(); _pattern = OOCAccessPattern.UNSET; } @@ -165,6 +168,18 @@ public final void tryStartExecution() { } } + protected final boolean fail(Throwable error) { + if(!_failed.compareAndSet(false, true)) + return false; + if(_context != null) + _context.failAll(DMLRuntimeException.of(error)); + return true; + } + + protected final boolean hasFailed() { + return _failed.get(); + } + public final void onComplete() { for(int i = 0; i < _inputs.size(); i++) discardInputHandle(i); diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/PlannableDataGenOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/PlannableDataGenOOCPrimitive.java index b604e66061a..55d2a318083 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/PlannableDataGenOOCPrimitive.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/PlannableDataGenOOCPrimitive.java @@ -21,7 +21,6 @@ import java.util.function.Function; -import org.apache.sysds.runtime.DMLRuntimeException; import org.apache.sysds.runtime.instructions.ooc.OOCStream; import org.apache.sysds.runtime.instructions.ooc.OOCStreamable; import org.apache.sysds.runtime.instructions.ooc.SubscribableTaskQueue; @@ -77,7 +76,7 @@ protected void startExecution() { budget.close(); } }, getContext()).thenRun(output::closeInput).exceptionally(error -> { - output.propagateFailure(DMLRuntimeException.of(error)); + fail(error); return null; }).thenRun(this::onComplete); 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 f86c2360608..e93876dfc03 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 @@ -119,6 +119,25 @@ public static CompletableFuture submitOOCTasks(OOCStream queue, return submitOOCTasks(List.of(queue), (i, callback) -> consumer.accept(callback), null, null, context); } + public static CompletableFuture submitCloseableOOCTasks(OOCStream queue, + Consumer consumer, StreamContext context) { + return submitOOCTasks(List.of(queue), (index, callback) -> { + try(T value = callback.get()) { + consumer.accept(value); + } + catch(Exception error) { + throw DMLRuntimeException.of(error); + } + }, null, (index, callback) -> { + try { + callback.get().close(); + } + catch(Exception error) { + throw DMLRuntimeException.of(error); + } + }, context); + } + public static CompletableFuture submitAdmittedOOCTasks(OOCStream in, OOCStream out, Function operation, MemoryAllowance allowance, StreamContext context) { 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 a889cc73e83..7737f45e9fd 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 @@ -64,6 +64,38 @@ public void testSubmitTasksClosesCallbacksAfterCompletion() throws Exception { Assert.assertEquals(1, released.get()); } + @Test + public void testSubmitCloseableOOCTasks() throws Exception { + SubscribableTaskQueue source = new SubscribableTaskQueue<>(); + AtomicInteger processed = new AtomicInteger(); + AtomicInteger closed = new AtomicInteger(); + CompletableFuture completion = OOCInstructionUtils.submitCloseableOOCTasks(source, + (OwnedTask work) -> processed.addAndGet(work._value), new StreamContext().addOutStream()); + + source.enqueue(new OwnedTask(1, closed)); + source.enqueue(new OwnedTask(2, closed)); + source.closeInput(); + completion.get(10, TimeUnit.SECONDS); + + Assert.assertEquals(3, processed.get()); + Assert.assertEquals(2, closed.get()); + } + + private static final class OwnedTask implements AutoCloseable { + private final int _value; + private final AtomicInteger _closed; + + private OwnedTask(int value, AtomicInteger closed) { + _value = value; + _closed = closed; + } + + @Override + public void close() { + _closed.incrementAndGet(); + } + } + @Test public void testSubmitTasksWaitsForAllStreams() throws Exception { SubscribableTaskQueue first = new SubscribableTaskQueue<>();