[CELEBORN-1792] MemoryManager resume should use pinnedDirectMemory instead of usedDirectMemory #3018
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -30,12 +30,14 @@ | |
| import com.google.common.base.Preconditions; | ||
| import io.netty.buffer.ByteBuf; | ||
| import io.netty.buffer.ByteBufAllocator; | ||
| import io.netty.buffer.PooledByteBufAllocator; | ||
| import io.netty.util.internal.PlatformDependent; | ||
| import org.slf4j.Logger; | ||
| import org.slf4j.LoggerFactory; | ||
|
|
||
| import org.apache.celeborn.common.CelebornConf; | ||
| import org.apache.celeborn.common.metrics.source.AbstractSource; | ||
| import org.apache.celeborn.common.network.util.NettyUtils; | ||
| import org.apache.celeborn.common.protocol.TransportModuleConstants; | ||
| import org.apache.celeborn.common.util.ThreadUtils; | ||
| import org.apache.celeborn.common.util.Utils; | ||
|
|
@@ -50,7 +52,8 @@ public class MemoryManager { | |
| @VisibleForTesting public long maxDirectMemory; | ||
| private final long pausePushDataThreshold; | ||
| private final long pauseReplicateThreshold; | ||
| private final double directMemoryResumeRatio; | ||
| private final double pinnedMemoryResumeRatio; | ||
| private final long maxSortMemory; | ||
| private final int forceAppendPauseSpentTimeThreshold; | ||
| private final List<MemoryPressureListener> memoryPressureListeners = new ArrayList<>(); | ||
|
|
@@ -93,6 +96,9 @@ public class MemoryManager { | |
| private long memoryFileStorageThreshold; | ||
| private final LongAdder memoryFileStorageCounter = new LongAdder(); | ||
| private final StorageManager storageManager; | ||
| private boolean pinnedMemoryCheckEnabled; | ||
| private long pinnedMemoryCheckInterval; | ||
|
There was a problem hiding this comment. To avoid frequently calling a pinned memory counter, I think you can cache the last pinned memory value and refresh it periodically. And exporting the pinned memory value to the metrics. There was a problem hiding this comment. Here is another PR introducing pinnedMemory metrics #3019 There was a problem hiding this comment. getPinnedMemory is not called very frequently. It is called once every pinnedMemoryCheckInterval. The default is 10 seconds. |
||
| private long pinnedMemoryLastCheckTime = 0L; | ||
|
There was a problem hiding this comment. The default value for this is 0. |
||
|
|
||
| @VisibleForTesting | ||
| public static MemoryManager initialize(CelebornConf conf) { | ||
|
|
@@ -120,11 +126,14 @@ public static MemoryManager instance() { | |
| private MemoryManager(CelebornConf conf, StorageManager storageManager, AbstractSource source) { | ||
| double pausePushDataRatio = conf.workerDirectMemoryRatioToPauseReceive(); | ||
| double pauseReplicateRatio = conf.workerDirectMemoryRatioToPauseReplicate(); | ||
| this.directMemoryResumeRatio = conf.workerDirectMemoryRatioToResume(); | ||
| this.pinnedMemoryResumeRatio = conf.workerPinnedMemoryRatioToResume(); | ||
| double maxSortMemRatio = conf.workerPartitionSorterDirectMemoryRatioThreshold(); | ||
| double readBufferRatio = conf.workerDirectMemoryRatioForReadBuffer(); | ||
| double memoryFileStorageRatio = conf.workerDirectMemoryRatioForMemoryFilesStorage(); | ||
| long checkInterval = conf.workerDirectMemoryPressureCheckIntervalMs(); | ||
| this.pinnedMemoryCheckEnabled = conf.workerPinnedMemoryCheckEnabled(); | ||
| this.pinnedMemoryCheckInterval = conf.workerPinnedMemoryCheckIntervalMs(); | ||
| long reportInterval = conf.workerDirectMemoryReportIntervalSecond(); | ||
| double readBufferTargetRatio = conf.readBufferTargetRatio(); | ||
| long readBufferTargetUpdateInterval = conf.readBufferTargetUpdateInterval(); | ||
|
|
@@ -148,9 +157,10 @@ private MemoryManager(CelebornConf conf, StorageManager storageManager, Abstract | |
| pauseReplicateRatio, | ||
| CelebornConf.WORKER_DIRECT_MEMORY_RATIO_PAUSE_RECEIVE().key(), | ||
| pausePushDataRatio)); | ||
| Preconditions.checkArgument(pausePushDataRatio > directMemoryResumeRatio); | ||
| if (memoryFileStorageRatio > 0) { | ||
| Preconditions.checkArgument( | ||
| directMemoryResumeRatio > (readBufferRatio + memoryFileStorageRatio)); | ||
| } | ||
|
|
||
| maxSortMemory = ((long) (maxDirectMemory * maxSortMemRatio)); | ||
|
|
@@ -282,7 +292,7 @@ private MemoryManager(CelebornConf conf, StorageManager storageManager, Abstract | |
| Utils.bytesToString(readBufferThreshold), | ||
| Utils.bytesToString(readBufferTarget), | ||
| Utils.bytesToString(memoryFileStorageThreshold), | ||
| directMemoryResumeRatio); | ||
|
There was a problem hiding this comment. You can add pinned memory resume ratio here. It is an important parameter for memory manager. |
||
| } | ||
|
|
||
| public boolean shouldEvict(boolean aggressiveMemoryFileEvictEnabled, double evictRatio) { | ||
|
|
@@ -305,7 +315,7 @@ public ServingState currentServingState() { | |
| return ServingState.PUSH_PAUSED; | ||
| } | ||
| // trigger resume | ||
| if (memoryUsage / (double) (maxDirectMemory) < directMemoryResumeRatio) { | ||
| isPaused = false; | ||
| return ServingState.NONE_PAUSED; | ||
| } | ||
|
|
@@ -315,69 +325,68 @@ public ServingState currentServingState() { | |
| } | ||
|
|
||
| @VisibleForTesting | ||
| public void switchServingState() { | ||
| ServingState lastState = servingState; | ||
| servingState = currentServingState(); | ||
| logger.info("Serving state transformed from {} to {}", lastState, servingState); | ||
| switch (servingState) { | ||
| case PUSH_PAUSED: | ||
| if (canResumeByPinnedMemory()) { | ||
| resumeByPinnedMemory(servingState); | ||
| } else { | ||
| pausePushDataCounter.increment(); | ||
| if (lastState == ServingState.PUSH_AND_REPLICATE_PAUSED) { | ||
| logger.info("Trigger action: RESUME REPLICATE"); | ||
| resumeReplicate(); | ||
| } else { | ||
| logger.info("Trigger action: PAUSE PUSH"); | ||
| pausePushDataStartTime = System.currentTimeMillis(); | ||
| memoryPressureListeners.forEach( | ||
| memoryPressureListener -> | ||
| memoryPressureListener.onPause(TransportModuleConstants.PUSH_MODULE)); | ||
| } | ||
| } | ||
| logger.debug("Trigger action: TRIM"); | ||
| trimAllListeners(); | ||
| if (trimCounter >= forceAppendPauseSpentTimeThreshold) { | ||
|
|
||
| logger.debug( | ||
| "Trigger action: TRIM for {} times, force to append pause spent time.", trimCounter); | ||
| appendPauseSpentTime(servingState); | ||
| } | ||
| break; | ||
| case PUSH_AND_REPLICATE_PAUSED: | ||
| if (canResumeByPinnedMemory()) { | ||
| resumeByPinnedMemory(servingState); | ||
| } else { | ||
| pausePushDataAndReplicateCounter.increment(); | ||
| logger.info("Trigger action: PAUSE PUSH"); | ||
| pausePushDataAndReplicateStartTime = System.currentTimeMillis(); | ||
| memoryPressureListeners.forEach( | ||
| memoryPressureListener -> | ||
| memoryPressureListener.onPause(TransportModuleConstants.PUSH_MODULE)); | ||
| logger.info("Trigger action: PAUSE REPLICATE"); | ||
| memoryPressureListeners.forEach( | ||
| memoryPressureListener -> | ||
| memoryPressureListener.onPause(TransportModuleConstants.REPLICATE_MODULE)); | ||
| } | ||
| logger.debug("Trigger action: TRIM"); | ||
| trimAllListeners(); | ||
|
|
||
| if (trimCounter >= forceAppendPauseSpentTimeThreshold) { | ||
| logger.debug( | ||
| "Trigger action: TRIM for {} times, force to append pause spent time.", trimCounter); | ||
| appendPauseSpentTime(servingState); | ||
| } | ||
| break; | ||
| case NONE_PAUSED: | ||
| // resume from paused mode, append pause spent time | ||
| if (lastState == ServingState.PUSH_AND_REPLICATE_PAUSED) { | ||
| resumeReplicate(); | ||
| resumePush(); | ||
| appendPauseSpentTime(lastState); | ||
| } else if (lastState == ServingState.PUSH_PAUSED) { | ||
| resumePush(); | ||
| appendPauseSpentTime(lastState); | ||
| } | ||
| } | ||
| } | ||
|
|
||
|
|
@@ -436,6 +445,16 @@ public long getMemoryUsage() { | |
| return getNettyUsedDirectMemory() + sortMemoryCounter.get(); | ||
| } | ||
|
|
||
| public long getAllocatedMemory() { | ||
|
There was a problem hiding this comment. This method should be renamed to getPinnedMemory. The allocated memory is the netty memory counter. |
||
| return getNettyPinnedDirectMemory() + sortMemoryCounter.get(); | ||
| } | ||
|
|
||
| public long getNettyPinnedDirectMemory() { | ||
| return NettyUtils.getAllPooledByteBufAllocators().stream() | ||
| .mapToLong(PooledByteBufAllocator::pinnedDirectMemory) | ||
| .sum(); | ||
| } | ||
|
|
||
| public AtomicLong getSortMemoryCounter() { | ||
| return sortMemoryCounter; | ||
| } | ||
|
|
@@ -557,6 +576,47 @@ public static void reset() { | |
| _INSTANCE = null; | ||
| } | ||
|
|
||
| private void resumeByPinnedMemory(ServingState servingState) { | ||
| switch (servingState) { | ||
| case PUSH_AND_REPLICATE_PAUSED: | ||
| logger.info( | ||
| "Serving State is PUSH_AND_REPLICATE_PAUSED, but resume by lower pinned memory {}", | ||
| getNettyPinnedDirectMemory()); | ||
| resumeReplicate(); | ||
| resumePush(); | ||
| case PUSH_PAUSED: | ||
| logger.info( | ||
| "Serving State is PUSH_PAUSED, but resume by lower pinned memory {}", | ||
| getNettyPinnedDirectMemory()); | ||
| resumePush(); | ||
| } | ||
| } | ||
|
|
||
| private boolean canResumeByPinnedMemory() { | ||
| if (pinnedMemoryCheckEnabled | ||
| && System.currentTimeMillis() - pinnedMemoryLastCheckTime >= pinnedMemoryCheckInterval | ||
| && getAllocatedMemory() / (double) (maxDirectMemory) < pinnedMemoryResumeRatio) { | ||
| pinnedMemoryLastCheckTime = System.currentTimeMillis(); | ||
| return true; | ||
| } else { | ||
| return false; | ||
| } | ||
| } | ||
|
|
||
| private void resumePush() { | ||
| logger.info("Trigger action: RESUME PUSH"); | ||
| memoryPressureListeners.forEach( | ||
| memoryPressureListener -> | ||
| memoryPressureListener.onResume(TransportModuleConstants.PUSH_MODULE)); | ||
| } | ||
|
|
||
| private void resumeReplicate() { | ||
| logger.info("Trigger action: RESUME REPLICATE"); | ||
| memoryPressureListeners.forEach( | ||
| memoryPressureListener -> | ||
| memoryPressureListener.onResume(TransportModuleConstants.REPLICATE_MODULE)); | ||
| } | ||
|
|
||
| public interface MemoryPressureListener { | ||
| void onPause(String moduleName); | ||
|
|
||
|
|
||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Maybe we can add a new conf for pinnedMemoryToResume and keep exist conf for directMemoryRatioToResume