diff --git a/README.md b/README.md index 3da83b9..1c3f5f7 100644 --- a/README.md +++ b/README.md @@ -31,5 +31,6 @@ Runnable companion code for the Spring AI articles on [ankurm.com](https://ankur | [`audio/`](audio) | A meeting-minutes generator: `OpenAiAudioTranscriptionModel` -> `ChatClient.entity(Minutes.class)` -> `OpenAiAudioSpeechModel`. Splitting a 30 MB recording at a silence under the 25 MB upload limit (and what a blind split does), SDK retries that re-upload the file, the 4096-character speech limit, three replies that fail to parse, and a cost model with formulas. The audio endpoints are a local mock and the chat model a stub; no real model was called. Spring Boot 4.1.2, Spring AI 2.0.2, Java 25. | [Speech-to-Text and Text-to-Speech with Spring AI: Build a Meeting-Minutes Generator](https://ankurm.com/spring-ai-speech-to-text-text-to-speech-meeting-minutes/) | | [`graphrag-neo4j/`](graphrag-neo4j) | Entities and relations extracted with `ChatClient.entity(...)` into a typed Neo4j graph with a vector index on entity names (Testcontainers, `neo4j:2026.10.1-community`; also run on 5.26), then multi-hop retrieval by span linking, variable-length Cypher paths and path ranking, against plain vector RAG over the same 35 chunks and 9 questions: complete evidence 7/9 for the graph, 5/9 for vector k=8, 2/9 for k=4, with the traps reproduced (duplicate entity under an alias, an embedding merge that joins the wrong incidents, a reply with a sentence before the JSON fence, a hostile relation type, hop explosion). **The extraction step is a rule-based stand-in, not a language model**; Neo4j, embeddings and Cypher are real. Spring Boot 4.1.2, Spring AI 2.0.2, Java 25. | [GraphRAG in Java with Neo4j and Spring AI](https://ankurm.com/graphrag-java-neo4j-spring-ai-knowledge-graph/) | | [`image-gen/`](image-gen) | `ImageModel` and `ImagePrompt` with `OpenAiImageOptions` (what the real `OpenAiImageModel` puts on the wire, a per-call options object that silently resets the model, a `List` prompt that throws), a prompt gate with `ModerationModel` that fails closed, and safe storage in MinIO: sniffed content type, byte and pixel caps, server-chosen keys, a fetcher for provider URLs, presigned reads. A per-image price table from the vendors' pages dated 2026-10-28. No image provider called: the `ImageModel` and `ModerationModel` are hand-written fakes and a mock server; the MinIO tests use a real server through Testcontainers. Spring Boot 4.1.2, Spring AI 2.0.2, Java 25. | [Image Generation and Vision in Spring AI: ImageModel, Prompt Options and Safe Storage](https://ankurm.com/spring-ai-image-generation-imagemodel-s3-minio-safe-storage/) | +| [`batch-llm/`](batch-llm) | 100,000 synthetic support tickets classified by a `ChatClient` inside a Spring Batch 6 chunk step: a Resilience4j rate limiter and retry around the model call (and where Spring Batch's own retry and skip belong instead), replies that are not valid JSON skipped and recorded, a restart after `kill -9` in a second JVM that ends with a byte-identical result, and a cost ledger row committed with every chunk. **The model is a deterministic simulated `ChatModel`, not a language model; token counts are estimates and prices are configurable placeholders.** Spring Boot 4.1.2, Spring Batch 6.0.6, Spring AI 2.0.2, Resilience4j 2.4.0, Java 25. | [Spring AI + Spring Batch: Classifying 100,000 Support Tickets with an LLM](https://ankurm.com/spring-ai-spring-batch-llm-ticket-classification/) | Upgrading from Spring AI 1.x: [migration guide](https://ankurm.com/spring-ai-1-to-2-migration-guide/). diff --git a/batch-llm/.gitignore b/batch-llm/.gitignore new file mode 100644 index 0000000..2f7896d --- /dev/null +++ b/batch-llm/.gitignore @@ -0,0 +1 @@ +target/ diff --git a/batch-llm/README.md b/batch-llm/README.md new file mode 100644 index 0000000..bfff830 --- /dev/null +++ b/batch-llm/README.md @@ -0,0 +1,29 @@ +# batch-llm + +Classify 100,000 synthetic support tickets with a `ChatClient` inside a Spring Batch 6 chunk step: a Resilience4j rate limiter and retry around the model call, skip-and-record for replies that are not valid JSON, restart after `kill -9`, and a cost ledger row committed with every chunk. + +**No real model was called.** The "provider" is `SimulatedProvider`, a hand-written deterministic `ChatModel` (same idea as the `ScriptedChatModel` in the other modules) that returns valid JSON, malformed JSON, 429s, 503s and slow answers from a seeded fault plan, and enforces its own calls-per-second quota. Tickets are generated in code from their number. Token counts are estimated (characters / 4) and prices in `application.yml` are **placeholders you configure**, not any provider's prices. Nothing here measures a real model's accuracy, latency or price. + +| | | +|---|---| +| Spring Boot | 4.1.2 | +| Spring Batch | 6.0.6 (managed by Boot; `ChunkOrientedStep`) | +| Spring AI | 2.0.2 (`spring-ai-client-chat`) | +| Resilience4j | 2.4.0 (`resilience4j-spring-boot4`) | +| Java | 25 | +| Run | `scripts/run-all.sh` (tests, then the 100,000-ticket job twice; `TICKETS=20000 scripts/run-all.sh` is faster) | + +## Source + +| File | What it is | +|---|---| +| [`TicketClassifier`](src/main/java/com/ankurm/batchllm/TicketClassifier.java) | the `ItemProcessor`: `Retry(RateLimiter(ChatClient call))`, then parse and validate the reply | +| [`LlmJobConfig`](src/main/java/com/ankurm/batchllm/LlmJobConfig.java) | the `ChunkOrientedStep`: fault tolerant, skip policy, virtual-thread executor, where Spring Batch retry would go | +| [`ClassificationWriter`](src/main/java/com/ankurm/batchllm/ClassificationWriter.java) | results and the `COST_LEDGER` row in one transaction | +| [`SkippedTicketListener`](src/main/java/com/ankurm/batchllm/SkippedTicketListener.java) | records skipped tickets in their own transaction | +| [`SimulatedProvider`](src/main/java/com/ankurm/batchllm/SimulatedProvider.java) | the fake model and its fault plan | +| [`LlmJobRunner`](src/main/java/com/ankurm/batchllm/LlmJobRunner.java) | start or restart, `JobOperator.recover`, final report | + +## Output (`output/`) + +`01-versions` classpath and builder methods, `02-simulated-provider` the fault mix, `03-throttle` with and without a limiter, `04-skip-and-ledger` every ticket accounted for and ledger == provider meter, `05-retry-layers` which layer retries what, `06-skip-limit-and-restart` a model regression stops the job, `08-clean-run` the full 100,000, `09-crash-restart` two JVMs and `kill -9`, `10-compare` the resumed database equals the clean one. diff --git a/batch-llm/output/01-versions.txt b/batch-llm/output/01-versions.txt new file mode 100644 index 0000000..ec1f5bb --- /dev/null +++ b/batch-llm/output/01-versions.txt @@ -0,0 +1,23 @@ +# What is on the classpath and what the builder offers + +Spring Boot 4.1.2 +Spring Framework 7.0.10 +Spring Batch (core) 6.0.6 +Spring AI (client) 2.0.2 +Resilience4j 2.4.0 +Java 25.0.5 + +Resilience4j config read from application.yml by resilience4j-spring-boot4: + ratelimiter 'llm': limitForPeriod=200 refreshPeriod=PT0.1S timeout=PT1M + retry 'llm': maxAttempts=4 + +ChunkOrientedStepBuilder methods that matter here (javap of spring-batch-core): + faultTolerant [] + taskExecutor [org.springframework.core.task.AsyncTaskExecutor] + retryPolicy [org.springframework.core.retry.RetryPolicy] + retry [[Ljava.lang.Class;] + retryLimit [long] + skip [[Ljava.lang.Class;] + skipLimit [long] + skipListener [org.springframework.batch.core.listener.SkipListener] + processorNonTransactional [] diff --git a/batch-llm/output/02-simulated-provider.txt b/batch-llm/output/02-simulated-provider.txt new file mode 100644 index 0000000..7d86cdd --- /dev/null +++ b/batch-llm/output/02-simulated-provider.txt @@ -0,0 +1,15 @@ +# The simulated provider: same input, same outcome + +Two separate provider instances, same seed, tickets 1..20,000, one call each: identical outcome lists = true + +Outcome of the first call per ticket (the plan: 1.5% 429, 0.5% 503, 0.3% never valid, 1.2% invalid on the first reply only): + 429 304 1.52% + 503 93 0.47% + bad JSON: category not in schema 76 0.38% + bad JSON: missing field 77 0.39% + bad JSON: not JSON 137 0.69% + valid JSON 17,361 86.81% + valid JSON (in a markdown fence, stripped) 1,952 9.76% + +Quota check: quota 100 calls/second, 300 calls fired back to back -> 201 rejected with 429 (quota rejections plus any injected 429) +Every number above is a property of the TEST PLAN, not a measurement of any provider. diff --git a/batch-llm/output/03-throttle.txt b/batch-llm/output/03-throttle.txt new file mode 100644 index 0000000..ee42e81 --- /dev/null +++ b/batch-llm/output/03-throttle.txt @@ -0,0 +1,11 @@ +# 4,000 tickets against a simulated quota of 800 calls per second + +Both runs: 64 virtual threads, Resilience4j retry max 4 attempts, 25 ms x2 backoff, the same fault plan. +Time is indicative only (shared 2-vCPU sandbox, other processes running). + +scenario ok bad JSON exhausted calls quota 429s retries seconds +no limiter 3066 13 921 7,596 4,396 3,543 3.5 +limiter 600/s 3986 14 0 4,146 0 85 6.9 + +'exhausted' = tickets whose 4 attempts all met a 429/503: with a step around this call, the step fails. +'calls' counts every attempt the provider received, rejected or not. diff --git a/batch-llm/output/04-skip-and-ledger.txt b/batch-llm/output/04-skip-and-ledger.txt new file mode 100644 index 0000000..035df9c --- /dev/null +++ b/batch-llm/output/04-skip-and-ledger.txt @@ -0,0 +1,33 @@ +# 3,000 tickets through the real job + +job status COMPLETED +step: read=3000 write=2987 processSkip=13 commits=30 rollbacks=0 +TICKET_RESULT rows 2987 +SKIPPED_TICKET rows 13 +classified + skipped 3000 (tickets: 3000) + +Skipped tickets and why (first one of each kind): + ticket 46 category not in schema raw reply: {"category":"Account & Payments","priority":1,"confidence":0.74} + ticket 234 not JSON raw reply: Sure! Here is the classification you asked for: {"category":"BILLING","priority":4,"confidence":0.7} + ticket 491 missing field raw reply: {"category":"OTHER","priority":2} + +First 3 ledger rows (one per committed chunk): + tickets written skipped attempts 429s 503s bad promptTok complTok cost + 1-100 99 1 104 1 1 3 10831 1469 0.007619000 + 101-200 100 0 100 0 0 0 10574 1453 0.007466500 + 201-300 99 1 105 2 1 3 10736 1495 0.007610500 + +Ledger totals vs the provider's own meter (they must agree when nothing was lost): + ledger provider + calls received 3,104 3,104 + 429 rejections 39 39 + 503 rejections 14 14 + billed calls 3,051 3,051 + prompt tokens 324,790 324,790 + completion tokens 44,107 44,107 + +Cost = prompt tokens x input price + completion tokens x output price, per million tokens. + table 'EXAMPLE-NOT-A-REAL-PRICE' (input 0.50, output 1.50): 0.228555500 <- the ledger's total + same tokens under table 'OTHER-EXAMPLE' (input 2.00, output 6.00): 0.914222000 + Both tables are invented inputs. Tokens are the measurement; money is tokens x a table you supply. + Tokens spent on replies that were then skipped: bad replies billed = 64 of 3051 billed calls. diff --git a/batch-llm/output/05-retry-layers.txt b/batch-llm/output/05-retry-layers.txt new file mode 100644 index 0000000..5bdcd96 --- /dev/null +++ b/batch-llm/output/05-retry-layers.txt @@ -0,0 +1,16 @@ +# Which layer retries what: counts taken from the provider's meter + +(a) A transient DB failure on the write of chunk 2 (tickets 101-200), 300 tickets, Spring Batch write retry = 2 + no failure status=COMPLETED billed calls=304 rollbacks=0 + failure + Spring Batch retry status=COMPLETED billed calls=304 rollbacks=0 + -> the write was retried inside the same transaction; the model was not asked again. + +(b) The provider answers 429 to everything, 1 ticket. Calls the provider received for that ticket: + Resilience4j retry only (max 4 attempts) status=FAILED attempts=4 + + Spring Batch retry of the same exception (2) status=FAILED attempts=12 + -> 4 x (1 + 2) = 12. Two retry layers on the same failure multiply, and each attempt also takes a limiter permit. + +(c) A poison row at WRITE time (ticket 250): chunk 3 (201-300) rolls back and is scanned one item at a time. + processorNonTransactional (cached) status=COMPLETED billed calls=304 (no-failure baseline: 304) rollbacks=2 ledger rows=100 skipped=3 + default (processor runs again per item) status=COMPLETED billed calls=403 (extra over baseline: 99) rollbacks=2 ledger rows=100 skipped=3 + -> without processorNonTransactional(), every item of the failed chunk is classified a second time and paid for twice. diff --git a/batch-llm/output/06-skip-limit-and-restart.txt b/batch-llm/output/06-skip-limit-and-restart.txt new file mode 100644 index 0000000..f5db787 --- /dev/null +++ b/batch-llm/output/06-skip-limit-and-restart.txt @@ -0,0 +1,13 @@ +# A model regression, the skip limit, and a restart with a fixed model + +Execution 1: simulated model returns invalid JSON for 50% of tickets, no re-ask, skip limit 50, 2,000 tickets. + status / reason FAILED / org.springframework.batch.core.step.FatalStepExecutionException: Unable to process chunk + tickets read 200 of 2000 + calls the provider billed before the step stopped: 200 (a full run needs about 2,000) + committed results 52 skipped rows 50 + +Execution 2: same job parameters, same database, model fixed, skip limit 1000. + status COMPLETED + tickets read this time 1900 (execution 1 had committed 100, so the reader resumed after them) + committed results 1942 skipped rows 58 classified + skipped = 2000 + tickets present in BOTH tables: 0 (stale skip rows from execution 1 are deleted when the ticket is classified) diff --git a/batch-llm/output/08-clean-run.txt b/batch-llm/output/08-clean-run.txt new file mode 100644 index 0000000..8beccdf --- /dev/null +++ b/batch-llm/output/08-clean-run.txt @@ -0,0 +1,20 @@ +# One uninterrupted run, 100000 tickets (java -jar, fresh H2 file, fresh provider meter) + +START run.id=clean tickets=100000 chunk=100 +JOB FINISHED: executionId=1 status=COMPLETED seconds=55.2 +---- REPORT ---- +tickets total=100000 classified=99691 skipped=309 unaccounted=0 +step executions: + {ID=1, JOB=1, STATUS=COMPLETED, READS=100000, WRITES=99691, SKIPS=309, COMMITS=1000, ROLLBACKS=0} +ledger (committed chunks only): {ROWS_=1000, WRITTEN=99691, SKIPPED=309, ATTEMPTS=103519, R429=1531, U503=458, RETRIES=1989, BILLED=101530, BAD=1839, PT=10809111, CT=1467192, COST=7.605343500} +price table 'EXAMPLE-NOT-A-REAL-PRICE': input 0.50 / output 1.50 per million tokens (configurable placeholder) +provider's own meter (survives kill -9): + received 103,519 + rejected_429 1,531 + window_429 0 + unavailable_503 458 + billed 101,530 + prompt_tokens 10,809,111 + completion_tokens 1,467,192 +billed calls the ledger does not know about (work lost in a crash or rollback): 0 +results checksum (sha-256 over ordered rows): 98bdd29feb6d6f8d diff --git a/batch-llm/output/09-crash-restart.txt b/batch-llm/output/09-crash-restart.txt new file mode 100644 index 0000000..7180c8a --- /dev/null +++ b/batch-llm/output/09-crash-restart.txt @@ -0,0 +1,30 @@ +# Crash and restart, 100000 tickets: JVM 1 killed with kill -9 after 400 committed chunks, JVM 2 resumes + +## JVM 1 (progress lines, then SIGKILL) +PROGRESS committed-chunks-this-jvm=250 last-ticket=25000 +PROGRESS committed-chunks-this-jvm=300 last-ticket=30000 +PROGRESS committed-chunks-this-jvm=350 last-ticket=35000 +PROGRESS committed-chunks-this-jvm=400 last-ticket=40000 +(process killed by SIGKILL; no shutdown hook ran) + +## JVM 2 +RECOVER: execution 1 is STARTED but its JVM is gone; marking it FAILED so the instance can restart +START run.id=crash-demo tickets=100000 chunk=100 +JOB FINISHED: executionId=33 status=COMPLETED seconds=34.2 +---- REPORT ---- +tickets total=100000 classified=99691 skipped=309 unaccounted=0 +step executions: + {ID=1, JOB=1, STATUS=FAILED, READS=39600, WRITES=39493, SKIPS=107, COMMITS=396, ROLLBACKS=0} + {ID=33, JOB=33, STATUS=COMPLETED, READS=60400, WRITES=60198, SKIPS=202, COMMITS=604, ROLLBACKS=0} +ledger (committed chunks only): {ROWS_=1000, WRITTEN=99691, SKIPPED=309, ATTEMPTS=103519, R429=1531, U503=458, RETRIES=1989, BILLED=101530, BAD=1839, PT=10809111, CT=1467192, COST=7.605343500} +price table 'EXAMPLE-NOT-A-REAL-PRICE': input 0.50 / output 1.50 per million tokens (configurable placeholder) +provider's own meter (survives kill -9): + received 104,350 + rejected_429 1,543 + window_429 0 + unavailable_503 462 + billed 102,302 + prompt_tokens 10,891,290 + completion_tokens 1,478,281 +billed calls the ledger does not know about (work lost in a crash or rollback): 772 +results checksum (sha-256 over ordered rows): 98bdd29feb6d6f8d diff --git a/batch-llm/output/10-compare.txt b/batch-llm/output/10-compare.txt new file mode 100644 index 0000000..0bca6f0 --- /dev/null +++ b/batch-llm/output/10-compare.txt @@ -0,0 +1,15 @@ +# Clean run vs crashed-and-resumed run + +## target/clean/run.log +tickets total=100000 classified=99691 skipped=309 unaccounted=0 +ledger (committed chunks only): {ROWS_=1000, WRITTEN=99691, SKIPPED=309, ATTEMPTS=103519, R429=1531, U503=458, RETRIES=1989, BILLED=101530, BAD=1839, PT=10809111, CT=1467192, COST=7.605343500} +billed calls the ledger does not know about (work lost in a crash or rollback): 0 +results checksum (sha-256 over ordered rows): 98bdd29feb6d6f8d + +## target/crash/jvm2.log +tickets total=100000 classified=99691 skipped=309 unaccounted=0 +ledger (committed chunks only): {ROWS_=1000, WRITTEN=99691, SKIPPED=309, ATTEMPTS=103519, R429=1531, U503=458, RETRIES=1989, BILLED=101530, BAD=1839, PT=10809111, CT=1467192, COST=7.605343500} +billed calls the ledger does not know about (work lost in a crash or rollback): 772 +results checksum (sha-256 over ordered rows): 98bdd29feb6d6f8d + +checksums identical: 98bdd29feb6d6f8d diff --git a/batch-llm/pom.xml b/batch-llm/pom.xml new file mode 100644 index 0000000..96f94da --- /dev/null +++ b/batch-llm/pom.xml @@ -0,0 +1,70 @@ + + + 4.0.0 + + + org.springframework.boot + spring-boot-starter-parent + 4.1.2 + + + + com.ankurm + batch-llm + 1.0.0 + batch-llm + Classify 100,000 support tickets with a ChatClient inside a Spring Batch chunk step: Resilience4j rate limiting and retry, skip on malformed JSON, restart after kill -9, and a per-chunk cost ledger. Runs against a deterministic simulated chat model. + + + 25 + 2.0.2 + 2.4.0 + + + + + + org.springframework.ai + spring-ai-bom + ${spring-ai.version} + pom + import + + + + + + org.springframework.bootspring-boot-starter-batch-jdbc + com.h2databaseh2runtime + org.springframework.aispring-ai-client-chat + io.github.resilience4jresilience4j-spring-boot4${resilience4j.version} + + org.springframework.boot + spring-boot-starter-test + test + + + org.springframework.boot + spring-boot-starter-batch-jdbc-test + test + + + + + + + org.springframework.boot + spring-boot-maven-plugin + + + org.apache.maven.plugins + maven-surefire-plugin + + -Duser.timezone=UTC -Dstdout.encoding=UTF-8 -Dfile.encoding=UTF-8 + + + + + diff --git a/batch-llm/scripts/common.sh b/batch-llm/scripts/common.sh new file mode 100755 index 0000000..ce11ce9 --- /dev/null +++ b/batch-llm/scripts/common.sh @@ -0,0 +1,8 @@ +# sourced by the other scripts +export JAVA_HOME="${JDK25:-/tmp/tools/jdk-25.0.5+8}" # Java 25 required +export PATH="$JAVA_HOME/bin:/opt/maven/bin:$PATH" +cd "$(dirname "${BASH_SOURCE[0]}")/.." +JAR=target/batch-llm-1.0.0.jar +TICKETS="${TICKETS:-100000}" +CHUNK=100 +clean_log() { grep -v "^Picked up JAVA_TOOL_OPTIONS" ; } diff --git a/batch-llm/scripts/compare.sh b/batch-llm/scripts/compare.sh new file mode 100755 index 0000000..bf5c71f --- /dev/null +++ b/batch-llm/scripts/compare.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +# Same tickets, same fault plan: the crashed-and-resumed database must hold exactly the clean run's results. +set -euo pipefail +source "$(dirname "$0")/common.sh" +{ echo "# Clean run vs crashed-and-resumed run"; echo + for f in target/clean/run.log target/crash/jvm2.log; do + echo "## $f" + grep -E "^tickets total|^ledger|results checksum|billed calls the ledger" "$f" + echo + done; } > output/10-compare.txt +cat output/10-compare.txt +a=$(grep checksum target/clean/run.log | awk '{print $NF}'); b=$(grep checksum target/crash/jvm2.log | awk '{print $NF}') +[ "$a" = "$b" ] && echo "checksums identical: $a" | tee -a output/10-compare.txt || { echo "CHECKSUMS DIFFER $a $b" | tee -a output/10-compare.txt; exit 1; } diff --git a/batch-llm/scripts/crash-restart.sh b/batch-llm/scripts/crash-restart.sh new file mode 100755 index 0000000..9244d49 --- /dev/null +++ b/batch-llm/scripts/crash-restart.sh @@ -0,0 +1,32 @@ +#!/usr/bin/env bash +# Two real JVMs. JVM 1 is killed with SIGKILL once it has committed KILL_AT chunks; JVM 2 starts the +# same job (same run.id, same H2 file) and finishes it. Nothing is shut down gracefully. +set -euo pipefail +source "$(dirname "$0")/common.sh" +[ -f "$JAR" ] || mvn -B -q -T1 package -DskipTests +CHUNKS=$((TICKETS / CHUNK)) +KILL_AT=$(( (CHUNKS * 4 / 10 / 50) * 50 )); [ "$KILL_AT" -ge 50 ] || KILL_AT=50 +D=target/crash +rm -rf $D && mkdir -p $D output +OPTS=(--batch.llm.tickets="$TICKETS" --batch.llm.run-id=crash-demo + --spring.datasource.url="jdbc:h2:file:./$D/batchdb;DB_CLOSE_ON_EXIT=FALSE" + --sim.provider.meter-file=$D/provider-meter.bin) + +java -jar "$JAR" "${OPTS[@]}" > $D/jvm1.log 2>&1 & +PID=$! +echo "JVM 1 started, pid $PID; waiting for $KILL_AT committed chunks" +while kill -0 "$PID" 2>/dev/null; do + if grep -q "committed-chunks-this-jvm=$KILL_AT " $D/jvm1.log; then + kill -9 "$PID"; echo "kill -9 sent to JVM 1 (pid $PID)"; break + fi + sleep 0.2 +done +wait "$PID" 2>/dev/null || true +echo "JVM 1 exit code: $? (137 = SIGKILL)" +echo "--- last lines JVM 1 printed:"; clean_log < $D/jvm1.log | tail -3 + +echo "--- JVM 2 (same run.id, same database file):" +java -jar "$JAR" "${OPTS[@]}" 2>&1 | clean_log | grep -v '^PROGRESS' | tee $D/jvm2.log +{ echo "# Crash and restart, $TICKETS tickets: JVM 1 killed with kill -9 after $KILL_AT committed chunks, JVM 2 resumes"; echo; + echo "## JVM 1 (progress lines, then SIGKILL)"; clean_log < $D/jvm1.log | tail -4; echo "(process killed by SIGKILL; no shutdown hook ran)"; echo; + echo "## JVM 2"; cat $D/jvm2.log; } > output/09-crash-restart.txt diff --git a/batch-llm/scripts/run-all.sh b/batch-llm/scripts/run-all.sh new file mode 100755 index 0000000..d54c0ed --- /dev/null +++ b/batch-llm/scripts/run-all.sh @@ -0,0 +1,12 @@ +#!/usr/bin/env bash +# Regenerates every file under output/. No API key, no Docker, no network beyond Maven Central. +# Takes several minutes on a busy 2-vCPU machine: the 100,000-ticket job runs twice (once clean, +# once crashed and resumed). TICKETS=20000 scripts/run-all.sh gives a quicker, smaller version. +set -euo pipefail +source "$(dirname "$0")/common.sh" +rm -rf output target/test-db target/test-meters +mvn -B -q -T1 package 2>&1 | grep -E "Tests run:|BUILD|FAIL|ERROR" || true +scripts/run-clean.sh +scripts/crash-restart.sh +scripts/compare.sh +ls output diff --git a/batch-llm/scripts/run-clean.sh b/batch-llm/scripts/run-clean.sh new file mode 100755 index 0000000..cecdef0 --- /dev/null +++ b/batch-llm/scripts/run-clean.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +# One uninterrupted run of the whole job against a fresh database and a fresh provider meter. +set -euo pipefail +source "$(dirname "$0")/common.sh" +[ -f "$JAR" ] || mvn -B -q -T1 package -DskipTests +rm -rf target/clean && mkdir -p target/clean output +java -jar "$JAR" --batch.llm.tickets="$TICKETS" --batch.llm.run-id=clean \ + --spring.datasource.url="jdbc:h2:file:./target/clean/batchdb;DB_CLOSE_ON_EXIT=FALSE" \ + --sim.provider.meter-file=target/clean/provider-meter.bin 2>&1 | clean_log | tee target/clean/run.log +{ echo "# One uninterrupted run, $TICKETS tickets (java -jar, fresh H2 file, fresh provider meter)"; echo; + grep -v '^PROGRESS' target/clean/run.log; } > output/08-clean-run.txt diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/BadModelOutputException.java b/batch-llm/src/main/java/com/ankurm/batchllm/BadModelOutputException.java new file mode 100644 index 0000000..3b2e98a --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/BadModelOutputException.java @@ -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; + } +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/BatchLlmApplication.java b/batch-llm/src/main/java/com/ankurm/batchllm/BatchLlmApplication.java new file mode 100644 index 0000000..26fab59 --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/BatchLlmApplication.java @@ -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))); + } +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/Classification.java b/batch-llm/src/main/java/com/ankurm/batchllm/Classification.java new file mode 100644 index 0000000..f5a6623 --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/Classification.java @@ -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 CATEGORIES = + Set.of("BILLING", "TECHNICAL", "ACCOUNT", "SHIPPING", "FEATURE_REQUEST", "OTHER"); +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/ClassificationWriter.java b/batch-llm/src/main/java/com/ankurm/batchllm/ClassificationWriter.java new file mode 100644 index 0000000..08614d9 --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/ClassificationWriter.java @@ -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 { + + 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 chunk) { + List 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(); + } +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/LlmJobConfig.java b/batch-llm/src/main/java/com/ankurm/batchllm/LlmJobConfig.java new file mode 100644 index 0000000..053dfd8 --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/LlmJobConfig.java @@ -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 writer, + SkippedTicketListener skipListener, AsyncTaskExecutor llmExecutor, LlmJobProperties props) { + ChunkOrientedStepBuilder step = new StepBuilder("classifyStep", jobRepository) + .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(); + } +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/LlmJobProperties.java b/batch-llm/src/main/java/com/ankurm/batchllm/LlmJobProperties.java new file mode 100644 index 0000000..3411272 --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/LlmJobProperties.java @@ -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) { +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/LlmJobRunner.java b/batch-llm/src/main/java/com/ankurm/batchllm/LlmJobRunner.java new file mode 100644 index 0000000..5a3637d --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/LlmJobRunner.java @@ -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. + * + *

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 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 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 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)); + } +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/Pricing.java b/batch-llm/src/main/java/com/ankurm/batchllm/Pricing.java new file mode 100644 index 0000000..a04239a --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/Pricing.java @@ -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); + } +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/ProviderFaults.java b/batch-llm/src/main/java/com/ankurm/batchllm/ProviderFaults.java new file mode 100644 index 0000000..793e2ca --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/ProviderFaults.java @@ -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) { +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/ProviderMeter.java b/batch-llm/src/main/java/com/ankurm/batchllm/ProviderMeter.java new file mode 100644 index 0000000..57a6a8b --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/ProviderMeter.java @@ -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(); + } +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/ProviderRateLimitException.java b/batch-llm/src/main/java/com/ankurm/batchllm/ProviderRateLimitException.java new file mode 100644 index 0000000..66fcbfe --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/ProviderRateLimitException.java @@ -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); + } +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/ProviderUnavailableException.java b/batch-llm/src/main/java/com/ankurm/batchllm/ProviderUnavailableException.java new file mode 100644 index 0000000..b364dcf --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/ProviderUnavailableException.java @@ -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); + } +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/SimulatedProvider.java b/batch-llm/src/main/java/com/ankurm/batchllm/SimulatedProvider.java new file mode 100644 index 0000000..e0b533e --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/SimulatedProvider.java @@ -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. + * + *

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. + * + *

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 CATEGORIES = List.of("BILLING", "TECHNICAL", "ACCOUNT", "SHIPPING", "FEATURE_REQUEST", "OTHER"); + + private final ProviderFaults faults; + private final ProviderMeter meter; + private final ConcurrentHashMap callsPerTicket = new ConcurrentHashMap<>(); + private final ConcurrentHashMap billedPerTicket = new ConcurrentHashMap<>(); + private final ArrayDeque 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 #' 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 + }; + } +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/SkippedTicketListener.java b/batch-llm/src/main/java/com/ankurm/batchllm/SkippedTicketListener.java new file mode 100644 index 0000000..6a6d04f --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/SkippedTicketListener.java @@ -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. + * + *

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 { + + 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)); + } +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/Ticket.java b/batch-llm/src/main/java/com/ankurm/batchllm/Ticket.java new file mode 100644 index 0000000..ab77e9f --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/Ticket.java @@ -0,0 +1,5 @@ +package com.ankurm.batchllm; + +/** One support ticket. Synthetic: see {@link TicketGenerator}. */ +public record Ticket(long id, String subject, String body) { +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/TicketClassifier.java b/batch-llm/src/main/java/com/ankurm/batchllm/TicketClassifier.java new file mode 100644 index 0000000..5eb4605 --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/TicketClassifier.java @@ -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. + * + *

Three layers wrap the single model call, innermost first: + *

    + *
  1. the call itself, through {@link ChatClient};
    2. + *
  2. a Resilience4j {@link RateLimiter} that takes a permit BEFORE every attempt;
    3. + *
  3. a Resilience4j {@link Retry} outside it, so a retried attempt takes a fresh permit.
    4. + *
+ * 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 { + + 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 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()); + } +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/TicketGenerator.java b/batch-llm/src/main/java/com/ankurm/batchllm/TicketGenerator.java new file mode 100644 index 0000000..4659dc1 --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/TicketGenerator.java @@ -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()); + } +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/TicketReader.java b/batch-llm/src/main/java/com/ankurm/batchllm/TicketReader.java new file mode 100644 index 0000000..a8f1de6 --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/TicketReader.java @@ -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 { + + 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++); + } +} diff --git a/batch-llm/src/main/java/com/ankurm/batchllm/UsageMeter.java b/batch-llm/src/main/java/com/ankurm/batchllm/UsageMeter.java new file mode 100644 index 0000000..d6edfb3 --- /dev/null +++ b/batch-llm/src/main/java/com/ankurm/batchllm/UsageMeter.java @@ -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()); + } +} diff --git a/batch-llm/src/main/resources/application.yml b/batch-llm/src/main/resources/application.yml new file mode 100644 index 0000000..d984cf0 --- /dev/null +++ b/batch-llm/src/main/resources/application.yml @@ -0,0 +1,52 @@ +spring: + main: + banner-mode: off + datasource: + url: jdbc:h2:file:./target/batchdb/llm;DB_CLOSE_ON_EXIT=FALSE + username: sa + password: "" + batch: + job: + enabled: false # LlmJobRunner starts the job itself, with an explicit run.id + jdbc: + initialize-schema: always + sql: + init: + mode: always + +logging: + level: + root: warn + org.springframework.batch.core.step.item: error + +batch: + llm: + tickets: 100000 + chunk-size: 100 + concurrency: 64 + reask: 1 + skip-limit: 1000 + # PLACEHOLDER prices so the arithmetic is easy to follow. Not any provider's price list. + pricing: + version: EXAMPLE-NOT-A-REAL-PRICE + input-per-million: "0.50" + output-per-million: "1.50" + +# The model call's two Resilience4j policies. The scale of time here is deliberately small: see README. +resilience4j: + ratelimiter: + instances: + llm: + limit-for-period: 200 + limit-refresh-period: 100ms + timeout-duration: 60s + retry: + instances: + llm: + max-attempts: 4 + wait-duration: 25ms + enable-exponential-backoff: true + exponential-backoff-multiplier: 2 + retry-exceptions: + - com.ankurm.batchllm.ProviderRateLimitException + - com.ankurm.batchllm.ProviderUnavailableException diff --git a/batch-llm/src/main/resources/schema.sql b/batch-llm/src/main/resources/schema.sql new file mode 100644 index 0000000..84468eb --- /dev/null +++ b/batch-llm/src/main/resources/schema.sql @@ -0,0 +1,29 @@ +CREATE TABLE IF NOT EXISTS TICKET_RESULT ( + TICKET_ID BIGINT PRIMARY KEY, + CATEGORY VARCHAR(32) NOT NULL, + PRIORITY INT NOT NULL, + CONFIDENCE DOUBLE NOT NULL +); +CREATE TABLE IF NOT EXISTS SKIPPED_TICKET ( + TICKET_ID BIGINT PRIMARY KEY, + REASON VARCHAR(200), + RAW_REPLY VARCHAR(400) +); +CREATE TABLE IF NOT EXISTS COST_LEDGER ( + LEDGER_ID BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, + FIRST_TICKET BIGINT NOT NULL, + LAST_TICKET BIGINT NOT NULL, + TICKETS_WRITTEN INT NOT NULL, + TICKETS_SKIPPED INT NOT NULL, + ATTEMPTS BIGINT NOT NULL, + REJECTED_429 BIGINT NOT NULL, + UNAVAILABLE_503 BIGINT NOT NULL, + RETRIES BIGINT NOT NULL, + BILLED_CALLS BIGINT NOT NULL, + BAD_REPLIES BIGINT NOT NULL, + PROMPT_TOKENS BIGINT NOT NULL, + COMPLETION_TOKENS BIGINT NOT NULL, + COST DECIMAL(24,9) NOT NULL, + PRICE_VERSION VARCHAR(64) NOT NULL, + WRITTEN_AT TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL +); diff --git a/batch-llm/src/test/java/com/ankurm/batchllm/JobSkipAndLedgerTest.java b/batch-llm/src/test/java/com/ankurm/batchllm/JobSkipAndLedgerTest.java new file mode 100644 index 0000000..dbcef29 --- /dev/null +++ b/batch-llm/src/test/java/com/ankurm/batchllm/JobSkipAndLedgerTest.java @@ -0,0 +1,94 @@ +package com.ankurm.batchllm; + +import static com.ankurm.batchllm.support.JobSupport.count; +import static com.ankurm.batchllm.support.JobSupport.ledger; +import static com.ankurm.batchllm.support.JobSupport.n; +import static org.assertj.core.api.Assertions.assertThat; + +import java.math.BigDecimal; +import java.util.List; +import java.util.Map; + +import org.junit.jupiter.api.Test; +import org.springframework.batch.core.BatchStatus; +import org.springframework.batch.core.job.Job; +import org.springframework.batch.core.job.JobExecution; +import org.springframework.batch.core.launch.JobOperator; +import org.springframework.context.ConfigurableApplicationContext; +import org.springframework.jdbc.core.JdbcTemplate; + +import com.ankurm.batchllm.ProviderMeter.Slot; +import com.ankurm.batchllm.support.JobSupport; +import com.ankurm.batchllm.support.Transcript; + +class JobSkipAndLedgerTest { + + @Test + void everyTicketIsAccountedForAndTheLedgerMatchesTheProvidersMeter() throws Exception { + try (ConfigurableApplicationContext ctx = SpringContexts.start("skipledger", Map.of("batch.llm.tickets", "3000"))) { + JobExecution exec = JobSupport.run(ctx.getBean(JobOperator.class), ctx.getBean(Job.class), "t1"); + JdbcTemplate jdbc = ctx.getBean(JdbcTemplate.class); + ProviderMeter provider = ctx.getBean(ProviderMeter.class); + Pricing pricing = ctx.getBean(Pricing.class); + + long classified = count(jdbc, "TICKET_RESULT"); + long skipped = count(jdbc, "SKIPPED_TICKET"); + Map l = ledger(jdbc); + var step = exec.getStepExecutions().iterator().next(); + + try (Transcript t = new Transcript("04-skip-and-ledger.txt", "3,000 tickets through the real job")) { + t.line("job status %s", exec.getStatus()); + t.line("step: read=%d write=%d processSkip=%d commits=%d rollbacks=%d", step.getReadCount(), step.getWriteCount(), + step.getProcessSkipCount(), step.getCommitCount(), step.getRollbackCount()); + t.line("TICKET_RESULT rows %d", classified); + t.line("SKIPPED_TICKET rows %d", skipped); + t.line("classified + skipped %d (tickets: 3000)", classified + skipped); + t.blank(); + t.line("Skipped tickets and why (first one of each kind):"); + List> rows = jdbc.queryForList("SELECT TICKET_ID, REASON, RAW_REPLY FROM SKIPPED_TICKET ORDER BY TICKET_ID"); + rows.stream().map(r -> r.get("REASON").toString().replaceAll(":.*", "")).distinct().forEach(kind -> { + var r = rows.stream().filter(x -> x.get("REASON").toString().startsWith(kind)).findFirst().orElseThrow(); + t.line(" ticket %-5s %-24s raw reply: %s", r.get("TICKET_ID"), kind, r.get("RAW_REPLY").toString().replace("\n", "\\n")); + }); + t.blank(); + t.line("First 3 ledger rows (one per committed chunk):"); + t.line(" %-12s %7s %7s %8s %6s %6s %6s %9s %8s %12s", "tickets", "written", "skipped", "attempts", "429s", "503s", "bad", "promptTok", "complTok", "cost"); + jdbc.queryForList("SELECT * FROM COST_LEDGER ORDER BY LEDGER_ID LIMIT 3").forEach(r -> t.line( + " %-12s %7s %7s %8s %6s %6s %6s %9s %8s %12s", r.get("FIRST_TICKET") + "-" + r.get("LAST_TICKET"), + r.get("TICKETS_WRITTEN"), r.get("TICKETS_SKIPPED"), r.get("ATTEMPTS"), r.get("REJECTED_429"), + r.get("UNAVAILABLE_503"), r.get("BAD_REPLIES"), r.get("PROMPT_TOKENS"), r.get("COMPLETION_TOKENS"), r.get("COST"))); + t.blank(); + t.line("Ledger totals vs the provider's own meter (they must agree when nothing was lost):"); + t.line(" %-22s %12s %12s", "", "ledger", "provider"); + t.line(" %-22s %,12d %,12d", "calls received", n(l.get("ATTEMPTS")), provider.get(Slot.RECEIVED)); + t.line(" %-22s %,12d %,12d", "429 rejections", n(l.get("R429")), provider.get(Slot.REJECTED_429)); + t.line(" %-22s %,12d %,12d", "503 rejections", n(l.get("U503")), provider.get(Slot.UNAVAILABLE_503)); + t.line(" %-22s %,12d %,12d", "billed calls", n(l.get("BILLED")), provider.get(Slot.BILLED)); + t.line(" %-22s %,12d %,12d", "prompt tokens", n(l.get("PT")), provider.get(Slot.PROMPT_TOKENS)); + t.line(" %-22s %,12d %,12d", "completion tokens", n(l.get("CT")), provider.get(Slot.COMPLETION_TOKENS)); + t.blank(); + BigDecimal alt = new Pricing("OTHER-EXAMPLE", new BigDecimal("2.00"), new BigDecimal("6.00")).cost(n(l.get("PT")), n(l.get("CT"))); + t.line("Cost = prompt tokens x input price + completion tokens x output price, per million tokens."); + t.line(" table '%s' (input %s, output %s): %s <- the ledger's total", pricing.version(), pricing.inputPerMillion(), + pricing.outputPerMillion(), l.get("COST")); + t.line(" same tokens under table 'OTHER-EXAMPLE' (input 2.00, output 6.00): %s", alt); + t.line(" Both tables are invented inputs. Tokens are the measurement; money is tokens x a table you supply."); + t.line(" Tokens spent on replies that were then skipped: bad replies billed = %d of %d billed calls.", n(l.get("BAD")), n(l.get("BILLED"))); + } + + assertThat(exec.getStatus()).isEqualTo(BatchStatus.COMPLETED); + assertThat(classified + skipped).isEqualTo(3000); + assertThat(step.getProcessSkipCount()).isEqualTo(skipped); + assertThat(skipped).isGreaterThan(0); + assertThat(n(l.get("WRITTEN"))).isEqualTo(classified); + assertThat(n(l.get("SKIPPED"))).isEqualTo(skipped); + assertThat(n(l.get("ATTEMPTS"))).isEqualTo(provider.get(Slot.RECEIVED)); + assertThat(n(l.get("R429"))).isEqualTo(provider.get(Slot.REJECTED_429)); + assertThat(n(l.get("BILLED"))).isEqualTo(provider.get(Slot.BILLED)); + assertThat(n(l.get("PT"))).isEqualTo(provider.get(Slot.PROMPT_TOKENS)); + assertThat(n(l.get("CT"))).isEqualTo(provider.get(Slot.COMPLETION_TOKENS)); + assertThat((BigDecimal) l.get("COST")).isEqualByComparingTo(pricing.cost(n(l.get("PT")), n(l.get("CT")))); + assertThat(jdbc.queryForObject("SELECT COUNT(*) FROM SKIPPED_TICKET s JOIN TICKET_RESULT r ON r.TICKET_ID = s.TICKET_ID", Long.class)).isZero(); + } + } +} diff --git a/batch-llm/src/test/java/com/ankurm/batchllm/RetryLayersTest.java b/batch-llm/src/test/java/com/ankurm/batchllm/RetryLayersTest.java new file mode 100644 index 0000000..bfba003 --- /dev/null +++ b/batch-llm/src/test/java/com/ankurm/batchllm/RetryLayersTest.java @@ -0,0 +1,127 @@ +package com.ankurm.batchllm; + +import static com.ankurm.batchllm.support.JobSupport.count; +import static org.assertj.core.api.Assertions.assertThat; + +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; + +import org.junit.jupiter.api.Test; +import org.springframework.batch.core.job.Job; +import org.springframework.batch.core.job.JobExecution; +import org.springframework.batch.core.launch.JobOperator; +import org.springframework.batch.infrastructure.item.ItemWriter; +import org.springframework.boot.test.context.TestConfiguration; +import org.springframework.context.ConfigurableApplicationContext; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Primary; +import org.springframework.dao.DataIntegrityViolationException; +import org.springframework.dao.TransientDataAccessResourceException; +import org.springframework.jdbc.core.JdbcTemplate; + +import com.ankurm.batchllm.ProviderMeter.Slot; +import com.ankurm.batchllm.support.JobSupport; +import com.ankurm.batchllm.support.Transcript; + +/** + * Where does each retry/skip mechanism belong? Four experiments, each counting what the provider + * received, because the provider's meter is where the money is. + */ +class RetryLayersTest { + + /** A writer in front of the real one that fails on demand. */ + @TestConfiguration + static class FaultyWriterConfig { + static final AtomicInteger transientFailuresLeft = new AtomicInteger(); + static volatile long poisonTicket = -1; + + @Bean + @Primary + ItemWriter faultyWriter(ClassificationWriter real) { + return chunk -> { + boolean hasFirstOfSecondChunk = chunk.getItems().stream().anyMatch(c -> c.ticketId() == 101); + if (hasFirstOfSecondChunk && transientFailuresLeft.getAndDecrement() > 0) { + throw new TransientDataAccessResourceException("simulated connection blip"); + } + if (poisonTicket > 0 && chunk.getItems().stream().anyMatch(c -> c.ticketId() == poisonTicket)) { + throw new DataIntegrityViolationException("simulated: value too long for column"); + } + real.write(chunk); + }; + } + } + + private record Outcome(String status, long billed, long received, long rollbacks, long commits, long ledgerRows, long classified, long skipped) { + } + + private static Outcome run(String name, Map props) throws Exception { + try (ConfigurableApplicationContext ctx = SpringContexts.start(name, props, FaultyWriterConfig.class)) { + JobExecution exec = JobSupport.run(ctx.getBean(JobOperator.class), ctx.getBean(Job.class), name); + var step = exec.getStepExecutions().iterator().next(); + JdbcTemplate jdbc = ctx.getBean(JdbcTemplate.class); + ProviderMeter m = ctx.getBean(ProviderMeter.class); + return new Outcome(exec.getStatus().toString(), m.get(Slot.BILLED), m.get(Slot.RECEIVED), step.getRollbackCount(), + step.getCommitCount(), count(jdbc, "COST_LEDGER"), count(jdbc, "TICKET_RESULT"), count(jdbc, "SKIPPED_TICKET")); + } + } + + @Test + void whereEachMechanismBelongs() throws Exception { + Map small = Map.of("batch.llm.tickets", "300"); + FaultyWriterConfig.transientFailuresLeft.set(0); + FaultyWriterConfig.poisonTicket = -1; + Outcome baseline = run("rl-base", small); + + // (a) a transient DATABASE failure on one chunk write; Spring Batch retries the write only + FaultyWriterConfig.transientFailuresLeft.set(1); + Outcome writeRetry = run("rl-write-retry", Map.of("batch.llm.tickets", "300", "batch.llm.write-retries", "2")); + FaultyWriterConfig.transientFailuresLeft.set(0); + + // (c) a poison row at write time: the chunk is rolled back and scanned item by item + FaultyWriterConfig.poisonTicket = 250; + Outcome scanReprocess = run("rl-scan-reprocess", Map.of("batch.llm.tickets", "300", "batch.llm.processor-transactional", "true")); + Outcome scanCached = run("rl-scan-cached", Map.of("batch.llm.tickets", "300", "batch.llm.processor-transactional", "false")); + FaultyWriterConfig.poisonTicket = -1; + + // (b) every call is a 429: one ticket, Resilience4j only vs Resilience4j plus Spring Batch retry on the same exception + Map dead = Map.of("batch.llm.tickets", "1", "batch.llm.chunk-size", "1", "sim.provider.random429", "1.0"); + Outcome r4jOnly = run("rl-dead-r4j", dead); + var both = new java.util.HashMap<>(dead); + both.put("batch.llm.batch-retries-calls", "true"); + both.put("batch.llm.write-retries", "2"); + Outcome stacked = run("rl-dead-both", both); + + try (Transcript t = new Transcript("05-retry-layers.txt", "Which layer retries what: counts taken from the provider's meter")) { + t.line("(a) A transient DB failure on the write of chunk 2 (tickets 101-200), 300 tickets, Spring Batch write retry = 2"); + t.line(" %-28s status=%-9s billed calls=%d rollbacks=%d", "no failure", baseline.status(), baseline.billed(), baseline.rollbacks()); + t.line(" %-28s status=%-9s billed calls=%d rollbacks=%d", "failure + Spring Batch retry", writeRetry.status(), writeRetry.billed(), writeRetry.rollbacks()); + t.line(" -> the write was retried inside the same transaction; the model was not asked again."); + t.blank(); + t.line("(b) The provider answers 429 to everything, 1 ticket. Calls the provider received for that ticket:"); + t.line(" %-44s status=%-7s attempts=%d", "Resilience4j retry only (max 4 attempts)", r4jOnly.status(), r4jOnly.received()); + t.line(" %-44s status=%-7s attempts=%d", "+ Spring Batch retry of the same exception (2)", stacked.status(), stacked.received()); + t.line(" -> 4 x (1 + 2) = 12. Two retry layers on the same failure multiply, and each attempt also takes a limiter permit."); + t.blank(); + t.line("(c) A poison row at WRITE time (ticket 250): chunk 3 (201-300) rolls back and is scanned one item at a time."); + t.line(" %-40s status=%-9s billed calls=%d (no-failure baseline: %d) rollbacks=%d ledger rows=%d skipped=%d", + "processorNonTransactional (cached)", scanCached.status(), scanCached.billed(), baseline.billed(), scanCached.rollbacks(), scanCached.ledgerRows(), scanCached.skipped()); + t.line(" %-40s status=%-9s billed calls=%d (extra over baseline: %d) rollbacks=%d ledger rows=%d skipped=%d", + "default (processor runs again per item)", scanReprocess.status(), scanReprocess.billed(), + scanReprocess.billed() - baseline.billed(), scanReprocess.rollbacks(), scanReprocess.ledgerRows(), scanReprocess.skipped()); + t.line(" -> without processorNonTransactional(), every item of the failed chunk is classified a second time and paid for twice."); + } + + assertThat(baseline.status()).isEqualTo("COMPLETED"); + assertThat(writeRetry.status()).isEqualTo("COMPLETED"); + assertThat(writeRetry.billed()).isEqualTo(baseline.billed()); + assertThat(writeRetry.rollbacks()).isZero(); + assertThat(r4jOnly.received()).isEqualTo(4); + assertThat(stacked.received()).isEqualTo(12); + assertThat(r4jOnly.status()).isEqualTo("FAILED"); + assertThat(scanCached.billed()).isEqualTo(baseline.billed()); + assertThat(scanReprocess.billed() - baseline.billed()).isGreaterThan(90); + assertThat(scanCached.skipped()).isEqualTo(3); // tickets 46 and 234 (bad JSON) and 250 (write) + assertThat(scanCached.classified() + scanCached.skipped()).isEqualTo(300); + assertThat(scanReprocess.classified() + scanReprocess.skipped()).isEqualTo(300); + } +} diff --git a/batch-llm/src/test/java/com/ankurm/batchllm/SimulatedProviderTest.java b/batch-llm/src/test/java/com/ankurm/batchllm/SimulatedProviderTest.java new file mode 100644 index 0000000..71b60dc --- /dev/null +++ b/batch-llm/src/test/java/com/ankurm/batchllm/SimulatedProviderTest.java @@ -0,0 +1,88 @@ +package com.ankurm.batchllm; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; +import java.util.TreeMap; + +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.prompt.Prompt; + +import com.ankurm.batchllm.support.Transcript; + +class SimulatedProviderTest { + + private static SimulatedProvider provider(String meterFile, int quota) throws Exception { + Files.createDirectories(Path.of("target/test-meters")); + ProviderMeter meter = new ProviderMeter(meterFile); + meter.reset(); + SimulatedProvider p = new SimulatedProvider(new ProviderFaults(7, quota, 0.015, 0.005, 0.003, 0.012, 0.01, 40, + meterFile), meter); + p.setLatencyDivisor(1000); + return p; + } + + private static String outcome(SimulatedProvider p, Ticket t, TicketClassifier parser) { + try { + String text = p.call(new Prompt("Ticket #" + t.id() + "\n" + t.body())).getResult().getOutput().getText(); + try { + parser.parse(t.id(), text); + return text.startsWith("```") ? "valid JSON (in a markdown fence, stripped)" : "valid JSON"; + } catch (BadModelOutputException e) { + return "bad JSON: " + e.getMessage().replaceAll(":.*", ""); + } + } catch (ProviderRateLimitException e) { + return "429"; + } catch (ProviderUnavailableException e) { + return "503"; + } + } + + @Test + void faultPlanIsDeterministicAndHasTheAdvertisedMix() throws Exception { + TicketClassifier parser = new TicketClassifier(null, io.github.resilience4j.ratelimiter.RateLimiterRegistry.ofDefaults(), + io.github.resilience4j.retry.RetryRegistry.ofDefaults(), new UsageMeter(), null); + int n = 20_000; + List first = new ArrayList<>(); + List second = new ArrayList<>(); + SimulatedProvider a = provider("target/test-meters/sim-a.bin", 1_000_000); + SimulatedProvider b = provider("target/test-meters/sim-b.bin", 1_000_000); + for (long id = 1; id <= n; id++) { + first.add(outcome(a, TicketGenerator.ticket(id), parser)); + second.add(outcome(b, TicketGenerator.ticket(id), parser)); + } + assertThat(second).isEqualTo(first); + + TreeMap mix = new TreeMap<>(); + first.forEach(o -> mix.merge(o, 1, Integer::sum)); + + // quota: 300 calls fired back to back against a quota of 100 per second + SimulatedProvider q = provider("target/test-meters/sim-q.bin", 100); + int rejected = 0; + for (long id = 1; id <= 300; id++) { + try { + q.call(new Prompt("Ticket #" + id)); + } catch (ProviderRateLimitException e) { + rejected++; + } catch (ProviderUnavailableException ignored) { + // an injected 503 is neither accepted nor over quota; ignore + } + } + + try (Transcript t = new Transcript("02-simulated-provider.txt", "The simulated provider: same input, same outcome")) { + t.line("Two separate provider instances, same seed, tickets 1..%,d, one call each: identical outcome lists = %s", + n, second.equals(first)); + t.blank(); + t.line("Outcome of the first call per ticket (the plan: 1.5% 429, 0.5% 503, 0.3% never valid, 1.2% invalid on the first reply only):"); + mix.forEach((k, v) -> t.line(" %-48s %,7d %5.2f%%", k, v, 100.0 * v / n)); + t.blank(); + t.line("Quota check: quota 100 calls/second, 300 calls fired back to back -> %d rejected with 429 (quota rejections plus any injected 429)", rejected); + t.line("Every number above is a property of the TEST PLAN, not a measurement of any provider."); + } + assertThat(mix.get("valid JSON") + mix.getOrDefault("valid JSON (in a markdown fence, stripped)", 0)).isGreaterThan(18_000); + assertThat(rejected).isBetween(195, 205); + } +} diff --git a/batch-llm/src/test/java/com/ankurm/batchllm/SkipLimitAndRestartTest.java b/batch-llm/src/test/java/com/ankurm/batchllm/SkipLimitAndRestartTest.java new file mode 100644 index 0000000..c8d7614 --- /dev/null +++ b/batch-llm/src/test/java/com/ankurm/batchllm/SkipLimitAndRestartTest.java @@ -0,0 +1,97 @@ +package com.ankurm.batchllm; + +import static com.ankurm.batchllm.support.JobSupport.count; +import static org.assertj.core.api.Assertions.assertThat; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.Map; + +import org.junit.jupiter.api.Test; +import org.springframework.batch.core.job.Job; +import org.springframework.batch.core.job.JobExecution; +import org.springframework.batch.core.launch.JobOperator; +import org.springframework.boot.WebApplicationType; +import org.springframework.boot.builder.SpringApplicationBuilder; +import org.springframework.context.ConfigurableApplicationContext; +import org.springframework.jdbc.core.JdbcTemplate; + +import com.ankurm.batchllm.ProviderMeter.Slot; +import com.ankurm.batchllm.support.JobSupport; +import com.ankurm.batchllm.support.Transcript; + +/** A model regression (half the replies are not JSON) must stop the job, and a fixed model must resume it. */ +class SkipLimitAndRestartTest { + + private static ConfigurableApplicationContext start(Map extra) { + Map p = new LinkedHashMap<>(); + p.put("batch.llm.autorun", "false"); + p.put("batch.llm.tickets", "2000"); + p.put("spring.datasource.url", "jdbc:h2:file:./target/test-db/skiplimit;DB_CLOSE_ON_EXIT=FALSE"); + p.put("sim.provider.meter-file", "target/test-meters/skiplimit.bin"); + p.putAll(extra); + return new SpringApplicationBuilder(BatchLlmApplication.class).web(WebApplicationType.NONE).run(p.entrySet().stream().map(e -> "--" + e.getKey() + "=" + e.getValue()).toArray(String[]::new)); + } + + @Test + void skipLimitStopsASpendingSpreeAndTheJobResumesWhenTheModelIsFixed() throws Exception { + Path db = Path.of("target/test-db"); + if (Files.exists(db)) { + try (var s = Files.walk(db)) { + s.sorted(Comparator.reverseOrder()).forEach(p -> p.toFile().delete()); + } + } + Files.deleteIfExists(Path.of("target/test-meters/skiplimit.bin")); + + String s1; + long billed1; + long read1; + long classified1; + long skippedRows1; + long calls1; + try (var ctx = start(Map.of("sim.provider.permanent-bad", "0.5", "sim.provider.flaky-bad", "0", "batch.llm.reask", "0", + "batch.llm.skip-limit", "50"))) { + JobExecution e1 = JobSupport.run(ctx.getBean(JobOperator.class), ctx.getBean(Job.class), "regression"); + var step = e1.getStepExecutions().iterator().next(); + JdbcTemplate jdbc = ctx.getBean(JdbcTemplate.class); + s1 = e1.getStatus() + " / " + step.getExitStatus().getExitDescription().lines().findFirst().orElse(""); + billed1 = ctx.getBean(ProviderMeter.class).get(Slot.BILLED); + calls1 = ctx.getBean(ProviderMeter.class).get(Slot.RECEIVED); + read1 = step.getReadCount(); + classified1 = count(jdbc, "TICKET_RESULT"); + skippedRows1 = count(jdbc, "SKIPPED_TICKET"); + assertThat(e1.getStatus().toString()).isEqualTo("FAILED"); + } + + try (var ctx = start(Map.of("batch.llm.skip-limit", "1000"))) { + JdbcTemplate jdbc = ctx.getBean(JdbcTemplate.class); + var operator = ctx.getBean(JobOperator.class); + JobExecution e2 = JobSupport.run(operator, ctx.getBean(Job.class), "regression"); + var step = e2.getStepExecutions().iterator().next(); + long classified = count(jdbc, "TICKET_RESULT"); + long skipped = count(jdbc, "SKIPPED_TICKET"); + long both = jdbc.queryForObject("SELECT COUNT(*) FROM SKIPPED_TICKET s JOIN TICKET_RESULT r ON r.TICKET_ID = s.TICKET_ID", Long.class); + try (Transcript t = new Transcript("06-skip-limit-and-restart.txt", "A model regression, the skip limit, and a restart with a fixed model")) { + t.line("Execution 1: simulated model returns invalid JSON for 50% of tickets, no re-ask, skip limit 50, 2,000 tickets."); + t.line(" status / reason %s", s1); + t.line(" tickets read %d of 2000", read1); + t.line(" calls the provider billed before the step stopped: %d (a full run needs about 2,000)", billed1); + t.line(" committed results %d skipped rows %d", classified1, skippedRows1); + t.blank(); + t.line("Execution 2: same job parameters, same database, model fixed, skip limit 1000."); + t.line(" status %s", e2.getStatus()); + t.line(" tickets read this time %d (execution 1 had committed %d, so the reader resumed after them)", step.getReadCount(), 2000 - step.getReadCount()); + t.line(" committed results %d skipped rows %d classified + skipped = %d", classified, skipped, classified + skipped); + t.line(" tickets present in BOTH tables: %d (stale skip rows from execution 1 are deleted when the ticket is classified)", both); + } + assertThat(e2.getStatus().toString()).isEqualTo("COMPLETED"); + assertThat(classified + skipped).isEqualTo(2000); + assertThat(both).isZero(); + assertThat(billed1).isLessThan(1000); + assertThat(step.getReadCount()).isLessThan(2000); + } + assertThat(s1).contains("FAILED"); + } +} diff --git a/batch-llm/src/test/java/com/ankurm/batchllm/SpringContexts.java b/batch-llm/src/test/java/com/ankurm/batchllm/SpringContexts.java new file mode 100644 index 0000000..60c72f8 --- /dev/null +++ b/batch-llm/src/test/java/com/ankurm/batchllm/SpringContexts.java @@ -0,0 +1,28 @@ +package com.ankurm.batchllm; + +import java.util.Map; + +import org.springframework.boot.WebApplicationType; +import org.springframework.boot.builder.SpringApplicationBuilder; +import org.springframework.context.ConfigurableApplicationContext; + +/** Starts the real application in a fresh context with the given properties and extra configuration classes. */ +final class SpringContexts { + + private SpringContexts() { + } + + static ConfigurableApplicationContext start(String name, Map props, Class... extra) { + var all = new java.util.LinkedHashMap(); + all.put("batch.llm.autorun", "false"); + all.put("spring.datasource.url", "jdbc:h2:mem:" + name + ";DB_CLOSE_DELAY=-1"); + all.put("sim.provider.meter-file", "target/test-meters/" + name + ".bin"); + all.putAll(props); + Class[] sources = new Class[extra.length + 1]; + sources[0] = BatchLlmApplication.class; + System.arraycopy(extra, 0, sources, 1, extra.length); + var ctx = new SpringApplicationBuilder(sources).web(WebApplicationType.NONE).run(all.entrySet().stream().map(e -> "--" + e.getKey() + "=" + e.getValue()).toArray(String[]::new)); + ctx.getBean(ProviderMeter.class).reset(); + return ctx; + } +} diff --git a/batch-llm/src/test/java/com/ankurm/batchllm/ThrottleTest.java b/batch-llm/src/test/java/com/ankurm/batchllm/ThrottleTest.java new file mode 100644 index 0000000..0815666 --- /dev/null +++ b/batch-llm/src/test/java/com/ankurm/batchllm/ThrottleTest.java @@ -0,0 +1,106 @@ +package com.ankurm.batchllm; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.Semaphore; +import java.util.concurrent.atomic.AtomicInteger; + +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.client.ChatClient; + +import com.ankurm.batchllm.ProviderMeter.Slot; +import com.ankurm.batchllm.support.Transcript; + +import io.github.resilience4j.core.IntervalFunction; +import io.github.resilience4j.ratelimiter.RateLimiterConfig; +import io.github.resilience4j.ratelimiter.RateLimiterRegistry; +import io.github.resilience4j.retry.RetryConfig; +import io.github.resilience4j.retry.RetryRegistry; + +/** The same 4,000 tickets against a simulated quota of 800 calls/second, with and without a client-side limiter. */ +class ThrottleTest { + + record Result(String name, int ok, int badJson, int exhausted, long received, long window429, long retries, double seconds) { + } + + private static Result run(String name, int permitsPer100ms, int quota, int tickets) throws Exception { + Files.createDirectories(Path.of("target/test-meters")); + String file = "target/test-meters/throttle-" + name.hashCode() + ".bin"; + ProviderMeter meter = new ProviderMeter(file); + meter.reset(); + var faults = new ProviderFaults(7, quota, 0.015, 0.005, 0.003, 0.012, 0.01, 40, file); + var sim = new SimulatedProvider(faults, meter); + var limiters = RateLimiterRegistry.of(RateLimiterConfig.custom().limitForPeriod(permitsPer100ms) + .limitRefreshPeriod(Duration.ofMillis(100)).timeoutDuration(Duration.ofSeconds(120)).build()); + var retries = RetryRegistry.of(RetryConfig.custom().maxAttempts(4) + .intervalFunction(IntervalFunction.ofExponentialBackoff(25, 2)) + .retryExceptions(ProviderRateLimitException.class, ProviderUnavailableException.class).build()); + var usage = new UsageMeter(); + var props = new LlmJobProperties(tickets, 100, 64, 1, 1000, 0, false, false, "x"); + var classifier = new TicketClassifier(ChatClient.builder(sim).defaultSystem(TicketClassifier.SYSTEM).build(), + limiters, retries, usage, props); + + AtomicInteger ok = new AtomicInteger(); + AtomicInteger bad = new AtomicInteger(); + AtomicInteger exhausted = new AtomicInteger(); + Semaphore inFlight = new Semaphore(64); + long t0 = System.nanoTime(); + try (var pool = Executors.newVirtualThreadPerTaskExecutor()) { + List> futures = new ArrayList<>(); + for (long id = 1; id <= tickets; id++) { + Ticket ticket = TicketGenerator.ticket(id); + inFlight.acquire(); + futures.add(pool.submit(() -> { + try { + classifier.process(ticket); + ok.incrementAndGet(); + } catch (BadModelOutputException e) { + bad.incrementAndGet(); + } catch (ProviderRateLimitException | ProviderUnavailableException e) { + exhausted.incrementAndGet(); + } finally { + inFlight.release(); + } + })); + } + for (Future f : futures) { + f.get(); + } + } + double seconds = (System.nanoTime() - t0) / 1e9; + return new Result(name, ok.get(), bad.get(), exhausted.get(), meter.get(Slot.RECEIVED), meter.get(Slot.WINDOW_429), + usage.retries.get(), seconds); + } + + @Test + void limiterKeepsTheClientUnderTheProvidersQuota() throws Exception { + int tickets = 4000; + Result unlimited = run("no limiter", 1_000_000, 800, tickets); + Result limited = run("limiter 600/s", 60, 800, tickets); + try (Transcript t = new Transcript("03-throttle.txt", "4,000 tickets against a simulated quota of 800 calls per second")) { + t.line("Both runs: 64 virtual threads, Resilience4j retry max 4 attempts, 25 ms x2 backoff, the same fault plan."); + t.line("Time is indicative only (shared 2-vCPU sandbox, other processes running)."); + t.blank(); + t.line("%-14s %6s %8s %10s %9s %11s %8s %8s", "scenario", "ok", "bad JSON", "exhausted", "calls", "quota 429s", "retries", "seconds"); + for (Result r : List.of(unlimited, limited)) { + t.line("%-14s %6d %8d %10d %,9d %,11d %,8d %8.1f", r.name(), r.ok(), r.badJson(), r.exhausted(), r.received(), + r.window429(), r.retries(), r.seconds()); + } + t.blank(); + t.line("'exhausted' = tickets whose 4 attempts all met a 429/503: with a step around this call, the step fails."); + t.line("'calls' counts every attempt the provider received, rejected or not."); + } + assertThat(unlimited.window429()).isGreaterThan(500); + assertThat(unlimited.exhausted()).isGreaterThan(100); + assertThat(limited.window429()).isZero(); + assertThat(limited.exhausted()).isZero(); + assertThat(limited.ok() + limited.badJson()).isEqualTo(tickets); + } +} diff --git a/batch-llm/src/test/java/com/ankurm/batchllm/VersionsTest.java b/batch-llm/src/test/java/com/ankurm/batchllm/VersionsTest.java new file mode 100644 index 0000000..e1b58f0 --- /dev/null +++ b/batch-llm/src/test/java/com/ankurm/batchllm/VersionsTest.java @@ -0,0 +1,71 @@ +package com.ankurm.batchllm; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.lang.reflect.Method; +import java.util.Arrays; + +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.client.ChatClient; +import org.springframework.batch.core.launch.JobOperator; +import org.springframework.batch.core.step.builder.ChunkOrientedStepBuilder; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.SpringBootVersion; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.core.SpringVersion; + +import com.ankurm.batchllm.support.Transcript; + +import io.github.resilience4j.ratelimiter.RateLimiter; +import io.github.resilience4j.ratelimiter.RateLimiterRegistry; +import io.github.resilience4j.retry.Retry; +import io.github.resilience4j.retry.RetryRegistry; + +@SpringBootTest(properties = {"batch.llm.autorun=false", "spring.datasource.url=jdbc:h2:mem:versions;DB_CLOSE_DELAY=-1", + "sim.provider.meter-file=target/test-meters/versions.bin"}) +class VersionsTest { + + @Autowired RateLimiterRegistry limiters; + @Autowired RetryRegistry retries; + + static String version(Class c) { + String v = c.getPackage().getImplementationVersion(); + if (v != null) { + return v; + } + String path = c.getProtectionDomain().getCodeSource().getLocation().getPath(); + return path.substring(path.lastIndexOf('/') + 1); + } + + @Test + void versionsAndApiShape() { + RateLimiter limiter = limiters.rateLimiter("llm"); + Retry retry = retries.retry("llm"); + try (Transcript t = new Transcript("01-versions.txt", "What is on the classpath and what the builder offers")) { + t.line("Spring Boot %s", SpringBootVersion.getVersion()); + t.line("Spring Framework %s", SpringVersion.getVersion()); + t.line("Spring Batch (core) %s", version(JobOperator.class)); + t.line("Spring AI (client) %s", version(ChatClient.class)); + t.line("Resilience4j %s", version(RateLimiter.class)); + t.line("Java %s", System.getProperty("java.version")); + t.blank(); + t.line("Resilience4j config read from application.yml by resilience4j-spring-boot4:"); + t.line(" ratelimiter 'llm': limitForPeriod=%d refreshPeriod=%s timeout=%s", + limiter.getRateLimiterConfig().getLimitForPeriod(), + limiter.getRateLimiterConfig().getLimitRefreshPeriod(), + limiter.getRateLimiterConfig().getTimeoutDuration()); + t.line(" retry 'llm': maxAttempts=%d", retry.getRetryConfig().getMaxAttempts()); + t.blank(); + t.line("ChunkOrientedStepBuilder methods that matter here (javap of spring-batch-core):"); + for (String name : new String[] {"faultTolerant", "taskExecutor", "retryPolicy", "retry", "retryLimit", "skip", + "skipLimit", "skipListener", "processorNonTransactional"}) { + Method m = Arrays.stream(ChunkOrientedStepBuilder.class.getMethods()) + .filter(x -> x.getName().equals(name)).findFirst().orElseThrow(); + t.line(" %-27s %s", name, Arrays.toString(m.getParameterTypes()).replace("class ", "").replace("interface ", "")); + } + assertThat(SpringBootVersion.getVersion()).isEqualTo("4.1.2"); + assertThat(limiter.getRateLimiterConfig().getLimitForPeriod()).isEqualTo(200); + assertThat(retry.getRetryConfig().getMaxAttempts()).isEqualTo(4); + } + } +} diff --git a/batch-llm/src/test/java/com/ankurm/batchllm/support/JobSupport.java b/batch-llm/src/test/java/com/ankurm/batchllm/support/JobSupport.java new file mode 100644 index 0000000..394f062 --- /dev/null +++ b/batch-llm/src/test/java/com/ankurm/batchllm/support/JobSupport.java @@ -0,0 +1,36 @@ +package com.ankurm.batchllm.support; + +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.JobParametersBuilder; +import org.springframework.batch.core.launch.JobOperator; +import org.springframework.jdbc.core.JdbcTemplate; + +public final class JobSupport { + + private JobSupport() { + } + + public static JobExecution run(JobOperator operator, Job job, String runId) throws Exception { + return operator.start(job, new JobParametersBuilder().addString("run.id", runId).toJobParameters()); + } + + public static Map ledger(JdbcTemplate jdbc) { + return 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"""); + } + + public static long n(Object o) { + return ((Number) o).longValue(); + } + + public static long count(JdbcTemplate jdbc, String table) { + return jdbc.queryForObject("SELECT COUNT(*) FROM " + table, Long.class); + } +} diff --git a/batch-llm/src/test/java/com/ankurm/batchllm/support/Transcript.java b/batch-llm/src/test/java/com/ankurm/batchllm/support/Transcript.java new file mode 100644 index 0000000..81b4918 --- /dev/null +++ b/batch-llm/src/test/java/com/ankurm/batchllm/support/Transcript.java @@ -0,0 +1,46 @@ +package com.ankurm.batchllm.support; + +import java.io.IOException; +import java.io.PrintWriter; +import java.io.StringWriter; +import java.nio.file.Files; +import java.nio.file.Path; + +/** + * Writes a numbered transcript under {@code output/} (module root, not {@code docs/}) and echoes + * it to the console. Every console block quoted in the article comes out of one of these files. + */ +public final class Transcript implements AutoCloseable { + + private final Path path; + private final StringWriter buffer = new StringWriter(); + private final PrintWriter out = new PrintWriter(buffer); + + public Transcript(String fileName, String title) { + this.path = Path.of("output", fileName); + out.println("# " + title); + out.println(); + } + + public Transcript line(String format, Object... args) { + out.println(args.length == 0 ? format : String.format(format, args)); + return this; + } + + public Transcript blank() { + out.println(); + return this; + } + + @Override + public void close() { + out.flush(); + try { + Files.createDirectories(path.getParent()); + Files.writeString(path, buffer.toString()); + } catch (IOException e) { + throw new IllegalStateException("could not write " + path, e); + } + System.out.print(buffer); + } +}