✨ Async Replicate prediction generation

MrGraversen · MrGraversen · commit efa769a2fd20 · 2024-10-28T14:41:53.000+01:00
With event-driven notifications
diff --git a/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/common/BasePredictionMapper.java b/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/common/BasePredictionMapper.java
@@ -0,0 +1,15 @@
+package io.graversen.replicate.common;
+
+import lombok.NonNull;
+
+import java.util.Objects;
+import java.util.Set;
+
+public abstract class BasePredictionMapper<T, R> implements PredictionMapper<T, R> {
+    protected abstract Set<ReplicateModel> supportedModels();
+
+    @Override
+    public boolean supportsModel(@NonNull ReplicateModel model) {
+        return Objects.requireNonNullElseGet(supportedModels(), Set::of).stream().anyMatch(model::equals);
+    }
+}
diff --git a/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/common/PredictionMapper.java b/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/common/PredictionMapper.java
@@ -0,0 +1,14 @@
+package io.graversen.replicate.common;
+
+import lombok.NonNull;
+
+import java.util.function.BiFunction;
+
+public interface PredictionMapper<T, R> extends BiFunction<ReplicateModel, T, R> {
+    @Override
+    R apply(@NonNull ReplicateModel model, @NonNull T createPrediction);
+
+    boolean supportsType(@NonNull PredictionTypes type);
+
+    boolean supportsModel(@NonNull ReplicateModel model);
+}
diff --git a/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/common/PredictionTypes.java b/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/common/PredictionTypes.java
@@ -0,0 +1,6 @@
+package io.graversen.replicate.common;
+
+public enum PredictionTypes {
+    TEXT,
+    IMAGE
+}
diff --git a/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/common/ReplicateModel.java b/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/common/ReplicateModel.java
@@ -0,0 +1,10 @@
+package io.graversen.replicate.common;
+
+import lombok.NonNull;
+import lombok.Value;
+
+@Value
+public class ReplicateModel {
+    @NonNull String owner;
+    @NonNull String name;
+}
diff --git a/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/configuration/ReplicateAutoConfiguration.java b/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/configuration/ReplicateAutoConfiguration.java
@@ -0,0 +1,10 @@
+package io.graversen.replicate.configuration;
+
+import org.springframework.context.annotation.Configuration;
+import org.springframework.context.annotation.Import;
+
+@Configuration
+@Import(ReplicateConfiguration.class)
+public class ReplicateAutoConfiguration {
+
+}
diff --git a/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/configuration/ReplicateConfiguration.java b/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/configuration/ReplicateConfiguration.java
@@ -0,0 +1,56 @@
+package io.graversen.replicate.configuration;
+
+import io.graversen.replicate.client.configuration.ReplicateClientProperties;
+import io.graversen.replicate.client.configuration.ReplicateClients;
+import io.graversen.replicate.client.feign.FeignUtils;
+import io.graversen.replicate.client.replicate.Replicate;
+import io.graversen.replicate.facade.PredictionRetryPolicy;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.boot.context.properties.EnableConfigurationProperties;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.ComponentScan;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.context.annotation.Import;
+
+import java.util.Objects;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.function.Function;
+
+@Configuration
+@EnableConfigurationProperties(ReplicateProperties.class)
+@Import({Llama3Configuration.class})
+@ComponentScan({"io.graversen.replicate.service", "io.graversen.replicate.facade"})
+public class ReplicateConfiguration {
+    @Bean
+    public Replicate replicate(ReplicateProperties properties) {
+        final var replicateClientProperties = mapReplicateClientProperties().apply(properties);
+        return ReplicateClients.v1(replicateClientProperties);
+    }
+
+    @Bean
+    public FeignUtils feignUtils() {
+        return new FeignUtils(ReplicateClients.objectMapper());
+    }
+
+    @Bean
+    @ConditionalOnMissingBean
+    public ExecutorService defaultReplicateExecutor() {
+        return Executors.newVirtualThreadPerTaskExecutor();
+    }
+
+    @Bean
+    public PredictionRetryPolicy predictionRetryPolicy(ReplicateProperties properties) {
+        final var defaultRetryPolicy = PredictionRetryPolicy.defaultPolicy();
+        final var pollDelay = Objects.requireNonNullElse(properties.getPredictionPollDelay(), defaultRetryPolicy.getDelay());
+        final var pollAttempts = Objects.requireNonNullElse(properties.getPredictionPollAttempts(), defaultRetryPolicy.getMaxAttempts());
+        return new PredictionRetryPolicy(pollAttempts, pollDelay);
+    }
+
+    Function<ReplicateProperties, ReplicateClientProperties> mapReplicateClientProperties() {
+        return replicateProperties -> new ReplicateClientProperties(
+                replicateProperties.getToken(),
+                replicateProperties.getApiUrl()
+        );
+    }
+}
diff --git a/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/configuration/ReplicateProperties.java b/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/configuration/ReplicateProperties.java
@@ -0,0 +1,18 @@
+package io.graversen.replicate.configuration;
+
+import lombok.Getter;
+import lombok.NonNull;
+import lombok.RequiredArgsConstructor;
+import org.springframework.boot.context.properties.ConfigurationProperties;
+
+import java.time.Duration;
+
+@Getter
+@RequiredArgsConstructor
+@ConfigurationProperties(prefix = "replicate")
+public class ReplicateProperties {
+    private final @NonNull String token;
+    private final String apiUrl;
+    private final Duration predictionPollDelay;
+    private final Integer predictionPollAttempts;
+}
diff --git a/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/CheckPredictionStateTask.java b/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/CheckPredictionStateTask.java
@@ -0,0 +1,23 @@
+package io.graversen.replicate.facade;
+
+import jakarta.annotation.Nullable;
+import lombok.SneakyThrows;
+import org.springframework.stereotype.Component;
+
+import java.util.Optional;
+import java.util.concurrent.CompletionException;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+@Component
+public class CheckPredictionStateTask implements Function<Optional<PredictionResponseAndModel>, PredictionResponseAndModel> {
+    @Override
+    @SneakyThrows
+    public PredictionResponseAndModel apply(@Nullable Optional<PredictionResponseAndModel> predictionResponse) {
+        return predictionResponse.orElseThrow(predictionFailedError());
+    }
+
+    private Supplier<Exception> predictionFailedError() {
+        return () -> new CompletionException(new RuntimeException("Could not initiate prediction"));
+    }
+}
diff --git a/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/EmitPredictionResponseTask.java b/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/EmitPredictionResponseTask.java
@@ -0,0 +1,47 @@
+package io.graversen.replicate.facade;
+
+import jakarta.annotation.Nullable;
+import lombok.NonNull;
+import lombok.RequiredArgsConstructor;
+import org.springframework.context.ApplicationEventPublisher;
+import org.springframework.stereotype.Component;
+
+import java.util.function.BiConsumer;
+import java.util.function.Function;
+
+@Component
+@RequiredArgsConstructor
+public class EmitPredictionResponseTask implements BiConsumer<PredictionResponseAndModel, Throwable> {
+    private final @NonNull ApplicationEventPublisher eventPublisher;
+
+    @Override
+    public void accept(@Nullable PredictionResponseAndModel predictionResponse, @Nullable Throwable throwable) {
+        if (throwable == null) {
+            final var event = mapPredictionUpdatedEvent().apply(predictionResponse);
+            eventPublisher.publishEvent(event);
+        } else {
+            final var event = mapPredictionFailedEvent(throwable).apply(predictionResponse);
+            eventPublisher.publishEvent(event);
+        }
+    }
+
+    Function<PredictionResponseAndModel, PredictionUpdatedEvent> mapPredictionUpdatedEvent() {
+        return predictionResponse -> {
+            final var status = PredictionStatus.fromString(predictionResponse.getPredictionResponse().getStatus());
+            return new PredictionUpdatedEvent(
+                    predictionResponse.getPredictionResponse().getId(),
+                    predictionResponse.getModel(),
+                    status,
+                    predictionResponse.getPredictionResponse()
+            );
+        };
+    }
+
+    Function<PredictionResponseAndModel, PredictionFailedEvent> mapPredictionFailedEvent(@NonNull Throwable throwable) {
+        return predictionResponse -> new PredictionFailedEvent(
+                predictionResponse != null ? predictionResponse.getPredictionResponse().getId() : null,
+                predictionResponse != null ? predictionResponse.getModel() : null,
+                throwable
+        );
+    }
+}
diff --git a/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/PollPredictionStatusTask.java b/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/PollPredictionStatusTask.java
@@ -0,0 +1,63 @@
+package io.graversen.replicate.facade;
+
+import io.graversen.replicate.client.feign.FeignUtils;
+import io.graversen.replicate.client.replicate.Replicate;
+import io.graversen.replicate.service.PredictionResponse;
+import io.graversen.replicate.util.ReplicateUtils;
+import jakarta.annotation.Nullable;
+import lombok.NonNull;
+import lombok.RequiredArgsConstructor;
+import lombok.SneakyThrows;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.stereotype.Component;
+
+import java.util.LinkedHashMap;
+import java.util.function.Function;
+
+@Slf4j
+@Component
+@RequiredArgsConstructor
+public class PollPredictionStatusTask implements Function<PredictionResponseAndModel, PredictionResponseAndModel> {
+    private final @NonNull PredictionRetryPolicy retryPolicy;
+    private final @NonNull Replicate replicate;
+    private final @NonNull FeignUtils feignUtils;
+
+    @Override
+    @SneakyThrows
+    public PredictionResponseAndModel apply(@Nullable PredictionResponseAndModel predictionResponse) {
+        if (predictionResponse == null) {
+            return null;
+        }
+
+        final var predictionId = predictionResponse.getPredictionResponse().getId();
+        log.debug("[Prediction: {}]: Prediction retrieved for processing", predictionId);
+
+        for (int attempt = 1; attempt <= retryPolicy.getMaxAttempts(); attempt++) {
+            log.debug("[Prediction: {}]: Attempt {} of {} to retrieve prediction status", predictionId, attempt, retryPolicy.getMaxAttempts());
+            final var getPredictionResponse = replicate.getPrediction(predictionId);
+            final LinkedHashMap<String, Object> convertedResponse = feignUtils.convert(getPredictionResponse, LinkedHashMap.class);
+            final var mappedResponse = ReplicateUtils.mapPredictionResponse(convertedResponse);
+
+            if (mappedResponse.isPresent() && PredictionStatus.SUCCEEDED.asString().equals(mappedResponse.get().getStatus())) {
+                log.info("[Prediction: {}]: Prediction succeeded on attempt {}/{}", predictionId, attempt, retryPolicy.getMaxAttempts());
+                return new PredictionResponseAndModel(mappedResponse.get(), predictionResponse.getModel());
+            }
+
+            if (mappedResponse.isPresent() && PredictionStatus.FAILED.asString().equals(mappedResponse.get().getStatus())) {
+                log.error("[Prediction: {}]: Prediction failed on attempt {}/{}", predictionId, attempt, retryPolicy.getMaxAttempts());
+                throw new IllegalStateException(String.format("Prediction '%s' failed: %s", predictionId, mappedResponse.get().getStatus()));
+            }
+
+            try {
+                log.debug("[Prediction: {}]: Prediction not yet succeeded (status: {}), retrying after {} ms", predictionId, mappedResponse.map(PredictionResponse::getStatus).orElse("unknown"), retryPolicy.getDelay().toMillis());
+                Thread.sleep(retryPolicy.getDelay().toMillis());
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                throw new IllegalStateException("Thread was interrupted during retry delay", e);
+            }
+        }
+
+        log.error("[Prediction: {}]: Max retry attempts reached without success", predictionId);
+        throw new IllegalStateException(String.format("Max retry attempts reached for prediction '%s' without success", predictionId));
+    }
+}
diff --git a/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/PredictionFailedEvent.java b/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/PredictionFailedEvent.java
@@ -0,0 +1,15 @@
+package io.graversen.replicate.facade;
+
+import io.graversen.replicate.common.ReplicateModel;
+import lombok.AccessLevel;
+import lombok.NonNull;
+import lombok.RequiredArgsConstructor;
+import lombok.Value;
+
+@Value
+@RequiredArgsConstructor(access = AccessLevel.PACKAGE)
+public class PredictionFailedEvent {
+    String id;
+    ReplicateModel model;
+    @NonNull Throwable error;
+}
diff --git a/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/PredictionResponseAndModel.java b/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/PredictionResponseAndModel.java
@@ -0,0 +1,17 @@
+package io.graversen.replicate.facade;
+
+import io.graversen.replicate.common.ReplicateModel;
+import io.graversen.replicate.service.PredictionResponse;
+import lombok.AccessLevel;
+import lombok.NonNull;
+import lombok.RequiredArgsConstructor;
+import lombok.Value;
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
+@Value
+@RequiredArgsConstructor(access = AccessLevel.PACKAGE)
+public class PredictionResponseAndModel {
+    @NonNull PredictionResponse predictionResponse;
+    @NonNull ReplicateModel model;
+}
diff --git a/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/PredictionRetryPolicy.java b/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/PredictionRetryPolicy.java
@@ -0,0 +1,27 @@
+package io.graversen.replicate.facade;
+
+import lombok.Getter;
+import lombok.NonNull;
+
+import java.time.Duration;
+
+@Getter
+public class PredictionRetryPolicy {
+    private final @NonNull Integer maxAttempts;
+    private final @NonNull Duration delay;
+
+    public PredictionRetryPolicy(@NonNull Integer maxAttempts, @NonNull Duration delay) {
+        if (maxAttempts <= 0) {
+            throw new IllegalArgumentException("maxAttempts must be greater than 0");
+        }
+        if (delay.isNegative()) {
+            throw new IllegalArgumentException("delay must be non-negative");
+        }
+        this.maxAttempts = maxAttempts;
+        this.delay = delay;
+    }
+
+    public static PredictionRetryPolicy defaultPolicy() {
+        return new PredictionRetryPolicy(50, Duration.ofMillis(200));
+    }
+}
diff --git a/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/PredictionStatus.java b/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/PredictionStatus.java
@@ -0,0 +1,23 @@
+package io.graversen.replicate.facade;
+
+import lombok.NonNull;
+
+import java.util.Arrays;
+
+public enum PredictionStatus {
+    PROCESSING,
+    SUCCEEDED,
+    FAILED,
+    UNKNOWN;
+
+    public String asString() {
+        return name().toLowerCase();
+    }
+
+    public static PredictionStatus fromString(@NonNull String status) {
+        return Arrays.stream(values())
+                .filter(predictionStatus -> predictionStatus.name().equalsIgnoreCase(status))
+                .findFirst()
+                .orElse(UNKNOWN);
+    }
+}
diff --git a/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/PredictionUpdatedEvent.java b/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/PredictionUpdatedEvent.java
@@ -0,0 +1,17 @@
+package io.graversen.replicate.facade;
+
+import io.graversen.replicate.common.ReplicateModel;
+import io.graversen.replicate.service.PredictionResponse;
+import lombok.AccessLevel;
+import lombok.NonNull;
+import lombok.RequiredArgsConstructor;
+import lombok.Value;
+
+@Value
+@RequiredArgsConstructor(access = AccessLevel.PACKAGE)
+public class PredictionUpdatedEvent {
+    @NonNull String id;
+    @NonNull ReplicateModel model;
+    @NonNull PredictionStatus status;
+    @NonNull PredictionResponse prediction;
+}
diff --git a/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/ReplicateFacade.java b/spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/ReplicateFacade.java

-Original file line number
+Diff line change
@@ @@ -0,0 +1,6 @@ @@
 +package io.graversen.replicate.common;
++
 +public enum PredictionTypes {
 +    TEXT,
 +    IMAGE
 +}