Add batch-llm module: 100,000 synthetic tickets classified by a ChatClient in a Spring Batch 6 chunk step with Resilience4j rate limiting and retry, skip-and-record for malformed JSON, restart after kill -9 in a second JVM, and a per-chunk cost ledger. The model is a deterministic simulated ChatModel; prices are configurable placeholders. No provider called.
Co-Authored-By: Claude Sonnet 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_017ic4b4VQVEJWetTT1dYzyJ
This commit is contained in:
Claude Sonnet 5.5
@@ -0,0 +1,20 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
/**
|
||||
* The model answered, and the answer is not usable JSON for our schema. Deterministic from the
|
||||
* pipeline's point of view: asking the same question again will usually give the same answer, so
|
||||
* this is a skip, not a retry. The tokens for the bad answer were still billed.
|
||||
*/
|
||||
public class BadModelOutputException extends RuntimeException {
|
||||
|
||||
private final String rawReply;
|
||||
|
||||
public BadModelOutputException(String reason, String rawReply) {
|
||||
super(reason);
|
||||
this.rawReply = rawReply;
|
||||
}
|
||||
|
||||
public String rawReply() {
|
||||
return rawReply;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
import org.springframework.boot.SpringApplication;
|
||||
import org.springframework.boot.autoconfigure.SpringBootApplication;
|
||||
import org.springframework.boot.context.properties.ConfigurationPropertiesScan;
|
||||
|
||||
@SpringBootApplication
|
||||
@ConfigurationPropertiesScan
|
||||
public class BatchLlmApplication {
|
||||
|
||||
public static void main(String[] args) {
|
||||
System.exit(SpringApplication.exit(SpringApplication.run(BatchLlmApplication.class, args)));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
import java.util.Set;
|
||||
|
||||
/** What the model is asked to return, and what the writer stores. */
|
||||
public record Classification(long ticketId, String category, int priority, double confidence) {
|
||||
|
||||
public static final Set<String> CATEGORIES =
|
||||
Set.of("BILLING", "TECHNICAL", "ACCOUNT", "SHIPPING", "FEATURE_REQUEST", "OTHER");
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
import java.math.BigDecimal;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.atomic.AtomicLong;
|
||||
|
||||
import org.springframework.batch.infrastructure.item.Chunk;
|
||||
import org.springframework.batch.infrastructure.item.ItemWriter;
|
||||
import org.springframework.jdbc.core.JdbcTemplate;
|
||||
import org.springframework.stereotype.Component;
|
||||
import org.springframework.transaction.support.TransactionSynchronization;
|
||||
import org.springframework.transaction.support.TransactionSynchronizationManager;
|
||||
|
||||
/**
|
||||
* Writes the results AND the cost ledger row for the chunk in one transaction, so a ledger row
|
||||
* exists if and only if its tickets do. The reader position, the results and the ledger either
|
||||
* all commit or all roll back.
|
||||
*/
|
||||
@Component
|
||||
public class ClassificationWriter implements ItemWriter<Classification> {
|
||||
|
||||
private final JdbcTemplate jdbc;
|
||||
private final UsageMeter meter;
|
||||
private final Pricing pricing;
|
||||
private final AtomicLong committedLedgerRows = new AtomicLong();
|
||||
|
||||
public ClassificationWriter(JdbcTemplate jdbc, UsageMeter meter, Pricing pricing) {
|
||||
this.jdbc = jdbc;
|
||||
this.meter = meter;
|
||||
this.pricing = pricing;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void write(Chunk<? extends Classification> chunk) {
|
||||
List<? extends Classification> items = chunk.getItems();
|
||||
jdbc.batchUpdate("INSERT INTO TICKET_RESULT (TICKET_ID, CATEGORY, PRIORITY, CONFIDENCE) VALUES (?,?,?,?)",
|
||||
items, items.size(), (ps, c) -> {
|
||||
ps.setLong(1, c.ticketId());
|
||||
ps.setString(2, c.category());
|
||||
ps.setInt(3, c.priority());
|
||||
ps.setDouble(4, c.confidence());
|
||||
});
|
||||
|
||||
// A ticket skipped in an earlier, uncommitted attempt (the skip listener commits on its own) and
|
||||
// classified now must not stay in the skipped table.
|
||||
jdbc.batchUpdate("DELETE FROM SKIPPED_TICKET WHERE TICKET_ID = ?", items, items.size(),
|
||||
(ps, c) -> ps.setLong(1, c.ticketId()));
|
||||
|
||||
long first = items.stream().mapToLong(Classification::ticketId).min().orElse(0);
|
||||
long last = items.stream().mapToLong(Classification::ticketId).max().orElse(0);
|
||||
UsageMeter.Snapshot spent = meter.drain();
|
||||
if (TransactionSynchronizationManager.isSynchronizationActive()) {
|
||||
TransactionSynchronizationManager.registerSynchronization(new TransactionSynchronization() {
|
||||
@Override
|
||||
public void afterCompletion(int status) {
|
||||
if (status == STATUS_COMMITTED) {
|
||||
long rows = committedLedgerRows.incrementAndGet();
|
||||
if (rows % 50 == 0) {
|
||||
System.out.println("PROGRESS committed-chunks-this-jvm=" + rows + " last-ticket=" + last);
|
||||
}
|
||||
} else {
|
||||
meter.restore(spent); // rolled back: the money was still spent
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
BigDecimal cost = pricing.cost(spent.promptTokens(), spent.completionTokens());
|
||||
jdbc.update("""
|
||||
INSERT INTO COST_LEDGER (FIRST_TICKET, LAST_TICKET, TICKETS_WRITTEN, TICKETS_SKIPPED, ATTEMPTS,
|
||||
REJECTED_429, UNAVAILABLE_503, RETRIES, BILLED_CALLS, BAD_REPLIES, PROMPT_TOKENS, COMPLETION_TOKENS,
|
||||
COST, PRICE_VERSION) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?)""",
|
||||
first, last, items.size(), spent.skipped(), spent.attempts(), spent.rejected429(),
|
||||
spent.unavailable503(), spent.retries(), spent.billedCalls(), spent.badReplies(),
|
||||
spent.promptTokens(), spent.completionTokens(), cost, pricing.version());
|
||||
}
|
||||
|
||||
public long committedLedgerRows() {
|
||||
return committedLedgerRows.get();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.concurrent.Executors;
|
||||
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
import org.springframework.ai.chat.model.ChatModel;
|
||||
import org.springframework.batch.core.job.Job;
|
||||
import org.springframework.batch.core.job.builder.JobBuilder;
|
||||
import org.springframework.batch.core.repository.JobRepository;
|
||||
import org.springframework.batch.core.step.Step;
|
||||
import org.springframework.batch.core.step.builder.ChunkOrientedStepBuilder;
|
||||
import org.springframework.batch.core.step.builder.StepBuilder;
|
||||
import org.springframework.batch.infrastructure.item.ItemWriter;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
import org.springframework.context.annotation.Configuration;
|
||||
import org.springframework.core.task.AsyncTaskExecutor;
|
||||
import org.springframework.core.task.support.TaskExecutorAdapter;
|
||||
import org.springframework.dao.TransientDataAccessException;
|
||||
import org.springframework.transaction.PlatformTransactionManager;
|
||||
|
||||
@Configuration
|
||||
public class LlmJobConfig {
|
||||
|
||||
@Bean
|
||||
ProviderMeter providerMeter(ProviderFaults faults) {
|
||||
return new ProviderMeter(faults.meterFile());
|
||||
}
|
||||
|
||||
/** The "provider". Replace this one bean with a real ChatModel and nothing else changes. */
|
||||
@Bean
|
||||
SimulatedProvider chatModel(ProviderFaults faults, ProviderMeter meter) {
|
||||
return new SimulatedProvider(faults, meter);
|
||||
}
|
||||
|
||||
@Bean
|
||||
ChatClient chatClient(ChatModel chatModel) {
|
||||
return ChatClient.builder(chatModel).defaultSystem(TicketClassifier.SYSTEM).build();
|
||||
}
|
||||
|
||||
@Bean
|
||||
TicketReader ticketReader(LlmJobProperties props) {
|
||||
return new TicketReader(props.tickets());
|
||||
}
|
||||
|
||||
/** Virtual threads, at most {@code concurrency} model calls in flight inside one chunk. */
|
||||
@Bean
|
||||
AsyncTaskExecutor llmExecutor(LlmJobProperties props) {
|
||||
var semaphore = new java.util.concurrent.Semaphore(props.concurrency());
|
||||
var virtual = Executors.newVirtualThreadPerTaskExecutor();
|
||||
return new TaskExecutorAdapter(task -> virtual.execute(() -> {
|
||||
semaphore.acquireUninterruptibly();
|
||||
try {
|
||||
task.run();
|
||||
} finally {
|
||||
semaphore.release();
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
@Bean
|
||||
Step classifyStep(JobRepository jobRepository, PlatformTransactionManager transactionManager,
|
||||
TicketReader reader, TicketClassifier classifier, ItemWriter<Classification> writer,
|
||||
SkippedTicketListener skipListener, AsyncTaskExecutor llmExecutor, LlmJobProperties props) {
|
||||
ChunkOrientedStepBuilder<Ticket, Classification> step = new StepBuilder("classifyStep", jobRepository)
|
||||
.<Ticket, Classification>chunk(props.chunkSize())
|
||||
.transactionManager(transactionManager)
|
||||
.reader(reader)
|
||||
.processor(classifier)
|
||||
.writer(writer)
|
||||
.taskExecutor(llmExecutor)
|
||||
.faultTolerant()
|
||||
// Deterministic bad output: skip it, remember it, keep going -- up to a limit.
|
||||
.skip(BadModelOutputException.class)
|
||||
// ... and a row the database refuses (too long, constraint) is just as deterministic.
|
||||
.skip(org.springframework.dao.DataIntegrityViolationException.class)
|
||||
.skipLimit(props.skipLimit())
|
||||
.skipListener(skipListener);
|
||||
if (!props.processorTransactional()) {
|
||||
step.processorNonTransactional();
|
||||
}
|
||||
// Spring Batch retry is for the DATABASE side of the chunk. The model call has its own
|
||||
// Resilience4j retry. Registering the 429 exception here as well multiplies the two.
|
||||
if (props.writeRetries() > 0 || props.batchRetriesCalls()) {
|
||||
var policy = org.springframework.core.retry.RetryPolicy.builder()
|
||||
.maxRetries(Math.max(1, props.writeRetries()))
|
||||
.delay(Duration.ofMillis(20))
|
||||
.includes(props.batchRetriesCalls()
|
||||
? java.util.Set.of(TransientDataAccessException.class, ProviderRateLimitException.class)
|
||||
: java.util.Set.of(TransientDataAccessException.class))
|
||||
.build();
|
||||
step.retryPolicy(policy);
|
||||
}
|
||||
return step.build();
|
||||
}
|
||||
|
||||
@Bean
|
||||
Job classifyJob(JobRepository jobRepository, Step classifyStep) {
|
||||
return new JobBuilder("classifyTickets", jobRepository).start(classifyStep).build();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
import org.springframework.boot.context.properties.ConfigurationProperties;
|
||||
import org.springframework.boot.context.properties.bind.DefaultValue;
|
||||
|
||||
/**
|
||||
* @param tickets number of synthetic tickets to classify
|
||||
* @param chunkSize items per transaction; also the most work a crash can lose
|
||||
* @param concurrency threads (virtual) calling the model at once inside one chunk
|
||||
* @param reask extra model calls allowed when a reply is not valid JSON
|
||||
* @param skipLimit the step fails when more than this many tickets are skipped
|
||||
* @param writeRetries Spring Batch retries of a failed WRITE (0 = none)
|
||||
* @param batchRetriesCalls also let Spring Batch retry the rate-limit exception (the mistake the post warns about)
|
||||
* @param processorTransactional re-run the model for every item when a chunk is scanned after a failed write
|
||||
* @param runId identifying job parameter: same value = same JobInstance = restart target
|
||||
*/
|
||||
@ConfigurationProperties("batch.llm")
|
||||
public record LlmJobProperties(
|
||||
@DefaultValue("100000") long tickets,
|
||||
@DefaultValue("100") int chunkSize,
|
||||
@DefaultValue("64") int concurrency,
|
||||
@DefaultValue("1") int reask,
|
||||
@DefaultValue("1000") long skipLimit,
|
||||
@DefaultValue("0") int writeRetries,
|
||||
@DefaultValue("false") boolean batchRetriesCalls,
|
||||
@DefaultValue("false") boolean processorTransactional,
|
||||
@DefaultValue("run-1") String runId) {
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
import java.security.MessageDigest;
|
||||
import java.util.HexFormat;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import org.springframework.batch.core.job.Job;
|
||||
import org.springframework.batch.core.job.JobExecution;
|
||||
import org.springframework.batch.core.job.parameters.JobParameters;
|
||||
import org.springframework.batch.core.job.parameters.JobParametersBuilder;
|
||||
import org.springframework.batch.core.launch.JobOperator;
|
||||
import org.springframework.batch.core.repository.JobRepository;
|
||||
import org.springframework.boot.ApplicationArguments;
|
||||
import org.springframework.boot.ApplicationRunner;
|
||||
import org.springframework.jdbc.core.JdbcTemplate;
|
||||
import org.springframework.stereotype.Component;
|
||||
|
||||
/**
|
||||
* {@code --mode=run} (default) starts or restarts the job; {@code --mode=report} only prints what
|
||||
* is in the database. Starting is the same call whether it is the first run or the tenth: the
|
||||
* identifying parameter {@code run.id} names the JobInstance.
|
||||
*
|
||||
* <p>If the previous JVM was killed, its JobExecution is still STARTED in the repository and
|
||||
* Spring Batch refuses to start another ("already running"). {@code JobOperator.recover} marks it
|
||||
* FAILED so the instance becomes restartable. Only do that when you know the old process is dead.
|
||||
*/
|
||||
@Component
|
||||
@org.springframework.boot.autoconfigure.condition.ConditionalOnProperty(name = "batch.llm.autorun", havingValue = "true", matchIfMissing = true)
|
||||
public class LlmJobRunner implements ApplicationRunner {
|
||||
|
||||
private final JobOperator operator;
|
||||
private final JobRepository repository;
|
||||
private final Job job;
|
||||
private final JdbcTemplate jdbc;
|
||||
private final LlmJobProperties props;
|
||||
private final Pricing pricing;
|
||||
private final ProviderMeter providerMeter;
|
||||
|
||||
public LlmJobRunner(JobOperator operator, JobRepository repository, Job job, JdbcTemplate jdbc,
|
||||
LlmJobProperties props, Pricing pricing, ProviderMeter providerMeter) {
|
||||
this.operator = operator;
|
||||
this.repository = repository;
|
||||
this.job = job;
|
||||
this.jdbc = jdbc;
|
||||
this.props = props;
|
||||
this.pricing = pricing;
|
||||
this.providerMeter = providerMeter;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void run(ApplicationArguments args) throws Exception {
|
||||
List<String> mode = args.getOptionValues("mode");
|
||||
if (mode != null && mode.contains("report")) {
|
||||
report();
|
||||
return;
|
||||
}
|
||||
JobParameters params = new JobParametersBuilder().addString("run.id", props.runId()).toJobParameters();
|
||||
var instance = repository.getJobInstance(job.getName(), params);
|
||||
if (instance != null) {
|
||||
var last = repository.getLastJobExecution(instance);
|
||||
if (last != null && last.getStatus().isRunning()) {
|
||||
System.out.println("RECOVER: execution " + last.getId() + " is " + last.getStatus()
|
||||
+ " but its JVM is gone; marking it FAILED so the instance can restart");
|
||||
operator.recover(last);
|
||||
}
|
||||
}
|
||||
long t0 = System.nanoTime();
|
||||
System.out.println("START run.id=" + props.runId() + " tickets=" + props.tickets() + " chunk=" + props.chunkSize());
|
||||
JobExecution execution = operator.start(job, params);
|
||||
System.out.printf("JOB FINISHED: executionId=%d status=%s seconds=%.1f%n", execution.getId(),
|
||||
execution.getStatus(), (System.nanoTime() - t0) / 1e9);
|
||||
report();
|
||||
}
|
||||
|
||||
void report() throws Exception {
|
||||
System.out.println("---- REPORT ----");
|
||||
long results = jdbc.queryForObject("SELECT COUNT(*) FROM TICKET_RESULT", Long.class);
|
||||
long skipped = jdbc.queryForObject("SELECT COUNT(*) FROM SKIPPED_TICKET", Long.class);
|
||||
System.out.printf("tickets total=%d classified=%d skipped=%d unaccounted=%d%n", props.tickets(), results, skipped,
|
||||
props.tickets() - results - skipped);
|
||||
System.out.println("step executions:");
|
||||
for (Map<String, Object> r : jdbc.queryForList("""
|
||||
SELECT se.STEP_EXECUTION_ID AS ID, se.JOB_EXECUTION_ID AS JOB, se.STATUS, se.READ_COUNT AS READS,
|
||||
se.WRITE_COUNT AS WRITES, se.PROCESS_SKIP_COUNT AS SKIPS, se.COMMIT_COUNT AS COMMITS, se.ROLLBACK_COUNT AS ROLLBACKS
|
||||
FROM BATCH_STEP_EXECUTION se ORDER BY se.STEP_EXECUTION_ID""")) {
|
||||
System.out.println(" " + r);
|
||||
}
|
||||
Map<String, Object> l = jdbc.queryForMap("""
|
||||
SELECT COUNT(*) AS ROWS_, COALESCE(SUM(TICKETS_WRITTEN),0) AS WRITTEN, COALESCE(SUM(TICKETS_SKIPPED),0) AS SKIPPED,
|
||||
COALESCE(SUM(ATTEMPTS),0) AS ATTEMPTS, COALESCE(SUM(REJECTED_429),0) AS R429, COALESCE(SUM(UNAVAILABLE_503),0) AS U503,
|
||||
COALESCE(SUM(RETRIES),0) AS RETRIES, COALESCE(SUM(BILLED_CALLS),0) AS BILLED, COALESCE(SUM(BAD_REPLIES),0) AS BAD,
|
||||
COALESCE(SUM(PROMPT_TOKENS),0) AS PT, COALESCE(SUM(COMPLETION_TOKENS),0) AS CT, COALESCE(SUM(COST),0) AS COST
|
||||
FROM COST_LEDGER""");
|
||||
System.out.println("ledger (committed chunks only): " + l);
|
||||
System.out.printf("price table '%s': input %s / output %s per million tokens (configurable placeholder)%n",
|
||||
pricing.version(), pricing.inputPerMillion(), pricing.outputPerMillion());
|
||||
long billed = providerMeter.get(ProviderMeter.Slot.BILLED);
|
||||
long ledgerBilled = ((Number) l.get("BILLED")).longValue();
|
||||
System.out.println("provider's own meter (survives kill -9):");
|
||||
System.out.print(providerMeter.summary().indent(2));
|
||||
System.out.printf("billed calls the ledger does not know about (work lost in a crash or rollback): %d%n", billed - ledgerBilled);
|
||||
var md = MessageDigest.getInstance("SHA-256");
|
||||
jdbc.query("SELECT TICKET_ID, CATEGORY, PRIORITY, CONFIDENCE FROM TICKET_RESULT ORDER BY TICKET_ID", rs -> {
|
||||
md.update((rs.getLong(1) + ":" + rs.getString(2) + ":" + rs.getInt(3) + ":" + rs.getDouble(4) + "\n").getBytes());
|
||||
});
|
||||
System.out.println("results checksum (sha-256 over ordered rows): " + HexFormat.of().formatHex(md.digest()).substring(0, 16));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
import java.math.BigDecimal;
|
||||
import java.math.RoundingMode;
|
||||
|
||||
import org.springframework.boot.context.properties.ConfigurationProperties;
|
||||
import org.springframework.boot.context.properties.bind.DefaultValue;
|
||||
|
||||
/**
|
||||
* A price table you configure. The defaults are PLACEHOLDERS chosen to make the arithmetic easy to
|
||||
* follow; they are not the price of any provider or model. Put your own contract prices in
|
||||
* {@code batch.pricing.*} and the ledger recomputes with them.
|
||||
*
|
||||
* @param version a label stored on every ledger row, so a price change is visible in the data
|
||||
* @param inputPerMillion currency units per million prompt tokens
|
||||
* @param outputPerMillion currency units per million completion tokens
|
||||
*/
|
||||
@ConfigurationProperties("batch.pricing")
|
||||
public record Pricing(
|
||||
@DefaultValue("EXAMPLE-NOT-A-REAL-PRICE") String version,
|
||||
@DefaultValue("0.50") BigDecimal inputPerMillion,
|
||||
@DefaultValue("1.50") BigDecimal outputPerMillion) {
|
||||
|
||||
private static final BigDecimal MILLION = BigDecimal.valueOf(1_000_000L);
|
||||
|
||||
public BigDecimal cost(long promptTokens, long completionTokens) {
|
||||
return inputPerMillion.multiply(BigDecimal.valueOf(promptTokens))
|
||||
.add(outputPerMillion.multiply(BigDecimal.valueOf(completionTokens)))
|
||||
.divide(MILLION, 9, RoundingMode.HALF_UP);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
import org.springframework.boot.context.properties.ConfigurationProperties;
|
||||
import org.springframework.boot.context.properties.bind.DefaultValue;
|
||||
|
||||
/**
|
||||
* Knobs of the SIMULATED provider. Every number here is an invented test input, not a measurement
|
||||
* of any real service.
|
||||
*
|
||||
* @param seed changes which tickets and calls are unlucky
|
||||
* @param limitPerSecond the simulated provider's own quota: more calls than this in any sliding
|
||||
* second get a 429
|
||||
* @param random429 chance that any single call is rejected with a 429 regardless of load
|
||||
* @param random503 chance that any single call gets a 503
|
||||
* @param permanentBad share of tickets whose reply is never valid JSON, however often asked
|
||||
* @param flakyBad share of tickets whose FIRST billed reply is invalid and the second is fine
|
||||
* @param slowShare share of calls that take {@code slowMillis} instead of 1-3 ms
|
||||
* @param slowMillis latency of a slow call
|
||||
* @param meterFile memory-mapped counters kept by the "provider", so they survive kill -9
|
||||
*/
|
||||
@ConfigurationProperties("sim.provider")
|
||||
public record ProviderFaults(
|
||||
@DefaultValue("7") long seed,
|
||||
@DefaultValue("2400") int limitPerSecond,
|
||||
@DefaultValue("0.015") double random429,
|
||||
@DefaultValue("0.005") double random503,
|
||||
@DefaultValue("0.003") double permanentBad,
|
||||
@DefaultValue("0.012") double flakyBad,
|
||||
@DefaultValue("0.01") double slowShare,
|
||||
@DefaultValue("40") int slowMillis,
|
||||
@DefaultValue("target/provider-meter.bin") String meterFile) {
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.RandomAccessFile;
|
||||
import java.nio.MappedByteBuffer;
|
||||
import java.nio.channels.FileChannel;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
|
||||
/**
|
||||
* The simulated provider's own invoice: counters in a memory-mapped file. The operating system
|
||||
* keeps a mapped page after the process is killed, so the counters survive {@code kill -9}, which
|
||||
* is what lets the crash demo say how many billed calls the dead JVM had made.
|
||||
*/
|
||||
public final class ProviderMeter {
|
||||
|
||||
public enum Slot { RECEIVED, REJECTED_429, WINDOW_429, UNAVAILABLE_503, BILLED, PROMPT_TOKENS, COMPLETION_TOKENS }
|
||||
|
||||
private final MappedByteBuffer buffer;
|
||||
|
||||
public ProviderMeter(String file) {
|
||||
try {
|
||||
Path path = Path.of(file);
|
||||
if (path.getParent() != null) {
|
||||
Files.createDirectories(path.getParent());
|
||||
}
|
||||
try (RandomAccessFile raf = new RandomAccessFile(path.toFile(), "rw")) {
|
||||
buffer = raf.getChannel().map(FileChannel.MapMode.READ_WRITE, 0, 8L * Slot.values().length);
|
||||
}
|
||||
} catch (IOException e) {
|
||||
throw new IllegalStateException("cannot map " + file, e);
|
||||
}
|
||||
}
|
||||
|
||||
public synchronized void add(Slot slot, long n) {
|
||||
int pos = slot.ordinal() * 8;
|
||||
buffer.putLong(pos, buffer.getLong(pos) + n);
|
||||
}
|
||||
|
||||
public synchronized long get(Slot slot) {
|
||||
return buffer.getLong(slot.ordinal() * 8);
|
||||
}
|
||||
|
||||
public synchronized void reset() {
|
||||
for (Slot s : Slot.values()) {
|
||||
buffer.putLong(s.ordinal() * 8, 0L);
|
||||
}
|
||||
}
|
||||
|
||||
public synchronized String summary() {
|
||||
StringBuilder sb = new StringBuilder();
|
||||
for (Slot s : Slot.values()) {
|
||||
sb.append(String.format("%-18s %,d%n", s.name().toLowerCase(), get(s)));
|
||||
}
|
||||
return sb.toString();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
/** The provider said 429 Too Many Requests. Transient: retrying later can succeed. */
|
||||
public class ProviderRateLimitException extends RuntimeException {
|
||||
public ProviderRateLimitException(String message) {
|
||||
super(message);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
/** The provider said 503. Transient, like {@link ProviderRateLimitException}. */
|
||||
public class ProviderUnavailableException extends RuntimeException {
|
||||
public ProviderUnavailableException(String message) {
|
||||
super(message);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,150 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
import java.util.ArrayDeque;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
import java.util.regex.Matcher;
|
||||
import java.util.regex.Pattern;
|
||||
|
||||
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
|
||||
import org.springframework.ai.chat.metadata.DefaultUsage;
|
||||
import org.springframework.ai.chat.model.ChatModel;
|
||||
import org.springframework.ai.chat.model.ChatResponse;
|
||||
import org.springframework.ai.chat.model.Generation;
|
||||
import org.springframework.ai.chat.prompt.Prompt;
|
||||
|
||||
import com.ankurm.batchllm.ProviderMeter.Slot;
|
||||
|
||||
/**
|
||||
* A hand-written, deterministic {@link ChatModel} -- the same idea as the ScriptedChatModel in the
|
||||
* other modules of this repository, but driven by a fault plan instead of a list of replies.
|
||||
* It is NOT a language model: it reads the ticket number out of the prompt and decides, from a
|
||||
* hash of (seed, ticket, call number), whether this call returns valid JSON, malformed JSON, a
|
||||
* 429, a 503 or a slow answer. No network, no key, no vendor.
|
||||
*
|
||||
* <p>It also enforces its own quota: more than {@code limitPerSecond} calls in any sliding second
|
||||
* get a 429, which is what a client-side rate limiter exists to avoid.
|
||||
*
|
||||
* <p>Token counts are ESTIMATED as characters / 4 and returned through the same
|
||||
* {@code ChatResponseMetadata.usage} field a real provider fills in.
|
||||
*/
|
||||
public class SimulatedProvider implements ChatModel {
|
||||
|
||||
private static final Pattern TICKET_ID = Pattern.compile("Ticket #(\\d+)");
|
||||
private static final List<String> CATEGORIES = List.of("BILLING", "TECHNICAL", "ACCOUNT", "SHIPPING", "FEATURE_REQUEST", "OTHER");
|
||||
|
||||
private final ProviderFaults faults;
|
||||
private final ProviderMeter meter;
|
||||
private final ConcurrentHashMap<Long, AtomicInteger> callsPerTicket = new ConcurrentHashMap<>();
|
||||
private final ConcurrentHashMap<Long, AtomicInteger> billedPerTicket = new ConcurrentHashMap<>();
|
||||
private final ArrayDeque<Long> window = new ArrayDeque<>();
|
||||
private volatile int latencyDivisor = 1;
|
||||
|
||||
public SimulatedProvider(ProviderFaults faults, ProviderMeter meter) {
|
||||
this.faults = faults;
|
||||
this.meter = meter;
|
||||
}
|
||||
|
||||
/** Test hook: make every sleep this many times shorter. */
|
||||
public void setLatencyDivisor(int divisor) {
|
||||
this.latencyDivisor = Math.max(1, divisor);
|
||||
}
|
||||
|
||||
public ProviderMeter meter() {
|
||||
return meter;
|
||||
}
|
||||
|
||||
/** Total calls the model has seen for one ticket, rejected ones included. */
|
||||
public int callsFor(long ticketId) {
|
||||
AtomicInteger n = callsPerTicket.get(ticketId);
|
||||
return n == null ? 0 : n.get();
|
||||
}
|
||||
|
||||
private double unit(long ticket, long salt) {
|
||||
long h = TicketGenerator.mix(TicketGenerator.mix(faults.seed() * 31 + salt) ^ ticket);
|
||||
return (h >>> 11) * 0x1.0p-53;
|
||||
}
|
||||
|
||||
@Override
|
||||
public ChatResponse call(Prompt prompt) {
|
||||
String text = prompt.getContents();
|
||||
Matcher m = TICKET_ID.matcher(text);
|
||||
if (!m.find()) {
|
||||
throw new IllegalArgumentException("simulated provider expects 'Ticket #<id>' in the prompt");
|
||||
}
|
||||
long id = Long.parseLong(m.group(1));
|
||||
int callNo = callsPerTicket.computeIfAbsent(id, k -> new AtomicInteger()).incrementAndGet();
|
||||
meter.add(Slot.RECEIVED, 1);
|
||||
|
||||
if (overQuota()) {
|
||||
meter.add(Slot.REJECTED_429, 1);
|
||||
meter.add(Slot.WINDOW_429, 1);
|
||||
throw new ProviderRateLimitException("429 Too Many Requests (quota: " + faults.limitPerSecond() + " calls per second)");
|
||||
}
|
||||
if (unit(id, 1000L + callNo) < faults.random429()) {
|
||||
meter.add(Slot.REJECTED_429, 1);
|
||||
throw new ProviderRateLimitException("429 Too Many Requests (injected)");
|
||||
}
|
||||
if (unit(id, 2000L + callNo) < faults.random503()) {
|
||||
meter.add(Slot.UNAVAILABLE_503, 1);
|
||||
throw new ProviderUnavailableException("503 Service Unavailable (injected)");
|
||||
}
|
||||
|
||||
boolean slow = unit(id, 3000L + callNo) < faults.slowShare();
|
||||
long millis = slow ? faults.slowMillis() : 1 + (long) (unit(id, 4000L + callNo) * 3);
|
||||
try {
|
||||
Thread.sleep(Math.max(0, millis / latencyDivisor));
|
||||
} catch (InterruptedException e) {
|
||||
Thread.currentThread().interrupt();
|
||||
throw new IllegalStateException(e);
|
||||
}
|
||||
|
||||
int billedNo = billedPerTicket.computeIfAbsent(id, k -> new AtomicInteger()).incrementAndGet();
|
||||
String reply = reply(id, billedNo);
|
||||
int promptTokens = (text.length() + 3) / 4;
|
||||
int completionTokens = (reply.length() + 3) / 4;
|
||||
meter.add(Slot.BILLED, 1);
|
||||
meter.add(Slot.PROMPT_TOKENS, promptTokens);
|
||||
meter.add(Slot.COMPLETION_TOKENS, completionTokens);
|
||||
|
||||
ChatResponseMetadata metadata = ChatResponseMetadata.builder()
|
||||
.model("simulated-classifier")
|
||||
.usage(new DefaultUsage(promptTokens, completionTokens))
|
||||
.build();
|
||||
return new ChatResponse(List.of(new Generation(new AssistantMessage(reply))), metadata);
|
||||
}
|
||||
|
||||
private synchronized boolean overQuota() {
|
||||
long now = System.nanoTime();
|
||||
while (!window.isEmpty() && now - window.peekFirst() > 1_000_000_000L) {
|
||||
window.pollFirst();
|
||||
}
|
||||
if (window.size() >= faults.limitPerSecond()) {
|
||||
return true;
|
||||
}
|
||||
window.addLast(now);
|
||||
return false;
|
||||
}
|
||||
|
||||
/** The reply for ticket {@code id} on its {@code billedNo}-th billed answer. */
|
||||
String reply(long id, int billedNo) {
|
||||
String category = CATEGORIES.get((int) Math.floorMod(TicketGenerator.mix(id * 1_000_003L + 17), 6L));
|
||||
int priority = 1 + (int) (unit(id, 5000L) * 4);
|
||||
double confidence = Math.round((0.60 + unit(id, 6000L) * 0.39) * 100) / 100.0;
|
||||
String good = "{\"category\":\"" + category + "\",\"priority\":" + priority + ",\"confidence\":" + confidence + "}";
|
||||
boolean permanent = unit(id, 7000L) < faults.permanentBad();
|
||||
boolean flaky = !permanent && billedNo == 1 && unit(id, 8000L) < faults.flakyBad();
|
||||
if (!permanent && !flaky) {
|
||||
// one reply in ten comes wrapped in a markdown fence: valid after the fence is stripped
|
||||
return unit(id, 9000L) < 0.10 ? "```json\n" + good + "\n```" : good;
|
||||
}
|
||||
return switch ((int) (unit(id, 10_000L) * 4)) {
|
||||
case 0 -> good.substring(0, good.length() - 14); // truncated
|
||||
case 1 -> "Sure! Here is the classification you asked for: " + good; // prose in front
|
||||
case 2 -> good.replace(category, category.charAt(0) + category.substring(1).toLowerCase() + " & Payments"); // enum value not in the schema
|
||||
default -> "{\"category\":\"" + category + "\",\"priority\":" + priority + "}"; // field missing
|
||||
};
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
import org.springframework.batch.core.listener.SkipListener;
|
||||
import org.springframework.jdbc.core.JdbcTemplate;
|
||||
import org.springframework.stereotype.Component;
|
||||
import org.springframework.transaction.PlatformTransactionManager;
|
||||
import org.springframework.transaction.TransactionDefinition;
|
||||
import org.springframework.transaction.support.TransactionTemplate;
|
||||
|
||||
/**
|
||||
* Records every skipped ticket with the reason and the raw reply, so a skipped ticket is a row
|
||||
* you can look at and re-run, not a number in a counter.
|
||||
*
|
||||
* <p>Both callbacks write in a REQUIRES_NEW transaction. {@code onSkipInWrite} is called while the
|
||||
* chunk's own transaction is about to be rolled back (that is how the scan finds the bad item); an
|
||||
* insert that joined that transaction would be rolled back with it and the skipped ticket would
|
||||
* leave no trace. This was found by a test that counted tickets: 299 of 300 were accounted for.
|
||||
* {@code MERGE} makes the insert idempotent, because after a crash the same ticket can be skipped twice.
|
||||
*/
|
||||
@Component
|
||||
public class SkippedTicketListener implements SkipListener<Ticket, Classification> {
|
||||
|
||||
private final JdbcTemplate jdbc;
|
||||
private final TransactionTemplate ownTransaction;
|
||||
|
||||
public SkippedTicketListener(JdbcTemplate jdbc, PlatformTransactionManager transactionManager) {
|
||||
this.jdbc = jdbc;
|
||||
this.ownTransaction = new TransactionTemplate(transactionManager);
|
||||
this.ownTransaction.setPropagationBehavior(TransactionDefinition.PROPAGATION_REQUIRES_NEW);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onSkipInProcess(Ticket ticket, Throwable t) {
|
||||
String raw = t instanceof BadModelOutputException b ? b.rawReply() : null;
|
||||
record(ticket.id(), t.getMessage(), raw);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onSkipInWrite(Classification item, Throwable t) {
|
||||
record(item.ticketId(), "write rejected: " + t.getClass().getSimpleName(), item.toString());
|
||||
}
|
||||
|
||||
private void record(long ticketId, String reason, String raw) {
|
||||
String r = reason == null ? null : reason.substring(0, Math.min(reason.length(), 200));
|
||||
String rawReply = raw == null ? null : raw.substring(0, Math.min(raw.length(), 400));
|
||||
ownTransaction.executeWithoutResult(s -> jdbc.update(
|
||||
"MERGE INTO SKIPPED_TICKET (TICKET_ID, REASON, RAW_REPLY) KEY (TICKET_ID) VALUES (?,?,?)", ticketId, r, rawReply));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
/** One support ticket. Synthetic: see {@link TicketGenerator}. */
|
||||
public record Ticket(long id, String subject, String body) {
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
import java.util.function.Supplier;
|
||||
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
import org.springframework.ai.chat.metadata.Usage;
|
||||
import org.springframework.ai.chat.model.ChatResponse;
|
||||
import org.springframework.batch.infrastructure.item.ItemProcessor;
|
||||
import org.springframework.stereotype.Component;
|
||||
|
||||
import io.github.resilience4j.ratelimiter.RateLimiter;
|
||||
import io.github.resilience4j.ratelimiter.RateLimiterRegistry;
|
||||
import io.github.resilience4j.retry.Retry;
|
||||
import io.github.resilience4j.retry.RetryRegistry;
|
||||
import tools.jackson.databind.JsonNode;
|
||||
import tools.jackson.databind.json.JsonMapper;
|
||||
|
||||
/**
|
||||
* The ItemProcessor. One ticket in, one {@link Classification} out.
|
||||
*
|
||||
* <p>Three layers wrap the single model call, innermost first:
|
||||
* <ol>
|
||||
* <li>the call itself, through {@link ChatClient};</li>
|
||||
* <li>a Resilience4j {@link RateLimiter} that takes a permit BEFORE every attempt;</li>
|
||||
* <li>a Resilience4j {@link Retry} outside it, so a retried attempt takes a fresh permit.</li>
|
||||
* </ol>
|
||||
* Parsing is outside the retry on purpose: an answer that arrived and is not JSON is not a
|
||||
* transient fault, and a Resilience4j retry would pay for the same bad answer again.
|
||||
*/
|
||||
@Component
|
||||
public class TicketClassifier implements ItemProcessor<Ticket, Classification> {
|
||||
|
||||
static final String SYSTEM = """
|
||||
You classify customer support tickets. Reply with ONLY a JSON object:
|
||||
{"category": one of BILLING, TECHNICAL, ACCOUNT, SHIPPING, FEATURE_REQUEST, OTHER,
|
||||
"priority": integer 1 (low) to 4 (urgent), "confidence": number between 0 and 1}
|
||||
""";
|
||||
|
||||
private final ChatClient chatClient;
|
||||
private final RateLimiter limiter;
|
||||
private final Retry retry;
|
||||
private final UsageMeter meter;
|
||||
private final LlmJobProperties props;
|
||||
private final JsonMapper json = JsonMapper.builder().build();
|
||||
|
||||
public TicketClassifier(ChatClient chatClient, RateLimiterRegistry limiters, RetryRegistry retries,
|
||||
UsageMeter meter, LlmJobProperties props) {
|
||||
this.chatClient = chatClient;
|
||||
this.limiter = limiters.rateLimiter("llm");
|
||||
this.retry = retries.retry("llm");
|
||||
this.meter = meter;
|
||||
this.props = props;
|
||||
this.retry.getEventPublisher().onRetry(e -> meter.retries.incrementAndGet());
|
||||
}
|
||||
|
||||
@Override
|
||||
public Classification process(Ticket ticket) {
|
||||
BadModelOutputException last = null;
|
||||
for (int ask = 0; ask <= props.reask(); ask++) {
|
||||
ChatResponse response = callModel(ticket);
|
||||
Usage usage = response.getMetadata().getUsage();
|
||||
meter.billedCalls.incrementAndGet();
|
||||
meter.promptTokens.addAndGet(usage.getPromptTokens());
|
||||
meter.completionTokens.addAndGet(usage.getCompletionTokens());
|
||||
String text = response.getResult().getOutput().getText();
|
||||
try {
|
||||
return parse(ticket.id(), text);
|
||||
} catch (BadModelOutputException e) {
|
||||
meter.badReplies.incrementAndGet();
|
||||
last = e;
|
||||
}
|
||||
}
|
||||
meter.skipped.incrementAndGet();
|
||||
throw last;
|
||||
}
|
||||
|
||||
private ChatResponse callModel(Ticket ticket) {
|
||||
Supplier<ChatResponse> attempt = () -> {
|
||||
meter.attempts.incrementAndGet();
|
||||
try {
|
||||
return chatClient.prompt()
|
||||
.user("Ticket #" + ticket.id() + "\nSubject: " + ticket.subject() + "\nBody: " + ticket.body())
|
||||
.call()
|
||||
.chatResponse();
|
||||
} catch (ProviderRateLimitException e) {
|
||||
meter.rejected429.incrementAndGet();
|
||||
throw e;
|
||||
} catch (ProviderUnavailableException e) {
|
||||
meter.unavailable503.incrementAndGet();
|
||||
throw e;
|
||||
}
|
||||
};
|
||||
// Retry( RateLimiter( call ) ): every retry goes back through the limiter.
|
||||
return Retry.decorateSupplier(retry, RateLimiter.decorateSupplier(limiter, attempt)).get();
|
||||
}
|
||||
|
||||
Classification parse(long ticketId, String raw) {
|
||||
String text = raw == null ? "" : raw.strip();
|
||||
if (text.startsWith("```")) {
|
||||
text = text.replaceFirst("^```[a-zA-Z]*\\s*", "").replaceFirst("\\s*```$", "").strip();
|
||||
}
|
||||
JsonNode node;
|
||||
try {
|
||||
node = json.readTree(text);
|
||||
} catch (RuntimeException e) {
|
||||
throw new BadModelOutputException("not JSON: " + e.getClass().getSimpleName(), raw);
|
||||
}
|
||||
JsonNode category = node.get("category");
|
||||
JsonNode priority = node.get("priority");
|
||||
JsonNode confidence = node.get("confidence");
|
||||
if (category == null || priority == null || confidence == null) {
|
||||
throw new BadModelOutputException("missing field", raw);
|
||||
}
|
||||
if (!Classification.CATEGORIES.contains(category.asString())) {
|
||||
throw new BadModelOutputException("category not in schema: " + category.asString(), raw);
|
||||
}
|
||||
if (!priority.isInt() || priority.asInt() < 1 || priority.asInt() > 4) {
|
||||
throw new BadModelOutputException("priority out of range", raw);
|
||||
}
|
||||
return new Classification(ticketId, category.asString(), priority.asInt(), confidence.asDouble());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
/**
|
||||
* Builds ticket number {@code id} from nothing but the number, so 100,000 tickets exist without a
|
||||
* file or a table, and the same id is the same ticket in every JVM. SYNTHETIC text assembled from
|
||||
* phrase lists; it is not real customer data.
|
||||
*/
|
||||
public final class TicketGenerator {
|
||||
|
||||
private static final String[] OPENERS = {
|
||||
"Hello,", "Hi team,", "Urgent:", "Quick question,", "Following up,", "To whom it may concern,", "Hey,"};
|
||||
private static final String[][] TOPICS = {
|
||||
{"I was charged twice for invoice", "my refund has not arrived for order", "the VAT on my last invoice looks wrong for plan",
|
||||
"please update the card on file for subscription", "the discount code was not applied to order"},
|
||||
{"the export job crashes with a timeout in workspace", "the API returns 500 for endpoint", "login loops forever on device",
|
||||
"the dashboard shows stale numbers for project", "webhook deliveries stopped for integration"},
|
||||
{"I cannot reset the password for user", "please merge the duplicate accounts for team", "the 2FA code never arrives for",
|
||||
"remove my former colleague from workspace", "I need to change the owner email on account"},
|
||||
{"my parcel is stuck in customs for order", "the tracking page says delivered but nothing arrived for order",
|
||||
"wrong size shipped on order", "the courier left order", "I need to change the delivery address of order"},
|
||||
{"it would be great to have dark mode in", "can you add CSV import to", "please support single sign-on for",
|
||||
"a bulk edit option for", "we would love an API endpoint for"},
|
||||
{"thanks for the quick help with", "I am writing about something else entirely regarding", "just checking the status of",
|
||||
"can someone call me about", "feedback on the onboarding of"}};
|
||||
private static final String[] TAILS = {
|
||||
"This is the third time I am asking.", "It started yesterday after the update.", "Our whole team is blocked.",
|
||||
"No rush, but please reply this week.", "I attached a screenshot earlier.", "Please escalate if needed.",
|
||||
"Happy to jump on a call.", "This affects roughly twenty users.", "I already tried the steps in your help centre."};
|
||||
|
||||
private TicketGenerator() {
|
||||
}
|
||||
|
||||
public static long mix(long x) {
|
||||
x += 0x9E3779B97F4A7C15L;
|
||||
x = (x ^ (x >>> 30)) * 0xBF58476D1CE4E5B9L;
|
||||
x = (x ^ (x >>> 27)) * 0x94D049BB133111EBL;
|
||||
return x ^ (x >>> 31);
|
||||
}
|
||||
|
||||
public static Ticket ticket(long id) {
|
||||
long h = mix(id * 1_000_003L + 17);
|
||||
int topic = (int) Math.floorMod(h, (long) TOPICS.length);
|
||||
String[] phrases = TOPICS[topic];
|
||||
String lead = phrases[(int) Math.floorMod(h >>> 8, (long) phrases.length)];
|
||||
String opener = OPENERS[(int) Math.floorMod(h >>> 16, (long) OPENERS.length)];
|
||||
StringBuilder body = new StringBuilder(opener).append(' ').append(lead).append(" #")
|
||||
.append(Math.floorMod(h >>> 24, 90_000L) + 10_000).append(". ");
|
||||
int tails = 1 + (int) Math.floorMod(h >>> 40, 4L);
|
||||
for (int i = 0; i < tails; i++) {
|
||||
body.append(TAILS[(int) Math.floorMod(mix(h + i) >>> 5, (long) TAILS.length)]).append(' ');
|
||||
}
|
||||
return new Ticket(id, "Ticket about " + lead.split(" ")[Math.min(3, lead.split(" ").length - 1)], body.toString().trim());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
import org.springframework.batch.infrastructure.item.ExecutionContext;
|
||||
import org.springframework.batch.infrastructure.item.ItemStreamReader;
|
||||
|
||||
/**
|
||||
* Reads ticket 1..total. It is restartable for the same reason the CSV reader in the earlier
|
||||
* Spring Batch article is: it saves its position in the step's ExecutionContext, and Spring Batch
|
||||
* persists that context in the SAME transaction as the chunk's writes.
|
||||
*/
|
||||
public class TicketReader implements ItemStreamReader<Ticket> {
|
||||
|
||||
static final String KEY = "ticket.reader.next";
|
||||
|
||||
private final long total;
|
||||
private long next = 1;
|
||||
|
||||
public TicketReader(long total) {
|
||||
this.total = total;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void open(ExecutionContext executionContext) {
|
||||
next = executionContext.getLong(KEY, 1L);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void update(ExecutionContext executionContext) {
|
||||
executionContext.putLong(KEY, next);
|
||||
}
|
||||
|
||||
@Override
|
||||
public Ticket read() {
|
||||
return next > total ? null : TicketGenerator.ticket(next++);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package com.ankurm.batchllm;
|
||||
|
||||
import java.util.concurrent.atomic.AtomicLong;
|
||||
|
||||
import org.springframework.stereotype.Component;
|
||||
|
||||
/**
|
||||
* Money and calls spent since the last ledger row was committed. The classifier adds to it from
|
||||
* many threads; the writer drains it inside the chunk transaction and, if that transaction rolls
|
||||
* back, puts the amounts back -- the calls were still made and billed.
|
||||
*/
|
||||
@Component
|
||||
public class UsageMeter {
|
||||
|
||||
public record Snapshot(long attempts, long rejected429, long unavailable503, long retries, long billedCalls,
|
||||
long promptTokens, long completionTokens, long badReplies, long skipped) {
|
||||
|
||||
public Snapshot plus(Snapshot o) {
|
||||
return new Snapshot(attempts + o.attempts, rejected429 + o.rejected429, unavailable503 + o.unavailable503,
|
||||
retries + o.retries, billedCalls + o.billedCalls, promptTokens + o.promptTokens,
|
||||
completionTokens + o.completionTokens, badReplies + o.badReplies, skipped + o.skipped);
|
||||
}
|
||||
}
|
||||
|
||||
final AtomicLong attempts = new AtomicLong();
|
||||
final AtomicLong rejected429 = new AtomicLong();
|
||||
final AtomicLong unavailable503 = new AtomicLong();
|
||||
final AtomicLong retries = new AtomicLong();
|
||||
final AtomicLong billedCalls = new AtomicLong();
|
||||
final AtomicLong promptTokens = new AtomicLong();
|
||||
final AtomicLong completionTokens = new AtomicLong();
|
||||
final AtomicLong badReplies = new AtomicLong();
|
||||
final AtomicLong skipped = new AtomicLong();
|
||||
|
||||
public Snapshot drain() {
|
||||
return new Snapshot(attempts.getAndSet(0), rejected429.getAndSet(0), unavailable503.getAndSet(0),
|
||||
retries.getAndSet(0), billedCalls.getAndSet(0), promptTokens.getAndSet(0),
|
||||
completionTokens.getAndSet(0), badReplies.getAndSet(0), skipped.getAndSet(0));
|
||||
}
|
||||
|
||||
public void restore(Snapshot s) {
|
||||
attempts.addAndGet(s.attempts());
|
||||
rejected429.addAndGet(s.rejected429());
|
||||
unavailable503.addAndGet(s.unavailable503());
|
||||
retries.addAndGet(s.retries());
|
||||
billedCalls.addAndGet(s.billedCalls());
|
||||
promptTokens.addAndGet(s.promptTokens());
|
||||
completionTokens.addAndGet(s.completionTokens());
|
||||
badReplies.addAndGet(s.badReplies());
|
||||
skipped.addAndGet(s.skipped());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user