From 8225d8345e385fff265d3bea3ddd15fe07b2ebdb Mon Sep 17 00:00:00 2001 From: Jannik Lindemann Date: Mon, 6 Jul 2026 15:42:38 +0200 Subject: [PATCH 1/3] [OOC] Add New OOCPackedCache Implementation --- .../sysds/runtime/ooc/cache/OOCCache.java | 33 +- .../ooc/cache/io/SpillableObjectRegistry.java | 5 + .../ooc/cache/packed/OOCPackedCache.java | 761 ++++++++++++++++++ .../runtime/ooc/cache/packed/PackBuilder.java | 136 ++++ .../runtime/ooc/cache/packed/PackedBlock.java | 72 ++ .../ooc/cache/packed/PackedCacheLocation.java | 73 ++ .../ooc/cache/packed/PackedPinState.java | 234 ++++++ .../ooc/cache/packed/PackedUnpinHandle.java | 80 ++ .../component/ooc/cache/OOCCacheImplTest.java | 121 +-- .../ooc/cache/OOCCacheTestUtils.java | 130 +++ .../ooc/cache/OOCPackedCacheTest.java | 226 ++++++ 11 files changed, 1762 insertions(+), 109 deletions(-) create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackBuilder.java create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackedBlock.java create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackedCacheLocation.java create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackedPinState.java create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackedUnpinHandle.java create mode 100644 src/test/java/org/apache/sysds/test/component/ooc/cache/OOCCacheTestUtils.java create mode 100644 src/test/java/org/apache/sysds/test/component/ooc/cache/OOCPackedCacheTest.java diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCache.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCache.java index 7f1fdb493d0..aec88891595 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCache.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCache.java @@ -24,10 +24,29 @@ import java.util.function.LongUnaryOperator; public interface OOCCache { + /** + * Pins an item backed by an allowance. A successful pin transfers memory ownership from the cache to the owner of + * the allowance and guarantees data availability. While pinned, the bytes of the entry are not counted as + * cache-owned memory. + * + * @param key + * @param allowance + * @return a non-null future of the pinned block entry; the future result is null if the required memory could not + * be reserved + */ default OOCFuture pin(BlockKey key, MemoryAllowance allowance) { return pin(key.getStreamId(), key.getSequenceNumber(), allowance); } + /** + * Pins an item backed by an allowance. If the allowance cannot reserve enough memory, this method will wait until + * memory is available. A successful pin transfers memory ownership from the cache to the owner of the allowance and + * guarantees data availability. While pinned, the bytes of the entry are not counted as cache-owned memory. + * + * @param key + * @param allowance + * @return a non-null future of the pinned block entry + */ default OOCFuture pinAdmitted(BlockKey key, MemoryAllowance allowance) { return pinAdmitted(key.getStreamId(), key.getSequenceNumber(), allowance); } @@ -59,9 +78,17 @@ default BlockEntry putPinned(BlockKey key, Object data, long size, MemoryAllowan */ OOCFuture pin(long sId, long tId, MemoryAllowance allowance); - default OOCFuture pinAdmitted(long sId, long tId, MemoryAllowance allowance) { - return pin(sId, tId, allowance); - } + /** + * Pins an item backed by an allowance. If the allowance cannot reserve enough memory, this method will wait until + * memory is available. A successful pin transfers memory ownership from the cache to the owner of the allowance and + * guarantees data availability. While pinned, the bytes of the entry are not counted as cache-owned memory. + * + * @param sId + * @param tId + * @param allowance + * @return a non-null future of the pinned block entry + */ + OOCFuture pinAdmitted(long sId, long tId, MemoryAllowance allowance); /** * Pins an item backed by an allowance if it is already live in cache. A successful pin transfers memory ownership diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/io/SpillableObjectRegistry.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/io/SpillableObjectRegistry.java index 6cd8fded9e7..3b10a10ec54 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/cache/io/SpillableObjectRegistry.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/io/SpillableObjectRegistry.java @@ -20,6 +20,7 @@ package org.apache.sysds.runtime.ooc.cache.io; import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue; +import org.apache.sysds.runtime.ooc.cache.packed.PackedBlock; import java.io.DataInput; import java.io.DataOutput; @@ -27,6 +28,7 @@ public final class SpillableObjectRegistry { private static final byte INDEXED_MATRIX_VALUE = 1; + private static final byte PACKED_BLOCK = 2; private SpillableObjectRegistry() { } @@ -41,6 +43,7 @@ public static SpillableObject read(DataInput in) throws IOException { byte type = in.readByte(); SpillableObject obj = switch(type) { case INDEXED_MATRIX_VALUE -> new IndexedMatrixValue(); + case PACKED_BLOCK -> new PackedBlock(); default -> throw new IOException("Unknown spillable object type: " + type); }; obj.read(in); @@ -50,6 +53,8 @@ public static SpillableObject read(DataInput in) throws IOException { private static byte typeOf(SpillableObject obj) throws IOException { if(obj instanceof IndexedMatrixValue) return INDEXED_MATRIX_VALUE; + if(obj instanceof PackedBlock) + return PACKED_BLOCK; throw new IOException("Unsupported spillable object type: " + obj.getClass().getName()); } } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java new file mode 100644 index 00000000000..9bbdcbeae5a --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java @@ -0,0 +1,761 @@ +/* + * 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.cache.packed; + +import org.apache.sysds.runtime.instructions.ooc.CachingStream; +import org.apache.sysds.runtime.ooc.cache.BlockEntry; +import org.apache.sysds.runtime.ooc.cache.BlockKey; +import org.apache.sysds.runtime.ooc.cache.BlockState; +import org.apache.sysds.runtime.ooc.cache.OOCCache; +import org.apache.sysds.runtime.ooc.cache.OOCCacheImpl; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; +import org.apache.sysds.runtime.ooc.cache.collections.MaskedOnceArrayList; +import org.apache.sysds.runtime.ooc.cache.collections.SegmentedStreamTableList; +import org.apache.sysds.runtime.ooc.cache.io.OOCIOHandler; +import org.apache.sysds.runtime.ooc.memory.MemoryAllowance; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.LockSupport; +import java.util.function.LongUnaryOperator; + +public final class OOCPackedCache implements OOCCache { + private static final long PACKED_STREAM_ID = CachingStream._streamSeq.getNextID(); + private static final long DEFAULT_PACK_THRESHOLD_BYTES = 1L << 18; + private static final long DEFAULT_PACK_TARGET_BYTES = 1L << 19; // 512 KB tile packing + private static final long DEFAULT_MAX_STAGING_BYTES = 1L << 26; + private static final int DEFAULT_MAX_OPEN_BUILDERS = 64; + private static final long DEFAULT_SEAL_DELAY_MS = 5; + private static final long DEFAULT_PACK_RELEASE_DELAY_MS = 5; + + private final OOCCacheImpl _physical; + private final long _packThresholdBytes; + private final long _packTargetBytes; + private final long _maxStagingBytes; + private final int _maxOpenBuilders; + private final long _sealDelayMs; + private final long _packReleaseDelayMs; + private final SegmentedStreamTableList _locations; + private final MaskedOnceArrayList _packedStates; + private final ScheduledExecutorService _sealExecutor; + private final ExecutorService _releaseExecutor; + private final ConcurrentLinkedQueue _releaseQueue; + private final AtomicBoolean _releaseRunning; + private final AtomicBoolean _packedPolicyInstalled; + private final AtomicInteger _nextPackedId; + private final ArrayList> _logicalEvictionPolicies; + + private PackBuilder[] _builders; + private long _stagingBytes; + private int _openBuilderCount; + private boolean _running; + + public OOCPackedCache(OOCIOHandler ioHandler, long hardLimit, long evictionLimit) { + this(new OOCCacheImpl(ioHandler, hardLimit, evictionLimit), DEFAULT_PACK_THRESHOLD_BYTES, + DEFAULT_PACK_TARGET_BYTES, DEFAULT_SEAL_DELAY_MS); + } + + public OOCPackedCache(OOCCacheImpl physical) { + this(physical, DEFAULT_PACK_THRESHOLD_BYTES, DEFAULT_PACK_TARGET_BYTES, DEFAULT_SEAL_DELAY_MS); + } + + public OOCPackedCache(OOCCacheImpl physical, long packThresholdBytes, long packTargetBytes, long sealDelayMs) { + this(physical, packThresholdBytes, packTargetBytes, sealDelayMs, DEFAULT_PACK_RELEASE_DELAY_MS); + } + + public OOCPackedCache(OOCCacheImpl physical, long packThresholdBytes, long packTargetBytes, long sealDelayMs, + long packReleaseDelayMs) { + this(physical, packThresholdBytes, packTargetBytes, DEFAULT_MAX_STAGING_BYTES, DEFAULT_MAX_OPEN_BUILDERS, + sealDelayMs, packReleaseDelayMs); + } + + public OOCPackedCache(OOCCacheImpl physical, long packThresholdBytes, long packTargetBytes, long maxStagingBytes, + int maxOpenBuilders, long sealDelayMs, long packReleaseDelayMs) { + if(packThresholdBytes <= 0 || packTargetBytes < packThresholdBytes) + throw new IllegalArgumentException( + "Invalid pack sizes: threshold=" + packThresholdBytes + ", target=" + packTargetBytes); + _physical = physical; + _packThresholdBytes = packThresholdBytes; + _packTargetBytes = packTargetBytes; + _maxStagingBytes = Math.max(packTargetBytes, maxStagingBytes); + _maxOpenBuilders = Math.max(1, maxOpenBuilders); + _sealDelayMs = sealDelayMs; + _packReleaseDelayMs = packReleaseDelayMs; + _locations = new SegmentedStreamTableList<>(); + _packedStates = new MaskedOnceArrayList<>(); + _nextPackedId = new AtomicInteger(); + _releaseQueue = new ConcurrentLinkedQueue<>(); + _releaseRunning = new AtomicBoolean(false); + _packedPolicyInstalled = new AtomicBoolean(false); + _logicalEvictionPolicies = new ArrayList<>(); + _builders = new PackBuilder[16]; + _stagingBytes = 0; + _openBuilderCount = 0; + _running = true; + _sealExecutor = Executors.newSingleThreadScheduledExecutor(r -> { + Thread t = new Thread(r, "ooc-pack-sealer"); + t.setDaemon(true); + return t; + }); + _releaseExecutor = Executors.newSingleThreadExecutor(r -> { + Thread t = new Thread(r, "ooc-pack-release"); + t.setDaemon(true); + return t; + }); + } + + @Override + public BlockEntry putPinned(long sId, long tId, Object data, long size, MemoryAllowance allowance) { + if(size >= _packThresholdBytes) + return _physical.putPinned(sId, tId, data, size, allowance); + + PackBuilder builder; + int slot; + synchronized(this) { + checkRunning(); + builder = getOpenBuilder(sId, allowance, size); + slot = appendToBuilder(builder, sId, tId, data, size); + } + + BlockEntry logical = new BlockEntry(new BlockKey(sId, tId), size, data, BlockState.REMOVED); + logical.pin(); + logical.setCacheMeta(new PendingLogicalPin(builder, slot)); + return logical; + } + + public BlockEntry[] putPackPinned(long sId, long[] tIds, Object[] data, long[] sizes, int off, int len, + MemoryAllowance allowance) { + BlockEntry[] entries = new BlockEntry[len]; + synchronized(this) { + checkRunning(); + for(int i = 0; i < len; i++) { + int p = off + i; + long tId = tIds[p]; + long size = sizes[p]; + if(size >= _packThresholdBytes) { + entries[i] = _physical.putPinned(sId, tId, data[p], size, allowance); + continue; + } + PackBuilder builder = getOpenBuilder(sId, allowance, size); + int slot = appendToBuilder(builder, sId, tId, data[p], size); + BlockEntry logical = new BlockEntry(new BlockKey(sId, tId), size, data[p], BlockState.REMOVED); + logical.pin(); + logical.setCacheMeta(new PendingLogicalPin(builder, slot)); + entries[i] = logical; + } + } + return entries; + } + + public BlockEntry putSealedPackPinned(long sId, long[] tIds, Object[] data, long[] sizes, int off, int len, + MemoryAllowance allowance) { + long totalSize = 0; + Object[] packedData = new Object[len]; + long[] packedSizes = new long[len]; + for(int i = 0; i < len; i++) { + int p = off + i; + packedData[i] = data[p]; + packedSizes[i] = sizes[p]; + totalSize += sizes[p]; + } + + synchronized(this) { + checkRunning(); + BlockEntry physicalEntry = putSealedBlockPinned(new PackedBlock(packedData, packedSizes, totalSize), + allowance); + PackedPinState state = new PackedPinState(physicalEntry, sId, + Arrays.stream(tIds).mapToInt(Math::toIntExact).toArray(), off, len, len); + registerPackedState(state); + for(int i = 0; i < len; i++) + putLocation(new BlockKey(sId, tIds[off + i]), new SealedPackLocation(state, i)); + return physicalEntry; + } + } + + public PackGroup getPackGroup(long sId, long tId) { + PackedCacheLocation location = getLocation(sId, tId); + if(location instanceof PendingPackLocation pending) + location = forceSeal(pending); + return location instanceof SealedPackLocation packed ? packed.state().group : null; + } + + public int getPackGroupCount() { + return _nextPackedId.get(); + } + + public OOCFuture pinPack(PackGroup group, MemoryAllowance allowance) { + return group.state.pin(_physical, allowance, false) + .map(entry -> entry == null ? null : new PackLease(this, group, allowance)); + } + + @Override + public OOCFuture pin(long sId, long tId, MemoryAllowance allowance) { + PackedCacheLocation location = getLocation(sId, tId); + if(location == null) + return _physical.pin(sId, tId, allowance); + if(location instanceof PendingPackLocation pending) + location = forceSeal(pending); + if(!(location instanceof SealedPackLocation packed)) + return _physical.pin(sId, tId, allowance); + + return packed.state().pin(_physical, allowance, false).map(physicalEntry -> { + if(physicalEntry == null) + return null; + return createLogicalPin(new BlockKey(sId, tId), packed); + }); + } + + @Override + public OOCFuture pinAdmitted(long sId, long tId, MemoryAllowance allowance) { + PackedCacheLocation location = getLocation(sId, tId); + if(location == null) + return _physical.pinAdmitted(sId, tId, allowance); + if(location instanceof PendingPackLocation pending) + location = forceSeal(pending); + if(!(location instanceof SealedPackLocation packed)) + return _physical.pinAdmitted(sId, tId, allowance); + + return packed.state().pinAdmitted(_physical, allowance).map(physicalEntry -> { + if(physicalEntry == null) + return null; + return createLogicalPin(new BlockKey(sId, tId), packed); + }); + } + + @Override + public BlockEntry pinIfLive(long sId, long tId, MemoryAllowance allowance) { + PackedCacheLocation location = getLocation(sId, tId); + if(location == null) + return _physical.pinIfLive(sId, tId, allowance); + if(location instanceof PendingPackLocation pending) + location = forceSeal(pending); + if(!(location instanceof SealedPackLocation packed)) + return _physical.pinIfLive(sId, tId, allowance); + + if(packed.state().pinIfLive(_physical, allowance) == null) + return null; + return createLogicalPin(new BlockKey(sId, tId), packed); + } + + @Override + public UnpinHandle unpin(BlockEntry entry, MemoryAllowance allowance) { + Object meta = entry.getCacheMeta(); + if(meta instanceof PendingLogicalPin pending) + return unpinPending(entry, pending, allowance); + if(meta instanceof PackedLogicalPin packed) + return unpinPacked(entry, packed, allowance); + return _physical.unpin(entry, allowance); + } + + @Override + public int reference(BlockEntry entry) { + Object meta = entry.getCacheMeta(); + if(meta instanceof PackedLogicalPin packed) + return packed.location().retain(); + if(meta instanceof PendingLogicalPin pending) + return referencePending(pending.builder(), pending.slot()); + return _physical.reference(entry); + } + + @Override + public int dereference(BlockEntry entry) { + Object meta = entry.getCacheMeta(); + if(meta instanceof PackedLogicalPin packed) + return releaseLocation(entry.getKey(), packed.location()); + if(meta instanceof PendingLogicalPin pending) + return dereferencePending(pending.builder(), pending.slot()); + return _physical.dereference(entry); + } + + @Override + public int dereference(BlockKey key) { + PackedCacheLocation location = getLocation(key.getStreamId(), key.getSequenceNumber()); + if(location == null) + return _physical.dereference(key); + if(location instanceof PendingPackLocation pending) + return dereferencePending(pending.builder(), pending.slot()); + if(!(location instanceof SealedPackLocation packed)) + return _physical.dereference(key); + return releaseLocation(key, packed); + } + + /** + * References/dereferences on tiles in open builders are counted on the builder slot instead of forcing a seal, so + * pipelined consumers that park references (state tables, store readers) do not fragment packs into per-tile + * physical entries. Slot counts carry over into the SealedPackLocation at seal time. Only physical access (pin, + * pack group) forces a seal. + */ + private synchronized int referencePending(PackBuilder builder, int slot) { + if(!builder.sealed) + return builder.retainSlot(slot); + PackedCacheLocation location = getLocation(builder.streamIds[slot], builder.tileIds[slot]); + if(!(location instanceof SealedPackLocation packed)) + throw new IllegalStateException("Cannot retain a forgotten packed location."); + return packed.retain(); + } + + private synchronized int dereferencePending(PackBuilder builder, int slot) { + BlockKey key = new BlockKey(builder.streamIds[slot], builder.tileIds[slot]); + if(builder.sealed) { + PackedCacheLocation location = getLocation(key.getStreamId(), key.getSequenceNumber()); + return location instanceof SealedPackLocation packed ? releaseLocation(key, packed) : 0; + } + int references = builder.releaseSlot(slot); + if(references == 0) + clearLocation(key); + return references; + } + + @Override + public void updateLimits(long hardLimit, long evictionLimit) { + _physical.updateLimits(hardLimit, evictionLimit); + } + + @Override + public void addEvictionPolicy(long streamId, LongUnaryOperator scoreFn) { + _physical.addEvictionPolicy(streamId, scoreFn); + addLogicalEvictionPolicy(streamId, scoreFn); + if(_packedPolicyInstalled.compareAndSet(false, true)) + _physical.addEvictionPolicy(PACKED_STREAM_ID, this::scorePackedBlock); + } + + @Override + public long getOwnedCacheSize() { + return _physical.getOwnedCacheSize(); + } + + @Override + public synchronized void shutdown() { + if(!_running) + return; + _running = false; + for(PackBuilder builder : _builders) + if(builder != null) + sealBuilder(builder); + _physical.updateLimits(Long.MAX_VALUE, Long.MAX_VALUE); + _sealExecutor.shutdownNow(); + _releaseExecutor.shutdown(); + awaitReleaseExecutor(); + drainReleaseQueue(); + _physical.shutdown(); + } + + private void awaitReleaseExecutor() { + try { + _releaseExecutor.awaitTermination(Math.max(100, _packReleaseDelayMs * 2), TimeUnit.MILLISECONDS); + } + catch(InterruptedException ex) { + Thread.currentThread().interrupt(); + } + } + + private void drainReleaseQueue() { + PackedPinState state; + while((state = _releaseQueue.poll()) != null) { + state.clearReleaseQueued(); + state.releaseDuePins(_physical, Long.MAX_VALUE); + } + } + + public synchronized void flushPacks() { + for(PackBuilder builder : _builders) + if(builder != null) + sealBuilder(builder); + } + + private UnpinHandle unpinPending(BlockEntry entry, PendingLogicalPin pin, MemoryAllowance allowance) { + if(entry.fastUnpin()) { + allowance.release(entry.getSize()); + return PackedUnpinHandle.committed(entry, allowance, entry.getSize()); + } + synchronized(this) { + if(entry.getPinCount() > 1) { + entry.unpin(); + allowance.release(entry.getSize()); + return PackedUnpinHandle.committed(entry, allowance, entry.getSize()); + } + entry.unpin(); + entry.setCacheMeta(null); + PackedUnpinHandle handle = pin.builder().unpinProducer(entry, pin.slot(), allowance); + if(pin.builder().sealed && pin.builder().activePins == 0) + pin.builder().transferProducerOwnership(_physical); + scheduleSeal(pin.builder()); + return handle; + } + } + + private UnpinHandle unpinPacked(BlockEntry entry, PackedLogicalPin pin, MemoryAllowance allowance) { + if(entry.fastUnpin()) + return PackedUnpinHandle.committed(entry, allowance, entry.getSize()); + if(entry.getPinCount() > 1) { + entry.unpin(); + return PackedUnpinHandle.committed(entry, allowance, entry.getSize()); + } + entry.unpin(); + entry.setCacheMeta(null); + return pin.location().state().unpin(this, _packReleaseDelayMs, allowance); + } + + void enqueueRelease(PackedPinState state) { + if(!_running) + return; + if(state.markReleaseQueued()) { + _releaseQueue.offer(state); + scheduleReleaseMaintenance(); + } + } + + private void enqueueReleaseNoSchedule(PackedPinState state) { + if(state.markReleaseQueued()) + _releaseQueue.offer(state); + } + + private void scheduleReleaseMaintenance() { + if(!_releaseRunning.compareAndSet(false, true)) + return; + _releaseExecutor.execute(this::runReleaseMaintenance); + } + + private void runReleaseMaintenance() { + try { + while(_running) { + long nextDueNanos = Long.MAX_VALUE; + ArrayList delayed = null; + PackedPinState state; + long nowNanos = System.nanoTime(); + while((state = _releaseQueue.poll()) != null) { + state.clearReleaseQueued(); + long stateNextDue = state.releaseDuePins(_physical, nowNanos); + if(stateNextDue != Long.MAX_VALUE) { + if(delayed == null) + delayed = new ArrayList<>(); + delayed.add(state); + nextDueNanos = Math.min(nextDueNanos, stateNextDue); + } + } + if(delayed != null) + for(PackedPinState delayedState : delayed) + enqueueReleaseNoSchedule(delayedState); + if(nextDueNanos == Long.MAX_VALUE) + return; + long waitNanos = nextDueNanos - System.nanoTime(); + if(waitNanos > 0) + LockSupport.parkNanos(waitNanos); + } + } + finally { + _releaseRunning.set(false); + if(_running && !_releaseQueue.isEmpty()) + scheduleReleaseMaintenance(); + } + } + + private SealedPackLocation forceSeal(PendingPackLocation pending) { + synchronized(this) { + sealBuilder(pending.builder()); + PackedCacheLocation location = getLocation(pending.builder().streamIds[pending.slot()], + pending.builder().tileIds[pending.slot()]); + return (SealedPackLocation) location; + } + } + + private static BlockEntry createLogicalPin(BlockKey logicalKey, SealedPackLocation location) { + PackedBlock block = (PackedBlock) location.state().physicalEntry.getDataUnsafe(); + Object data = block.values[location.slot()]; + long size = block.sizes[location.slot()]; + BlockEntry logical = new BlockEntry(logicalKey, size, data, BlockState.REMOVED); + logical.pin(); + logical.setCacheMeta(new PackedLogicalPin(location)); + return logical; + } + + private int releaseLocation(BlockKey key, SealedPackLocation location) { + int references = location.release(); + if(references > 0) + return references; + if(!clearLocation(key)) + return 0; + if(location.state().forgetLocation()) { + _packedStates.clear((int) location.state().physicalEntry.getKey().getSequenceNumber()); + return _physical.dereference(location.state().physicalEntry); + } + return 0; + } + + private PackBuilder getOpenBuilder(long streamId, MemoryAllowance allowance, long nextSize) { + int sid = (int) streamId; + PackBuilder builder = sid < _builders.length ? _builders[sid] : null; + if(builder != null && (builder.sealed || builder.allowance != allowance)) { + sealBuilder(builder); + builder = null; + } + if(builder != null) + return builder; + while(!canOpenBuilder(nextSize)) { + PackBuilder largest = findLargestOpenBuilder(); + if(largest == null) + break; + sealBuilder(largest); + } + ensureBuilderCapacity(sid); + builder = new PackBuilder(sid, allowance, _packTargetBytes); + _builders[sid] = builder; + _openBuilderCount++; + return builder; + } + + private boolean canOpenBuilder(long nextSize) { + return _openBuilderCount < _maxOpenBuilders && _stagingBytes + nextSize <= _maxStagingBytes; + } + + private int appendToBuilder(PackBuilder builder, long streamId, long tileId, Object data, long size) { + int slot = builder.append(streamId, tileId, data, size); + _stagingBytes += size; + putLocation(new BlockKey(streamId, tileId), new PendingPackLocation(builder, slot)); + if(builder.getBytes() >= builder.packTargetBytes) + sealBuilder(builder); + else + enforceStagingBudget(); + return slot; + } + + private void enforceStagingBudget() { + while(_stagingBytes > _maxStagingBytes || _openBuilderCount > _maxOpenBuilders) { + PackBuilder builder = findLargestOpenBuilder(); + if(builder == null) + return; + sealBuilder(builder); + } + } + + private PackBuilder findLargestOpenBuilder() { + PackBuilder largest = null; + for(PackBuilder builder : _builders) + if(builder != null && !builder.sealed && (largest == null || builder.getBytes() > largest.getBytes())) + largest = builder; + return largest; + } + + private void sealBuilder(PackBuilder builder) { + if(builder.sealed || builder.count == 0) + return; + builder.sealed = true; + _stagingBytes -= builder.getBytes(); + _openBuilderCount--; + if(builder.streamSlot >= 0 && builder.streamSlot < _builders.length && _builders[builder.streamSlot] == builder) + _builders[builder.streamSlot] = null; + + PackedBlock block = builder.createBlock(); + BlockEntry physicalEntry = putSealedBlockPinned(block, builder.allowance); + int liveSlots = builder.countLiveSlots(); + PackedPinState state = new PackedPinState(physicalEntry, builder.streamIds[0], + Arrays.stream(builder.tileIds).mapToInt(Math::toIntExact).toArray(), 0, builder.count, liveSlots); + builder.state = state; + if(liveSlots > 0) + registerPackedState(state); + + // slots forgotten while pending stay in the physical pack but get no location + for(int i = 0; i < builder.count; i++) + if(builder.refCounts[i] > 0) + putLocation(new BlockKey(builder.streamIds[i], builder.tileIds[i]), + new SealedPackLocation(state, i, builder.refCounts[i])); + + if(builder.activePins == 0) + builder.transferProducerOwnership(_physical); + if(liveSlots == 0) + _physical.dereference(physicalEntry); + } + + private BlockEntry putSealedBlockPinned(PackedBlock block, MemoryAllowance allowance) { + BlockKey packedKey = new BlockKey(PACKED_STREAM_ID, _nextPackedId.getAndIncrement()); + return _physical.putPinned(packedKey, block, block.totalSize, allowance); + } + + private void registerPackedState(PackedPinState state) { + _packedStates.put((int) state.physicalEntry.getKey().getSequenceNumber(), state); + } + + private long scorePackedBlock(long packId) { + PackedPinState state = _packedStates.get((int) packId); + if(state == null) + return packId; + PackGroup group = state.group; + CopyOnWriteArrayList policies = group.streamId < _logicalEvictionPolicies + .size() ? _logicalEvictionPolicies.get((int) group.streamId) : null; + if(policies == null || policies.isEmpty()) + return packId; + long score = Long.MAX_VALUE; + for(int i = 0; i < group.size(); i++) { + long tileId = group.index(i); + for(LongUnaryOperator policy : policies) + score = Math.min(score, policy.applyAsLong(tileId)); + } + return score; + } + + private synchronized void addLogicalEvictionPolicy(long streamId, LongUnaryOperator scoreFn) { + int sid = (int) streamId; + while(sid >= _logicalEvictionPolicies.size()) + _logicalEvictionPolicies.add(null); + CopyOnWriteArrayList policies = _logicalEvictionPolicies.get(sid); + if(policies == null) { + policies = new CopyOnWriteArrayList<>(); + _logicalEvictionPolicies.set(sid, policies); + } + policies.add(scoreFn); + } + + private void scheduleSeal(PackBuilder builder) { + if(builder.sealScheduled || builder.sealed || _sealDelayMs < 0) + return; + builder.sealScheduled = true; + _sealExecutor.schedule(() -> { + synchronized(OOCPackedCache.this) { + builder.sealScheduled = false; + sealBuilder(builder); + } + }, _sealDelayMs, TimeUnit.MILLISECONDS); + } + + private PackedCacheLocation getLocation(long sId, long tId) { + MaskedOnceArrayList stream = _locations.get(sId); + return stream == null ? null : stream.get((int) tId); + } + + private void putLocation(BlockKey key, PackedCacheLocation location) { + _locations.getOrCreate(key.getStreamId()).put((int) key.getSequenceNumber(), location); + } + + private boolean clearLocation(BlockKey key) { + MaskedOnceArrayList stream = _locations.get(key.getStreamId()); + return stream != null && stream.clear((int) key.getSequenceNumber()); + } + + private void ensureBuilderCapacity(int streamId) { + if(streamId < _builders.length) + return; + int len = _builders.length; + while(streamId >= len) + len <<= 1; + PackBuilder[] bigger = new PackBuilder[len]; + System.arraycopy(_builders, 0, bigger, 0, _builders.length); + _builders = bigger; + } + + private void checkRunning() { + if(!_running) + throw new IllegalStateException("Cache has been shut down."); + } + + public static final class PackGroup { + private final PackedPinState state; + private final long streamId; + private final int firstIndex; + private final int[] indices; + private final int size; + + PackGroup(PackedPinState state, long streamId, int[] tileIds, int off, int size) { + this.state = state; + this.streamId = streamId; + this.size = size; + firstIndex = tileIds[off]; + boolean contiguous = true; + for(int i = 1; i < size; i++) { + if(tileIds[off + i] != firstIndex + i) { + contiguous = false; + break; + } + } + if(contiguous) + indices = null; + else { + indices = new int[size]; + System.arraycopy(tileIds, off, indices, 0, size); + } + } + + public int id() { + return (int) state.physicalEntry.getKey().getSequenceNumber(); + } + + public long streamId() { + return streamId; + } + + public int size() { + return size; + } + + public int index(int slot) { + if(slot < 0 || slot >= size) + throw new IndexOutOfBoundsException("Invalid pack slot: " + slot); + return indices == null ? firstIndex + slot : indices[slot]; + } + } + + public static final class PackLease implements AutoCloseable { + private final OOCPackedCache owner; + private final PackGroup group; + private final MemoryAllowance allowance; + private boolean open; + + private PackLease(OOCPackedCache owner, PackGroup group, MemoryAllowance allowance) { + this.owner = owner; + this.group = group; + this.allowance = allowance; + open = true; + } + + public PackGroup group() { + return group; + } + + public int size() { + return group.size(); + } + + public int index(int slot) { + return group.index(slot); + } + + public Object value(int slot) { + if(!open) + throw new IllegalStateException("Pack lease is closed"); + PackedBlock block = (PackedBlock) group.state.physicalEntry.getData(); + return block.values[slot]; + } + + @Override + public void close() { + if(!open) + return; + open = false; + group.state.unpin(owner, owner._packReleaseDelayMs, allowance); + } + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackBuilder.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackBuilder.java new file mode 100644 index 00000000000..9a5f998e17d --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackBuilder.java @@ -0,0 +1,136 @@ +/* + * 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.cache.packed; + +import org.apache.sysds.runtime.ooc.cache.BlockEntry; +import org.apache.sysds.runtime.ooc.cache.OOCCache; +import org.apache.sysds.runtime.ooc.cache.OOCCacheImpl; +import org.apache.sysds.runtime.ooc.memory.MemoryAllowance; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +final class PackBuilder { + final int streamSlot; + final MemoryAllowance allowance; + final long packTargetBytes; + final List deferredUnpins = new ArrayList<>(); + long[] streamIds = new long[16]; + long[] tileIds = new long[16]; + private Object[] values = new Object[16]; + long[] sizes = new long[16]; + int[] refCounts = new int[16]; + long bytes; + int count; + int activePins; + boolean sealed; + boolean sealScheduled; + PackedPinState state; + private boolean _producerTransferred; + + PackBuilder(int streamSlot, MemoryAllowance allowance, long packTargetBytes) { + this.streamSlot = streamSlot; + this.allowance = allowance; + this.packTargetBytes = packTargetBytes; + } + + int append(long streamId, long tileId, Object value, long size) { + ensureCapacity(count + 1); + int slot = count++; + streamIds[slot] = streamId; + tileIds[slot] = tileId; + values[slot] = value; + sizes[slot] = size; + refCounts[slot] = 1; + bytes += size; + activePins++; + return slot; + } + + int retainSlot(int slot) { + int references = refCounts[slot]; + if(references <= 0) + throw new IllegalStateException("Cannot retain a forgotten packed location."); + return refCounts[slot] = references + 1; + } + + int releaseSlot(int slot) { + int references = refCounts[slot]; + if(references <= 0) + return 0; + return refCounts[slot] = references - 1; + } + + int countLiveSlots() { + int live = 0; + for(int i = 0; i < count; i++) + if(refCounts[i] > 0) + live++; + return live; + } + + long getBytes() { + return bytes; + } + + PackedBlock createBlock() { + return new PackedBlock(Arrays.copyOf(values, count), Arrays.copyOf(sizes, count), bytes); + } + + PackedUnpinHandle unpinProducer(BlockEntry entry, int slot, MemoryAllowance owner) { + activePins--; + PackedUnpinHandle handle = PackedUnpinHandle.pendingProducerTransfer(entry, owner, sizes[slot]); + deferredUnpins.add(handle); + return handle; + } + + void transferProducerOwnership(OOCCacheImpl physical) { + if(state == null || physical == null || _producerTransferred) + return; + _producerTransferred = true; + OOCCache.UnpinHandle physicalUnpin = physical.unpin(state.physicalEntry, allowance); + if(physicalUnpin.isCommitted()) { + completeDeferredUnpins(true); + return; + } + physicalUnpin.getCompletionFuture() + .whenComplete((committed, ex) -> completeDeferredUnpins(ex == null && committed)); + } + + private void ensureCapacity(int minSize) { + if(minSize <= values.length) + return; + int len = values.length; + while(minSize > len) + len <<= 1; + streamIds = Arrays.copyOf(streamIds, len); + tileIds = Arrays.copyOf(tileIds, len); + values = Arrays.copyOf(values, len); + sizes = Arrays.copyOf(sizes, len); + refCounts = Arrays.copyOf(refCounts, len); + } + + private void completeDeferredUnpins(boolean committed) { + for(PackedUnpinHandle handle : deferredUnpins) + handle.complete(committed); + deferredUnpins.clear(); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackedBlock.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackedBlock.java new file mode 100644 index 00000000000..5727e2eea63 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackedBlock.java @@ -0,0 +1,72 @@ +/* + * 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.cache.packed; + +import org.apache.sysds.runtime.ooc.cache.io.SpillableObject; +import org.apache.sysds.runtime.ooc.cache.io.SpillableObjectRegistry; + +import java.io.DataInput; +import java.io.DataOutput; +import java.io.IOException; + +public final class PackedBlock implements SpillableObject { + Object[] values; + long[] sizes; + long totalSize; + + public PackedBlock() { + values = null; + sizes = null; + totalSize = 0; + } + + PackedBlock(Object[] values, long[] sizes, long totalSize) { + this.values = values; + this.sizes = sizes; + this.totalSize = totalSize; + } + + @Override + public boolean tryWrite(DataOutput out) throws IOException { + out.writeInt(values.length); + for(int i = 0; i < values.length; i++) { + out.writeLong(sizes[i]); + Object value = values[i]; + if(!(value instanceof SpillableObject spillable)) + return false; + if(!SpillableObjectRegistry.tryWrite(out, spillable)) + return false; + } + return true; + } + + @Override + public void read(DataInput in) throws IOException { + int count = in.readInt(); + values = new Object[count]; + sizes = new long[count]; + totalSize = 0; + for(int i = 0; i < count; i++) { + sizes[i] = in.readLong(); + values[i] = SpillableObjectRegistry.read(in); + totalSize += sizes[i]; + } + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackedCacheLocation.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackedCacheLocation.java new file mode 100644 index 00000000000..20f429cb0bd --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackedCacheLocation.java @@ -0,0 +1,73 @@ +/* + * 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.cache.packed; + +interface PackedCacheLocation { +} + +record PendingPackLocation(PackBuilder builder, int slot) implements PackedCacheLocation { +} + +final class SealedPackLocation implements PackedCacheLocation { + private final PackedPinState _state; + private final int _slot; + private int _references; + + SealedPackLocation(PackedPinState state, int slot) { + this(state, slot, 1); + } + + SealedPackLocation(PackedPinState state, int slot, int references) { + if(references <= 0) + throw new IllegalArgumentException("Sealed location requires a positive reference count."); + _state = state; + _slot = slot; + _references = references; + } + + PackedPinState state() { + return _state; + } + + int slot() { + return _slot; + } + + synchronized int retain() { + if(_references <= 0) + throw new IllegalStateException("Cannot retain a forgotten packed location."); + return ++_references; + } + + synchronized int release() { + if(_references <= 0) { + // tolerated for legacy double-forget callers; assertion surfaces it in debug runs + assert false : "Packed location slot " + _slot + " dereferenced below zero."; + return 0; + } + return --_references; + } +} + +record PendingLogicalPin(PackBuilder builder, int slot) { +} + +record PackedLogicalPin(SealedPackLocation location) { +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackedPinState.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackedPinState.java new file mode 100644 index 00000000000..43b94c8e63e --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackedPinState.java @@ -0,0 +1,234 @@ +/* + * 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.cache.packed; + +import org.apache.sysds.runtime.ooc.memory.MemoryAllowance; +import org.apache.sysds.runtime.ooc.cache.BlockEntry; +import org.apache.sysds.runtime.ooc.cache.OOCCache; +import org.apache.sysds.runtime.ooc.cache.OOCCacheImpl; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; + +import java.util.ArrayList; +import java.util.concurrent.TimeUnit; + +final class PackedPinState { + final BlockEntry physicalEntry; + final OOCPackedCache.PackGroup group; + private MemoryAllowance[] _allowances; + private int[] _counts; + private OOCFuture[] _futures; + private long[] _releaseDueNanos; + private PackedUnpinHandle[] _releaseHandles; + private int _size; + private boolean _releaseQueued; + private int _liveLocations; + + @SuppressWarnings("unchecked") + PackedPinState(BlockEntry physicalEntry, long streamId, int[] tileIds, int off, int count, int liveLocations) { + this.physicalEntry = physicalEntry; + group = new OOCPackedCache.PackGroup(this, streamId, tileIds, off, count); + this._liveLocations = liveLocations; + _allowances = new MemoryAllowance[2]; + _counts = new int[2]; + _futures = new OOCFuture[2]; + _releaseDueNanos = new long[2]; + _releaseHandles = new PackedUnpinHandle[2]; + } + + synchronized OOCFuture pin(OOCCacheImpl physical, MemoryAllowance allowance, boolean liveOnly) { + int ix = indexOf(allowance); + if(ix >= 0) { + cancelRelease(ix); + _counts[ix]++; + return _futures[ix]; + } + OOCFuture future = liveOnly ? OOCFuture.completed( + physical.pinIfLive(physicalEntry.getKey().getStreamId(), physicalEntry.getKey().getSequenceNumber(), + allowance)) : physical.pin(physicalEntry.getKey(), allowance); + addAllowance(allowance, future); + future.whenComplete((entry, ex) -> { + if(entry == null || ex != null) + removeFailedAllowance(allowance, future); + }); + return future; + } + + synchronized OOCFuture pinAdmitted(OOCCacheImpl physical, MemoryAllowance allowance) { + int ix = indexOf(allowance); + if(ix >= 0) { + cancelRelease(ix); + _counts[ix]++; + return _futures[ix]; + } + OOCFuture future = physical.pinAdmitted(physicalEntry.getKey(), allowance); + addAllowance(allowance, future); + future.whenComplete((entry, ex) -> { + if(entry == null || ex != null) + removeFailedAllowance(allowance, future); + }); + return future; + } + + BlockEntry pinIfLive(OOCCacheImpl physical, MemoryAllowance allowance) { + try { + return pin(physical, allowance, true).getNow(null); + } + catch(RuntimeException ex) { + return null; + } + } + + synchronized OOCCache.UnpinHandle unpin(OOCPackedCache owner, long releaseDelayMs, MemoryAllowance allowance) { + int ix = indexOf(allowance); + if(ix < 0) + return PackedUnpinHandle.committed(physicalEntry, allowance, physicalEntry.getSize()); + _counts[ix]--; + if(_counts[ix] > 0) + return PackedUnpinHandle.committed(physicalEntry, allowance, physicalEntry.getSize()); + PackedUnpinHandle handle = PackedUnpinHandle.delayedPhysicalRelease(physicalEntry, allowance); + _releaseHandles[ix] = handle; + _releaseDueNanos[ix] = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(Math.max(0, releaseDelayMs)); + owner.enqueueRelease(this); + return handle; + } + + long releaseDuePins(OOCCacheImpl physical, long nowNanos) { + ArrayList due = null; + long nextDueNanos = Long.MAX_VALUE; + synchronized(this) { + for(int i = 0; i < _size;) { + PackedUnpinHandle handle = _releaseHandles[i]; + if(handle == null || _counts[i] > 0) { + i++; + continue; + } + long dueNanos = _releaseDueNanos[i]; + if(dueNanos > nowNanos) { + nextDueNanos = Math.min(nextDueNanos, dueNanos); + i++; + continue; + } + if(due == null) + due = new ArrayList<>(); + due.add(new PackedRelease(_allowances[i], handle)); + removeAt(i); + } + } + if(due != null) + for(PackedRelease release : due) + releasePhysicalPin(physical, release.allowance, release.handle); + return nextDueNanos; + } + + synchronized boolean markReleaseQueued() { + if(_releaseQueued) + return false; + _releaseQueued = true; + return true; + } + + synchronized void clearReleaseQueued() { + _releaseQueued = false; + } + + synchronized boolean forgetLocation() { + if(_liveLocations <= 0) + return false; + return --_liveLocations == 0; + } + + private int indexOf(MemoryAllowance allowance) { + for(int i = 0; i < _size; i++) + if(_allowances[i] == allowance) + return i; + return -1; + } + + private synchronized void addAllowance(MemoryAllowance allowance, OOCFuture future) { + if(_size == _allowances.length) + grow(); + _allowances[_size] = allowance; + _counts[_size] = 1; + _futures[_size] = future; + _size++; + } + + @SuppressWarnings("unchecked") + private void grow() { + int nextSize = _size * 2; + MemoryAllowance[] biggerAllowances = new MemoryAllowance[nextSize]; + int[] biggerCounts = new int[nextSize]; + OOCFuture[] biggerFutures = new OOCFuture[nextSize]; + long[] biggerReleaseDueNanos = new long[nextSize]; + PackedUnpinHandle[] biggerReleaseHandles = new PackedUnpinHandle[nextSize]; + System.arraycopy(_allowances, 0, biggerAllowances, 0, _size); + System.arraycopy(_counts, 0, biggerCounts, 0, _size); + System.arraycopy(_futures, 0, biggerFutures, 0, _size); + System.arraycopy(_releaseDueNanos, 0, biggerReleaseDueNanos, 0, _size); + System.arraycopy(_releaseHandles, 0, biggerReleaseHandles, 0, _size); + _allowances = biggerAllowances; + _counts = biggerCounts; + _futures = biggerFutures; + _releaseDueNanos = biggerReleaseDueNanos; + _releaseHandles = biggerReleaseHandles; + } + + private void cancelRelease(int ix) { + _releaseDueNanos[ix] = 0; + PackedUnpinHandle handle = _releaseHandles[ix]; + if(handle != null) { + _releaseHandles[ix] = null; + handle.complete(false); + } + } + + private void releasePhysicalPin(OOCCacheImpl physical, MemoryAllowance allowance, PackedUnpinHandle handle) { + OOCCache.UnpinHandle physicalHandle = physical.unpin(physicalEntry, allowance); + if(physicalHandle.isCommitted()) { + handle.complete(true); + return; + } + physicalHandle.getCompletionFuture() + .whenComplete((committed, ex) -> handle.complete(ex == null && Boolean.TRUE.equals(committed))); + } + + private synchronized void removeFailedAllowance(MemoryAllowance allowance, OOCFuture future) { + int ix = indexOf(allowance); + if(ix >= 0 && _futures[ix] == future) + removeAt(ix); + } + + private void removeAt(int ix) { + int last = --_size; + _allowances[ix] = _allowances[last]; + _counts[ix] = _counts[last]; + _futures[ix] = _futures[last]; + _releaseDueNanos[ix] = _releaseDueNanos[last]; + _releaseHandles[ix] = _releaseHandles[last]; + _allowances[last] = null; + _counts[last] = 0; + _futures[last] = null; + _releaseDueNanos[last] = 0; + _releaseHandles[last] = null; + } + + private record PackedRelease(MemoryAllowance allowance, PackedUnpinHandle handle) { + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackedUnpinHandle.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackedUnpinHandle.java new file mode 100644 index 00000000000..cb6037dc1fe --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackedUnpinHandle.java @@ -0,0 +1,80 @@ +/* + * 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.cache.packed; + +import org.apache.sysds.runtime.ooc.cache.BlockEntry; +import org.apache.sysds.runtime.ooc.cache.OOCCache; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; +import org.apache.sysds.runtime.ooc.memory.MemoryAllowance; + +final class PackedUnpinHandle implements OOCCache.UnpinHandle { + final BlockEntry entry; + final MemoryAllowance allowance; + final long bytes; + final OOCFuture future; + + static PackedUnpinHandle committed(BlockEntry entry, MemoryAllowance allowance, long bytes) { + return new PackedUnpinHandle(entry, allowance, bytes, true); + } + + static PackedUnpinHandle pendingProducerTransfer(BlockEntry entry, MemoryAllowance allowance, long bytes) { + return new PackedUnpinHandle(entry, allowance, bytes, false); + } + + static PackedUnpinHandle delayedPhysicalRelease(BlockEntry entry, MemoryAllowance allowance) { + return new PackedUnpinHandle(entry, allowance, entry.getSize(), false); + } + + private PackedUnpinHandle(BlockEntry entry, MemoryAllowance allowance, long bytes, boolean committed) { + this.entry = entry; + this.allowance = allowance; + this.bytes = bytes; + future = committed ? OOCFuture.completed(true) : new OOCFuture<>(); + } + + @Override + public BlockEntry entry() { + return entry; + } + + @Override + public MemoryAllowance allowance() { + return allowance; + } + + @Override + public long bytes() { + return bytes; + } + + @Override + public boolean isCommitted() { + return Boolean.TRUE.equals(future.getNow(false)); + } + + @Override + public OOCFuture getCompletionFuture() { + return future; + } + + void complete(boolean committed) { + future.complete(committed); + } +} diff --git a/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCCacheImplTest.java b/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCCacheImplTest.java index 283cebe4667..5219c12b241 100644 --- a/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCCacheImplTest.java +++ b/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCCacheImplTest.java @@ -19,22 +19,18 @@ package org.apache.sysds.test.component.ooc.cache; -import java.util.Map; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ConcurrentHashMap; +import static org.apache.sysds.test.component.ooc.cache.OOCCacheTestUtils.await; + import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.function.BooleanSupplier; import org.apache.sysds.runtime.ooc.cache.BlockEntry; import org.apache.sysds.runtime.ooc.cache.BlockKey; import org.apache.sysds.runtime.ooc.cache.OOCCache; import org.apache.sysds.runtime.ooc.cache.OOCCacheImpl; -import org.apache.sysds.runtime.ooc.cache.OOCFuture; -import org.apache.sysds.runtime.ooc.cache.io.OOCIOHandler; import org.apache.sysds.runtime.ooc.memory.GlobalMemoryBroker; import org.apache.sysds.runtime.ooc.memory.SyncMemoryAllowance; +import org.apache.sysds.test.component.ooc.cache.OOCCacheTestUtils.RecordingOOCIOHandler; import org.junit.After; import org.junit.Assert; import org.junit.Before; @@ -46,7 +42,7 @@ public class OOCCacheImplTest { private static final long BYTES = 1_000; private static final long WAIT_TIMEOUT_SEC = 10; - private RecordingIOHandler _io; + private RecordingOOCIOHandler _io; private GlobalMemoryBroker _broker; private SyncMemoryAllowance _producer; private SyncMemoryAllowance _reader; @@ -54,7 +50,7 @@ public class OOCCacheImplTest { @Before public void setUp() { - _io = new RecordingIOHandler(); + _io = new RecordingOOCIOHandler(); _broker = new GlobalMemoryBroker(8 * BYTES); _producer = new SyncMemoryAllowance(_broker, 4 * BYTES); _reader = new SyncMemoryAllowance(_broker, 4 * BYTES); @@ -92,7 +88,7 @@ public void testResidentPinTransfersOwnershipBetweenCacheAndAllowance() throws E Assert.assertEquals(0, _cache.getOwnedCacheSize()); Assert.assertEquals(BYTES, _producer.getUsedMemory()); - await(_cache.unpin(entry, _producer)); + await(_cache.unpin(entry, _producer), WAIT_TIMEOUT_SEC); Assert.assertEquals(BYTES, _cache.getOwnedCacheSize()); Assert.assertEquals(0, _producer.getUsedMemory()); @@ -103,7 +99,7 @@ public void testResidentPinTransfersOwnershipBetweenCacheAndAllowance() throws E Assert.assertEquals(BYTES, _reader.getUsedMemory()); Assert.assertEquals(0, _io.readCount()); - await(_cache.unpin(pinned, _reader)); + await(_cache.unpin(pinned, _reader), WAIT_TIMEOUT_SEC); Assert.assertEquals(BYTES, _cache.getOwnedCacheSize()); Assert.assertEquals(0, _reader.getUsedMemory()); } @@ -116,8 +112,8 @@ public void testPinReloadsColdBackedEntry() throws Exception { _producer.reserveBlocking(BYTES); BlockEntry entry = _cache.putPinned(key, payload, BYTES, _producer); - await(_cache.unpin(entry, _producer)); - waitFor(() -> _io.evictionCount() == 1 && BlockEntryTestAccess.getDataUnsafe(entry) == null); + await(_cache.unpin(entry, _producer), WAIT_TIMEOUT_SEC); + await(() -> _io.evictionCount() == 1 && BlockEntryTestAccess.getDataUnsafe(entry) == null, WAIT_TIMEOUT_SEC); Assert.assertEquals(0, _producer.getUsedMemory()); BlockEntry pinned = _cache.pin(key, _reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); @@ -128,7 +124,7 @@ public void testPinReloadsColdBackedEntry() throws Exception { Assert.assertEquals(1, _io.readCount()); Assert.assertEquals(BYTES, _reader.getUsedMemory()); - await(_cache.unpin(pinned, _reader)); + await(_cache.unpin(pinned, _reader), WAIT_TIMEOUT_SEC); Assert.assertEquals(0, _reader.getUsedMemory()); } @@ -140,8 +136,8 @@ public void testPinIfLiveDoesNotReadColdBackedEntry() throws Exception { _producer.reserveBlocking(BYTES); BlockEntry entry = _cache.putPinned(key, payload, BYTES, _producer); - await(_cache.unpin(entry, _producer)); - waitFor(() -> _io.evictionCount() == 1 && BlockEntryTestAccess.getDataUnsafe(entry) == null); + await(_cache.unpin(entry, _producer), WAIT_TIMEOUT_SEC); + await(() -> _io.evictionCount() == 1 && BlockEntryTestAccess.getDataUnsafe(entry) == null, WAIT_TIMEOUT_SEC); BlockEntry pinned = _cache.pinIfLive(STREAM_ID, BLOCK_ID, _reader); @@ -203,7 +199,7 @@ public void testDereferenceRemovesEntryAfterLastUnpin() throws Exception { BlockEntry entry = _cache.putPinned(key, "drop", BYTES, _producer); Assert.assertEquals(0, _cache.dereference(entry)); - await(_cache.unpin(entry, _producer)); + await(_cache.unpin(entry, _producer), WAIT_TIMEOUT_SEC); Assert.assertEquals(0, _producer.getUsedMemory()); Assert.assertNull(_cache.pin(key, _reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS)); @@ -217,8 +213,8 @@ public void testBackingReadFailureReleasesReservedBytes() throws Exception { _producer.reserveBlocking(BYTES); BlockEntry entry = _cache.putPinned(key, "fail-read", BYTES, _producer); - await(_cache.unpin(entry, _producer)); - waitFor(() -> _io.evictionCount() == 1 && BlockEntryTestAccess.getDataUnsafe(entry) == null); + await(_cache.unpin(entry, _producer), WAIT_TIMEOUT_SEC); + await(() -> _io.evictionCount() == 1 && BlockEntryTestAccess.getDataUnsafe(entry) == null, WAIT_TIMEOUT_SEC); _io.failReads(true); try { @@ -244,91 +240,4 @@ private void useZeroHardLimitCache() { _io.reset(); _cache = new OOCCacheImpl(_io, 0, 0); } - - private static void await(OOCCache.UnpinHandle handle) throws Exception { - if(!handle.isCommitted()) - handle.getCompletionFuture().get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); - } - - private static void waitFor(BooleanSupplier condition) throws Exception { - long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(WAIT_TIMEOUT_SEC); - while(!condition.getAsBoolean() && System.nanoTime() < deadline) - Thread.sleep(1); - Assert.assertTrue(condition.getAsBoolean()); - } - - private static final class RecordingIOHandler implements OOCIOHandler { - private final Map _spilled = new ConcurrentHashMap<>(); - private final AtomicInteger _evictions = new AtomicInteger(); - private final AtomicInteger _reads = new AtomicInteger(); - private volatile boolean _failReads; - - @Override - public void shutdown() { - _spilled.clear(); - } - - @Override - public CompletableFuture scheduleEviction(BlockEntry block) { - _spilled.put(block.getKey(), BlockEntryTestAccess.getDataUnsafe(block)); - _evictions.incrementAndGet(); - return CompletableFuture.completedFuture(null); - } - - @Override - public OOCFuture scheduleRead(BlockEntry block) { - _reads.incrementAndGet(); - if(_failReads) - return OOCFuture.failed(new IllegalStateException("read failed")); - Object data = _spilled.get(block.getKey()); - if(data == null) - return OOCFuture.completed(null); - BlockEntryTestAccess.setDataUnsafe(block, data); - return OOCFuture.completed(block); - } - - @Override - public void prioritizeRead(BlockKey key, double priority) { - } - - @Override - public CompletableFuture scheduleDeletion(BlockEntry block) { - _spilled.remove(block.getKey()); - return CompletableFuture.completedFuture(true); - } - - @Override - public void registerSourceLocation(BlockKey key, SourceBlockDescriptor descriptor) { - } - - @Override - public CompletableFuture scheduleSourceRead(SourceReadRequest request) { - return CompletableFuture.failedFuture(new UnsupportedOperationException()); - } - - @Override - public CompletableFuture continueSourceRead(SourceReadContinuation continuation, - long maxBytesInFlight) { - return CompletableFuture.failedFuture(new UnsupportedOperationException()); - } - - private int evictionCount() { - return _evictions.get(); - } - - private int readCount() { - return _reads.get(); - } - - private void failReads(boolean failReads) { - _failReads = failReads; - } - - private void reset() { - _spilled.clear(); - _evictions.set(0); - _reads.set(0); - _failReads = false; - } - } } diff --git a/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCCacheTestUtils.java b/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCCacheTestUtils.java new file mode 100644 index 00000000000..8fee37646e9 --- /dev/null +++ b/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCCacheTestUtils.java @@ -0,0 +1,130 @@ +/* + * 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.cache; + +import org.apache.sysds.runtime.ooc.cache.BlockEntry; +import org.apache.sysds.runtime.ooc.cache.BlockKey; +import org.apache.sysds.runtime.ooc.cache.OOCCache; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; +import org.apache.sysds.runtime.ooc.cache.io.OOCIOHandler; +import org.apache.sysds.runtime.ooc.memory.SyncMemoryAllowance; +import org.junit.Assert; + +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BooleanSupplier; + +public class OOCCacheTestUtils { + + public static void await(OOCCache.UnpinHandle handle, long timeout) throws Exception { + if(!handle.isCommitted()) + handle.getCompletionFuture().get(timeout, TimeUnit.SECONDS); + } + + public static void await(BooleanSupplier condition, long timeout) throws Exception { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(timeout); + while(!condition.getAsBoolean() && System.nanoTime() < deadline) + Thread.sleep(1); + Assert.assertTrue(condition.getAsBoolean()); + } + + public static void awaitUsedMemory(SyncMemoryAllowance allowance, long expected, long timeout) throws Exception { + await(() -> allowance.getUsedMemory() == expected, timeout); + } + + public static class RecordingOOCIOHandler implements OOCIOHandler { + private final Map _spilled = new ConcurrentHashMap<>(); + private final AtomicInteger _evictions = new AtomicInteger(); + private final AtomicInteger _reads = new AtomicInteger(); + private volatile boolean _failReads; + + @Override + public void shutdown() { + _spilled.clear(); + } + + @Override + public CompletableFuture scheduleEviction(BlockEntry block) { + _spilled.put(block.getKey(), BlockEntryTestAccess.getDataUnsafe(block)); + _evictions.incrementAndGet(); + return CompletableFuture.completedFuture(null); + } + + @Override + public OOCFuture scheduleRead(BlockEntry block) { + _reads.incrementAndGet(); + if(_failReads) + return OOCFuture.failed(new RuntimeException("Injected read failure")); + Object data = _spilled.get(block.getKey()); + if(data == null) + return OOCFuture.completed(null); + BlockEntryTestAccess.setDataUnsafe(block, data); + return OOCFuture.completed(block); + } + + @Override + public void prioritizeRead(BlockKey key, double priority) { + } + + @Override + public CompletableFuture scheduleDeletion(BlockEntry block) { + _spilled.remove(block.getKey()); + return CompletableFuture.completedFuture(true); + } + + @Override + public void registerSourceLocation(BlockKey key, OOCIOHandler.SourceBlockDescriptor descriptor) { + } + + @Override + public CompletableFuture scheduleSourceRead( + OOCIOHandler.SourceReadRequest request) { + return CompletableFuture.failedFuture(new UnsupportedOperationException()); + } + + @Override + public CompletableFuture continueSourceRead( + OOCIOHandler.SourceReadContinuation continuation, long maxBytesInFlight) { + return CompletableFuture.failedFuture(new UnsupportedOperationException()); + } + + public int evictionCount() { + return _evictions.get(); + } + + public int readCount() { + return _reads.get(); + } + + public void failReads(boolean failReads) { + _failReads = failReads; + } + + public void reset() { + _spilled.clear(); + _evictions.set(0); + _reads.set(0); + _failReads = false; + } + } +} diff --git a/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCPackedCacheTest.java b/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCPackedCacheTest.java new file mode 100644 index 00000000000..0ffb5210b5e --- /dev/null +++ b/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCPackedCacheTest.java @@ -0,0 +1,226 @@ +/* + * 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.cache; + +import static org.apache.sysds.test.component.ooc.cache.OOCCacheTestUtils.await; +import static org.apache.sysds.test.component.ooc.cache.OOCCacheTestUtils.awaitUsedMemory; + +import java.util.concurrent.TimeUnit; + +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.BlockEntry; +import org.apache.sysds.runtime.ooc.cache.OOCCache; +import org.apache.sysds.runtime.ooc.cache.OOCCacheImpl; +import org.apache.sysds.runtime.ooc.cache.io.OOCMatrixIOHandler; +import org.apache.sysds.runtime.ooc.cache.packed.OOCPackedCache; +import org.apache.sysds.runtime.ooc.memory.GlobalMemoryBroker; +import org.apache.sysds.runtime.ooc.memory.SyncMemoryAllowance; +import org.apache.sysds.test.component.ooc.cache.OOCCacheTestUtils.RecordingOOCIOHandler; +import org.junit.Assert; +import org.junit.Test; + +public class OOCPackedCacheTest { + private static final long STREAM_ID = 41; + private static final long BYTES = 1000; + private static final long WAIT_TIMEOUT_SEC = 10; + + @Test + public void testSmallTilesShareOnePhysicalPack() throws Exception { + GlobalMemoryBroker broker = new GlobalMemoryBroker(1L << 32); + SyncMemoryAllowance producer = new SyncMemoryAllowance(broker); + producer.setTargetMemory(1L << 30); + SyncMemoryAllowance reader = new SyncMemoryAllowance(broker); + reader.setTargetMemory(1L << 30); + OOCPackedCache cache = new OOCPackedCache(new OOCCacheImpl(new OOCMatrixIOHandler(), 1L << 30, 1L << 30), + 2 * BYTES, 10 * BYTES, -1, 0); + try { + BlockEntry[] entries = publishSmallTiles(cache, producer, STREAM_ID, 3); + unpinAndFlush(cache, producer, entries); + awaitUsedMemory(producer, 0, WAIT_TIMEOUT_SEC); + + Assert.assertEquals(1, cache.getPackGroupCount()); + OOCPackedCache.PackGroup group = cache.getPackGroup(STREAM_ID, 0); + Assert.assertNotNull(group); + Assert.assertEquals(3, group.size()); + + BlockEntry first = cache.pin(STREAM_ID, 0, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + BlockEntry second = cache.pin(STREAM_ID, 1, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + Assert.assertEquals(1.0, scalar(first), 0.0); + Assert.assertEquals(2.0, scalar(second), 0.0); + Assert.assertEquals("Multiple logical pins in one pack should charge the physical pack once.", 3 * BYTES, + reader.getUsedMemory()); + + await(cache.unpin(first, reader), WAIT_TIMEOUT_SEC); + Assert.assertEquals("The physical pack stays pinned while another logical pin remains.", 3 * BYTES, + reader.getUsedMemory()); + await(cache.unpin(second, reader), WAIT_TIMEOUT_SEC); + awaitUsedMemory(reader, 0, WAIT_TIMEOUT_SEC); + } + finally { + cache.shutdown(); + producer.destroy(); + reader.destroy(); + } + } + + @Test + public void testLargeBlockBypassesPacking() throws Exception { + GlobalMemoryBroker broker = new GlobalMemoryBroker(1L << 32); + SyncMemoryAllowance producer = new SyncMemoryAllowance(broker); + producer.setTargetMemory(1L << 30); + SyncMemoryAllowance reader = new SyncMemoryAllowance(broker); + reader.setTargetMemory(1L << 30); + OOCPackedCache cache = new OOCPackedCache(new OOCCacheImpl(new OOCMatrixIOHandler(), 1L << 30, 1L << 30), + 2 * BYTES, 10 * BYTES, -1, 0); + long largeBytes = 2 * BYTES; + try { + producer.reserveBlocking(largeBytes); + BlockEntry entry = cache.putPinned(STREAM_ID, 0, value(5.0), largeBytes, producer); + await(cache.unpin(entry, producer), WAIT_TIMEOUT_SEC); + + Assert.assertEquals(0, cache.getPackGroupCount()); + Assert.assertNull(cache.getPackGroup(STREAM_ID, 0)); + + BlockEntry pinned = cache.pin(STREAM_ID, 0, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + Assert.assertNotNull(pinned); + Assert.assertEquals(5.0, scalar(pinned), 0.0); + Assert.assertEquals(largeBytes, reader.getUsedMemory()); + await(cache.unpin(pinned, reader), WAIT_TIMEOUT_SEC); + awaitUsedMemory(reader, 0, WAIT_TIMEOUT_SEC); + } + finally { + cache.shutdown(); + producer.destroy(); + reader.destroy(); + } + } + + @Test + public void testPinPackExposesWholePack() throws Exception { + GlobalMemoryBroker broker = new GlobalMemoryBroker(1L << 32); + SyncMemoryAllowance producer = new SyncMemoryAllowance(broker); + producer.setTargetMemory(1L << 30); + SyncMemoryAllowance reader = new SyncMemoryAllowance(broker); + reader.setTargetMemory(1L << 30); + OOCPackedCache cache = new OOCPackedCache(new OOCCacheImpl(new OOCMatrixIOHandler(), 1L << 30, 1L << 30), + 2 * BYTES, 10 * BYTES, -1, 0); + try { + long[] tileIds = new long[] {2, 5}; + Object[] values = new Object[] {value(7.0), value(11.0)}; + long[] sizes = new long[] {BYTES, BYTES}; + producer.reserveBlocking(2 * BYTES); + BlockEntry physical = cache.putSealedPackPinned(STREAM_ID, tileIds, values, sizes, 0, tileIds.length, + producer); + await(cache.unpin(physical, producer), WAIT_TIMEOUT_SEC); + awaitUsedMemory(producer, 0, WAIT_TIMEOUT_SEC); + + OOCPackedCache.PackGroup group = cache.getPackGroup(STREAM_ID, 5); + Assert.assertNotNull(group); + Assert.assertEquals(2, group.size()); + Assert.assertEquals(2, group.index(0)); + Assert.assertEquals(5, group.index(1)); + + OOCPackedCache.PackLease lease = cache.pinPack(group, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + Assert.assertNotNull(lease); + try { + Assert.assertEquals(7.0, scalar((IndexedMatrixValue) lease.value(0)), 0.0); + Assert.assertEquals(11.0, scalar((IndexedMatrixValue) lease.value(1)), 0.0); + Assert.assertEquals(2 * BYTES, reader.getUsedMemory()); + } + finally { + lease.close(); + } + awaitUsedMemory(reader, 0, WAIT_TIMEOUT_SEC); + } + finally { + cache.shutdown(); + producer.destroy(); + reader.destroy(); + } + } + + @Test + public void testEvictedPackReplaysThroughLogicalPin() throws Exception { + RecordingOOCIOHandler io = new RecordingOOCIOHandler(); + GlobalMemoryBroker broker = new GlobalMemoryBroker(1L << 32); + SyncMemoryAllowance producer = new SyncMemoryAllowance(broker); + producer.setTargetMemory(1L << 30); + SyncMemoryAllowance reader = new SyncMemoryAllowance(broker); + reader.setTargetMemory(1L << 30); + OOCPackedCache cache = new OOCPackedCache(new OOCCacheImpl(io, 4 * BYTES, 0), 2 * BYTES, 10 * BYTES, -1, 0); + try { + BlockEntry[] entries = publishSmallTiles(cache, producer, STREAM_ID, 4); + unpinAndFlush(cache, producer, entries); + awaitUsedMemory(producer, 0, WAIT_TIMEOUT_SEC); + await(() -> io.evictionCount() > 0 && cache.getOwnedCacheSize() == 0, WAIT_TIMEOUT_SEC); + + int readsBefore = io.readCount(); + BlockEntry pinned = cache.pin(STREAM_ID, 3, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + + Assert.assertNotNull(pinned); + Assert.assertEquals(4.0, scalar(pinned), 0.0); + Assert.assertTrue("Pinning an evicted logical tile should read the physical pack.", + io.readCount() > readsBefore); + Assert.assertEquals(4 * BYTES, reader.getUsedMemory()); + + await(cache.unpin(pinned, reader), WAIT_TIMEOUT_SEC); + awaitUsedMemory(reader, 0, WAIT_TIMEOUT_SEC); + } + finally { + cache.shutdown(); + producer.destroy(); + reader.destroy(); + } + } + + private static BlockEntry[] publishSmallTiles(OOCPackedCache cache, SyncMemoryAllowance producer, long streamId, + int count) { + BlockEntry[] entries = new BlockEntry[count]; + for(int i = 0; i < count; i++) { + producer.reserveBlocking(BYTES); + entries[i] = cache.putPinned(streamId, i, value(i + 1.0), BYTES, producer); + } + return entries; + } + + private static void unpinAndFlush(OOCPackedCache cache, SyncMemoryAllowance producer, BlockEntry[] entries) + throws Exception { + OOCCache.UnpinHandle[] handles = new OOCCache.UnpinHandle[entries.length]; + for(int i = 0; i < entries.length; i++) + handles[i] = cache.unpin(entries[i], producer); + cache.flushPacks(); + for(OOCCache.UnpinHandle handle : handles) + await(handle, WAIT_TIMEOUT_SEC); + } + + private static IndexedMatrixValue value(double scalar) { + return new IndexedMatrixValue(new MatrixIndexes(1, 1), new MatrixBlock(1, 1, scalar)); + } + + private static double scalar(BlockEntry entry) { + return scalar((IndexedMatrixValue) entry.getData()); + } + + private static double scalar(IndexedMatrixValue value) { + return value.getValue().get(0, 0); + } +} From 59b37686b4aa8d8f90e8eb3acd724c33b3830c2e Mon Sep 17 00:00:00 2001 From: Jannik Lindemann Date: Tue, 7 Jul 2026 09:51:47 +0200 Subject: [PATCH 2/3] [OOC] Extend Packed Cache Test for Better Coverage --- .../ooc/cache/OOCPackedCacheTest.java | 169 +++++++++++++++++- 1 file changed, 165 insertions(+), 4 deletions(-) diff --git a/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCPackedCacheTest.java b/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCPackedCacheTest.java index 0ffb5210b5e..415f4bd1dda 100644 --- a/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCPackedCacheTest.java +++ b/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCPackedCacheTest.java @@ -28,6 +28,7 @@ import org.apache.sysds.runtime.matrix.data.MatrixBlock; import org.apache.sysds.runtime.matrix.data.MatrixIndexes; import org.apache.sysds.runtime.ooc.cache.BlockEntry; +import org.apache.sysds.runtime.ooc.cache.BlockKey; import org.apache.sysds.runtime.ooc.cache.OOCCache; import org.apache.sysds.runtime.ooc.cache.OOCCacheImpl; import org.apache.sysds.runtime.ooc.cache.io.OOCMatrixIOHandler; @@ -82,6 +83,128 @@ public void testSmallTilesShareOnePhysicalPack() throws Exception { } } + @Test + public void testPutPackPinnedPacksSmallTilesAndBypassesLargeTile() throws Exception { + GlobalMemoryBroker broker = new GlobalMemoryBroker(1L << 32); + SyncMemoryAllowance producer = new SyncMemoryAllowance(broker); + producer.setTargetMemory(1L << 30); + SyncMemoryAllowance reader = new SyncMemoryAllowance(broker); + reader.setTargetMemory(1L << 30); + OOCPackedCache cache = new OOCPackedCache(new OOCCacheImpl(new OOCMatrixIOHandler(), 1L << 30, 1L << 30), + 2 * BYTES, 10 * BYTES, -1, 0); + try { + long largeBytes = 2 * BYTES; + long[] tileIds = new long[] {0, 1, 2}; + Object[] values = new Object[] {value(3.0), value(9.0), value(5.0)}; + long[] sizes = new long[] {BYTES, largeBytes, BYTES}; + producer.reserveBlocking(2 * BYTES + largeBytes); + BlockEntry[] entries = cache.putPackPinned(STREAM_ID, tileIds, values, sizes, 0, tileIds.length, + producer); + unpinAndFlush(cache, producer, entries); + awaitUsedMemory(producer, 0, WAIT_TIMEOUT_SEC); + + OOCPackedCache.PackGroup group = cache.getPackGroup(STREAM_ID, 0); + Assert.assertNotNull(group); + Assert.assertEquals(2, group.size()); + Assert.assertEquals(0, group.index(0)); + Assert.assertEquals(2, group.index(1)); + Assert.assertNull("Large tiles in putPackPinned should bypass packing.", + cache.getPackGroup(STREAM_ID, 1)); + + BlockEntry packed = cache.pin(STREAM_ID, 2, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + BlockEntry large = cache.pin(STREAM_ID, 1, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + Assert.assertEquals(5.0, scalar(packed), 0.0); + Assert.assertEquals(9.0, scalar(large), 0.0); + Assert.assertEquals(2 * BYTES + largeBytes, reader.getUsedMemory()); + + await(cache.unpin(packed, reader), WAIT_TIMEOUT_SEC); + await(cache.unpin(large, reader), WAIT_TIMEOUT_SEC); + awaitUsedMemory(reader, 0, WAIT_TIMEOUT_SEC); + } + finally { + cache.shutdown(); + producer.destroy(); + reader.destroy(); + } + } + + @Test + public void testPinAdmittedReusesPackedPhysicalPin() throws Exception { + GlobalMemoryBroker broker = new GlobalMemoryBroker(1L << 32); + SyncMemoryAllowance producer = new SyncMemoryAllowance(broker); + producer.setTargetMemory(1L << 30); + SyncMemoryAllowance reader = new SyncMemoryAllowance(broker); + reader.setTargetMemory(1L << 30); + OOCPackedCache cache = new OOCPackedCache(new OOCCacheImpl(new OOCMatrixIOHandler(), 1L << 30, 1L << 30), + 2 * BYTES, 10 * BYTES, -1, 1000); + try { + BlockEntry[] entries = publishSmallTiles(cache, producer, STREAM_ID, 2); + unpinAndFlush(cache, producer, entries); + awaitUsedMemory(producer, 0, WAIT_TIMEOUT_SEC); + + BlockEntry first = cache.pinAdmitted(STREAM_ID, 0, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + Assert.assertNotNull(first); + Assert.assertEquals(1.0, scalar(first), 0.0); + Assert.assertEquals(2 * BYTES, reader.getUsedMemory()); + + OOCCache.UnpinHandle delayed = cache.unpin(first, reader); + Assert.assertFalse(delayed.isCommitted()); + Assert.assertEquals("Delayed packed release keeps the physical pack charged.", 2 * BYTES, + reader.getUsedMemory()); + + BlockEntry second = cache.pinAdmitted(STREAM_ID, 1, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + Assert.assertNotNull(second); + Assert.assertEquals(2.0, scalar(second), 0.0); + Assert.assertFalse("Re-pinning with the same allowance should cancel the pending physical release.", + delayed.getCompletionFuture().get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS)); + Assert.assertEquals(2 * BYTES, reader.getUsedMemory()); + + await(cache.unpin(second, reader), WAIT_TIMEOUT_SEC); + awaitUsedMemory(reader, 0, WAIT_TIMEOUT_SEC); + } + finally { + cache.shutdown(); + producer.destroy(); + reader.destroy(); + } + } + + @Test + public void testReferenceAndDereferencePackedLocations() throws Exception { + GlobalMemoryBroker broker = new GlobalMemoryBroker(1L << 32); + SyncMemoryAllowance producer = new SyncMemoryAllowance(broker); + producer.setTargetMemory(1L << 30); + SyncMemoryAllowance reader = new SyncMemoryAllowance(broker); + reader.setTargetMemory(1L << 30); + OOCPackedCache cache = new OOCPackedCache(new OOCCacheImpl(new OOCMatrixIOHandler(), 1L << 30, 1L << 30), + 2 * BYTES, 10 * BYTES, -1, 0); + try { + producer.reserveBlocking(BYTES); + BlockEntry pending = cache.putPinned(STREAM_ID, 0, value(13.0), BYTES, producer); + Assert.assertEquals(2, cache.reference(pending)); + Assert.assertEquals(1, cache.dereference(pending)); + + unpinAndFlush(cache, producer, new BlockEntry[] {pending}); + awaitUsedMemory(producer, 0, WAIT_TIMEOUT_SEC); + + BlockEntry pinned = cache.pin(STREAM_ID, 0, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + Assert.assertNotNull(pinned); + Assert.assertEquals(13.0, scalar(pinned), 0.0); + Assert.assertEquals(2, cache.reference(pinned)); + Assert.assertEquals(1, cache.dereference(pinned)); + Assert.assertEquals(0, cache.dereference(new BlockKey(STREAM_ID, 0))); + + await(cache.unpin(pinned, reader), WAIT_TIMEOUT_SEC); + awaitUsedMemory(reader, 0, WAIT_TIMEOUT_SEC); + Assert.assertNull(cache.pin(STREAM_ID, 0, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS)); + } + finally { + cache.shutdown(); + producer.destroy(); + reader.destroy(); + } + } + @Test public void testLargeBlockBypassesPacking() throws Exception { GlobalMemoryBroker broker = new GlobalMemoryBroker(1L << 32); @@ -141,14 +264,52 @@ public void testPinPackExposesWholePack() throws Exception { OOCPackedCache.PackLease lease = cache.pinPack(group, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); Assert.assertNotNull(lease); - try { + try(lease) { Assert.assertEquals(7.0, scalar((IndexedMatrixValue) lease.value(0)), 0.0); Assert.assertEquals(11.0, scalar((IndexedMatrixValue) lease.value(1)), 0.0); Assert.assertEquals(2 * BYTES, reader.getUsedMemory()); } - finally { - lease.close(); - } + awaitUsedMemory(reader, 0, WAIT_TIMEOUT_SEC); + } + finally { + cache.shutdown(); + producer.destroy(); + reader.destroy(); + } + } + + @Test + public void testLogicalEvictionPolicyScoresPackedBlocks() throws Exception { + RecordingOOCIOHandler io = new RecordingOOCIOHandler(); + GlobalMemoryBroker broker = new GlobalMemoryBroker(1L << 32); + SyncMemoryAllowance producer = new SyncMemoryAllowance(broker); + producer.setTargetMemory(1L << 30); + SyncMemoryAllowance reader = new SyncMemoryAllowance(broker); + reader.setTargetMemory(1L << 30); + OOCPackedCache cache = new OOCPackedCache(new OOCCacheImpl(io, 4 * BYTES, 2 * BYTES), 2 * BYTES, + 2 * BYTES, -1, 0); + try { + cache.addEvictionPolicy(STREAM_ID, tileId -> tileId < 2 ? 100 : 0); + BlockEntry[] entries = publishSmallTiles(cache, producer, STREAM_ID, 4); + unpinAndFlush(cache, producer, entries); + awaitUsedMemory(producer, 0, WAIT_TIMEOUT_SEC); + await(() -> io.evictionCount() == 1 && cache.getOwnedCacheSize() == 2 * BYTES, WAIT_TIMEOUT_SEC); + + int readsBefore = io.readCount(); + BlockEntry retained = cache.pin(STREAM_ID, 2, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + Assert.assertNotNull(retained); + Assert.assertEquals(3.0, scalar(retained), 0.0); + Assert.assertEquals("The lower-scored packed group should remain resident.", readsBefore, io.readCount()); + await(cache.unpin(retained, reader), WAIT_TIMEOUT_SEC); + awaitUsedMemory(reader, 0, WAIT_TIMEOUT_SEC); + + readsBefore = io.readCount(); + BlockEntry evicted = cache.pin(STREAM_ID, 0, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + Assert.assertNotNull(evicted); + Assert.assertEquals(1.0, scalar(evicted), 0.0); + Assert.assertTrue("The higher-scored packed group should be evicted first.", + io.readCount() > readsBefore); + await(cache.unpin(evicted, reader), WAIT_TIMEOUT_SEC); awaitUsedMemory(reader, 0, WAIT_TIMEOUT_SEC); } finally { From 3539bd941d877296a5671a32c1af8193f91ec6f4 Mon Sep 17 00:00:00 2001 From: Jannik Lindemann Date: Tue, 7 Jul 2026 09:59:13 +0200 Subject: [PATCH 3/3] [OOC] Reformat --- .../component/ooc/cache/OOCPackedCacheTest.java | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCPackedCacheTest.java b/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCPackedCacheTest.java index 415f4bd1dda..7a29f8781b4 100644 --- a/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCPackedCacheTest.java +++ b/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCPackedCacheTest.java @@ -98,8 +98,7 @@ public void testPutPackPinnedPacksSmallTilesAndBypassesLargeTile() throws Except Object[] values = new Object[] {value(3.0), value(9.0), value(5.0)}; long[] sizes = new long[] {BYTES, largeBytes, BYTES}; producer.reserveBlocking(2 * BYTES + largeBytes); - BlockEntry[] entries = cache.putPackPinned(STREAM_ID, tileIds, values, sizes, 0, tileIds.length, - producer); + BlockEntry[] entries = cache.putPackPinned(STREAM_ID, tileIds, values, sizes, 0, tileIds.length, producer); unpinAndFlush(cache, producer, entries); awaitUsedMemory(producer, 0, WAIT_TIMEOUT_SEC); @@ -108,8 +107,7 @@ public void testPutPackPinnedPacksSmallTilesAndBypassesLargeTile() throws Except Assert.assertEquals(2, group.size()); Assert.assertEquals(0, group.index(0)); Assert.assertEquals(2, group.index(1)); - Assert.assertNull("Large tiles in putPackPinned should bypass packing.", - cache.getPackGroup(STREAM_ID, 1)); + Assert.assertNull("Large tiles in putPackPinned should bypass packing.", cache.getPackGroup(STREAM_ID, 1)); BlockEntry packed = cache.pin(STREAM_ID, 2, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); BlockEntry large = cache.pin(STREAM_ID, 1, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); @@ -286,8 +284,8 @@ public void testLogicalEvictionPolicyScoresPackedBlocks() throws Exception { producer.setTargetMemory(1L << 30); SyncMemoryAllowance reader = new SyncMemoryAllowance(broker); reader.setTargetMemory(1L << 30); - OOCPackedCache cache = new OOCPackedCache(new OOCCacheImpl(io, 4 * BYTES, 2 * BYTES), 2 * BYTES, - 2 * BYTES, -1, 0); + OOCPackedCache cache = new OOCPackedCache(new OOCCacheImpl(io, 4 * BYTES, 2 * BYTES), 2 * BYTES, 2 * BYTES, -1, + 0); try { cache.addEvictionPolicy(STREAM_ID, tileId -> tileId < 2 ? 100 : 0); BlockEntry[] entries = publishSmallTiles(cache, producer, STREAM_ID, 4); @@ -307,8 +305,7 @@ public void testLogicalEvictionPolicyScoresPackedBlocks() throws Exception { BlockEntry evicted = cache.pin(STREAM_ID, 0, reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); Assert.assertNotNull(evicted); Assert.assertEquals(1.0, scalar(evicted), 0.0); - Assert.assertTrue("The higher-scored packed group should be evicted first.", - io.readCount() > readsBefore); + Assert.assertTrue("The higher-scored packed group should be evicted first.", io.readCount() > readsBefore); await(cache.unpin(evicted, reader), WAIT_TIMEOUT_SEC); awaitUsedMemory(reader, 0, WAIT_TIMEOUT_SEC); }