diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/CountingLiveness.java b/src/main/java/org/apache/sysds/runtime/ooc/store/CountingLiveness.java new file mode 100644 index 00000000000..04e86f90a84 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/CountingLiveness.java @@ -0,0 +1,68 @@ +/* + * 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.AtomicIntegerArray; + +public final class CountingLiveness implements MaterializedStore.Liveness { + private final AtomicIntegerArray _remaining; + private final AtomicIntegerArray _reservable; + + public CountingLiveness(int size, int count) { + if(size < 0 || count < 0) + throw new IllegalArgumentException("Invalid args: size=" + size + ", count=" + count); + _remaining = new AtomicIntegerArray(size); + _reservable = new AtomicIntegerArray(size); + for(int i = 0; i < size; i++) { + _remaining.set(i, count); + _reservable.set(i, count); + } + } + + @Override + public boolean needs(int index) { + return index >= 0 && index < _remaining.length() && _remaining.get(index) > 0; + } + + @Override + public void consumed(int index) { + decrement(_remaining, index); + } + + @Override + public boolean reserve(int index) { + return index >= 0 && index < _reservable.length() && decrement(_reservable, index); + } + + @Override + public void unreserve(int index) { + _reservable.incrementAndGet(index); + } + + private static boolean decrement(AtomicIntegerArray counters, int index) { + while(true) { + int current = counters.get(index); + if(current <= 0) + return false; + if(counters.compareAndSet(index, current, current - 1)) + return true; + } + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedCallback.java b/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedCallback.java index 829e60b32ea..226dd5c2b68 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedCallback.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedCallback.java @@ -21,21 +21,21 @@ import org.apache.sysds.runtime.DMLRuntimeException; import org.apache.sysds.runtime.instructions.ooc.OOCStream; -import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue; import org.apache.sysds.runtime.ooc.cache.BlockEntry; +import org.apache.sysds.runtime.ooc.cache.io.SpillableObject; import java.util.concurrent.atomic.AtomicReference; -public final class MaterializedCallback implements OOCStream.QueueCallback { - private final StoreLease _lease; +public final class MaterializedCallback implements OOCStream.QueueCallback { + private final StoreLease _lease; private final AtomicReference _failure; private boolean _closed; - public MaterializedCallback(StoreLease lease) { + public MaterializedCallback(StoreLease lease) { this(lease, new AtomicReference<>()); } - private MaterializedCallback(StoreLease lease, AtomicReference failure) { + private MaterializedCallback(StoreLease lease, AtomicReference failure) { _lease = lease; _failure = failure; } @@ -45,7 +45,7 @@ public BlockEntry pinnedEntry() { } @Override - public IndexedMatrixValue get() { + public T get() { DMLRuntimeException failure = _failure.get(); if(failure != null) throw failure; @@ -53,10 +53,10 @@ public IndexedMatrixValue get() { } @Override - public synchronized OOCStream.QueueCallback keepOpen() { + public synchronized OOCStream.QueueCallback keepOpen() { if(_closed) throw new IllegalStateException("Cannot keep open a closed callback"); - return new MaterializedCallback(_lease.retain(), _failure); + return new MaterializedCallback<>(_lease.retain(), _failure); } @Override diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/SequentialAccessPattern.java b/src/main/java/org/apache/sysds/runtime/ooc/store/SequentialAccessPattern.java new file mode 100644 index 00000000000..18da6bc1452 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/SequentialAccessPattern.java @@ -0,0 +1,64 @@ +/* + * 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 org.apache.sysds.runtime.ooc.cache.collections.ConcurrentBitSet; + +public final class SequentialAccessPattern implements MaterializedStore.AccessPattern { + private final int _size; + private final ConcurrentBitSet _consumed; + private int _next; + private volatile int _consumedThrough; + + public SequentialAccessPattern(int size) { + if(size < 0) + throw new IllegalArgumentException("Size must not be negative: " + size); + _size = size; + _consumed = new ConcurrentBitSet(Math.max(1, size)); + _next = 0; + _consumedThrough = -1; + } + + @Override + public boolean hasNext() { + return _next < _size; + } + + @Override + public int next() { + if(!hasNext()) + throw new IllegalStateException("No remaining index"); + return _next++; + } + + @Override + public boolean needs(int index) { + return index >= 0 && index < _size && index > _consumedThrough && !_consumed.get(index); + } + + @Override + public synchronized void consumed(int index) { + if(index < 0 || index >= _size) + throw new IndexOutOfBoundsException("Invalid consumed index: " + index); + _consumed.set(index); + while(_consumedThrough + 1 < _size && _consumed.get(_consumedThrough + 1)) + _consumedThrough++; + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/StoreBackedStream.java b/src/main/java/org/apache/sysds/runtime/ooc/store/StoreBackedStream.java new file mode 100644 index 00000000000..7fad1105c5b --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/StoreBackedStream.java @@ -0,0 +1,185 @@ +/* + * 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.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.meta.DataCharacteristics; +import org.apache.sysds.runtime.ooc.cache.io.SpillableObject; + +public final class StoreBackedStream implements OOCStream { + private final OrderedMaterializedStoreReader _reader; + private volatile DMLRuntimeException _failure; + private boolean _subscriberSet; + private OOCStream.QueueCallback _lastDequeue; + private boolean _exhausted; + private CacheableData _data; + + public StoreBackedStream(OrderedMaterializedStoreReader reader) { + _reader = reader; + } + + @Override + public void enqueue(T value) { + throw new DMLRuntimeException("Cannot enqueue to a store-backed stream"); + } + + @Override + public void enqueue(QueueCallback callback) { + throw new DMLRuntimeException("Cannot enqueue to a store-backed stream"); + } + + @Override + public void closeInput() { + throw new DMLRuntimeException("Cannot close the input of a store-backed stream"); + } + + @Override + public synchronized T dequeue() { + QueueCallback callback = dequeueInternal(); + return callback == null ? null : callback.get(); + } + + @Override + public synchronized QueueCallback dequeueCB() { + return dequeueInternal(); + } + + private QueueCallback dequeueInternal() { + if(_subscriberSet) + throw new IllegalStateException("Cannot dequeue after setting a subscriber"); + if(_lastDequeue != null) { + _lastDequeue.close(); + _lastDequeue = null; + } + if(_failure != null) + throw _failure; + if(_exhausted) + return null; + try { + if(!_reader.hasNext()) { + _exhausted = true; + _reader.close(); + return null; + } + _lastDequeue = new MaterializedCallback<>(_reader.next()); + return _lastDequeue; + } + catch(InterruptedException e) { + Thread.currentThread().interrupt(); + throw recordFailure(e); + } + catch(RuntimeException e) { + throw recordFailure(e); + } + } + + @Override + public synchronized void setSubscriber(Consumer> subscriber) { + if(subscriber == null) + throw new IllegalArgumentException("Cannot set subscriber to null"); + if(_subscriberSet) + throw new IllegalStateException("Subscriber cannot be set multiple times"); + _subscriberSet = true; + Thread driver = new Thread(() -> drive(subscriber), "ooc-store-replay"); + driver.setDaemon(true); + driver.start(); + } + + private void drive(Consumer> subscriber) { + DMLRuntimeException failure = _failure; + if(failure != null) { + subscriber.accept(OOCStream.eos(failure)); + return; + } + try { + while(_reader.hasNext()) { + try(QueueCallback callback = new MaterializedCallback<>(_reader.next())) { + subscriber.accept(callback); + } + } + _reader.close(); + subscriber.accept(OOCStream.eos(null)); + } + catch(InterruptedException e) { + Thread.currentThread().interrupt(); + subscriber.accept(OOCStream.eos(recordFailure(e))); + } + catch(RuntimeException e) { + subscriber.accept(OOCStream.eos(recordFailure(e))); + } + } + + private synchronized DMLRuntimeException recordFailure(Exception error) { + if(_failure == null) + _failure = DMLRuntimeException.of(error); + _reader.close(); + return _failure; + } + + @Override + public void propagateFailure(DMLRuntimeException failure) { + recordFailure(failure); + } + + @Override + public OOCStream getReadStream() { + return this; + } + + @Override + public OOCStream getWriteStream() { + throw new UnsupportedOperationException("A store-backed stream has no write stream"); + } + + @Override + public boolean hasStreamCache() { + return false; + } + + @Override + public CachingStream getStreamCache() { + return null; + } + + @Override + public boolean isProcessed() { + return false; + } + + @Override + public DataCharacteristics getDataCharacteristics() { + return _data == null ? null : _data.getDataCharacteristics(); + } + + @Override + public CacheableData getData() { + return _data; + } + + @Override + public void setData(CacheableData data) { + _data = data; + } +} diff --git a/src/test/java/org/apache/sysds/test/component/ooc/MaterializedStoreTest.java b/src/test/java/org/apache/sysds/test/component/ooc/MaterializedStoreTest.java index d8ea0e3986f..14d26db3ae2 100644 --- a/src/test/java/org/apache/sysds/test/component/ooc/MaterializedStoreTest.java +++ b/src/test/java/org/apache/sysds/test/component/ooc/MaterializedStoreTest.java @@ -20,6 +20,7 @@ package org.apache.sysds.test.component.ooc; import java.util.List; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; @@ -37,8 +38,11 @@ import org.apache.sysds.runtime.ooc.memory.SyncMemoryAllowance; import org.apache.sysds.runtime.ooc.store.IndexedMaterializedStoreReader; import org.apache.sysds.runtime.ooc.store.MaterializedStore; +import org.apache.sysds.runtime.ooc.store.CountingLiveness; import org.apache.sysds.runtime.ooc.store.OOCStreamMaterializer; import org.apache.sysds.runtime.ooc.store.OrderedMaterializedStoreReader; +import org.apache.sysds.runtime.ooc.store.SequentialAccessPattern; +import org.apache.sysds.runtime.ooc.store.StoreBackedStream; import org.apache.sysds.runtime.ooc.store.StoreLease; import org.apache.sysds.test.component.ooc.cache.OOCCacheTestUtils; import org.junit.After; @@ -96,32 +100,16 @@ public void testMaterializationReadersAndForgetting() throws Exception { Assert.assertEquals(0, _materializerAllowance.getUsedMemory()); Assert.assertEquals(3, _store.size()); - MaterializedStore.AccessPattern pattern = sequentialPattern(3); - boolean[] indexedConsumed = new boolean[3]; - MaterializedStore.Liveness liveness = new MaterializedStore.Liveness() { - @Override - public boolean needs(int index) { - return !indexedConsumed[index]; - } - - @Override - public void consumed(int index) { - indexedConsumed[index] = true; - } - }; - - OrderedMaterializedStoreReader ordered = _store.openReader(pattern, _readerAllowance, 2); - IndexedMaterializedStoreReader indexed = _store.openIndexedReader(liveness); + StoreBackedStream ordered = new StoreBackedStream<>( + _store.openReader(new SequentialAccessPattern(3), _readerAllowance, 2, false)); + IndexedMaterializedStoreReader indexed = _store + .openIndexedReader(new CountingLiveness(3, 1)); _store.sealReaders(); int index = 0; - while(ordered.hasNext()) { - try(StoreLease lease = ordered.next()) { - Assert.assertEquals(index + 1L, lease.value().getIndexes().getRowIndex()); - index++; - } - } - ordered.close(); + IndexedMatrixValue value; + while((value = ordered.dequeue()) != null) + Assert.assertEquals(++index, value.getIndexes().getRowIndex()); Assert.assertEquals(3, index); Assert.assertTrue(_cache.getOwnedCacheSize() > 0); @@ -151,7 +139,7 @@ public void testOrderedReaderRetries() throws Exception { _readerAllowance.destroy(); _readerAllowance = new SyncMemoryAllowance(_broker, TILE_BYTES); _readerAllowance.setTargetMemory(TILE_BYTES); - OrderedMaterializedStoreReader reader = _store.openReader(sequentialPattern(2), + OrderedMaterializedStoreReader reader = _store.openReader(new SequentialAccessPattern(2), _readerAllowance, 2, false); _store.sealReaders(); @@ -185,7 +173,7 @@ public void testSoftOrderingReturnsReadyRequestFirst() throws Exception { _readerAllowance.setTargetMemory(largeBytes); long heldBytes = largeBytes - TILE_BYTES; _readerAllowance.reserveBlocking(heldBytes); - OrderedMaterializedStoreReader reader = _store.openReader(sequentialPattern(2), + OrderedMaterializedStoreReader reader = _store.openReader(new SequentialAccessPattern(2), _readerAllowance, 2); _store.sealReaders(); @@ -211,19 +199,8 @@ public void testDirectRequests() throws Exception { materializer.accept(OOCStream.eos(null)); materializer.completion().get(WAIT_SECONDS, TimeUnit.SECONDS); - boolean[] needed = {true}; IndexedMaterializedStoreReader reader = _store - .openIndexedReader(new MaterializedStore.Liveness() { - @Override - public boolean needs(int index) { - return needed[index]; - } - - @Override - public void consumed(int index) { - needed[index] = false; - } - }); + .openIndexedReader(new CountingLiveness(1, 1)); _store.sealReaders(); try(StoreLease published = _store.requestPublished(0, _readerAllowance).get(WAIT_SECONDS, @@ -292,6 +269,46 @@ public void testLiveCallbackKeepsPublicationPinned() throws Exception { Assert.assertEquals(1, eos.get()); } + @Test + public void testStoreBackedStreamSubscriber() throws Exception { + OOCStreamMaterializer materializer = new OOCStreamMaterializer(_store, + indexes -> (int) indexes.getRowIndex() - 1, _materializerAllowance); + for(int i = 0; i < 2; i++) + materializer.accept(new OOCStream.SimpleQueueCallback<>(tile(i, i + 1.0), null)); + materializer.accept(OOCStream.eos(null)); + materializer.completion().get(WAIT_SECONDS, TimeUnit.SECONDS); + + StoreBackedStream stream = new StoreBackedStream<>( + _store.openReader(new SequentialAccessPattern(2), _readerAllowance, 1)); + _store.sealReaders(); + AtomicInteger count = new AtomicInteger(); + CountDownLatch complete = new CountDownLatch(1); + stream.setSubscriber(callback -> { + if(callback.isEos()) + complete.countDown(); + else + Assert.assertEquals(count.incrementAndGet(), callback.get().getIndexes().getRowIndex()); + }); + + Assert.assertTrue(complete.await(WAIT_SECONDS, TimeUnit.SECONDS)); + Assert.assertEquals(2, count.get()); + Assert.assertEquals(0, _readerAllowance.getUsedMemory()); + } + + @Test + public void testCountingLiveness() { + CountingLiveness liveness = new CountingLiveness(1, 2); + Assert.assertTrue(liveness.reserve(0)); + Assert.assertTrue(liveness.reserve(0)); + Assert.assertFalse(liveness.reserve(0)); + liveness.unreserve(0); + Assert.assertTrue(liveness.reserve(0)); + liveness.consumed(0); + Assert.assertTrue(liveness.needs(0)); + liveness.consumed(0); + Assert.assertFalse(liveness.needs(0)); + } + @Test public void testFailurePropagation() throws Exception { DMLRuntimeException sourceFailure = new DMLRuntimeException("injected failure"); @@ -332,28 +349,4 @@ private static IndexedMatrixValue tile(int index, double value) { return new IndexedMatrixValue(new MatrixIndexes(index + 1L, 1), new MatrixBlock(4, 4, value)); } - private static MaterializedStore.AccessPattern sequentialPattern(int to) { - return new MaterializedStore.AccessPattern() { - private int _next; - - @Override - public boolean hasNext() { - return _next < to; - } - - @Override - public int next() { - return _next++; - } - - @Override - public boolean needs(int index) { - return true; - } - - @Override - public void consumed(int index) { - } - }; - } }