Java实现Consul/Nacos根据GPU型号、显存余量执行负载均衡
步骤一:服务端获取GPU元数据
1. 添加依赖
在pom.xml
中引入Apache Commons Exec用于执行命令:
<dependency><groupId>org.apache.commons</groupId><artifactId>commons-exec</artifactId><version>1.3</version>
</dependency>
<dependency><groupId>com.google.code.gson</groupId><artifactId>gson</artifactId><version>2.8.9</version>
</dependency>
2. 实现GPU信息采集
import org.apache.commons.exec.CommandLine;
import org.apache.commons.exec.DefaultExecutor;
import org.apache.commons.exec.PumpStreamHandler;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import com.google.gson.Gson;public class GpuInfoUtil {public static List<GpuMeta> getGpuMetadata() throws IOException {CommandLine cmd = CommandLine.parse("nvidia-smi --query-gpu=name,memory.total,memory.free --format=csv,noheader,nounits");ByteArrayOutputStream outputStream = new ByteArrayOutputStream();PumpStreamHandler streamHandler = new PumpStreamHandler(outputStream);DefaultExecutor executor = new DefaultExecutor();executor.setStreamHandler(streamHandler);executor.execute(cmd);String output = outputStream.toString();return parseOutput(output);}private static List<GpuMeta> parseOutput(String output) {List<GpuMeta> gpus = new ArrayList<>();for (String line : output.split("\\r?\\n")) {String[] parts = line.split(",");if (parts.length >= 3) {String name = parts[0].trim();long total = Long.parseLong(parts[1].trim()) * 1024 * 1024; // MB -> byteslong free = Long.parseLong(parts[2].trim()) * 1024 * 1024;gpus.add(new GpuMeta(name, total, free));}}return gpus;}public static class GpuMeta {private String name;private long totalMem;private long freeMem;// 构造方法、getters、setters省略}
}
步骤二:服务注册到Consul/Nacos
1. Consul注册实现
import com.ecwid.consul.v1.ConsulClient;
import com.ecwid.consul.v1.agent.model.NewService;public class ConsulRegistrar {public void register(String serviceName, String ip, int port) throws Exception {ConsulClient consul = new ConsulClient("localhost", 8500);List<GpuMeta> gpus = GpuInfoUtil.getGpuMetadata();NewService service = new NewService();service.setId(serviceName + "-" + ip + ":" + port);service.setName(serviceName);service.setAddress(ip);service.setPort(port);// 序列化GPU元数据Gson gson = new Gson();service.getMeta().put("gpus", gson.toJson(gpus));consul.agentServiceRegister(service);}
}
2. Nacos注册实现
import com.alibaba.nacos.api.naming.NamingFactory;
import com.alibaba.nacos.api.naming.NamingService;
import com.alibaba.nacos.api.naming.pojo.Instance;public class NacosRegistrar {public void register(String serviceName, String ip, int port) throws Exception {NamingService naming = NamingFactory.createNamingService("localhost:8848");List<GpuMeta> gpus = GpuInfoUtil.getGpuMetadata();Instance instance = new Instance();instance.setIp(ip);instance.setPort(port);instance.setServiceName(serviceName);instance.getMetadata().put("gpus", new Gson().toJson(gpus));naming.registerInstance(serviceName, instance);}
}
步骤三:动态更新元数据
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;public class MetadataUpdater {private ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();private ConsulClient consulClient;private String serviceId;public void startUpdating() {scheduler.scheduleAtFixedRate(() -> {try {List<GpuMeta> gpus = GpuInfoUtil.getGpuMetadata();String gpuJson = new Gson().toJson(gpus);// 重新注册以更新元数据NewService service = new NewService();service.setId(serviceId);service.setMeta(Collections.singletonMap("gpus", gpuJson));consulClient.agentServiceRegister(service);} catch (Exception e) {e.printStackTrace();}}, 0, 10, TimeUnit.SECONDS);}
}
步骤四:客户端负载均衡(Spring Cloud示例)
1. 自定义负载均衡器
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Flux;public class GpuAwareServiceSupplier implements ServiceInstanceListSupplier {private final ServiceInstanceListSupplier delegate;private final Gson gson = new Gson();public GpuAwareServiceSupplier(ServiceInstanceListSupplier delegate) {this.delegate = delegate;}@Overridepublic Flux<List<ServiceInstance>> get() {return delegate.get().map(instances -> instances.stream().filter(instance -> {String gpuJson = instance.getMetadata().get("gpus");List<GpuMeta> gpus = gson.fromJson(gpuJson, new TypeToken<List<GpuMeta>>(){}.getType());return gpus.stream().anyMatch(g -> g.getFreeMem() > 2 * 1024 * 1024 * 1024L); // 2GB}).collect(Collectors.toList()));}
}
2. 配置负载均衡策略
@Configuration
public class LoadBalancerConfig {@Beanpublic ServiceInstanceListSupplier discoveryClientSupplier(ConfigurableApplicationContext context) {return ServiceInstanceListSupplier.builder().withDiscoveryClient().withCaching().withHealthChecks().withBlockingDiscoveryClient().build(context);}
}
最终验证
-
检查注册中心元数据
curl http://localhost:8500/v1/catalog/service/my-service | jq .
输出应包含类似:
{"ServiceMeta": {"gpus": "[{\"name\":\"Tesla T4\",\"totalMem\":17179869184,\"freeMem\":8589934592}]"} }
-
客户端调用验证
客户端会自动选择显存充足的节点,日志输出示例:INFO Selected instance 192.168.1.101:8080 with 8GB free GPU memory
通过以上步骤,即可在Java中实现基于GPU元数据的服务注册与负载均衡。