Java虚拟线程AI推理实战

// 虚拟线程执行HTTP调用的伪代码
voidhandleRequest(VirtualThreadvt,Requestreq){
// 虚拟线程执行这段代码
Connectionconn=socket.connect(remote);// 阻塞I/O

// —- JVM在这里将vt挂起 —-
// 虚拟线程A的栈帧被压缩，存到堆中
// Carrier线程去执行虚拟线程B、C、D…

// —- I/O完成，JVM恢复vt —-
Responseresp=conn.read();// 继续执行
}

调度器

阻塞检测

CPU密集任务

调试

栈大小

# Python asyncio协程内存测试
# 测量10000个”休眠1小时”协程的实际内存
importasyncio,sys,tracemalloc
tracemalloc.start()
asyncdefdummy():
awaitasyncio.sleep(3600)
asyncdefmain():
tasks=[asyncio.create_task(dummy())for_inrange(10_000)]
awaitasyncio.sleep(1)# 让所有协程进入等待
current,peak=tracemalloc.get_traced_memory()
print(f“10万asyncio协程: {current/1024/1024:.1f} MB 当前, {peak/1024/1024:.1f} MB 峰值“)
# 10万协程 ≈ 200MB+
asyncio.run(main())

// Java虚拟线程内存测试
importjava.time.*;
importjava.util.concurrent.*;
importjava.lang.management.*;
publicclassVThreadMemoryTest{
publicstaticvoidmain(String[]args)throwsException{
MemoryMXBeanmemoryBean=ManagementFactory.getMemoryMXBean();

try(varexecutor=Executors.newVirtualThreadPerTaskExecutor()){
for(inti=0;i<100_000;i++){
executor.submit(()–>{
try{Thread.sleep(Duration.ofHours(1));}
catch(InterruptedExceptione){}
});
}

Thread.sleep(500);
longheapUsed=memoryBean.getHeapMemoryUsage().getUsed();
System.out.printf(“10万虚拟线程: %.1f MB 堆占用%n“,
heapUsed/1024.0/1024.0);
// 10万虚拟线程 ≈ 80MB（堆占用几乎不增加）
}
}
}

src/main/java/com/ai/inference/
├── config/
│   ├── VirtualThreadConfig.java      # 虚拟线程执行器配置
│   └── TomcatVirtualThreadConfig.java
├── controller/
│   └── InferenceController.java     # REST API（同步/异步/批处理）
├── service/
│   ├── InferenceService.java        # 推理服务接口
│   ├── OnnxInferenceService.java    # ONNX Runtime实现
│   └── HttpClientService.java       # 下游服务调用（RestTemplate）
├── model/
│   ├── InferenceRequest.java        # 请求DTO
│   ├── InferenceResult.java         # 结果DTO
│   └── BatchTask.java               # 批处理任务
└── scheduler/
   └── DynamicBatchScheduler.java   # 动态batch聚合调度器

