diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/reranking/configuration/RerankingProvidersConfig.java b/src/main/java/io/stargate/sgv2/jsonapi/service/reranking/configuration/RerankingProvidersConfig.java index 4978debeff..1a9172cf4e 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/reranking/configuration/RerankingProvidersConfig.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/reranking/configuration/RerankingProvidersConfig.java @@ -4,6 +4,7 @@ import io.smallrye.config.WithDefault; import io.stargate.sgv2.jsonapi.service.provider.ApiModelSupport; import io.stargate.sgv2.jsonapi.service.schema.collections.CollectionRerankDef; +import jakarta.validation.constraints.Positive; import java.util.List; import java.util.Map; import java.util.Objects; @@ -73,6 +74,8 @@ interface ModelConfig { RequestProperties properties(); interface RequestProperties { + int DEFAULT_CONNECTION_POOL_SIZE = 50; + /** * Specifies the maximum number of attempts before failing. Default is 3 (1 request + 2 * retries). @@ -100,6 +103,11 @@ interface RequestProperties { @WithDefault("5000") int readTimeoutMillis(); + /** Maximum number of HTTP/1.x connections in this model's shared REST client pool. */ + @Positive + @WithDefault("50") + int connectionPoolSize(); + /** * The maximum delay between retries in milliseconds. * diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/reranking/configuration/RerankingProvidersConfigImpl.java b/src/main/java/io/stargate/sgv2/jsonapi/service/reranking/configuration/RerankingProvidersConfigImpl.java index cc5e0897f9..5eb6fb246c 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/reranking/configuration/RerankingProvidersConfigImpl.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/reranking/configuration/RerankingProvidersConfigImpl.java @@ -35,8 +35,26 @@ public record RequestPropertiesImpl( int readTimeoutMillis, int maxBackOffMillis, double jitter, + int connectionPoolSize, int maxBatchSize) - implements RerankingProviderConfig.ModelConfig.RequestProperties {} + implements RerankingProviderConfig.ModelConfig.RequestProperties { + public RequestPropertiesImpl( + int atMostRetries, + int initialBackOffMillis, + int readTimeoutMillis, + int maxBackOffMillis, + double jitter, + int maxBatchSize) { + this( + atMostRetries, + initialBackOffMillis, + readTimeoutMillis, + maxBackOffMillis, + jitter, + DEFAULT_CONNECTION_POOL_SIZE, + maxBatchSize); + } + } } } } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/reranking/operation/NvidiaRerankingProvider.java b/src/main/java/io/stargate/sgv2/jsonapi/service/reranking/operation/NvidiaRerankingProvider.java index e2a9ff2979..6cd8f2648b 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/reranking/operation/NvidiaRerankingProvider.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/reranking/operation/NvidiaRerankingProvider.java @@ -10,6 +10,7 @@ import io.stargate.sgv2.jsonapi.service.provider.ModelProvider; import io.stargate.sgv2.jsonapi.service.provider.ProviderBillingFilter; import io.stargate.sgv2.jsonapi.service.reranking.configuration.RerankingProvidersConfig; +import io.vertx.core.http.HttpClientOptions; import jakarta.ws.rs.HeaderParam; import jakarta.ws.rs.POST; import jakarta.ws.rs.core.HttpHeaders; @@ -18,6 +19,7 @@ import java.net.URI; import java.util.*; import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; import org.eclipse.microprofile.rest.client.annotation.ClientHeaderParam; import org.eclipse.microprofile.rest.client.annotation.RegisterProvider; import org.eclipse.microprofile.rest.client.inject.RegisterRestClient; @@ -56,7 +58,7 @@ * } * } */ -public class NvidiaRerankingProvider extends RerankingProvider { +public class NvidiaRerankingProvider extends RerankingProvider implements AutoCloseable { private final NvidiaRerankingClient nvidiaClient; @@ -73,13 +75,28 @@ public class NvidiaRerankingProvider extends RerankingProvider { public NvidiaRerankingProvider( RerankingProvidersConfig.RerankingProviderConfig.ModelConfig modelConfig) { + this(modelConfig, createClient(modelConfig)); + } + + NvidiaRerankingProvider( + RerankingProvidersConfig.RerankingProviderConfig.ModelConfig modelConfig, + NvidiaRerankingClient nvidiaClient) { super(ModelProvider.NVIDIA, modelConfig); + this.nvidiaClient = nvidiaClient; + } + + private static NvidiaRerankingClient createClient( + RerankingProvidersConfig.RerankingProviderConfig.ModelConfig modelConfig) { + return QuarkusRestClientBuilder.newBuilder() + .baseUri(URI.create(modelConfig.url())) + .readTimeout(modelConfig.properties().readTimeoutMillis(), TimeUnit.MILLISECONDS) + .httpClientOptionsCustomizer(clientOptionsCustomizer(modelConfig)) + .build(NvidiaRerankingClient.class); + } - nvidiaClient = - QuarkusRestClientBuilder.newBuilder() - .baseUri(URI.create(modelConfig.url())) - .readTimeout(modelConfig.properties().readTimeoutMillis(), TimeUnit.MILLISECONDS) - .build(NvidiaRerankingClient.class); + static Consumer clientOptionsCustomizer( + RerankingProvidersConfig.RerankingProviderConfig.ModelConfig modelConfig) { + return options -> options.setMaxPoolSize(modelConfig.properties().connectionPoolSize()); } @Override @@ -132,6 +149,11 @@ public Uni rerank( }); } + @Override + public void close() throws Exception { + nvidiaClient.close(); + } + /** * REST client interface for the Nvidia Reranking Service. * @@ -140,7 +162,7 @@ public Uni rerank( @RegisterRestClient @RegisterProvider(RerankingProviderContentTypeFilter.class) @RegisterProvider(ProviderBillingFilter.class) - public interface NvidiaRerankingClient { + public interface NvidiaRerankingClient extends AutoCloseable { @POST @ClientHeaderParam(name = HttpHeaders.CONTENT_TYPE, value = MediaType.APPLICATION_JSON) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/reranking/operation/RerankingProviderFactory.java b/src/main/java/io/stargate/sgv2/jsonapi/service/reranking/operation/RerankingProviderFactory.java index c5ffa4dfab..b78b7dd39a 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/reranking/operation/RerankingProviderFactory.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/reranking/operation/RerankingProviderFactory.java @@ -8,9 +8,12 @@ import io.stargate.sgv2.jsonapi.service.provider.ModelProvider; import io.stargate.sgv2.jsonapi.service.reranking.configuration.RerankingProvidersConfig; import io.stargate.sgv2.jsonapi.service.reranking.gateway.RerankingEGWClient; +import jakarta.annotation.PreDestroy; import jakarta.enterprise.context.ApplicationScoped; import jakarta.inject.Inject; import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -33,6 +36,18 @@ RerankingProvider create( private static final Map RERANKING_PROVIDER_CTORS = Map.ofEntries(Map.entry(ModelProvider.NVIDIA, NvidiaRerankingProvider::new)); + private final Map providerConstructors; + private final ConcurrentMap directProviders = + new ConcurrentHashMap<>(); + + public RerankingProviderFactory() { + this(RERANKING_PROVIDER_CTORS); + } + + RerankingProviderFactory(Map providerConstructors) { + this.providerConstructors = Map.copyOf(providerConstructors); + } + public RerankingProvider create( Tenant tenant, String authToken, @@ -59,7 +74,7 @@ public RerankingProvider create( return create(tenant, authToken, modelProvider, modelName, authentication, commandName); } - private synchronized RerankingProvider create( + private RerankingProvider create( Tenant tenant, String authToken, ModelProvider modelProvider, @@ -105,16 +120,35 @@ private synchronized RerankingProvider create( commandName); } - RerankingProviderFactory.ProviderConstructor ctor = RERANKING_PROVIDER_CTORS.get(modelProvider); + RerankingProviderFactory.ProviderConstructor ctor = providerConstructors.get(modelProvider); if (ctor == null) { throw SchemaException.Code.RERANKING_SERVICE_TYPE_UNAVAILABLE.get( Map.of( "errorMessage", "unknown service provider '%s'".formatted(modelProvider.apiName()))); } - return ctor.create(modelConfig); + return directProviders.computeIfAbsent( + new ProviderKey(modelProvider, modelConfig.name()), ignored -> ctor.create(modelConfig)); + } + + @PreDestroy + void close() { + directProviders.values().stream() + .filter(AutoCloseable.class::isInstance) + .map(AutoCloseable.class::cast) + .forEach( + provider -> { + try { + provider.close(); + } catch (Exception exception) { + LOGGER.warn("Failed to close a cached reranking provider", exception); + } + }); + directProviders.clear(); } public RerankingProvidersConfig getRerankingConfig() { return rerankingConfig; } + + private record ProviderKey(ModelProvider modelProvider, String modelName) {} } diff --git a/src/main/resources/reranking-providers-config.yaml b/src/main/resources/reranking-providers-config.yaml index 67739fe0a2..f04fca2a21 100644 --- a/src/main/resources/reranking-providers-config.yaml +++ b/src/main/resources/reranking-providers-config.yaml @@ -14,4 +14,5 @@ stargate: is-default: true url: https://us-west-2.api-dev.ai.datastax.com/nvidia/v1/ranking properties: - max-batch-size: 10 \ No newline at end of file + connection-pool-size: 50 + max-batch-size: 10 diff --git a/src/main/resources/test-reranking-providers-config.yaml b/src/main/resources/test-reranking-providers-config.yaml index 7f1350470a..b706fe4888 100644 --- a/src/main/resources/test-reranking-providers-config.yaml +++ b/src/main/resources/test-reranking-providers-config.yaml @@ -14,6 +14,7 @@ stargate: is-default: true url: https://us-west-2.api-dev.ai.datastax.com/nvidia/v1/ranking properties: + connection-pool-size: 50 max-batch-size: 10 - name: nvidia/a-random-deprecated-model api-model-support: @@ -21,6 +22,7 @@ stargate: message: This model has been deprecated, it will be removed in a future release. It is not supported for new Collections or Tables. url: https://us-west-2.api-dev.ai.datastax.com/nvidia/v1/ranking properties: + connection-pool-size: 50 max-batch-size: 10 - name: nvidia/a-random-EOL-model api-model-support: @@ -28,4 +30,5 @@ stargate: message: This model is at END_OF_LIFE status, it is not supported. url: https://us-west-2.api-dev.ai.datastax.com/nvidia/v1/ranking properties: - max-batch-size: 10 \ No newline at end of file + connection-pool-size: 50 + max-batch-size: 10 diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/reranking/operation/NvidiaRerankingProviderTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/reranking/operation/NvidiaRerankingProviderTest.java index 8abeec3a4d..f405482f8c 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/reranking/operation/NvidiaRerankingProviderTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/reranking/operation/NvidiaRerankingProviderTest.java @@ -2,19 +2,32 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import io.quarkus.test.junit.QuarkusTest; import io.quarkus.test.junit.TestProfile; +import io.smallrye.mutiny.Uni; import io.smallrye.mutiny.helpers.test.UniAssertSubscriber; import io.stargate.sgv2.jsonapi.TestConstants; import io.stargate.sgv2.jsonapi.api.request.RerankingCredentials; +import io.stargate.sgv2.jsonapi.api.request.tenant.Tenant; +import io.stargate.sgv2.jsonapi.config.DatabaseType; +import io.stargate.sgv2.jsonapi.config.constants.HttpConstants; import io.stargate.sgv2.jsonapi.exception.SchemaException; import io.stargate.sgv2.jsonapi.service.provider.ApiModelSupport; import io.stargate.sgv2.jsonapi.service.reranking.configuration.RerankingProvidersConfig; import io.stargate.sgv2.jsonapi.service.reranking.configuration.RerankingProvidersConfigImpl; import io.stargate.sgv2.jsonapi.testresource.NoGlobalResourcesTestProfile; +import io.vertx.core.http.HttpClientOptions; +import jakarta.inject.Inject; import java.util.List; import java.util.Optional; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.Executors; import org.junit.jupiter.api.Test; /** Tests for {@link NvidiaRerankingProvider} */ @@ -28,7 +41,7 @@ public class NvidiaRerankingProviderTest { .RequestPropertiesImpl REQUEST_PROPERTIES = new RerankingProvidersConfigImpl.RerankingProviderConfigImpl.ModelConfigImpl - .RequestPropertiesImpl(3, 10, 100, 100, 0.5, 10); + .RequestPropertiesImpl(3, 10, 100, 100, 0.5, 37, 10); private static final RerankingProvidersConfig.RerankingProviderConfig.ModelConfig MODEL_CONFIG = new RerankingProvidersConfigImpl.RerankingProviderConfigImpl.ModelConfigImpl( @@ -39,6 +52,79 @@ public class NvidiaRerankingProviderTest { "https://us-west-2.api-dev.ai.datastax.com/nvidia/v1/ranking", REQUEST_PROPERTIES); + @Inject RerankingProvidersConfig rerankingProvidersConfig; + + @Test + void configuresConnectionPoolSizeFromModel() { + var options = new HttpClientOptions(); + + NvidiaRerankingProvider.clientOptionsCustomizer(MODEL_CONFIG).accept(options); + + assertThat(options.getMaxPoolSize()).isEqualTo(37); + } + + @Test + void loadsExplicitConnectionPoolSize() { + var configuredModel = + rerankingProvidersConfig.providers().values().stream() + .flatMap(provider -> provider.models().stream()) + .filter(RerankingProvidersConfig.RerankingProviderConfig.ModelConfig::isDefault) + .findFirst() + .orElseThrow(); + + assertThat(configuredModel.properties().connectionPoolSize()).isEqualTo(50); + } + + @Test + void sharedProviderUsesCredentialsFromConcurrentCalls() throws Exception { + var nvidiaClient = mock(NvidiaRerankingProvider.NvidiaRerankingClient.class); + when(nvidiaClient.rerank(any(), any(), any())).thenReturn(Uni.createFrom().nothing()); + var provider = new NvidiaRerankingProvider(MODEL_CONFIG, nvidiaClient); + var firstTenant = Tenant.create(DatabaseType.ASTRA, "first-tenant"); + var secondTenant = Tenant.create(DatabaseType.ASTRA, "second-tenant"); + var barrier = new CyclicBarrier(2); + + try (var executor = Executors.newFixedThreadPool(2)) { + var calls = + List.of( + executor.submit( + () -> { + barrier.await(); + provider.rerank( + 0, + "query", + List.of("first"), + new RerankingCredentials(firstTenant, "first-key")); + return null; + }), + executor.submit( + () -> { + barrier.await(); + provider.rerank( + 1, + "query", + List.of("second"), + new RerankingCredentials(secondTenant, "second-key")); + return null; + })); + + for (var call : calls) { + call.get(); + } + } + + verify(nvidiaClient) + .rerank( + eq(HttpConstants.BEARER_PREFIX_FOR_API_KEY + "first-key"), + eq(firstTenant.toString()), + any(NvidiaRerankingProvider.NvidiaRerankingRequest.class)); + verify(nvidiaClient) + .rerank( + eq(HttpConstants.BEARER_PREFIX_FOR_API_KEY + "second-key"), + eq(secondTenant.toString()), + any(NvidiaRerankingProvider.NvidiaRerankingRequest.class)); + } + @Test void testEmptyApiKeyThrowsException() { NvidiaRerankingProvider provider = new NvidiaRerankingProvider(MODEL_CONFIG); diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/reranking/operation/RerankingProviderFactoryTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/reranking/operation/RerankingProviderFactoryTest.java new file mode 100644 index 0000000000..494997146c --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/reranking/operation/RerankingProviderFactoryTest.java @@ -0,0 +1,160 @@ +package io.stargate.sgv2.jsonapi.service.reranking.operation; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; + +import io.stargate.embedding.gateway.RerankingService; +import io.stargate.sgv2.jsonapi.api.request.tenant.Tenant; +import io.stargate.sgv2.jsonapi.config.DatabaseType; +import io.stargate.sgv2.jsonapi.config.OperationsConfig; +import io.stargate.sgv2.jsonapi.service.provider.ModelProvider; +import io.stargate.sgv2.jsonapi.service.reranking.configuration.RerankingProvidersConfig; +import io.stargate.sgv2.jsonapi.service.reranking.gateway.RerankingEGWClient; +import java.util.List; +import java.util.Map; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class RerankingProviderFactoryTest { + + private static final String MODEL_NAME = "nvidia/test-model"; + private static final Tenant TENANT = Tenant.create(DatabaseType.ASTRA, "test-tenant"); + + private final RerankingProvidersConfig.RerankingProviderConfig.ModelConfig modelConfig = + mock(RerankingProvidersConfig.RerankingProviderConfig.ModelConfig.class); + private final RerankingProvidersConfig.RerankingProviderConfig providerConfig = + mock(RerankingProvidersConfig.RerankingProviderConfig.class); + private final RerankingProvidersConfig rerankingConfig = mock(RerankingProvidersConfig.class); + private final OperationsConfig operationsConfig = mock(OperationsConfig.class); + + @BeforeEach + void setUp() { + var requestProperties = + mock(RerankingProvidersConfig.RerankingProviderConfig.ModelConfig.RequestProperties.class); + when(requestProperties.initialBackOffMillis()).thenReturn(100); + when(requestProperties.maxBackOffMillis()).thenReturn(500); + when(modelConfig.name()).thenReturn(MODEL_NAME); + when(modelConfig.properties()).thenReturn(requestProperties); + when(providerConfig.models()).thenReturn(List.of(modelConfig)); + when(rerankingConfig.providers()) + .thenReturn(Map.of(ModelProvider.NVIDIA.apiName(), providerConfig)); + } + + @Test + void reusesDirectProviderAcrossRequests() { + var creations = new AtomicInteger(); + var directProvider = mock(RerankingProvider.class); + var factory = + factory( + config -> { + creations.incrementAndGet(); + return directProvider; + }); + + var first = create(factory, TENANT, "first-token"); + var second = + create(factory, Tenant.create(DatabaseType.ASTRA, "another-tenant"), "second-token"); + + assertThat(first).isSameAs(second); + assertThat(creations).hasValue(1); + } + + @Test + void createsOneDirectProviderUnderConcurrency() throws Exception { + var creations = new AtomicInteger(); + var factory = + factory( + config -> { + creations.incrementAndGet(); + return mock(RerankingProvider.class); + }); + + try (var executor = Executors.newFixedThreadPool(16)) { + var requests = + java.util.stream.IntStream.range(0, 64) + .>mapToObj( + index -> () -> create(factory, TENANT, "token-" + index)) + .toList(); + + var providers = + executor.invokeAll(requests).stream() + .map( + future -> { + try { + return future.get(); + } catch (Exception exception) { + throw new AssertionError(exception); + } + }) + .toList(); + + assertThat(providers).allMatch(provider -> provider == providers.getFirst()); + assertThat(creations).hasValue(1); + } + } + + @Test + void doesNotShareDirectProvidersAcrossModels() { + var otherModel = mock(RerankingProvidersConfig.RerankingProviderConfig.ModelConfig.class); + when(otherModel.name()).thenReturn("nvidia/other-model"); + when(providerConfig.models()).thenReturn(List.of(modelConfig, otherModel)); + var factory = factory(config -> mock(RerankingProvider.class)); + + var first = create(factory, TENANT, "token"); + var second = + factory.create( + TENANT, + "token", + ModelProvider.NVIDIA.apiName(), + otherModel.name(), + Map.of(), + "findAndRerank"); + + assertThat(first).isNotSameAs(second); + } + + @Test + void keepsGatewayProvidersRequestScoped() { + when(operationsConfig.enableEmbeddingGateway()).thenReturn(true); + var factory = factory(config -> mock(RerankingProvider.class)); + factory.grpcGatewayService = mock(RerankingService.class); + + var first = create(factory, TENANT, "first-token"); + var second = create(factory, TENANT, "second-token"); + + assertThat(first).isNotSameAs(second); + assertThat(first).isInstanceOf(RerankingEGWClient.class); + assertThat(second).isInstanceOf(RerankingEGWClient.class); + } + + @Test + void closesCachedDirectProvidersAtShutdown() throws Exception { + var directProvider = + mock(RerankingProvider.class, withSettings().extraInterfaces(AutoCloseable.class)); + var factory = factory(config -> directProvider); + create(factory, TENANT, "token"); + + factory.close(); + + verify((AutoCloseable) directProvider).close(); + } + + private RerankingProviderFactory factory( + RerankingProviderFactory.ProviderConstructor providerConstructor) { + var factory = new RerankingProviderFactory(Map.of(ModelProvider.NVIDIA, providerConstructor)); + factory.rerankingConfig = rerankingConfig; + factory.operationsConfig = operationsConfig; + return factory; + } + + private RerankingProvider create( + RerankingProviderFactory factory, Tenant tenant, String authToken) { + return factory.create( + tenant, authToken, ModelProvider.NVIDIA.apiName(), MODEL_NAME, Map.of(), "findAndRerank"); + } +}