Update Chroma Warmup
This commit is contained in:
@@ -5,20 +5,24 @@ import org.slf4j.LoggerFactory;
|
||||
import org.springframework.beans.factory.annotation.Value;
|
||||
import org.springframework.boot.context.event.ApplicationReadyEvent;
|
||||
import org.springframework.context.event.EventListener;
|
||||
import org.springframework.scheduling.annotation.Async;
|
||||
import org.springframework.scheduling.annotation.Scheduled;
|
||||
import org.springframework.stereotype.Component;
|
||||
import org.springframework.web.client.RestTemplate;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
|
||||
/**
|
||||
* Service per mantenere Chroma "sveglio" ed evitare cold start
|
||||
* Esegue un warmup all'avvio e periodicamente ogni 5 minuti
|
||||
* Esegue un warmup all'avvio con retry e periodicamente ogni 5 minuti
|
||||
*/
|
||||
@Component
|
||||
public class ChromaWarmupService {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(ChromaWarmupService.class);
|
||||
private static final int MAX_RETRY_ATTEMPTS = 5;
|
||||
private static final int INITIAL_BACKOFF_MS = 2000;
|
||||
|
||||
@Value("${spring.ai.vectorstore.chroma.client.host}")
|
||||
private String chromaHost;
|
||||
@@ -26,47 +30,118 @@ public class ChromaWarmupService {
|
||||
@Value("${spring.ai.vectorstore.chroma.client.port}")
|
||||
private String chromaPort;
|
||||
|
||||
@Value("${spring.ai.vectorstore.chroma.collection-name:olympus_collection}")
|
||||
private String collectionName;
|
||||
|
||||
private final RestTemplate restTemplate;
|
||||
private volatile boolean chromaReady = false;
|
||||
|
||||
public ChromaWarmupService() {
|
||||
this.restTemplate = new RestTemplate();
|
||||
// Timeout generosi per il warmup
|
||||
// Timeout più lunghi per gestire cold start
|
||||
this.restTemplate.setRequestFactory(new org.springframework.http.client.SimpleClientHttpRequestFactory() {{
|
||||
setConnectTimeout((int) Duration.ofSeconds(30).toMillis());
|
||||
setReadTimeout((int) Duration.ofSeconds(60).toMillis());
|
||||
setConnectTimeout((int) Duration.ofSeconds(60).toMillis());
|
||||
setReadTimeout((int) Duration.ofMinutes(2).toMillis());
|
||||
}});
|
||||
}
|
||||
|
||||
/**
|
||||
* Esegue il warmup all'avvio dell'applicazione
|
||||
* Esegue il warmup all'avvio dell'applicazione con retry
|
||||
* Eseguito in modo asincrono per non bloccare lo startup
|
||||
*/
|
||||
@Async
|
||||
@EventListener(ApplicationReadyEvent.class)
|
||||
public void warmupOnStartup() {
|
||||
logger.info("Starting Chroma warmup on application startup...");
|
||||
performWarmup();
|
||||
logger.info("Starting aggressive Chroma warmup on application startup...");
|
||||
|
||||
// Primo warmup con retry aggressivo
|
||||
for (int attempt = 1; attempt <= MAX_RETRY_ATTEMPTS; attempt++) {
|
||||
try {
|
||||
logger.info("Warmup attempt {}/{}", attempt, MAX_RETRY_ATTEMPTS);
|
||||
|
||||
if (performCompleteWarmup()) {
|
||||
chromaReady = true;
|
||||
logger.info("Chroma is now ready after {} attempt(s)", attempt);
|
||||
|
||||
// Esegui ancora 2-3 warmup aggiuntivi per stabilizzare
|
||||
for (int i = 1; i <= 3; i++) {
|
||||
TimeUnit.SECONDS.sleep(5);
|
||||
logger.info("Stabilization warmup {}/3", i);
|
||||
performCompleteWarmup();
|
||||
}
|
||||
return;
|
||||
}
|
||||
} catch (Exception e) {
|
||||
logger.warn("Warmup attempt {} failed: {}", attempt, e.getMessage());
|
||||
}
|
||||
|
||||
// Backoff esponenziale
|
||||
if (attempt < MAX_RETRY_ATTEMPTS) {
|
||||
try {
|
||||
int backoffTime = INITIAL_BACKOFF_MS * (int) Math.pow(2, attempt - 1);
|
||||
logger.info("Waiting {}ms before next attempt...", backoffTime);
|
||||
TimeUnit.MILLISECONDS.sleep(backoffTime);
|
||||
} catch (InterruptedException ie) {
|
||||
Thread.currentThread().interrupt();
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger.error("Failed to warmup Chroma after {} attempts", MAX_RETRY_ATTEMPTS);
|
||||
}
|
||||
|
||||
/**
|
||||
* Esegue il warmup ogni 5 minuti per mantenere Chroma attivo
|
||||
*/
|
||||
@Scheduled(fixedRate = 300000) // TEST: 30 secondi (cambiare a 300000 per production)
|
||||
@Scheduled(fixedRate = 300000) // 5 minuti
|
||||
public void scheduledWarmup() {
|
||||
logger.info("Performing scheduled Chroma warmup...");
|
||||
performWarmup();
|
||||
}
|
||||
|
||||
private void performWarmup() {
|
||||
String healthUrl = chromaHost + ":" + chromaPort + "/api/v2/heartbeat";
|
||||
|
||||
try {
|
||||
long startTime = System.currentTimeMillis();
|
||||
restTemplate.getForObject(healthUrl, String.class);
|
||||
long duration = System.currentTimeMillis() - startTime;
|
||||
|
||||
logger.info("Chroma warmup successful ({}ms)", duration);
|
||||
} catch (Exception e) {
|
||||
logger.warn("Chroma warmup failed: {}", e.getMessage());
|
||||
// Non blocchiamo l'applicazione se il warmup fallisce
|
||||
if (chromaReady) {
|
||||
logger.debug("Performing scheduled Chroma warmup...");
|
||||
performCompleteWarmup();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Esegue un warmup completo: heartbeat + operazioni sulla collezione
|
||||
*/
|
||||
private boolean performCompleteWarmup() {
|
||||
long startTime = System.currentTimeMillis();
|
||||
boolean success = true;
|
||||
|
||||
try {
|
||||
// 1. Health check
|
||||
String healthUrl = chromaHost + ":" + chromaPort + "/api/v1/heartbeat";
|
||||
restTemplate.getForObject(healthUrl, String.class);
|
||||
logger.debug("Heartbeat successful");
|
||||
|
||||
// 2. Lista collezioni (operazione più pesante)
|
||||
String collectionsUrl = chromaHost + ":" + chromaPort + "/api/v1/collections";
|
||||
restTemplate.getForObject(collectionsUrl, String.class);
|
||||
logger.debug("Collections list successful");
|
||||
|
||||
// 3. Prova a contare elementi nella collezione principale
|
||||
try {
|
||||
String countUrl = chromaHost + ":" + chromaPort + "/api/v1/collections/" + collectionName + "/count";
|
||||
restTemplate.getForObject(countUrl, String.class);
|
||||
logger.debug("Collection count successful");
|
||||
} catch (Exception e) {
|
||||
// La collezione potrebbe non esistere ancora
|
||||
logger.debug("Collection count failed (collection may not exist): {}", e.getMessage());
|
||||
}
|
||||
|
||||
long duration = System.currentTimeMillis() - startTime;
|
||||
logger.info("Chroma warmup completed successfully ({}ms)", duration);
|
||||
|
||||
} catch (Exception e) {
|
||||
logger.warn("Chroma warmup failed: {}", e.getMessage());
|
||||
success = false;
|
||||
}
|
||||
|
||||
return success;
|
||||
}
|
||||
|
||||
public boolean isChromaReady() {
|
||||
return chromaReady;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user