packagecom.ai.inference.config;
importorg.apache.coyote.ProtocolHandler;
importorg.apache.coyote.http11.Http11NioProtocol;
importorg.springframework.boot.web.embedded.tomcat.TomcatProtocolHandlerCustomizer;
importorg.springframework.context.annotation.Bean;
importorg.springframework.context.annotation.Configuration;
importjava.util.concurrent.ExecutorService;
importjava.util.concurrent.Executors;
importjava.util.concurrent.TimeUnit;
/**
*SpringBoot3.2虚拟线程配置
*
*核心策略：
*1.Tomcat请求处理：使用虚拟线程（每个请求一个虚拟线程）
*2.推理任务提交：使用固定线程池（保护GPU资源）
*3.下游服务调用：RestTemplate，天然支持虚拟线程
*/
@Configuration
publicclassVirtualThreadConfig{
/**
*虚拟线程执行器：用于IO密集型任务
*每个任务创建一个虚拟线程，任务完成后虚拟线程销毁
*适合：HTTP调用、数据库查询、文件IO
*/
@Bean
publicExecutorServicevirtualThreadExecutor(){
returnExecutors.newVirtualThreadPerTaskExecutor();
}
/**
*Tomcat协议处理器自定义：使用虚拟线程处理请求
*SpringBoot3.2+支持此配置
*/
@Bean
publicTomcatProtocolHandlerCustomizer<Http11NioProtocol>tomcatVirtualThreadCustomizer(){
returnprotocolHandler–>{
// 将Tomcat的工作线程池替换为虚拟线程执行器
// 每个HTTP请求会在一个虚拟线程中处理
ExecutorServiceexecutor=Executors.newVirtualThreadPerTaskExecutor();
protocolHandler.setExecutor(executor);
};
}
/**
*推理执行器：固定线程池，保护GPU资源
*
*注意：这里是平台线程池，不是虚拟线程
*因为GPU推理是CPU密集型，虚拟线程在CPU密集场景没有优势
*/
@Bean(name=“inferenceExecutor“)
publicExecutorServiceinferenceExecutor(){
// 核心数 * 2：GPU推理虽然消耗GPU，但CPU侧需要做数据拷贝、结果解析
intthreads=Runtime.getRuntime().availableProcessors()*2;
returnnewThreadPoolExecutor(
threads,threads,
0L,TimeUnit.MILLISECONDS,
newLinkedBlockingQueue<>(1024),
newThreadPoolExecutor.CallerRunsPolicy()// 拒绝时由调用方执行
);
}
}

