Skip to content

[SYSTEMDS-3891] OOC Memory Allowance/Broker Improvements #2551

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
janniklinde wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[SYSTEMDS-3891] OOC Memory Allowance/Broker Improvements #2551

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
162 changes: 93 additions & 69 deletions src/main/java/org/apache/sysds/runtime/ooc/memory/GlobalMemoryBroker.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand All @@ -35,18 +42,28 @@ public static GlobalMemoryBroker get() {
}

private final long _allowedBytes;
private final List<MemoryAllowance> _allowances;
private final LinkedList<MemoryAllowance> _overconsumers;
private final CopyOnWriteArrayList<MemoryAllowance> _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
Expand All @@ -56,76 +73,41 @@ 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)
applyTargetUpdates(updates);
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<TargetUpdate> updates = null;
List<TargetUpdate> 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<TargetUpdate> updates;
synchronized(this) {
_overconsumers.remove(allowance);
updates = rebalance(true);
}
applyTargetUpdates(updates);
notifyReservationWaiters();
}

@Override
Expand All @@ -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
Expand All @@ -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<TargetUpdate> 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<TargetUpdate> rebalance(boolean force) {
long free = _allowedBytes - _usedBytes;
if(force)
Expand All @@ -158,6 +189,13 @@ private List<TargetUpdate> rebalance(boolean force) {
return switchBrokerMode(BrokerMode.STRICT);
}

private List<TargetUpdate> rebalanceAfterFree() {
long free = _allowedBytes - _usedBytes;
if(_brokerMode == BrokerMode.RELAXED && free > _allowedBytes / 5)
return rebalanceToRelaxed();
return rebalance(false);
}

private List<TargetUpdate> switchBrokerMode(BrokerMode newMode) {
if(newMode == _brokerMode)
return null;
Expand All @@ -181,7 +219,6 @@ private List<TargetUpdate> rebalanceToStrict() {
Math.min(allowance.getTargetMemory(), share + (long) ((allowance.getUsedMemory() - share) * 0.9))));
}
}
refreshOverconsumers(updates);
return updates;
}

Expand All @@ -193,34 +230,21 @@ private List<TargetUpdate> 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<TargetUpdate> 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();
}
}

Expand Down
4 changes: 4 additions & 0 deletions src/main/java/org/apache/sysds/runtime/ooc/memory/MemoryAllowance.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,4 +53,8 @@ default long getFreeMemory() {
default boolean isUnderPressure() {
return getGrantedMemory() > getTargetMemory();
}

default long reclaimUnused() {
return 0;
}
}
6 changes: 6 additions & 0 deletions src/main/java/org/apache/sysds/runtime/ooc/memory/MemoryBroker.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Loading
Loading