Skip to content

Commit c5b58b5

Browse files
committed
[SPARK-54023][CORE] Support AUTO IO Mode
### What changes were proposed in this pull request? This PR aims to support a new Netty IO Mode, `AUTO`, on top of the existing `NIO`, `EPOLL`, and `KQUEUE`. `AUTO` mode prefers to use native `EPOLL` mode on Linux and `KQUEUE` mode on MacOS if available. Then, it fallbacks to `NIO` mode. ### Why are the changes needed? To help a user to try to use native IO mode more easily. ### Does this PR introduce _any_ user-facing change? No, this is a new IO mode. ### How was this patch tested? Pass the CIs with newly added test suite, `ShuffleNettyAutoSuite`. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #52724 from dongjoon-hyun/SPARK-54023. Authored-by: Dongjoon Hyun <dongjoon@apache.org> Signed-off-by: Dongjoon Hyun <dongjoon@apache.org>
1 parent 06674ba commit c5b58b5

File tree

4 files changed

+40
-3
lines changed

4 files changed

+40
-3
lines changed

common/network-common/src/main/java/org/apache/spark/network/util/IOMode.java

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,5 +32,9 @@ public enum IOMode {
3232
/**
3333
* Native KQUEUE via JNI, MacOS/BSD only
3434
*/
35-
KQUEUE
35+
KQUEUE,
36+
/**
37+
* Prefer to use native EPOLL on Linux (or KQUEUE on MacOS) if available. Then, fallback to NIO.
38+
*/
39+
AUTO
3640
}

common/network-common/src/main/java/org/apache/spark/network/util/NettyUtils.java

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,11 @@
2121

2222
import io.netty.buffer.PooledByteBufAllocator;
2323
import io.netty.channel.*;
24+
import io.netty.channel.epoll.Epoll;
2425
import io.netty.channel.epoll.EpollIoHandler;
2526
import io.netty.channel.epoll.EpollServerSocketChannel;
2627
import io.netty.channel.epoll.EpollSocketChannel;
28+
import io.netty.channel.kqueue.KQueue;
2729
import io.netty.channel.kqueue.KQueueIoHandler;
2830
import io.netty.channel.kqueue.KQueueServerSocketChannel;
2931
import io.netty.channel.kqueue.KQueueSocketChannel;
@@ -35,7 +37,7 @@
3537

3638
/**
3739
* Utilities for creating various Netty constructs based on whether we're using NIO, EPOLL,
38-
* or KQUEUE.
40+
* , KQUEUE, or AUTO.
3941
*/
4042
public class NettyUtils {
4143

@@ -71,6 +73,15 @@ public static EventLoopGroup createEventLoop(IOMode mode, int numThreads, String
7173
case NIO -> NioIoHandler.newFactory();
7274
case EPOLL -> EpollIoHandler.newFactory();
7375
case KQUEUE -> KQueueIoHandler.newFactory();
76+
case AUTO -> {
77+
if (JavaUtils.isLinux && Epoll.isAvailable()) {
78+
yield EpollIoHandler.newFactory();
79+
} else if (JavaUtils.isMac && KQueue.isAvailable()) {
80+
yield KQueueIoHandler.newFactory();
81+
} else {
82+
yield NioIoHandler.newFactory();
83+
}
84+
}
7485
};
7586
return new MultiThreadIoEventLoopGroup(numThreads, threadFactory, handlerFactory);
7687
}
@@ -81,6 +92,15 @@ public static Class<? extends Channel> getClientChannelClass(IOMode mode) {
8192
case NIO -> NioSocketChannel.class;
8293
case EPOLL -> EpollSocketChannel.class;
8394
case KQUEUE -> KQueueSocketChannel.class;
95+
case AUTO -> {
96+
if (JavaUtils.isLinux && Epoll.isAvailable()) {
97+
yield EpollSocketChannel.class;
98+
} else if (JavaUtils.isMac && KQueue.isAvailable()) {
99+
yield KQueueSocketChannel.class;
100+
} else {
101+
yield NioSocketChannel.class;
102+
}
103+
}
84104
};
85105
}
86106

@@ -90,6 +110,15 @@ public static Class<? extends ServerChannel> getServerChannelClass(IOMode mode)
90110
case NIO -> NioServerSocketChannel.class;
91111
case EPOLL -> EpollServerSocketChannel.class;
92112
case KQUEUE -> KQueueServerSocketChannel.class;
113+
case AUTO -> {
114+
if (JavaUtils.isLinux && Epoll.isAvailable()) {
115+
yield EpollServerSocketChannel.class;
116+
} else if (JavaUtils.isMac && KQueue.isAvailable()) {
117+
yield KQueueServerSocketChannel.class;
118+
} else {
119+
yield NioServerSocketChannel.class;
120+
}
121+
}
93122
};
94123
}
95124

common/network-common/src/main/java/org/apache/spark/network/util/TransportConf.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ public String getModuleName() {
8787
return module;
8888
}
8989

90-
/** IO mode: NIO, EPOLL, or KQUEUE */
90+
/** IO mode: NIO, EPOLL, KQUEUE, or AUTO */
9191
public String ioMode() {
9292
String defaultIOMode = conf.get(SPARK_NETWORK_DEFAULT_IO_MODE_KEY, "NIO");
9393
return conf.get(SPARK_NETWORK_IO_MODE_KEY, defaultIOMode).toUpperCase(Locale.ROOT);

core/src/test/scala/org/apache/spark/ShuffleNettySuite.scala

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,3 +55,7 @@ class ShuffleNettyKQueueSuite extends ShuffleNettySuite {
5555
override def shouldRunTests: Boolean = Utils.isMac
5656
override def ioMode: IOMode = IOMode.KQUEUE
5757
}
58+
59+
class ShuffleNettyAutoSuite extends ShuffleNettySuite {
60+
override def ioMode: IOMode = IOMode.AUTO
61+
}

0 commit comments

Comments
 (0)