packagecom.ai.inference.service;
importai.onnxruntime.OnnxTensor;
importai.onnxruntime.OrtEnvironment;
importai.onnxruntime.OrtException;
importai.onnxruntime.OrtSession;
importcom.ai.inference.model.InferenceRequest;
importcom.ai.inference.model.InferenceResult;
importlombok.extern.slf4j.Slf4j;
importorg.springframework.beans.factory.annotation.Qualifier;
importorg.springframework.stereotype.Service;
importjava.time.Duration;
importjava.time.Instant;
importjava.util.*;
importjava.util.concurrent.*;
importjava.util.concurrent.atomic.AtomicLong;
/**
*基于虚拟线程的高性能AI推理服务
*
*设计理念：
*–HTTP请求处理：虚拟线程（高并发、零成本创建）
*–模型推理执行：平台线程池（保护GPU，防止过载）
*–下游服务调用：虚拟线程（RestTemplate）
*/
@Service
@Slf4j
publicclassOnnxInferenceService{
privatefinalOrtEnvironmentortEnv;// ONNX全局环境
privatefinalConcurrentHashMap<String,OrtSession>sessionCache;// 模型会话缓存
privatefinalExecutorServiceinferenceExecutor;// 推理专用线程池

// 批处理队列：收集请求后批量推理
privatefinalBlockingQueue<BatchTask>batchQueue;
privatefinalScheduledExecutorServicebatchScheduler;

// 指标
privatefinalAtomicLongtotalRequests=newAtomicLong(0);
privatefinalAtomicLongtotalLatencyNs=newAtomicLong(0);
publicOnnxInferenceService(
@Qualifier(“inferenceExecutor“)ExecutorServiceinferenceExecutor){
this.ortEnv=OrtEnvironment.getEnvironment();
this.sessionCache=newConcurrentHashMap<>();
this.inferenceExecutor=inferenceExecutor;
this.batchQueue=newLinkedBlockingQueue<>(2048);
this.batchScheduler=Executors.newSingleThreadScheduledExecutor(
r–>newThread(r,“batch-scheduler“)
);

// 启动批处理调度器：每50ms尝试收集一批请求
startBatchScheduler();

log.info(“OnnxInferenceService initialized with virtual thread support“);
}
/**
*同步推理：虚拟线程执行
*
*这个方法在虚拟线程中运行
*当调用onnxSession.run()发起GPU推理时，
*虚拟线程会被挂起，Carrier线程去执行其他虚拟线程
*GPU推理完成后，虚拟线程自动恢复，继续处理结果
*/
publicInferenceResultinfer(InferenceRequestrequest){
Instantstart=Instant.now();
totalRequests.incrementAndGet();
try{
// 模型会话获取（可能涉及磁盘加载，首次慢）
OrtSessionsession=getOrCreateSession(request.getModelName());

// 输入张量构造
OnnxTensorinputTensor=createInputTensor(request);

// ⚡ 这里是关键：推理是同步调用
// 虚拟线程在这里被挂起，直到GPU完成推理
// 但Carrier线程此时可以处理其他虚拟线程的请求！
OrtSession.RunResultrunResult=session.run(
Collections.singletonMap(“input“,inputTensor)
);
// 解析输出
float[]output=(float[])runResult.getValue(0).get().getValue();

InferenceResultresult=InferenceResult.builder()
.modelName(request.getModelName())
.predictions(output)
.latencyMs(Duration.between(start,Instant.now()).toMillis())
.timestamp(Instant.now())
.build();
totalLatencyNs.addAndGet(
Duration.between(start,Instant.now()).toNanos()
);

returnresult;
}catch(OrtExceptione){
log.error(“推理失败: model={}, error={}“,request.getModelName(),e.getMessage());
thrownewInferenceException(“ONNX推理执行失败“,e);
}
}
/**
*异步推理：CompletableFuture+虚拟线程
*
*底层仍然是同步推理，但包装了Future
*适合需要非阻塞返回的场景
*/
publicCompletableFuture<InferenceResult>inferAsync(InferenceRequestrequest){
// supplyAsync 默认使用 ForkJoinPool.common()
// 但这里我们显式用虚拟线程执行器
returnCompletableFuture.supplyAsync(
()–>infer(request),
CompletableFuture.delayedExecutor(0,TimeUnit.MILLISECONDS)
);
}
/**
*批处理推理：动态收集请求，批量推理
*
*批处理的核心价值：GPU一次处理多个样本，
*利用矩阵运算的并行性，提升吞吐量
*/
publicCompletableFuture<InferenceResult>inferBatch(InferenceRequestrequest){
CompletableFuture<InferenceResult>resultFuture=newCompletableFuture<>();

batchQueue.offer(newBatchTask(request,resultFuture));

returnresultFuture;
}
privatevoidstartBatchScheduler(){
// 每50ms调度一次
batchScheduler.scheduleAtFixedRate(()–>{
List<BatchTask>batch=newArrayList<>();
batchQueue.drainTo(batch,64);// 最多64个一批

if(batch.size()<2)return;// 少于2个不划算

try{
// 构造batch输入
float[][]batchInput=newfloat[batch.size()][];
for(inti=0;i<batch.size();i++){
batchInput[i]=batch.get(i).getRequest().getFeatures();
}

// 批量推理（平台线程池）
Instantstart=Instant.now();
OnnxTensorbatchTensor=OnnxTensor.createTensor(ortEnv,batchInput);

// 复用第一个请求的模型
OrtSessionsession=getOrCreateSession(
batch.get(0).getRequest().getModelName()
);

OrtSession.RunResultresult=session.run(
Collections.singletonMap(“input“,batchTensor)
);

float[][]outputs=(float[][])result.getValue(0).get().getValue();

longbatchLatency=Duration.between(start,Instant.now()).toMillis();

// 分发结果
for(inti=0;i<batch.size();i++){
InferenceResultir=InferenceResult.builder()
.modelName(batch.get(i).getRequest().getModelName())
.predictions(outputs[i])
.latencyMs(batchLatency)
.batchSize(batch.size())
.timestamp(Instant.now())
.build();
batch.get(i).getFuture().complete(ir);
}

totalRequests.addAndGet(batch.size());
totalLatencyNs.addAndGet(batchLatency*1_000_000);

}catch(Exceptione){
log.error(“批处理失败: {}“,e.getMessage());
batch.forEach(t–>t.getFuture().completeExceptionally(e));
}
},10,50,TimeUnit.MILLISECONDS);
}
privatesynchronizedOrtSessiongetOrCreateSession(StringmodelName){
returnsessionCache.computeIfAbsent(modelName,name–>{
try{
OrtSession.SessionOptionsoptions=newOrtSession.SessionOptions();
// CPU执行提供者
options.add_CPU();
// 多线程配置：推理是CPU密集型的，需要合理配置线程数
options.setInterOpNumThreads(4);// 算子间并行
options.setIntraOpNumThreads(16);// 算子内并行

log.info(“加载模型: {} -> {}“,modelName,name);
returnortEnv.createSession(“/models/“+name+“.onnx“,options);
}catch(OrtExceptione){
thrownewRuntimeException(“模型加载失败: “+name,e);
}
});
}
privateOnnxTensorcreateInputTensor(InferenceRequestrequest)throwsOrtException{
float[][]input={request.getFeatures()};
returnOnnxTensor.createTensor(ortEnv,input);
}
// Getters for metrics
publicStatsgetStats(){
longcount=totalRequests.get();
returnnewStats(
count,
count>0?totalLatencyNs.get()/count/1_000_000:0,
sessionCache.size(),
batchQueue.size()
);
}
@lombok.Data
@lombok.AllArgsConstructor
publicstaticclassStats{
publiclongtotalRequests;
publiclongavgLatencyMs;
publicintloadedModels;
publicintpendingBatch;
}
}

