VertexAITextEmbeddingModel not closing PredictionServiceClient causing thread leak #2609
Description
We faced an issue with VertexAITextEmbeddingModel when generating more than 500 text embeddings from a text file being ingested.
We managed to overcome the issue by copying the entire source code of the VertexAITextEmbeddingModel class and making a small change that will be explained below.
The proposed fix PR: #2611
Bug description
VertexAITextEmbeddingModel implementation is using the PredictionServiceClient to generate embeddings.
The client is not being properly closed, causing the following exception to be thrown
2025-03-19T15:06:14.197-04:00 ERROR 1 — [embedding-proxy-service] [istenerId-0-C-1] i.g.i.ManagedChannelOrphanWrapper : *~ Previous channel ManagedChannelImpl
{logId=685, target=us-central1-aiplatform.googleapis.com:443}
was garbage collected without being shut down! ~*
Make sure to call shutdown()/shutdownNow()
java.lang.RuntimeException: ManagedChannel allocation site
at io.grpc.internal.ManagedChannelOrphanWrapper$ManagedChannelReference.<init>(ManagedChannelOrphanWrapper.java:102) ~[grpc-core-1.68.1.jar:1.68.1]
at io.grpc.internal.ManagedChannelOrphanWrapper.<init>(ManagedChannelOrphanWrapper.java:60) ~[grpc-core-1.68.1.jar:1.68.1]
at io.grpc.internal.ManagedChannelOrphanWrapper.<init>(ManagedChannelOrphanWrapper.java:51) ~[grpc-core-1.68.1.jar:1.68.1]
at io.grpc.internal.ManagedChannelImplBuilder.build(ManagedChannelImplBuilder.java:710) ~[grpc-core-1.68.1.jar:1.68.1]
at io.grpc.ForwardingChannelBuilder2.build(ForwardingChannelBuilder2.java:272) ~[grpc-api-1.68.1.jar:1.68.1]
at com.google.api.gax.grpc.InstantiatingGrpcChannelProvider.createSingleChannel(InstantiatingGrpcChannelProvider.java:636) ~[gax-grpc-2.58.0.jar:2.58.0]
at com.google.api.gax.grpc.ChannelPool.<init>(ChannelPool.java:106) ~[gax-grpc-2.58.0.jar:2.58.0]
at com.google.api.gax.grpc.ChannelPool.create(ChannelPool.java:84) ~[gax-grpc-2.58.0.jar:2.58.0]
at com.google.api.gax.grpc.InstantiatingGrpcChannelProvider.createChannel(InstantiatingGrpcChannelProvider.java:296) ~[gax-grpc-2.58.0.jar:2.58.0]
at com.google.api.gax.grpc.InstantiatingGrpcChannelProvider.getTransportChannel(InstantiatingGrpcChannelProvider.java:289) ~[gax-grpc-2.58.0.jar:2.58.0]
at com.google.api.gax.rpc.ClientContext.create(ClientContext.java:226) ~[gax-2.58.0.jar:2.58.0]
at com.google.cloud.aiplatform.v1.stub.GrpcPredictionServiceStub.create(GrpcPredictionServiceStub.java:295) ~[google-cloud-aiplatform-3.52.0.jar:3.52.0]
at com.google.cloud.aiplatform.v1.stub.PredictionServiceStubSettings.createStub(PredictionServiceStubSettings.java:330) ~[google-cloud-aiplatform-3.52.0.jar:3.52.0]
at com.google.cloud.aiplatform.v1.PredictionServiceClient.<init>(PredictionServiceClient.java:412) ~[google-cloud-aiplatform-3.52.0.jar:3.52.0]
at com.google.cloud.aiplatform.v1.PredictionServiceClient.create(PredictionServiceClient.java:394) ~[google-cloud-aiplatform-3.52.0.jar:3.52.0]
at org.springframework.ai.vertexai.embedding.text.VertexAiTextEmbeddingModel.createPredictionServiceClient(VertexAiTextEmbeddingModel.java:209) ~[spring-ai-vertex-ai-embedding-1.0.0-M6.jar:1.0.0-M6]
The issue above is causing threads to be kept by the JVM, causing a JVM crash after a few minutes running. Attached are some Sumologic screenshots displaying how threads grow and are kept high until the container crashes.
Environment
Spring AI 1.0.0-M6 running on kubernetes cluster
Steps to reproduce
Expected behavior
The client should be closed for a proper disposal of resources.
Minimal Complete Reproducible example
Clone https://github.com/rodrigomalara/spring-ai-embedding-issue
Run the spring boot application - the exception will be thrown.
Suggested fix
Open a try with resources in line 131 of VertexAiTextEmbeddingModel class and close it properly and the end of the block.
try (PredictionServiceClient client = createPredictionServiceClient()) {
That took care of the issue for us.