Skip to content
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

[CELEBORN-1815] Support UnpooledByteBufAllocator #3043

Closed
wants to merge 3 commits into from
Closed

[CELEBORN-1815] Support UnpooledByteBufAllocator #3043

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e2de372
[CELEBORN-1815] Support UnpooledByteBufAllocator
pan3793 Dec 31, 2024
7dc3531
fix flink
pan3793 Dec 31, 2024
b3d569a
comment
pan3793 Jan 1, 2025
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
9 changes: 4 additions & 5 deletions common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ private static class ClientPool {
private final int sendBuf;
private final Class<? extends Channel> socketChannelClass;
private EventLoopGroup workerGroup;
protected ByteBufAllocator pooledAllocator;
protected ByteBufAllocator allocator;
private final int maxClientConnectRetries;
private final int maxClientConnectRetryWaitTimeMs;

Expand All @@ -115,9 +115,8 @@ public TransportClientFactory(
logger.info("Module {} mode {} threads {}", conf.getModuleName(), ioMode, conf.clientThreads());
this.workerGroup =
NettyUtils.createEventLoop(ioMode, conf.clientThreads(), conf.getModuleName() + "-client");
this.pooledAllocator =
NettyUtils.getPooledByteBufAllocator(
conf, context.getSource(), false, conf.clientThreads());
this.allocator =
NettyUtils.getByteBufAllocator(conf, context.getSource(), false, conf.clientThreads());
this.maxClientConnectRetries = conf.maxIORetries();
this.maxClientConnectRetryWaitTimeMs = conf.ioRetryWaitTimeMs();
}
Expand Down Expand Up @@ -268,7 +267,7 @@ private TransportClient internalCreateClient(
.option(ChannelOption.TCP_NODELAY, true)
.option(ChannelOption.SO_KEEPALIVE, true)
.option(ChannelOption.CONNECT_TIMEOUT_MILLIS, connectTimeoutMs)
.option(ChannelOption.ALLOCATOR, pooledAllocator);
.option(ChannelOption.ALLOCATOR, allocator);

if (receiveBuf > 0) {
bootstrap.option(ChannelOption.SO_RCVBUF, receiveBuf);
Expand Down
6 changes: 3 additions & 3 deletions common/src/main/java/org/apache/celeborn/common/network/server/TransportServer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
import io.netty.bootstrap.ServerBootstrap;
import io.netty.buffer.PooledByteBufAllocator;
import io.netty.buffer.ByteBufAllocator;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelInitializer;
import io.netty.channel.ChannelOption;
Expand Down Expand Up @@ -91,8 +91,8 @@ private void init(String hostToBind, int portToBind) {
EventLoopGroup workerGroup =
NettyUtils.createEventLoop(ioMode, conf.serverThreads(), conf.getModuleName() + "-server");

PooledByteBufAllocator allocator =
NettyUtils.getPooledByteBufAllocator(conf, source, true, conf.serverThreads());
ByteBufAllocator allocator =
NettyUtils.getByteBufAllocator(conf, source, true, conf.serverThreads());

bootstrap =
new ServerBootstrap()
Expand Down
126 changes: 70 additions & 56 deletions common/src/main/java/org/apache/celeborn/common/network/util/NettyMemoryMetrics.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,19 +23,17 @@

import com.codahale.metrics.MetricRegistry;
import com.google.common.annotations.VisibleForTesting;
import io.netty.buffer.PoolArenaMetric;
import io.netty.buffer.PooledByteBufAllocator;
import io.netty.buffer.PooledByteBufAllocatorMetric;
import io.netty.buffer.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.common.metrics.source.AbstractSource;

/** A Netty memory metrics class to collect metrics from Netty PooledByteBufAllocator. */
/** A Netty memory metrics class to collect metrics from Netty ByteBufAllocator. */
public class NettyMemoryMetrics {
private final Logger logger = LoggerFactory.getLogger(this.getClass());

private final PooledByteBufAllocator pooledAllocator;
private final ByteBufAllocator allocator;

private final boolean verboseMetricsEnabled;

Expand Down Expand Up @@ -69,71 +67,87 @@ public class NettyMemoryMetrics {
}

public NettyMemoryMetrics(
PooledByteBufAllocator pooledAllocator,
ByteBufAllocator allocator,
String metricPrefix,
boolean verboseMetricsEnabled,
AbstractSource source,
Map<String, String> labels) {
this.pooledAllocator = pooledAllocator;
this.allocator = allocator;
this.metricPrefix = metricPrefix;
this.verboseMetricsEnabled = verboseMetricsEnabled;
this.source = source;
this.labels = labels;

registerMetrics(this.pooledAllocator);
registerMetrics();
}

private void registerMetrics(PooledByteBufAllocator allocator) {
PooledByteBufAllocatorMetric pooledAllocatorMetric = allocator.metric();

private void registerMetrics() {
// Register general metrics.
if (source != null) {
logger.debug("setup netty metrics");
source.addGauge(
MetricRegistry.name(metricPrefix, "usedHeapMemory"),
labels,
pooledAllocatorMetric::usedHeapMemory);
source.addGauge(
MetricRegistry.name(metricPrefix, "usedDirectMemory"),
labels,
pooledAllocatorMetric::usedDirectMemory);
source.addGauge(
MetricRegistry.name(metricPrefix, "numHeapArenas"),
labels,
pooledAllocatorMetric::numHeapArenas);
source.addGauge(
MetricRegistry.name(metricPrefix, "numDirectArenas"),
labels,
pooledAllocatorMetric::numDirectArenas);
source.addGauge(
MetricRegistry.name(metricPrefix, "tinyCacheSize"),
labels,
pooledAllocatorMetric::tinyCacheSize);
source.addGauge(
MetricRegistry.name(metricPrefix, "smallCacheSize"),
labels,
pooledAllocatorMetric::smallCacheSize);
source.addGauge(
MetricRegistry.name(metricPrefix, "normalCacheSize"),
labels,
pooledAllocatorMetric::normalCacheSize);
source.addGauge(
MetricRegistry.name(metricPrefix, "numThreadLocalCaches"),
labels,
pooledAllocatorMetric::numThreadLocalCaches);
source.addGauge(
MetricRegistry.name(metricPrefix, "chunkSize"), labels, pooledAllocatorMetric::chunkSize);
if (verboseMetricsEnabled) {
int directArenaIndex = 0;
for (PoolArenaMetric metric : pooledAllocatorMetric.directArenas()) {
registerArenaMetric(metric, "directArena" + directArenaIndex);
directArenaIndex++;
}
if (allocator instanceof UnpooledByteBufAllocator) {
logger.debug("Setup netty metrics for UnpooledByteBufAllocator");
ByteBufAllocatorMetric unpooledMetric = ((UnpooledByteBufAllocator) allocator).metric();
source.addGauge(
MetricRegistry.name(metricPrefix, "usedHeapMemory"),
labels,
unpooledMetric::usedHeapMemory);
source.addGauge(
MetricRegistry.name(metricPrefix, "usedDirectMemory"),
labels,
unpooledMetric::usedDirectMemory);
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

UnpooledByteBufAllocator only supports these two metrics

} else if (allocator instanceof PooledByteBufAllocator) {
logger.debug("Setup netty metrics for PooledByteBufAllocator");
PooledByteBufAllocatorMetric pooledAllocatorMetric =
((PooledByteBufAllocator) allocator).metric();
source.addGauge(
MetricRegistry.name(metricPrefix, "usedHeapMemory"),
labels,
pooledAllocatorMetric::usedHeapMemory);
source.addGauge(
MetricRegistry.name(metricPrefix, "usedDirectMemory"),
labels,
pooledAllocatorMetric::usedDirectMemory);

int heapArenaIndex = 0;
for (PoolArenaMetric metric : pooledAllocatorMetric.heapArenas()) {
registerArenaMetric(metric, "heapArena" + heapArenaIndex);
heapArenaIndex++;
source.addGauge(
MetricRegistry.name(metricPrefix, "numHeapArenas"),
labels,
pooledAllocatorMetric::numHeapArenas);
source.addGauge(
MetricRegistry.name(metricPrefix, "numDirectArenas"),
labels,
pooledAllocatorMetric::numDirectArenas);
source.addGauge(
MetricRegistry.name(metricPrefix, "tinyCacheSize"),
labels,
pooledAllocatorMetric::tinyCacheSize);
source.addGauge(
MetricRegistry.name(metricPrefix, "smallCacheSize"),
labels,
pooledAllocatorMetric::smallCacheSize);
source.addGauge(
MetricRegistry.name(metricPrefix, "normalCacheSize"),
labels,
pooledAllocatorMetric::normalCacheSize);
source.addGauge(
MetricRegistry.name(metricPrefix, "numThreadLocalCaches"),
labels,
pooledAllocatorMetric::numThreadLocalCaches);
source.addGauge(
MetricRegistry.name(metricPrefix, "chunkSize"),
labels,
pooledAllocatorMetric::chunkSize);
if (verboseMetricsEnabled) {
int directArenaIndex = 0;
for (PoolArenaMetric metric : pooledAllocatorMetric.directArenas()) {
registerArenaMetric(metric, "directArena" + directArenaIndex);
directArenaIndex++;
}

int heapArenaIndex = 0;
for (PoolArenaMetric metric : pooledAllocatorMetric.heapArenas()) {
registerArenaMetric(metric, "heapArena" + heapArenaIndex);
heapArenaIndex++;
}
}
}
}
Expand Down
67 changes: 39 additions & 28 deletions common/src/main/java/org/apache/celeborn/common/network/util/NettyUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,9 @@
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ThreadFactory;

import io.netty.buffer.ByteBufAllocator;
import io.netty.buffer.PooledByteBufAllocator;
import io.netty.buffer.UnpooledByteBufAllocator;
import io.netty.channel.Channel;
import io.netty.channel.EventLoopGroup;
import io.netty.channel.ServerChannel;
Expand All @@ -42,9 +44,8 @@

/** Utilities for creating various Netty constructs based on whether we're using EPOLL or NIO. */
public class NettyUtils {
private static final PooledByteBufAllocator[] _sharedPooledByteBufAllocator =
new PooledByteBufAllocator[2];
private static ConcurrentHashMap<String, Integer> allocatorsIndex =
private static final ByteBufAllocator[] _sharedByteBufAllocator = new ByteBufAllocator[2];
private static final ConcurrentHashMap<String, Integer> allocatorsIndex =
JavaUtils.newConcurrentHashMap();
/** Creates a new ThreadFactory which prefixes each thread with the given name. */
public static ThreadFactory createThreadFactory(String threadPoolPrefix) {
Expand Down Expand Up @@ -103,53 +104,59 @@ public static String getRemoteAddress(Channel channel) {
* released by the executor thread rather than the event loop thread. Those thread-local caches
* actually delay the recycling of buffers, leading to larger memory usage.
*/
private static PooledByteBufAllocator createPooledByteBufAllocator(
boolean allowDirectBufs, boolean allowCache, int numCores) {
if (numCores == 0) {
numCores = Runtime.getRuntime().availableProcessors();
private static ByteBufAllocator createByteBufAllocator(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The method comments should also be updated

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

update: move the existing comment to the caller side, add new comments to explain each parameters

boolean pooled, boolean allowDirectBufs, boolean allowCache, int numCores) {
if (pooled) {
if (numCores == 0) {
numCores = Runtime.getRuntime().availableProcessors();
}
return new PooledByteBufAllocator(
allowDirectBufs && PlatformDependent.directBufferPreferred(),
Math.min(PooledByteBufAllocator.defaultNumHeapArena(), numCores),
Math.min(PooledByteBufAllocator.defaultNumDirectArena(), allowDirectBufs ? numCores : 0),
PooledByteBufAllocator.defaultPageSize(),
PooledByteBufAllocator.defaultMaxOrder(),
allowCache ? PooledByteBufAllocator.defaultSmallCacheSize() : 0,
allowCache ? PooledByteBufAllocator.defaultNormalCacheSize() : 0,
allowCache && PooledByteBufAllocator.defaultUseCacheForAllThreads());
} else {
return new UnpooledByteBufAllocator(
allowDirectBufs && PlatformDependent.directBufferPreferred());
}
return new PooledByteBufAllocator(
allowDirectBufs && PlatformDependent.directBufferPreferred(),
Math.min(PooledByteBufAllocator.defaultNumHeapArena(), numCores),
Math.min(PooledByteBufAllocator.defaultNumDirectArena(), allowDirectBufs ? numCores : 0),
PooledByteBufAllocator.defaultPageSize(),
PooledByteBufAllocator.defaultMaxOrder(),
allowCache ? PooledByteBufAllocator.defaultSmallCacheSize() : 0,
allowCache ? PooledByteBufAllocator.defaultNormalCacheSize() : 0,
allowCache && PooledByteBufAllocator.defaultUseCacheForAllThreads());
}

/**
* Returns the lazily created shared pooled ByteBuf allocator for the specified allowCache
* parameter value.
*/
public static synchronized PooledByteBufAllocator getSharedPooledByteBufAllocator(
public static synchronized ByteBufAllocator getSharedByteBufAllocator(
CelebornConf conf, AbstractSource source, boolean allowCache) {
final int index = allowCache ? 0 : 1;
if (_sharedPooledByteBufAllocator[index] == null) {
_sharedPooledByteBufAllocator[index] =
createPooledByteBufAllocator(true, allowCache, conf.networkAllocatorArenas());
if (_sharedByteBufAllocator[index] == null) {
_sharedByteBufAllocator[index] =
createByteBufAllocator(
conf.networkMemoryAllocatorPooled(), true, allowCache, conf.networkAllocatorArenas());
if (source != null) {
new NettyMemoryMetrics(
_sharedPooledByteBufAllocator[index],
_sharedByteBufAllocator[index],
"shared-pool-" + index,
conf.networkAllocatorVerboseMetric(),
source,
Collections.emptyMap());
}
}
return _sharedPooledByteBufAllocator[index];
return _sharedByteBufAllocator[index];
}

public static PooledByteBufAllocator getPooledByteBufAllocator(
public static ByteBufAllocator getByteBufAllocator(
TransportConf conf, AbstractSource source, boolean allowCache) {
return getPooledByteBufAllocator(conf, source, allowCache, 0);
return getByteBufAllocator(conf, source, allowCache, 0);
}

public static PooledByteBufAllocator getPooledByteBufAllocator(
public static ByteBufAllocator getByteBufAllocator(
TransportConf conf, AbstractSource source, boolean allowCache, int coreNum) {
if (conf.getCelebornConf().networkShareMemoryAllocator()) {
return getSharedPooledByteBufAllocator(
return getSharedByteBufAllocator(
conf.getCelebornConf(),
source,
allowCache && conf.getCelebornConf().networkMemoryAllocatorAllowCache());
Expand All @@ -160,8 +167,12 @@ public static PooledByteBufAllocator getPooledByteBufAllocator(
} else {
arenas = conf.getCelebornConf().networkAllocatorArenas();
}
PooledByteBufAllocator allocator =
createPooledByteBufAllocator(conf.preferDirectBufs(), allowCache, arenas);
ByteBufAllocator allocator =
createByteBufAllocator(
conf.getCelebornConf().networkMemoryAllocatorPooled(),
conf.preferDirectBufs(),
allowCache,
arenas);
if (source != null) {
String poolName = "default-netty-pool";
Map<String, String> labels = new HashMap<>();
Expand Down
14 changes: 14 additions & 0 deletions common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -603,6 +603,9 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
def networkMemoryAllocatorAllowCache: Boolean =
get(NETWORK_MEMORY_ALLOCATOR_ALLOW_CACHE)

def networkMemoryAllocatorPooled: Boolean =
get(NETWORK_MEMORY_ALLOCATOR_POOLED)

def networkAllocatorArenas: Int = get(NETWORK_MEMORY_ALLOCATOR_ARENAS).getOrElse(Math.max(
Runtime.getRuntime.availableProcessors(),
2))
Expand Down Expand Up @@ -1782,6 +1785,17 @@ object CelebornConf extends Logging {
.booleanConf
.createWithDefault(false)

val NETWORK_MEMORY_ALLOCATOR_POOLED: ConfigEntry[Boolean] =
buildConf("celeborn.network.memory.allocator.pooled")
.categories("network")
.internal
.version("0.6.0")
.doc("If disabled, always use UnpooledByteBufAllocator for aggressive memory reclamation, " +
"this is helpful for cases that worker has high memory usage even after triming. " +
"Disabling would cause performace degression and higher CPU usage.")
.booleanConf
.createWithDefault(true)

val NETWORK_MEMORY_ALLOCATOR_SHARE: ConfigEntry[Boolean] =
buildConf("celeborn.network.memory.allocator.share")
.categories("network")
Expand Down
4 changes: 2 additions & 2 deletions worker/src/main/java/org/apache/celeborn/service/deploy/worker/memory/MemoryManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.PooledByteBufAllocator;
import io.netty.buffer.ByteBufAllocator;
import io.netty.util.internal.PlatformDependent;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
Expand Down Expand Up @@ -546,7 +546,7 @@ public void close() {
readBufferDispatcher.close();
}

public PooledByteBufAllocator getStoragePooledByteBufAllocator() {
public ByteBufAllocator getStorageByteBufAllocator() {
return storageManager.storageBufferAllocator();
}

Expand Down
Loading
Loading