packagecom.ai.inference.controller;
importcom.ai.inference.model.InferenceRequest;
importcom.ai.inference.model.InferenceResult;
importcom.ai.inference.service.OnnxInferenceService;
importlombok.RequiredArgsConstructor;
importorg.springframework.http.ResponseEntity;
importorg.springframework.web.bind.annotation.*;
importjava.util.List;
importjava.util.concurrent.CompletableFuture;
@RestController
@RequestMapping(“/api/v1/inference“)
@RequiredArgsConstructor
publicclassInferenceController{
privatefinalOnnxInferenceServiceinferenceService;
/**
*同步推理
*虚拟线程在这里阻塞等待推理结果
*但不会阻塞其他请求！
*/
@PostMapping(“/sync“)
publicResponseEntity<InferenceResult>syncInfer(
@RequestBodyInferenceRequestrequest){
InferenceResultresult=inferenceService.infer(request);
returnResponseEntity.ok(result);
}
/**
*异步推理
*立即返回202Accepted，附任务ID
*客户端通过/status/{taskId}轮询结果
*/
@PostMapping(“/async“)
publicResponseEntity<TaskResponse>asyncInfer(
@RequestBodyInferenceRequestrequest){
CompletableFuture<InferenceResult>future=
inferenceService.inferAsync(request);

StringtaskId=“task-“+System.currentTimeMillis();
TaskResponseresp=newTaskResponse(taskId,“PENDING“);
returnResponseEntity.accepted().body(resp);
}
/**
*批处理推理
*适合离线批量预测
*/
@PostMapping(“/batch“)
publicResponseEntity<List<InferenceResult>>batchInfer(
@RequestBodyList<InferenceRequest>requests){
// 逐个提交，全部返回Future
List<CompletableFuture<InferenceResult>>futures=requests.stream()
.map(inferenceService::inferBatch)
.toList();

// 等待所有结果
CompletableFuture.allOf(futures.toArray(newCompletableFuture[0]))
.join();

List<InferenceResult>results=futures.stream()
.map(CompletableFuture::join)
.toList();

returnResponseEntity.ok(results);
}
@GetMapping(“/stats“)
publicResponseEntity<OnnxInferenceService.Stats>stats(){
returnResponseEntity.ok(inferenceService.getStats());
}
@lombok.Data
@lombok.AllArgsConstructor
staticclassTaskResponse{
privateStringtaskId;
privateStringstatus;
}
}

spring:
application:
name:ai-inference-service-vt

# Spring Boot 3.2 虚拟线程配置
threads:
virtual:
enabled:true# ⚠️ 这个配置实际上不需要了，3.2默认开启
# Tomcat配置：控制平台线程数
server:
tomcat:
threads:
max:200# 最大平台线程数（负载高时扩展）
min-spare:20# 最小保留（应对突发流量）
max-connections:10000
accept-count:500
# RestTemplate配置
http:
client:
connect-timeout:5s
read-timeout:30s
management:
endpoints:
web:
exposure:
include:health,metrics,prometheus
metrics:
export:
prometheus:
enabled:true
logging:
level:
com.ai.inference:INFO
ai.onnxruntime:WARN

