From 16bca828f99f0eef93388fcee221f3a773a476c6 Mon Sep 17 00:00:00 2001 From: Jannik Lindemann Date: Thu, 16 Jul 2026 15:15:59 +0200 Subject: [PATCH] [OOC] Add OOCInstructionUtils --- .../ooc/CSVReblockOOCInstruction.java | 5 +- .../ooc/DataGenOOCInstruction.java | 6 +- .../instructions/ooc/OOCInstruction.java | 322 ++---------------- .../runtime/ooc/stream/StreamContext.java | 52 ++- .../runtime/ooc/util/OOCInstructionUtils.java | 313 +++++++++++++++++ .../ooc/OOCInstructionUtilsTest.java | 110 ++++++ 6 files changed, 493 insertions(+), 315 deletions(-) create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/util/OOCInstructionUtils.java create mode 100644 src/test/java/org/apache/sysds/test/component/ooc/OOCInstructionUtilsTest.java diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/CSVReblockOOCInstruction.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/CSVReblockOOCInstruction.java index 4ac54ee3a57..2df1388c8ab 100644 --- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/CSVReblockOOCInstruction.java +++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/CSVReblockOOCInstruction.java @@ -32,6 +32,7 @@ import org.apache.sysds.runtime.matrix.operators.Operator; import org.apache.sysds.runtime.meta.DataCharacteristics; import org.apache.sysds.runtime.ooc.stream.StreamContext; +import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; public class CSVReblockOOCInstruction extends ComputationOOCInstruction { private final int blen; @@ -74,14 +75,14 @@ public void processInstruction(ExecutionContext ec) { final long cols = mc.getCols(); final long nnz = mc.getNonZeros(); - submitOOCTask(() -> { + OOCInstructionUtils.submitOOCTask(() -> { try { reader.readMatrixAsStream(qOut, fileName, rows, cols, blen, nnz); } catch(Exception ex) { throw (ex instanceof DMLRuntimeException) ? (DMLRuntimeException) ex : new DMLRuntimeException(ex); } - }, new StreamContext().addOutStream(qOut)); + }, new StreamContext(_callerId, getExtendedOpcode()).addOutStream(qOut)); MatrixObject mout = ec.getMatrixObject(output); mout.setStreamHandle(qOut); 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 f44eb79dc98..348d6f6930c 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 @@ -38,6 +38,7 @@ import org.apache.sysds.runtime.matrix.data.RandomMatrixGenerator; import org.apache.sysds.runtime.matrix.operators.UnaryOperator; import org.apache.sysds.runtime.ooc.stream.StreamContext; +import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; import org.apache.sysds.runtime.util.UtilFunctions; public class DataGenOOCInstruction extends UnaryOOCInstruction { @@ -259,8 +260,7 @@ else if(method == Types.OpOpDG.SEQ) { final int maxK = (int) UtilFunctions.getSeqLength(lfrom, lto, lincr); final double finalLincr = lincr; - - submitOOCTask(() -> { + OOCInstructionUtils.submitOOCTask(() -> { int k = 0; double curFrom = lfrom; double curTo; @@ -286,7 +286,7 @@ else if(method == Types.OpOpDG.SEQ) { } qOut.closeInput(); - }, new StreamContext().addOutStream(qOut)); + }, new StreamContext(_callerId, getExtendedOpcode()).addOutStream(qOut)); } else throw new NotImplementedException(); 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 80d71231646..805f8723486 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 @@ -41,8 +41,7 @@ import org.apache.sysds.runtime.ooc.stream.SplittingOOCStream; import org.apache.sysds.runtime.ooc.stream.StreamContext; import org.apache.sysds.runtime.ooc.stream.TaskContext; -import org.apache.sysds.runtime.util.CommonThreadPool; -import org.apache.sysds.utils.Statistics; +import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; import scala.Tuple2; import scala.Tuple4; import scala.Tuple5; @@ -60,22 +59,17 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.ForkJoinTask; -import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; -import java.util.concurrent.atomic.LongAdder; import java.util.function.BiConsumer; import java.util.function.BiFunction; import java.util.function.Consumer; import java.util.function.Function; public abstract class OOCInstruction extends Instruction { - public static final boolean ALLOW_PIPELINING = true; - public static final ExecutorService COMPUTE_EXECUTOR = CommonThreadPool.get(); - private static final AtomicInteger COMPUTE_IN_FLIGHT = new AtomicInteger(0); - private static final int COMPUTE_BACKPRESSURE_THRESHOLD = 100; + public static final boolean ALLOW_PIPELINING = OOCInstructionUtils.ALLOW_PIPELINING; + public static final ExecutorService COMPUTE_EXECUTOR = OOCInstructionUtils.COMPUTE_EXECUTOR; protected static final Log LOG = LogFactory.getLog(OOCInstruction.class.getName()); - private static final AtomicInteger nextStreamId = new AtomicInteger(0); private long nanoTime; public enum OOCType { @@ -86,7 +80,6 @@ public enum OOCType { protected final OOCInstruction.OOCType _ooctype; protected final boolean _requiresLabelUpdate; protected StreamContext _streamContext; - private LongAdder _localStatisticsAdder; public final int _callerId; protected OOCInstruction(OOCInstruction.OOCType type, String opcode, String istr) { @@ -101,17 +94,15 @@ protected OOCInstruction(OOCInstruction.OOCType type, Operator op, String opcode _requiresLabelUpdate = super.requiresLabelUpdate(); - if (DMLScript.STATISTICS) - _localStatisticsAdder = new LongAdder(); _callerId = DMLScript.OOC_LOG_EVENTS ? OOCEventLog.registerCaller(getExtendedOpcode() + "_" + hashCode()) : 0; } public static int getComputeInFlight() { - return COMPUTE_IN_FLIGHT.get(); + return OOCInstructionUtils.getComputeInFlight(); } public static int getComputeBackpressureThreshold() { - return COMPUTE_BACKPRESSURE_THRESHOLD; + return OOCInstructionUtils.getComputeBackpressureThreshold(); } @Override @@ -164,13 +155,13 @@ public void postprocessInstruction(ExecutionContext ec) { protected void addInStream(OOCStream... queue) { if(_streamContext == null) - _streamContext = new StreamContext(); + _streamContext = new StreamContext(_callerId, getExtendedOpcode()); _streamContext.addInStream(queue); } protected void addOutStream(OOCStream... queue) { if(_streamContext == null) - _streamContext = new StreamContext(); + _streamContext = new StreamContext(_callerId, getExtendedOpcode()); _streamContext.addOutStream(queue); } @@ -982,297 +973,30 @@ protected CompletableFuture scanOOC(OOCStream q }); } - protected CompletableFuture scanOOC(OOCStream qIn, OOCStream qOut, - Function seqFn, BiFunction scanner, - Function carryFn, long sequenceSize) { - return scanOOC(qIn, qOut, seqFn, (IndexedMatrixValue item, C carry) -> { - R out = scanner.apply(item, carry); - if(out == null) - throw new DMLRuntimeException("Ordered scan output must not be null."); - return new ScanStep<>(out, carryFn.apply(out)); - }, sequenceSize); - } - - protected CompletableFuture submitOOCTasks(final List> queues, BiConsumer> consumer) { + protected CompletableFuture submitOOCTasks(final List> queues, + BiConsumer> consumer) { return submitOOCTasks(queues, consumer, null, null); } - protected CompletableFuture submitOOCTasks(final List> queues, BiConsumer> consumer, BiFunction, Boolean> predicate, BiConsumer> onNotProcessed) { - addInStream(queues.toArray(OOCStream[]::new)); - if(!outStreamsDefined()) - throw new IllegalArgumentException("Explicit specification of all output streams is required before submitting tasks. If no output streams are present use addOutStream()."); - - final List activeTaskCtrs = new ArrayList<>(queues.size()); - final List> futures = new ArrayList<>(queues.size()); - - for(int i = 0; i < queues.size(); i++) { - activeTaskCtrs.add(new AtomicInteger(1)); - futures.add(new CompletableFuture<>()); - } - - final CompletableFuture globalFuture = CompletableFuture.allOf(futures.toArray(CompletableFuture[]::new)); - final StreamContext streamContext = _streamContext.copy(); // Snapshot of the current stream context - if(streamContext == null || !streamContext.inStreamsDefined() || !streamContext.outStreamsDefined()) - throw new IllegalArgumentException("Explicit specification of all output streams is required before submitting tasks. If no output streams are present use addOutStream()."); - - int i = 0; - @SuppressWarnings("unused") - final int streamId = nextStreamId.getAndIncrement(); - - for (OOCStream queue : queues) { - final int k = i; - final AtomicInteger localTaskCtr = activeTaskCtrs.get(k); - final CompletableFuture localFuture = futures.get(k); - final AtomicBoolean closeRaceWatchdog = new AtomicBoolean(false); - - queue.setSubscriber(oocTask(callback -> { - long startTime = DMLScript.STATISTICS ? System.nanoTime() : 0; - try(callback) { - if(callback.isEos()) { - if(!closeRaceWatchdog.compareAndSet(false, true)) - throw new DMLRuntimeException( - "Race condition observed: NO_MORE_TASKS callback has been triggered more than once"); - - if(localTaskCtr.decrementAndGet() == 0) { - // Then we can run the finalization procedure already - localFuture.complete(null); - } - return; - } - - Consumer> process = cb -> { - if(predicate != null && !predicate.apply(k, cb)) { // Can get closed due to cancellation - if(onNotProcessed != null) - onNotProcessed.accept(k, cb); - return; - } - - if(localFuture.isDone()) { - if(onNotProcessed != null) - onNotProcessed.accept(k, cb); - return; - } - else { - localTaskCtr.incrementAndGet(); - } - - // The item needs to be pinned in memory to be accessible in the executor thread - final OOCStream.QueueCallback pinned = cb.keepOpen(); - - COMPUTE_IN_FLIGHT.incrementAndGet(); - try { - Runnable oocTask = oocTask(() -> { - long taskStartTime = DMLScript.STATISTICS || DMLScript.OOC_LOG_EVENTS ? System.nanoTime() : 0; - try(pinned) { - consumer.accept(k, pinned); - - if(localTaskCtr.decrementAndGet() == 0) { - TaskContext.defer(() -> localFuture.complete(null)); - } - } - finally { - COMPUTE_IN_FLIGHT.decrementAndGet(); - if (DMLScript.STATISTICS) { - _localStatisticsAdder.add(System.nanoTime() - taskStartTime); - if (globalFuture.isDone()) { - Statistics.maintainOOCHeavyHitter(getExtendedOpcode(), _localStatisticsAdder.sum()); - _localStatisticsAdder.reset(); - } - } - if (DMLScript.OOC_LOG_EVENTS) - OOCEventLog.onComputeEvent(_callerId, taskStartTime, System.nanoTime()); - } - }, localFuture, streamContext); - COMPUTE_EXECUTOR.submit(oocTask); - } - catch (Exception e) { - COMPUTE_IN_FLIGHT.decrementAndGet(); - throw e; - } - }; - - if(callback instanceof OOCStream.GroupQueueCallback) { - OOCStream.GroupQueueCallback group = (OOCStream.GroupQueueCallback) callback; - - if(localFuture.isDone()) { - for(int idx = 0; idx < group.size(); idx++) { - OOCStream.QueueCallback sub = group.getCallback(idx); - try(sub) { - if(onNotProcessed != null) - onNotProcessed.accept(k, sub); - } - } - return; - } - - localTaskCtr.incrementAndGet(); - final OOCStream.GroupQueueCallback pinnedGroup = - (OOCStream.GroupQueueCallback) group.keepOpen(); - - COMPUTE_IN_FLIGHT.incrementAndGet(); - try { - Runnable oocTask = oocTask(() -> { - long taskStartTime = DMLScript.STATISTICS || DMLScript.OOC_LOG_EVENTS ? System.nanoTime() : 0; - try(pinnedGroup) { - for(int idx = 0; idx < pinnedGroup.size(); idx++) { - OOCStream.QueueCallback sub = pinnedGroup.getCallback(idx); - try(sub) { - process.accept(sub); - } - } - - if(localTaskCtr.decrementAndGet() == 0) { - TaskContext.defer(() -> localFuture.complete(null)); - } - } - finally { - COMPUTE_IN_FLIGHT.decrementAndGet(); - if (DMLScript.STATISTICS) { - _localStatisticsAdder.add(System.nanoTime() - taskStartTime); - if (globalFuture.isDone()) { - Statistics.maintainOOCHeavyHitter(getExtendedOpcode(), _localStatisticsAdder.sum()); - _localStatisticsAdder.reset(); - } - } - if (DMLScript.OOC_LOG_EVENTS) - OOCEventLog.onComputeEvent(_callerId, taskStartTime, System.nanoTime()); - } - }, localFuture, streamContext); - COMPUTE_EXECUTOR.submit(oocTask); - } - catch (Exception e) { - COMPUTE_IN_FLIGHT.decrementAndGet(); - throw e; - } - } - else { - process.accept(callback); - } - - if(closeRaceWatchdog.get()) // Sanity check - throw new DMLRuntimeException("Race condition observed"); - } - catch(Throwable t) { - streamContext.failAll(DMLRuntimeException.of(t)); - throw t; - } - finally { - if (DMLScript.STATISTICS) { - _localStatisticsAdder.add(System.nanoTime() - startTime); - if (globalFuture.isDone()) { - Statistics.maintainOOCHeavyHitter(getExtendedOpcode(), _localStatisticsAdder.sum()); - _localStatisticsAdder.reset(); - } - } - } - }, null, streamContext)); - - i++; - } - - return globalFuture.handle((res, e) -> { - if (globalFuture.isCancelled() || globalFuture.isCompletedExceptionally()) { - futures.forEach(f -> { - if(!f.isDone()) { - if(globalFuture.isCancelled() || globalFuture.isCompletedExceptionally()) - f.cancel(true); - else - f.complete(null); - } - }); - } - - streamContext.clear(); - return null; - }); + protected CompletableFuture submitOOCTasks(final List> queues, + BiConsumer> consumer, + BiFunction, Boolean> predicate, + BiConsumer> onNotProcessed) { + if(_streamContext == null) + _streamContext = new StreamContext(_callerId, getExtendedOpcode()); + return OOCInstructionUtils.submitOOCTasks(queues, consumer, predicate, onNotProcessed, _streamContext); } - protected CompletableFuture submitOOCTasks(OOCStream queue, Consumer> consumer) { + protected CompletableFuture submitOOCTasks(OOCStream queue, + Consumer> consumer) { return submitOOCTasks(List.of(queue), (i, tmp) -> consumer.accept(tmp), null, null); } - protected CompletableFuture submitOOCTasks(OOCStream queue, Consumer> consumer, Function, Boolean> predicate, BiConsumer> onNotProcessed) { - return submitOOCTasks(List.of(queue), (i, tmp) -> consumer.accept(tmp), (i, tmp) -> predicate.apply(tmp), onNotProcessed); - } - - protected CompletableFuture submitOOCTask(Runnable r, StreamContext ctx) { - ExecutorService pool = CommonThreadPool.getDynamicPool(); - final CompletableFuture future = new CompletableFuture<>(); - try { - COMPUTE_IN_FLIGHT.incrementAndGet(); - pool.submit(oocTask(() -> { - long startTime = DMLScript.STATISTICS || DMLScript.OOC_LOG_EVENTS ? System.nanoTime() : 0; - try { - r.run(); - future.complete(null); - ctx.clear(); - if (DMLScript.STATISTICS) - Statistics.maintainOOCHeavyHitter(getExtendedOpcode(), System.nanoTime() - startTime); - if (DMLScript.OOC_LOG_EVENTS) - OOCEventLog.onComputeEvent(_callerId, startTime, System.nanoTime()); - } - finally { - COMPUTE_IN_FLIGHT.decrementAndGet(); - } - }, future, ctx)); - } - catch (Exception ex) { - COMPUTE_IN_FLIGHT.decrementAndGet(); - throw new DMLRuntimeException(ex); - } - - return future; - } - - private Runnable oocTask(Runnable r, CompletableFuture future, StreamContext ctx) { - return () -> { - boolean setContext = TaskContext.getContext() == null; - if(setContext) - TaskContext.setContext(new TaskContext()); - long startTime = DMLScript.STATISTICS ? System.nanoTime() : 0; - try { - r.run(); - if(setContext) { - while(TaskContext.runDeferred()) { - } - } - } - catch (Exception ex) { - DMLRuntimeException re = DMLRuntimeException.of(ex); - - ctx.failAll(re); - - if (future != null) - future.completeExceptionally(re); - - // Rethrow to ensure proper future handling - throw re; - } finally { - if(setContext) - TaskContext.clearContext(); - if (DMLScript.STATISTICS) - _localStatisticsAdder.add(System.nanoTime() - startTime); - } - }; - } - - private Consumer> oocTask(Consumer> c, CompletableFuture future, StreamContext ctx) { - return callback -> { - try { - c.accept(callback); - } - catch (Exception ex) { - DMLRuntimeException re = DMLRuntimeException.of(ex); - - ctx.failAll(re); - - if (future != null) - future.completeExceptionally(re); - - // Rethrow to ensure proper future handling - throw re; - } - }; + protected CompletableFuture submitOOCTasks(OOCStream queue, + Consumer> consumer, Function, Boolean> predicate, + BiConsumer> onNotProcessed) { + return submitOOCTasks(List.of(queue), (i, tmp) -> consumer.accept(tmp), (i, tmp) -> predicate.apply(tmp), + onNotProcessed); } /** diff --git a/src/main/java/org/apache/sysds/runtime/ooc/stream/StreamContext.java b/src/main/java/org/apache/sysds/runtime/ooc/stream/StreamContext.java index 9c9f2e3fc0e..4565d1768b7 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/stream/StreamContext.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/stream/StreamContext.java @@ -27,12 +27,38 @@ import java.util.List; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.LongAdder; public class StreamContext { + private final int _callerId; + private final String _opcode; + private final LongAdder _statistics; private Set> _inStreams; private Set> _outStreams; private DMLRuntimeException _failure; + public StreamContext() { + this(0, null); + } + + public StreamContext(int callerId, String opcode) { + _callerId = callerId; + _opcode = opcode; + _statistics = new LongAdder(); + } + + public int getCallerId() { + return _callerId; + } + + public String getExtendedOpcode() { + return _opcode; + } + + public LongAdder getLocalStatisticsLongAdder() { + return _statistics; + } + public boolean inStreamsDefined() { return _inStreams != null; } @@ -73,19 +99,23 @@ public void failAll(DMLRuntimeException e) { return; _failure = e; - for(OOCStream stream : _outStreams) { - try { - stream.propagateFailure(e); + if(_outStreams != null) + for(OOCStream stream : _outStreams) { + try { + stream.propagateFailure(e); + } + catch(Throwable ignored) { + } } - catch(Throwable ignored) {} - } - for(OOCStream stream : _inStreams) { - try { - stream.propagateFailure(e); + if(_inStreams != null) + for(OOCStream stream : _inStreams) { + try { + stream.propagateFailure(e); + } + catch(Throwable ignored) { + } } - catch(Throwable ignored) {} - } } public void clear() { @@ -94,7 +124,7 @@ public void clear() { } public StreamContext copy() { - StreamContext cpy = new StreamContext(); + StreamContext cpy = new StreamContext(_callerId, _opcode); cpy._inStreams = _inStreams; cpy._outStreams = _outStreams; return cpy; 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 new file mode 100644 index 00000000000..75f6f9db2e5 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/util/OOCInstructionUtils.java @@ -0,0 +1,313 @@ +/* + * 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.util; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BiConsumer; +import java.util.function.BiFunction; +import java.util.function.Consumer; +import java.util.function.Function; + +import org.apache.sysds.api.DMLScript; +import org.apache.sysds.runtime.DMLRuntimeException; +import org.apache.sysds.runtime.instructions.ooc.OOCStream; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; +import org.apache.sysds.runtime.ooc.stats.OOCEventLog; +import org.apache.sysds.runtime.ooc.stream.StreamContext; +import org.apache.sysds.runtime.ooc.stream.TaskContext; +import org.apache.sysds.runtime.util.CommonThreadPool; +import org.apache.sysds.utils.Statistics; + +public final class OOCInstructionUtils { + public static final boolean ALLOW_PIPELINING = true; + public static final ExecutorService COMPUTE_EXECUTOR = CommonThreadPool.get(); + private static final AtomicInteger COMPUTE_IN_FLIGHT = new AtomicInteger(); + private static final int COMPUTE_BACKPRESSURE_THRESHOLD = 100; + + public static int getComputeInFlight() { + return COMPUTE_IN_FLIGHT.get(); + } + + public static int getComputeBackpressureThreshold() { + return COMPUTE_BACKPRESSURE_THRESHOLD; + } + + public static CompletableFuture submitOOCTasks(OOCStream queue, + Consumer> consumer, StreamContext context) { + return submitOOCTasks(List.of(queue), (i, callback) -> consumer.accept(callback), null, null, context); + } + + public static CompletableFuture submitOOCTasks(OOCStream queue, + Consumer> consumer, Function, Boolean> predicate, + BiConsumer> onNotProcessed, StreamContext context) { + return submitOOCTasks(List.of(queue), (i, callback) -> consumer.accept(callback), + (i, callback) -> predicate.apply(callback), onNotProcessed, context); + } + + public static CompletableFuture submitOOCTasks(List> queues, + BiConsumer> consumer, StreamContext context) { + return submitOOCTasks(queues, consumer, null, null, context); + } + + public static CompletableFuture submitOOCTasks(List> queues, + BiConsumer> consumer, + BiFunction, Boolean> predicate, + BiConsumer> onNotProcessed, StreamContext context) { + context.addInStream(queues.toArray(OOCStream[]::new)); + if(!context.outStreamsDefined()) + throw new IllegalArgumentException("Explicit specification of all output streams is required before " + + "submitting tasks. If no output streams are present use addOutStream()."); + + List activeTaskCounters = new ArrayList<>(queues.size()); + List> futures = new ArrayList<>(queues.size()); + for(int i = 0; i < queues.size(); i++) { + activeTaskCounters.add(new AtomicInteger(1)); + futures.add(new OOCFuture<>()); + } + + CompletableFuture globalFuture = new CompletableFuture<>(); + AtomicInteger remaining = new AtomicInteger(futures.size()); + if(futures.isEmpty()) + globalFuture.complete(null); + for(OOCFuture future : futures) + future.whenComplete((result, error) -> { + if(error != null) + globalFuture.completeExceptionally(error); + else if(remaining.decrementAndGet() == 0) + globalFuture.complete(null); + }); + StreamContext streamContext = context.copy(); + for(int i = 0; i < queues.size(); i++) + subscribe(queues.get(i), i, consumer, predicate, onNotProcessed, activeTaskCounters.get(i), futures.get(i), + globalFuture, streamContext); + + return globalFuture.handle((result, error) -> { + if(error != null) { + for(OOCFuture future : futures) + if(!future.isDone()) + future.completeExceptionally(error); + } + streamContext.clear(); + return null; + }); + } + + private static void subscribe(OOCStream queue, int streamIndex, + BiConsumer> consumer, + BiFunction, Boolean> predicate, + BiConsumer> onNotProcessed, AtomicInteger activeTaskCounter, + OOCFuture future, CompletableFuture globalFuture, StreamContext context) { + AtomicBoolean closed = new AtomicBoolean(); + queue.setSubscriber(guard(callback -> { + long startTime = DMLScript.STATISTICS ? System.nanoTime() : 0; + try(callback) { + if(callback.isEos()) { + if(!closed.compareAndSet(false, true)) + throw new DMLRuntimeException( + "Race condition observed: NO_MORE_TASKS callback has been triggered more than once"); + if(activeTaskCounter.decrementAndGet() == 0) + future.complete(null); + return; + } + + Consumer> process = item -> { + if(predicate != null && !predicate.apply(streamIndex, item)) { + if(onNotProcessed != null) + onNotProcessed.accept(streamIndex, item); + return; + } + if(future.isDone()) { + if(onNotProcessed != null) + onNotProcessed.accept(streamIndex, item); + return; + } + + activeTaskCounter.incrementAndGet(); + OOCStream.QueueCallback pinned = item.keepOpen(); + submit(() -> { + long taskStartTime = DMLScript.STATISTICS || DMLScript.OOC_LOG_EVENTS ? System.nanoTime() : 0; + try(pinned) { + consumer.accept(streamIndex, pinned); + if(activeTaskCounter.decrementAndGet() == 0) + TaskContext.defer(() -> future.complete(null)); + } + finally { + recordStatistics(context, globalFuture, taskStartTime); + recordEvent(context, taskStartTime); + } + }, future, context); + }; + + if(callback instanceof OOCStream.GroupQueueCallback) { + OOCStream.GroupQueueCallback group = (OOCStream.GroupQueueCallback) callback; + if(future.isDone()) { + for(int index = 0; index < group.size(); index++) { + try(OOCStream.QueueCallback item = group.getCallback(index)) { + if(onNotProcessed != null) + onNotProcessed.accept(streamIndex, item); + } + } + return; + } + + activeTaskCounter.incrementAndGet(); + OOCStream.GroupQueueCallback pinned = (OOCStream.GroupQueueCallback) group.keepOpen(); + submit(() -> { + long taskStartTime = DMLScript.STATISTICS || DMLScript.OOC_LOG_EVENTS ? System.nanoTime() : 0; + try(pinned) { + for(int index = 0; index < pinned.size(); index++) { + try(OOCStream.QueueCallback item = pinned.getCallback(index)) { + process.accept(item); + } + } + if(activeTaskCounter.decrementAndGet() == 0) + TaskContext.defer(() -> future.complete(null)); + } + finally { + recordStatistics(context, globalFuture, taskStartTime); + recordEvent(context, taskStartTime); + } + }, future, context); + } + else + process.accept(callback); + + if(closed.get()) + throw new DMLRuntimeException("Race condition observed"); + } + catch(RuntimeException error) { + context.failAll(DMLRuntimeException.of(error)); + throw error; + } + finally { + recordStatistics(context, globalFuture, startTime); + } + }, context)); + } + + public static OOCFuture submitOOCTask(Runnable task, StreamContext context) { + // May be blocking tasks, thus should not run on default executor pool + ExecutorService pool = CommonThreadPool.getDynamicPool(); + OOCFuture future = new OOCFuture<>(); + COMPUTE_IN_FLIGHT.incrementAndGet(); + try { + pool.submit(task(() -> { + long startTime = DMLScript.STATISTICS || DMLScript.OOC_LOG_EVENTS ? System.nanoTime() : 0; + try { + task.run(); + future.complete(null); + context.clear(); + if(DMLScript.STATISTICS && context.getExtendedOpcode() != null) + Statistics.maintainOOCHeavyHitter(context.getExtendedOpcode(), System.nanoTime() - startTime); + recordEvent(context, startTime); + } + finally { + COMPUTE_IN_FLIGHT.decrementAndGet(); + } + }, future, context)); + } + catch(RuntimeException error) { + COMPUTE_IN_FLIGHT.decrementAndGet(); + throw DMLRuntimeException.of(error); + } + return future; + } + + private static void submit(Runnable runnable, OOCFuture future, StreamContext context) { + COMPUTE_IN_FLIGHT.incrementAndGet(); + try { + COMPUTE_EXECUTOR.submit(task(() -> { + try { + runnable.run(); + } + finally { + COMPUTE_IN_FLIGHT.decrementAndGet(); + } + }, future, context)); + } + catch(RuntimeException error) { + COMPUTE_IN_FLIGHT.decrementAndGet(); + throw error; + } + } + + private static Runnable task(Runnable runnable, OOCFuture future, StreamContext context) { + return () -> { + boolean setContext = TaskContext.getContext() == null; + if(setContext) + TaskContext.setContext(new TaskContext()); + long startTime = DMLScript.STATISTICS ? System.nanoTime() : 0; + try { + runnable.run(); + if(setContext) { + while(TaskContext.runDeferred()) { + } + } + } + catch(RuntimeException error) { + DMLRuntimeException failure = DMLRuntimeException.of(error); + context.failAll(failure); + if(future != null) + future.completeExceptionally(failure); + throw failure; + } + finally { + if(setContext) + TaskContext.clearContext(); + if(DMLScript.STATISTICS) + context.getLocalStatisticsLongAdder().add(System.nanoTime() - startTime); + } + }; + } + + private static Consumer> guard(Consumer> consumer, + StreamContext context) { + return callback -> { + try { + consumer.accept(callback); + } + catch(RuntimeException error) { + DMLRuntimeException failure = DMLRuntimeException.of(error); + context.failAll(failure); + throw failure; + } + }; + } + + private static void recordStatistics(StreamContext context, CompletableFuture globalFuture, long startTime) { + if(!DMLScript.STATISTICS) + return; + context.getLocalStatisticsLongAdder().add(System.nanoTime() - startTime); + if(globalFuture.isDone() && context.getExtendedOpcode() != null) { + Statistics.maintainOOCHeavyHitter(context.getExtendedOpcode(), context.getLocalStatisticsLongAdder().sum()); + context.getLocalStatisticsLongAdder().reset(); + } + } + + private static void recordEvent(StreamContext context, long startTime) { + if(DMLScript.OOC_LOG_EVENTS && context.getCallerId() != 0) + OOCEventLog.onComputeEvent(context.getCallerId(), startTime, System.nanoTime()); + } +} 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 new file mode 100644 index 00000000000..f4bab5d9813 --- /dev/null +++ b/src/test/java/org/apache/sysds/test/component/ooc/OOCInstructionUtilsTest.java @@ -0,0 +1,110 @@ +/* + * 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.test.component.ooc; + +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import org.apache.sysds.runtime.DMLRuntimeException; +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.ooc.cache.OOCFuture; +import org.apache.sysds.runtime.ooc.store.MaterializedCallback; +import org.apache.sysds.runtime.ooc.store.StoreLease; +import org.apache.sysds.runtime.ooc.stream.StreamContext; +import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; +import org.junit.Assert; +import org.junit.Test; + +public class OOCInstructionUtilsTest { + @Test + public void testSubmitTasksClosesCallbacksAfterCompletion() throws Exception { + SubscribableTaskQueue source = new SubscribableTaskQueue<>(); + AtomicInteger processed = new AtomicInteger(); + AtomicInteger released = new AtomicInteger(); + CompletableFuture completion = OOCInstructionUtils.submitOOCTasks(source, callback -> { + Assert.assertEquals(1, callback.get().getIndexes().getRowIndex()); + processed.incrementAndGet(); + }, 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.closeInput(); + completion.get(10, TimeUnit.SECONDS); + + Assert.assertEquals(1, processed.get()); + Assert.assertEquals(1, released.get()); + } + + @Test + public void testSubmitTasksWaitsForAllStreams() throws Exception { + SubscribableTaskQueue first = new SubscribableTaskQueue<>(); + SubscribableTaskQueue second = new SubscribableTaskQueue<>(); + AtomicInteger processed = new AtomicInteger(); + CompletableFuture completion = OOCInstructionUtils.submitOOCTasks(List.of(first, second), + (index, callback) -> processed.addAndGet(callback.get()), new StreamContext().addOutStream()); + + first.enqueue(1); + first.closeInput(); + second.enqueue(2); + Assert.assertFalse(completion.isDone()); + second.closeInput(); + completion.get(10, TimeUnit.SECONDS); + Assert.assertEquals(3, processed.get()); + } + + @Test + public void testSubmitTaskPropagatesFailure() throws Exception { + SubscribableTaskQueue output = new SubscribableTaskQueue<>(); + AtomicReference propagated = new AtomicReference<>(); + output.setSubscriber(callback -> { + try(callback) { + if(callback.isFailure()) { + try { + callback.get(); + } + catch(DMLRuntimeException failure) { + propagated.compareAndSet(null, failure); + } + } + } + }); + + OOCFuture completion = OOCInstructionUtils.submitOOCTask(() -> { + throw new DMLRuntimeException("injected failure"); + }, new StreamContext().addOutStream(output)); + try { + completion.get(10, TimeUnit.SECONDS); + Assert.fail("Expected task failure"); + } + catch(ExecutionException expected) { + Assert.assertTrue(expected.getCause() instanceof DMLRuntimeException); + } + output.closeInput(); + Assert.assertNotNull(propagated.get()); + Assert.assertEquals("injected failure", propagated.get().getMessage()); + } +}