Java实现Consul/Nacos根据GPU型号、显存余量执行负载均衡的步骤详解

<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>

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 -> bytes
                long 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省略
    }
}

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);
    }
}

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);
    }
}

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;
    }
    @Override
    public 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())
        );
    }
}

@Configuration
public class LoadBalancerConfig {
    @Bean
    public 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