服务器

AMD EPYC 7K62 64-Core / 128GB DDR4

模型

BERT-base ONNX格式（110M参数）

并发工具

JMeter + 1000并发连接

测试时长

预热100次 + 稳定压测5分钟

Java版本

OpenJDK 21.0.2

Python版本

Python 3.11 + uvloop

延迟百分位对比：

Python asyncio:

 P50   ████████████                      35ms

 P90   ████████████████████████         85ms  

 P99   ██████████████████████████████████████████████████ 120ms

Java虚拟线程:

 P50   ████████████                      18ms

 P90   ██████████████████████           55ms

 P99   ██████████████████████████████████████████████      80ms

Java虚拟线程+Batch:

 P50   ████████████                      12ms ⚡

 P90   ████████████████████               45ms

 P99   ██████████████████████████████████████████          65ms

// ❌ 错误：在线程池中嵌套提交虚拟线程（画蛇添足）
ExecutorServiceexecutor=Executors.newVirtualThreadPerTaskExecutor();
Future<?>f=executor.submit(()–>{
// 虚拟线程嵌套虚拟线程 = 没有任何意义
// 反而增加调度开销
});
// ✅ 正确：直接创建虚拟线程处理IO任务
try(varexecutor=Executors.newVirtualThreadPerTaskExecutor()){
List<Future<String>>futures=IntStream.range(0,10_000)
.mapToObj(i–>executor.submit(()–>{
// 发起HTTP请求、数据库查询、文件读写
returnrestTemplate.getForObject(“http://service/api“,String.class);
}))
.toList();
// join所有future
}

// ❌ 错误：synchronized会pin住Carrier线程
publicclassBadService{
privatefinalObjectlock=newObject();

publicvoidprocess(Requestreq){
synchronized(lock){
// 虚拟线程在synchronized块内会被pin到Carrier线程
// 导致Carrier线程无法调度其他虚拟线程
// 相当于退化为传统线程！
doExpensiveWork();
}
}
}
// ✅ 正确：使用ReentrantLock
publicclassGoodService{
privatefinalReentrantLocklock=newReentrantLock();

publicvoidprocess(Requestreq){
lock.lock();
try{
doExpensiveWork();
}finally{
lock.unlock();
}
}
}

// 虚拟线程会重复使用，所以ThreadLocal需要特殊处理
// ❌ 危险：ThreadLocal的值会泄漏到其他请求
publicclassBadController{
privatestaticfinalThreadLocal<Request>currentRequest=newThreadLocal<>();

@GetMapping(“/bad“)
publicResponsehandle(){
currentRequest.set(request);// ⚠️ 虚拟线程复用时会读到脏数据
// …
}
}
// ✅ 正确：使用ScopedValue（Java 21+）
publicclassGoodController{
privatestaticfinalScopedValue<Request>currentRequest=ScopedValue.newInstance();

@GetMapping(“/good“)
publicResponsehandle(){
ScopedValue.runWhere(currentRequest,request,()–>{
// 作用域内的请求绑定，方法结束后自动清理
doSomething();
});
}
}

你的团队技术栈？
│
├─ Python为主
│  └─ 优先考虑 FastAPI + asyncio + ONNX Runtime
│      （复用现有Python生态，避免切换成本）
│
└─ Java为主
  │
  ├─ AI推理吞吐 < 2000 QPS？
  │  └─ 虚拟线程 + ONNX：开箱即用
  │
  └─ AI推理吞吐 > 5000 QPS？
     └─ 虚拟线程 + 动态Batch + 连接池优化
        （实测6800 QPS，接近硬件极限）

虚拟线程≠更快

对于纯CPU计算，和平台线程一样快

虚拟线程=更高吞吐

对于IO密集型，能用更少资源处理更多请求

Batch是性能倍增器

动态batch是AI推理的专属优化，吞吐提升2-3倍

不要all in

Python负责模型训练/实验，Java负责生产推理服务