From aa56b05eae53ee628a3e1a1f98ea6cd6f4c5eb73 Mon Sep 17 00:00:00 2001 From: Jannik Lindemann Date: Wed, 15 Jul 2026 17:42:45 +0200 Subject: [PATCH] [OOC] Improvements of Memory Allowance/Broker --- .../ooc/memory/GlobalMemoryBroker.java | 162 ++++++++++-------- .../runtime/ooc/memory/MemoryAllowance.java | 4 + .../runtime/ooc/memory/MemoryBroker.java | 6 + .../ooc/memory/SyncMemoryAllowance.java | 86 ++++++---- .../ooc/memory/OOCMemoryAllowanceTest.java | 54 ++++++ 5 files changed, 211 insertions(+), 101 deletions(-) diff --git a/src/main/java/org/apache/sysds/runtime/ooc/memory/GlobalMemoryBroker.java b/src/main/java/org/apache/sysds/runtime/ooc/memory/GlobalMemoryBroker.java index 6009182f156..f7ad7b28577 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/memory/GlobalMemoryBroker.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/memory/GlobalMemoryBroker.java @@ -20,10 +20,17 @@ package org.apache.sysds.runtime.ooc.memory; import java.util.ArrayList; -import java.util.LinkedList; import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; public class GlobalMemoryBroker implements MemoryBroker { + private static final long RECLAIM_RETRY_DELAY_MS = 5; + private static final double RECLAIM_PRESSURE = 0.85; + private static final ScheduledThreadPoolExecutor RECLAIM_EXECUTOR = createReclaimExecutor(); + private enum BrokerMode { RELAXED, STRICT } @@ -35,18 +42,28 @@ public static GlobalMemoryBroker get() { } private final long _allowedBytes; - private final List _allowances; - private final LinkedList _overconsumers; + private final CopyOnWriteArrayList _allowances; + private final AtomicBoolean _reclaimRunning; private long _usedBytes; private BrokerMode _brokerMode; private record TargetUpdate(MemoryAllowance _allowance, long _target) {} + private static ScheduledThreadPoolExecutor createReclaimExecutor() { + ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(1, runnable -> { + Thread thread = new Thread(runnable, "ooc-memory-broker-reclaimer"); + thread.setDaemon(true); + return thread; + }); + executor.setRemoveOnCancelPolicy(true); + return executor; + } + public GlobalMemoryBroker(long allowedBytes) { _allowedBytes = allowedBytes; _usedBytes = 0; - _allowances = new ArrayList<>(); - _overconsumers = new LinkedList<>(); + _allowances = new CopyOnWriteArrayList<>(); + _reclaimRunning = new AtomicBoolean(false); } @Override @@ -56,29 +73,11 @@ public long requestMemory(MemoryAllowance allowance, long minSize, long maxSize) synchronized(this) { if(minSize < 0 || maxSize < minSize) throw new IllegalArgumentException(); - long share = getEqualShare(); long free = _allowedBytes - _usedBytes; - if(free < minSize) { - if(allowance.getGrantedMemory() > share && allowance.getTargetMemory() > allowance.getGrantedMemory()) - updates = List.of(new TargetUpdate(allowance, allowance.getUsedMemory())); - else { - MemoryAllowance largestConsumer = findAndRemoveLargestConsumer(); - if(largestConsumer != null) { - long newTarget = (long) (largestConsumer.getGrantedMemory() * 0.8); - if(newTarget <= share) - newTarget = share; - else - addOverconsumer(largestConsumer); - updates = List.of(new TargetUpdate(largestConsumer, newTarget)); - } - } - } - else { + if(free >= minSize) { allow = Math.min(free, maxSize); _usedBytes += allow; updates = rebalance(false); - if(allowance.getGrantedMemory() <= share && allowance.getGrantedMemory() + allow > share) - addOverconsumer(allowance); } } if(updates != null) @@ -86,46 +85,29 @@ public long requestMemory(MemoryAllowance allowance, long minSize, long maxSize) return allow; } - private MemoryAllowance findAndRemoveLargestConsumer() { - long largest = Long.MIN_VALUE; - MemoryAllowance allowance = null; - for(MemoryAllowance largestConsumer : _overconsumers) { - if(largestConsumer.getGrantedMemory() > largest) { - largest = largestConsumer.getGrantedMemory(); - allowance = largestConsumer; - } - } - _overconsumers.remove(allowance); - return allowance; - } - @Override public void freeMemory(MemoryAllowance allowance, long freedMemory) { - List updates = null; + List updates; synchronized(this) { if(freedMemory < 0) throw new IllegalArgumentException(); _usedBytes -= freedMemory; - if(allowance.isShutdown()) - updates = rebalance(false); - long share = getEqualShare(); - if(allowance.getGrantedMemory() <= share && allowance.getGrantedMemory() + freedMemory > share) - _overconsumers.remove(allowance); - else if(allowance.getGrantedMemory() <= allowance.getTargetMemory() && allowance.getGrantedMemory() > share) - addOverconsumer(allowance); + updates = rebalanceAfterFree(); } if(updates != null) applyTargetUpdates(updates); + if(freedMemory > 0) + notifyReservationWaiters(); } @Override public void shutdownAllowance(MemoryAllowance allowance) { List updates; synchronized(this) { - _overconsumers.remove(allowance); updates = rebalance(true); } applyTargetUpdates(updates); + notifyReservationWaiters(); } @Override @@ -135,11 +117,11 @@ public void destroyAllowance(MemoryAllowance allowance, long freedMemory) { if(freedMemory < 0) throw new IllegalArgumentException(); _allowances.remove(allowance); - _overconsumers.remove(allowance); _usedBytes -= freedMemory; updates = rebalance(true); } applyTargetUpdates(updates); + notifyReservationWaiters(); } @Override @@ -148,6 +130,55 @@ public synchronized void attachAllowance(MemoryAllowance allowance) { allowance.setTargetMemory(_allowedBytes); } + @Override + public void reservationBlocked(MemoryAllowance allowance, long bytes) { + if(_reclaimRunning.compareAndSet(false, true)) + RECLAIM_EXECUTOR.execute(this::runReclaim); + } + + private void runReclaim() { + try { + long reclaimed = 0; + for(MemoryAllowance allowance : _allowances) + if(!allowance.isShutdown()) + reclaimed += allowance.reclaimUnused(); + if(reclaimed == 0) + return; + + List updates; + synchronized(this) { + _usedBytes = Math.max(0, _usedBytes - reclaimed); + updates = rebalanceAfterFree(); + } + if(updates != null) + applyTargetUpdates(updates); + notifyReservationWaiters(); + } + finally { + if(shouldRetryReclaim()) + RECLAIM_EXECUTOR.schedule(this::runReclaim, RECLAIM_RETRY_DELAY_MS, TimeUnit.MILLISECONDS); + else { + _reclaimRunning.set(false); + if(shouldRetryReclaim() && _reclaimRunning.compareAndSet(false, true)) + RECLAIM_EXECUTOR.execute(this::runReclaim); + } + } + } + + private boolean shouldRetryReclaim() { + if(!hasReclaimPressure()) + return false; + for(MemoryAllowance allowance : _allowances) { + if(allowance instanceof SyncMemoryAllowance sync && sync.hasReservationWaiters()) + return true; + } + return false; + } + + private synchronized boolean hasReclaimPressure() { + return _usedBytes >= _allowedBytes * RECLAIM_PRESSURE; + } + private List rebalance(boolean force) { long free = _allowedBytes - _usedBytes; if(force) @@ -158,6 +189,13 @@ private List rebalance(boolean force) { return switchBrokerMode(BrokerMode.STRICT); } + private List rebalanceAfterFree() { + long free = _allowedBytes - _usedBytes; + if(_brokerMode == BrokerMode.RELAXED && free > _allowedBytes / 5) + return rebalanceToRelaxed(); + return rebalance(false); + } + private List switchBrokerMode(BrokerMode newMode) { if(newMode == _brokerMode) return null; @@ -181,7 +219,6 @@ private List rebalanceToStrict() { Math.min(allowance.getTargetMemory(), share + (long) ((allowance.getUsedMemory() - share) * 0.9)))); } } - refreshOverconsumers(updates); return updates; } @@ -193,34 +230,21 @@ private List rebalanceToRelaxed() { continue; updates.add(new TargetUpdate(allowance, allowance.getGrantedMemory() + free)); } - refreshOverconsumers(updates); return updates; } private long getEqualShare() { - return _allowances.isEmpty() ? _allowedBytes : _allowedBytes / _allowances.size(); - } - - private void addOverconsumer(MemoryAllowance allowance) { - if(!_overconsumers.contains(allowance)) - _overconsumers.add(allowance); + int active = 0; + for(MemoryAllowance allowance : _allowances) + if(!allowance.isShutdown()) + active++; + return active == 0 ? _allowedBytes : _allowedBytes / active; } - private void refreshOverconsumers(List updates) { - _overconsumers.clear(); - long share = getEqualShare(); + private void notifyReservationWaiters() { for(MemoryAllowance allowance : _allowances) { - if(allowance.isShutdown()) - continue; - long target = allowance.getTargetMemory(); - for(TargetUpdate update : updates) { - if(update._allowance == allowance) { - target = update._target; - break; - } - } - if(allowance.getGrantedMemory() > share && allowance.getGrantedMemory() <= target) - _overconsumers.add(allowance); + if(allowance instanceof SyncMemoryAllowance sync) + sync.onBrokerMemoryAvailable(); } } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/memory/MemoryAllowance.java b/src/main/java/org/apache/sysds/runtime/ooc/memory/MemoryAllowance.java index b2db5ea7533..10a8bbe7d69 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/memory/MemoryAllowance.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/memory/MemoryAllowance.java @@ -53,4 +53,8 @@ default long getFreeMemory() { default boolean isUnderPressure() { return getGrantedMemory() > getTargetMemory(); } + + default long reclaimUnused() { + return 0; + } } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/memory/MemoryBroker.java b/src/main/java/org/apache/sysds/runtime/ooc/memory/MemoryBroker.java index fb4d6ae182d..2c82cc156da 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/memory/MemoryBroker.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/memory/MemoryBroker.java @@ -21,8 +21,14 @@ public interface MemoryBroker { long requestMemory(MemoryAllowance allowance, long minSize, long maxSize); + void freeMemory(MemoryAllowance allowance, long freedMemory); + void shutdownAllowance(MemoryAllowance allowance); + void destroyAllowance(MemoryAllowance allowance, long freedMemory); + void attachAllowance(MemoryAllowance allowance); + + void reservationBlocked(MemoryAllowance allowance, long bytes); } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/memory/SyncMemoryAllowance.java b/src/main/java/org/apache/sysds/runtime/ooc/memory/SyncMemoryAllowance.java index 2c4a1b7a0fa..900549819a0 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/memory/SyncMemoryAllowance.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/memory/SyncMemoryAllowance.java @@ -22,7 +22,8 @@ import org.apache.sysds.runtime.DMLRuntimeException; import org.apache.sysds.runtime.ooc.cache.OOCFuture; -import java.util.ArrayDeque; +import java.util.Queue; +import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.ExecutionException; public class SyncMemoryAllowance implements MemoryAllowance { @@ -36,7 +37,7 @@ public class SyncMemoryAllowance implements MemoryAllowance { protected volatile long _targetBytes; protected volatile boolean _shutdown; protected volatile boolean _destroyed; - private final ArrayDeque _reservationWaiters; + private final Queue _reservationWaiters; private boolean _drainingReservationWaiters; private boolean _reservationDrainRequested; @@ -62,7 +63,7 @@ public SyncMemoryAllowance(MemoryBroker broker, long consumptionLimit, long mini _targetBytes = 0; _shutdown = false; _destroyed = false; - _reservationWaiters = new ArrayDeque<>(); + _reservationWaiters = new ConcurrentLinkedQueue<>(); _drainingReservationWaiters = false; _reservationDrainRequested = false; broker.attachAllowance(this); @@ -70,12 +71,16 @@ public SyncMemoryAllowance(MemoryBroker broker, long consumptionLimit, long mini @Override public boolean tryReserve(long bytes) { + if(bytes < 0) + throw new IllegalArgumentException("Cannot reserve negative bytes: " + bytes); + if(bytes > _consumptionLimit) + throw new IllegalArgumentException("Cannot reserve more memory than the consumption limit"); long minRequest; long maxRequest; synchronized(this) { if(_shutdown || _destroyed) return false; - if(_usedBytes + bytes <= _grantedBytes) { + if(_usedBytes + bytes <= _grantedBytes && _usedBytes + bytes <= _targetBytes) { _usedBytes += bytes; return true; } @@ -85,9 +90,6 @@ public boolean tryReserve(long bytes) { maxRequest = Math.max(minRequest, Math.max(_grantedBytes, bytes) * 2); } - if(bytes > _consumptionLimit) - throw new IllegalArgumentException("Cannot reserve more memory than the consumption limit"); - long granted = _broker.requestMemory(this, minRequest, maxRequest); long refund = 0; boolean success = false; @@ -135,17 +137,20 @@ public OOCFuture reserveAsync(long bytes) { if(bytes > _consumptionLimit) return OOCFuture .failed(new IllegalArgumentException("Cannot reserve more memory than the consumption limit")); - if(tryReserve(bytes)) + if(_shutdown || _destroyed) + return OOCFuture.failed(new IllegalStateException("Cannot reserve memory on closed allowance.")); + if(_reservationWaiters.isEmpty() && tryReserve(bytes)) return OOCFuture.completed(null); OOCFuture future = new OOCFuture<>(); - synchronized(this) { - if(_shutdown || _destroyed) { - future.completeExceptionally(new IllegalStateException("Cannot reserve memory on closed allowance.")); - return future; - } - _reservationWaiters.addLast(new ReservationWaiter(bytes, future)); + ReservationWaiter waiter = new ReservationWaiter(bytes, future); + _reservationWaiters.add(waiter); + if((_shutdown || _destroyed) && _reservationWaiters.remove(waiter)) { + future.completeExceptionally(new IllegalStateException("Cannot reserve memory on closed allowance.")); + return future; } requestReservationDrain(); + if(!future.isDone()) + _broker.reservationBlocked(this, bytes); return future; } @@ -218,6 +223,8 @@ public long getTargetMemory() { @Override public void setTargetMemory(long targetMemory) { + if(targetMemory < 0) + throw new IllegalArgumentException("Target memory must not be negative: " + targetMemory); long freedMemory = 0; boolean drainWaiters = false; synchronized(this) { @@ -238,12 +245,21 @@ public void setTargetMemory(long targetMemory) { requestReservationDrain(); } + @Override + public synchronized long reclaimUnused() { + if(_shutdown || _destroyed || _grantedBytes <= _usedBytes) + return 0; + long reclaimed = _grantedBytes - _usedBytes; + _grantedBytes = _usedBytes; + notifyAll(); + return reclaimed; + } + @Override public void shutdown() { long freedMemory = 0; long destroyFreedMemory = 0; boolean destroy = false; - ArrayDeque waiters; synchronized(this) { if(_shutdown || _destroyed) return; @@ -259,8 +275,6 @@ public void shutdown() { else { freedMemory = oldGrantedBytes - _grantedBytes; } - waiters = new ArrayDeque<>(_reservationWaiters); - _reservationWaiters.clear(); notifyAll(); } _broker.shutdownAllowance(this); @@ -269,8 +283,9 @@ public void shutdown() { else if(freedMemory > 0) _broker.freeMemory(this, freedMemory); IllegalStateException ex = new IllegalStateException("Cannot reserve memory on closed allowance."); - while(!waiters.isEmpty()) - waiters.removeFirst().future.completeExceptionally(ex); + ReservationWaiter waiter; + while((waiter = _reservationWaiters.poll()) != null) + waiter.future.completeExceptionally(ex); } @Override @@ -278,6 +293,20 @@ public boolean isShutdown() { return _shutdown || _destroyed; } + void onBrokerMemoryAvailable() { + boolean drainWaiters; + synchronized(this) { + drainWaiters = !_reservationWaiters.isEmpty() && !_shutdown && !_destroyed; + notifyAll(); + } + if(drainWaiters) + requestReservationDrain(); + } + + boolean hasReservationWaiters() { + return !_reservationWaiters.isEmpty(); + } + private void requestReservationDrain() { synchronized(this) { _reservationDrainRequested = true; @@ -309,14 +338,11 @@ private void requestReservationDrain() { private void drainReservationWaitersOnce() { while(true) { - ReservationWaiter waiter; - synchronized(this) { - if(_shutdown || _destroyed) - return; - waiter = _reservationWaiters.peekFirst(); - if(waiter == null) - return; - } + if(_shutdown || _destroyed) + return; + ReservationWaiter waiter = _reservationWaiters.peek(); + if(waiter == null) + return; boolean admitted; try { admitted = tryReserve(waiter.bytes); @@ -335,11 +361,7 @@ private void drainReservationWaitersOnce() { } } - private synchronized boolean removeReservationWaiter(ReservationWaiter waiter) { - if(_reservationWaiters.peekFirst() == waiter) { - _reservationWaiters.removeFirst(); - return true; - } + private boolean removeReservationWaiter(ReservationWaiter waiter) { return _reservationWaiters.remove(waiter); } diff --git a/src/test/java/org/apache/sysds/test/component/ooc/memory/OOCMemoryAllowanceTest.java b/src/test/java/org/apache/sysds/test/component/ooc/memory/OOCMemoryAllowanceTest.java index 6ae3cd62ed8..979498502d1 100644 --- a/src/test/java/org/apache/sysds/test/component/ooc/memory/OOCMemoryAllowanceTest.java +++ b/src/test/java/org/apache/sysds/test/component/ooc/memory/OOCMemoryAllowanceTest.java @@ -30,6 +30,7 @@ import org.apache.sysds.runtime.matrix.operators.BinaryOperator; import org.apache.sysds.runtime.matrix.operators.RightScalarOperator; import org.apache.sysds.runtime.ooc.cache.OOCCacheManager; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; import org.apache.sysds.runtime.ooc.memory.CachedAllowance; import org.apache.sysds.runtime.ooc.memory.GlobalMemoryBroker; import org.apache.sysds.runtime.ooc.memory.InMemoryQueueCallback; @@ -45,6 +46,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiFunction; import java.util.function.Function; @@ -62,6 +64,54 @@ public void testWorstCase() { test(false, 0, 1); } + @Test + public void testBlockedReservationReclaims() throws Exception { + GlobalMemoryBroker broker = new GlobalMemoryBroker(100); + SyncMemoryAllowance holder = new SyncMemoryAllowance(broker); + SyncMemoryAllowance waiter = new SyncMemoryAllowance(broker); + try { + holder.reserveBlocking(100); + OOCFuture reservation = waiter.reserveAsync(50); + Assert.assertFalse(reservation.isDone()); + + holder.release(50); + reservation.get(10, TimeUnit.SECONDS); + Assert.assertEquals(50, holder.getGrantedMemory()); + Assert.assertEquals(50, waiter.getUsedMemory()); + } + finally { + holder.release(holder.getUsedMemory()); + waiter.release(waiter.getUsedMemory()); + holder.destroy(); + waiter.destroy(); + } + } + + @Test + public void testReservationWaiters() throws Exception { + CoordinatedBroker broker = new CoordinatedBroker(new GlobalMemoryBroker(100)); + SyncMemoryAllowance allowance = new SyncMemoryAllowance(broker); + try { + allowance.reserveBlocking(100); + OOCFuture first = allowance.reserveAsync(60); + Assert.assertFalse(first.isDone()); + + allowance.release(20); + OOCFuture second = allowance.reserveAsync(20); + Assert.assertFalse(second.isDone()); + + allowance.release(60); + first.get(10, TimeUnit.SECONDS); + second.get(10, TimeUnit.SECONDS); + Assert.assertEquals(100, allowance.getUsedMemory()); + } + finally { + allowance.release(allowance.getUsedMemory()); + allowance.destroy(); + broker.destroy(); + } + } + public void test(boolean optimal, int nWarmup, int nMeasure) { //DMLScript.OOC_STATISTICS = true; long millis; @@ -363,6 +413,10 @@ public void attachAllowance(MemoryAllowance allowance) { applyTargetUpdates(updates); } + @Override + public void reservationBlocked(MemoryAllowance allowance, long bytes) { + } + @Override public long requestMemory(MemoryAllowance allowance, long minSize, long maxSize) { if(!_credits.containsKey(allowance))