Add image generation via Draw Things gRPCServerCLI#9
Merged
Conversation
…patch Implements the full request pipeline for image generation (Phases 1, 3, 4): Protocol (Go + Rust): - ImageGenerationRequest/Complete message types with E2E encryption support - ImageGenerationRequestBody (model, prompt, negative_prompt, n, size, steps, seed) - ImageGenerationCompleteMessage with base64 images and per-image usage billing - Round-trip serialization tests for all new message types Coordinator: - POST /v1/images/generations endpoint (OpenAI-compatible) - E2E encrypted request forwarding to providers via WebSocket - handleImageGenerationComplete handler with job tracking - ImageGenerationCh channel on PendingRequest for non-streaming response Provider: - CoordinatorEvent::ImageGenerationRequest in event loop - Decrypt + parse + dispatch in coordinator.rs - handle_image_generation_request() in proxy.rs (POST to bridge on port be_port+2) - 5-minute timeout for image generation requests The provider proxy forwards requests to a local image bridge server (dginf_image_bridge) on port 8102, which will wrap Draw Things gRPCServerCLI for Metal FlashAttention-accelerated image generation. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Phase 2: Python image bridge (dginf_image_bridge) - FastAPI server exposing OpenAI-compatible /v1/images/generations - Pluggable backend: MfluxBackend (MLX-native FLUX) with Draw Things upgrade path - /health endpoint for provider health monitoring - Configurable model, port, quantization via CLI args - Entry point: python -m dginf_image_bridge --port 8102 --model flux-schnell Tests (22 Python + 3 Rust): - 18 unit tests: health, generation, validation, error handling, response format - 4 integration tests: real HTTP server with mock backend - 3 Rust proxy tests: mock bridge server, error response, cancellation Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…tier Phase 5 + 6: Billing and model catalog Billing: - CalculateImageCost() in payments/pricing.go: per-image pricing scaled by resolution relative to 1024x1024 base ($0.002/image) - handleImageGenerationComplete now charges consumer, credits provider, records usage, and distributes platform fee + referral rewards - Same 95/5 provider/platform split as text inference Install script: - 8-15 GB machines now get FLUX.2 Klein 4B (image generation) instead of falling through to the tiny 0.5B chat model - <8 GB machines get Cohere Transcribe (speech-to-text) - MODEL_TYPE variable tracks text/image/transcription for backend selection Tests: - TestCalculateImageCost: base pricing, multi-image, resolution scaling, minimum charge, small/large images Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
The mflux API changed in v0.17 — Flux1 moved to mflux.models.common.cli.save and generate_image returns a GeneratedImage with a .image (PIL) attribute. Tested end-to-end: downloaded FLUX.1-schnell, started bridge server, generated real images via HTTP API. 256x256 in 2.18s on Apple Silicon. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…uite DrawThingsBackend connects to gRPCServerCLI via gRPC and translates OpenAI-format requests into Draw Things' FlatBuffers + protobuf format: - Proto: vendored imageService.proto from kcjerrell/dt-grpc-ts - gRPC stubs: vendored from drawthingsai/draw-things-comfyui - FlatBuffers: vendored config_generated.py (86-field GenerationConfigurationT) - build_config_bytes(): constructs serialized FlatBuffers config for text-to-image - convert_response_image(): decodes raw float16 tensor responses to PIL Images - DrawThingsBackend: manages gRPC channel, builds requests, streams responses - CLI: --backend drawthings (default) or --backend mflux (fallback) Tests (15 new, 37 total): - FlatBuffers config: construction, serialization, deserialization roundtrip - Image decoding: RGB, RGBA, pixel value mapping, edge cases - Mock gRPC server: single/batch generation, health check, wrong port - Full stack: HTTP POST → FastAPI → DrawThingsBackend → mock gRPC → PNG Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Simplified to a single backend path: - Removed MfluxBackend class and all mflux imports - Removed --backend CLI flag (Draw Things is the only option) - gRPCServerCLI spawned as subprocess via --model-path flag - atexit handler ensures clean shutdown - Integration tests now use mock gRPC server (not mock ImageBackend) - requirements.txt: removed mflux, kept grpcio + flatbuffers Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…all, validation P1 fixes: - Raise WebSocket read limit to 10 MB (from 32 KB default) so base64 image responses don't disconnect the provider with StatusMessageTooBig - Start image bridge subprocess in cmd_serve() (DGINF_IMAGE_MODEL env var) following the STT backend pattern — TCP poll until ready, advertise model - Install script sets DGINF_IMAGE_MODEL / DGINF_STT_MODEL env vars so the provider starts the correct backend for image/transcription model types P2 fixes: - Validate n (1-4) and steps (> 0) at the coordinator API edge to prevent silent timeouts when negative values fail Rust u32 deserialization - Pass --port to gRPCServerCLI subprocess so non-default grpc_port works Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…afety Redesigned image transport — images uploaded via HTTP, not WebSocket: - Provider POSTs generated PNG bytes to coordinator's new POST /v1/provider/image-upload endpoint (no size limits) - WebSocket image_generation_complete message now carries only usage metadata (tiny JSON), eliminating the 10MB WebSocket limit problem - Coordinator stores uploaded images in memory, serves to consumer when completion message arrives - upload_url included in image_generation_request so provider knows where to POST Dimension validation: - Coordinator rejects sizes > 2048x2048 at the API edge - Validates n in [1,4] and steps > 0 (prevents silent Rust u32 parse failures) Deploy safety: - Install script reverted to safe model selection (no auto-selecting image models until gRPCServerCLI provisioning is in the installer) - Image generation is opt-in via DGINF_IMAGE_MODEL env var - Provider sets PYTHONPATH when spawning image bridge so the package is importable from bundled locations Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Gajesh2007
force-pushed
the
worktree-image-generation
branch
from
April 1, 2026 21:04
9aed5d3 to
d7ac445
Compare
- Klein 4B (8.2 GB, MinRAM 8 GB): fast image gen for smallest machines - Klein 9B (13.4 GB, MinRAM 16 GB): higher quality for 16 GB+ machines - Both use Draw Things gRPCServerCLI with Metal FlashAttention - Model type "image" enables catalog filtering via ?type=image Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Gajesh2007
force-pushed
the
worktree-image-generation
branch
from
April 1, 2026 21:41
1b56b61 to
994c2ee
Compare
Gajesh2007
added a commit
that referenced
this pull request
Jun 17, 2026
The #1 rotation cleanup cleared the reuse record and outstanding nonce but left lastPush keyed by SE key. Since the rearm loop sets CodeAttested=false and the old token's background cooldown can be up to 20 min, the forced re-challenge could not push to the NEW token for that whole window — derouting the provider even though the new token has its own (untouched) Apple push budget. Add codeAttestThrottle.clearPushBudget(seKey) and call it in maybeRearmCodeAttest alongside invalidateReuse/clearChallenge. Regression test TestCodeAttestThrottleClearPushBudget.
Gajesh2007
added a commit
that referenced
this pull request
Jun 17, 2026
…— threat-model The #9 fix cleared the per-device push cooldown on every APNs token change, and maybeRearmCodeAttest runs on every heartbeat. A provider can put any string in the heartbeat APNs-token field, so it could 'rotate' its token each heartbeat to reset the budget every time and force one APNs push per heartbeat — far past Apple's ~3/hr/device budget (flagged by the threat-model review). Bound clearPushBudget with a per-device meta-cooldown (budgetClearCooldown, 20m): a budget reset is honored at most ~3x/hour/device. A genuine token rotation is rare and still clears promptly; a flood falls back to the normal cooldown so no APNs spam. clearPushBudget now returns whether it cleared. Regression assertions added to TestCodeAttestThrottleClearPushBudget. (Verified not-a-bug, no change: code-attest nonces are single-use — cleared on verify via clearChallengeIf — and the SE signature is verified against the registration-bound key, never a key from the response.) Full coordinator go test ./... green.
This was referenced Jun 26, 2026
Gajesh2007
added a commit
that referenced
this pull request
Jul 2, 2026
…probs buffer, sanity guards Security/hardening findings from PR #499 review (ethenotethan): #4 request-id validation: a caller-supplied request-id is validated (non-empty, ≤256 chars, no ASCII control chars) before it becomes a dictionary key / cancel-correlation handle; nil/invalid ids normalize to a fresh req-<uuid>. Static normalizedRequestId/isValidRequestId are pure and unit-tested. This also subsumes #11 (duplicate-id): prod ids are provider-minted UUIDs, and any malformed id is now normalized to a fresh UUID before the existing duplicate guard, so an attacker cannot induce a collision — no separate change needed. #5 nextRawId overflow: += → &+= (wrapping) with a comment; 2^64 is unreachable in practice, this makes the theoretical overflow explicitly defined rather than a trap. #6 engine.submit on the actor: verified O(1) non-blocking enqueue against the CBv2Engine contract (admission check + register + enqueue, no forward pass; tokenization already off-actor) and documented it at the call site. No code change. #7 pump task lifecycle: per-request pump Tasks are now tracked in pumpTasks and cancelled by shutdown() (before the engine drain), so none outlives the bridge; AsyncStream cancellation unwinds each pump through its teardown path, releasing per-request KV. Tested: shutdown cancels a live pump and its reservation returns to 0. #8 bounded logprobs buffer: EngineV2LogprobsChannel caps at 4096 entries with drop-oldest eviction + a droppedCount and a one-line WARN on first drop, so a stalled/never-draining consumer can't grow it unboundedly. Dropping logprobs never affects content/billing/hash (out-of-band). #9 dropped logit_bias keys: parseLogitBiasCountingDropped returns the count of dropped (non-numeric/negative) keys; samplingParams emits a count-only WARN (never the keys — request content) instead of dropping silently. #10 kvBytesCapacity sanity: clampKVBytesCapacity caps the v2 admission ceiling at physical RAM so a mis-derived budget degrades to all-of-memory rather than an absurd ceiling. Tests: id validation matrix + normalization-in-submit, wrapping id, shutdown-cancels-live-pumps, logprobs cap drop-oldest + droppedCount, dropped-key count, and the capacity clamp. Full suite green (1245).
Gajesh2007
added a commit
that referenced
this pull request
Jul 2, 2026
…probs buffer, sanity guards Security/hardening findings from PR #499 review (ethenotethan): #4 request-id validation: a caller-supplied request-id is validated (non-empty, ≤256 chars, no ASCII control chars) before it becomes a dictionary key / cancel-correlation handle; nil/invalid ids normalize to a fresh req-<uuid>. Static normalizedRequestId/isValidRequestId are pure and unit-tested. This also subsumes #11 (duplicate-id): prod ids are provider-minted UUIDs, and any malformed id is now normalized to a fresh UUID before the existing duplicate guard, so an attacker cannot induce a collision — no separate change needed. #5 nextRawId overflow: += → &+= (wrapping) with a comment; 2^64 is unreachable in practice, this makes the theoretical overflow explicitly defined rather than a trap. #6 engine.submit on the actor: verified O(1) non-blocking enqueue against the CBv2Engine contract (admission check + register + enqueue, no forward pass; tokenization already off-actor) and documented it at the call site. No code change. #7 pump task lifecycle: per-request pump Tasks are now tracked in pumpTasks and cancelled by shutdown() (before the engine drain), so none outlives the bridge; AsyncStream cancellation unwinds each pump through its teardown path, releasing per-request KV. Tested: shutdown cancels a live pump and its reservation returns to 0. #8 bounded logprobs buffer: EngineV2LogprobsChannel caps at 4096 entries with drop-oldest eviction + a droppedCount and a one-line WARN on first drop, so a stalled/never-draining consumer can't grow it unboundedly. Dropping logprobs never affects content/billing/hash (out-of-band). #9 dropped logit_bias keys: parseLogitBiasCountingDropped returns the count of dropped (non-numeric/negative) keys; samplingParams emits a count-only WARN (never the keys — request content) instead of dropping silently. #10 kvBytesCapacity sanity: clampKVBytesCapacity caps the v2 admission ceiling at physical RAM so a mis-derived budget degrades to all-of-memory rather than an absurd ceiling. Tests: id validation matrix + normalization-in-submit, wrapping id, shutdown-cancels-live-pumps, logprobs cap drop-oldest + droppedCount, dropped-key count, and the capacity clamp. Full suite green (1245).
Gajesh2007
added a commit
that referenced
this pull request
Jul 3, 2026
…ma 4 + GPT-OSS 20B) (#499) * feat(engine-v2): provider bridge for ContinuousBatchingV2 behind DARKBLOOM_ENGINE_V2 - EngineV2Config: env DARKBLOOM_ENGINE_V2 / [backend] engine_v2 key selection, gemma-4*/gpt-oss* per-model allowlist (DARKBLOOM_ENGINE_V2_MODELS override), EngineV2Factory with safe legacy fallback + WARN engine_health telemetry (operation=engine_v2_fallback) on v2 init failure - EngineV2Bridge: adapts CBv2Engine to the existing GenerationEvent surface; event framing mirrors BatchScheduler+EngineBridge (chunk/info/error, abort usage-before-error, teardown sentinel); billing-zero max() defense and the legacy decode-TPS methodology/EWMA - EngineV2Bridge+Translation: pure ChatRequest -> CBv2Request translation (temperature ?? 0.0 legacy default, logit_bias string-key id parsing, logprobs/top_logprobs mapping, seed passthrough, buildStopTokenIds-semantics stop resolution, max_tokens defaulting); CBv2KVError.capacityExhausted -> canonical token_budget_exhausted: string (429/503 classification unchanged) - EngineV2Bridge+Capacity: CBv2CapacitySnapshot -> existing BackendSlotCapacity fields (no protocol change), truthful bytes-derived token budgets, WedgeMonitor reuse + engine_v2 step_wedge transition telemetry - EngineV2Runtime: process-wide bridge registry; ProviderLoop capacity heartbeat folds v2 slots, handleCancellation fans request-id out to CBv2Engine.cancel; both no-ops when the flag is off - ChatCompletionRequest: additive optional logit_bias/logprobs/top_logprobs (v2 path only; legacy ignores them) - docs/engine-v2/CONTRACT-ISSUES-H-provider.md: contract gaps + chosen shapes - submodule: pin libs/mlx-swift-lm at the CBv2 contracts commit (fe60e17) * test(engine-v2): bridge unit suite — scripted CBv2Engine stub, no models/network 34 tests in 7 suites (swift test --filter EngineV2): - translation: field-by-field CBv2Request mapping, legacy greedy defaults, logit_bias string-key parsing, logprobs/top_logprobs mapping, buildStopTokenIds-semantics stop resolution + construction-time stamping - event framing: fixture-compare against the recorded legacy BatchScheduler+EngineBridge stream shape (chunks -> single info -> finish; abort usage-before-error; teardown sentinel; billing-zero max() defense; tokenize failure short-circuits before the engine) - error mapping: capacityExhausted -> token_budget_exhausted (retryable .tokenBudgetExhausted class), backendIneligible/unknown -> generationFailed - cancellation: provider id -> minted CBv2RequestID, unknown-id no-op, EngineV2Runtime fan-out (ProviderLoop hook path) - capacity: snapshot -> BackendSlotCapacity (bytes-derived budgets, idle/ running states, wedge counters), runtime aggregate summary - config gating: env/config flag matrix, allowlist + env override, factory legacy selection without invoking the builder, init-failure fallback with engine_v2_fallback telemetry, engine_v2 TOML key decode * fix(engine-v2): duplicate request-id guard + deterministic fan-out test - EngineV2Bridge.submitTokenized: reject a second submit under an active request id with the legacy planner's canonical message ('token_budget_exhausted: duplicate request ID' -> .requestRejected, 400) before minting an engine id, so two pumps can never share bookkeeping - test: duplicate-id rejection classification + engine isolation - test: runtimeFanOut holds the submitted stream alive for the whole test; dropping it fired onTermination(.cancelled) and raced a second (idempotent) engine cancel into the count assertions * feat(engine-v2): wire ContinuousBatchingV2 into production serving paths Closes the P1 integration gap: the EngineV2 bridge (DARKBLOOM_ENGINE_V2 / engine_v2) was implemented and tested but never instantiated outside tests — model loading always built only the legacy BatchScheduler, both request paths captured slot.scheduler directly, and the EngineV2Runtime registry stayed permanently empty. - Model load (ProviderLoop+ModelLoading -> new ProviderLoop+EngineV2): when EngineV2Config selects v2 for a (non-VLM) model, assemble the real CBv2 engine over the already-loaded container via the new EngineV2Factory.makeProductionEngine (Gemma4/GPT-OSS cbv2LayerKinds + newCacheV2, contiguous KV backend sized from the unified-memory KV budget, text detokenizer over the model tokenizer), register the bridge with EngineV2Runtime, and store it on ModelSlot alongside (never replacing) the legacy scheduler. - Routing: MultiModelBatchSchedulerEngine registry entries and AcquiredModel carry the optional bridge; streamChatCompletion submits the same tokenized prompt through it when present (translated sampling/ stop params — v2 additionally honors penalties + stop strings the legacy submit drops) and cancels through the owning engine. The coordinator inference handler and the unified local endpoint both pass slot.engineV2; standalone --local mode intentionally stays legacy. - Zero-overhead flag-off: capacity heartbeats and cancellation fan-out consult the runtime only when at least one v2 slot exists (previously an unconditional empty-registry actor hop). A v2-served slot reports capacity through its bridge ONLY (the dormant legacy scheduler is skipped so the model is never advertised twice). unloadModel unregisters and drains the bridge. - Failure semantics now exercised in production: v2 init throw -> WARN engine_health (engine_v2_fallback) + legacy fallback; VLM slots skip selection silently (no per-load WARN for permanently-unsupported shapes). - Pin libs/mlx-swift-lm to b2b6d53 (integrated CBv2 engine) — the public EngineV2 assembly the production factory constructs did not exist at the previously pinned SHA. Tests (EngineV2ProductionWiringTests, live-isolated, no weights/network): slot-factory flag-off/allowlist/VLM-gate/init-failure + runtime registration, both request-routing shapes through a scripted CBv2Engine, cancellation propagation to the engine-minted id, runtime-guard zero-consult, heartbeat fold-in with the no-double-report regression, unload retirement, and the production factory throw paths. 50 EngineV2 tests green; full suite 1214 tests / 96 suites green. * chore(submodule): bump mlx-swift-lm to engine-v2 final tip 8662159 Pulls the CBv2 contract updates the provider bridge adopts next: - CBv2Engine is now Sendable (contract-issue H§1 fix) - CBv2CapacitySnapshot.stepsExecuted monotonic step counter (H§3 fix) - CBv2Event.delta logprobs populated when sampling.topLogprobs > 0 (E§1) - CBv2Request.cacheSalt per-request prefix-cache scope (TB-007) - prefix-cache donation skip for quantized KV; shutdown timeout hardening * feat(engine-v2): adopt final CBv2 contract — drop engine box, real step counter, logprobs + cacheSalt plumbing Follow-ups on the engine-v2 bridge against mlx-swift-lm tip 8662159: - Drop EngineV2EngineBox: CBv2Engine is Sendable in the contract now, so the bridge actor holds and calls the engine directly (H§1 resolved). - Wedge signal from the engine's own counter: WedgeMonitor now samples CBv2CapacitySnapshot.stepsExecuted (published every step) instead of the bridge's event-count proxy (H§3 resolved). New capacity test pins the passthrough plus the flatline-under-hanging-admits wedge verdict. - Logprobs passthrough (coordinator path): engine delta logprobs are converted to the OpenAI streaming entry shape and ride out-of-band via a per-request EngineV2LogprobsChannel (bridge pump appends before the chunk yield); the inference handler splices them into the next content-bearing SSE chunk as choices[0].logprobs = {content: [...]} before encryption. The legacy engine NEVER emitted logprobs, so this is the OpenAI-standard shape, not a legacy match; the standalone --local path (frames served inside the upstream router) still emits none. logprobs/top_logprobs are decoded from the sealed body (the upstream request type does not model them) like reasoning_effort. - cacheSalt forward plumbing (inert): the per-tenant cache scope (SHA256(prompt_cache_key)/SHA256(user), the same value the legacy scheduler threads into the checkpoint cache) now maps onto CBv2Request.cacheSalt; "" maps to nil (cache-level salt fallback). Production still constructs the v2 engine with prefixCache: nil. Tests: scripted-stub coverage for channel ordering/conversion, SSE splice shape + skip rules, extractLogprobsSpec, cacheSalt mapping on both bridge entry points, and engine wiring threading scope+plumbing. * chore(submodule): bump mlx-swift-lm to engine-v2 final tip a0ae73b (paged eligibility guard + bench driver) * chore(submodule): bump mlx-swift-lm to engine-v2 composed tip 255b223 (compiled decode + kernel-opt + profiler) * fix(engine-v2): thread sealed-body logit_bias + seed through the coordinator v2 translation The upstream OpenAIChatCompletionRequest models neither logit_bias nor seed, so the coordinator path's internal ChatCompletionRequest always carried nil for both and EngineV2Translation silently dropped them (PR #499 review, Codex P2 #1). Recover them the same way logprobs / reasoning_effort / cache scope are recovered: decode straight from the E2E-sealed body (ProviderLoop.extractSamplingOverrides, field-independent probes) and overlay onto the v2 translation via a new EngineV2SamplingOverrides plumbing struct. v2 engine path only — the legacy engine never honored either knob and its submit call is byte-identical; the --local path still drops both (no provider seam, documented in the narrowed KNOWN DEVIATION). Tests: sealed-body extraction (both fields, absence, malformed-field independence), translate overlay → CBv2SamplingParams, and a coordinator-path wiring test proving the parsed bias + seed reach the stub CBv2Engine. * fix(engine-v2): truthful shared-KV accounting + fp16 sizing for kv_quant models Two coupled memory-accounting gaps from PR #499 review (Codex P2 #2 + #3): Fix #2 (shared-budget KV accounting): the v2 bridge bypassed the shared GlobalKVCacheBudget, so model-load admission (availableMemoryGb minus outstanding reservations) and the legacy live-KV gate under-counted when a v2 slot was active alongside a legacy slot — a concurrent legacy load could over-commit unified memory. The v2 engine runs its own optimistic + preemption admission against a private byte ledger, so a per-request GATING reservation would fight that design (and double-count against the engine's ledger). Instead each in-flight v2 request records its worst-case KV footprint (prompt + maxTokens, fp16 rate) in the shared ledger via a new bookkeeping-only recordEngineKV (never rejects — the v2 engine stays authoritative for v2 acceptance), released on every terminal/teardown path in the pump. outstandingReservedBytes() now truthfully includes v2 in-flight KV; the v2 engine's own ledger is unchanged. Fix #3 (kv_quant fp16 sizing): engine_v2 builds UNQUANTIZED CBv2LayerCache regardless of kv_quant (makeProductionEngine never builds a quantized cache), but the slot factory sized the bridge from the scheduler's live kvBytesPerToken — the QUANTIZED rate when kv_quant=true — overstating v2 token budgets/heartbeats 2–4x and under-sizing the fix-#2 reservation. EngineV2KVSizing.resolve now picks the fp16 rate (scheduler fp16KVBytesPerToken) for v2 sizing and emits a WARN engine_health telemetry (engine_v2_kv_quant_unsupported) once per load when kv_quant actually engaged for the model. Tests: sizing decision (both kv_quant configs), record+release invariant (finish and no-terminal teardown), no-op when rate unknown, and the slot factory's fp16 sizing + WARN end-to-end via a new BatchScheduler rate seam. * harden(engine-v2): request-id validation, pump lifecycle, bounded logprobs buffer, sanity guards Security/hardening findings from PR #499 review (ethenotethan): #4 request-id validation: a caller-supplied request-id is validated (non-empty, ≤256 chars, no ASCII control chars) before it becomes a dictionary key / cancel-correlation handle; nil/invalid ids normalize to a fresh req-<uuid>. Static normalizedRequestId/isValidRequestId are pure and unit-tested. This also subsumes #11 (duplicate-id): prod ids are provider-minted UUIDs, and any malformed id is now normalized to a fresh UUID before the existing duplicate guard, so an attacker cannot induce a collision — no separate change needed. #5 nextRawId overflow: += → &+= (wrapping) with a comment; 2^64 is unreachable in practice, this makes the theoretical overflow explicitly defined rather than a trap. #6 engine.submit on the actor: verified O(1) non-blocking enqueue against the CBv2Engine contract (admission check + register + enqueue, no forward pass; tokenization already off-actor) and documented it at the call site. No code change. #7 pump task lifecycle: per-request pump Tasks are now tracked in pumpTasks and cancelled by shutdown() (before the engine drain), so none outlives the bridge; AsyncStream cancellation unwinds each pump through its teardown path, releasing per-request KV. Tested: shutdown cancels a live pump and its reservation returns to 0. #8 bounded logprobs buffer: EngineV2LogprobsChannel caps at 4096 entries with drop-oldest eviction + a droppedCount and a one-line WARN on first drop, so a stalled/never-draining consumer can't grow it unboundedly. Dropping logprobs never affects content/billing/hash (out-of-band). #9 dropped logit_bias keys: parseLogitBiasCountingDropped returns the count of dropped (non-numeric/negative) keys; samplingParams emits a count-only WARN (never the keys — request content) instead of dropping silently. #10 kvBytesCapacity sanity: clampKVBytesCapacity caps the v2 admission ceiling at physical RAM so a mis-derived budget degrades to all-of-memory rather than an absurd ceiling. Tests: id validation matrix + normalization-in-submit, wrapping id, shutdown-cancels-live-pumps, logprobs cap drop-oldest + droppedCount, dropped-key count, and the capacity clamp. Full suite green (1245). * fix(engine-v2): record shared-KV reservation atomically with admission; truthful kv_quant WARN Addresses both reviewers of the PR #499 fixes: Codex (race): the shared-budget KV reservation was recorded in the detached pump task, which runs AFTER submitTokenized returns — leaving a window where a concurrent model-load gate could observe zero v2 KV between engine admission and the pump starting. Make submit/submitTokenized async (all call sites already await) and record the reservation SYNCHRONOUSLY right after engine.submit succeeds, before returning the stream. The pump now only releases. Tests assert the no-gap invariant (reservation visible the instant submit returns, no polling). Claude (truthfulness): move the engine_v2_kv_quant_unsupported WARN below the makeBridgeIfSelected guard so it fires only when the v2 bridge actually built. On a v2 init-failure fallback the model serves via legacy, which DOES honor kv_quant, so warning there was untruthful. Full suite green (1245). * chore(submodule): bump mlx-swift-lm to engine-v2 tip e4158bc (PR#62 review fixes) * fix(engine-v2): gate v2 admission on the shared KV budget before engine submission The bridge now RESERVES each request's worst-case KV footprint in the process-wide GlobalKVCacheBudget before engine.submit — atomically on the budget actor (check + record in one hop) — and rejects with the canonical token_budget_exhausted capacity error when the shared pool has no headroom, exactly like the legacy scheduler's KV-reserve rejection. Engine rejections after the gate release the reservation. The non-gating recordEngineKV is removed (the gating reserve subsumes it). Degenerate maxTokens<=0 requests skip the gate (the engine finishes them without allocating KV). Round-2 PR#499 P2: with another slot's live KV reserved, record-after-accept could overcommit unified memory on multi-slot providers. * fix(engine-v2): report committed worst-case tokens in heartbeat budget fields The coordinator's token-budget gate (activeTokenBudgetUsed + queuedTokenBudget + requestTokens <= activeTokenBudgetMax) expects 'used' to carry the committed worst-case reservation of accepted requests — the legacy scheduler's Σ(prompt + maxTokens) — and never consults max_tokens_potential. Feeding it materialized KV bytes let the coordinator over-route while accepted long-max_tokens requests still had unprotected growth. Mapping (existing wire fields only, documented in backendSlotCapacity): activeTokens stays engine truth (real KV-resident tokens); activeTokenBudgetUsed and maxTokensPotential carry Σ(prompt+maxTokens); activeTokenBudgetMax is the engine byte capacity / fp16 rate (0 when the rate is unknown, disengaging the gate). Round-2 PR#499 P2. * fix(engine-v2): stable (seed, prompt)-derived engine ids for reproducible seeded sampling The v2 sampler keys its RNG on (seed, requestID.raw, stepIndex), so minting a fresh monotonic id per submission made identical seeded requests sample differently depending on prior traffic. Seeded submissions now derive their engine id from a SplitMix64 chain over (seed, promptTokens), tagged with bit 63 to stay disjoint from the monotonic family; the id is minted in the same synchronous stretch as engine.submit so the liveness collision check cannot race a concurrent identical submission. An identical seeded request that is still live falls back to a fresh id (documented waiver); batching reproducibility remains best-effort per the engine contract. Test: stochastic sampler stub keyed exactly like SamplerV2 proves same-seed + same-prompt submissions reproduce identical output across interleaved traffic, and diverge on different seed/prompt. Round-2 PR#499 P2. * fix(engine-v2): size v2 KV admission ceilings against fleet-wide residency A v2 engine's kvBytesCapacity was derived as if only its own slot's weights were resident, so multi-slot providers could grant Σ(ceilings) far past the unified-memory KV budget. The slot factory now derives the ceiling with EngineV2KVSizing.engineKVBytesCapacity: kvBudgetBytes over ALL resident models' weights (co-resident slots + the new model) minus the construction-fixed ceilings already granted to existing v2 engines. When nothing remains the capacity is 0 and makeProductionEngine throws noKVHeadroom, so the slot serves via the legacy scheduler (whose per-request shared-budget gate — and now the v2 bridge's too — needs no static ceiling). Test hooks now receive the computed capacity ((modelId, kvBytesCapacity)) so wiring tests assert the two-slot derivation end-to-end; pure sizing tests pin the sum-within-budget invariant and degenerate clamps. Round-2 PR#499 P2. * docs(engine-v2): hardening-review rationales — pump-task bound, logprobs lock, cancel scan - pumpTasks: bounded by the engine's own admission (maxConcurrent + maxWaiting, ~68 in production) — pump tasks exist only for engine-accepted requests and self-clear on terminal; cite the contract. - EngineV2LogprobsChannel: why an uncontended NSLock beats actor isolation for a two-party bounded append/swap on the delta path. - EngineV2Runtime.cancel: O(n) scan is over v2-served MODELS (<= slot cap, single digits), not requests; an id->bridge index would add hot-path register/unregister traffic to save a 3-entry scan on a rare event. * fix(engine-v2): duplicate re-check across the gate suspension; scope KV release to the reserving request The shared-budget reserve introduced the bridge's first suspension between the duplicate-id guard and the bookkeeping install. A concurrent same-id submission that SKIPS the gate (degenerate maxTokens / unknown rate) could slip in across that gap and be overwritten; and the pump's unconditional release could drop a same-keyed reservation owned by a different submission. Now: re-check active[id] after the reserve (roll back + reject as duplicate), and the pump releases only when its own submission reserved (holdsSharedReservation). Degenerate maxTokens<=0 requests are covered by a test proving they skip the gate even on an exhausted pool. * chore(submodule): bump mlx-swift-lm to engine-v2 tip ef1ad12 (bench cells on top of e4158bc) * chore(submodule): bump mlx-swift-lm to engine-v2 tip ec541f5 (PR#62 round-2 review fixes) * feat(provider): config keys for startup preload + auto-update rollover jitter New provider.toml keys (all additive; defaults chosen to be fleet-safe): [backend] - startup_preload (default true): preload models before coordinator registration on boot. A fresh install has no persisted serving set, so out-of-the-box behavior is unchanged until a model has been served. - preload_models (default []): explicit preload list, operator order; empty means "the models served before the last restart". - startup_preload_timeout_secs (default 120): upper bound on how long registration is deferred while the preload runs. - startup_selftest (default true): 1-token greedy decode through the serving path after each preload. - startup_selftest_fail_closed (default false): retire a model whose self-test failed instead of keeping it advertised (fail-open). [provider] - update_jitter_seconds (default 300): max random delay between staging a verified auto-update bundle and beginning the drain+restart. * feat(provider): startup model preload + registration readiness gate Fixes the v0.6.30-class cold-restart failure mode: after a release restart the provider registered (and attracted routing) with NOTHING loaded, so first requests paid the full ~26 GB weight load + engine build inside a live request — and at a fleet rollover every box was cold at once (first_chunk_timeout storm). - LoadedModelsStore: persists the live serving set to ~/.darkbloom/loaded-models.json (data dir — survives auto-updates) on every load and every NON-shutdown unload (idle timeout, eviction, retirement). Shutdown teardown deliberately preserves it so a stop/update/restart remembers what was being served. - StartupPreloader: sequential, deps-injected preload driver. Memory admission never evicts — a candidate that does not fit what is free is skipped with a WARN and left to the lazy-load path; failed loads and self-tests degrade, never crash. - ProviderLoop.runStartupPreloadGate(): runs in run() BEFORE the coordinator client exists. Plan = preload_models (operator order) or the persisted set (biggest first), filtered to advertised models and capped at max_model_slots. Each preload is the full ensureModelLoaded (weights + legacy scheduler + v2 bridge/warmup when flagged). Readiness vs availability: registration is deferred until warm, at most startup_preload_timeout_secs (default 120s). On timeout the provider registers anyway — a lone provider for a model must still serve it cold, and the lazy-load path is unchanged as the fallback — while the remaining loads continue in the background. The heartbeat warm_models field (which the coordinator's warm-model bonus scores) stays truthful throughout: it is derived from live slots only. - Startup self-test (startup_selftest, default on): after each preload, a 1-token greedy decode through the REAL serving path (v2 bridge when flagged, else legacy scheduler) forces end-to-end warmth (Metal JIT, compiled buckets, chat-template render). Failure emits WARN engine_health telemetry and is fail-open by default; startup_selftest_fail_closed retires the model from this run's advertised set (registration filters through advertisedModels). Tests (live-isolated, scripted loaders — no weights, no network): persistence round-trip incl. unload removal + the shutdown-preserve guard; preloader ordering/admission-skip/failure-continuation/ fail-open/fail-closed/cancellation; plan order + unknown-id filter + dedup + slot cap; gate defers-until-warm, proceeds-at-timeout with background continuation, and fail-closed retirement. * feat(provider): rollover jitter before background auto-update install Behavior change at defaults (update_jitter_seconds = 300) — justified as the fleet-rollover fix: providers discover a release on aligned 30-minute update ticks (ticks correlate after any previous synchronized restart) and drained + restarted in near-unison, putting the entire fleet cold at the same moment (first_chunk_timeout storm). A uniform random 0-300s delay between staging the verified bundle and beginning the drain decorrelates the restarts; combined with the startup preload gate, warm boxes absorb traffic while each cold box re-warms. Placement: the delay sits strictly AFTER download + signature verification (no security-critical check is deferred; the bundle is already verified and staged, and the provider keeps serving while it waits) and BEFORE beginDraining. Forced/manual updates — the darkbloom-update command and the startup update check — do not go through AutoUpdateController and stay immediate. A 1-hour sanity cap bounds a misconfigured value. Tests: delay bounds with a seeded RNG, zero-max and forced bypasses, the sanity cap, and controller sequencing (jitter strictly between stage and drain; never invoked on up-to-date / check-failed / stage-failed paths). * fix(provider): review fixes — cancel-safe jitter, no-evict preload loads, retirement + persistence hygiene Review round (Codex + Claude) on the startup-preload/jitter feature found four real gaps; all fixed with regression tests: 1. Cancel-safe rollover jitter. The jitter closure slept with `try? await Task.sleep`, so a shutdown cancelling autoUpdateTask inside the (up to 300s) jitter window returned early and the controller proceeded to drain + commit + RESTART a shutting-down provider. AutoUpdateController now checks Task.isCancelled after waitBeforeInstall and aborts with a new .cancelled outcome (resumeServing discards the staged bundle; the next tick retries). 2. No-evict preload loads, enforced in the real load path. The preloader's freeMemoryGb pre-check was advisory only — the actual ensureModelLoaded could still slot-cap-evict or memory-evict an EARLIER preloaded (idle) model to fit a later candidate, and a local-endpoint load interleaving with the driver could make the pre-check stale. ensureModelLoaded/evictUntilAvailable now take allowEviction (default true = live-traffic behavior unchanged); the startup preload passes false, and the check runs inside the serialized isLoadingAny critical section so it cannot go stale. Tested against the production load path (fake HF snapshot + real admission gates): no-evict refuses at the slot cap and the memory gate while the default policy still evicts. 3. Post-registration fail-closed retirement now reaches the coordinator. When the gate times out, registration has already announced the model; a later self-test retirement only edited local state, so the coordinator kept routing it. Retirement now mirrors the hard-swap drop: unadvertise on the client store + forceReconnect (new provider-side CoordinatorClient method; a fresh register is the existing wire mechanism for advertised-set removals — models_update is additive, and NO protocol changes are introduced). End-to-end test: real CoordinatorClient against MockCoordinator observes the second register without the retired model. 4. Loaded-models persistence is inert until run() arms it. The persist-on-load/unload write points fired from ANY ProviderLoop — including unit tests that install/unload slots — clobbering the operator's real ~/.darkbloom/loaded-models.json (the next boot's preload plan). Persistence is now gated on loadedModelsPersistenceEnabled, set only by run() (and explicitly by the test seam with a temp path). * fix(engine-v2): honor memory_reserve_gb in v2 KV sizing; clamp heartbeat budget max to live fleet residency (PR#499 r3 P2) Two coupled capacity-truth fixes on the v2 path: * memory_reserve_gb in v2 ceilings: UnifiedMemoryCap.kvBudgetBytes gains configReserveBytes (effective cap = min(hardCap, physical - reserve), the same max(configReserve, capImplied) hold-back the load gate and the shared KV gate apply); EngineV2KVSizing.engineKVBytesCapacity threads it and the slot factory passes the operator reserve. On 16/32 GiB boxes (default 4 GiB reserve > cap-implied reserve) the bridge no longer advertises/privately admits capacity the shared gate rejects post-acceptance. * Stale activeTokenBudgetMax after later loads: v2 admission ceilings are construction-fixed, so a model loaded later (especially a legacy/non-allowlisted slot) left existing bridges advertising a stale budget the coordinator kept routing against. The heartbeat now CLAMPS the reported max to the sizing function's current answer: ProviderLoop.updateAggregateCapacity snapshots live fleet residency (all slots' weights + operator reserve) into EngineV2Runtime.capacitySummary(fleetKV:), which recomputes each bridge's budget via EngineV2KVSizing.liveEngineKVBytesBudget (min(construction grant, current fleet answer)) and passes it to backendSlotCapacity(kvBytesBudgetClamp:). The engine's private ledger keeps its grant; only the coordinator-visible figure shrinks. Heartbeat cadence only - never the submit path. Tests: kvBudgetBytes reserve regimes (binding / no-op / degenerate), pure sizing with reserve both regimes, live-clamp pure math (later legacy load shrinks by exactly its weights; co-resident v2 grants subtract; never inflates past the grant; clamps to 0), and an end-to-end wiring test: load a v2 slot, register a second slot's weights, heartbeat max shrinks by exactly weights/rate while the engine keeps its construction grant. * fix(engine-v2): preserve queue-full classification on v2 admission rejections (PR#499 r3 P2) EngineV2.submit surfaces BOTH byte-ledger capacity exhaustion and its waiting-queue-full / draining-for-shutdown guards as CBv2KVError.capacityExhausted; the bridge mapped every shape to the generic token_budget_exhausted string, so v2 queue saturation surfaced as a 503 token-budget failure instead of the legacy 429 + Retry-After queue-full path. The two families are distinguishable by shape: the slot-rejection guards throw the sentinel (needed: 1, available: 0) (no byte estimate exists), while byte-ledger rejections carry needed = a multiple of the per-token KV cost (never 1) and available = admissibleBytesCapacity (> 0 for any constructible engine). EngineV2Translation .admissionErrorMessage now maps the sentinel to the exact legacy queue-full string ('token_budget_exhausted: request queue full', BatchSchedulerTypes.RejectionReason.queueFull), which fromSchedulerMessage classifies as .queueFull -> 429. Shutdown-drain riding the same sentinel lands on queue-full too, matching the legacy handler's .queueFull("provider shutting down"). No new classification strings, no wire changes. Tests: sentinel -> queue-full string -> .queueFull -> 429; byte-figure rejection stays .tokenBudgetExhausted -> 503; near-sentinel byte figures (needed 2/available 0, needed 1/available 512) never classify as queue-full. Also carries the bridge/runtime heartbeat-clamp tests for the previous commit (same test file). * fix(engine-v2): bound pending logprobs before the first content frame (PR#499 r3 P2) The inference handler drains the capped EngineV2LogprobsChannel into a local pendingLogprobs array every frame but only clears it when injectLogprobs finds a content-bearing frame - so a stream with a long reasoning-only/tool-only prefix (GPT-OSS reasoning) re-accumulated everything the channel cap had just bounded, growing without limit and eventually serializing thousands of entries into one huge first content chunk. EngineV2LogprobsChannel.capPending mirrors the channel's own policy on the caller-held buffer: cap at maxEntries (4096), drop-OLDEST so the freshest window (the entries nearest the content that eventually renders) is kept, return the evicted count. The handler re-caps after every drain, WARNs once on first drop and logs the request's total at stream end (counts only - never token text; entries are response content and stay inside the E2E stream). Tests: drop-oldest at the real 4096 cap with count reporting and no-op below the cap; frames-loop simulation where a capped window survives reasoning-only frames and the first content chunk carries exactly the retained entries. * chore(submodule): bump mlx-swift-lm to engine-v2 tip 0c5d3a9 (round-3 review fixes) * chore(submodule): pin mlx-swift-lm to main d63b1c8 (squash-merge of engine-v2 PR #62; content-identical to 0c5d3a9) * fix(provider): startup self-test no longer poisons Qwen3.5 slots — stale VLM mrope position ids + e2e state isolation Root cause of the CI e2e regression (test 1 passes, every later test's provider dies on its first inference with '[broadcast_shapes] Shapes (1,8,L,64) and (1,1,12,64) cannot be broadcast'): 1. Qwen3.5-0.8B declares a vision_config, so the provider loads it via VLMModelFactory. The MLXVLM Qwen35 LanguageModel cached the full-prompt mrope position ids (precomputedPositionIds) on the long-lived model module and reused them for ANY later forward at cache offset 0. The startup self-test decode ('Hi', 12 prompt tokens, maxTokens=1) was the first request in every preloading provider, so the first routed request with a longer prompt sliced the stale 12-position ids (MLX clamps the out-of-range slice), producing cos/sin frozen at (1,1,12,64) against the new request's (1,8,L,64) rotary block — process-fatal in the first full-attention layer. Fixed in the mlx-swift-lm submodule (d3b1ae9): always recompute position ids at sequence start, exactly like Qwen3VL and HF transformers; precomputedPositionIds removed. Also covers Qwen35MoE (subclass). NOT self-test-specific: ANY second, longer text-only request on one scheduler crashed the same way. 2. The e2e testbed launched providers with the real user HOME, so ~/.darkbloom/loaded-models.json leaked across testbed instances: test 1's provider persisted its serving set, and test 2's freshly-booted provider startup-preloaded + self-tested it — the trigger that made every post-first test hit (1). Providers now get a per-instance temp state dir (DARKBLOOM_STATE_FILE + DARKBLOOM_LOADED_MODELS_FILE), removed on Stop. Regression test: StartupSelfTestDecodeLiveTests (weight-gated, DARKBLOOM_LIVE_MLX_TESTS) loads Qwen3.5-0.8B through the SAME VLMModelFactory selection production uses and drives the self-test-shaped request then a longer request through MultiModelBatchSchedulerEngine + MLXOpenAIService — crashes with the exact production signature without the submodule fix, passes with it. Validated end-to-end: 4 sequential e2e integration tests green locally (previously test 2+ crashed 502). * feat(provider): v0.7.0 engine posture — v2 default-on for exact prod checkpoints, legacy compiled decode opt-in, finish_reason 'length' threaded end-to-end Three coordinated changes for the v0.7.0 release posture: 1. engine_v2 default TRUE, allowlist narrowed to the EXACT production checkpoint ids (mlx-community/gpt-oss-20b-MXFP4-Q8, mlx-community/gemma-4-26b-a4b-it-8bit) — default-on scope equals the real-model-validated scope; other quantizations/families keep legacy until widened via DARKBLOOM_ENGINE_V2_MODELS. DARKBLOOM_ENGINE_V2=0 stays the absolute per-box kill switch (beats config, no release). The v2-init-failure fallback telemetry (engine_v2_fallback, WARN engine_health) now also carries the human-readable 'error' reason next to error_class + model — it is load-bearing now that v2 is the default. Exact-pattern allowlist matching compares last path components on both sides so bare ids match org-qualified patterns. Retirement (failed self-test) and hard-swap/idle/eviction paths were audited for default-on v2: every unload funnels through ProviderLoop.unloadModel, which unregisters the bridge from the runtime and drains it gracefully before scheduler teardown; the hard-swap dropAdvertisedBuild is a lazy advertised-set drop that leaves in-flight requests on the resident slot untouched. 2. Legacy compiled decode is now OPT-IN (new [backend] key legacy_compiled_decode, default false). The bumped mlx-swift-lm pin ships CompiledDecode default-ON; LegacyCompiledDecodeGate forces DARKBLOOM_COMPILED_DECODE=0 into the process env at both serve entry points (ProviderLoop.run, local standalone) BEFORE any engine/model code latches the library's static — unless the operator opted in via config or set the env var explicitly (explicit env always wins, both directions). Keeps release behavior identical to prod v0.6.30 (the compiled-decode rollback); single-stream speed ships via the v2 canary instead. 3. finish_reason 'length' is threaded end-to-end instead of being flattened to 'stop'. GenerationEvent.info gains finishReason; the legacy engine bridge passes RequestOutput.finishReason through, the v2 bridge maps CBv2FinishReason.length/.stop, the B=1 greedy fast path infers length from a full-budget decode, and MultiModelBatchSchedulerEngine forwards it as ServerGenerationInfo.stopReason (tool_calls still wins). Clients now see finish_reason 'length' on max_tokens truncation from BOTH engines. Tests: LegacyCompiledDecodeGateTests (gate matrix + real-env apply, snapshot/restore), EngineV2 gating suite updated for the exact-id allowlist + default-on selection matrix (kill switch beats default-on), fallback telemetry asserts the new 'error' field, v2 bridge test for .length/.stop/cancelled reason threading. Full provider suite green (1327 tests); live engine-path tests green (weight-gated). * docs(provider): review-nit comment fixes — GenerationEvent.info arity in docs, gate empty-env rationale
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
POST /v1/images/generationsendpoint with E2E encryption, per-image billing, and cancellation supportArchitecture
The Python bridge translates OpenAI REST requests into Draw Things' gRPC + FlatBuffers protocol. gRPCServerCLI runs as a subprocess managed by the bridge.
Changes
Protocol (Go + Rust):
image_generation_request/image_generation_completemessage types with E2E encryption, round-trip testsCoordinator:
POST /v1/images/generationsconsumer endpoint,handleImageGenerationCompleteprovider handler with charge/credit/referral billing pipelineProvider: Event dispatch, decrypt/parse in coordinator.rs, proxy handler POSTs to bridge on
be_port+2, 5-min timeout, cancellation via CancellationTokenImage Bridge: FastAPI server with DrawThingsBackend — builds FlatBuffers configs, calls GenerateImage gRPC, decodes float16 tensor responses to PNG
Billing:
CalculateImageCost()per-image pricing scaled by resolution ($0.002/image at 1024x1024 base), same 95/5 provider/platform splitInstall Script: 8-15 GB machines now get image model (FLUX.2 Klein 4B) instead of falling through to tiny chat model
Test plan
TestCalculateImageCost)🤖 Generated with Claude Code