diff --git a/pom.xml b/pom.xml index e8c2b94..0e8c3d2 100644 --- a/pom.xml +++ b/pom.xml @@ -6,7 +6,7 @@ org.springframework.boot spring-boot-starter-parent - 3.1.3 + 3.2.0 com.apolloconfig.apollo.ai @@ -16,10 +16,13 @@ a smart qa bot 17 - 0.16.0 - 32.1.2-jre + 0.18.2 + 32.1.3-jre 0.64.8 - 2.3.0 + 2.3.3 + + 4.1.100.Final + 1.4.0 @@ -44,6 +47,16 @@ milvus-sdk-java ${milvus.version} + + io.netty + netty-codec-http2 + ${netty.codec.http2.version} + + + org.crac + crac + ${crac.version} + @@ -82,6 +95,10 @@ + + org.crac + crac + org.springframework.boot diff --git a/src/main/java/com/apolloconfig/apollo/ai/qabot/milvus/MilvusClientFactory.java b/src/main/java/com/apolloconfig/apollo/ai/qabot/milvus/MilvusClientFactory.java index 24d685e..43be515 100644 --- a/src/main/java/com/apolloconfig/apollo/ai/qabot/milvus/MilvusClientFactory.java +++ b/src/main/java/com/apolloconfig/apollo/ai/qabot/milvus/MilvusClientFactory.java @@ -4,8 +4,10 @@ import io.milvus.client.MilvusServiceClient; import io.milvus.param.ConnectParam; import java.util.Map; +import org.crac.Context; +import org.crac.Resource; -public class MilvusClientFactory { +public class MilvusClientFactory implements Resource { private static final MilvusClientFactory INSTANCE = new MilvusClientFactory(); private static final Map clients = Maps.newConcurrentMap(); @@ -43,4 +45,14 @@ private MilvusServiceClient createClient(String host, int port) { private MilvusServiceClient createCloudClient(String uri, String token) { return new MilvusServiceClient(ConnectParam.newBuilder().withUri(uri).withToken(token).build()); } + + @Override + public void beforeCheckpoint(Context context) throws Exception { + clients.clear(); + } + + @Override + public void afterRestore(Context context) throws Exception { + + } } diff --git a/src/main/java/com/apolloconfig/apollo/ai/qabot/milvus/MilvusService.java b/src/main/java/com/apolloconfig/apollo/ai/qabot/milvus/MilvusService.java index 6becc6e..e5c541b 100644 --- a/src/main/java/com/apolloconfig/apollo/ai/qabot/milvus/MilvusService.java +++ b/src/main/java/com/apolloconfig/apollo/ai/qabot/milvus/MilvusService.java @@ -34,20 +34,30 @@ import java.util.List; import java.util.Random; import java.util.stream.Collectors; +import org.crac.Context; +import org.crac.Resource; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import org.springframework.context.annotation.Profile; import org.springframework.stereotype.Service; import org.springframework.util.CollectionUtils; @Profile("milvus") @Service -class MilvusService implements VectorDBService { +class MilvusService implements VectorDBService, Resource { - private final MilvusServiceClient milvusServiceClient; + private static final Logger LOGGER = LoggerFactory.getLogger(MilvusService.class); + + private MilvusServiceClient milvusServiceClient; private final MilvusConfig milvusConfig; private final List dummyEmbeddings = Lists.newArrayList(); public MilvusService(MilvusConfig milvusConfig) { this.milvusConfig = milvusConfig; + this.init(); + } + + private void init() { if (milvusConfig.isUseZillzCloud()) { this.milvusServiceClient = MilvusClientFactory.getCloudClient( milvusConfig.getZillizCloudUri(), @@ -413,4 +423,16 @@ private void ensureFileCollection() { ); } + @Override + public void beforeCheckpoint(Context context) throws Exception { + LOGGER.info("beforeCheckpoint"); + this.milvusServiceClient = null; + } + + @Override + public void afterRestore(Context context) throws Exception { + LOGGER.info("afterRestore"); + this.init(); + LOGGER.info("afterRestore done"); + } } diff --git a/src/main/java/com/apolloconfig/apollo/ai/qabot/openai/OpenAiService.java b/src/main/java/com/apolloconfig/apollo/ai/qabot/openai/OpenAiService.java index f02028c..c8e2344 100644 --- a/src/main/java/com/apolloconfig/apollo/ai/qabot/openai/OpenAiService.java +++ b/src/main/java/com/apolloconfig/apollo/ai/qabot/openai/OpenAiService.java @@ -10,19 +10,25 @@ import com.theokanning.openai.embedding.EmbeddingRequest; import io.reactivex.Flowable; import java.util.List; +import org.crac.Context; +import org.crac.Resource; import org.springframework.context.annotation.Profile; import org.springframework.stereotype.Component; @Profile("openai") @Component -class OpenAiService implements AiService { +class OpenAiService implements AiService, Resource { private static final String DEFAULT_MODEL = "gpt-3.5-turbo"; private static final String DEFAULT_EMBEDDING_MODEL = "text-embedding-ada-002"; - private final com.theokanning.openai.service.OpenAiService service; + private com.theokanning.openai.service.OpenAiService service; public OpenAiService() { + init(); + } + + private void init() { service = OpenAiServiceFactory.getService(System.getenv("OPENAI_API_KEY")); } @@ -60,4 +66,14 @@ public List getEmbeddings(List chunks) { return service.createEmbeddings(embeddingRequest).getData(); } + + @Override + public void beforeCheckpoint(Context context) throws Exception { + this.service = null; + } + + @Override + public void afterRestore(Context context) throws Exception { + this.init(); + } } diff --git a/src/main/java/com/apolloconfig/apollo/ai/qabot/openai/OpenAiServiceFactory.java b/src/main/java/com/apolloconfig/apollo/ai/qabot/openai/OpenAiServiceFactory.java index 1d67662..770b979 100644 --- a/src/main/java/com/apolloconfig/apollo/ai/qabot/openai/OpenAiServiceFactory.java +++ b/src/main/java/com/apolloconfig/apollo/ai/qabot/openai/OpenAiServiceFactory.java @@ -23,9 +23,11 @@ import okhttp3.Authenticator; import okhttp3.Credentials; import okhttp3.OkHttpClient; +import org.crac.Context; +import org.crac.Resource; import retrofit2.Retrofit; -public class OpenAiServiceFactory { +public class OpenAiServiceFactory implements Resource { private static final Duration DEFAULT_TIMEOUT = Duration.ofSeconds(60); @@ -85,6 +87,16 @@ private OkHttpClient client(String apiKey) { .build(); } + @Override + public void beforeCheckpoint(Context context) throws Exception { + SERVICES.clear(); + } + + @Override + public void afterRestore(Context context) throws Exception { + + } + private static class DelegatingSocketFactory extends SocketFactory { private final SocketFactory delegate; diff --git a/src/main/scripts/startup.sh b/src/main/scripts/startup.sh index 77f1b54..00a2130 100644 --- a/src/main/scripts/startup.sh +++ b/src/main/scripts/startup.sh @@ -19,13 +19,27 @@ SERVICE_NAME=qa-bot LOG_DIR=/opt/logs ## Adjust server port if necessary SERVER_PORT=${SERVER_PORT:=9090} +## Adjust crac files dir if necessary +CRAC_FILES_DIR=/opt/crac ## Create log directory if not existed because JDK 8+ won't do that mkdir -p $LOG_DIR +mkdir -p $CRAC_FILES_DIR + ## Adjust memory settings if necessary #export JAVA_OPTS="-Xms2560m -Xmx2560m -Xss256k -XX:MetaspaceSize=128m -XX:MaxMetaspaceSize=384m -XX:NewSize=1536m -XX:MaxNewSize=1536m -XX:SurvivorRatio=8" +# Check for 'checkpoint' argument +if [ "$1" = "checkpoint" ]; then + export JAVA_OPTS="$JAVA_OPTS -Dspring.context.checkpoint=onRefresh -XX:CRaCCheckpointTo=$CRAC_FILES_DIR" +fi + +# Check for 'restore' argument +if [ "$1" = "restore" ]; then + export JAVA_OPTS="$JAVA_OPTS -XX:CRaCRestoreFrom=$CRAC_FILES_DIR" +fi + ## Only uncomment the following when you are using server jvm export JAVA_OPTS="$JAVA_OPTS -server -XX:-ReduceInitialCardMarks"