Compare commits
37
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3ed172e8b7
|
||
|
|
106345495a
|
||
|
|
d033e8591b | ||
|
|
9c78a2254c
|
||
|
|
5a160ac70f
|
||
|
|
9264eeaaa8
|
||
|
|
1a5fd0d793 | ||
|
|
40146c4dd9
|
||
|
|
530084beef
|
||
|
|
939343b34c
|
||
|
|
d6c49217b0
|
||
|
|
3aaf523caf
|
||
|
|
1e7256ff61
|
||
|
|
6731f74ae6 | ||
|
|
992c819e27
|
||
|
|
cdcd5bb092
|
||
|
|
c7505a9344 | ||
|
|
8ed8d04d3a | ||
|
|
6a0d4fbb0c | ||
|
|
73a376e57c
|
||
|
|
06cf081c47
|
||
|
|
0a122fb26f
|
||
|
|
e3248cc186
|
||
|
|
057cb72806
|
||
|
|
5264f78b2d
|
||
|
|
2cbfd06e40 | ||
|
|
816c8dc96f
|
||
|
|
042d708da3 | ||
|
|
421bf55697
|
||
|
|
e242514121
|
||
|
|
1ba160491c
|
||
|
|
047005203e
|
||
|
|
86a6f25a6b
|
||
|
|
ec315a7f58
|
||
|
|
7c17a9fcd4 | ||
|
|
c6502e1519
|
||
|
|
21a1083308 |
@@ -525,7 +525,14 @@ jobs:
|
||||
cuda_home: /usr/local/cuda-13.0
|
||||
build_jobs: 8
|
||||
nvcc_threads: 4
|
||||
cargo_features: "cuda cudnn"
|
||||
# flash-attn on ampere (#337): Laya's encoder runs every
|
||||
# attention through flash_attn_varlen_windowed, matching the
|
||||
# PyTorch reference on a 3060 (36 vs 39 ms at 512 tokens) where
|
||||
# the eager path is 1.4-1.7x slower. This also moves the
|
||||
# flavour's other flash-capable paths onto flash -- qwen3_5
|
||||
# full attention and Z-Image -- which validate-neuron.sh and
|
||||
# validate-image.sh gate. NEURON_FLASH_ATTN=0 rolls back.
|
||||
cargo_features: "cuda cudnn flash-attn"
|
||||
- flavour: ada
|
||||
compute_cap: "89"
|
||||
runner: cuda-13.0
|
||||
@@ -540,8 +547,8 @@ jobs:
|
||||
build_jobs: 8
|
||||
nvcc_threads: 4
|
||||
# flash-attn on blackwell first (#95): beast carries the
|
||||
# agentic prefill pain; ada/ampere follow once the win is
|
||||
# measured. NEURON_FLASH_ATTN=0 is the runtime rollback.
|
||||
# agentic prefill pain. NEURON_FLASH_ATTN=0 is the runtime
|
||||
# rollback.
|
||||
cargo_features: "cuda cudnn flash-attn"
|
||||
runs-on: ${{ matrix.runner }}
|
||||
env:
|
||||
|
||||
@@ -244,10 +244,14 @@ jobs:
|
||||
flavour: blackwell
|
||||
load_timeout: 900
|
||||
max_prompt_tokens: 131072
|
||||
- host: benjy.hanzalova.internal
|
||||
flavour: ada
|
||||
load_timeout: 300
|
||||
max_prompt_tokens: 16384
|
||||
# benjy is out of the matrix while its boot NVMe is replaced
|
||||
# (dead 2026-09-29): every deploy failed at SSH init against a
|
||||
# host that is not there. Restore these four lines when it is
|
||||
# back on the mesh.
|
||||
# - host: benjy.hanzalova.internal
|
||||
# flavour: ada
|
||||
# load_timeout: 300
|
||||
# max_prompt_tokens: 16384
|
||||
- host: quadbrat.hanzalova.internal
|
||||
flavour: ampere
|
||||
load_timeout: 300
|
||||
|
||||
@@ -28,3 +28,6 @@ doc/plan/*
|
||||
.claude/
|
||||
helexa-upstream.toml
|
||||
helexa-angels.toml
|
||||
|
||||
# Local scratch space for playground example checks (never committed).
|
||||
.pg-scratch/
|
||||
|
||||
@@ -1143,3 +1143,102 @@ constructors on the MoE and MLP.
|
||||
budget matches a naive oracle — are closed.
|
||||
- **MTP (#313), vision (#314), quantisation and rollout (#315)** are
|
||||
untouched.
|
||||
|
||||
## 2026-09-29 addendum: System-One decisions (Laya, epic #333)
|
||||
|
||||
A third modality beside text and images: **typed decisions**.
|
||||
`convaiinnovations/laya` is a 421M non-autoregressive decision model, a
|
||||
ModernBERT encoder plus an option-scoring head. It answers `choice` /
|
||||
`score` / `noul` questions over a *state* in one encoder forward, with
|
||||
calibrated probabilities, on the TypeSafe Jev wire protocol:
|
||||
`POST /v1/systemone`. It never generates text. Resident on quadbrat's
|
||||
3060 since #351, replacing the Qwen3-1.7B pre-warm.
|
||||
|
||||
**The reference is the `laya` pip SDK** (`common.py`, `agent.py`,
|
||||
`router.py`, `serve.py`), **not** the `rl_common.py` in the model repo,
|
||||
which has diverged. Fixtures recorded from SDK 0.3.21 at model revision
|
||||
`55cf4c4e` live in `crates/neuron/src/harness/testdata/laya/`:
|
||||
- `reference.json`: 28 cases, f32 CPU.
|
||||
- `reference_cuda_bf16.json`: the SDK's own GPU bf16 run, the envelope
|
||||
our bf16 is judged against.
|
||||
- `lang_corpus.json`: 924 detector cases.
|
||||
|
||||
Recorders are `script/laya-reference.py` and `script/laya-lang-reference.py`.
|
||||
|
||||
**Where it lives:**
|
||||
- `cortex-core/src/decisions.rs`: wire types, and `Json`, an
|
||||
order-preserving, correctly rounded, `json.loads`-compatible value.
|
||||
Request data must not go through `serde_json::Value`: the workspace
|
||||
has no `preserve_order`, and without `float_roundtrip` serde_json
|
||||
misreads floats (`…445` → `…443`). Key order changes token ids and
|
||||
what the language detector reads.
|
||||
- `neuron/src/harness/decision/`: question validation, sequence
|
||||
assembly under every budget rule, calibration in numpy's float32
|
||||
order, answer shaping, and `route.rs` (checkpoint choice, ported from
|
||||
upstream `Router._route`, reason strings included).
|
||||
- `neuron/src/harness/lang_detect.rs`: the script/language detector.
|
||||
Python-semantics Unicode classes: `[^\W\d_]` never matches a combining
|
||||
mark.
|
||||
- `neuron/src/harness/arch/laya/`: the engine. The ModernBERT backbone
|
||||
is **vendored**, because candle's only runs in f32.
|
||||
- `decision_model.rs`, plus `CandleHarness::systemone` in `candle.rs`.
|
||||
|
||||
**One repo, three checkpoints:** English at the root, `multilingual/`
|
||||
and `typed-decisions/`, all resident. Precedence: a `model` pin (a
|
||||
checkpoint name, an alias, or a published id such as
|
||||
`convaiinnovations/laya-multilingual`), else the detected script and
|
||||
language, else English. The English checkpoint collapses outside
|
||||
English (Khmer: 0.000 accuracy at 0.952 confidence upstream), so
|
||||
language routing is correctness, not polish.
|
||||
|
||||
**Surface semantics:**
|
||||
- On this endpoint only, a missing or unknown `model` resolves to
|
||||
**`helexa/one`** at every hop, because Jev clients send their own
|
||||
service's name (`jev-1`).
|
||||
- The body is never rewritten, since `model` may be a checkpoint pin;
|
||||
the routed model travels in `x-helexa-decision-model`.
|
||||
- Metering follows Jev: **input tokens only**. That is the sum of
|
||||
encoded sequence lengths, and the state is encoded once per question.
|
||||
Reserved as questions × `max_len`, settled to `usage.input_tokens`,
|
||||
priced at $0.042/M.
|
||||
- Errors carry the #63 envelope plus FastAPI's `detail`.
|
||||
|
||||
**GPU performance (3060, #337).** Laya needs the ampere flavour's
|
||||
`flash-attn` (enabled in #350; `NEURON_FLASH_ATTN=0` rolls back). What
|
||||
the profile found:
|
||||
- candle's `LayerNorm` only fuses when a bias exists, and ModernBERT's
|
||||
norms have none: six kernels per norm, about 25% of GPU time. Fixed
|
||||
with an explicit zero bias.
|
||||
- Mask add and scale over the full score tensor. Elided when nothing is
|
||||
masked; scale folded into q.
|
||||
- Eager attention replaced by `flash_attn_varlen_windowed` over an
|
||||
unpadded pack (sliding layers ±64).
|
||||
|
||||
Measured, bf16: 36 ms for one question over a 512-token state (reference
|
||||
39), 32 seq/s batched. Live: 10.6 ms median at the neuron over the mesh.
|
||||
bf16 drift is autocast-style: residual, norms and attention scores stay
|
||||
in f32. Beyond that, per-case drift is chaotic; judge it by
|
||||
mean-vs-reference-envelope, not single cases.
|
||||
|
||||
**Validation:**
|
||||
- `script/validate-decision.py <url>` replays the corpus against any
|
||||
hop: checkpoint and usage exact, answers except on reference
|
||||
near-ties, probabilities within a bf16 tolerance.
|
||||
- The public edge is `https://helexa.ai/v1/systemone`. It is rate-limited
|
||||
(10 r/min per IP), so replay against the router (gallumbits:8088).
|
||||
Probing the public names has network constraints; see
|
||||
`doc/infra/edge-ingress.md` in the private helexa/security repo.
|
||||
|
||||
**Sharp edges:**
|
||||
- **rustingface#6:** a neuron cold load of anything rf.internal hasn't
|
||||
cached fails on a 503 to hf-hub's bounded-range size probe. A Laya
|
||||
family member then loads *without* that checkpoint, logged but
|
||||
silent. Warm with full GETs, then reload.
|
||||
- helexa-bench chat scenarios must gate on `serves_chat()`, which
|
||||
excludes both `image` and `decision` models.
|
||||
|
||||
**Open:** #353 (revision pinning, per-language temperatures, fetch
|
||||
retries, a bench `decision:` scenario, a dashboard panel,
|
||||
micro-batching); #342 (serving fine-tuned checkpoints); #352 (Z-Image on
|
||||
the ampere flash build, unrelated to Laya but enabled by the same
|
||||
flavour change).
|
||||
|
||||
Generated
+65
-50
@@ -40,7 +40,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cabdc9d845d08ec7ed2d0c9de1ae4a1b198301407d55855261572761be90ec9f"
|
||||
dependencies = [
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -261,7 +261,7 @@ checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -272,13 +272,13 @@ checksum = "8b75356056920673b02621b35afd0f7dda9306d03c79a30f5c56c44cf256e3de"
|
||||
|
||||
[[package]]
|
||||
name = "async-trait"
|
||||
version = "0.1.89"
|
||||
version = "0.1.92"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb"
|
||||
checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 3.0.9",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -394,7 +394,7 @@ checksum = "7aa268c23bfbbd2c4363b9cd302a4f504fb2a9dfe7e3451d66f35dd392e20aca"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -508,7 +508,7 @@ checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -697,7 +697,7 @@ dependencies = [
|
||||
"heck",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -853,6 +853,7 @@ dependencies = [
|
||||
"reqwest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tempfile",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
@@ -861,6 +862,7 @@ dependencies = [
|
||||
"tracing",
|
||||
"url",
|
||||
"urlencoding",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -966,7 +968,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "10a2a99df6e410a8ff4245aa2006499ea662245f967cc7c0a38c83ef8eb44dbf"
|
||||
dependencies = [
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1039,7 +1041,7 @@ dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"strsim",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1052,7 +1054,7 @@ dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"strsim",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1063,7 +1065,7 @@ checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead"
|
||||
dependencies = [
|
||||
"darling_core 0.20.11",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1074,7 +1076,7 @@ checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d"
|
||||
dependencies = [
|
||||
"darling_core 0.23.0",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1125,7 +1127,7 @@ dependencies = [
|
||||
"darling 0.20.11",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1135,7 +1137,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c"
|
||||
dependencies = [
|
||||
"derive_builder_core",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1157,7 +1159,7 @@ dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"rustc_version",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
"unicode-xid",
|
||||
]
|
||||
|
||||
@@ -1202,7 +1204,7 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1337,7 +1339,7 @@ dependencies = [
|
||||
"heck",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1667,7 +1669,7 @@ checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3096,7 +3098,7 @@ checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3166,6 +3168,8 @@ dependencies = [
|
||||
"toml",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"unicode-normalization",
|
||||
"unicode-properties",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3372,7 +3376,7 @@ checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3471,7 +3475,7 @@ dependencies = [
|
||||
"proc-macro2",
|
||||
"proc-macro2-diagnostics",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3540,7 +3544,7 @@ dependencies = [
|
||||
"phf_shared",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3569,7 +3573,7 @@ checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3692,7 +3696,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3712,7 +3716,7 @@ checksum = "af066a9c399a26e020ada66a034357a868728e72cd426f3adcd35f80d88d88c8"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
"version_check",
|
||||
"yansi",
|
||||
]
|
||||
@@ -4064,7 +4068,7 @@ checksum = "b7186006dcb21920990093f30e3dea63b7d6e977bf1256be20c3563a5db070da"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4189,7 +4193,7 @@ dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"serde_json",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4392,7 +4396,7 @@ dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"serde_derive_internals",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4482,7 +4486,7 @@ checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4493,7 +4497,7 @@ checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4579,7 +4583,7 @@ dependencies = [
|
||||
"darling 0.23.0",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4812,7 +4816,7 @@ dependencies = [
|
||||
"quote",
|
||||
"sqlx-core",
|
||||
"sqlx-macros-core",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4835,7 +4839,7 @@ dependencies = [
|
||||
"sqlx-mysql",
|
||||
"sqlx-postgres",
|
||||
"sqlx-sqlite",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
"tokio",
|
||||
"url",
|
||||
]
|
||||
@@ -5020,7 +5024,7 @@ dependencies = [
|
||||
"heck",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -5040,6 +5044,17 @@ dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "syn"
|
||||
version = "3.0.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sync_wrapper"
|
||||
version = "1.0.2"
|
||||
@@ -5057,7 +5072,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -5143,7 +5158,7 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -5154,7 +5169,7 @@ checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -5280,7 +5295,7 @@ checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -5438,7 +5453,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -5798,7 +5813,7 @@ dependencies = [
|
||||
"bumpalo",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
"wasm-bindgen-shared",
|
||||
]
|
||||
|
||||
@@ -5988,7 +6003,7 @@ checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -5999,7 +6014,7 @@ checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -6313,7 +6328,7 @@ dependencies = [
|
||||
"heck",
|
||||
"indexmap 2.14.0",
|
||||
"prettyplease",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
"wasm-metadata",
|
||||
"wit-bindgen-core",
|
||||
"wit-component",
|
||||
@@ -6329,7 +6344,7 @@ dependencies = [
|
||||
"prettyplease",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
"wit-bindgen-core",
|
||||
"wit-bindgen-rust",
|
||||
]
|
||||
@@ -6423,7 +6438,7 @@ checksum = "2380878cad4ac9aac1e2435f3eb4020e8374b5f13c296cb75b4620ff8e229154"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
"synstructure",
|
||||
]
|
||||
|
||||
@@ -6435,7 +6450,7 @@ checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
"synstructure",
|
||||
]
|
||||
|
||||
@@ -6456,7 +6471,7 @@ checksum = "70e3cd084b1788766f53af483dd21f93881ff30d7320490ec3ef7526d203bad4"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -6476,7 +6491,7 @@ checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
"synstructure",
|
||||
]
|
||||
|
||||
@@ -6516,7 +6531,7 @@ checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -52,21 +52,10 @@ curl -s 'http://localhost:26559/api/v1/targets' | jq '.data.activeTargets[]|sele
|
||||
|
||||
`grafana-helexa-fleet.json` is a raw dashboard model with a `datasource`
|
||||
template variable, so it binds to whichever Prometheus data source you pick
|
||||
at import. Import via the API (creds in `/etc/grafana/grafana.env`):
|
||||
|
||||
```sh
|
||||
# on golgafrinchans; GF_SECURITY_ADMIN_USER/PASSWORD live in the root-only
|
||||
# env file, Grafana is published on :28767
|
||||
set -a; source <(sudo cat /etc/grafana/grafana.env); set +a
|
||||
jq -n --slurpfile d grafana-helexa-fleet.json \
|
||||
'{dashboard: $d[0], overwrite: true, folderUid: null}' \
|
||||
| curl -s -u "$GF_SECURITY_ADMIN_USER:$GF_SECURITY_ADMIN_PASSWORD" \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d @- http://localhost:28767/api/dashboards/db | jq '{status,uid,version}'
|
||||
```
|
||||
|
||||
The dashboard lands at uid `helexa-fleet`. Re-importing with `overwrite:true`
|
||||
updates it in place.
|
||||
at import. Import it through the Grafana dashboards API with admin
|
||||
credentials (`overwrite: true` updates it in place; it lands at uid
|
||||
`helexa-fleet`). The exact command and where the admin credentials live are
|
||||
in `doc/infra/monitoring-access.md` in the private helexa/security repo.
|
||||
|
||||
## What the dashboard shows
|
||||
|
||||
|
||||
@@ -7,9 +7,10 @@
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# neuron.toml for quadbrat.hanzalova.internal
|
||||
#
|
||||
# 1x RTX 3060 (12 GB) — small / quantised tier. Pre-warms Qwen3-1.7B
|
||||
# (bf16, ~4 GB), leaving ~7 GB for KV cache so long contexts on a small
|
||||
# model still have plenty of room.
|
||||
# 1x RTX 3060 (12 GB). Resident: Laya, the System-One decision model
|
||||
# (#333, ~2.3 GB bf16 for its three checkpoints). It replaced the
|
||||
# Qwen3-1.7B pre-warm (operator decision 2026-09-29, #341); the 1.7B is
|
||||
# still cold-loadable beside it for helexa/small.
|
||||
#
|
||||
# Synced to /etc/neuron/neuron.toml by script/infra-setup.sh.
|
||||
|
||||
@@ -21,10 +22,21 @@ name = "candle"
|
||||
[harness.candle]
|
||||
|
||||
[[default_models]]
|
||||
model_id = "Qwen/Qwen3-1.7B"
|
||||
model_id = "convaiinnovations/laya"
|
||||
harness = "candle"
|
||||
devices = [0]
|
||||
|
||||
# A decision request is one forward of tens of milliseconds (36 ms for
|
||||
# one question over a 512-token state, 162 ms for five), so a one-slot
|
||||
# queue of 8 rejects ordinary bursts. Admit
|
||||
# several at once -- the device worker still runs forwards one at a
|
||||
# time -- and let a deep, short-deadline queue absorb the rest. Applies
|
||||
# to every model on this host, including a cold-loaded helexa/small.
|
||||
[harness.candle.admission]
|
||||
max_in_flight = 4
|
||||
max_queue_depth = 64
|
||||
max_wait_secs = 10
|
||||
|
||||
# Text-to-image (#204): the 12 GB tier serves the q8_0-quantized DiT
|
||||
# (8.7 GB resident incl. VAE). Ceiling is 768² — naive f32 attention
|
||||
# transients OOM 1024² beside the resident model; lifts to 1024² if
|
||||
|
||||
@@ -34,8 +34,8 @@ server {
|
||||
|
||||
server {
|
||||
# TCP 443 on the edge is owned by the stream SNI router
|
||||
# (streams-enabled/sni-router.conf, architecture reverse-proxies.md
|
||||
# §5), which preads the SNI and hands every local name to this http
|
||||
# (streams-enabled/sni-router.conf; see doc/infra/edge-ingress.md in
|
||||
# the private helexa/security repo), which preads the SNI and hands every local name to this http
|
||||
# tier over PROXY protocol. Listening on :443 directly here would
|
||||
# never be reached — the router would answer first with whichever
|
||||
# vhost happened to be default, which is exactly the symptom seen
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# Internal bench UI vhost — https://bench.internal, reachable from inside
|
||||
# the WireGuard mesh (the public bench.helexa.ai dead-ends at the OPNsense
|
||||
# LAN interface, which only port-forwards :443 from the WAN). Same SPA +
|
||||
# Internal bench UI vhost — https://bench.internal, for clients inside
|
||||
# the WireGuard mesh, where the public bench.helexa.ai doesn't resolve
|
||||
# to a usable path (see doc/infra/edge-ingress.md in the private
|
||||
# helexa/security repo). Same SPA +
|
||||
# /api→bob proxy as bench.helexa.ai, but with an internal-CA cert
|
||||
# (smallstep "lair", renewed by step@bench.timer). Mirrors the
|
||||
# *.internal vhost convention on oolon.kosherinata.internal.
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
# Installed on BOTH edge proxies (oolon.kosherinata.internal and
|
||||
# hanzalova.internal); Cloudflare DNS load balancing routes the public
|
||||
# name to the two site WAN IPs. Mesh clients use helexa.internal instead
|
||||
# (public names don't hairpin — see architecture/reverse-proxies.md §2).
|
||||
# (see doc/infra/edge-ingress.md in the private helexa/security repo).
|
||||
#
|
||||
# Same-origin backends (#167): /v1 (+ /health) reverse-proxies to
|
||||
# helexa-router and /api to helexa-upstream, both on gallumbits — so the
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# Internal helexa website vhost — https://helexa.internal, reachable from
|
||||
# inside the WireGuard mesh (the public helexa.ai dead-ends at the
|
||||
# OPNsense LAN interface, which only port-forwards :443 from the WAN —
|
||||
# see architecture/reverse-proxies.md §2). Same SPA webroot as the
|
||||
# Internal helexa website vhost — https://helexa.internal, for clients
|
||||
# inside the WireGuard mesh, where the public name doesn't resolve to a
|
||||
# usable path (see doc/infra/edge-ingress.md in the private
|
||||
# helexa/security repo). Same SPA webroot as the
|
||||
# public vhost, but with an internal-CA cert (smallstep "lair", renewed
|
||||
# by step@helexa.timer). Installed on hanzalova.internal only; the
|
||||
# split-horizon helexa.internal record points at hanzalova's mesh IP on
|
||||
|
||||
@@ -113,3 +113,19 @@ enabled = false
|
||||
# timeout_secs = 5
|
||||
# How often to flush served-usage counters to upstream for reconciliation (#58).
|
||||
# served_usage_report_interval_secs = 60
|
||||
|
||||
# -- Async audio generation jobs (#382) -----------------------------------
|
||||
# Music renders take seconds to minutes, so /v1/audio/generations is a job
|
||||
# API: submit, poll, fetch the audio. Job records and rendered audio live
|
||||
# under `dir` and survive a restart (the packaged unit's StateDirectory
|
||||
# creates /var/lib/cortex). All values below are the defaults.
|
||||
[audio_jobs]
|
||||
# dir = "/var/lib/cortex/audio-jobs"
|
||||
# Finished jobs and their audio are deleted after this many hours.
|
||||
# retention_hours = 72
|
||||
# Renders running at once across the fleet; the rest queue in cortex.
|
||||
# max_concurrent = 4
|
||||
# Dispatch attempts per job; a busy (429/503) model does not use one up.
|
||||
# max_attempts = 2
|
||||
# How long a job may wait on a busy model before failing, in seconds.
|
||||
# queue_timeout_secs = 3600
|
||||
|
||||
@@ -0,0 +1,342 @@
|
||||
//! Music / audio generation surface (#382).
|
||||
//!
|
||||
//! OpenAI has no music endpoint, so this is a helexa extension. It follows
|
||||
//! the images surface where that makes sense (#197) but is **asynchronous**:
|
||||
//! a song takes seconds to minutes to render, so the public API is a job API.
|
||||
//!
|
||||
//! - `POST /v1/audio/generations` submit → `202` + [`AudioJob`]
|
||||
//! - `GET /v1/audio/generations` list the caller's jobs
|
||||
//! - `GET /v1/audio/generations/{id}` poll → [`AudioJob`]
|
||||
//! - `GET /v1/audio/generations/{id}/content` the rendered audio bytes
|
||||
//! - `DELETE /v1/audio/generations/{id}` cancel
|
||||
//!
|
||||
//! cortex owns the job record. A neuron renders exactly one generation per
|
||||
//! dispatch, synchronously, on `POST /v1/audio/generations`: the response
|
||||
//! body is the audio (`audio/wav`) and the result metadata travels in the
|
||||
//! [`RESULT_HEADER`] header as JSON ([`NeuronAudioResult`]). Keeping the
|
||||
//! bytes out of a JSON envelope matters here — ten minutes of 48 kHz stereo
|
||||
//! is over 100 MB of PCM.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Header carrying [`NeuronAudioResult`] JSON on a neuron's audio response.
|
||||
pub const RESULT_HEADER: &str = "x-helexa-audio-result";
|
||||
|
||||
/// Capability string a model advertises when it serves this surface.
|
||||
pub const CAPABILITY: &str = "audio";
|
||||
|
||||
/// `POST /v1/audio/generations` request body. The same body is accepted by
|
||||
/// cortex (which turns it into a job) and by neuron (which renders it).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct AudioGenerationRequest {
|
||||
pub model: String,
|
||||
/// Style / caption: what the music should sound like.
|
||||
pub prompt: String,
|
||||
/// Lyrics for models that sing. Absent, empty, or `[Instrumental]`
|
||||
/// means an instrumental.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub lyrics: Option<String>,
|
||||
/// Requested length in seconds. Defaults to the model's advertised
|
||||
/// default; validated against its advertised range.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub duration: Option<f64>,
|
||||
/// Fixed RNG seed for reproducible renders. The seed actually used is
|
||||
/// always reported back, so a random render can be reproduced too.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub seed: Option<u64>,
|
||||
/// Container for the delivered audio. Only `"wav"` today.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub response_format: Option<String>,
|
||||
|
||||
// ── model-specific extensions ───────────────────────────────
|
||||
/// Diffusion steps (diffusion models only).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub steps: Option<u32>,
|
||||
/// Guidance / CFG scale.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub guidance_scale: Option<f64>,
|
||||
/// Song planner: `"off"`, `"auto"` (model default) or `"on"`, for
|
||||
/// models that have one (ACE-Step's LM).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub planner: Option<String>,
|
||||
/// Sampling temperature (autoregressive models only).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub temperature: Option<f64>,
|
||||
/// Top-k sampling (autoregressive models only).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub top_k: Option<u32>,
|
||||
}
|
||||
|
||||
/// Response formats the surface delivers.
|
||||
pub const SUPPORTED_FORMATS: &[&str] = &["wav"];
|
||||
|
||||
impl AudioGenerationRequest {
|
||||
/// True when the request asks for an instrumental (no singing).
|
||||
pub fn is_instrumental(&self) -> bool {
|
||||
match self.lyrics.as_deref().map(str::trim) {
|
||||
None | Some("") => true,
|
||||
Some(l) => l.eq_ignore_ascii_case("[instrumental]"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Transport-level validation every hop can do without knowing the
|
||||
/// model: required fields, sane numbers, a known format. Model-specific
|
||||
/// limits ([`AudioLimits`]) are checked separately.
|
||||
pub fn validate_shape(&self) -> Result<(), String> {
|
||||
if self.model.trim().is_empty() {
|
||||
return Err("'model' must be non-empty".into());
|
||||
}
|
||||
if self.prompt.trim().is_empty() {
|
||||
return Err("'prompt' must be non-empty".into());
|
||||
}
|
||||
if let Some(d) = self.duration
|
||||
&& !(d.is_finite() && d > 0.0)
|
||||
{
|
||||
return Err(format!("'duration' must be a positive number, got {d}"));
|
||||
}
|
||||
if let Some(f) = self.response_format.as_deref()
|
||||
&& !SUPPORTED_FORMATS.contains(&f)
|
||||
{
|
||||
return Err(format!(
|
||||
"response_format '{f}' is not supported; use one of {SUPPORTED_FORMATS:?}"
|
||||
));
|
||||
}
|
||||
if let Some(p) = self.planner.as_deref()
|
||||
&& !matches!(p, "off" | "auto" | "on")
|
||||
{
|
||||
return Err(format!("planner must be 'off', 'auto' or 'on', got '{p}'"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Validate against a model's advertised limits.
|
||||
pub fn validate_limits(&self, limits: &AudioLimits) -> Result<(), String> {
|
||||
let d = self.duration.unwrap_or(limits.default_duration);
|
||||
if d < limits.min_duration || d > limits.max_duration {
|
||||
return Err(format!(
|
||||
"duration {d}s is outside this model's range {}..={}s",
|
||||
limits.min_duration, limits.max_duration
|
||||
));
|
||||
}
|
||||
if !limits.vocals && !self.is_instrumental() {
|
||||
return Err("this model does not sing; omit 'lyrics' or use [Instrumental]".into());
|
||||
}
|
||||
if let (Some(s), Some(max)) = (self.steps, limits.max_steps)
|
||||
&& (s == 0 || s > max)
|
||||
{
|
||||
return Err(format!("steps must be in 1..={max}, got {s}"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// What a music model can do, advertised in `/v1/models` metadata so
|
||||
/// clients (and the `/music` page) can build their forms from it.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct AudioLimits {
|
||||
pub min_duration: f64,
|
||||
pub max_duration: f64,
|
||||
pub default_duration: f64,
|
||||
/// Whether the model sings (accepts lyrics).
|
||||
pub vocals: bool,
|
||||
pub sample_rate: u32,
|
||||
pub channels: u16,
|
||||
/// Upper bound for `steps`, when the model takes one.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub max_steps: Option<u32>,
|
||||
}
|
||||
|
||||
/// Lifecycle of a job.
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum JobStatus {
|
||||
Queued,
|
||||
Running,
|
||||
Succeeded,
|
||||
Failed,
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
impl JobStatus {
|
||||
pub fn is_terminal(self) -> bool {
|
||||
matches!(self, Self::Succeeded | Self::Failed | Self::Cancelled)
|
||||
}
|
||||
}
|
||||
|
||||
/// The public job object returned by submit, poll, list and cancel.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct AudioJob {
|
||||
pub id: String,
|
||||
/// Always `"audio.generation"`.
|
||||
pub object: String,
|
||||
pub status: JobStatus,
|
||||
pub model: String,
|
||||
/// Unix seconds.
|
||||
pub created_at: u64,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub started_at: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub completed_at: Option<u64>,
|
||||
/// Requested duration, echoed so a client can show progress context.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub duration: Option<f64>,
|
||||
/// Present when `status == succeeded`.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub result: Option<AudioResult>,
|
||||
/// Present when `status == succeeded`.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub usage: Option<AudioUsage>,
|
||||
/// Present when `status == failed`: the #63 envelope's inner `error`.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub error: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
pub const JOB_OBJECT: &str = "audio.generation";
|
||||
|
||||
/// The rendered audio's properties.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct AudioResult {
|
||||
/// Actual length in seconds.
|
||||
pub duration: f64,
|
||||
pub sample_rate: u32,
|
||||
pub channels: u16,
|
||||
/// The seed the render used (random when the request gave none).
|
||||
pub seed: u64,
|
||||
/// Container of the stored audio (`"wav"`).
|
||||
pub format: String,
|
||||
/// Size of the stored audio in bytes.
|
||||
pub bytes: u64,
|
||||
}
|
||||
|
||||
/// Metering block. `helexa_audio_units` is audio-seconds generated (#386).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct AudioUsage {
|
||||
pub helexa_audio_units: f64,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub helexa_timing: Option<AudioTiming>,
|
||||
}
|
||||
|
||||
/// Phase timing for one render, measured on the neuron.
|
||||
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq)]
|
||||
pub struct AudioTiming {
|
||||
/// Text / lyric encoding.
|
||||
pub encode_ms: u64,
|
||||
/// Song planner (0 when it did not run).
|
||||
pub plan_ms: u64,
|
||||
/// Denoising, or autoregressive token generation.
|
||||
pub generate_ms: u64,
|
||||
/// Latent / codec decode to waveform.
|
||||
pub decode_ms: u64,
|
||||
}
|
||||
|
||||
/// What a neuron reports, in [`RESULT_HEADER`], beside the audio bytes.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct NeuronAudioResult {
|
||||
pub duration: f64,
|
||||
pub sample_rate: u32,
|
||||
pub channels: u16,
|
||||
pub seed: u64,
|
||||
pub usage: AudioUsage,
|
||||
}
|
||||
|
||||
/// `GET /v1/audio/generations` response.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AudioJobList {
|
||||
/// Always `"list"`.
|
||||
pub object: String,
|
||||
pub data: Vec<AudioJob>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn req() -> AudioGenerationRequest {
|
||||
serde_json::from_str(r#"{"model": "m", "prompt": "lofi"}"#).unwrap()
|
||||
}
|
||||
|
||||
fn limits() -> AudioLimits {
|
||||
AudioLimits {
|
||||
min_duration: 10.0,
|
||||
max_duration: 600.0,
|
||||
default_duration: 30.0,
|
||||
vocals: true,
|
||||
sample_rate: 48_000,
|
||||
channels: 2,
|
||||
max_steps: Some(20),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn minimal_request_deserializes_and_validates() {
|
||||
let r = req();
|
||||
assert!(r.validate_shape().is_ok());
|
||||
assert!(r.validate_limits(&limits()).is_ok());
|
||||
assert!(r.is_instrumental());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn instrumental_detection() {
|
||||
let mut r = req();
|
||||
r.lyrics = Some(" [Instrumental] ".into());
|
||||
assert!(r.is_instrumental());
|
||||
r.lyrics = Some("[verse]\nla la".into());
|
||||
assert!(!r.is_instrumental());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shape_rejects_bad_values() {
|
||||
let mut r = req();
|
||||
r.prompt = " ".into();
|
||||
assert!(r.validate_shape().is_err());
|
||||
let mut r = req();
|
||||
r.duration = Some(f64::NAN);
|
||||
assert!(r.validate_shape().is_err());
|
||||
let mut r = req();
|
||||
r.duration = Some(-1.0);
|
||||
assert!(r.validate_shape().is_err());
|
||||
let mut r = req();
|
||||
r.response_format = Some("ogg".into());
|
||||
assert!(r.validate_shape().is_err());
|
||||
let mut r = req();
|
||||
r.planner = Some("maybe".into());
|
||||
assert!(r.validate_shape().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn limits_reject_out_of_range() {
|
||||
let l = limits();
|
||||
let mut r = req();
|
||||
r.duration = Some(5.0);
|
||||
assert!(r.validate_limits(&l).is_err());
|
||||
r.duration = Some(601.0);
|
||||
assert!(r.validate_limits(&l).is_err());
|
||||
r.duration = Some(600.0);
|
||||
assert!(r.validate_limits(&l).is_ok());
|
||||
r.steps = Some(21);
|
||||
assert!(r.validate_limits(&l).is_err());
|
||||
r.steps = Some(0);
|
||||
assert!(r.validate_limits(&l).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lyrics_refused_by_non_singing_model() {
|
||||
let mut l = limits();
|
||||
l.vocals = false;
|
||||
let mut r = req();
|
||||
r.lyrics = Some("sing this".into());
|
||||
assert!(r.validate_limits(&l).is_err());
|
||||
r.lyrics = Some("[Instrumental]".into());
|
||||
assert!(r.validate_limits(&l).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn job_status_serializes_snake_case() {
|
||||
assert_eq!(
|
||||
serde_json::to_string(&JobStatus::Succeeded).unwrap(),
|
||||
"\"succeeded\""
|
||||
);
|
||||
assert!(JobStatus::Cancelled.is_terminal());
|
||||
assert!(!JobStatus::Running.is_terminal());
|
||||
}
|
||||
}
|
||||
@@ -89,6 +89,11 @@ pub struct ModelProfile {
|
||||
/// are unioned with this set in the gateway's `/v1/models` response.
|
||||
#[serde(default)]
|
||||
pub capabilities: Vec<String>,
|
||||
/// What an audio model can do (#382): duration range, vocals, sample
|
||||
/// rate. Advertised in `/v1/models` and checked by cortex before a job
|
||||
/// is accepted. `[models.audio]` in TOML.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub audio: Option<crate::audio::AudioLimits>,
|
||||
}
|
||||
|
||||
fn default_min_devices() -> u32 {
|
||||
@@ -283,6 +288,7 @@ mod tests {
|
||||
limit: None,
|
||||
cost: None,
|
||||
capabilities: vec![],
|
||||
audio: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -28,6 +28,63 @@ pub struct GatewayConfig {
|
||||
/// — a single operator runs purely local.
|
||||
#[serde(default)]
|
||||
pub upstream: UpstreamClientConfig,
|
||||
/// Async music/audio generation jobs (#382).
|
||||
#[serde(default)]
|
||||
pub audio_jobs: AudioJobsConfig,
|
||||
}
|
||||
|
||||
/// `[audio_jobs]` — where cortex keeps async audio generation jobs (#382).
|
||||
///
|
||||
/// A job record and its rendered audio are written under `dir`, so finished
|
||||
/// results and the job list survive a cortex restart. If `dir` cannot be
|
||||
/// created the store runs in memory only and says so at startup.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AudioJobsConfig {
|
||||
#[serde(default = "default_audio_jobs_dir")]
|
||||
pub dir: String,
|
||||
/// How long a finished job and its audio are kept, in hours.
|
||||
#[serde(default = "default_audio_retention_hours")]
|
||||
pub retention_hours: u64,
|
||||
/// How many renders cortex runs at once across the fleet. Each one
|
||||
/// occupies a neuron device for its whole duration; queued jobs wait
|
||||
/// in cortex rather than piling onto neuron admission queues.
|
||||
#[serde(default = "default_audio_max_concurrent")]
|
||||
pub max_concurrent: usize,
|
||||
/// Dispatch attempts per job before it fails. A busy neuron (429/503)
|
||||
/// does not consume an attempt; an unreachable or failing one does.
|
||||
#[serde(default = "default_audio_max_attempts")]
|
||||
pub max_attempts: u32,
|
||||
/// How long a job may wait for a busy model before it fails, in seconds.
|
||||
#[serde(default = "default_audio_queue_timeout")]
|
||||
pub queue_timeout_secs: u64,
|
||||
}
|
||||
|
||||
impl Default for AudioJobsConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
dir: default_audio_jobs_dir(),
|
||||
retention_hours: default_audio_retention_hours(),
|
||||
max_concurrent: default_audio_max_concurrent(),
|
||||
max_attempts: default_audio_max_attempts(),
|
||||
queue_timeout_secs: default_audio_queue_timeout(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn default_audio_jobs_dir() -> String {
|
||||
"/var/lib/cortex/audio-jobs".into()
|
||||
}
|
||||
fn default_audio_retention_hours() -> u64 {
|
||||
72
|
||||
}
|
||||
fn default_audio_max_concurrent() -> usize {
|
||||
4
|
||||
}
|
||||
fn default_audio_max_attempts() -> u32 {
|
||||
2
|
||||
}
|
||||
fn default_audio_queue_timeout() -> u64 {
|
||||
3600
|
||||
}
|
||||
|
||||
/// `[upstream]` — the helexa-upstream authority client (#57). Locally
|
||||
@@ -200,6 +257,7 @@ impl Default for GatewayConfig {
|
||||
models_config: default_models_path(),
|
||||
entitlements: EntitlementsConfig::default(),
|
||||
upstream: UpstreamClientConfig::default(),
|
||||
audio_jobs: AudioJobsConfig::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,7 +1,9 @@
|
||||
pub mod anthropic;
|
||||
pub mod audio;
|
||||
pub mod build_info;
|
||||
pub mod catalogue;
|
||||
pub mod config;
|
||||
pub mod decisions;
|
||||
pub mod discovery;
|
||||
pub mod entitlements;
|
||||
pub mod error_envelope;
|
||||
|
||||
@@ -186,6 +186,14 @@ pub struct CortexModelEntry {
|
||||
/// the absent (not priced) vs `0.0` (intentionally free) distinction.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub cost: Option<ModelCost>,
|
||||
/// What an audio model can do (#382), from its catalogue profile:
|
||||
/// duration range, vocals, sample rate. Clients build forms from it.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub audio: Option<crate::audio::AudioLimits>,
|
||||
/// Ledger tokens charged per second of generated audio (#386), for
|
||||
/// audio models with a configured rate.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub audio_tokens_per_second: Option<u64>,
|
||||
/// `true` when any neuron reports this model supports tool calls.
|
||||
#[serde(default)]
|
||||
pub tool_call: bool,
|
||||
|
||||
@@ -27,6 +27,8 @@ eventsource-stream.workspace = true
|
||||
bytes = "1"
|
||||
urlencoding = "2"
|
||||
url = "2"
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
|
||||
[dev-dependencies]
|
||||
tokio = { workspace = true, features = ["test-util"] }
|
||||
tempfile = "3"
|
||||
|
||||
@@ -0,0 +1,884 @@
|
||||
//! Async music / audio generation jobs (#382).
|
||||
//!
|
||||
//! A render takes seconds to minutes, so the public surface is a job API
|
||||
//! (see [`cortex_core::audio`] for the routes). cortex owns the job: it
|
||||
//! records it, queues it, dispatches it to a neuron serving the model, keeps
|
||||
//! the rendered audio, and meters it. A neuron renders one generation per
|
||||
//! dispatch, synchronously, and knows nothing about jobs.
|
||||
//!
|
||||
//! Records and audio are written under `[audio_jobs] dir`, so the job list
|
||||
//! and finished results survive a cortex restart. Jobs that were queued or
|
||||
//! running when cortex stopped are re-queued at startup. A job whose neuron
|
||||
//! fails is retried on whatever node the router picks next, up to
|
||||
//! `max_attempts`; a neuron that is merely busy (429/503) is waited on, up
|
||||
//! to `queue_timeout_secs`, without spending an attempt.
|
||||
//!
|
||||
//! Ownership: a job belongs to the account that submitted it, and only that
|
||||
//! account can see, fetch or cancel it. Anonymous jobs (no principal) are
|
||||
//! reachable by id alone — the id is an unguessable capability — and never
|
||||
//! appear in a listing.
|
||||
|
||||
use crate::metering::ReservationGuard;
|
||||
use crate::state::CortexState;
|
||||
use axum::Router;
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::extract::{Path, State};
|
||||
use axum::http::{HeaderMap, HeaderValue, StatusCode, header};
|
||||
use axum::response::{IntoResponse, Json, Response};
|
||||
use axum::routing::get;
|
||||
use cortex_core::audio::{
|
||||
AudioGenerationRequest, AudioJob, AudioJobList, AudioResult, JOB_OBJECT, JobStatus,
|
||||
NeuronAudioResult, RESULT_HEADER,
|
||||
};
|
||||
use cortex_core::config::AudioJobsConfig;
|
||||
use cortex_core::entitlements::{HEADER_ACCOUNT_ID, HEADER_KEY_ID, Principal};
|
||||
use cortex_core::error_envelope::OpenAiError;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
|
||||
use tokio::sync::{RwLock, Semaphore};
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
/// Path of the job collection; `{id}` and `{id}/content` hang off it.
|
||||
pub const JOBS_PATH: &str = "/v1/audio/generations";
|
||||
|
||||
/// Wait between dispatches to a busy model when the neuron sends no
|
||||
/// `Retry-After`.
|
||||
const DEFAULT_BUSY_WAIT: Duration = Duration::from_secs(5);
|
||||
|
||||
/// How often expired jobs are swept.
|
||||
const SWEEP_INTERVAL: Duration = Duration::from_secs(600);
|
||||
|
||||
/// A job as cortex stores it: the public object plus what is needed to run
|
||||
/// (or re-run) it.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
struct JobRecord {
|
||||
job: AudioJob,
|
||||
request: AudioGenerationRequest,
|
||||
/// The submitting principal; `None` for an anonymous submission.
|
||||
owner: Option<Principal>,
|
||||
/// Ledger tokens per audio-second this job was reserved at (#386), so
|
||||
/// it settles at the price it was accepted at even if the rate changes
|
||||
/// while it runs. `None` for anonymous (unmetered) jobs.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
tokens_per_second: Option<u64>,
|
||||
}
|
||||
|
||||
/// The job store and dispatcher, held in [`CortexState`].
|
||||
pub struct AudioJobs {
|
||||
config: AudioJobsConfig,
|
||||
/// See [`AudioJobs::dir`].
|
||||
dir: std::sync::OnceLock<Option<PathBuf>>,
|
||||
jobs: RwLock<HashMap<String, JobRecord>>,
|
||||
/// Worker tasks for jobs that are queued or running, for cancellation.
|
||||
workers: Mutex<HashMap<String, JoinHandle<()>>>,
|
||||
/// Bounds concurrent renders across the fleet.
|
||||
slots: Arc<Semaphore>,
|
||||
/// Rendered audio, when there is no job directory to write it to.
|
||||
memory_audio: Mutex<HashMap<String, Bytes>>,
|
||||
/// Client for neuron dispatches. A neuron sends nothing until the render
|
||||
/// is finished, so the read timeout has to cover a whole render.
|
||||
http: reqwest::Client,
|
||||
}
|
||||
|
||||
impl AudioJobs {
|
||||
pub fn new(config: &AudioJobsConfig) -> Self {
|
||||
let http = reqwest::Client::builder()
|
||||
.read_timeout(Duration::from_secs(config.queue_timeout_secs.max(1800)))
|
||||
.build()
|
||||
.expect("failed to build audio HTTP client");
|
||||
Self {
|
||||
config: config.clone(),
|
||||
dir: std::sync::OnceLock::new(),
|
||||
jobs: RwLock::new(HashMap::new()),
|
||||
workers: Mutex::new(HashMap::new()),
|
||||
slots: Arc::new(Semaphore::new(config.max_concurrent.max(1))),
|
||||
memory_audio: Mutex::new(HashMap::new()),
|
||||
http,
|
||||
}
|
||||
}
|
||||
|
||||
/// The job directory, created on first use. `None` (memory only) when it
|
||||
/// cannot be created. Lazy so that building a gateway does not touch the
|
||||
/// filesystem until audio jobs are actually used.
|
||||
fn dir(&self) -> Option<&PathBuf> {
|
||||
self.dir
|
||||
.get_or_init(|| {
|
||||
let dir = PathBuf::from(&self.config.dir);
|
||||
match std::fs::create_dir_all(&dir) {
|
||||
Ok(()) => Some(dir),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
dir = %self.config.dir,
|
||||
error = %e,
|
||||
"audio job directory unavailable; jobs will not survive a restart"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
})
|
||||
.as_ref()
|
||||
}
|
||||
|
||||
fn record_path(&self, id: &str) -> Option<PathBuf> {
|
||||
self.dir().map(|d| d.join(format!("{id}.json")))
|
||||
}
|
||||
|
||||
fn audio_path(&self, id: &str) -> Option<PathBuf> {
|
||||
self.dir().map(|d| d.join(format!("{id}.wav")))
|
||||
}
|
||||
|
||||
async fn persist(&self, record: &JobRecord) {
|
||||
let Some(path) = self.record_path(&record.job.id) else {
|
||||
return;
|
||||
};
|
||||
let bytes = match serde_json::to_vec_pretty(record) {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
tracing::warn!(job = %record.job.id, error = %e, "audio job serialise failed");
|
||||
return;
|
||||
}
|
||||
};
|
||||
// Write-then-rename so a crash never leaves a torn record.
|
||||
let tmp = path.with_extension("json.tmp");
|
||||
if let Err(e) = tokio::fs::write(&tmp, &bytes).await {
|
||||
tracing::warn!(job = %record.job.id, error = %e, "audio job write failed");
|
||||
return;
|
||||
}
|
||||
if let Err(e) = tokio::fs::rename(&tmp, &path).await {
|
||||
tracing::warn!(job = %record.job.id, error = %e, "audio job rename failed");
|
||||
}
|
||||
}
|
||||
|
||||
async fn update<F: FnOnce(&mut JobRecord)>(&self, id: &str, f: F) -> Option<JobRecord> {
|
||||
let snapshot = {
|
||||
let mut jobs = self.jobs.write().await;
|
||||
let record = jobs.get_mut(id)?;
|
||||
f(record);
|
||||
record.clone()
|
||||
};
|
||||
self.persist(&snapshot).await;
|
||||
Some(snapshot)
|
||||
}
|
||||
|
||||
async fn get(&self, id: &str) -> Option<JobRecord> {
|
||||
self.jobs.read().await.get(id).cloned()
|
||||
}
|
||||
|
||||
fn forget_worker(&self, id: &str) {
|
||||
if let Ok(mut w) = self.workers.lock() {
|
||||
w.remove(id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Load persisted jobs, re-queue unfinished ones, and start the retention
|
||||
/// sweep. Called once at startup.
|
||||
pub async fn start(fleet: Arc<CortexState>) {
|
||||
let Some(dir) = fleet.audio_jobs.dir().cloned() else {
|
||||
spawn_sweeper(fleet);
|
||||
return;
|
||||
};
|
||||
let mut resumed = Vec::new();
|
||||
if let Ok(mut entries) = tokio::fs::read_dir(&dir).await {
|
||||
while let Ok(Some(entry)) = entries.next_entry().await {
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|e| e.to_str()) != Some("json") {
|
||||
continue;
|
||||
}
|
||||
let record: JobRecord = match tokio::fs::read(&path)
|
||||
.await
|
||||
.ok()
|
||||
.and_then(|b| serde_json::from_slice(&b).ok())
|
||||
{
|
||||
Some(r) => r,
|
||||
None => {
|
||||
tracing::warn!(path = %path.display(), "skipping unreadable audio job record");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
if !record.job.status.is_terminal() {
|
||||
resumed.push(record.job.id.clone());
|
||||
}
|
||||
fleet
|
||||
.audio_jobs
|
||||
.jobs
|
||||
.write()
|
||||
.await
|
||||
.insert(record.job.id.clone(), record);
|
||||
}
|
||||
}
|
||||
for id in resumed {
|
||||
let Some(record) = fleet
|
||||
.audio_jobs
|
||||
.update(&id, |r| {
|
||||
r.job.status = JobStatus::Queued;
|
||||
r.job.started_at = None;
|
||||
})
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
// The reservation died with the old process: take a fresh one, and
|
||||
// fail the job if the budget no longer allows it.
|
||||
let guard = match reserve(&fleet, record.owner.as_ref(), &record.request).await {
|
||||
Ok((g, rate)) => {
|
||||
fleet
|
||||
.audio_jobs
|
||||
.update(&id, |r| r.tokens_per_second = rate)
|
||||
.await;
|
||||
g
|
||||
}
|
||||
Err(env) => {
|
||||
fail(&fleet, &id, envelope_error(&env)).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
tracing::info!(job = %id, "re-queued audio job after restart");
|
||||
spawn_worker(Arc::clone(&fleet), id, guard);
|
||||
}
|
||||
spawn_sweeper(fleet);
|
||||
}
|
||||
|
||||
fn spawn_sweeper(fleet: Arc<CortexState>) {
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::time::sleep(SWEEP_INTERVAL).await;
|
||||
sweep_expired(&fleet).await;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Delete finished jobs older than the retention window, with their audio.
|
||||
pub async fn sweep_expired(fleet: &CortexState) {
|
||||
let jobs = &fleet.audio_jobs;
|
||||
let cutoff = now_secs().saturating_sub(jobs.config.retention_hours * 3600);
|
||||
let expired: Vec<String> = jobs
|
||||
.jobs
|
||||
.read()
|
||||
.await
|
||||
.values()
|
||||
.filter(|r| r.job.status.is_terminal())
|
||||
.filter(|r| r.job.completed_at.unwrap_or(r.job.created_at) < cutoff)
|
||||
.map(|r| r.job.id.clone())
|
||||
.collect();
|
||||
for id in expired {
|
||||
jobs.jobs.write().await.remove(&id);
|
||||
if let Ok(mut m) = jobs.memory_audio.lock() {
|
||||
m.remove(&id);
|
||||
}
|
||||
for path in [jobs.record_path(&id), jobs.audio_path(&id)]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
{
|
||||
let _ = tokio::fs::remove_file(path).await;
|
||||
}
|
||||
tracing::debug!(job = %id, "expired audio job removed");
|
||||
}
|
||||
}
|
||||
|
||||
/// Routes for the job API.
|
||||
pub fn routes() -> Router<Arc<CortexState>> {
|
||||
Router::new()
|
||||
.route(JOBS_PATH, get(list_jobs).post(submit))
|
||||
.route(
|
||||
"/v1/audio/generations/{id}",
|
||||
get(get_job).delete(cancel_job),
|
||||
)
|
||||
.route("/v1/audio/generations/{id}/content", get(get_content))
|
||||
}
|
||||
|
||||
/// `POST /v1/audio/generations` — validate, reserve budget, queue.
|
||||
async fn submit(
|
||||
State(fleet): State<Arc<CortexState>>,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Response {
|
||||
let request: AudioGenerationRequest = match serde_json::from_slice(&body) {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
return invalid_params(format!("request body is not a valid audio request: {e}"));
|
||||
}
|
||||
};
|
||||
if let Err(msg) = request.validate_shape() {
|
||||
return invalid_params(msg);
|
||||
}
|
||||
|
||||
// The model must be one this fleet knows. Resolving it fully would
|
||||
// cold-load it at submit time, so only check that it exists here; the
|
||||
// worker resolves (and loads) when the job actually runs.
|
||||
let model = fleet.catalogue.resolve_alias(&request.model).to_string();
|
||||
let known = fleet.catalogue.get(&model).is_some()
|
||||
|| fleet
|
||||
.nodes
|
||||
.read()
|
||||
.await
|
||||
.values()
|
||||
.any(|n| n.models.contains_key(&model));
|
||||
if !known {
|
||||
return crate::error::envelope_response(OpenAiError::new(
|
||||
404,
|
||||
"invalid_request_error",
|
||||
"model_not_found",
|
||||
format!(
|
||||
"model '{}' not found on any node and not in catalogue",
|
||||
request.model
|
||||
),
|
||||
));
|
||||
}
|
||||
if let Some(profile) = fleet.catalogue.get(&model)
|
||||
&& let Some(limits) = &profile.audio
|
||||
&& let Err(msg) = request.validate_limits(limits)
|
||||
{
|
||||
return invalid_params(msg);
|
||||
}
|
||||
|
||||
let owner = crate::metering::principal_from_headers(&headers);
|
||||
let mut request = request;
|
||||
request.model = model.clone();
|
||||
let (guard, tokens_per_second) = match reserve(&fleet, owner.as_ref(), &request).await {
|
||||
Ok(r) => r,
|
||||
Err(env) => return crate::error::envelope_response(env),
|
||||
};
|
||||
|
||||
let id = format!("audiogen_{}", uuid::Uuid::new_v4().simple());
|
||||
let job = AudioJob {
|
||||
id: id.clone(),
|
||||
object: JOB_OBJECT.into(),
|
||||
status: JobStatus::Queued,
|
||||
model,
|
||||
created_at: now_secs(),
|
||||
started_at: None,
|
||||
completed_at: None,
|
||||
duration: request.duration,
|
||||
result: None,
|
||||
usage: None,
|
||||
error: None,
|
||||
};
|
||||
let record = JobRecord {
|
||||
job: job.clone(),
|
||||
request,
|
||||
owner,
|
||||
tokens_per_second,
|
||||
};
|
||||
fleet
|
||||
.audio_jobs
|
||||
.jobs
|
||||
.write()
|
||||
.await
|
||||
.insert(id.clone(), record.clone());
|
||||
fleet.audio_jobs.persist(&record).await;
|
||||
metrics::counter!("cortex_audio_jobs_submitted_total", "model" => job.model.clone())
|
||||
.increment(1);
|
||||
tracing::info!(job = %id, model = %job.model, "audio job queued");
|
||||
|
||||
spawn_worker(Arc::clone(&fleet), id, guard);
|
||||
(StatusCode::ACCEPTED, Json(job)).into_response()
|
||||
}
|
||||
|
||||
/// `GET /v1/audio/generations` — the caller's jobs, newest first.
|
||||
async fn list_jobs(State(fleet): State<Arc<CortexState>>, headers: HeaderMap) -> Response {
|
||||
let caller = crate::metering::principal_from_headers(&headers);
|
||||
let mut data: Vec<AudioJob> = match &caller {
|
||||
None => Vec::new(),
|
||||
Some(p) => fleet
|
||||
.audio_jobs
|
||||
.jobs
|
||||
.read()
|
||||
.await
|
||||
.values()
|
||||
.filter(|r| {
|
||||
r.owner
|
||||
.as_ref()
|
||||
.is_some_and(|o| o.account_id == p.account_id)
|
||||
})
|
||||
.map(|r| r.job.clone())
|
||||
.collect(),
|
||||
};
|
||||
data.sort_by(|a, b| b.created_at.cmp(&a.created_at).then(b.id.cmp(&a.id)));
|
||||
Json(AudioJobList {
|
||||
object: "list".into(),
|
||||
data,
|
||||
})
|
||||
.into_response()
|
||||
}
|
||||
|
||||
/// `GET /v1/audio/generations/{id}`.
|
||||
async fn get_job(
|
||||
State(fleet): State<Arc<CortexState>>,
|
||||
headers: HeaderMap,
|
||||
Path(id): Path<String>,
|
||||
) -> Response {
|
||||
match visible_record(&fleet, &headers, &id).await {
|
||||
Some(r) => Json(r.job).into_response(),
|
||||
None => job_not_found(&id),
|
||||
}
|
||||
}
|
||||
|
||||
/// `GET /v1/audio/generations/{id}/content` — the rendered audio.
|
||||
async fn get_content(
|
||||
State(fleet): State<Arc<CortexState>>,
|
||||
headers: HeaderMap,
|
||||
Path(id): Path<String>,
|
||||
) -> Response {
|
||||
let Some(record) = visible_record(&fleet, &headers, &id).await else {
|
||||
return job_not_found(&id);
|
||||
};
|
||||
if record.job.status != JobStatus::Succeeded {
|
||||
return crate::error::envelope_response(OpenAiError::new(
|
||||
409,
|
||||
"invalid_request_error",
|
||||
"job_not_complete",
|
||||
format!(
|
||||
"audio job '{id}' is {}; content is available once it has succeeded",
|
||||
status_word(record.job.status)
|
||||
),
|
||||
));
|
||||
}
|
||||
let bytes = match fleet.audio_jobs.audio_path(&id) {
|
||||
Some(path) => tokio::fs::read(&path).await.ok(),
|
||||
None => fleet.audio_jobs.memory_audio(&id),
|
||||
};
|
||||
let Some(bytes) = bytes else {
|
||||
return crate::error::envelope_response(OpenAiError::new(
|
||||
410,
|
||||
"invalid_request_error",
|
||||
"content_expired",
|
||||
format!("the audio for job '{id}' is no longer stored"),
|
||||
));
|
||||
};
|
||||
let mut resp = Response::new(Body::from(bytes));
|
||||
resp.headers_mut()
|
||||
.insert(header::CONTENT_TYPE, HeaderValue::from_static("audio/wav"));
|
||||
if let Ok(v) = HeaderValue::from_str(&format!("attachment; filename=\"{id}.wav\"")) {
|
||||
resp.headers_mut().insert(header::CONTENT_DISPOSITION, v);
|
||||
}
|
||||
resp
|
||||
}
|
||||
|
||||
/// `DELETE /v1/audio/generations/{id}` — cancel a queued or running job.
|
||||
/// Cancelling a finished job is a no-op that returns it unchanged.
|
||||
async fn cancel_job(
|
||||
State(fleet): State<Arc<CortexState>>,
|
||||
headers: HeaderMap,
|
||||
Path(id): Path<String>,
|
||||
) -> Response {
|
||||
let Some(record) = visible_record(&fleet, &headers, &id).await else {
|
||||
return job_not_found(&id);
|
||||
};
|
||||
if record.job.status.is_terminal() {
|
||||
return Json(record.job).into_response();
|
||||
}
|
||||
// Aborting the worker drops its reservation guard, which releases the
|
||||
// budget. A render already running on a neuron finishes there unseen.
|
||||
let handle = fleet
|
||||
.audio_jobs
|
||||
.workers
|
||||
.lock()
|
||||
.ok()
|
||||
.and_then(|mut w| w.remove(&id));
|
||||
if let Some(h) = handle {
|
||||
h.abort();
|
||||
}
|
||||
let updated = fleet
|
||||
.audio_jobs
|
||||
.update(&id, |r| {
|
||||
r.job.status = JobStatus::Cancelled;
|
||||
r.job.completed_at = Some(now_secs());
|
||||
})
|
||||
.await;
|
||||
metrics::counter!("cortex_audio_jobs_cancelled_total", "model" => record.job.model.clone())
|
||||
.increment(1);
|
||||
tracing::info!(job = %id, "audio job cancelled");
|
||||
match updated {
|
||||
Some(r) => Json(r.job).into_response(),
|
||||
None => job_not_found(&id),
|
||||
}
|
||||
}
|
||||
|
||||
/// A job the caller may see: owned by the caller's account, or anonymous.
|
||||
async fn visible_record(fleet: &CortexState, headers: &HeaderMap, id: &str) -> Option<JobRecord> {
|
||||
let record = fleet.audio_jobs.get(id).await?;
|
||||
let caller = crate::metering::principal_from_headers(headers);
|
||||
match (&record.owner, &caller) {
|
||||
(None, _) => Some(record),
|
||||
(Some(owner), Some(caller)) if owner.account_id == caller.account_id => Some(record),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn spawn_worker(fleet: Arc<CortexState>, id: String, guard: ReservationGuard) {
|
||||
let jobs_fleet = Arc::clone(&fleet);
|
||||
let task_id = id.clone();
|
||||
let handle = tokio::spawn(async move {
|
||||
run_job(&jobs_fleet, &task_id, guard).await;
|
||||
jobs_fleet.audio_jobs.forget_worker(&task_id);
|
||||
});
|
||||
if let Ok(mut w) = fleet.audio_jobs.workers.lock() {
|
||||
// The task may already have finished (and tried to forget itself)
|
||||
// before this insert; only track it while it is still running.
|
||||
if !handle.is_finished() {
|
||||
w.insert(id, handle);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Outcome of one dispatch to a neuron.
|
||||
enum Dispatch {
|
||||
Done(Bytes, NeuronAudioResult),
|
||||
/// The model is busy; try again after the given wait.
|
||||
Busy(Duration),
|
||||
/// The request itself was refused; retrying will not help.
|
||||
Rejected(Value),
|
||||
/// The node failed or was unreachable; another attempt may succeed.
|
||||
Failed(Value),
|
||||
}
|
||||
|
||||
async fn run_job(fleet: &Arc<CortexState>, id: &str, guard: ReservationGuard) {
|
||||
let jobs = &fleet.audio_jobs;
|
||||
let Ok(_slot) = Arc::clone(&jobs.slots).acquire_owned().await else {
|
||||
return;
|
||||
};
|
||||
let Some(record) = jobs
|
||||
.update(id, |r| {
|
||||
if !r.job.status.is_terminal() {
|
||||
r.job.status = JobStatus::Running;
|
||||
r.job.started_at = Some(now_secs());
|
||||
}
|
||||
})
|
||||
.await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
// Cancelled while it waited for a slot.
|
||||
if record.job.status != JobStatus::Running {
|
||||
return;
|
||||
}
|
||||
let model = record.job.model.clone();
|
||||
let deadline = Instant::now() + Duration::from_secs(jobs.config.queue_timeout_secs);
|
||||
let mut attempts = 0u32;
|
||||
let mut last_error = error_value(
|
||||
"api_error",
|
||||
"audio_generation_failed",
|
||||
"the audio job could not be dispatched",
|
||||
);
|
||||
|
||||
while attempts < jobs.config.max_attempts.max(1) {
|
||||
let outcome = dispatch_once(fleet, &record).await;
|
||||
match outcome {
|
||||
Dispatch::Done(bytes, result) => {
|
||||
complete(fleet, id, &record, bytes, result, guard).await;
|
||||
return;
|
||||
}
|
||||
Dispatch::Busy(wait) => {
|
||||
if Instant::now() + wait > deadline {
|
||||
last_error = error_value(
|
||||
"rate_limit_error",
|
||||
"queue_timeout",
|
||||
"the model stayed busy for longer than the job queue timeout",
|
||||
);
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(wait).await;
|
||||
}
|
||||
Dispatch::Rejected(err) => {
|
||||
last_error = err;
|
||||
break;
|
||||
}
|
||||
Dispatch::Failed(err) => {
|
||||
attempts += 1;
|
||||
tracing::warn!(job = %id, model = %model, attempt = attempts, error = %err, "audio dispatch failed");
|
||||
last_error = err;
|
||||
if attempts < jobs.config.max_attempts {
|
||||
tokio::time::sleep(Duration::from_secs(2)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
drop(guard); // releases the reservation: nothing was rendered
|
||||
fail(fleet, id, last_error).await;
|
||||
}
|
||||
|
||||
async fn dispatch_once(fleet: &Arc<CortexState>, record: &JobRecord) -> Dispatch {
|
||||
let route = match crate::router::resolve(fleet, &record.job.model).await {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
let err = error_value(e.broad_type(), e.code(), &e.to_string());
|
||||
return match e.retry_after_secs() {
|
||||
Some(secs) => Dispatch::Busy(Duration::from_secs(secs.max(1))),
|
||||
None if e.http_status() < 500 => Dispatch::Rejected(err),
|
||||
None => Dispatch::Failed(err),
|
||||
};
|
||||
}
|
||||
};
|
||||
crate::handlers::touch_model(fleet, &route.node_name, &route.resolved_model_id).await;
|
||||
|
||||
let mut body = record.request.clone();
|
||||
body.model = route.resolved_model_id.clone();
|
||||
let url = format!("{}{JOBS_PATH}", route.endpoint);
|
||||
let mut req = fleet.audio_jobs.http.post(&url).json(&body);
|
||||
if let Some(owner) = &record.owner {
|
||||
req = req
|
||||
.header(HEADER_ACCOUNT_ID, &owner.account_id)
|
||||
.header(HEADER_KEY_ID, &owner.key_id);
|
||||
}
|
||||
let labels = [
|
||||
("model", route.resolved_model_id.clone()),
|
||||
("node", route.node_name.clone()),
|
||||
];
|
||||
metrics::counter!("cortex_requests_total", &labels).increment(1);
|
||||
if route.cold_start {
|
||||
metrics::counter!("cortex_cold_starts_total", &labels).increment(1);
|
||||
}
|
||||
let start = Instant::now();
|
||||
let resp = match req.send().await {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
metrics::counter!("cortex_request_errors_total", &labels).increment(1);
|
||||
return Dispatch::Failed(error_value(
|
||||
"api_error",
|
||||
"upstream_unreachable",
|
||||
&format!("failed to reach the serving node: {e}"),
|
||||
));
|
||||
}
|
||||
};
|
||||
let status = resp.status();
|
||||
let retry_after = resp
|
||||
.headers()
|
||||
.get(header::RETRY_AFTER)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|v| v.trim().parse::<u64>().ok())
|
||||
.map(|s| Duration::from_secs(s.max(1)));
|
||||
let result_header = resp
|
||||
.headers()
|
||||
.get(RESULT_HEADER)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(str::to_string);
|
||||
let bytes = match resp.bytes().await {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
metrics::counter!("cortex_request_errors_total", &labels).increment(1);
|
||||
return Dispatch::Failed(error_value(
|
||||
"api_error",
|
||||
"upstream_read_failed",
|
||||
&format!("failed to read the serving node's response: {e}"),
|
||||
));
|
||||
}
|
||||
};
|
||||
metrics::histogram!("cortex_audio_render_seconds", &labels)
|
||||
.record(start.elapsed().as_secs_f64());
|
||||
|
||||
if status.is_success() {
|
||||
let parsed = result_header
|
||||
.as_deref()
|
||||
.and_then(|h| serde_json::from_str::<NeuronAudioResult>(h).ok());
|
||||
return match parsed {
|
||||
Some(result) => Dispatch::Done(bytes, result),
|
||||
None => Dispatch::Failed(error_value(
|
||||
"api_error",
|
||||
"upstream_bad_response",
|
||||
"the serving node returned audio without its result metadata",
|
||||
)),
|
||||
};
|
||||
}
|
||||
metrics::counter!("cortex_request_errors_total", &labels).increment(1);
|
||||
if status.as_u16() == 429 || status.as_u16() == 503 {
|
||||
return Dispatch::Busy(retry_after.unwrap_or(DEFAULT_BUSY_WAIT));
|
||||
}
|
||||
let err = serde_json::from_slice::<Value>(&bytes)
|
||||
.ok()
|
||||
.and_then(|v| v.get("error").cloned())
|
||||
.unwrap_or_else(|| {
|
||||
error_value(
|
||||
"api_error",
|
||||
"upstream_error",
|
||||
&format!("the serving node answered {status}"),
|
||||
)
|
||||
});
|
||||
if status.is_client_error() {
|
||||
Dispatch::Rejected(err)
|
||||
} else {
|
||||
Dispatch::Failed(err)
|
||||
}
|
||||
}
|
||||
|
||||
async fn complete(
|
||||
fleet: &Arc<CortexState>,
|
||||
id: &str,
|
||||
record: &JobRecord,
|
||||
bytes: Bytes,
|
||||
result: NeuronAudioResult,
|
||||
guard: ReservationGuard,
|
||||
) {
|
||||
let jobs = &fleet.audio_jobs;
|
||||
// Cancelled while the render was in flight: drop the result, release
|
||||
// the budget, and leave the job cancelled.
|
||||
if jobs.get(id).await.map(|r| r.job.status) == Some(JobStatus::Cancelled) {
|
||||
drop(guard);
|
||||
return;
|
||||
}
|
||||
let size = bytes.len() as u64;
|
||||
match jobs.audio_path(id) {
|
||||
Some(path) => {
|
||||
if let Err(e) = tokio::fs::write(&path, &bytes).await {
|
||||
drop(guard);
|
||||
fail(
|
||||
fleet,
|
||||
id,
|
||||
error_value(
|
||||
"api_error",
|
||||
"storage_failed",
|
||||
&format!("could not store the rendered audio: {e}"),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
None => jobs.keep_in_memory(id, bytes),
|
||||
}
|
||||
|
||||
let units = result.usage.helexa_audio_units;
|
||||
match &record.owner {
|
||||
Some(owner) => {
|
||||
// An owned job always carries the rate it was reserved at.
|
||||
let rate = record
|
||||
.tokens_per_second
|
||||
.unwrap_or(crate::rates::LOCAL_DEFAULT_TOKENS_PER_AUDIO_SECOND);
|
||||
crate::metering::settle_audio_usage(owner, guard, &fleet.served_usage, units, rate)
|
||||
}
|
||||
None => drop(guard),
|
||||
}
|
||||
jobs.update(id, |r| {
|
||||
r.job.status = JobStatus::Succeeded;
|
||||
r.job.completed_at = Some(now_secs());
|
||||
r.job.result = Some(AudioResult {
|
||||
duration: result.duration,
|
||||
sample_rate: result.sample_rate,
|
||||
channels: result.channels,
|
||||
seed: result.seed,
|
||||
format: "wav".into(),
|
||||
bytes: size,
|
||||
});
|
||||
r.job.usage = Some(result.usage.clone());
|
||||
})
|
||||
.await;
|
||||
metrics::counter!("cortex_audio_jobs_succeeded_total", "model" => record.job.model.clone())
|
||||
.increment(1);
|
||||
tracing::info!(job = %id, model = %record.job.model, seconds = result.duration, "audio job succeeded");
|
||||
}
|
||||
|
||||
async fn fail(fleet: &CortexState, id: &str, error: Value) {
|
||||
let updated = fleet
|
||||
.audio_jobs
|
||||
.update(id, |r| {
|
||||
r.job.status = JobStatus::Failed;
|
||||
r.job.completed_at = Some(now_secs());
|
||||
r.job.error = Some(error.clone());
|
||||
})
|
||||
.await;
|
||||
if let Some(r) = updated {
|
||||
metrics::counter!("cortex_audio_jobs_failed_total", "model" => r.job.model.clone())
|
||||
.increment(1);
|
||||
tracing::warn!(job = %id, model = %r.job.model, error = %error, "audio job failed");
|
||||
}
|
||||
}
|
||||
|
||||
/// Reserve budget for a job at its model's current rate (#386). Returns the
|
||||
/// guard and the rate it reserved at; anonymous jobs are unmetered.
|
||||
#[allow(clippy::result_large_err)] // the Err is the #63 envelope; see metering::reserve_or_reject
|
||||
async fn reserve(
|
||||
fleet: &Arc<CortexState>,
|
||||
owner: Option<&Principal>,
|
||||
request: &AudioGenerationRequest,
|
||||
) -> Result<(ReservationGuard, Option<u64>), OpenAiError> {
|
||||
let Some(p) = owner else {
|
||||
return Ok((
|
||||
ReservationGuard::anonymous(Arc::clone(&fleet.entitlements)),
|
||||
None,
|
||||
));
|
||||
};
|
||||
let Some(rate) = fleet.rates.audio_tokens_per_second(&request.model).await else {
|
||||
return Err(OpenAiError::service_unavailable(
|
||||
format!("no audio rate is configured for model '{}'", request.model),
|
||||
Some(5),
|
||||
));
|
||||
};
|
||||
let default = fleet
|
||||
.catalogue
|
||||
.get(&request.model)
|
||||
.and_then(|m| m.audio.as_ref())
|
||||
.map(|l| l.default_duration);
|
||||
let seconds = crate::metering::audio_reservation_seconds(request.duration, default);
|
||||
let guard = crate::metering::reserve_or_reject(
|
||||
Arc::clone(&fleet.entitlements),
|
||||
p,
|
||||
crate::metering::audio_seconds_to_tokens(seconds, rate),
|
||||
)
|
||||
.await?;
|
||||
Ok((guard, Some(rate)))
|
||||
}
|
||||
|
||||
impl AudioJobs {
|
||||
fn keep_in_memory(&self, id: &str, bytes: Bytes) {
|
||||
if let Ok(mut m) = self.memory_audio.lock() {
|
||||
m.insert(id.to_string(), bytes);
|
||||
}
|
||||
}
|
||||
|
||||
fn memory_audio(&self, id: &str) -> Option<Vec<u8>> {
|
||||
self.memory_audio
|
||||
.lock()
|
||||
.ok()
|
||||
.and_then(|m| m.get(id).map(|b| b.to_vec()))
|
||||
}
|
||||
}
|
||||
|
||||
fn now_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn status_word(s: JobStatus) -> &'static str {
|
||||
match s {
|
||||
JobStatus::Queued => "queued",
|
||||
JobStatus::Running => "running",
|
||||
JobStatus::Succeeded => "succeeded",
|
||||
JobStatus::Failed => "failed",
|
||||
JobStatus::Cancelled => "cancelled",
|
||||
}
|
||||
}
|
||||
|
||||
fn error_value(typ: &str, code: &str, message: &str) -> Value {
|
||||
serde_json::json!({"type": typ, "code": code, "message": message})
|
||||
}
|
||||
|
||||
fn envelope_error(env: &OpenAiError) -> Value {
|
||||
env.body()
|
||||
.get("error")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| error_value("api_error", "budget_rejected", "budget reservation failed"))
|
||||
}
|
||||
|
||||
fn invalid_params(message: impl Into<String>) -> Response {
|
||||
crate::error::envelope_response(OpenAiError::new(
|
||||
400,
|
||||
"invalid_request_error",
|
||||
"invalid_audio_params",
|
||||
format!("invalid audio generation parameters: {}", message.into()),
|
||||
))
|
||||
}
|
||||
|
||||
fn job_not_found(id: &str) -> Response {
|
||||
crate::error::envelope_response(OpenAiError::new(
|
||||
404,
|
||||
"invalid_request_error",
|
||||
"job_not_found",
|
||||
format!("no audio job '{id}'"),
|
||||
))
|
||||
}
|
||||
@@ -24,6 +24,7 @@ pub fn api_routes() -> Router<Arc<CortexState>> {
|
||||
.route("/v1/completions", post(completions))
|
||||
.route("/v1/responses", post(responses))
|
||||
.route("/v1/images/generations", post(images_generations))
|
||||
.merge(crate::audio_jobs::routes())
|
||||
.route("/v1/systemone", post(systemone))
|
||||
.route("/v1/models", get(list_models))
|
||||
.route("/v1/messages", post(anthropic_messages))
|
||||
@@ -977,6 +978,8 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
// `cost` stays operator-set and flows from the catalogue.
|
||||
limit: None,
|
||||
cost: profile.cost.clone(),
|
||||
audio: profile.audio.clone(),
|
||||
audio_tokens_per_second: None,
|
||||
// Runtime-detected — will be OR-ed in Pass 2 from neuron data.
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
@@ -1059,6 +1062,8 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
capabilities: entry.capabilities.clone(),
|
||||
limit: entry.limit.clone(),
|
||||
cost: None,
|
||||
audio: None,
|
||||
audio_tokens_per_second: None,
|
||||
tool_call: entry.tool_call,
|
||||
reasoning: entry.reasoning,
|
||||
max_model_len: None,
|
||||
@@ -1128,6 +1133,8 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
capabilities: discovered.get(model_id).cloned().unwrap_or_default(),
|
||||
limit: None,
|
||||
cost: None,
|
||||
audio: None,
|
||||
audio_tokens_per_second: None,
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
max_model_len: None,
|
||||
@@ -1169,6 +1176,8 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
capabilities: target_entry.capabilities,
|
||||
limit: target_entry.limit.clone(),
|
||||
cost: target_entry.cost.clone(),
|
||||
audio: target_entry.audio.clone(),
|
||||
audio_tokens_per_second: None,
|
||||
tool_call: target_entry.tool_call,
|
||||
reasoning: target_entry.reasoning,
|
||||
max_model_len: None,
|
||||
@@ -1186,10 +1195,19 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
// from each entry's now-settled `limit`, so vLLM-convention clients
|
||||
// (Hermes Agent et al.) can read the window without knowing helexa's
|
||||
// `limit` schema.
|
||||
//
|
||||
// Audio models also carry their current rate (#386), so a client can
|
||||
// estimate a job's cost before submitting it. An alias is priced as
|
||||
// the model it names.
|
||||
let rates = fleet.rates.snapshot().await;
|
||||
let data: Vec<Value> = entries
|
||||
.values_mut()
|
||||
.map(|e| {
|
||||
e.sync_flat_limit();
|
||||
if e.audio.is_some() {
|
||||
let priced_as = catalogue.resolve_alias(&e.id);
|
||||
e.audio_tokens_per_second = rates.audio_tokens_per_second(priced_as);
|
||||
}
|
||||
json!(e)
|
||||
})
|
||||
.collect();
|
||||
@@ -1455,7 +1473,7 @@ fn context_length_exceeded_response(
|
||||
}
|
||||
|
||||
/// Update `last_accessed` timestamp for a model on a node (drives LRU eviction).
|
||||
async fn touch_model(fleet: &CortexState, node_name: &str, model_id: &str) {
|
||||
pub(crate) async fn touch_model(fleet: &CortexState, node_name: &str, model_id: &str) {
|
||||
let mut nodes = fleet.nodes.write().await;
|
||||
if let Some(node) = nodes.get_mut(node_name)
|
||||
&& let Some(entry) = node.models.get_mut(model_id)
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
pub mod anthropic_sse;
|
||||
pub mod audio_jobs;
|
||||
pub mod auth;
|
||||
pub mod entitlements_chain;
|
||||
pub mod entitlements_local;
|
||||
@@ -10,6 +11,7 @@ pub mod metering;
|
||||
pub mod metrics;
|
||||
pub mod poller;
|
||||
pub mod proxy;
|
||||
pub mod rates;
|
||||
pub mod router;
|
||||
pub mod served_usage;
|
||||
pub mod state;
|
||||
@@ -52,6 +54,15 @@ pub async fn run(config: GatewayConfig) -> Result<()> {
|
||||
poller::poll_loop(poller_fleet).await;
|
||||
});
|
||||
|
||||
// Per-model rates (#386): keep the cached table in step with upstream.
|
||||
if config.upstream.enabled {
|
||||
rates::spawn_refresh(Arc::clone(&fleet.rates), &config.upstream);
|
||||
}
|
||||
|
||||
// Async audio jobs (#382): reload persisted jobs, re-queue unfinished
|
||||
// ones, start the retention sweep.
|
||||
audio_jobs::start(Arc::clone(&fleet)).await;
|
||||
|
||||
// Spawn the evictor (reacts to VRAM pressure events from the router).
|
||||
let evictor_fleet = Arc::clone(&fleet);
|
||||
tokio::spawn(async move {
|
||||
|
||||
@@ -228,6 +228,42 @@ pub fn settle_image_usage(
|
||||
guard.settle(tokens);
|
||||
}
|
||||
|
||||
/// Audio-seconds to reserve for a job: the requested duration, else the
|
||||
/// model's advertised default, else a minute. Settle corrects to the actual
|
||||
/// length the neuron reports.
|
||||
pub fn audio_reservation_seconds(requested: Option<f64>, model_default: Option<f64>) -> f64 {
|
||||
requested
|
||||
.or(model_default)
|
||||
.filter(|d| d.is_finite() && *d > 0.0)
|
||||
.unwrap_or(60.0)
|
||||
}
|
||||
|
||||
/// Convert audio-seconds to budget tokens at a per-model rate (#386),
|
||||
/// rounding up.
|
||||
pub fn audio_seconds_to_tokens(seconds: f64, tokens_per_second: u64) -> u64 {
|
||||
(seconds * tokens_per_second as f64).ceil() as u64
|
||||
}
|
||||
|
||||
/// Record per-principal audio spend and settle the reservation with the
|
||||
/// audio-seconds actually generated.
|
||||
pub fn settle_audio_usage(
|
||||
principal: &Principal,
|
||||
guard: ReservationGuard,
|
||||
served_usage: &crate::served_usage::ServedUsage,
|
||||
audio_seconds: f64,
|
||||
tokens_per_second: u64,
|
||||
) {
|
||||
let labels = [
|
||||
("account", principal.account_id.clone()),
|
||||
("key", principal.key_id.clone()),
|
||||
];
|
||||
metrics::counter!("cortex_spend_audio_milliseconds_total", &labels)
|
||||
.increment((audio_seconds * 1000.0).round() as u64);
|
||||
let tokens = audio_seconds_to_tokens(audio_seconds, tokens_per_second);
|
||||
served_usage.add(&principal.account_id, &principal.key_id, tokens);
|
||||
guard.settle(tokens);
|
||||
}
|
||||
|
||||
/// Map a [`BudgetError`] to the #63 envelope. The provider chose the window
|
||||
/// semantics; this only translates them to HTTP.
|
||||
fn budget_error_to_envelope(err: BudgetError) -> OpenAiError {
|
||||
|
||||
@@ -0,0 +1,208 @@
|
||||
//! Per-model rates for work metered in units other than tokens (#386).
|
||||
//!
|
||||
//! The budget ledger counts tokens. Audio jobs are metered in seconds of
|
||||
//! generated audio, converted to ledger tokens at a per-model rate. The
|
||||
//! rates are configuration held by the clearing house (helexa-upstream's
|
||||
//! `model_rates` table), so pricing changes without a deploy: cortex
|
||||
//! fetches the table from `/authz/v1/rates`, caches it, and refreshes it in
|
||||
//! the background.
|
||||
//!
|
||||
//! Resolution: the model's own rate, else the unit's `"*"` default. A model
|
||||
//! with neither is not priced, and cortex refuses its metered work rather
|
||||
//! than serving it free. Until the first fetch succeeds nothing is priced,
|
||||
//! for the same reason.
|
||||
//!
|
||||
//! A cortex running without upstream has no clearing house to ask. It
|
||||
//! prices every audio-second at [`LOCAL_DEFAULT_TOKENS_PER_AUDIO_SECOND`],
|
||||
//! which only matters for the operator's own local key caps.
|
||||
|
||||
use cortex_core::config::UpstreamClientConfig;
|
||||
use serde::Deserialize;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
/// Unit name for one second of generated audio, as stored upstream.
|
||||
pub const AUDIO_SECOND: &str = "audio_second";
|
||||
|
||||
/// The default-rate model id.
|
||||
pub const ANY_MODEL: &str = "*";
|
||||
|
||||
/// Audio rate used when this cortex runs without an upstream authority.
|
||||
pub const LOCAL_DEFAULT_TOKENS_PER_AUDIO_SECOND: u64 = 1000;
|
||||
|
||||
/// How often the rate table is refreshed from upstream.
|
||||
const REFRESH_INTERVAL: Duration = Duration::from_secs(60);
|
||||
|
||||
/// Retry interval while no table has loaded yet.
|
||||
const INITIAL_RETRY: Duration = Duration::from_secs(5);
|
||||
|
||||
/// A snapshot of the rate table: `(unit, model_id) → tokens per unit`.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct RateTable {
|
||||
rates: HashMap<(String, String), u64>,
|
||||
}
|
||||
|
||||
impl RateTable {
|
||||
/// The table a cortex without upstream uses.
|
||||
pub fn local_default() -> Self {
|
||||
let mut t = Self::default();
|
||||
t.set(
|
||||
AUDIO_SECOND,
|
||||
ANY_MODEL,
|
||||
LOCAL_DEFAULT_TOKENS_PER_AUDIO_SECOND,
|
||||
);
|
||||
t
|
||||
}
|
||||
|
||||
pub fn set(&mut self, unit: &str, model_id: &str, tokens_per_unit: u64) {
|
||||
self.rates
|
||||
.insert((unit.to_string(), model_id.to_string()), tokens_per_unit);
|
||||
}
|
||||
|
||||
/// The model's rate for `unit`, else the unit's default, else `None`.
|
||||
pub fn rate(&self, unit: &str, model_id: &str) -> Option<u64> {
|
||||
self.rates
|
||||
.get(&(unit.to_string(), model_id.to_string()))
|
||||
.or_else(|| self.rates.get(&(unit.to_string(), ANY_MODEL.to_string())))
|
||||
.copied()
|
||||
}
|
||||
|
||||
/// Ledger tokens per second of audio for `model_id`.
|
||||
pub fn audio_tokens_per_second(&self, model_id: &str) -> Option<u64> {
|
||||
self.rate(AUDIO_SECOND, model_id)
|
||||
}
|
||||
}
|
||||
|
||||
/// The live rate table, held in `CortexState`.
|
||||
pub struct Rates {
|
||||
table: RwLock<RateTable>,
|
||||
}
|
||||
|
||||
impl Rates {
|
||||
/// Rates for a cortex with no upstream: the local default.
|
||||
pub fn local() -> Self {
|
||||
Self {
|
||||
table: RwLock::new(RateTable::local_default()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Rates for a cortex with upstream: empty (nothing priced) until the
|
||||
/// first fetch lands.
|
||||
pub fn upstream() -> Self {
|
||||
Self {
|
||||
table: RwLock::new(RateTable::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn snapshot(&self) -> RateTable {
|
||||
self.table.read().await.clone()
|
||||
}
|
||||
|
||||
pub async fn replace(&self, table: RateTable) {
|
||||
*self.table.write().await = table;
|
||||
}
|
||||
|
||||
pub async fn audio_tokens_per_second(&self, model_id: &str) -> Option<u64> {
|
||||
self.table.read().await.audio_tokens_per_second(model_id)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RatesResp {
|
||||
data: Vec<RateRow>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RateRow {
|
||||
model_id: String,
|
||||
unit: String,
|
||||
tokens_per_unit: i64,
|
||||
}
|
||||
|
||||
/// Fetch the full rate table from upstream's `/authz/v1/rates`.
|
||||
pub async fn fetch(
|
||||
client: &reqwest::Client,
|
||||
base_url: &str,
|
||||
bearer: &str,
|
||||
) -> anyhow::Result<RateTable> {
|
||||
let url = format!("{}/authz/v1/rates", base_url.trim_end_matches('/'));
|
||||
let resp = client
|
||||
.get(&url)
|
||||
.bearer_auth(bearer)
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
let body: RatesResp = resp.json().await?;
|
||||
let mut table = RateTable::default();
|
||||
for row in body.data {
|
||||
// The database forbids non-positive rates; skip one rather than
|
||||
// trust it if that ever changes.
|
||||
if let Ok(tokens) = u64::try_from(row.tokens_per_unit)
|
||||
&& tokens > 0
|
||||
{
|
||||
table.set(&row.unit, &row.model_id, tokens);
|
||||
}
|
||||
}
|
||||
Ok(table)
|
||||
}
|
||||
|
||||
/// Keep `rates` in step with upstream. Retries quickly until the first
|
||||
/// table loads, then refreshes every minute. A failed refresh keeps the
|
||||
/// last good table.
|
||||
pub fn spawn_refresh(rates: Arc<Rates>, cfg: &UpstreamClientConfig) {
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(cfg.timeout_secs.max(1)))
|
||||
.build()
|
||||
.expect("failed to build rates HTTP client");
|
||||
let url = cfg.url.clone();
|
||||
let bearer = cfg.bearer.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut loaded = false;
|
||||
loop {
|
||||
match fetch(&client, &url, &bearer).await {
|
||||
Ok(table) => {
|
||||
if !loaded {
|
||||
tracing::info!("model rates loaded from upstream");
|
||||
}
|
||||
loaded = true;
|
||||
rates.replace(table).await;
|
||||
}
|
||||
Err(e) => tracing::warn!(error = %e, "model rate refresh failed"),
|
||||
}
|
||||
tokio::time::sleep(if loaded {
|
||||
REFRESH_INTERVAL
|
||||
} else {
|
||||
INITIAL_RETRY
|
||||
})
|
||||
.await;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn model_rate_beats_default_and_absent_unit_is_unpriced() {
|
||||
let mut t = RateTable::default();
|
||||
assert_eq!(t.audio_tokens_per_second("m"), None);
|
||||
t.set(AUDIO_SECOND, ANY_MODEL, 1000);
|
||||
assert_eq!(t.audio_tokens_per_second("m"), Some(1000));
|
||||
t.set(AUDIO_SECOND, "m", 250);
|
||||
assert_eq!(t.audio_tokens_per_second("m"), Some(250));
|
||||
assert_eq!(t.audio_tokens_per_second("other"), Some(1000));
|
||||
assert_eq!(t.rate("some_other_unit", "m"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_default_prices_every_audio_model() {
|
||||
let t = RateTable::local_default();
|
||||
assert_eq!(
|
||||
t.audio_tokens_per_second("anything"),
|
||||
Some(LOCAL_DEFAULT_TOKENS_PER_AUDIO_SECOND)
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -663,6 +663,7 @@ mod tests {
|
||||
limit: None,
|
||||
cost: None,
|
||||
capabilities: vec![],
|
||||
audio: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -52,6 +52,11 @@ pub struct CortexState {
|
||||
/// rather than an edge — it hid a GGUF that two of three nodes could
|
||||
/// describe perfectly well.
|
||||
pub capability_probe_attempts: RwLock<HashMap<(String, String), std::time::Instant>>,
|
||||
/// Async audio generation jobs (#382).
|
||||
pub audio_jobs: crate::audio_jobs::AudioJobs,
|
||||
/// Per-model rates for non-token units (#386), from upstream when it
|
||||
/// is enabled, else the local default.
|
||||
pub rates: Arc<crate::rates::Rates>,
|
||||
}
|
||||
|
||||
impl CortexState {
|
||||
@@ -131,6 +136,12 @@ impl CortexState {
|
||||
served_usage: Arc::new(crate::served_usage::ServedUsage::new()),
|
||||
discovered_capabilities: RwLock::new(HashMap::new()),
|
||||
capability_probe_attempts: RwLock::new(HashMap::new()),
|
||||
audio_jobs: crate::audio_jobs::AudioJobs::new(&config.audio_jobs),
|
||||
rates: Arc::new(if config.upstream.enabled {
|
||||
crate::rates::Rates::upstream()
|
||||
} else {
|
||||
crate::rates::Rates::local()
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -87,6 +87,7 @@ async fn test_alias_resolves_in_chat_completions() {
|
||||
models_config: models_path.to_string_lossy().to_string(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
@@ -176,6 +177,7 @@ async fn test_aliases_surface_in_v1_models() {
|
||||
models_config: models_path.to_string_lossy().to_string(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
@@ -268,6 +270,7 @@ async fn test_alias_falls_through_for_unmapped_model() {
|
||||
models_config: models_path.to_string_lossy().to_string(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
|
||||
@@ -0,0 +1,820 @@
|
||||
//! Async audio generation jobs (#382): submit → poll → content, cancel,
|
||||
//! ownership, retries, busy back-off, and restart recovery, against a
|
||||
//! scriptable mock neuron.
|
||||
|
||||
use axum::Router;
|
||||
use axum::extract::Path;
|
||||
use axum::http::{HeaderMap, StatusCode};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::{get, post};
|
||||
use axum::{Json, body::Bytes};
|
||||
use cortex_core::audio::{AudioUsage, NeuronAudioResult, RESULT_HEADER};
|
||||
use cortex_core::config::{
|
||||
ApiKeyConfig, AudioJobsConfig, EntitlementsConfig, EvictionSettings, EvictionStrategy,
|
||||
GatewayConfig, GatewaySettings, NeuronEndpoint,
|
||||
};
|
||||
use cortex_core::node::{ModelEntry, ModelStatus};
|
||||
use cortex_gateway::state::CortexState;
|
||||
use serde_json::{Value, json};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::time::Duration;
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
const MODEL: &str = "music-model";
|
||||
const WAV: &[u8] = b"RIFF\x24\x00\x00\x00WAVEfmt fake-audio-bytes";
|
||||
|
||||
/// How the mock neuron answers `POST /v1/audio/generations`.
|
||||
#[derive(Clone, Copy)]
|
||||
enum Mode {
|
||||
/// Render immediately.
|
||||
Ok,
|
||||
/// Answer 503 + Retry-After for the first N calls, then render.
|
||||
BusyThenOk(usize),
|
||||
/// Answer 500 for the first N calls, then render.
|
||||
FailThenOk(usize),
|
||||
/// Always 500.
|
||||
AlwaysFail,
|
||||
/// Always 422 wrong_modality.
|
||||
Reject,
|
||||
/// Sleep this long, then render.
|
||||
Slow(Duration),
|
||||
/// Always 503.
|
||||
AlwaysBusy,
|
||||
}
|
||||
|
||||
struct Mock {
|
||||
url: String,
|
||||
calls: Arc<AtomicUsize>,
|
||||
/// Account header seen on each render call.
|
||||
accounts: Arc<std::sync::Mutex<Vec<Option<String>>>>,
|
||||
}
|
||||
|
||||
async fn spawn_mock(mode: Mode) -> Mock {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let calls = Arc::new(AtomicUsize::new(0));
|
||||
let accounts = Arc::new(std::sync::Mutex::new(Vec::new()));
|
||||
let endpoint = url.clone();
|
||||
let (c, a) = (Arc::clone(&calls), Arc::clone(&accounts));
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/models/{id}/endpoint",
|
||||
get(move |Path(_id): Path<String>| {
|
||||
let url = endpoint.clone();
|
||||
async move { Json(json!({"url": url})) }
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/audio/generations",
|
||||
post(move |headers: HeaderMap, body: Bytes| {
|
||||
let (c, a) = (Arc::clone(&c), Arc::clone(&a));
|
||||
async move { render(mode, &c, &a, &headers, &body).await }
|
||||
}),
|
||||
);
|
||||
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
|
||||
Mock {
|
||||
url,
|
||||
calls,
|
||||
accounts,
|
||||
}
|
||||
}
|
||||
|
||||
async fn render(
|
||||
mode: Mode,
|
||||
calls: &AtomicUsize,
|
||||
accounts: &std::sync::Mutex<Vec<Option<String>>>,
|
||||
headers: &HeaderMap,
|
||||
body: &[u8],
|
||||
) -> Response {
|
||||
let n = calls.fetch_add(1, Ordering::SeqCst);
|
||||
accounts.lock().unwrap().push(
|
||||
headers
|
||||
.get(cortex_core::entitlements::HEADER_ACCOUNT_ID)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(str::to_string),
|
||||
);
|
||||
let busy = || {
|
||||
(
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
[("retry-after", "1")],
|
||||
Json(json!({"error": {"type": "rate_limit_error", "code": "rate_limit_exceeded", "message": "busy"}})),
|
||||
)
|
||||
.into_response()
|
||||
};
|
||||
let fail = || {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(json!({"error": {"type": "api_error", "code": "inference_failed", "message": "boom"}})),
|
||||
)
|
||||
.into_response()
|
||||
};
|
||||
match mode {
|
||||
Mode::BusyThenOk(k) if n < k => return busy(),
|
||||
Mode::AlwaysBusy => return busy(),
|
||||
Mode::FailThenOk(k) if n < k => return fail(),
|
||||
Mode::AlwaysFail => return fail(),
|
||||
Mode::Reject => {
|
||||
return (
|
||||
StatusCode::UNPROCESSABLE_ENTITY,
|
||||
Json(json!({"error": {"type": "invalid_request_error", "code": "wrong_modality", "message": "not an audio model"}})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
Mode::Slow(d) => tokio::time::sleep(d).await,
|
||||
_ => {}
|
||||
}
|
||||
let req: Value = serde_json::from_slice(body).unwrap();
|
||||
let duration = req["duration"].as_f64().unwrap_or(30.0);
|
||||
let result = NeuronAudioResult {
|
||||
duration,
|
||||
sample_rate: 48_000,
|
||||
channels: 2,
|
||||
seed: req["seed"].as_u64().unwrap_or(7),
|
||||
usage: AudioUsage {
|
||||
helexa_audio_units: duration,
|
||||
helexa_timing: None,
|
||||
},
|
||||
};
|
||||
(
|
||||
StatusCode::OK,
|
||||
[
|
||||
("content-type", "audio/wav".to_string()),
|
||||
(RESULT_HEADER, serde_json::to_string(&result).unwrap()),
|
||||
],
|
||||
WAV,
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
fn config(neuron: &str, dir: &std::path::Path) -> GatewayConfig {
|
||||
GatewayConfig {
|
||||
gateway: GatewaySettings {
|
||||
listen: "127.0.0.1:0".into(),
|
||||
metrics_listen: "127.0.0.1:0".into(),
|
||||
},
|
||||
eviction: EvictionSettings {
|
||||
strategy: EvictionStrategy::Lru,
|
||||
defrag_after_cycles: 0,
|
||||
},
|
||||
neurons: vec![NeuronEndpoint {
|
||||
name: "mock-node".into(),
|
||||
endpoint: neuron.to_string(),
|
||||
}],
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: EntitlementsConfig {
|
||||
require_auth: false,
|
||||
keys: ["alice", "bob"]
|
||||
.iter()
|
||||
.map(|who| ApiKeyConfig {
|
||||
key: format!("key-{who}"),
|
||||
account_id: format!("acct-{who}"),
|
||||
key_id: None,
|
||||
hard_cap: None,
|
||||
window: Default::default(),
|
||||
})
|
||||
.collect(),
|
||||
},
|
||||
upstream: Default::default(),
|
||||
audio_jobs: AudioJobsConfig {
|
||||
dir: dir.to_string_lossy().into_owned(),
|
||||
queue_timeout_secs: 20,
|
||||
..Default::default()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
async fn gateway(neuron: &str, dir: &std::path::Path) -> (Arc<CortexState>, String) {
|
||||
let fleet = Arc::new(CortexState::from_config(&config(neuron, dir)));
|
||||
seed_model(&fleet).await;
|
||||
let url = serve(&fleet).await;
|
||||
(fleet, url)
|
||||
}
|
||||
|
||||
async fn seed_model(fleet: &CortexState) {
|
||||
let mut nodes = fleet.nodes.write().await;
|
||||
let node = nodes.get_mut("mock-node").unwrap();
|
||||
node.healthy = true;
|
||||
node.models.insert(
|
||||
MODEL.into(),
|
||||
ModelEntry {
|
||||
id: MODEL.into(),
|
||||
status: ModelStatus::Loaded,
|
||||
last_accessed: None,
|
||||
vram_estimate_mb: Some(8000),
|
||||
capabilities: vec!["audio".into()],
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
limit: None,
|
||||
servable: None,
|
||||
reasoning_budget: Vec::new(),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
async fn serve(fleet: &Arc<CortexState>) -> String {
|
||||
let app = cortex_gateway::build_app(Arc::clone(fleet));
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
|
||||
format!("http://{addr}")
|
||||
}
|
||||
|
||||
async fn submit(gw: &str, key: Option<&str>, body: Value) -> reqwest::Response {
|
||||
let mut req = reqwest::Client::new()
|
||||
.post(format!("{gw}/v1/audio/generations"))
|
||||
.json(&body);
|
||||
if let Some(k) = key {
|
||||
req = req.bearer_auth(k);
|
||||
}
|
||||
req.send().await.unwrap()
|
||||
}
|
||||
|
||||
async fn get_job(gw: &str, key: Option<&str>, id: &str) -> reqwest::Response {
|
||||
let mut req = reqwest::Client::new().get(format!("{gw}/v1/audio/generations/{id}"));
|
||||
if let Some(k) = key {
|
||||
req = req.bearer_auth(k);
|
||||
}
|
||||
req.send().await.unwrap()
|
||||
}
|
||||
|
||||
/// Poll until the job reaches a terminal status (or panic after `secs`).
|
||||
async fn wait_terminal(gw: &str, key: Option<&str>, id: &str, secs: u64) -> Value {
|
||||
let deadline = std::time::Instant::now() + Duration::from_secs(secs);
|
||||
loop {
|
||||
let job: Value = get_job(gw, key, id).await.json().await.unwrap();
|
||||
if matches!(
|
||||
job["status"].as_str(),
|
||||
Some("succeeded" | "failed" | "cancelled")
|
||||
) {
|
||||
return job;
|
||||
}
|
||||
assert!(
|
||||
std::time::Instant::now() < deadline,
|
||||
"job {id} never finished: {job}"
|
||||
);
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
}
|
||||
}
|
||||
|
||||
fn song() -> Value {
|
||||
json!({"model": MODEL, "prompt": "warm lofi", "duration": 12.5, "seed": 42})
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn submit_poll_fetch_content() {
|
||||
let mock = spawn_mock(Mode::Ok).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
|
||||
|
||||
let resp = submit(&gw, Some("key-alice"), song()).await;
|
||||
assert_eq!(resp.status(), 202);
|
||||
let job: Value = resp.json().await.unwrap();
|
||||
assert_eq!(job["object"], "audio.generation");
|
||||
assert_eq!(job["status"], "queued");
|
||||
assert_eq!(job["model"], MODEL);
|
||||
let id = job["id"].as_str().unwrap().to_string();
|
||||
assert!(id.starts_with("audiogen_"));
|
||||
|
||||
let done = wait_terminal(&gw, Some("key-alice"), &id, 10).await;
|
||||
assert_eq!(done["status"], "succeeded", "{done}");
|
||||
assert_eq!(done["result"]["duration"], 12.5);
|
||||
assert_eq!(done["result"]["sample_rate"], 48_000);
|
||||
assert_eq!(done["result"]["seed"], 42);
|
||||
assert_eq!(done["result"]["bytes"], WAV.len());
|
||||
assert_eq!(done["usage"]["helexa_audio_units"], 12.5);
|
||||
|
||||
let content = reqwest::Client::new()
|
||||
.get(format!("{gw}/v1/audio/generations/{id}/content"))
|
||||
.bearer_auth("key-alice")
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(content.status(), 200);
|
||||
assert_eq!(content.headers()["content-type"], "audio/wav");
|
||||
assert_eq!(content.bytes().await.unwrap().as_ref(), WAV);
|
||||
|
||||
// The neuron saw the submitting account, for fair-share admission.
|
||||
assert_eq!(
|
||||
mock.accounts.lock().unwrap().as_slice(),
|
||||
&[Some("acct-alice".to_string())]
|
||||
);
|
||||
// Record and audio are on disk.
|
||||
assert!(dir.path().join(format!("{id}.json")).exists());
|
||||
assert!(dir.path().join(format!("{id}.wav")).exists());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn content_before_completion_is_409() {
|
||||
let mock = spawn_mock(Mode::Slow(Duration::from_secs(3))).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
|
||||
let job: Value = submit(&gw, None, song()).await.json().await.unwrap();
|
||||
let id = job["id"].as_str().unwrap();
|
||||
let resp = reqwest::get(format!("{gw}/v1/audio/generations/{id}/content"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 409);
|
||||
let body: Value = resp.json().await.unwrap();
|
||||
assert_eq!(body["error"]["code"], "job_not_complete");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn jobs_are_private_to_their_account() {
|
||||
let mock = spawn_mock(Mode::Ok).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
|
||||
|
||||
let job: Value = submit(&gw, Some("key-alice"), song())
|
||||
.await
|
||||
.json()
|
||||
.await
|
||||
.unwrap();
|
||||
let id = job["id"].as_str().unwrap().to_string();
|
||||
wait_terminal(&gw, Some("key-alice"), &id, 10).await;
|
||||
|
||||
// Another account, and an anonymous caller, cannot see it.
|
||||
assert_eq!(get_job(&gw, Some("key-bob"), &id).await.status(), 404);
|
||||
assert_eq!(get_job(&gw, None, &id).await.status(), 404);
|
||||
let content = reqwest::Client::new()
|
||||
.get(format!("{gw}/v1/audio/generations/{id}/content"))
|
||||
.bearer_auth("key-bob")
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(content.status(), 404);
|
||||
let cancel = reqwest::Client::new()
|
||||
.delete(format!("{gw}/v1/audio/generations/{id}"))
|
||||
.bearer_auth("key-bob")
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(cancel.status(), 404);
|
||||
|
||||
// Listings are per account.
|
||||
let list = |key: &'static str| {
|
||||
let gw = gw.clone();
|
||||
async move {
|
||||
reqwest::Client::new()
|
||||
.get(format!("{gw}/v1/audio/generations"))
|
||||
.bearer_auth(key)
|
||||
.send()
|
||||
.await
|
||||
.unwrap()
|
||||
.json::<Value>()
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
};
|
||||
let alice = list("key-alice").await;
|
||||
assert_eq!(alice["object"], "list");
|
||||
assert_eq!(alice["data"].as_array().unwrap().len(), 1);
|
||||
assert_eq!(alice["data"][0]["id"], id.as_str());
|
||||
assert_eq!(list("key-bob").await["data"].as_array().unwrap().len(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn anonymous_jobs_are_reachable_by_id_but_never_listed() {
|
||||
let mock = spawn_mock(Mode::Ok).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
|
||||
let job: Value = submit(&gw, None, song()).await.json().await.unwrap();
|
||||
let id = job["id"].as_str().unwrap().to_string();
|
||||
let done = wait_terminal(&gw, None, &id, 10).await;
|
||||
assert_eq!(done["status"], "succeeded");
|
||||
let list: Value = reqwest::get(format!("{gw}/v1/audio/generations"))
|
||||
.await
|
||||
.unwrap()
|
||||
.json()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(list["data"].as_array().unwrap().len(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancel_running_job() {
|
||||
let mock = spawn_mock(Mode::Slow(Duration::from_secs(5))).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
|
||||
let job: Value = submit(&gw, Some("key-alice"), song())
|
||||
.await
|
||||
.json()
|
||||
.await
|
||||
.unwrap();
|
||||
let id = job["id"].as_str().unwrap().to_string();
|
||||
|
||||
// Wait until it is running on the neuron.
|
||||
let deadline = std::time::Instant::now() + Duration::from_secs(5);
|
||||
while mock.calls.load(Ordering::SeqCst) == 0 {
|
||||
assert!(std::time::Instant::now() < deadline);
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
}
|
||||
let resp = reqwest::Client::new()
|
||||
.delete(format!("{gw}/v1/audio/generations/{id}"))
|
||||
.bearer_auth("key-alice")
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 200);
|
||||
let cancelled: Value = resp.json().await.unwrap();
|
||||
assert_eq!(cancelled["status"], "cancelled");
|
||||
|
||||
// It stays cancelled: the slow render finishing must not resurrect it.
|
||||
tokio::time::sleep(Duration::from_secs(6)).await;
|
||||
let job: Value = get_job(&gw, Some("key-alice"), &id)
|
||||
.await
|
||||
.json()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(job["status"], "cancelled");
|
||||
assert!(job.get("result").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn busy_neuron_is_waited_on_without_spending_attempts() {
|
||||
// Three busy answers exceed max_attempts (2), so this only succeeds if
|
||||
// busy answers don't count as attempts.
|
||||
let mock = spawn_mock(Mode::BusyThenOk(3)).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
|
||||
let job: Value = submit(&gw, None, song()).await.json().await.unwrap();
|
||||
let id = job["id"].as_str().unwrap().to_string();
|
||||
let done = wait_terminal(&gw, None, &id, 15).await;
|
||||
assert_eq!(done["status"], "succeeded", "{done}");
|
||||
assert_eq!(mock.calls.load(Ordering::SeqCst), 4);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn failing_node_is_retried_then_succeeds() {
|
||||
let mock = spawn_mock(Mode::FailThenOk(1)).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
|
||||
let job: Value = submit(&gw, None, song()).await.json().await.unwrap();
|
||||
let done = wait_terminal(&gw, None, job["id"].as_str().unwrap(), 15).await;
|
||||
assert_eq!(done["status"], "succeeded", "{done}");
|
||||
assert_eq!(mock.calls.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn failing_node_exhausts_attempts() {
|
||||
let mock = spawn_mock(Mode::AlwaysFail).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
|
||||
let job: Value = submit(&gw, None, song()).await.json().await.unwrap();
|
||||
let done = wait_terminal(&gw, None, job["id"].as_str().unwrap(), 15).await;
|
||||
assert_eq!(done["status"], "failed");
|
||||
assert_eq!(done["error"]["code"], "inference_failed");
|
||||
assert_eq!(mock.calls.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn neuron_rejection_fails_without_retry() {
|
||||
let mock = spawn_mock(Mode::Reject).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
|
||||
let job: Value = submit(&gw, None, song()).await.json().await.unwrap();
|
||||
let done = wait_terminal(&gw, None, job["id"].as_str().unwrap(), 10).await;
|
||||
assert_eq!(done["status"], "failed");
|
||||
assert_eq!(done["error"]["code"], "wrong_modality");
|
||||
assert_eq!(mock.calls.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn submit_validation() {
|
||||
let mock = spawn_mock(Mode::Ok).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
|
||||
|
||||
let resp = submit(&gw, None, json!({"model": MODEL, "prompt": ""})).await;
|
||||
assert_eq!(resp.status(), 400);
|
||||
let body: Value = resp.json().await.unwrap();
|
||||
assert_eq!(body["error"]["code"], "invalid_audio_params");
|
||||
|
||||
let resp = submit(
|
||||
&gw,
|
||||
None,
|
||||
json!({"model": MODEL, "prompt": "x", "response_format": "ogg"}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(resp.status(), 400);
|
||||
|
||||
let resp = submit(&gw, None, json!({"model": "nope", "prompt": "x"})).await;
|
||||
assert_eq!(resp.status(), 404);
|
||||
let body: Value = resp.json().await.unwrap();
|
||||
assert_eq!(body["error"]["code"], "model_not_found");
|
||||
|
||||
let resp = submit(&gw, None, json!({"prompt": "x"})).await;
|
||||
assert_eq!(resp.status(), 400);
|
||||
|
||||
// Nothing reached the neuron.
|
||||
assert_eq!(mock.calls.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unknown_job_is_404() {
|
||||
let mock = spawn_mock(Mode::Ok).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
|
||||
let resp = get_job(&gw, None, "audiogen_doesnotexist").await;
|
||||
assert_eq!(resp.status(), 404);
|
||||
let body: Value = resp.json().await.unwrap();
|
||||
assert_eq!(body["error"]["code"], "job_not_found");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn jobs_survive_a_restart() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
// First cortex: one job finishes, one is stuck behind a busy model.
|
||||
let ok = spawn_mock(Mode::Ok).await;
|
||||
let (_a, gw_a) = gateway(&ok.url, dir.path()).await;
|
||||
let done: Value = submit(&gw_a, Some("key-alice"), song())
|
||||
.await
|
||||
.json()
|
||||
.await
|
||||
.unwrap();
|
||||
let done_id = done["id"].as_str().unwrap().to_string();
|
||||
wait_terminal(&gw_a, Some("key-alice"), &done_id, 10).await;
|
||||
|
||||
let busy = spawn_mock(Mode::AlwaysBusy).await;
|
||||
let (_b, gw_b) = gateway(&busy.url, dir.path()).await;
|
||||
let stuck: Value = submit(&gw_b, Some("key-alice"), song())
|
||||
.await
|
||||
.json()
|
||||
.await
|
||||
.unwrap();
|
||||
let stuck_id = stuck["id"].as_str().unwrap().to_string();
|
||||
let deadline = std::time::Instant::now() + Duration::from_secs(5);
|
||||
while busy.calls.load(Ordering::SeqCst) == 0 {
|
||||
assert!(std::time::Instant::now() < deadline);
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
}
|
||||
|
||||
// A fresh cortex over the same directory, with a healthy neuron.
|
||||
let healthy = spawn_mock(Mode::Ok).await;
|
||||
let fleet = Arc::new(CortexState::from_config(&config(&healthy.url, dir.path())));
|
||||
seed_model(&fleet).await;
|
||||
cortex_gateway::audio_jobs::start(Arc::clone(&fleet)).await;
|
||||
let gw = serve(&fleet).await;
|
||||
|
||||
// The finished job and its audio are still there.
|
||||
let job: Value = get_job(&gw, Some("key-alice"), &done_id)
|
||||
.await
|
||||
.json()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(job["status"], "succeeded");
|
||||
let content = reqwest::Client::new()
|
||||
.get(format!("{gw}/v1/audio/generations/{done_id}/content"))
|
||||
.bearer_auth("key-alice")
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(content.bytes().await.unwrap().as_ref(), WAV);
|
||||
|
||||
// The unfinished one was re-queued and now completes.
|
||||
let resumed = wait_terminal(&gw, Some("key-alice"), &stuck_id, 10).await;
|
||||
assert_eq!(resumed["status"], "succeeded", "{resumed}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn expired_jobs_are_swept() {
|
||||
let mock = spawn_mock(Mode::Ok).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mut cfg = config(&mock.url, dir.path());
|
||||
cfg.audio_jobs.retention_hours = 0;
|
||||
let fleet = Arc::new(CortexState::from_config(&cfg));
|
||||
seed_model(&fleet).await;
|
||||
let gw = serve(&fleet).await;
|
||||
|
||||
let job: Value = submit(&gw, None, song()).await.json().await.unwrap();
|
||||
let id = job["id"].as_str().unwrap().to_string();
|
||||
wait_terminal(&gw, None, &id, 10).await;
|
||||
// completed_at has one-second resolution; let it fall behind the cutoff.
|
||||
tokio::time::sleep(Duration::from_millis(1100)).await;
|
||||
|
||||
cortex_gateway::audio_jobs::sweep_expired(&fleet).await;
|
||||
assert_eq!(get_job(&gw, None, &id).await.status(), 404);
|
||||
assert!(!dir.path().join(format!("{id}.json")).exists());
|
||||
assert!(!dir.path().join(format!("{id}.wav")).exists());
|
||||
}
|
||||
|
||||
// ── Per-model rates (#386) ───────────────────────────────────────────
|
||||
|
||||
/// A gateway whose keys are capped at `cap` tokens and whose audio rate
|
||||
/// table is `rates` (model id → tokens per second; "*" is the default).
|
||||
async fn priced_gateway(
|
||||
neuron: &str,
|
||||
dir: &std::path::Path,
|
||||
cap: Option<u64>,
|
||||
rates: &[(&str, u64)],
|
||||
) -> (Arc<CortexState>, String) {
|
||||
let mut cfg = config(neuron, dir);
|
||||
for k in &mut cfg.entitlements.keys {
|
||||
k.hard_cap = cap;
|
||||
}
|
||||
let fleet = Arc::new(CortexState::from_config(&cfg));
|
||||
let mut table = cortex_gateway::rates::RateTable::default();
|
||||
for (model, rate) in rates {
|
||||
table.set(cortex_gateway::rates::AUDIO_SECOND, model, *rate);
|
||||
}
|
||||
fleet.rates.replace(table).await;
|
||||
seed_model(&fleet).await;
|
||||
let url = serve(&fleet).await;
|
||||
(fleet, url)
|
||||
}
|
||||
|
||||
async fn spent(fleet: &CortexState, who: &str) -> u64 {
|
||||
let principal = cortex_core::entitlements::Principal {
|
||||
account_id: format!("acct-{who}"),
|
||||
key_id: format!("acct-{who}"),
|
||||
};
|
||||
fleet
|
||||
.entitlements
|
||||
.snapshot(&principal)
|
||||
.await
|
||||
.map(|s| s.spent)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reservation_uses_the_models_rate() {
|
||||
// A 12.5 s song against a 10,000-token cap: refused at 1000 tokens/s
|
||||
// (12,500), accepted at 500 tokens/s (6,250).
|
||||
let mock = spawn_mock(Mode::Ok).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let (_f, gw) = priced_gateway(&mock.url, dir.path(), Some(10_000), &[("*", 1000)]).await;
|
||||
let resp = submit(&gw, Some("key-alice"), song()).await;
|
||||
assert_eq!(resp.status(), 429, "{:?}", resp.text().await);
|
||||
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let (_f, gw) = priced_gateway(
|
||||
&mock.url,
|
||||
dir.path(),
|
||||
Some(10_000),
|
||||
&[("*", 1000), (MODEL, 500)],
|
||||
)
|
||||
.await;
|
||||
assert_eq!(submit(&gw, Some("key-alice"), song()).await.status(), 202);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unpriced_model_is_refused_for_metered_callers_only() {
|
||||
let mock = spawn_mock(Mode::Ok).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
// An empty table: what a cortex with upstream sees before its first fetch.
|
||||
let (_f, gw) = priced_gateway(&mock.url, dir.path(), None, &[]).await;
|
||||
let resp = submit(&gw, Some("key-alice"), song()).await;
|
||||
assert_eq!(resp.status(), 503);
|
||||
let body: Value = resp.json().await.unwrap();
|
||||
assert!(
|
||||
body["error"]["message"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("no audio rate"),
|
||||
"{body}"
|
||||
);
|
||||
// Anonymous work is not metered, so it does not need a price.
|
||||
assert_eq!(submit(&gw, None, song()).await.status(), 202);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn settles_at_the_rate_reserved_even_if_it_changes_mid_job() {
|
||||
let mock = spawn_mock(Mode::Slow(Duration::from_millis(800))).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let (fleet, gw) = priced_gateway(&mock.url, dir.path(), None, &[(MODEL, 300)]).await;
|
||||
let job: Value = submit(&gw, Some("key-alice"), song())
|
||||
.await
|
||||
.json()
|
||||
.await
|
||||
.unwrap();
|
||||
// Reprice while the render is in flight.
|
||||
let mut table = cortex_gateway::rates::RateTable::default();
|
||||
table.set(cortex_gateway::rates::AUDIO_SECOND, MODEL, 1);
|
||||
fleet.rates.replace(table).await;
|
||||
|
||||
let done = wait_terminal(&gw, Some("key-alice"), job["id"].as_str().unwrap(), 10).await;
|
||||
assert_eq!(done["status"], "succeeded");
|
||||
// 12.5 s × 300 tokens/s, settled asynchronously.
|
||||
let deadline = std::time::Instant::now() + Duration::from_secs(5);
|
||||
loop {
|
||||
let s = spent(&fleet, "alice").await;
|
||||
if s == 3750 {
|
||||
break;
|
||||
}
|
||||
assert!(std::time::Instant::now() < deadline, "spent {s}, want 3750");
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rates_are_fetched_from_upstream_with_the_client_bearer() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let app = Router::new().route(
|
||||
"/authz/v1/rates",
|
||||
get(|headers: HeaderMap| async move {
|
||||
if headers.get("authorization").and_then(|v| v.to_str().ok())
|
||||
!= Some("Bearer op-secret")
|
||||
{
|
||||
return StatusCode::UNAUTHORIZED.into_response();
|
||||
}
|
||||
Json(json!({"data": [
|
||||
{"model_id": "*", "unit": "audio_second", "tokens_per_unit": 1000},
|
||||
{"model_id": MODEL, "unit": "audio_second", "tokens_per_unit": 40},
|
||||
{"model_id": "bad", "unit": "audio_second", "tokens_per_unit": 0}
|
||||
]}))
|
||||
.into_response()
|
||||
}),
|
||||
);
|
||||
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let table = cortex_gateway::rates::fetch(&client, &url, "op-secret")
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(table.audio_tokens_per_second(MODEL), Some(40));
|
||||
assert_eq!(table.audio_tokens_per_second("other"), Some(1000));
|
||||
// A non-positive rate is skipped, falling back to the default.
|
||||
assert_eq!(table.audio_tokens_per_second("bad"), Some(1000));
|
||||
assert!(
|
||||
cortex_gateway::rates::fetch(&client, &url, "wrong")
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn models_endpoint_advertises_audio_limits_and_rate() {
|
||||
let mock = spawn_mock(Mode::Ok).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let catalogue = dir.path().join("models.toml");
|
||||
std::fs::write(
|
||||
&catalogue,
|
||||
format!(
|
||||
r#"
|
||||
[[models]]
|
||||
id = "{MODEL}"
|
||||
harness = "candle"
|
||||
quant = "q8_0"
|
||||
|
||||
[models.audio]
|
||||
min_duration = 10.0
|
||||
max_duration = 600.0
|
||||
default_duration = 30.0
|
||||
vocals = true
|
||||
sample_rate = 48000
|
||||
channels = 2
|
||||
max_steps = 20
|
||||
|
||||
[aliases]
|
||||
"helexa/music" = "{MODEL}"
|
||||
"#
|
||||
),
|
||||
)
|
||||
.unwrap();
|
||||
let mut cfg = config(&mock.url, dir.path());
|
||||
cfg.models_config = catalogue.to_string_lossy().into_owned();
|
||||
let fleet = Arc::new(CortexState::from_config(&cfg));
|
||||
let mut table = cortex_gateway::rates::RateTable::default();
|
||||
table.set(cortex_gateway::rates::AUDIO_SECOND, MODEL, 77);
|
||||
fleet.rates.replace(table).await;
|
||||
seed_model(&fleet).await;
|
||||
let gw = serve(&fleet).await;
|
||||
|
||||
let models: Value = reqwest::get(format!("{gw}/v1/models"))
|
||||
.await
|
||||
.unwrap()
|
||||
.json()
|
||||
.await
|
||||
.unwrap();
|
||||
let find = |id: &str| {
|
||||
models["data"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.find(|m| m["id"] == id)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| panic!("{id} missing from {models}"))
|
||||
};
|
||||
for id in [MODEL, "helexa/music"] {
|
||||
let m = find(id);
|
||||
assert_eq!(m["audio"]["max_duration"], 600.0, "{m}");
|
||||
assert_eq!(m["audio"]["vocals"], true);
|
||||
assert_eq!(m["audio_tokens_per_second"], 77, "{m}");
|
||||
}
|
||||
|
||||
// The catalogue's limits are enforced at submit.
|
||||
let resp = submit(
|
||||
&gw,
|
||||
None,
|
||||
json!({"model": MODEL, "prompt": "x", "duration": 5.0}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(resp.status(), 400);
|
||||
// An alias is accepted and resolved to the model.
|
||||
let resp = submit(&gw, None, json!({"model": "helexa/music", "prompt": "x"})).await;
|
||||
assert_eq!(resp.status(), 202);
|
||||
let job: Value = resp.json().await.unwrap();
|
||||
assert_eq!(job["model"], MODEL);
|
||||
}
|
||||
@@ -106,6 +106,7 @@ async fn spawn_gateway(neuron_url: &str, entitlements: EntitlementsConfig) -> St
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements,
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
|
||||
@@ -82,6 +82,7 @@ async fn spawn_gateway(neuron_url: &str, key: ApiKeyConfig) -> (Arc<CortexState>
|
||||
keys: vec![key],
|
||||
},
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
{
|
||||
|
||||
@@ -498,6 +498,7 @@ pub async fn spawn_gateway_with_state(mock_url: &str) -> (Arc<CortexState>, Stri
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
@@ -559,6 +560,7 @@ pub async fn two_neuron_fleet(endpoint_a: &str, endpoint_b: &str) -> Arc<CortexS
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
Arc::new(CortexState::from_config(&config))
|
||||
}
|
||||
|
||||
@@ -90,6 +90,7 @@ async fn error_response_no_healthy_nodes() {
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(cortex_gateway::state::CortexState::from_config(&config));
|
||||
|
||||
@@ -82,6 +82,7 @@ fn make_fleet_with_catalogue(
|
||||
models_config: path.to_string_lossy().into_owned(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
(Arc::new(CortexState::from_config(&config)), path)
|
||||
}
|
||||
@@ -144,6 +145,7 @@ fn make_fleet(endpoint: &str, defrag_after: u32) -> Arc<CortexState> {
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
Arc::new(CortexState::from_config(&config))
|
||||
}
|
||||
|
||||
@@ -88,6 +88,7 @@ async fn fleet_with(big_healthy: bool, big_devices: usize) -> Arc<CortexState> {
|
||||
models_config: cat.to_string_lossy().into_owned(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
{
|
||||
|
||||
@@ -165,6 +165,7 @@ async fn spawn_keyed_gateway(neuron_url: &str, hard_cap: u64) -> String {
|
||||
}],
|
||||
},
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
{
|
||||
|
||||
@@ -89,6 +89,7 @@ async fn two_neuron_fleet(endpoint_a: &str, endpoint_b: &str) -> Arc<CortexState
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
Arc::new(CortexState::from_config(&config))
|
||||
}
|
||||
|
||||
@@ -54,6 +54,7 @@ async fn spawn_metered_gateway(neuron_url: &str) -> (Arc<CortexState>, String) {
|
||||
}],
|
||||
},
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
@@ -162,6 +163,7 @@ async fn anonymous_request_records_no_spend() {
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: EntitlementsConfig::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
{
|
||||
|
||||
@@ -204,6 +204,7 @@ async fn test_capacity_gauges_exported_from_health_poll() {
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
let fleet = std::sync::Arc::new(cortex_gateway::state::CortexState::from_config(&config));
|
||||
cortex_gateway::poller::poll_once(&fleet).await;
|
||||
|
||||
@@ -67,6 +67,7 @@ harness = "candle"
|
||||
models_config: cat_path.to_string_lossy().into_owned(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
|
||||
@@ -59,6 +59,7 @@ capabilities = ["text"]
|
||||
models_config: cat_path.to_string_lossy().into_owned(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
|
||||
@@ -33,6 +33,7 @@ async fn test_poller_discovers_models() {
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
@@ -86,6 +87,7 @@ async fn test_poller_updates_gateway_models_endpoint() {
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
@@ -159,6 +161,7 @@ async fn test_models_endpoint_unions_capabilities_across_nodes() {
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
@@ -223,6 +226,7 @@ async fn test_poller_marks_unreachable_node_unhealthy() {
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
@@ -278,6 +282,7 @@ async fn test_poller_removes_stale_models() {
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
@@ -310,6 +315,7 @@ async fn test_poller_removes_stale_models() {
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
|
||||
let fleet2 = Arc::new(CortexState::from_config(&config2));
|
||||
@@ -397,6 +403,7 @@ async fn test_poller_captures_activation_from_health() {
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
@@ -443,6 +450,7 @@ async fn test_poller_parses_recovering_status() {
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
|
||||
@@ -76,6 +76,7 @@ async fn spawn_gateway(neuron: &str, context: usize) -> String {
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
{
|
||||
|
||||
@@ -119,6 +119,7 @@ async fn test_no_healthy_nodes() {
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
let fleet = std::sync::Arc::new(cortex_gateway::state::CortexState::from_config(&config));
|
||||
|
||||
|
||||
@@ -119,6 +119,7 @@ async fn spawn_gateway(neuron_url: &str, default_alias: bool, hard_cap: Option<u
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements,
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
let mut state = CortexState::from_config(&config);
|
||||
if default_alias {
|
||||
|
||||
@@ -153,17 +153,29 @@ pub trait Scenario: Send + Sync {
|
||||
fn prompt_size(&self) -> u32;
|
||||
|
||||
/// Whether this scenario should run against the given model. The
|
||||
/// default runs against everything except image-generation models
|
||||
/// (a chat request at an image model is a 422 `wrong_modality`);
|
||||
/// modality-specific scenarios override this.
|
||||
/// default runs against chat models only: an image-generation model
|
||||
/// or a decision model (#333) answers a chat request with a 422
|
||||
/// `wrong_modality`. Modality-specific scenarios override this.
|
||||
fn applies_to(&self, model: &ModelInfo) -> bool {
|
||||
!model.capabilities.iter().any(|c| c == "image")
|
||||
serves_chat(model)
|
||||
}
|
||||
|
||||
/// Issue one shaped request and measure it.
|
||||
async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics>;
|
||||
}
|
||||
|
||||
/// Capabilities that mark a model as serving an endpoint other than
|
||||
/// chat completions.
|
||||
const NON_CHAT_MODALITIES: [&str; 3] = ["image", "decision", "audio"];
|
||||
|
||||
/// Whether a model answers `/v1/chat/completions`.
|
||||
pub fn serves_chat(model: &ModelInfo) -> bool {
|
||||
!model
|
||||
.capabilities
|
||||
.iter()
|
||||
.any(|c| NON_CHAT_MODALITIES.contains(&c.as_str()))
|
||||
}
|
||||
|
||||
/// Fixed-seed image-latency scenario (#203): one 9-step generation at a
|
||||
/// configured square size against `/v1/images/generations`, recording
|
||||
/// the server-side phase timing (`usage.helexa_timing`) and the metered
|
||||
@@ -883,6 +895,35 @@ async fn stream_and_measure_inner(
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn model(capabilities: &[&str]) -> ModelInfo {
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"id": "m",
|
||||
"harness": "candle",
|
||||
"status": "loaded",
|
||||
"devices": [0],
|
||||
"vram_used_mb": null,
|
||||
"capabilities": capabilities,
|
||||
}))
|
||||
.expect("model info")
|
||||
}
|
||||
|
||||
/// Chat scenarios skip every model that answers chat with
|
||||
/// `wrong_modality`: image generators and decision models alike.
|
||||
#[test]
|
||||
fn chat_scenarios_skip_image_and_decision_models() {
|
||||
assert!(serves_chat(&model(&[])));
|
||||
assert!(serves_chat(&model(&["text", "vision"])));
|
||||
assert!(!serves_chat(&model(&["image"])));
|
||||
assert!(!serves_chat(&model(&["decision"])));
|
||||
assert!(!serves_chat(&model(&["audio"])));
|
||||
let image = ImageLatencyScenario {
|
||||
id: "image:512".into(),
|
||||
side_px: 512,
|
||||
};
|
||||
assert!(image.applies_to(&model(&["image"])));
|
||||
assert!(!image.applies_to(&model(&["decision"])));
|
||||
}
|
||||
|
||||
/// A burst must report server-side phase timing, or `queue_wait_ms`
|
||||
/// (TTFT minus server prefill, #85) cannot be computed for the very
|
||||
/// cells it was designed for, and the prefill-tok/s panel stays
|
||||
|
||||
@@ -0,0 +1,173 @@
|
||||
//! Async audio generation jobs through the router (#382).
|
||||
//!
|
||||
//! Submitting a job is ordinary capacity dispatch: the router picks a cortex
|
||||
//! that serves the model and forwards the body, and that cortex owns the job
|
||||
//! from then on. Every later call names only a job id, which carries no hint
|
||||
//! of its cortex, so those calls are asked of each cortex in turn and the
|
||||
//! first answer that is not `404` wins. Listing merges every cortex's list.
|
||||
//!
|
||||
//! Ownership is enforced by the cortex holding the job; the router forwards
|
||||
//! the caller's bearer unchanged, exactly as for inference.
|
||||
|
||||
use crate::dispatch;
|
||||
use crate::error::envelope_response;
|
||||
use crate::state::RouterState;
|
||||
use axum::Router;
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::extract::{Path, State};
|
||||
use axum::http::{HeaderMap, Method, StatusCode};
|
||||
use axum::response::{IntoResponse, Json, Response};
|
||||
use axum::routing::get;
|
||||
use cortex_core::audio::{AudioJob, AudioJobList};
|
||||
use cortex_core::error_envelope::OpenAiError;
|
||||
use std::sync::Arc;
|
||||
|
||||
const PATH: &str = "/v1/audio/generations";
|
||||
|
||||
pub fn routes() -> Router<Arc<RouterState>> {
|
||||
Router::new()
|
||||
.route(PATH, get(list).post(submit))
|
||||
.route("/v1/audio/generations/{id}", get(get_job).delete(cancel))
|
||||
.route("/v1/audio/generations/{id}/content", get(content))
|
||||
}
|
||||
|
||||
async fn submit(
|
||||
State(state): State<Arc<RouterState>>,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Response {
|
||||
dispatch::dispatch(&state, PATH, headers, body).await
|
||||
}
|
||||
|
||||
async fn get_job(
|
||||
State(state): State<Arc<RouterState>>,
|
||||
headers: HeaderMap,
|
||||
Path(id): Path<String>,
|
||||
) -> Response {
|
||||
find(&state, Method::GET, &format!("{PATH}/{id}"), &headers, &id).await
|
||||
}
|
||||
|
||||
async fn cancel(
|
||||
State(state): State<Arc<RouterState>>,
|
||||
headers: HeaderMap,
|
||||
Path(id): Path<String>,
|
||||
) -> Response {
|
||||
find(
|
||||
&state,
|
||||
Method::DELETE,
|
||||
&format!("{PATH}/{id}"),
|
||||
&headers,
|
||||
&id,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn content(
|
||||
State(state): State<Arc<RouterState>>,
|
||||
headers: HeaderMap,
|
||||
Path(id): Path<String>,
|
||||
) -> Response {
|
||||
find(
|
||||
&state,
|
||||
Method::GET,
|
||||
&format!("{PATH}/{id}/content"),
|
||||
&headers,
|
||||
&id,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Ask each cortex for the job until one answers with something other than
|
||||
/// `404`, and stream that answer back verbatim.
|
||||
async fn find(
|
||||
state: &RouterState,
|
||||
method: Method,
|
||||
path: &str,
|
||||
headers: &HeaderMap,
|
||||
id: &str,
|
||||
) -> Response {
|
||||
for ep in &state.cortexes {
|
||||
let Some(client) = state.client_for(&ep.name) else {
|
||||
continue;
|
||||
};
|
||||
let url = format!("{}{path}", ep.endpoint);
|
||||
let resp = match client
|
||||
.request(method.clone(), &url)
|
||||
.headers(forward_headers(headers))
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
tracing::warn!(cortex = %ep.name, url = %url, error = %e, "cortex unreachable");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
if resp.status() == reqwest::StatusCode::NOT_FOUND {
|
||||
continue;
|
||||
}
|
||||
let status =
|
||||
StatusCode::from_u16(resp.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
|
||||
let mut out = Response::builder().status(status);
|
||||
for (k, v) in resp.headers() {
|
||||
if k == "transfer-encoding" || k == "connection" {
|
||||
continue;
|
||||
}
|
||||
out = out.header(k, v);
|
||||
}
|
||||
return out
|
||||
.body(Body::from_stream(resp.bytes_stream()))
|
||||
.unwrap_or_else(|_| StatusCode::BAD_GATEWAY.into_response());
|
||||
}
|
||||
envelope_response(OpenAiError::new(
|
||||
404,
|
||||
"invalid_request_error",
|
||||
"job_not_found",
|
||||
format!("no audio job '{id}'"),
|
||||
))
|
||||
}
|
||||
|
||||
/// Merge every reachable cortex's job list, newest first.
|
||||
async fn list(State(state): State<Arc<RouterState>>, headers: HeaderMap) -> Response {
|
||||
let mut data: Vec<AudioJob> = Vec::new();
|
||||
for ep in &state.cortexes {
|
||||
let Some(client) = state.client_for(&ep.name) else {
|
||||
continue;
|
||||
};
|
||||
let url = format!("{}{PATH}", ep.endpoint);
|
||||
let page = client
|
||||
.get(&url)
|
||||
.headers(forward_headers(&headers))
|
||||
.send()
|
||||
.await
|
||||
.ok()
|
||||
.filter(|r| r.status().is_success());
|
||||
let Some(page) = page else {
|
||||
tracing::warn!(cortex = %ep.name, "audio job list unavailable from cortex");
|
||||
continue;
|
||||
};
|
||||
if let Ok(list) = page.json::<AudioJobList>().await {
|
||||
data.extend(list.data);
|
||||
}
|
||||
}
|
||||
data.sort_by(|a, b| b.created_at.cmp(&a.created_at).then(b.id.cmp(&a.id)));
|
||||
Json(AudioJobList {
|
||||
object: "list".into(),
|
||||
data,
|
||||
})
|
||||
.into_response()
|
||||
}
|
||||
|
||||
fn forward_headers(headers: &HeaderMap) -> HeaderMap {
|
||||
let mut out = HeaderMap::new();
|
||||
for (k, v) in headers {
|
||||
if matches!(
|
||||
k.as_str(),
|
||||
"host" | "content-length" | "connection" | "transfer-encoding"
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
out.append(k.clone(), v.clone());
|
||||
}
|
||||
out
|
||||
}
|
||||
@@ -168,6 +168,8 @@ mod tests {
|
||||
capabilities: vec![],
|
||||
limit: None,
|
||||
cost: None,
|
||||
audio: None,
|
||||
audio_tokens_per_second: None,
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
max_model_len: None,
|
||||
|
||||
@@ -14,6 +14,7 @@ pub fn api_routes() -> Router<Arc<RouterState>> {
|
||||
.route("/v1/chat/completions", post(chat_completions))
|
||||
.route("/v1/completions", post(completions))
|
||||
.route("/v1/images/generations", post(images_generations))
|
||||
.merge(crate::audio::routes())
|
||||
.route("/v1/systemone", post(systemone))
|
||||
.route("/v1/responses", post(responses))
|
||||
.route("/v1/messages", post(messages))
|
||||
|
||||
@@ -11,6 +11,7 @@
|
||||
//! on capacity (epic #69). A background [`poller`] keeps a live
|
||||
//! per-cortex topology (#72) that the dispatcher (#73) will route on.
|
||||
|
||||
pub mod audio;
|
||||
pub mod catalogue;
|
||||
pub mod config;
|
||||
pub mod dispatch;
|
||||
|
||||
@@ -0,0 +1,187 @@
|
||||
//! Audio job lookups through the router (#382): a job id carries no hint of
|
||||
//! its cortex, so job calls are asked of each cortex until one answers with
|
||||
//! something other than 404, and listings merge every cortex's jobs.
|
||||
|
||||
use axum::extract::Path;
|
||||
use axum::http::{HeaderMap, StatusCode};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::get;
|
||||
use axum::{Json, Router};
|
||||
use helexa_router::config::{CortexEndpoint, RouterConfig};
|
||||
use helexa_router::state::RouterState;
|
||||
use serde_json::{Value, json};
|
||||
use std::sync::Arc;
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
/// A cortex that holds exactly one job, `job_id`, created at `created`.
|
||||
async fn spawn_cortex(job_id: &'static str, created: u64) -> String {
|
||||
let job = move || {
|
||||
json!({"id": job_id, "object": "audio.generation", "status": "succeeded",
|
||||
"model": "m", "created_at": created})
|
||||
};
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/v1/audio/generations",
|
||||
get(move |headers: HeaderMap| async move {
|
||||
// Only the bearer's own jobs are listed.
|
||||
if headers.get("authorization").is_none() {
|
||||
return Json(json!({"object": "list", "data": []})).into_response();
|
||||
}
|
||||
Json(json!({"object": "list", "data": [job()]})).into_response()
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/audio/generations/{id}",
|
||||
get(move |Path(id): Path<String>| async move {
|
||||
if id == job_id {
|
||||
Json(job()).into_response()
|
||||
} else {
|
||||
not_found()
|
||||
}
|
||||
})
|
||||
.delete(move |Path(id): Path<String>| async move {
|
||||
if id == job_id {
|
||||
let mut j = job();
|
||||
j["status"] = json!("cancelled");
|
||||
Json(j).into_response()
|
||||
} else {
|
||||
not_found()
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/audio/generations/{id}/content",
|
||||
get(move |Path(id): Path<String>| async move {
|
||||
if id == job_id {
|
||||
(
|
||||
StatusCode::OK,
|
||||
[("content-type", "audio/wav")],
|
||||
format!("audio-of-{job_id}"),
|
||||
)
|
||||
.into_response()
|
||||
} else {
|
||||
not_found()
|
||||
}
|
||||
}),
|
||||
);
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
|
||||
format!("http://{addr}")
|
||||
}
|
||||
|
||||
fn not_found() -> Response {
|
||||
(
|
||||
StatusCode::NOT_FOUND,
|
||||
Json(json!({"error": {"type": "invalid_request_error", "code": "job_not_found", "message": "no"}})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
async fn router_over(cortexes: Vec<(&str, String)>) -> String {
|
||||
let cfg = RouterConfig {
|
||||
cortexes: cortexes
|
||||
.into_iter()
|
||||
.map(|(name, endpoint)| CortexEndpoint {
|
||||
name: name.into(),
|
||||
endpoint,
|
||||
region: None,
|
||||
tls_ca: None,
|
||||
})
|
||||
.collect(),
|
||||
..Default::default()
|
||||
};
|
||||
let state = Arc::new(RouterState::from_config(&cfg));
|
||||
let app = helexa_router::handlers::api_routes().with_state(state);
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
|
||||
format!("http://{addr}")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn job_calls_find_the_cortex_that_holds_the_job() {
|
||||
let a = spawn_cortex("audiogen_a", 100).await;
|
||||
let b = spawn_cortex("audiogen_b", 200).await;
|
||||
// An unreachable cortex in the list must be skipped, not fatal.
|
||||
let router = router_over(vec![
|
||||
("dead", "http://127.0.0.1:1".into()),
|
||||
("a", a),
|
||||
("b", b),
|
||||
])
|
||||
.await;
|
||||
let http = reqwest::Client::new();
|
||||
|
||||
let job: Value = http
|
||||
.get(format!("{router}/v1/audio/generations/audiogen_b"))
|
||||
.send()
|
||||
.await
|
||||
.unwrap()
|
||||
.json()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(job["id"], "audiogen_b");
|
||||
|
||||
let content = http
|
||||
.get(format!("{router}/v1/audio/generations/audiogen_b/content"))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(content.status(), 200);
|
||||
assert_eq!(content.headers()["content-type"], "audio/wav");
|
||||
assert_eq!(content.text().await.unwrap(), "audio-of-audiogen_b");
|
||||
|
||||
let cancelled: Value = http
|
||||
.delete(format!("{router}/v1/audio/generations/audiogen_a"))
|
||||
.send()
|
||||
.await
|
||||
.unwrap()
|
||||
.json()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(cancelled["status"], "cancelled");
|
||||
|
||||
let missing = http
|
||||
.get(format!("{router}/v1/audio/generations/audiogen_nope"))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(missing.status(), 404);
|
||||
let body: Value = missing.json().await.unwrap();
|
||||
assert_eq!(body["error"]["code"], "job_not_found");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn listing_merges_cortexes_newest_first_and_forwards_the_bearer() {
|
||||
let a = spawn_cortex("audiogen_a", 100).await;
|
||||
let b = spawn_cortex("audiogen_b", 200).await;
|
||||
let router = router_over(vec![("a", a), ("b", b)]).await;
|
||||
let http = reqwest::Client::new();
|
||||
|
||||
let list: Value = http
|
||||
.get(format!("{router}/v1/audio/generations"))
|
||||
.bearer_auth("k")
|
||||
.send()
|
||||
.await
|
||||
.unwrap()
|
||||
.json()
|
||||
.await
|
||||
.unwrap();
|
||||
let ids: Vec<&str> = list["data"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|j| j["id"].as_str().unwrap())
|
||||
.collect();
|
||||
assert_eq!(ids, ["audiogen_b", "audiogen_a"]);
|
||||
|
||||
let anon: Value = http
|
||||
.get(format!("{router}/v1/audio/generations"))
|
||||
.send()
|
||||
.await
|
||||
.unwrap()
|
||||
.json()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(anon["data"].as_array().unwrap().len(), 0);
|
||||
}
|
||||
@@ -37,6 +37,8 @@ fn model_entry(loaded: bool, feasible: bool) -> CortexModelEntry {
|
||||
capabilities: vec![],
|
||||
limit: None,
|
||||
cost: None,
|
||||
audio: None,
|
||||
audio_tokens_per_second: None,
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
max_model_len: None,
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
-- ── Per-model rates for non-token units (#386) ───────────────────────
|
||||
--
|
||||
-- The ledger counts tokens. Work metered in another unit — audio-seconds
|
||||
-- for music generation today — is converted into ledger tokens at a
|
||||
-- per-model rate kept here, so an operator can price a model, or change
|
||||
-- its price, without a deploy. cortex fetches the table from
|
||||
-- `/authz/v1/rates` and reserves and settles at the rate it read.
|
||||
--
|
||||
-- `model_id = '*'` is the default for a unit: it applies to every model
|
||||
-- without a row of its own. A unit with no row at all, for the model or
|
||||
-- for '*', is not priced, and cortex refuses that work rather than serving
|
||||
-- it free.
|
||||
CREATE TABLE model_rates (
|
||||
model_id TEXT NOT NULL,
|
||||
-- 'audio_second' — one second of generated audio.
|
||||
unit TEXT NOT NULL,
|
||||
-- Ledger tokens charged per unit.
|
||||
tokens_per_unit BIGINT NOT NULL,
|
||||
description TEXT,
|
||||
-- Audit: who changed it and when (NULL updated_by = shipped default).
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_by UUID REFERENCES users(id) ON DELETE SET NULL,
|
||||
|
||||
PRIMARY KEY (model_id, unit),
|
||||
CONSTRAINT model_rates_unit_known CHECK (unit IN ('audio_second')),
|
||||
CONSTRAINT model_rates_positive CHECK (tokens_per_unit > 0)
|
||||
);
|
||||
|
||||
-- Placeholder default while music generation is being built; set real
|
||||
-- per-model prices with an UPDATE or per-model INSERT.
|
||||
INSERT INTO model_rates (model_id, unit, tokens_per_unit, description)
|
||||
VALUES ('*', 'audio_second', 1000, 'default audio rate (placeholder)');
|
||||
@@ -21,7 +21,7 @@ use axum::extract::{Request, State};
|
||||
use axum::http::{StatusCode, header};
|
||||
use axum::middleware::Next;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::post;
|
||||
use axum::routing::{get, post};
|
||||
use axum::{Extension, Json, Router};
|
||||
use cortex_core::error_envelope::OpenAiError;
|
||||
use serde::{Deserialize, Serialize};
|
||||
@@ -42,6 +42,7 @@ pub fn router(state: &AppState) -> Router<AppState> {
|
||||
.route("/authz/v1/release", post(release))
|
||||
.route("/authz/v1/snapshot", post(snapshot))
|
||||
.route("/authz/v1/served-usage", post(served_usage))
|
||||
.route("/authz/v1/rates", get(rates))
|
||||
.layer(axum::middleware::from_fn_with_state(
|
||||
state.clone(),
|
||||
client_auth,
|
||||
@@ -338,3 +339,23 @@ fn bad_request(msg: &str) -> Response {
|
||||
msg,
|
||||
))
|
||||
}
|
||||
|
||||
// ── rates (#386) ────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct RatesResp {
|
||||
data: Vec<crate::rates::ModelRate>,
|
||||
}
|
||||
|
||||
/// `GET /authz/v1/rates` — every per-model rate, for cortex to cache and
|
||||
/// resolve locally. A database error is a 500: cortex keeps its last good
|
||||
/// table and refuses unpriced work rather than guessing.
|
||||
async fn rates(State(state): State<AppState>) -> Response {
|
||||
match crate::rates::list(&state.pool).await {
|
||||
Ok(data) => Json(RatesResp { data }).into_response(),
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "rates query failed");
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, "rates unavailable").into_response()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -21,6 +21,7 @@ pub mod email;
|
||||
pub mod error;
|
||||
pub mod handlers;
|
||||
pub mod ledger;
|
||||
pub mod rates;
|
||||
pub mod reconcile;
|
||||
pub mod state;
|
||||
pub mod topup;
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
//! Per-model rates for work metered in units other than tokens (#386).
|
||||
//!
|
||||
//! The ledger only counts tokens; a rate converts another unit (one second
|
||||
//! of generated audio) into ledger tokens. Rates live in `model_rates` so
|
||||
//! they can change without a deploy. cortex reads the whole table from
|
||||
//! `/authz/v1/rates` and resolves a model's rate itself: the model's own
|
||||
//! row, else the unit's `'*'` default, else the work is not priced.
|
||||
|
||||
use serde::Serialize;
|
||||
use sqlx::postgres::PgPool;
|
||||
|
||||
/// One row of `model_rates`, as served to cortex.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct ModelRate {
|
||||
pub model_id: String,
|
||||
pub unit: String,
|
||||
pub tokens_per_unit: i64,
|
||||
}
|
||||
|
||||
/// Every configured rate.
|
||||
pub async fn list(pool: &PgPool) -> Result<Vec<ModelRate>, sqlx::Error> {
|
||||
let rows = sqlx::query(
|
||||
"SELECT model_id, unit, tokens_per_unit FROM model_rates ORDER BY unit, model_id",
|
||||
)
|
||||
.fetch_all(pool)
|
||||
.await?;
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.map(|r| ModelRate {
|
||||
model_id: sqlx::Row::get(&r, "model_id"),
|
||||
unit: sqlx::Row::get(&r, "unit"),
|
||||
tokens_per_unit: sqlx::Row::get(&r, "tokens_per_unit"),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
//! `model_rates` (#386) against a real PostgreSQL. Gated on
|
||||
//! `UPSTREAM_TEST_DATABASE_URL`, like the ledger tests: unset, every test
|
||||
//! logs a skip and passes.
|
||||
|
||||
use helexa_upstream::db::connect_and_migrate;
|
||||
use helexa_upstream::rates;
|
||||
use sqlx::postgres::PgPool;
|
||||
|
||||
async fn pool_or_skip(test: &str) -> Option<PgPool> {
|
||||
let Ok(url) = std::env::var("UPSTREAM_TEST_DATABASE_URL") else {
|
||||
eprintln!("skipping {test}: UPSTREAM_TEST_DATABASE_URL not set");
|
||||
return None;
|
||||
};
|
||||
Some(
|
||||
connect_and_migrate(&url, 4)
|
||||
.await
|
||||
.expect("connect + migrate"),
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn default_audio_rate_is_seeded_and_model_rows_are_listed() {
|
||||
let Some(pool) = pool_or_skip("default_audio_rate_is_seeded").await else {
|
||||
return;
|
||||
};
|
||||
let model = format!("test/model-{}", uuid::Uuid::new_v4());
|
||||
sqlx::query(
|
||||
"INSERT INTO model_rates (model_id, unit, tokens_per_unit) VALUES ($1, 'audio_second', 2500)",
|
||||
)
|
||||
.bind(&model)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let all = rates::list(&pool).await.unwrap();
|
||||
assert!(
|
||||
all.iter()
|
||||
.any(|r| r.model_id == "*" && r.unit == "audio_second" && r.tokens_per_unit > 0),
|
||||
"seeded default missing: {all:?}"
|
||||
);
|
||||
assert!(
|
||||
all.iter()
|
||||
.any(|r| r.model_id == model && r.tokens_per_unit == 2500)
|
||||
);
|
||||
|
||||
sqlx::query("DELETE FROM model_rates WHERE model_id = $1")
|
||||
.bind(&model)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rates_must_be_positive_and_of_a_known_unit() {
|
||||
let Some(pool) = pool_or_skip("rates_must_be_positive").await else {
|
||||
return;
|
||||
};
|
||||
let insert = |unit: &'static str, tokens: i64| {
|
||||
let pool = pool.clone();
|
||||
async move {
|
||||
sqlx::query(
|
||||
"INSERT INTO model_rates (model_id, unit, tokens_per_unit) VALUES ('test/bad', $1, $2)",
|
||||
)
|
||||
.bind(unit)
|
||||
.bind(tokens)
|
||||
.execute(&pool)
|
||||
.await
|
||||
}
|
||||
};
|
||||
assert!(insert("audio_second", 0).await.is_err());
|
||||
assert!(insert("audio_second", -5).await.is_err());
|
||||
assert!(insert("furlongs", 10).await.is_err());
|
||||
}
|
||||
@@ -76,6 +76,12 @@ toml.workspace = true
|
||||
# transitively via candle-core.
|
||||
rayon = "1"
|
||||
|
||||
# Unicode tables for the decision-model language detector (#339), which
|
||||
# must classify characters exactly as Python's `str.isalpha` / `re` and
|
||||
# `unicodedata.combining` do. Both already in the tree via tokenizers.
|
||||
unicode-normalization = "0.1"
|
||||
unicode-properties = "0.1"
|
||||
|
||||
# candle for in-process inference. CUDA support is gated behind the
|
||||
# crate's `cuda` feature (default off) so the workspace builds on
|
||||
# non-CUDA hosts and CI runners.
|
||||
|
||||
@@ -55,6 +55,8 @@ pub fn neuron_routes() -> Router<Arc<NeuronState>> {
|
||||
.route("/models/{model_id}/capabilities", get(model_capabilities))
|
||||
.route("/v1/chat/completions", post(chat_completions))
|
||||
.route("/v1/images/generations", post(images_generations))
|
||||
.route("/v1/audio/generations", post(audio_generations))
|
||||
.route("/v1/systemone", post(systemone))
|
||||
.route("/v1/responses", post(responses))
|
||||
}
|
||||
|
||||
@@ -189,6 +191,7 @@ fn preflight_kind(err: &PreflightError) -> &'static str {
|
||||
PreflightError::EmptyRepo { .. } => "empty_repo",
|
||||
PreflightError::TpRequiresSafetensors { .. } => "tp_requires_safetensors",
|
||||
PreflightError::QuantNotFound { .. } => "quant_not_found",
|
||||
PreflightError::Unsupported { .. } => "unsupported_load_option",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -668,6 +671,138 @@ async fn images_generations(
|
||||
}
|
||||
}
|
||||
|
||||
/// `POST /v1/audio/generations` (#382): render one audio generation and
|
||||
/// return it. This is the neuron half of the async job API — cortex owns
|
||||
/// the job and calls this once per job, waiting for the whole render. The
|
||||
/// body is the audio (`audio/wav`); the result metadata travels in the
|
||||
/// [`cortex_core::audio::RESULT_HEADER`] header so the bytes need no
|
||||
/// base64 envelope.
|
||||
async fn audio_generations(
|
||||
State(state): State<Arc<NeuronState>>,
|
||||
headers: axum::http::HeaderMap,
|
||||
ApiJson(req): ApiJson<cortex_core::audio::AudioGenerationRequest>,
|
||||
) -> axum::response::Response {
|
||||
use cortex_core::audio::RESULT_HEADER;
|
||||
use cortex_core::error_envelope::OpenAiError;
|
||||
|
||||
let Some(candle) = state.candle.as_ref().map(Arc::clone) else {
|
||||
return (
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
Json(json!({"error": "candle harness not enabled on this neuron"})),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
if let Err(msg) = req.validate_shape() {
|
||||
return envelope_response(OpenAiError::new(
|
||||
400,
|
||||
"invalid_request_error",
|
||||
"invalid_audio_params",
|
||||
format!("invalid audio generation parameters: {msg}"),
|
||||
));
|
||||
}
|
||||
let principal = principal_key(&headers);
|
||||
match candle.audio_generation(&req, principal.as_deref()).await {
|
||||
Ok((wav, result)) => {
|
||||
let meta = match serde_json::to_string(&result) {
|
||||
Ok(m) => m,
|
||||
Err(e) => {
|
||||
return inference_error_response(InferenceError::Other(anyhow::anyhow!(
|
||||
"serialise audio result: {e}"
|
||||
)));
|
||||
}
|
||||
};
|
||||
let mut resp = axum::response::Response::new(axum::body::Body::from(wav));
|
||||
let h = resp.headers_mut();
|
||||
h.insert(
|
||||
axum::http::header::CONTENT_TYPE,
|
||||
axum::http::HeaderValue::from_static("audio/wav"),
|
||||
);
|
||||
if let Ok(v) = axum::http::HeaderValue::from_str(&meta) {
|
||||
h.insert(RESULT_HEADER, v);
|
||||
}
|
||||
resp
|
||||
}
|
||||
Err(e) => inference_error_response(e),
|
||||
}
|
||||
}
|
||||
|
||||
/// `POST /v1/systemone` (#338): typed decisions over the TypeSafe Jev
|
||||
/// wire protocol, from a loaded decision model (Laya). Parsing and the
|
||||
/// transport-level limits happen here; routing, validation, admission,
|
||||
/// scoring and shaping in `CandleHarness::systemone`.
|
||||
///
|
||||
/// Errors carry the #63 envelope *and* the reference server's top-level
|
||||
/// `detail` string, so OpenAI-style and Jev/laya clients both read them.
|
||||
async fn systemone(
|
||||
State(state): State<Arc<NeuronState>>,
|
||||
headers: axum::http::HeaderMap,
|
||||
body: axum::body::Bytes,
|
||||
) -> axum::response::Response {
|
||||
use crate::harness::candle::SystemOneError;
|
||||
use cortex_core::decisions::{SystemOneRequest, limits};
|
||||
|
||||
let Some(candle) = state.candle.as_ref().map(Arc::clone) else {
|
||||
return decision_error_response(503, "candle harness not enabled on this neuron");
|
||||
};
|
||||
if body.len() > limits::MAX_BODY_BYTES {
|
||||
return decision_error_response(413, "request body too large");
|
||||
}
|
||||
let request = match SystemOneRequest::parse(&body, limits::DEFAULT_MAX_TOKEN_BUDGET) {
|
||||
Ok(r) => r,
|
||||
Err(e) => return decision_error_response(e.status, &e.detail),
|
||||
};
|
||||
let served = headers
|
||||
.get(cortex_core::catalogue::DECISION_MODEL_HEADER)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(str::to_string);
|
||||
let principal = principal_key(&headers);
|
||||
|
||||
let started = std::time::Instant::now();
|
||||
match candle
|
||||
.systemone(served.as_deref(), request, principal.as_deref())
|
||||
.await
|
||||
{
|
||||
Ok(response) => {
|
||||
let ms = started.elapsed().as_secs_f64() * 1e3;
|
||||
let mut out = Json(response).into_response();
|
||||
for (name, value) in [
|
||||
("server-timing", format!("inference;dur={ms:.2}")),
|
||||
("x-inference-time-ms", format!("{ms:.2}")),
|
||||
] {
|
||||
if let Ok(v) = axum::http::HeaderValue::from_str(&value) {
|
||||
out.headers_mut().insert(name, v);
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
Err(SystemOneError::Request(e)) => decision_error_response(e.status, &e.detail),
|
||||
Err(SystemOneError::Inference(e)) => {
|
||||
tracing::warn!(error = %e, "systemone request failed");
|
||||
inference_error_response(e)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A rejected decision request: the #63 envelope, plus the `detail` field
|
||||
/// the reference (FastAPI) server answers with.
|
||||
fn decision_error_response(status: u16, detail: &str) -> axum::response::Response {
|
||||
use cortex_core::error_envelope::OpenAiError;
|
||||
let (kind, code) = match status {
|
||||
404 => ("invalid_request_error", "model_not_found"),
|
||||
413 => ("invalid_request_error", "request_too_large"),
|
||||
422 => ("invalid_request_error", "invalid_decision_request"),
|
||||
400..=499 => ("invalid_request_error", "invalid_request"),
|
||||
_ => ("api_error", "service_unavailable"),
|
||||
};
|
||||
let env = OpenAiError::new(status, kind, code, detail.to_string());
|
||||
let mut body = env.body();
|
||||
if let Some(obj) = body.as_object_mut() {
|
||||
obj.insert("detail".into(), json!(detail));
|
||||
}
|
||||
let status = StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_REQUEST);
|
||||
(status, Json(body)).into_response()
|
||||
}
|
||||
|
||||
fn finish_reason_from_str(s: &str) -> crate::wire::FinishReason {
|
||||
use crate::wire::FinishReason;
|
||||
match s {
|
||||
|
||||
@@ -0,0 +1,478 @@
|
||||
//! ModernBERT, the bidirectional encoder under every Laya checkpoint.
|
||||
//!
|
||||
//! Follows the structure of `candle_transformers::models::modernbert`
|
||||
//! (itself a port of transformers' `ModernBertModel`), vendored rather
|
||||
//! than used as-is because that module only runs correctly in f32:
|
||||
//!
|
||||
//! - its attention masks are built in F32 whatever the model dtype, so
|
||||
//! an f16/bf16 load fails at the first `broadcast_add`;
|
||||
//! - it computes the rotary angles `position × inv_freq` in the model
|
||||
//! dtype, which in f16 already loses a quarter-radian of phase by
|
||||
//! position 500;
|
||||
//! - it loads each layer's `attn_norm` with `.ok()`, so a weight-prefix
|
||||
//! mistake silently yields a model with no pre-attention norms.
|
||||
//!
|
||||
//! Here the angles are computed in f32 and only the resulting cos/sin
|
||||
//! tables are cast; attention scores are softmaxed in f32 against f32
|
||||
//! masks; and `attn_norm` is required on every layer except layer 0,
|
||||
//! which ModernBERT builds as an identity.
|
||||
//!
|
||||
//! Precision follows PyTorch autocast, which is how the reference runs on
|
||||
//! a GPU: the linear layers (and so the attention inputs) run in the
|
||||
//! load dtype, while the residual stream, the embeddings and every
|
||||
//! LayerNorm stay in f32. Casting the residual stream to bf16 as well
|
||||
//! drifts the released English checkpoint's logits by ~0.25 over its 28
|
||||
//! layers; keeping it in f32 is what autocast does and costs little.
|
||||
//!
|
||||
//! Masked positions use a large *finite* negative rather than `-inf`.
|
||||
//! A padded query near the end of a short row can find every key in
|
||||
//! its sliding window padded; with `-inf` that row's softmax is NaN,
|
||||
//! and NaN in a padded row's value vector poisons the valid rows that
|
||||
//! give it zero weight (`0 × NaN = NaN`). A finite value makes such a
|
||||
//! row uniform instead, and valid rows never look at it.
|
||||
|
||||
use anyhow::{Context, Result, ensure};
|
||||
use candle_core::{D, DType, Device, Module, Tensor};
|
||||
use candle_nn::{Embedding, LayerNorm, Linear, VarBuilder};
|
||||
|
||||
use super::config::{EncoderConfig, LayerKind};
|
||||
|
||||
/// Additive mask value for a position that must not be attended.
|
||||
pub(crate) const MASKED: f32 = -1e9;
|
||||
|
||||
/// Apply `lin` in its weight's dtype, casting `x` in if needed. The
|
||||
/// result stays in that dtype; callers cast back to f32 where it joins
|
||||
/// the residual stream.
|
||||
pub(crate) fn project(lin: &Linear, x: &Tensor) -> candle_core::Result<Tensor> {
|
||||
x.to_dtype(lin.weight().dtype())?.apply(lin)
|
||||
}
|
||||
|
||||
/// A bias-free LayerNorm that takes candle's fused kernel.
|
||||
///
|
||||
/// `candle_nn::LayerNorm` dispatches to the fused `layer_norm` op only
|
||||
/// when it has a bias; without one it falls back to six elementwise
|
||||
/// kernels (mean, subtract, square, mean, divide, scale), which was a
|
||||
/// quarter of the encoder's GPU time. ModernBERT's norms have no bias,
|
||||
/// so they get an explicit zero one: same result, one kernel.
|
||||
fn fused_layer_norm_no_bias(size: usize, eps: f64, vb: VarBuilder) -> Result<LayerNorm> {
|
||||
let weight = vb.get(size, "weight")?;
|
||||
let bias = weight.zeros_like()?;
|
||||
Ok(LayerNorm::new(weight, bias, eps))
|
||||
}
|
||||
|
||||
/// Precomputed rotary tables for one base, `(max_pos, head_dim / 2)`.
|
||||
struct Rotary {
|
||||
cos: Tensor,
|
||||
sin: Tensor,
|
||||
}
|
||||
|
||||
impl Rotary {
|
||||
fn new(cfg: &EncoderConfig, theta: f64, dtype: DType, device: &Device) -> Result<Self> {
|
||||
let dim = cfg.head_dim();
|
||||
let inv_freq: Vec<f32> = (0..dim)
|
||||
.step_by(2)
|
||||
.map(|i| (1.0 / theta.powf(i as f64 / dim as f64)) as f32)
|
||||
.collect();
|
||||
let n = inv_freq.len();
|
||||
let inv_freq = Tensor::from_vec(inv_freq, (1, n), device)?;
|
||||
let max_pos = cfg.max_position_embeddings;
|
||||
let pos = Tensor::arange(0u32, max_pos as u32, device)?
|
||||
.to_dtype(DType::F32)?
|
||||
.reshape((max_pos, 1))?;
|
||||
let freqs = pos.matmul(&inv_freq)?;
|
||||
Ok(Self {
|
||||
cos: freqs.cos()?.to_dtype(dtype)?,
|
||||
sin: freqs.sin()?.to_dtype(dtype)?,
|
||||
})
|
||||
}
|
||||
|
||||
fn apply(&self, q: &Tensor, k: &Tensor, seq_len: usize) -> Result<(Tensor, Tensor)> {
|
||||
let cos = self.cos.narrow(0, 0, seq_len)?;
|
||||
let sin = self.sin.narrow(0, 0, seq_len)?;
|
||||
Ok((
|
||||
candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?,
|
||||
candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?,
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
struct Attention {
|
||||
wqkv: Linear,
|
||||
wo: Linear,
|
||||
n_heads: usize,
|
||||
head_dim: usize,
|
||||
kind: LayerKind,
|
||||
}
|
||||
|
||||
impl Attention {
|
||||
fn forward(&self, xs: &Tensor, rotary: &Rotary, mask: Option<&Tensor>) -> Result<Tensor> {
|
||||
let (b, l, d) = xs.dims3()?;
|
||||
let qkv = project(&self.wqkv, xs)?
|
||||
.reshape((b, l, 3, self.n_heads, self.head_dim))?
|
||||
.permute((2, 0, 3, 1, 4))?;
|
||||
// Rotary in the projection dtype, as under autocast; the scores
|
||||
// are then taken in f32 (see `scaled_attention`).
|
||||
let q = qkv.get(0)?;
|
||||
let k = qkv.get(1)?;
|
||||
let v = qkv.get(2)?.contiguous()?;
|
||||
let (q, k) = rotary.apply(&q, &k, l)?;
|
||||
let ctx = scaled_attention(&q, &k, &v, mask, self.head_dim)?;
|
||||
let ctx = ctx.transpose(1, 2)?.reshape((b, l, d))?;
|
||||
Ok(ctx.apply(&self.wo)?.to_dtype(DType::F32)?)
|
||||
}
|
||||
}
|
||||
|
||||
/// A batch packed without padding for the flash-attention path: every
|
||||
/// row's tokens back to back, `(total, …)`. Linear layers then do no work
|
||||
/// for padding, and flash varlen keeps rows from attending to each other.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
pub struct Packed {
|
||||
/// `(b + 1,)` u32 cumulative row lengths, as flash varlen takes them.
|
||||
pub cu_seqlens: Tensor,
|
||||
/// Longest row.
|
||||
pub max_len: usize,
|
||||
/// `(total,)` u32: each token's position within its own row.
|
||||
pub positions: Tensor,
|
||||
/// Where each row starts in the packed order.
|
||||
pub offsets: Vec<usize>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "flash-attn")]
|
||||
impl Packed {
|
||||
pub fn new(lens: &[usize], device: &Device) -> Result<Self> {
|
||||
let mut cu = Vec::with_capacity(lens.len() + 1);
|
||||
let mut offsets = Vec::with_capacity(lens.len());
|
||||
let mut positions = Vec::with_capacity(lens.iter().sum());
|
||||
let mut total = 0u32;
|
||||
cu.push(0u32);
|
||||
for &n in lens {
|
||||
offsets.push(total as usize);
|
||||
positions.extend(0..n as u32);
|
||||
total += n as u32;
|
||||
cu.push(total);
|
||||
}
|
||||
let n = cu.len();
|
||||
Ok(Self {
|
||||
cu_seqlens: Tensor::from_vec(cu, n, device)?,
|
||||
max_len: lens.iter().copied().max().unwrap_or(0),
|
||||
positions: Tensor::from_vec(positions, total as usize, device)?,
|
||||
offsets,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Bidirectional flash attention over a packed batch. `q`, `k`, `v` are
|
||||
/// `(total, h, dh)`; `half_window` limits each query to keys within that
|
||||
/// distance on either side (ModernBERT's sliding layers), `None` attends
|
||||
/// the whole row. Scores are accumulated in f32 inside the kernel.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
pub(crate) fn flash_attention(
|
||||
q: &Tensor,
|
||||
k: &Tensor,
|
||||
v: &Tensor,
|
||||
packed: &Packed,
|
||||
half_window: Option<usize>,
|
||||
) -> Result<Tensor> {
|
||||
let head_dim = q.dim(2)?;
|
||||
let scale = (head_dim as f64).powf(-0.5) as f32;
|
||||
Ok(candle_flash_attn::flash_attn_varlen_windowed(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
&packed.cu_seqlens,
|
||||
&packed.cu_seqlens,
|
||||
packed.max_len,
|
||||
packed.max_len,
|
||||
scale,
|
||||
half_window,
|
||||
half_window,
|
||||
)?)
|
||||
}
|
||||
|
||||
/// Split a packed `(total, 3·d)` qkv projection into contiguous
|
||||
/// `(total, h, dh)` q, k and v.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
pub(crate) fn split_qkv(
|
||||
qkv: &Tensor,
|
||||
n_heads: usize,
|
||||
head_dim: usize,
|
||||
) -> Result<(Tensor, Tensor, Tensor)> {
|
||||
let t = qkv.dim(0)?;
|
||||
let qkv = qkv.reshape((t, 3, n_heads, head_dim))?;
|
||||
let part = |i: usize| -> Result<Tensor> { Ok(qkv.narrow(1, i, 1)?.squeeze(1)?.contiguous()?) };
|
||||
Ok((part(0)?, part(1)?, part(2)?))
|
||||
}
|
||||
|
||||
#[cfg(feature = "flash-attn")]
|
||||
impl Rotary {
|
||||
/// Rotate packed `(total, h, dh)` q and k by each token's position.
|
||||
fn apply_packed(&self, q: &Tensor, k: &Tensor, positions: &Tensor) -> Result<(Tensor, Tensor)> {
|
||||
let cos = self.cos.index_select(positions, 0)?;
|
||||
let sin = self.sin.index_select(positions, 0)?;
|
||||
let rot = |x: &Tensor| -> Result<Tensor> {
|
||||
Ok(candle_nn::rotary_emb::rope_thd(&x.unsqueeze(0)?, &cos, &sin)?.squeeze(0)?)
|
||||
};
|
||||
Ok((rot(q)?, rot(k)?))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "flash-attn")]
|
||||
impl Attention {
|
||||
/// `xs` is the packed f32 residual stream `(total, d)`.
|
||||
fn forward_packed(
|
||||
&self,
|
||||
xs: &Tensor,
|
||||
rotary: &Rotary,
|
||||
packed: &Packed,
|
||||
half_window: usize,
|
||||
) -> Result<Tensor> {
|
||||
let (t, d) = xs.dims2()?;
|
||||
let (q, k, v) = split_qkv(&project(&self.wqkv, xs)?, self.n_heads, self.head_dim)?;
|
||||
let (q, k) = rotary.apply_packed(&q, &k, &packed.positions)?;
|
||||
let window = match self.kind {
|
||||
LayerKind::Sliding => Some(half_window),
|
||||
LayerKind::Global => None,
|
||||
};
|
||||
let ctx = flash_attention(&q, &k, &v, packed, window)?.reshape((t, d))?;
|
||||
Ok(ctx.apply(&self.wo)?.to_dtype(DType::F32)?)
|
||||
}
|
||||
}
|
||||
|
||||
/// `softmax(q kᵀ / √d + mask) v`, scores and softmax in f32.
|
||||
///
|
||||
/// `q`, `k`, `v` are `(b, h, l, d)`; `mask` is an f32 additive mask
|
||||
/// broadcastable to `(b, h, l, l)`, or `None` when nothing is masked —
|
||||
/// which saves a bandwidth-bound pass over the `(b, h, l, l)` scores.
|
||||
pub(crate) fn scaled_attention(
|
||||
q: &Tensor,
|
||||
k: &Tensor,
|
||||
v: &Tensor,
|
||||
mask: Option<&Tensor>,
|
||||
head_dim: usize,
|
||||
) -> Result<Tensor> {
|
||||
let dtype = v.dtype();
|
||||
let scale = (head_dim as f64).powf(-0.5);
|
||||
// Scores in f32 end to end: a bf16 `q kᵀ` is rounded to 8 mantissa
|
||||
// bits before the softmax ever sees it, which on this model's
|
||||
// logit-sized scores is most of the drift from the reference's GPU
|
||||
// path (whose SDPA kernels keep scores in f32).
|
||||
// The scale goes on q, not on the `(b, h, l, l)` scores: the same
|
||||
// product, on a tensor `l / d` times smaller.
|
||||
let scores = (q.to_dtype(DType::F32)? * scale)?.matmul(&k.to_dtype(DType::F32)?.t()?)?;
|
||||
let scores = match mask {
|
||||
Some(mask) => scores.broadcast_add(mask)?,
|
||||
None => scores,
|
||||
};
|
||||
let probs = candle_nn::ops::softmax_last_dim(&scores)?;
|
||||
Ok(probs.to_dtype(dtype)?.matmul(v)?)
|
||||
}
|
||||
|
||||
struct Mlp {
|
||||
wi: Linear,
|
||||
wo: Linear,
|
||||
}
|
||||
|
||||
impl Module for Mlp {
|
||||
/// f32 in, f32 out; the projections run in the weights' dtype.
|
||||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||||
let xs = project(&self.wi, xs)?;
|
||||
let parts = xs.chunk(2, D::Minus1)?;
|
||||
(parts[0].gelu_erf()? * &parts[1])?
|
||||
.apply(&self.wo)?
|
||||
.to_dtype(DType::F32)
|
||||
}
|
||||
}
|
||||
|
||||
struct Layer {
|
||||
attn_norm: Option<LayerNorm>,
|
||||
attn: Attention,
|
||||
mlp_norm: LayerNorm,
|
||||
mlp: Mlp,
|
||||
}
|
||||
|
||||
/// The two additive masks one forward needs, both f32. Either is `None`
|
||||
/// when it would mask nothing: `global` for a batch with no padding (every
|
||||
/// single-question request), `sliding` additionally when no row is longer
|
||||
/// than the window reaches.
|
||||
pub struct Masks {
|
||||
/// `(b, 1, 1, l)`: padded keys masked. Global layers and the
|
||||
/// decision head use this one.
|
||||
pub global: Option<Tensor>,
|
||||
/// `(b, 1, l, l)`: padded keys, plus keys outside the window.
|
||||
pub sliding: Option<Tensor>,
|
||||
}
|
||||
|
||||
impl Masks {
|
||||
/// Build both masks from per-row valid lengths. Rows are
|
||||
/// right-padded: positions `>= lens[i]` of row `i` are padding.
|
||||
pub fn new(
|
||||
lens: &[usize],
|
||||
seq_len: usize,
|
||||
half_window: usize,
|
||||
device: &Device,
|
||||
) -> Result<Self> {
|
||||
let b = lens.len();
|
||||
let padded = lens.iter().any(|&n| n < seq_len);
|
||||
let windowed = seq_len > half_window + 1;
|
||||
let global = if padded {
|
||||
let mut pad = Vec::with_capacity(b * seq_len);
|
||||
for &n in lens {
|
||||
pad.extend((0..seq_len).map(|j| if j < n { 0f32 } else { MASKED }));
|
||||
}
|
||||
Some(Tensor::from_vec(pad, (b, 1, 1, seq_len), device)?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if !windowed {
|
||||
let sliding = global.clone();
|
||||
return Ok(Self { global, sliding });
|
||||
}
|
||||
// Built on the device from two aranges, not on the host: at
|
||||
// 1024 tokens a host-built window is a 4 MB upload per forward.
|
||||
let pos = Tensor::arange(0u32, seq_len as u32, device)?.to_dtype(DType::F32)?;
|
||||
let dist = pos
|
||||
.reshape((seq_len, 1))?
|
||||
.broadcast_sub(&pos.reshape((1, seq_len))?)?
|
||||
.abs()?;
|
||||
let outside = dist.gt(half_window as f64)?;
|
||||
let zeros = Tensor::zeros((seq_len, seq_len), DType::F32, device)?;
|
||||
let masked = Tensor::full(MASKED, (seq_len, seq_len), device)?;
|
||||
let window = outside
|
||||
.where_cond(&masked, &zeros)?
|
||||
.reshape((1, 1, seq_len, seq_len))?;
|
||||
let sliding = match &global {
|
||||
Some(g) => g.broadcast_add(&window)?,
|
||||
None => window,
|
||||
};
|
||||
Ok(Self {
|
||||
global,
|
||||
sliding: Some(sliding),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ModernBert {
|
||||
embeddings: Embedding,
|
||||
embed_norm: LayerNorm,
|
||||
layers: Vec<Layer>,
|
||||
final_norm: LayerNorm,
|
||||
global_rotary: Rotary,
|
||||
sliding_rotary: Rotary,
|
||||
half_window: usize,
|
||||
}
|
||||
|
||||
impl ModernBert {
|
||||
/// Load from a VarBuilder rooted at the encoder's prefix (`encoder`
|
||||
/// in a Laya checkpoint, `model` in a stock ModernBERT one). Linear
|
||||
/// weights load in the VarBuilder's dtype; embeddings and norms in
|
||||
/// f32.
|
||||
pub fn load(vb: VarBuilder, cfg: &EncoderConfig) -> Result<Self> {
|
||||
let d = cfg.hidden_size;
|
||||
let f32_vb = vb.clone().set_dtype(DType::F32);
|
||||
let embeddings = candle_nn::embedding(
|
||||
cfg.vocab_size,
|
||||
d,
|
||||
f32_vb.pp("embeddings").pp("tok_embeddings"),
|
||||
)
|
||||
.context("encoder.embeddings.tok_embeddings")?;
|
||||
let embed_norm = fused_layer_norm_no_bias(d, cfg.norm_eps, f32_vb.pp("embeddings.norm"))?;
|
||||
let mut layers = Vec::with_capacity(cfg.num_hidden_layers);
|
||||
for (i, kind) in cfg.layer_kinds.iter().copied().enumerate() {
|
||||
let lvb = vb.pp(format!("layers.{i}"));
|
||||
let norm_vb = f32_vb.pp(format!("layers.{i}"));
|
||||
// ModernBERT builds layer 0's attn_norm as an identity and
|
||||
// saves nothing for it; every other layer must have one.
|
||||
let attn_norm = if lvb.contains_tensor("attn_norm.weight") {
|
||||
Some(fused_layer_norm_no_bias(
|
||||
d,
|
||||
cfg.norm_eps,
|
||||
norm_vb.pp("attn_norm"),
|
||||
)?)
|
||||
} else {
|
||||
ensure!(
|
||||
i == 0,
|
||||
"encoder layer {i} has no attn_norm.weight — wrong weight prefix, \
|
||||
or not a ModernBERT checkpoint"
|
||||
);
|
||||
None
|
||||
};
|
||||
let attn = Attention {
|
||||
wqkv: candle_nn::linear_no_bias(d, 3 * d, lvb.pp("attn.Wqkv"))?,
|
||||
wo: candle_nn::linear_no_bias(d, d, lvb.pp("attn.Wo"))?,
|
||||
n_heads: cfg.num_attention_heads,
|
||||
head_dim: cfg.head_dim(),
|
||||
kind,
|
||||
};
|
||||
let mlp_norm = fused_layer_norm_no_bias(d, cfg.norm_eps, norm_vb.pp("mlp_norm"))?;
|
||||
let mlp = Mlp {
|
||||
wi: candle_nn::linear_no_bias(d, 2 * cfg.intermediate_size, lvb.pp("mlp.Wi"))?,
|
||||
wo: candle_nn::linear_no_bias(cfg.intermediate_size, d, lvb.pp("mlp.Wo"))?,
|
||||
};
|
||||
layers.push(Layer {
|
||||
attn_norm,
|
||||
attn,
|
||||
mlp_norm,
|
||||
mlp,
|
||||
});
|
||||
}
|
||||
let final_norm = fused_layer_norm_no_bias(d, cfg.norm_eps, f32_vb.pp("final_norm"))?;
|
||||
let device = vb.device();
|
||||
Ok(Self {
|
||||
embeddings,
|
||||
embed_norm,
|
||||
layers,
|
||||
final_norm,
|
||||
global_rotary: Rotary::new(cfg, cfg.global_rope_theta, vb.dtype(), device)?,
|
||||
sliding_rotary: Rotary::new(cfg, cfg.local_rope_theta, vb.dtype(), device)?,
|
||||
half_window: cfg.local_attention / 2,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn half_window(&self) -> usize {
|
||||
self.half_window
|
||||
}
|
||||
|
||||
/// `ids` is `(b, l)` u32; returns the final hidden states `(b, l, d)`
|
||||
/// in f32.
|
||||
pub fn forward(&self, ids: &Tensor, masks: &Masks) -> Result<Tensor> {
|
||||
let mut xs = ids.apply(&self.embeddings)?.apply(&self.embed_norm)?;
|
||||
for layer in &self.layers {
|
||||
let normed = match &layer.attn_norm {
|
||||
Some(norm) => xs.apply(norm)?,
|
||||
None => xs.clone(),
|
||||
};
|
||||
let (rotary, mask) = match layer.attn.kind {
|
||||
LayerKind::Global => (&self.global_rotary, masks.global.as_ref()),
|
||||
LayerKind::Sliding => (&self.sliding_rotary, masks.sliding.as_ref()),
|
||||
};
|
||||
xs = (&xs + layer.attn.forward(&normed, rotary, mask)?)?;
|
||||
let mlp_out = xs.apply(&layer.mlp_norm)?.apply(&layer.mlp)?;
|
||||
xs = (xs + mlp_out)?;
|
||||
}
|
||||
Ok(xs.apply(&self.final_norm)?)
|
||||
}
|
||||
|
||||
/// The flash path: `ids` is the packed `(total,)` u32 token stream;
|
||||
/// returns the packed final hidden states `(total, d)` in f32.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
pub fn forward_packed(&self, ids: &Tensor, packed: &Packed) -> Result<Tensor> {
|
||||
let mut xs = ids.apply(&self.embeddings)?.apply(&self.embed_norm)?;
|
||||
for layer in &self.layers {
|
||||
let normed = match &layer.attn_norm {
|
||||
Some(norm) => xs.apply(norm)?,
|
||||
None => xs.clone(),
|
||||
};
|
||||
let rotary = match layer.attn.kind {
|
||||
LayerKind::Global => &self.global_rotary,
|
||||
LayerKind::Sliding => &self.sliding_rotary,
|
||||
};
|
||||
xs = (&xs
|
||||
+ layer
|
||||
.attn
|
||||
.forward_packed(&normed, rotary, packed, self.half_window)?)?;
|
||||
let mlp_out = xs.apply(&layer.mlp_norm)?.apply(&layer.mlp)?;
|
||||
xs = (xs + mlp_out)?;
|
||||
}
|
||||
Ok(xs.apply(&self.final_norm)?)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,358 @@
|
||||
//! The two config files a Laya checkpoint ships, and what they mean.
|
||||
//!
|
||||
//! A checkpoint directory holds `rl_agent_config.json` (the decision
|
||||
//! model: head depth, token budgets, calibration temperatures) and
|
||||
//! `encoder/config.json` (the ModernBERT backbone, written by
|
||||
//! transformers 5). There is no top-level `config.json`.
|
||||
//!
|
||||
//! The encoder config is parsed here rather than through
|
||||
//! `candle_transformers::models::modernbert::Config` because
|
||||
//! transformers 5 moved the rotary bases into a nested
|
||||
//! `rope_parameters.{full_attention,sliding_attention}.rope_theta`
|
||||
//! block, which that struct cannot read; the older flat
|
||||
//! `global_rope_theta` / `local_rope_theta` fields are still honoured
|
||||
//! when a checkpoint carries them. mmBERT (the multilingual
|
||||
//! checkpoint's backbone) is the case that matters: both of its bases
|
||||
//! are 160,000, where the flat-field defaults would silently give the
|
||||
//! sliding layers 10,000.
|
||||
|
||||
use anyhow::{Context, Result, bail, ensure};
|
||||
use serde::Deserialize;
|
||||
use std::collections::BTreeMap;
|
||||
use std::path::Path;
|
||||
|
||||
/// `rl_agent_config.json`.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct LayaConfig {
|
||||
/// The backbone the head was trained on, e.g.
|
||||
/// `answerdotai/ModernBERT-large`. Informational: the weights come
|
||||
/// from the checkpoint, the shape from `encoder/config.json`.
|
||||
#[serde(default)]
|
||||
pub encoder: Option<String>,
|
||||
#[serde(default = "default_head_layers")]
|
||||
pub head_layers: usize,
|
||||
#[serde(default = "default_max_len")]
|
||||
pub max_len: usize,
|
||||
#[serde(default = "default_head_max_len")]
|
||||
pub head_max_len: usize,
|
||||
/// One act head output per entry, plus one for "answer". The act
|
||||
/// head's width is `act_costs.len() + 1`.
|
||||
#[serde(default)]
|
||||
pub act_costs: BTreeMap<String, f64>,
|
||||
/// Per-question-type temperatures, indexed choice / score / noul.
|
||||
#[serde(default = "default_temperature")]
|
||||
pub temperature: Vec<f64>,
|
||||
/// Temperatures fitted per `"<type>:<cardinality bucket>"`.
|
||||
#[serde(default)]
|
||||
pub temperature_by_options: BTreeMap<String, f64>,
|
||||
/// Autocast dtype the reference runs on GPU (`"bf16"` or `"fp16"`).
|
||||
#[serde(default)]
|
||||
pub amp_dtype: Option<String>,
|
||||
}
|
||||
|
||||
fn default_head_layers() -> usize {
|
||||
2
|
||||
}
|
||||
fn default_max_len() -> usize {
|
||||
512
|
||||
}
|
||||
fn default_head_max_len() -> usize {
|
||||
192
|
||||
}
|
||||
fn default_temperature() -> Vec<f64> {
|
||||
vec![1.0, 1.0, 1.0]
|
||||
}
|
||||
|
||||
impl LayaConfig {
|
||||
pub fn n_act(&self) -> usize {
|
||||
self.act_costs.len() + 1
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether an attention layer sees the whole sequence or a window.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum LayerKind {
|
||||
Global,
|
||||
Sliding,
|
||||
}
|
||||
|
||||
/// The ModernBERT backbone shape, resolved from `encoder/config.json`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct EncoderConfig {
|
||||
pub vocab_size: usize,
|
||||
pub hidden_size: usize,
|
||||
pub num_hidden_layers: usize,
|
||||
pub num_attention_heads: usize,
|
||||
pub intermediate_size: usize,
|
||||
pub max_position_embeddings: usize,
|
||||
pub norm_eps: f64,
|
||||
pub pad_token_id: u32,
|
||||
/// Full sliding-window width. A token attends to keys at distance
|
||||
/// `<= local_attention / 2` on either side.
|
||||
pub local_attention: usize,
|
||||
pub global_rope_theta: f64,
|
||||
pub local_rope_theta: f64,
|
||||
pub layer_kinds: Vec<LayerKind>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RawEncoderConfig {
|
||||
vocab_size: usize,
|
||||
hidden_size: usize,
|
||||
num_hidden_layers: usize,
|
||||
num_attention_heads: usize,
|
||||
intermediate_size: usize,
|
||||
max_position_embeddings: usize,
|
||||
#[serde(default)]
|
||||
norm_eps: Option<f64>,
|
||||
#[serde(default)]
|
||||
layer_norm_eps: Option<f64>,
|
||||
pad_token_id: u32,
|
||||
#[serde(default = "default_global_every")]
|
||||
global_attn_every_n_layers: usize,
|
||||
local_attention: usize,
|
||||
#[serde(default)]
|
||||
layer_types: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
rope_parameters: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
global_rope_theta: Option<f64>,
|
||||
#[serde(default)]
|
||||
local_rope_theta: Option<f64>,
|
||||
#[serde(default)]
|
||||
norm_bias: bool,
|
||||
#[serde(default)]
|
||||
attention_bias: bool,
|
||||
#[serde(default)]
|
||||
mlp_bias: bool,
|
||||
#[serde(default)]
|
||||
hidden_activation: Option<String>,
|
||||
}
|
||||
|
||||
fn default_global_every() -> usize {
|
||||
3
|
||||
}
|
||||
|
||||
/// transformers' ModernBERT defaults for checkpoints that predate
|
||||
/// `rope_parameters` and also omit the flat fields.
|
||||
const DEFAULT_GLOBAL_ROPE_THETA: f64 = 160_000.0;
|
||||
const DEFAULT_LOCAL_ROPE_THETA: f64 = 10_000.0;
|
||||
|
||||
impl EncoderConfig {
|
||||
pub fn from_json(text: &str) -> Result<Self> {
|
||||
let raw: RawEncoderConfig =
|
||||
serde_json::from_str(text).context("parse ModernBERT encoder config")?;
|
||||
|
||||
// Variants this port does not implement are refused, not
|
||||
// approximated: each would load without error and answer wrong.
|
||||
ensure!(
|
||||
!raw.norm_bias && !raw.attention_bias && !raw.mlp_bias,
|
||||
"ModernBERT with norm/attention/mlp biases is not supported"
|
||||
);
|
||||
if let Some(act) = raw.hidden_activation.as_deref() {
|
||||
ensure!(
|
||||
act == "gelu",
|
||||
"ModernBERT hidden_activation {act:?} is not supported (only exact \"gelu\")"
|
||||
);
|
||||
}
|
||||
ensure!(
|
||||
raw.hidden_size.is_multiple_of(raw.num_attention_heads),
|
||||
"hidden_size {} is not divisible by num_attention_heads {}",
|
||||
raw.hidden_size,
|
||||
raw.num_attention_heads
|
||||
);
|
||||
|
||||
let (global_rope_theta, local_rope_theta) = rope_thetas(&raw)?;
|
||||
|
||||
let layer_kinds = match &raw.layer_types {
|
||||
Some(types) => {
|
||||
ensure!(
|
||||
types.len() == raw.num_hidden_layers,
|
||||
"layer_types has {} entries for {} layers",
|
||||
types.len(),
|
||||
raw.num_hidden_layers
|
||||
);
|
||||
types
|
||||
.iter()
|
||||
.map(|t| match t.as_str() {
|
||||
"full_attention" => Ok(LayerKind::Global),
|
||||
"sliding_attention" => Ok(LayerKind::Sliding),
|
||||
other => bail!("unknown ModernBERT layer type {other:?}"),
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?
|
||||
}
|
||||
None => (0..raw.num_hidden_layers)
|
||||
.map(|i| {
|
||||
if i % raw.global_attn_every_n_layers == 0 {
|
||||
LayerKind::Global
|
||||
} else {
|
||||
LayerKind::Sliding
|
||||
}
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
vocab_size: raw.vocab_size,
|
||||
hidden_size: raw.hidden_size,
|
||||
num_hidden_layers: raw.num_hidden_layers,
|
||||
num_attention_heads: raw.num_attention_heads,
|
||||
intermediate_size: raw.intermediate_size,
|
||||
max_position_embeddings: raw.max_position_embeddings,
|
||||
norm_eps: raw.norm_eps.or(raw.layer_norm_eps).unwrap_or(1e-5),
|
||||
pad_token_id: raw.pad_token_id,
|
||||
local_attention: raw.local_attention,
|
||||
global_rope_theta,
|
||||
local_rope_theta,
|
||||
layer_kinds,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn head_dim(&self) -> usize {
|
||||
self.hidden_size / self.num_attention_heads
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve the two rotary bases. `rope_parameters` (transformers 5)
|
||||
/// wins over the flat fields; within it, a per-layer-type block wins
|
||||
/// over a shared top-level `rope_theta`.
|
||||
fn rope_thetas(raw: &RawEncoderConfig) -> Result<(f64, f64)> {
|
||||
let mut global = raw.global_rope_theta.unwrap_or(DEFAULT_GLOBAL_ROPE_THETA);
|
||||
let mut local = raw.local_rope_theta.unwrap_or(DEFAULT_LOCAL_ROPE_THETA);
|
||||
if let Some(rope) = &raw.rope_parameters {
|
||||
let shared = rope.get("rope_theta").and_then(|v| v.as_f64());
|
||||
for (key, slot) in [
|
||||
("full_attention", &mut global),
|
||||
("sliding_attention", &mut local),
|
||||
] {
|
||||
let block = rope.get(key);
|
||||
if let Some(rope_type) = block
|
||||
.and_then(|b| b.get("rope_type"))
|
||||
.and_then(|v| v.as_str())
|
||||
{
|
||||
ensure!(
|
||||
rope_type == "default",
|
||||
"rope_type {rope_type:?} for {key} is not supported"
|
||||
);
|
||||
}
|
||||
let theta = block
|
||||
.and_then(|b| b.get("rope_theta"))
|
||||
.and_then(|v| v.as_f64())
|
||||
.or(shared);
|
||||
if let Some(theta) = theta {
|
||||
*slot = theta;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok((global, local))
|
||||
}
|
||||
|
||||
/// Both configs of one checkpoint directory.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CheckpointConfig {
|
||||
pub laya: LayaConfig,
|
||||
pub encoder: EncoderConfig,
|
||||
}
|
||||
|
||||
impl CheckpointConfig {
|
||||
/// Read `rl_agent_config.json` and `encoder/config.json` from a
|
||||
/// checkpoint directory (the repo root, or one of its subfolders).
|
||||
pub fn from_dir(dir: &Path) -> Result<Self> {
|
||||
let laya_path = dir.join("rl_agent_config.json");
|
||||
let laya: LayaConfig = serde_json::from_str(
|
||||
&std::fs::read_to_string(&laya_path)
|
||||
.with_context(|| format!("read {}", laya_path.display()))?,
|
||||
)
|
||||
.with_context(|| format!("parse {}", laya_path.display()))?;
|
||||
let enc_path = dir.join("encoder").join("config.json");
|
||||
let encoder = EncoderConfig::from_json(
|
||||
&std::fs::read_to_string(&enc_path)
|
||||
.with_context(|| format!("read {}", enc_path.display()))?,
|
||||
)
|
||||
.with_context(|| format!("parse {}", enc_path.display()))?;
|
||||
ensure!(
|
||||
laya.temperature.len() == 3,
|
||||
"rl_agent_config temperature must have 3 entries (choice, score, noul), got {}",
|
||||
laya.temperature.len()
|
||||
);
|
||||
Ok(Self { laya, encoder })
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn base() -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"vocab_size": 100, "hidden_size": 64, "num_hidden_layers": 4,
|
||||
"num_attention_heads": 4, "intermediate_size": 96,
|
||||
"max_position_embeddings": 512, "norm_eps": 1e-5, "pad_token_id": 0,
|
||||
"global_attn_every_n_layers": 3, "local_attention": 16,
|
||||
})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn nested_rope_parameters_override_flat_defaults() {
|
||||
let mut v = base();
|
||||
v["rope_parameters"] = serde_json::json!({
|
||||
"full_attention": {"rope_theta": 160000.0, "rope_type": "default"},
|
||||
"sliding_attention": {"rope_theta": 160000.0, "rope_type": "default"},
|
||||
});
|
||||
let c = EncoderConfig::from_json(&v.to_string()).unwrap();
|
||||
assert_eq!(c.global_rope_theta, 160_000.0);
|
||||
// The mmBERT case: the flat default would have been 10,000.
|
||||
assert_eq!(c.local_rope_theta, 160_000.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flat_fields_are_read_without_rope_parameters() {
|
||||
let mut v = base();
|
||||
v["global_rope_theta"] = serde_json::json!(80000.0);
|
||||
v["local_rope_theta"] = serde_json::json!(5000.0);
|
||||
let c = EncoderConfig::from_json(&v.to_string()).unwrap();
|
||||
assert_eq!((c.global_rope_theta, c.local_rope_theta), (80000.0, 5000.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn layer_kinds_follow_every_n_when_layer_types_absent() {
|
||||
let c = EncoderConfig::from_json(&base().to_string()).unwrap();
|
||||
assert_eq!(
|
||||
c.layer_kinds,
|
||||
vec![
|
||||
LayerKind::Global,
|
||||
LayerKind::Sliding,
|
||||
LayerKind::Sliding,
|
||||
LayerKind::Global
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_layer_types_win() {
|
||||
let mut v = base();
|
||||
v["layer_types"] = serde_json::json!([
|
||||
"sliding_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"sliding_attention"
|
||||
]);
|
||||
let c = EncoderConfig::from_json(&v.to_string()).unwrap();
|
||||
assert_eq!(c.layer_kinds[0], LayerKind::Sliding);
|
||||
assert_eq!(c.layer_kinds[1], LayerKind::Global);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unsupported_variants_are_refused() {
|
||||
let mut v = base();
|
||||
v["attention_bias"] = serde_json::json!(true);
|
||||
assert!(EncoderConfig::from_json(&v.to_string()).is_err());
|
||||
let mut v = base();
|
||||
v["hidden_activation"] = serde_json::json!("gelu_pytorch_tanh");
|
||||
assert!(EncoderConfig::from_json(&v.to_string()).is_err());
|
||||
let mut v = base();
|
||||
v["rope_parameters"] =
|
||||
serde_json::json!({"full_attention": {"rope_type": "yarn", "rope_theta": 1.0}});
|
||||
assert!(EncoderConfig::from_json(&v.to_string()).is_err());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,180 @@
|
||||
//! The decision head that sits on top of the encoder.
|
||||
//!
|
||||
//! Mirrors the reference `DecisionModel` (the `laya` SDK's
|
||||
//! `common.py`):
|
||||
//!
|
||||
//! 1. a per-question-type embedding added to every position;
|
||||
//! 2. `head_layers` × PyTorch `nn.TransformerEncoderLayer` with
|
||||
//! `norm_first=True`: pre-norm self-attention through a fused
|
||||
//! `in_proj` with bias, then a pre-norm feed-forward whose
|
||||
//! activation is **ReLU** — PyTorch's default, which the reference
|
||||
//! never overrides;
|
||||
//! 3. `scorer`, applied at each option's `[MASK]` marker: LayerNorm →
|
||||
//! Linear → exact GELU → Linear to one logit;
|
||||
//! 4. `act_head`, over the first position's hidden state plus four
|
||||
//! summary features of the (uncalibrated) answer distribution.
|
||||
|
||||
use anyhow::Result;
|
||||
use candle_core::{DType, Tensor};
|
||||
use candle_nn::{Embedding, LayerNorm, Linear, VarBuilder};
|
||||
|
||||
use super::backbone::{project, scaled_attention};
|
||||
|
||||
/// PyTorch's `LayerNorm` default.
|
||||
const TORCH_LN_EPS: f64 = 1e-5;
|
||||
|
||||
struct HeadLayer {
|
||||
norm1: LayerNorm,
|
||||
in_proj: Linear,
|
||||
out_proj: Linear,
|
||||
norm2: LayerNorm,
|
||||
linear1: Linear,
|
||||
linear2: Linear,
|
||||
n_heads: usize,
|
||||
head_dim: usize,
|
||||
}
|
||||
|
||||
impl HeadLayer {
|
||||
fn load(vb: VarBuilder, d: usize) -> Result<Self> {
|
||||
// `nn.TransformerEncoder` is built with `nhead = max(1, d // 64)`
|
||||
// and `dim_feedforward = 4 * d`.
|
||||
let n_heads = (d / 64).max(1);
|
||||
let f32_vb = vb.clone().set_dtype(DType::F32);
|
||||
Ok(Self {
|
||||
norm1: candle_nn::layer_norm(d, TORCH_LN_EPS, f32_vb.pp("norm1"))?,
|
||||
in_proj: linear_torch_mha_in_proj(d, vb.pp("self_attn"))?,
|
||||
out_proj: candle_nn::linear(d, d, vb.pp("self_attn.out_proj"))?,
|
||||
norm2: candle_nn::layer_norm(d, TORCH_LN_EPS, f32_vb.pp("norm2"))?,
|
||||
linear1: candle_nn::linear(d, 4 * d, vb.pp("linear1"))?,
|
||||
linear2: candle_nn::linear(4 * d, d, vb.pp("linear2"))?,
|
||||
n_heads,
|
||||
head_dim: d / n_heads,
|
||||
})
|
||||
}
|
||||
|
||||
/// `xs` is the f32 residual stream; projections run in the weights'
|
||||
/// dtype and join it back in f32, as under autocast.
|
||||
fn forward(&self, xs: &Tensor, key_mask: Option<&Tensor>) -> Result<Tensor> {
|
||||
let (b, l, d) = xs.dims3()?;
|
||||
let qkv = project(&self.in_proj, &xs.apply(&self.norm1)?)?
|
||||
.reshape((b, l, 3, self.n_heads, self.head_dim))?
|
||||
.permute((2, 0, 3, 1, 4))?;
|
||||
let q = qkv.get(0)?.contiguous()?;
|
||||
let k = qkv.get(1)?.contiguous()?;
|
||||
let v = qkv.get(2)?.contiguous()?;
|
||||
let ctx = scaled_attention(&q, &k, &v, key_mask, self.head_dim)?
|
||||
.transpose(1, 2)?
|
||||
.reshape((b, l, d))?;
|
||||
let xs = (xs + ctx.apply(&self.out_proj)?.to_dtype(DType::F32)?)?;
|
||||
let ff = project(&self.linear1, &xs.apply(&self.norm2)?)?
|
||||
.relu()?
|
||||
.apply(&self.linear2)?
|
||||
.to_dtype(DType::F32)?;
|
||||
Ok((xs + ff)?)
|
||||
}
|
||||
|
||||
/// The flash path over a packed `(total, d)` f32 stream.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
fn forward_packed(&self, xs: &Tensor, packed: &super::backbone::Packed) -> Result<Tensor> {
|
||||
use super::backbone::{flash_attention, split_qkv};
|
||||
let (t, d) = xs.dims2()?;
|
||||
let qkv = project(&self.in_proj, &xs.apply(&self.norm1)?)?;
|
||||
let (q, k, v) = split_qkv(&qkv, self.n_heads, self.head_dim)?;
|
||||
let ctx = flash_attention(&q, &k, &v, packed, None)?.reshape((t, d))?;
|
||||
let xs = (xs + ctx.apply(&self.out_proj)?.to_dtype(DType::F32)?)?;
|
||||
let ff = project(&self.linear1, &xs.apply(&self.norm2)?)?
|
||||
.relu()?
|
||||
.apply(&self.linear2)?
|
||||
.to_dtype(DType::F32)?;
|
||||
Ok((xs + ff)?)
|
||||
}
|
||||
}
|
||||
|
||||
/// `nn.MultiheadAttention` stores q/k/v as one `in_proj_weight`
|
||||
/// `(3d, d)` and `in_proj_bias` `(3d)`, rows ordered q then k then v —
|
||||
/// exactly the layout one fused `Linear` wants.
|
||||
fn linear_torch_mha_in_proj(d: usize, vb: VarBuilder) -> Result<Linear> {
|
||||
let w = vb.get((3 * d, d), "in_proj_weight")?;
|
||||
let b = vb.get(3 * d, "in_proj_bias")?;
|
||||
Ok(Linear::new(w, Some(b)))
|
||||
}
|
||||
|
||||
pub struct DecisionHead {
|
||||
type_emb: Embedding,
|
||||
layers: Vec<HeadLayer>,
|
||||
scorer_norm: LayerNorm,
|
||||
scorer_fc: Linear,
|
||||
scorer_out: Linear,
|
||||
act_fc: Linear,
|
||||
act_out: Linear,
|
||||
}
|
||||
|
||||
impl DecisionHead {
|
||||
/// `vb` is the checkpoint root; the head's tensors are top-level.
|
||||
/// Linear weights load in `vb`'s dtype; the type embedding and norms
|
||||
/// in f32.
|
||||
pub fn load(vb: VarBuilder, d: usize, head_layers: usize, n_act: usize) -> Result<Self> {
|
||||
let layers = (0..head_layers)
|
||||
.map(|i| HeadLayer::load(vb.pp(format!("head.layers.{i}")), d))
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
let f32_vb = vb.clone().set_dtype(DType::F32);
|
||||
Ok(Self {
|
||||
type_emb: candle_nn::embedding(3, d, f32_vb.pp("type_emb"))?,
|
||||
layers,
|
||||
scorer_norm: candle_nn::layer_norm(d, TORCH_LN_EPS, f32_vb.pp("scorer.0"))?,
|
||||
scorer_fc: candle_nn::linear(d, d, vb.pp("scorer.1"))?,
|
||||
scorer_out: candle_nn::linear(d, 1, vb.pp("scorer.3"))?,
|
||||
act_fc: candle_nn::linear(d + 4, 256, vb.pp("act_head.0"))?,
|
||||
act_out: candle_nn::linear(256, n_act, vb.pp("act_head.2"))?,
|
||||
})
|
||||
}
|
||||
|
||||
/// Type embedding plus the transformer layers. `h` is the encoder
|
||||
/// output `(b, l, d)`, `qtype` is `(b,)` u32.
|
||||
pub fn contextualise(
|
||||
&self,
|
||||
h: &Tensor,
|
||||
qtype: &Tensor,
|
||||
key_mask: Option<&Tensor>,
|
||||
) -> Result<Tensor> {
|
||||
let mut h = h.broadcast_add(&qtype.apply(&self.type_emb)?.unsqueeze(1)?)?;
|
||||
for layer in &self.layers {
|
||||
h = layer.forward(&h, key_mask)?;
|
||||
}
|
||||
Ok(h)
|
||||
}
|
||||
|
||||
/// The flash path: `h` is the packed encoder output `(total, d)` and
|
||||
/// `qtype` the per-token question type `(total,)` u32.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
pub fn contextualise_packed(
|
||||
&self,
|
||||
h: &Tensor,
|
||||
qtype: &Tensor,
|
||||
packed: &super::backbone::Packed,
|
||||
) -> Result<Tensor> {
|
||||
let mut h = (h + qtype.apply(&self.type_emb)?)?;
|
||||
for layer in &self.layers {
|
||||
h = layer.forward_packed(&h, packed)?;
|
||||
}
|
||||
Ok(h)
|
||||
}
|
||||
|
||||
/// Score gathered marker states `(b, k, d)` → `(b, k)` f32 logits.
|
||||
pub fn score(&self, markers: &Tensor) -> Result<Tensor> {
|
||||
let x = project(&self.scorer_fc, &markers.apply(&self.scorer_norm)?)?
|
||||
.gelu_erf()?
|
||||
.apply(&self.scorer_out)?;
|
||||
Ok(x.squeeze(2)?.to_dtype(DType::F32)?)
|
||||
}
|
||||
|
||||
/// Act logits `(b, n_act)` f32 from the pooled first position
|
||||
/// `(b, d)` and the four answer-distribution features `(b, 4)`.
|
||||
pub fn act(&self, pooled: &Tensor, feats: &Tensor) -> Result<Tensor> {
|
||||
let x = Tensor::cat(&[pooled, &feats.to_dtype(pooled.dtype())?], 1)?;
|
||||
let x = project(&self.act_fc, &x)?
|
||||
.gelu_erf()?
|
||||
.apply(&self.act_out)?;
|
||||
Ok(x.to_dtype(DType::F32)?)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
//! Laya — a non-autoregressive decision model (`convaiinnovations/laya`).
|
||||
//!
|
||||
//! Not a language model: a ModernBERT encoder plus a small decision
|
||||
//! head that scores the options of a typed question in one forward
|
||||
//! pass. The caller assembles each question into one token sequence
|
||||
//! (`[CLS] <type> question: <instructions> [SEP] [MASK] opt0 [MASK]
|
||||
//! opt1 … [SEP] state [SEP]`) and says where each option's `[MASK]`
|
||||
//! marker sits; this module returns one raw score per marker plus the
|
||||
//! act head's distribution. Calibration (temperature scaling) and
|
||||
//! answer shaping happen after, CPU-side, from the checkpoint's
|
||||
//! `rl_agent_config.json`.
|
||||
//!
|
||||
//! The reference implementation is the `laya` SDK's `DecisionModel`
|
||||
//! (`laya/common.py`), not the `rl_common.py` shipped in the model
|
||||
//! repository, which has diverged from it.
|
||||
//!
|
||||
//! - [`config`] — the two config files a checkpoint ships.
|
||||
//! - [`backbone`] — ModernBERT, vendored so it runs outside f32.
|
||||
//! - [`head`] — type embedding, the head's transformer layers, the
|
||||
//! marker scorer and the act head.
|
||||
//! - [`model`] — the composition, batch in, [`DecisionRow`]s out.
|
||||
|
||||
pub mod backbone;
|
||||
pub mod config;
|
||||
pub mod head;
|
||||
pub mod model;
|
||||
|
||||
pub use config::CheckpointConfig;
|
||||
pub use model::LayaModel;
|
||||
|
||||
/// Question type, as the checkpoint's type embedding indexes it.
|
||||
pub const QTYPE_CHOICE: u8 = 0;
|
||||
pub const QTYPE_SCORE: u8 = 1;
|
||||
pub const QTYPE_NOUL: u8 = 2;
|
||||
|
||||
/// One assembled question: its token ids, where each option's
|
||||
/// `[MASK]` marker sits, and its type.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct DecisionSeq {
|
||||
pub input_ids: Vec<u32>,
|
||||
pub markers: Vec<usize>,
|
||||
/// [`QTYPE_CHOICE`], [`QTYPE_SCORE`] or [`QTYPE_NOUL`].
|
||||
pub qtype: u8,
|
||||
}
|
||||
|
||||
/// Sequences scored in one forward. One request's questions normally
|
||||
/// travel together, so the state they share is encoded in one pass.
|
||||
#[derive(Debug, Clone, Default, PartialEq)]
|
||||
pub struct DecisionBatch {
|
||||
pub seqs: Vec<DecisionSeq>,
|
||||
}
|
||||
|
||||
/// One sequence's result, CPU-side.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct DecisionRow {
|
||||
/// One uncalibrated score per marker, in marker order.
|
||||
pub logits: Vec<f32>,
|
||||
/// Softmax over the act head's outputs. Index 0 is what the
|
||||
/// reference reports as `action.act_probability`.
|
||||
pub act: Vec<f32>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
@@ -0,0 +1,260 @@
|
||||
//! A whole Laya checkpoint: encoder + decision head, batch in, raw
|
||||
//! scores out.
|
||||
|
||||
use anyhow::{Context, Result, ensure};
|
||||
use candle_core::{DType, Device, Tensor};
|
||||
use candle_nn::VarBuilder;
|
||||
use std::path::Path;
|
||||
|
||||
use super::backbone::{Masks, ModernBert};
|
||||
use super::config::CheckpointConfig;
|
||||
use super::head::DecisionHead;
|
||||
use super::{DecisionBatch, DecisionRow};
|
||||
|
||||
pub struct LayaModel {
|
||||
encoder: ModernBert,
|
||||
head: DecisionHead,
|
||||
config: CheckpointConfig,
|
||||
device: Device,
|
||||
/// Precision of the linear layers (the load dtype).
|
||||
dtype: DType,
|
||||
}
|
||||
|
||||
/// A validated batch, flattened: what both forward paths start from.
|
||||
struct Flat {
|
||||
lens: Vec<usize>,
|
||||
qtypes: Vec<u32>,
|
||||
/// Each row's marker positions, padded to `k_max` with 0.
|
||||
marker_idx: Vec<u32>,
|
||||
k_max: usize,
|
||||
seq_len: usize,
|
||||
}
|
||||
|
||||
impl LayaModel {
|
||||
/// Load one checkpoint directory (the repo root, or one of its
|
||||
/// subfolders): `rl_agent_config.json`, `encoder/config.json` and
|
||||
/// `model.safetensors`.
|
||||
pub fn load(dir: &Path, device: &Device, dtype: DType) -> Result<Self> {
|
||||
let config = CheckpointConfig::from_dir(dir)?;
|
||||
let weights = dir.join("model.safetensors");
|
||||
// SAFETY: the file is memory-mapped read-only for the lifetime
|
||||
// of the VarBuilder; nothing writes to a checkpoint in place.
|
||||
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[&weights], dtype, device) }
|
||||
.with_context(|| format!("mmap {}", weights.display()))?;
|
||||
Self::from_var_builder(vb, config, device)
|
||||
}
|
||||
|
||||
pub fn from_var_builder(
|
||||
vb: VarBuilder,
|
||||
config: CheckpointConfig,
|
||||
device: &Device,
|
||||
) -> Result<Self> {
|
||||
let dtype = vb.dtype();
|
||||
let encoder = ModernBert::load(vb.pp("encoder"), &config.encoder)?;
|
||||
let head = DecisionHead::load(
|
||||
vb.clone(),
|
||||
config.encoder.hidden_size,
|
||||
config.laya.head_layers,
|
||||
config.laya.n_act(),
|
||||
)?;
|
||||
Ok(Self {
|
||||
encoder,
|
||||
head,
|
||||
config,
|
||||
device: device.clone(),
|
||||
dtype,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn config(&self) -> &CheckpointConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
/// Precision the linear layers were loaded in.
|
||||
pub fn dtype(&self) -> DType {
|
||||
self.dtype
|
||||
}
|
||||
|
||||
/// One forward over every sequence in `batch`.
|
||||
///
|
||||
/// On a CUDA device in reduced precision, in a build with the
|
||||
/// `flash-attn` feature, the batch is packed without padding and
|
||||
/// attention runs in flash varlen kernels. Otherwise rows are
|
||||
/// right-padded to the longest, as the reference collates them, and
|
||||
/// attention is eager; padding changes nothing about an unpadded
|
||||
/// row's result.
|
||||
pub fn forward(&self, batch: &DecisionBatch) -> Result<Vec<DecisionRow>> {
|
||||
let flat = self.validate(batch)?;
|
||||
#[cfg(feature = "flash-attn")]
|
||||
if self.device.is_cuda() && self.dtype != DType::F32 && flash_enabled() {
|
||||
return self.forward_packed(batch, &flat);
|
||||
}
|
||||
self.forward_padded(batch, &flat)
|
||||
}
|
||||
|
||||
fn validate(&self, batch: &DecisionBatch) -> Result<Flat> {
|
||||
let seqs = &batch.seqs;
|
||||
ensure!(!seqs.is_empty(), "empty decision batch");
|
||||
let max_pos = self.config.encoder.max_position_embeddings;
|
||||
let seq_len = seqs.iter().map(|s| s.input_ids.len()).max().unwrap_or(0);
|
||||
ensure!(seq_len > 0, "decision batch has an empty sequence");
|
||||
ensure!(
|
||||
seq_len <= max_pos,
|
||||
"sequence of {seq_len} tokens exceeds max_position_embeddings {max_pos}"
|
||||
);
|
||||
let vocab = self.config.encoder.vocab_size as u32;
|
||||
let k_max = seqs.iter().map(|s| s.markers.len()).max().unwrap_or(0);
|
||||
ensure!(
|
||||
k_max > 0,
|
||||
"decision batch has a sequence with no option markers"
|
||||
);
|
||||
let mut lens = Vec::with_capacity(seqs.len());
|
||||
let mut qtypes = Vec::with_capacity(seqs.len());
|
||||
let mut marker_idx = Vec::with_capacity(seqs.len() * k_max);
|
||||
for s in seqs {
|
||||
let n = s.input_ids.len();
|
||||
ensure!(n > 0, "decision batch has an empty sequence");
|
||||
ensure!(!s.markers.is_empty(), "sequence has no option markers");
|
||||
ensure!(s.qtype < 3, "question type {} out of range", s.qtype);
|
||||
if let Some(&bad) = s.input_ids.iter().find(|&&t| t >= vocab) {
|
||||
anyhow::bail!("token id {bad} outside vocabulary of {vocab}");
|
||||
}
|
||||
if let Some(&bad) = s.markers.iter().find(|&&m| m >= n) {
|
||||
anyhow::bail!("marker position {bad} outside a {n}-token sequence");
|
||||
}
|
||||
lens.push(n);
|
||||
qtypes.push(s.qtype as u32);
|
||||
marker_idx.extend(s.markers.iter().map(|&m| m as u32));
|
||||
// Absent slots gather position 0; their scores are discarded.
|
||||
marker_idx.extend(std::iter::repeat_n(0u32, k_max - s.markers.len()));
|
||||
}
|
||||
Ok(Flat {
|
||||
lens,
|
||||
qtypes,
|
||||
marker_idx,
|
||||
k_max,
|
||||
seq_len,
|
||||
})
|
||||
}
|
||||
|
||||
/// The eager path: rows right-padded to `(b, seq_len)`.
|
||||
fn forward_padded(&self, batch: &DecisionBatch, flat: &Flat) -> Result<Vec<DecisionRow>> {
|
||||
let b = batch.seqs.len();
|
||||
let seq_len = flat.seq_len;
|
||||
let pad = self.config.encoder.pad_token_id;
|
||||
let mut ids = Vec::with_capacity(b * seq_len);
|
||||
for s in &batch.seqs {
|
||||
ids.extend_from_slice(&s.input_ids);
|
||||
ids.extend(std::iter::repeat_n(pad, seq_len - s.input_ids.len()));
|
||||
}
|
||||
let ids = Tensor::from_vec(ids, (b, seq_len), &self.device)?;
|
||||
let masks = Masks::new(
|
||||
&flat.lens,
|
||||
seq_len,
|
||||
self.encoder.half_window(),
|
||||
&self.device,
|
||||
)?;
|
||||
let qtype = Tensor::from_vec(flat.qtypes.clone(), b, &self.device)?;
|
||||
|
||||
let h = self.encoder.forward(&ids, &masks)?;
|
||||
let h = self.head.contextualise(&h, &qtype, masks.global.as_ref())?;
|
||||
let d = self.config.encoder.hidden_size;
|
||||
let idx = Tensor::from_vec(flat.marker_idx.clone(), (b, flat.k_max, 1), &self.device)?
|
||||
.broadcast_as((b, flat.k_max, d))?
|
||||
.contiguous()?;
|
||||
let marker_h = h.gather(&idx, 1)?;
|
||||
let pooled = h.narrow(1, 0, 1)?.squeeze(1)?;
|
||||
self.finish(batch, &marker_h, &pooled)
|
||||
}
|
||||
|
||||
/// The flash path: tokens packed without padding, `(total, d)`.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
fn forward_packed(&self, batch: &DecisionBatch, flat: &Flat) -> Result<Vec<DecisionRow>> {
|
||||
use super::backbone::Packed;
|
||||
let b = batch.seqs.len();
|
||||
let packed = Packed::new(&flat.lens, &self.device)?;
|
||||
let total: usize = flat.lens.iter().sum();
|
||||
let mut ids = Vec::with_capacity(total);
|
||||
let mut qtype_tok = Vec::with_capacity(total);
|
||||
for (s, &q) in batch.seqs.iter().zip(&flat.qtypes) {
|
||||
ids.extend_from_slice(&s.input_ids);
|
||||
qtype_tok.extend(std::iter::repeat_n(q, s.input_ids.len()));
|
||||
}
|
||||
let ids = Tensor::from_vec(ids, total, &self.device)?;
|
||||
let qtype_tok = Tensor::from_vec(qtype_tok, total, &self.device)?;
|
||||
|
||||
let h = self.encoder.forward_packed(&ids, &packed)?;
|
||||
let h = self.head.contextualise_packed(&h, &qtype_tok, &packed)?;
|
||||
let d = self.config.encoder.hidden_size;
|
||||
// Marker positions, and each row's first token, in packed order.
|
||||
let global: Vec<u32> = flat
|
||||
.marker_idx
|
||||
.chunks(flat.k_max)
|
||||
.zip(&packed.offsets)
|
||||
.flat_map(|(ms, &off)| ms.iter().map(move |&m| off as u32 + m))
|
||||
.collect();
|
||||
let global = Tensor::from_vec(global, b * flat.k_max, &self.device)?;
|
||||
let marker_h = h.index_select(&global, 0)?.reshape((b, flat.k_max, d))?;
|
||||
let firsts: Vec<u32> = packed.offsets.iter().map(|&o| o as u32).collect();
|
||||
let pooled = h.index_select(&Tensor::from_vec(firsts, b, &self.device)?, 0)?;
|
||||
self.finish(batch, &marker_h, &pooled)
|
||||
}
|
||||
|
||||
/// Score the gathered marker states `(b, k_max, d)` and run the act
|
||||
/// head over the pooled first positions `(b, d)`.
|
||||
fn finish(
|
||||
&self,
|
||||
batch: &DecisionBatch,
|
||||
marker_h: &Tensor,
|
||||
pooled: &Tensor,
|
||||
) -> Result<Vec<DecisionRow>> {
|
||||
let b = batch.seqs.len();
|
||||
let logits: Vec<Vec<f32>> = self.head.score(marker_h)?.to_vec2()?;
|
||||
let mut feats = Vec::with_capacity(b * 4);
|
||||
let mut row_logits = Vec::with_capacity(b);
|
||||
for (s, full) in batch.seqs.iter().zip(logits) {
|
||||
let row: Vec<f32> = full[..s.markers.len()].to_vec();
|
||||
feats.extend_from_slice(&act_features(&row));
|
||||
row_logits.push(row);
|
||||
}
|
||||
let feats = Tensor::from_vec(feats, (b, 4), &self.device)?;
|
||||
let act_logits = self.head.act(pooled, &feats)?;
|
||||
let act = candle_nn::ops::softmax_last_dim(&act_logits)?.to_vec2::<f32>()?;
|
||||
Ok(row_logits
|
||||
.into_iter()
|
||||
.zip(act)
|
||||
.map(|(logits, act)| DecisionRow { logits, act })
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
/// Runtime kill switch for the flash path: `NEURON_FLASH_ATTN=0` (or
|
||||
/// `false`) forces eager attention without a rebuild — the same switch
|
||||
/// the text architectures read. Read once.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
fn flash_enabled() -> bool {
|
||||
static ENABLED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
|
||||
*ENABLED
|
||||
.get_or_init(|| !std::env::var("NEURON_FLASH_ATTN").is_ok_and(|v| v == "0" || v == "false"))
|
||||
}
|
||||
|
||||
/// The act head's four inputs, from one row's uncalibrated logits:
|
||||
/// top-1 probability, top-1 minus top-2, entropy normalised by
|
||||
/// `ln(max(k, 2))`, and `k / 255`.
|
||||
///
|
||||
/// The reference computes these over the batch's padded width, with
|
||||
/// absent option slots filled with a logit of -1e4. Those slots carry
|
||||
/// exactly zero probability in f32, so computing per row over only the
|
||||
/// real options gives the same result.
|
||||
pub fn act_features(logits: &[f32]) -> [f32; 4] {
|
||||
let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
|
||||
let exps: Vec<f32> = logits.iter().map(|&z| (z - max).exp()).collect();
|
||||
let sum: f32 = exps.iter().sum();
|
||||
let mut p: Vec<f32> = exps.iter().map(|e| e / sum).collect();
|
||||
let k = (logits.len().max(2)) as f32;
|
||||
let ent = -p.iter().map(|&x| x * x.max(1e-9).ln()).sum::<f32>() / k.ln();
|
||||
p.sort_by(|a, b| b.total_cmp(a));
|
||||
let top1 = p[0];
|
||||
let top2 = p.get(1).copied().unwrap_or(0.0);
|
||||
[top1, top1 - top2, ent, k / 255.0]
|
||||
}
|
||||
@@ -0,0 +1,450 @@
|
||||
//! Parity against the upstream reference.
|
||||
//!
|
||||
//! `testdata/laya/tiny/` is a toy-sized checkpoint built and scored by
|
||||
//! the SDK's own `DecisionModel` (`script/generate_laya_fixture.py`);
|
||||
//! it runs in CI. `testdata/laya/reference.json` holds the released
|
||||
//! checkpoints' outputs over a real corpus (`script/laya-reference.py`);
|
||||
//! scoring it needs the weights, so that test is ignored unless
|
||||
//! `LAYA_REFERENCE_DIR` names a snapshot of `convaiinnovations/laya`.
|
||||
|
||||
use super::*;
|
||||
use candle_core::{DType, Device};
|
||||
use serde::Deserialize;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
fn testdata() -> PathBuf {
|
||||
Path::new(env!("CARGO_MANIFEST_DIR")).join("src/harness/testdata/laya")
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct TinyExpected {
|
||||
seqs: Vec<TinySeq>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct TinySeq {
|
||||
input_ids: Vec<u32>,
|
||||
markers: Vec<usize>,
|
||||
qtype: u8,
|
||||
logits: Vec<f32>,
|
||||
act: Vec<f32>,
|
||||
}
|
||||
|
||||
fn load_tiny() -> (LayaModel, TinyExpected) {
|
||||
let dir = testdata().join("tiny");
|
||||
let model = LayaModel::load(&dir, &Device::Cpu, DType::F32).expect("load tiny checkpoint");
|
||||
let expected: TinyExpected = serde_json::from_str(
|
||||
&std::fs::read_to_string(dir.join("expected.json")).expect("read expected.json"),
|
||||
)
|
||||
.expect("parse expected.json");
|
||||
(model, expected)
|
||||
}
|
||||
|
||||
fn batch_of(seqs: &[TinySeq]) -> DecisionBatch {
|
||||
DecisionBatch {
|
||||
seqs: seqs
|
||||
.iter()
|
||||
.map(|s| DecisionSeq {
|
||||
input_ids: s.input_ids.clone(),
|
||||
markers: s.markers.clone(),
|
||||
qtype: s.qtype,
|
||||
})
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Device and precision for the weight-backed tests: `LAYA_DEVICE`
|
||||
/// (`cpu`, or `cuda:N` in a build with the `cuda` feature) and
|
||||
/// `LAYA_DTYPE` (`f32`, `bf16` or `f16`). Defaults: CPU, f32.
|
||||
fn test_device() -> (Device, DType) {
|
||||
let device = match std::env::var("LAYA_DEVICE").as_deref() {
|
||||
Ok("cpu") | Err(_) => Device::Cpu,
|
||||
Ok(s) => {
|
||||
let ordinal = s
|
||||
.strip_prefix("cuda:")
|
||||
.and_then(|n| n.parse().ok())
|
||||
.unwrap_or_else(|| panic!("LAYA_DEVICE={s}: expected cpu or cuda:N"));
|
||||
Device::new_cuda(ordinal).unwrap_or_else(|e| panic!("LAYA_DEVICE={s}: {e}"))
|
||||
}
|
||||
};
|
||||
let dtype = match std::env::var("LAYA_DTYPE").as_deref() {
|
||||
Ok("f32") | Err(_) => DType::F32,
|
||||
Ok("bf16") => DType::BF16,
|
||||
Ok("f16") => DType::F16,
|
||||
Ok(s) => panic!("LAYA_DTYPE={s}: expected f32, bf16 or f16"),
|
||||
};
|
||||
(device, dtype)
|
||||
}
|
||||
|
||||
/// Logit tolerance against the f32 reference for a given precision.
|
||||
/// Reduced precision is judged mainly on argmax agreement; the bound
|
||||
/// only catches a broken kernel, not rounding.
|
||||
fn logit_tolerance(dtype: DType) -> f32 {
|
||||
if dtype == DType::F32 { 1e-3 } else { 1.0 }
|
||||
}
|
||||
|
||||
fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 {
|
||||
assert_eq!(a.len(), b.len(), "length mismatch");
|
||||
a.iter()
|
||||
.zip(b)
|
||||
.map(|(x, y)| (x - y).abs())
|
||||
.fold(0.0, f32::max)
|
||||
}
|
||||
|
||||
/// The whole forward — encoder with both attention kinds and rotary
|
||||
/// bases, padding, type embedding, head layers, marker gather, scorer
|
||||
/// and act head — against the reference, in f32.
|
||||
#[test]
|
||||
fn tiny_checkpoint_matches_reference() {
|
||||
let (model, expected) = load_tiny();
|
||||
let rows = model.forward(&batch_of(&expected.seqs)).expect("forward");
|
||||
assert_eq!(rows.len(), expected.seqs.len());
|
||||
for (i, (got, want)) in rows.iter().zip(&expected.seqs).enumerate() {
|
||||
let dl = max_abs_diff(&got.logits, &want.logits);
|
||||
let da = max_abs_diff(&got.act, &want.act);
|
||||
assert!(
|
||||
dl < 1e-5 && da < 1e-5,
|
||||
"row {i}: logit diff {dl:e}, act diff {da:e}\n got {:?} / {:?}\nwant {:?} / {:?}",
|
||||
got.logits,
|
||||
got.act,
|
||||
want.logits,
|
||||
want.act
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// A row's result must not depend on what it was batched with: padded
|
||||
/// keys are masked out of every attention, including the head's.
|
||||
#[test]
|
||||
fn padding_does_not_change_a_row() {
|
||||
let (model, expected) = load_tiny();
|
||||
let batched = model.forward(&batch_of(&expected.seqs)).expect("batched");
|
||||
for (i, seq) in expected.seqs.iter().enumerate() {
|
||||
let alone = model
|
||||
.forward(&batch_of(std::slice::from_ref(seq)))
|
||||
.expect("single");
|
||||
let dl = max_abs_diff(&alone[0].logits, &batched[i].logits);
|
||||
let da = max_abs_diff(&alone[0].act, &batched[i].act);
|
||||
assert!(dl < 1e-5 && da < 1e-5, "row {i}: {dl:e} / {da:e}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_batches_are_refused() {
|
||||
let (model, _) = load_tiny();
|
||||
let seq = |ids: Vec<u32>, markers: Vec<usize>, qtype| DecisionBatch {
|
||||
seqs: vec![DecisionSeq {
|
||||
input_ids: ids,
|
||||
markers,
|
||||
qtype,
|
||||
}],
|
||||
};
|
||||
assert!(model.forward(&DecisionBatch::default()).is_err());
|
||||
assert!(model.forward(&seq(vec![1, 2, 3], vec![], 0)).is_err());
|
||||
assert!(model.forward(&seq(vec![1, 2, 3], vec![3], 0)).is_err());
|
||||
assert!(model.forward(&seq(vec![1, 2, 3], vec![1], 3)).is_err());
|
||||
assert!(model.forward(&seq(vec![1, 500, 3], vec![1], 0)).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn act_features_match_the_reference_definition() {
|
||||
// Two equal logits: top1 = top2 = 0.5, entropy exactly 1.
|
||||
let f = model::act_features(&[0.3, 0.3]);
|
||||
assert!((f[0] - 0.5).abs() < 1e-6 && f[1].abs() < 1e-6 && (f[2] - 1.0).abs() < 1e-6);
|
||||
assert!((f[3] - 2.0 / 255.0).abs() < 1e-7);
|
||||
// One option: probability 1, no second, k clamps to 2.
|
||||
let f = model::act_features(&[4.0]);
|
||||
assert_eq!(f[0], 1.0);
|
||||
assert_eq!(f[1], 1.0);
|
||||
assert!(f[2].abs() < 1e-6);
|
||||
assert!((f[3] - 2.0 / 255.0).abs() < 1e-7);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- released checkpoints
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Reference {
|
||||
cases: Vec<RefCase>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RefCase {
|
||||
name: String,
|
||||
items: Vec<RefItem>,
|
||||
response: serde_json::Value,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RefItem {
|
||||
input_ids: Vec<u32>,
|
||||
markers: Vec<usize>,
|
||||
qtype: u8,
|
||||
logits: Vec<f32>,
|
||||
act: Vec<f32>,
|
||||
}
|
||||
|
||||
/// Every case in `reference.json`, scored by the checkpoint the
|
||||
/// reference routed it to, on the device and precision
|
||||
/// [`test_device`] selects (f32 on the CPU by default).
|
||||
///
|
||||
/// `LAYA_REFERENCE_DIR` is a snapshot of `convaiinnovations/laya` at the
|
||||
/// revision recorded in the fixture header: the root holds the English
|
||||
/// checkpoint, `multilingual/` and `typed-decisions/` the other two.
|
||||
#[test]
|
||||
#[ignore = "needs the released weights: set LAYA_REFERENCE_DIR"]
|
||||
fn released_checkpoints_match_reference() {
|
||||
let Ok(root) = std::env::var("LAYA_REFERENCE_DIR") else {
|
||||
panic!("LAYA_REFERENCE_DIR is not set");
|
||||
};
|
||||
let root = PathBuf::from(root);
|
||||
let (device, dtype) = test_device();
|
||||
let tolerance = logit_tolerance(dtype);
|
||||
let reference: Reference = serde_json::from_str(
|
||||
&std::fs::read_to_string(testdata().join("reference.json")).expect("read reference.json"),
|
||||
)
|
||||
.expect("parse reference.json");
|
||||
|
||||
// The reference's own GPU run (bf16 autocast), for judging reduced
|
||||
// precision against the drift upstream itself serves with.
|
||||
let torch_bf16: Option<Reference> =
|
||||
std::fs::read_to_string(testdata().join("reference_cuda_bf16.json"))
|
||||
.ok()
|
||||
.map(|t| serde_json::from_str(&t).expect("parse reference_cuda_bf16.json"));
|
||||
// Per checkpoint: summed centred drift of ours and of the reference's
|
||||
// bf16 run, their maxima, and the row count.
|
||||
let mut drift: std::collections::BTreeMap<String, (f32, f32, f32, f32, usize)> =
|
||||
Default::default();
|
||||
|
||||
let mut models = std::collections::BTreeMap::new();
|
||||
let mut stats: std::collections::BTreeMap<String, (usize, f32, f32, usize)> =
|
||||
Default::default();
|
||||
// Forward wall time and tokens encoded, per checkpoint: a first
|
||||
// CPU number, not a benchmark (debug builds are far slower).
|
||||
let mut timing: std::collections::BTreeMap<String, (f64, usize)> = Default::default();
|
||||
for case in &reference.cases {
|
||||
let ckpt = case.response["routing"]["model"]
|
||||
.as_str()
|
||||
.expect("routing.model")
|
||||
.to_string();
|
||||
let model = models.entry(ckpt.clone()).or_insert_with(|| {
|
||||
let dir = match ckpt.as_str() {
|
||||
"english" => root.clone(),
|
||||
other => root.join(other),
|
||||
};
|
||||
LayaModel::load(&dir, &device, dtype).unwrap_or_else(|e| panic!("load {ckpt}: {e:#}"))
|
||||
});
|
||||
let batch = DecisionBatch {
|
||||
seqs: case
|
||||
.items
|
||||
.iter()
|
||||
.map(|it| DecisionSeq {
|
||||
input_ids: it.input_ids.clone(),
|
||||
markers: it.markers.clone(),
|
||||
qtype: it.qtype,
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
let started = std::time::Instant::now();
|
||||
let rows = model
|
||||
.forward(&batch)
|
||||
.unwrap_or_else(|e| panic!("{}: {e:#}", case.name));
|
||||
let t = timing.entry(ckpt.clone()).or_default();
|
||||
t.0 += started.elapsed().as_secs_f64() * 1e3;
|
||||
t.1 += batch.seqs.iter().map(|s| s.input_ids.len()).sum::<usize>();
|
||||
let torch_case = torch_bf16
|
||||
.as_ref()
|
||||
.and_then(|r| r.cases.iter().find(|c| c.name == case.name));
|
||||
for (i, (got, want)) in rows.iter().zip(&case.items).enumerate() {
|
||||
let d = drift.entry(ckpt.clone()).or_default();
|
||||
let ours = centred_diff(&got.logits, &want.logits);
|
||||
d.0 += ours;
|
||||
d.2 = d.2.max(ours);
|
||||
if let Some(t) = torch_case.and_then(|c| c.items.get(i)) {
|
||||
let theirs = centred_diff(&t.logits, &want.logits);
|
||||
d.1 += theirs;
|
||||
d.3 = d.3.max(theirs);
|
||||
}
|
||||
d.4 += 1;
|
||||
}
|
||||
let entry = stats.entry(ckpt.clone()).or_default();
|
||||
for (got, want) in rows.iter().zip(&case.items) {
|
||||
let dl = max_abs_diff(&got.logits, &want.logits);
|
||||
let da = max_abs_diff(&got.act, &want.act);
|
||||
let argmax = |v: &[f32]| {
|
||||
v.iter()
|
||||
.enumerate()
|
||||
.max_by(|a, b| a.1.total_cmp(b.1))
|
||||
.map(|(i, _)| i)
|
||||
};
|
||||
entry.0 += 1;
|
||||
entry.1 = entry.1.max(dl);
|
||||
entry.2 = entry.2.max(da);
|
||||
if argmax(&got.logits) == argmax(&want.logits) {
|
||||
entry.3 += 1;
|
||||
}
|
||||
if dtype != DType::F32 && std::env::var_os("LAYA_VERBOSE").is_some() {
|
||||
eprintln!(" {:<40} {ckpt:<16} |Δlogit| {dl:.3}", case.name);
|
||||
}
|
||||
assert!(
|
||||
dl < tolerance,
|
||||
"{} ({ckpt}): logit diff {dl:e}\n got {:?}\nwant {:?}",
|
||||
case.name,
|
||||
got.logits,
|
||||
want.logits
|
||||
);
|
||||
assert!(da < tolerance, "{} ({ckpt}): act diff {da:e}", case.name);
|
||||
}
|
||||
}
|
||||
for (ckpt, (n, dl, da, agree)) in &stats {
|
||||
let (ms, tokens) = timing[ckpt];
|
||||
eprintln!(
|
||||
"[{device:?} {dtype:?}] {ckpt}: {n} questions, max |Δlogit| {dl:.2e}, max |Δact| {da:.2e}, argmax {agree}/{n}; \
|
||||
{tokens} tokens in {ms:.0} ms ({:.2} ms/token)",
|
||||
ms / tokens as f64
|
||||
);
|
||||
assert_eq!(agree, n, "{ckpt}: argmax disagreement");
|
||||
}
|
||||
// Reduced precision: centred drift (softmax ignores a constant shift)
|
||||
// against f32, next to the reference's own bf16 drift. Per-case drift
|
||||
// is chaotic, so the gate is the mean: no worse than twice upstream's,
|
||||
// plus 0.01 — about one bf16 ulp at these logits' magnitude — so a
|
||||
// checkpoint with a single fixture case whose reference drift happens
|
||||
// to round to nothing is not held to zero.
|
||||
for (ckpt, (ours, theirs, ours_max, theirs_max, n)) in &drift {
|
||||
let (ours, theirs) = (ours / *n as f32, theirs / *n as f32);
|
||||
eprintln!(
|
||||
"[{device:?} {dtype:?}] {ckpt}: centred drift mean {ours:.4} (reference bf16 {theirs:.4}), \
|
||||
max {ours_max:.3} (reference bf16 {theirs_max:.3})"
|
||||
);
|
||||
if dtype != DType::F32 && torch_bf16.is_some() {
|
||||
assert!(
|
||||
ours <= 2.0 * theirs + 0.01,
|
||||
"{ckpt}: mean drift {ours:.4} exceeds twice the reference's own bf16 drift {theirs:.4}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Largest difference between two logit rows after removing each row's
|
||||
/// mean — the part of a difference a softmax can see.
|
||||
fn centred_diff(a: &[f32], b: &[f32]) -> f32 {
|
||||
let mean = |v: &[f32]| v.iter().sum::<f32>() / v.len() as f32;
|
||||
let (ma, mb) = (mean(a), mean(b));
|
||||
a.iter()
|
||||
.zip(b)
|
||||
.map(|(x, y)| ((x - ma) - (y - mb)).abs())
|
||||
.fold(0.0, f32::max)
|
||||
}
|
||||
|
||||
/// Forward latency for the four workloads #337 benchmarks against the
|
||||
/// PyTorch reference, plus batch throughput, on the English checkpoint.
|
||||
///
|
||||
/// Sequences are real ones from `reference.json`: `choice_basic` is the
|
||||
/// short state (41 tokens), `multi_question` five questions over it, and
|
||||
/// `state_right_truncated` a full 512-token sequence (repeated for the
|
||||
/// five-question and throughput batches). The timed span is the whole
|
||||
/// `forward` call — host-to-device copy, forward, and the CPU readback
|
||||
/// that ends it — so it is directly comparable to the reference's
|
||||
/// end-to-end `predict()` minus tokenisation.
|
||||
///
|
||||
/// Run with `cargo test --release -p neuron [--features cuda] --
|
||||
/// --ignored --nocapture laya_forward_timing`, with `LAYA_REFERENCE_DIR`
|
||||
/// and optionally `LAYA_DEVICE` / `LAYA_DTYPE` set.
|
||||
#[test]
|
||||
#[ignore = "benchmark: needs the released weights (LAYA_REFERENCE_DIR)"]
|
||||
fn laya_forward_timing() {
|
||||
let root = PathBuf::from(std::env::var("LAYA_REFERENCE_DIR").expect("LAYA_REFERENCE_DIR"));
|
||||
let (device, dtype) = test_device();
|
||||
let reference: Reference = serde_json::from_str(
|
||||
&std::fs::read_to_string(testdata().join("reference.json")).expect("read reference.json"),
|
||||
)
|
||||
.expect("parse reference.json");
|
||||
let seqs_of = |name: &str| -> Vec<DecisionSeq> {
|
||||
reference
|
||||
.cases
|
||||
.iter()
|
||||
.find(|c| c.name == name)
|
||||
.unwrap_or_else(|| panic!("no case {name}"))
|
||||
.items
|
||||
.iter()
|
||||
.map(|it| DecisionSeq {
|
||||
input_ids: it.input_ids.clone(),
|
||||
markers: it.markers.clone(),
|
||||
qtype: it.qtype,
|
||||
})
|
||||
.collect()
|
||||
};
|
||||
let long = seqs_of("state_right_truncated").remove(0);
|
||||
let workloads: Vec<(String, DecisionBatch)> = vec![
|
||||
(
|
||||
"short x1q".into(),
|
||||
DecisionBatch {
|
||||
seqs: seqs_of("choice_basic"),
|
||||
},
|
||||
),
|
||||
(
|
||||
"short x5q".into(),
|
||||
DecisionBatch {
|
||||
seqs: seqs_of("multi_question"),
|
||||
},
|
||||
),
|
||||
(
|
||||
"512tok x1q".into(),
|
||||
DecisionBatch {
|
||||
seqs: vec![long.clone()],
|
||||
},
|
||||
),
|
||||
(
|
||||
"512tok x5q".into(),
|
||||
DecisionBatch {
|
||||
seqs: vec![long.clone(); 5],
|
||||
},
|
||||
),
|
||||
(
|
||||
"512tok x16".into(),
|
||||
DecisionBatch {
|
||||
seqs: vec![long.clone(); 16],
|
||||
},
|
||||
),
|
||||
(
|
||||
"512tok x64".into(),
|
||||
DecisionBatch {
|
||||
seqs: vec![long.clone(); 64],
|
||||
},
|
||||
),
|
||||
];
|
||||
|
||||
let model = LayaModel::load(&root, &device, dtype).expect("load English checkpoint");
|
||||
// `LAYA_TIMING_ONLY` keeps the workloads whose name contains it, and
|
||||
// `LAYA_TIMING_ITERS` sets the timed iterations — for profiling one
|
||||
// workload under a tracer.
|
||||
let only = std::env::var("LAYA_TIMING_ONLY").ok();
|
||||
let iters: usize = std::env::var("LAYA_TIMING_ITERS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(20);
|
||||
for (name, batch) in &workloads {
|
||||
if only.as_deref().is_some_and(|o| !name.contains(o)) {
|
||||
continue;
|
||||
}
|
||||
for _ in 0..3 {
|
||||
model.forward(batch).expect("warm-up forward");
|
||||
}
|
||||
let mut ms: Vec<f64> = (0..iters)
|
||||
.map(|_| {
|
||||
let started = std::time::Instant::now();
|
||||
model.forward(batch).expect("forward");
|
||||
started.elapsed().as_secs_f64() * 1e3
|
||||
})
|
||||
.collect();
|
||||
ms.sort_by(f64::total_cmp);
|
||||
let median = ms[ms.len() / 2];
|
||||
let tokens: usize = batch.seqs.iter().map(|s| s.input_ids.len()).sum();
|
||||
eprintln!(
|
||||
"[{device:?} {dtype:?}] {name:>11}: median {median:7.1} ms, min {:7.1} ms, \
|
||||
{tokens} tokens, {:.0} seq/s",
|
||||
ms[0],
|
||||
batch.seqs.len() as f64 / (median / 1e3)
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -20,6 +20,7 @@
|
||||
//! TP-aware analogues live in `harness/tp/tp_<family>.rs` and follow
|
||||
//! the pattern set by `tp_qwen3.rs`.
|
||||
|
||||
pub mod laya;
|
||||
pub mod phase_probe;
|
||||
pub mod qwen3_5;
|
||||
pub mod qwen4_exp;
|
||||
|
||||
@@ -143,6 +143,9 @@ pub enum LoadedHandle {
|
||||
Tp(Arc<TpLoadedModel>),
|
||||
/// Z-Image text-to-image pipeline (#198).
|
||||
Image(Arc<LoadedImageModel>),
|
||||
/// Laya decision model (#336): typed questions in, calibrated
|
||||
/// option scores out, no text generation.
|
||||
Decision(Arc<super::decision_model::LoadedDecisionModel>),
|
||||
}
|
||||
|
||||
impl LoadedHandle {
|
||||
@@ -152,6 +155,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => &m.model_id,
|
||||
LoadedHandle::Image(m) => &m.model_id,
|
||||
LoadedHandle::Decision(m) => &m.model_id,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -162,6 +166,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => &m.spec,
|
||||
LoadedHandle::Image(m) => &m.spec,
|
||||
LoadedHandle::Decision(m) => &m.spec,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -171,6 +176,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.devices.clone(),
|
||||
LoadedHandle::Image(m) => m.devices.clone(),
|
||||
LoadedHandle::Decision(m) => m.devices.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -184,6 +190,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.poisoned.load(Ordering::Acquire),
|
||||
LoadedHandle::Image(m) => m.poisoned.load(Ordering::Acquire),
|
||||
LoadedHandle::Decision(m) => m.poisoned.load(Ordering::Acquire),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -195,6 +202,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => (m.admission.in_flight(), m.admission.queue_depth()),
|
||||
LoadedHandle::Image(m) => (m.admission.in_flight(), m.admission.queue_depth()),
|
||||
LoadedHandle::Decision(m) => (m.admission.in_flight(), m.admission.queue_depth()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -207,6 +215,9 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => (m.admission.max_in_flight(), m.admission.max_queue_depth()),
|
||||
LoadedHandle::Image(m) => (m.admission.max_in_flight(), m.admission.max_queue_depth()),
|
||||
LoadedHandle::Decision(m) => {
|
||||
(m.admission.max_in_flight(), m.admission.max_queue_depth())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -230,6 +241,10 @@ impl LoadedHandle {
|
||||
m.admission.anon_in_flight(),
|
||||
m.admission.anon_max_in_flight(),
|
||||
),
|
||||
LoadedHandle::Decision(m) => (
|
||||
m.admission.anon_in_flight(),
|
||||
m.admission.anon_max_in_flight(),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -240,6 +255,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.admission.close(),
|
||||
LoadedHandle::Image(m) => m.admission.close(),
|
||||
LoadedHandle::Decision(m) => m.admission.close(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -251,6 +267,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.admission.rejections(),
|
||||
LoadedHandle::Image(m) => m.admission.rejections(),
|
||||
LoadedHandle::Decision(m) => m.admission.rejections(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -262,6 +279,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => &m.admission,
|
||||
LoadedHandle::Image(m) => &m.admission,
|
||||
LoadedHandle::Decision(m) => &m.admission,
|
||||
};
|
||||
(admission.kv_budget_mb(), admission.kv_available_mb())
|
||||
}
|
||||
@@ -275,7 +293,7 @@ impl LoadedHandle {
|
||||
LoadedHandle::Tp(m) => &m.prefill_rate,
|
||||
// Image models have no token throughput; steps/sec lives in
|
||||
// the per-request timing instead.
|
||||
LoadedHandle::Image(_) => return (0.0, 0.0),
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return (0.0, 0.0),
|
||||
};
|
||||
(ema.get().unwrap_or(0.0), ema.decode().unwrap_or(0.0))
|
||||
}
|
||||
@@ -287,12 +305,15 @@ impl LoadedHandle {
|
||||
if matches!(self, LoadedHandle::Image(_)) {
|
||||
return vec!["image".to_string()];
|
||||
}
|
||||
if matches!(self, LoadedHandle::Decision(_)) {
|
||||
return vec!["decision".to_string()];
|
||||
}
|
||||
let mut caps = vec!["text".to_string()];
|
||||
let has_vision = match self {
|
||||
LoadedHandle::Single(m) => m.has_vision,
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.has_vision,
|
||||
LoadedHandle::Image(_) => unreachable!("handled above"),
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => unreachable!("handled above"),
|
||||
};
|
||||
if has_vision {
|
||||
caps.push("vision".to_string());
|
||||
@@ -338,7 +359,7 @@ impl LoadedHandle {
|
||||
LoadedHandle::Tp(m) => (m.last_free_mb.load(Ordering::Acquire), m.context_profile),
|
||||
// Image generation has its own VRAM path and no prefill
|
||||
// floor, so this check would say nothing useful about it.
|
||||
LoadedHandle::Image(_) => return None,
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return None,
|
||||
};
|
||||
// `0` is the CPU-load and missing-worker sentinel — the same one
|
||||
// `check_vram` skips on. Nothing to assert either way.
|
||||
@@ -419,7 +440,7 @@ impl LoadedHandle {
|
||||
LoadedHandle::Single(m) => m.tool_call_tokens.is_some(),
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.tool_call_tokens.is_some(),
|
||||
LoadedHandle::Image(_) => false,
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -430,7 +451,7 @@ impl LoadedHandle {
|
||||
LoadedHandle::Single(m) => m.reasoning_tokens.is_some(),
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.reasoning_tokens.is_some(),
|
||||
LoadedHandle::Image(_) => false,
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -441,7 +462,7 @@ impl LoadedHandle {
|
||||
LoadedHandle::Single(m) => m.preserve_thinking,
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.preserve_thinking,
|
||||
LoadedHandle::Image(_) => None,
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -451,7 +472,7 @@ impl LoadedHandle {
|
||||
LoadedHandle::Single(m) => m.reasoning_efforts.clone(),
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.reasoning_efforts.clone(),
|
||||
LoadedHandle::Image(_) => Default::default(),
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -477,7 +498,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.query_vram_tightest_free_mb().await,
|
||||
// No context-limit derivation for image models; nothing to cache.
|
||||
LoadedHandle::Image(_) => return,
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return,
|
||||
};
|
||||
// Don't clobber a good cached value with a transient `0`
|
||||
// (worker gone/poisoned sentinel).
|
||||
@@ -486,7 +507,7 @@ impl LoadedHandle {
|
||||
LoadedHandle::Single(m) => m.last_free_mb.store(free, Ordering::Release),
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.last_free_mb.store(free, Ordering::Release),
|
||||
LoadedHandle::Image(_) => {}
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -515,7 +536,7 @@ impl LoadedHandle {
|
||||
m.prefill_rate.get(),
|
||||
),
|
||||
// Token limits are meaningless for image generation.
|
||||
LoadedHandle::Image(_) => return None,
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return None,
|
||||
};
|
||||
let rate = rate.unwrap_or(cfg.bootstrap_prefill_tok_per_sec);
|
||||
// The KV budget admission will actually grant. Advertising past
|
||||
@@ -525,7 +546,9 @@ impl LoadedHandle {
|
||||
LoadedHandle::Single(m) => m.admission.kv_budget_mb(),
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.admission.kv_budget_mb(),
|
||||
LoadedHandle::Image(_) => unreachable!("derived_limit returned above"),
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {
|
||||
unreachable!("derived_limit returned above")
|
||||
}
|
||||
};
|
||||
let limit = super::context_limit::derive_limit(
|
||||
&profile,
|
||||
@@ -544,7 +567,9 @@ impl LoadedHandle {
|
||||
LoadedHandle::Single(m) => &m.derived_input_cap,
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => &m.derived_input_cap,
|
||||
LoadedHandle::Image(_) => unreachable!("derived_limit returned above"),
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {
|
||||
unreachable!("derived_limit returned above")
|
||||
}
|
||||
};
|
||||
cap.store(input, Ordering::Release);
|
||||
}
|
||||
@@ -3282,7 +3307,7 @@ impl CandleHarness {
|
||||
LoadedHandle::Tp(m) => {
|
||||
return self.inference_tp_stream(m, request, principal, wire).await;
|
||||
}
|
||||
LoadedHandle::Image(_) => {
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {
|
||||
return Err(InferenceError::WrongModality {
|
||||
model_id: request.model.clone(),
|
||||
});
|
||||
@@ -4001,6 +4026,11 @@ impl Harness for CandleHarness {
|
||||
if matches!(plan.format, super::preflight::SourceFormat::Diffusers) {
|
||||
return self.load_image_model(spec, &source_id).await;
|
||||
}
|
||||
// Laya decision models (#336) take the decision path: one
|
||||
// device, unquantized (preflight rejected tp>1 and quant).
|
||||
if matches!(plan.format, super::preflight::SourceFormat::DecisionModel) {
|
||||
return self.load_decision_model(spec, &source_id).await;
|
||||
}
|
||||
|
||||
let tp_size = spec.tensor_parallel.unwrap_or(1);
|
||||
if tp_size > 1 {
|
||||
@@ -4367,6 +4397,11 @@ impl Harness for CandleHarness {
|
||||
);
|
||||
}
|
||||
}
|
||||
LoadedHandle::Decision(dm) => {
|
||||
// Each checkpoint's weights drop on the thread that owns
|
||||
// them (the device worker, for a CUDA load).
|
||||
dm.release().await;
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(tp) => {
|
||||
// Try to recover the inner TpLoadedModel so we can move
|
||||
@@ -4541,6 +4576,132 @@ impl CandleHarness {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Load a Laya decision model (#336): resolve the checkpoints' files,
|
||||
/// build them on the device worker (or the CPU), and register the
|
||||
/// model with `capabilities: ["decision"]`.
|
||||
///
|
||||
/// The repo root is required. A released family also carries
|
||||
/// `multilingual/` and `typed-decisions/` beside it (#339); each one
|
||||
/// the repo has is loaded too, so the `/v1/systemone` handler can
|
||||
/// route non-English states to the multilingual checkpoint and honour
|
||||
/// a request that pins one. A repo without them (a single fine-tune)
|
||||
/// serves its root alone.
|
||||
async fn load_decision_model(
|
||||
&self,
|
||||
spec: &ModelSpec,
|
||||
source_id: &cortex_core::source::ModelSourceId,
|
||||
) -> Result<()> {
|
||||
let devices = spec.devices.clone().unwrap_or_else(|| vec![0]);
|
||||
let device = Self::pick_device(&devices)?;
|
||||
#[cfg(feature = "cuda")]
|
||||
let worker = if matches!(device, Device::Cuda(_)) {
|
||||
Some(self.ensure_device_worker(devices[0]).await?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
// CPU build: always the CPU backend. Correct, but at hundreds of
|
||||
// milliseconds per question it is a fallback, not a serving tier.
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
let worker = {
|
||||
let _ = &device;
|
||||
None
|
||||
};
|
||||
|
||||
let files = self.resolve_decision_files(source_id, None).await?;
|
||||
let root = super::decision_model::DecisionCheckpoint::load(
|
||||
super::decision_model::ROOT_CHECKPOINT.to_string(),
|
||||
files,
|
||||
&spec.model_id,
|
||||
worker.clone(),
|
||||
)
|
||||
.await?;
|
||||
let dtype = root.dtype;
|
||||
let mut checkpoints = vec![root];
|
||||
for sub in super::decision::route::FAMILY
|
||||
.iter()
|
||||
.filter(|c| **c != super::decision_model::ROOT_CHECKPOINT)
|
||||
{
|
||||
let files = match self.resolve_decision_files(source_id, Some(sub)).await {
|
||||
Ok(files) => files,
|
||||
Err(e) => {
|
||||
tracing::info!(
|
||||
model = %spec.model_id,
|
||||
checkpoint = sub,
|
||||
error = %format!("{e:#}"),
|
||||
"decision checkpoint not available; serving without it"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
checkpoints.push(
|
||||
super::decision_model::DecisionCheckpoint::load(
|
||||
sub.to_string(),
|
||||
files,
|
||||
&spec.model_id,
|
||||
worker.clone(),
|
||||
)
|
||||
.await?,
|
||||
);
|
||||
}
|
||||
let names: Vec<&str> = checkpoints.iter().map(|c| c.name.as_str()).collect();
|
||||
tracing::info!(model = %spec.model_id, checkpoints = ?names, "decision checkpoints loaded");
|
||||
|
||||
let loaded = super::decision_model::LoadedDecisionModel {
|
||||
model_id: spec.model_id.clone(),
|
||||
spec: spec.clone(),
|
||||
devices,
|
||||
checkpoints,
|
||||
poisoned: Arc::new(AtomicBool::new(false)),
|
||||
admission: crate::harness::admission::AdmissionController::new(&self.admission_cfg),
|
||||
};
|
||||
let mut models = self.models.write().await;
|
||||
models.insert(
|
||||
spec.model_id.clone(),
|
||||
LoadedHandle::Decision(Arc::new(loaded)),
|
||||
);
|
||||
tracing::info!(model = %spec.model_id, dtype = ?dtype, "decision model loaded");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Fetch one Laya checkpoint's files through hf-hub: the repo root
|
||||
/// for `subfolder = None`, else that subfolder. Only the files the
|
||||
/// checkpoint needs are fetched, not the whole repo.
|
||||
async fn resolve_decision_files(
|
||||
&self,
|
||||
source_id: &cortex_core::source::ModelSourceId,
|
||||
subfolder: Option<&str>,
|
||||
) -> Result<super::decision_model::CheckpointFiles> {
|
||||
let api = self.hf_api_for(&source_id.scheme)?;
|
||||
let repo = api.model(source_id.repo_path());
|
||||
let display_id = source_id.to_string();
|
||||
let prefix = subfolder.map(|s| format!("{s}/")).unwrap_or_default();
|
||||
let mut fetched = Vec::new();
|
||||
for file in [
|
||||
"rl_agent_config.json",
|
||||
"encoder/config.json",
|
||||
"model.safetensors",
|
||||
"tokenizer/tokenizer.json",
|
||||
"tokenizer/tokenizer_config.json",
|
||||
] {
|
||||
let path = format!("{prefix}{file}");
|
||||
fetched.push(
|
||||
repo.get(&path)
|
||||
.await
|
||||
.with_context(|| format!("fetch {path} from {display_id}"))?,
|
||||
);
|
||||
}
|
||||
// hf-hub materialises a snapshot as a directory tree mirroring
|
||||
// the repo, so the checkpoint directory is where its config sits.
|
||||
let dir = fetched[0]
|
||||
.parent()
|
||||
.context("checkpoint config has no parent directory")?
|
||||
.to_path_buf();
|
||||
Ok(super::decision_model::CheckpointFiles {
|
||||
dir,
|
||||
tokenizer: fetched[3].clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Resolve a diffusers-layout Z-Image repo into local component
|
||||
/// paths, downloading through hf-hub as needed. Verifies the
|
||||
/// pipeline class before fetching multi-GB weight shards.
|
||||
@@ -4654,6 +4815,212 @@ impl CandleHarness {
|
||||
}
|
||||
}
|
||||
|
||||
/// The loaded decision model `model_id` (#336), for the
|
||||
/// `/v1/systemone` handler: it takes the admission permit, scores
|
||||
/// through [`super::decision_model::DecisionCheckpoint::decide`],
|
||||
/// and shapes the answers. A model of another kind is a
|
||||
/// `WrongModality`, as for images.
|
||||
pub async fn decision_model(
|
||||
&self,
|
||||
model_id: &str,
|
||||
) -> Result<Arc<super::decision_model::LoadedDecisionModel>, InferenceError> {
|
||||
let handle = {
|
||||
let models = self.models.read().await;
|
||||
models.get(model_id).cloned()
|
||||
};
|
||||
match handle {
|
||||
Some(LoadedHandle::Decision(m)) => Ok(m),
|
||||
Some(_) => Err(InferenceError::WrongModality {
|
||||
model_id: model_id.to_string(),
|
||||
}),
|
||||
None => Err(InferenceError::ModelNotLoaded(model_id.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Answer one `/v1/systemone` request (#338): pick the served decision
|
||||
/// model and the checkpoint of it that answers, assemble the question
|
||||
/// sequences, score them in one forward pass, and shape the answers.
|
||||
///
|
||||
/// The served model is `served` when the gateway routed the request
|
||||
/// (the `x-helexa-decision-model` header), else the model the body's
|
||||
/// `model` names if one is loaded, else the one decision model this
|
||||
/// node serves. The body's `model` is then read a second time, as the
|
||||
/// reference server reads it: as an optional pin to a checkpoint of
|
||||
/// the family (`multilingual`, `typed-decisions`, …).
|
||||
///
|
||||
/// Questions are validated and tokenised before admission, so a
|
||||
/// malformed request never takes a slot.
|
||||
pub async fn systemone(
|
||||
&self,
|
||||
served: Option<&str>,
|
||||
request: cortex_core::decisions::SystemOneRequest,
|
||||
principal: Option<&str>,
|
||||
) -> Result<cortex_core::decisions::SystemOneResponse, SystemOneError> {
|
||||
use super::decision::{self, route};
|
||||
|
||||
let model = match served {
|
||||
Some(id) => self.decision_model(id).await?,
|
||||
None => {
|
||||
self.default_decision_model(request.model.as_deref())
|
||||
.await?
|
||||
}
|
||||
};
|
||||
if model.poisoned.load(Ordering::Acquire) {
|
||||
return Err(InferenceError::Other(anyhow::anyhow!(
|
||||
"decision model '{}' is poisoned; unload and reload it",
|
||||
model.model_id
|
||||
))
|
||||
.into());
|
||||
}
|
||||
|
||||
let pin = route::resolve_pin(request.model.as_deref());
|
||||
let mut routing = route::route(&model.model_id, &request.state, &request.questions, pin);
|
||||
let checkpoint = match model.checkpoint(Some(&routing.model)) {
|
||||
Some(c) => c,
|
||||
None if pin.is_some() => {
|
||||
return Err(decision::DecisionRequestError::new(
|
||||
404,
|
||||
format!(
|
||||
"checkpoint '{}' is not served by '{}' on this node",
|
||||
routing.model, model.model_id
|
||||
),
|
||||
)
|
||||
.into());
|
||||
}
|
||||
// Routed (not pinned) to a checkpoint this repo does not
|
||||
// carry: answer with the default one, and say so.
|
||||
None => {
|
||||
let fallback = model.checkpoint(None).ok_or_else(|| {
|
||||
InferenceError::Other(anyhow::anyhow!(
|
||||
"decision model '{}' has no checkpoints",
|
||||
model.model_id
|
||||
))
|
||||
})?;
|
||||
routing.reason = format!(
|
||||
"{}; the {} checkpoint is not loaded, so {} answered",
|
||||
routing.reason, routing.model, fallback.name
|
||||
);
|
||||
routing.model = fallback.name.clone();
|
||||
routing.repo = route::repo_of(&model.model_id, &fallback.name);
|
||||
fallback
|
||||
}
|
||||
};
|
||||
let lang = route::detected_language(&routing).map(str::to_string);
|
||||
|
||||
let prepared = decision::prepare(
|
||||
&request.state,
|
||||
&request.questions,
|
||||
&checkpoint.calibration,
|
||||
&checkpoint.tokenizer,
|
||||
decision::Budget {
|
||||
max_len: request.max_len,
|
||||
head_max_len: request.head_max_len,
|
||||
},
|
||||
)?;
|
||||
|
||||
let _permit = model.admission.enter(principal).await?;
|
||||
let outputs: Vec<decision::RowOutput> = if prepared.rows.is_empty() {
|
||||
Vec::new()
|
||||
} else {
|
||||
let batch = super::arch::laya::DecisionBatch {
|
||||
seqs: prepared
|
||||
.rows
|
||||
.iter()
|
||||
.map(|r| super::arch::laya::DecisionSeq {
|
||||
input_ids: r.input_ids.clone(),
|
||||
markers: r.markers.clone(),
|
||||
qtype: r.qtype as u8,
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
checkpoint
|
||||
.decide(batch)
|
||||
.await
|
||||
.map_err(|e| InferenceError::Other(anyhow::anyhow!("decision forward: {e:#}")))?
|
||||
.into_iter()
|
||||
.map(|row| decision::RowOutput {
|
||||
act_probability: row.act.first().copied().unwrap_or(0.0),
|
||||
logits: row.logits,
|
||||
})
|
||||
.collect()
|
||||
};
|
||||
let (answers, usage) = decision::shape(
|
||||
&prepared,
|
||||
&outputs,
|
||||
&checkpoint.calibration,
|
||||
lang.as_deref(),
|
||||
);
|
||||
Ok(decision::shape::response(answers, usage, Some(routing)))
|
||||
}
|
||||
|
||||
/// Register an already-built decision model, bypassing the hub fetch
|
||||
/// `load_model` does — for tests that serve local checkpoints.
|
||||
#[cfg(test)]
|
||||
pub(crate) async fn register_decision_model(
|
||||
&self,
|
||||
model: super::decision_model::LoadedDecisionModel,
|
||||
) {
|
||||
self.models.write().await.insert(
|
||||
model.model_id.clone(),
|
||||
LoadedHandle::Decision(Arc::new(model)),
|
||||
);
|
||||
}
|
||||
|
||||
/// The decision model to serve a request no gateway routed: the one
|
||||
/// `requested` names if it is loaded, else the only (first) loaded
|
||||
/// decision model.
|
||||
async fn default_decision_model(
|
||||
&self,
|
||||
requested: Option<&str>,
|
||||
) -> Result<Arc<super::decision_model::LoadedDecisionModel>, InferenceError> {
|
||||
let models = self.models.read().await;
|
||||
if let Some(id) = requested
|
||||
&& let Some(handle) = models.get(id)
|
||||
{
|
||||
return match handle {
|
||||
LoadedHandle::Decision(m) => Ok(Arc::clone(m)),
|
||||
_ => Err(InferenceError::WrongModality {
|
||||
model_id: id.to_string(),
|
||||
}),
|
||||
};
|
||||
}
|
||||
let mut decision: Vec<&Arc<super::decision_model::LoadedDecisionModel>> = models
|
||||
.values()
|
||||
.filter_map(|h| match h {
|
||||
LoadedHandle::Decision(m) => Some(m),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
decision.sort_by(|a, b| a.model_id.cmp(&b.model_id));
|
||||
decision
|
||||
.first()
|
||||
.map(|m| Arc::clone(m))
|
||||
.ok_or_else(|| InferenceError::ModelNotLoaded("no decision model is loaded".into()))
|
||||
}
|
||||
|
||||
/// Render one audio generation (#382). The cortex job manager sends one
|
||||
/// of these per job and waits for the whole render; the reply is the
|
||||
/// WAV bytes plus the result metadata cortex records on the job.
|
||||
///
|
||||
/// No audio architecture is loadable yet, so every model reachable
|
||||
/// here is the wrong modality or not loaded. The ACE-Step engine
|
||||
/// (#383) adds a `LoadedHandle::Audio` arm that validates against the
|
||||
/// model's limits, takes an admission permit and runs the device job.
|
||||
pub async fn audio_generation(
|
||||
&self,
|
||||
request: &cortex_core::audio::AudioGenerationRequest,
|
||||
_principal: Option<&str>,
|
||||
) -> Result<(Vec<u8>, cortex_core::audio::NeuronAudioResult), InferenceError> {
|
||||
let loaded = self.models.read().await.contains_key(&request.model);
|
||||
if loaded {
|
||||
Err(InferenceError::WrongModality {
|
||||
model_id: request.model.clone(),
|
||||
})
|
||||
} else {
|
||||
Err(InferenceError::ModelNotLoaded(request.model.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
/// Run one text-to-image generation (#198). This is the internal
|
||||
/// path the `/v1/images/generations` handler (#200) fronts;
|
||||
/// admission, validation, tokenization, and PNG encoding happen
|
||||
@@ -6278,6 +6645,33 @@ pub enum InferenceError {
|
||||
Other(#[from] anyhow::Error),
|
||||
}
|
||||
|
||||
/// Why a `/v1/systemone` request was not answered (#338): a request the
|
||||
/// reference server would also reject (its status and `detail` text), or
|
||||
/// a serving failure shared with the other endpoints.
|
||||
#[derive(Debug)]
|
||||
pub enum SystemOneError {
|
||||
Request(super::decision::DecisionRequestError),
|
||||
Inference(InferenceError),
|
||||
}
|
||||
|
||||
impl From<super::decision::DecisionRequestError> for SystemOneError {
|
||||
fn from(e: super::decision::DecisionRequestError) -> Self {
|
||||
SystemOneError::Request(e)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<InferenceError> for SystemOneError {
|
||||
fn from(e: InferenceError) -> Self {
|
||||
SystemOneError::Inference(e)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<super::admission::AdmissionRejection> for SystemOneError {
|
||||
fn from(e: super::admission::AdmissionRejection) -> Self {
|
||||
SystemOneError::Inference(e.into())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<super::admission::AdmissionRejection> for InferenceError {
|
||||
fn from(rejection: super::admission::AdmissionRejection) -> Self {
|
||||
use super::admission::AdmissionRejection;
|
||||
|
||||
@@ -57,6 +57,9 @@ pub fn classify_files(
|
||||
) -> Vec<String> {
|
||||
match preflight::classify(filenames) {
|
||||
SourceFormat::Diffusers => vec!["image".to_string()],
|
||||
// A decision model answers typed questions (`/v1/systemone`);
|
||||
// it generates no text, so it must not claim "text".
|
||||
SourceFormat::DecisionModel => vec!["decision".to_string()],
|
||||
// A repo with no recognised weights (tokenizer-only, or an empty
|
||||
// entry) is not servable, so claiming "text" would be a lie of
|
||||
// the same kind this module exists to remove.
|
||||
@@ -101,6 +104,22 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn laya_repo_is_a_decision_model() {
|
||||
let files = [
|
||||
"rl_agent_config.json",
|
||||
"encoder/config.json",
|
||||
"model.safetensors",
|
||||
"tokenizer/tokenizer.json",
|
||||
];
|
||||
assert_eq!(
|
||||
classify_files(&files, || panic!(
|
||||
"config must not be read for a decision model"
|
||||
)),
|
||||
vec!["decision".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dense_repo_without_vision_config_is_text_only() {
|
||||
let files = ["config.json", "tokenizer.json", "model.safetensors"];
|
||||
|
||||
@@ -0,0 +1,233 @@
|
||||
//! Per-checkpoint decision settings, read from `rl_agent_config.json`.
|
||||
//!
|
||||
//! Parsed with [`Json`], not `serde_json`: the fitted temperatures are
|
||||
//! written with 17 significant digits, and `serde_json`'s default float
|
||||
//! parser is not correctly rounded (it reads `1.7601518630981445` as
|
||||
//! `…443`).
|
||||
|
||||
use cortex_core::decisions::Json;
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Default token budgets when a checkpoint's config omits them.
|
||||
pub const DEFAULT_MAX_LEN: usize = 512;
|
||||
pub const DEFAULT_HEAD_MAX_LEN: usize = 192;
|
||||
|
||||
/// Bounds on a usable calibration temperature. A fitted value below 1
|
||||
/// sharpens instead of softening; the shipped `choice:11+` bucket is 0.1006,
|
||||
/// which would publish a coin flip as a certainty, so the reference refuses
|
||||
/// to apply anything outside this range.
|
||||
pub const TEMP_MIN: f64 = 0.5;
|
||||
pub const TEMP_MAX: f64 = 5.0;
|
||||
|
||||
/// A temperature confined to [`TEMP_MIN`, `TEMP_MAX`]; a non-finite or
|
||||
/// non-numeric value falls back to 1.0.
|
||||
pub fn clamp_temperature(t: Option<f64>) -> f64 {
|
||||
match t {
|
||||
Some(t) if t.is_finite() => t.clamp(TEMP_MIN, TEMP_MAX),
|
||||
_ => 1.0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Python's `float(v)` for a JSON value: numbers, bools, and numeric strings.
|
||||
fn py_float(v: &Json) -> Option<f64> {
|
||||
match v {
|
||||
Json::Int(d) => d.parse().ok(),
|
||||
Json::Float(x) => Some(*x),
|
||||
Json::Bool(b) => Some(if *b { 1.0 } else { 0.0 }),
|
||||
Json::String(s) => s.trim().parse().ok(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// One set of calibration temperatures: per question type, overridden per
|
||||
/// `(type, option-count bucket)` (see [`temp_bucket`]). Already clamped.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct Temperatures {
|
||||
/// Indexed by [`super::QType`] (`choice`, `score`, `noul`).
|
||||
pub by_type: [f64; 3],
|
||||
/// `"choice:3-5"` → temperature, and so on.
|
||||
pub by_options: HashMap<String, f64>,
|
||||
}
|
||||
|
||||
impl Temperatures {
|
||||
/// The temperature for a question of type index `qtype` with `k` options.
|
||||
pub fn for_question(&self, qtype: usize, k: usize) -> f64 {
|
||||
self.by_options
|
||||
.get(&temp_bucket(qtype, k))
|
||||
.copied()
|
||||
.unwrap_or(self.by_type[qtype])
|
||||
}
|
||||
|
||||
fn from_json(
|
||||
temperature: Option<&Json>,
|
||||
by_options: Option<&Json>,
|
||||
fallback: Option<[f64; 3]>,
|
||||
) -> Result<Self, String> {
|
||||
let raw = match temperature {
|
||||
None => fallback.unwrap_or([1.0; 3]).map(Some),
|
||||
Some(Json::Array(a)) if a.len() == 3 => {
|
||||
[py_float(&a[0]), py_float(&a[1]), py_float(&a[2])]
|
||||
}
|
||||
Some(other) => {
|
||||
return Err(format!(
|
||||
"temperature must be a list of 3 floats, got {}",
|
||||
other.py_repr()
|
||||
));
|
||||
}
|
||||
};
|
||||
let by_options = match by_options {
|
||||
Some(Json::Object(m)) => m
|
||||
.iter()
|
||||
.map(|(k, v)| (k.clone(), clamp_temperature(py_float(v))))
|
||||
.collect(),
|
||||
_ => HashMap::new(),
|
||||
};
|
||||
Ok(Self {
|
||||
by_type: raw.map(clamp_temperature),
|
||||
by_options,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// The calibration bucket of a question: its type and option count.
|
||||
pub fn temp_bucket(qtype: usize, k: usize) -> String {
|
||||
let size = match k {
|
||||
0..=2 => "2",
|
||||
3..=5 => "3-5",
|
||||
6..=10 => "6-10",
|
||||
_ => "11+",
|
||||
};
|
||||
format!("{}:{size}", super::QType::NAMES[qtype])
|
||||
}
|
||||
|
||||
/// What the sequence builder and the calibration need from a checkpoint.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct DecisionConfig {
|
||||
/// Total sequence budget.
|
||||
pub max_len: usize,
|
||||
/// Budget for the question text and every option together.
|
||||
pub head_max_len: usize,
|
||||
pub temperatures: Temperatures,
|
||||
/// Per-language overrides, keyed by the lower-cased primary subtag
|
||||
/// (`"de"` for `de-AT`). Empty for the shipped checkpoints; an operator
|
||||
/// may fit them per language.
|
||||
pub lang_temperatures: HashMap<String, Temperatures>,
|
||||
}
|
||||
|
||||
impl DecisionConfig {
|
||||
/// Parse a checkpoint's `rl_agent_config.json`.
|
||||
pub fn from_json(text: &str) -> Result<Self, String> {
|
||||
let v = Json::parse(text.as_bytes()).map_err(|e| format!("rl_agent_config.json: {e}"))?;
|
||||
let usize_or = |key: &str, default: usize| match v.get(key) {
|
||||
Some(Json::Int(d)) => d.parse().unwrap_or(default),
|
||||
_ => default,
|
||||
};
|
||||
Ok(Self {
|
||||
max_len: usize_or("max_len", DEFAULT_MAX_LEN),
|
||||
head_max_len: usize_or("head_max_len", DEFAULT_HEAD_MAX_LEN),
|
||||
temperatures: Temperatures::from_json(
|
||||
v.get("temperature"),
|
||||
v.get("temperature_by_options"),
|
||||
None,
|
||||
)?,
|
||||
lang_temperatures: HashMap::new(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Add per-language temperatures, the shape the reference takes them in:
|
||||
/// `{"de": {"temperature": [..3], "temperature_by_options": {..}}}`. A
|
||||
/// language that omits `temperature` inherits the checkpoint's raw
|
||||
/// per-type values.
|
||||
pub fn with_lang_temperatures(mut self, spec: &Json) -> Result<Self, String> {
|
||||
let Json::Object(m) = spec else {
|
||||
return Err("lang_temperatures must be an object".into());
|
||||
};
|
||||
for (lang, cfg) in m {
|
||||
let t = Temperatures::from_json(
|
||||
cfg.get("temperature"),
|
||||
cfg.get("temperature_by_options"),
|
||||
Some(self.temperatures.by_type),
|
||||
)
|
||||
.map_err(|e| format!("language override {lang:?}: {e}"))?;
|
||||
let primary = lang.split('-').next().unwrap_or(lang).to_lowercase();
|
||||
self.lang_temperatures.insert(primary, t);
|
||||
}
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
/// The temperatures to apply for a request detected (or declared) as
|
||||
/// `lang`, which may carry a region (`de-AT`).
|
||||
pub fn temperatures_for(&self, lang: Option<&str>) -> &Temperatures {
|
||||
lang.and_then(|l| {
|
||||
let primary = l.split('-').next().unwrap_or(l).to_lowercase();
|
||||
self.lang_temperatures.get(&primary)
|
||||
})
|
||||
.unwrap_or(&self.temperatures)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
const ENGLISH: &str = r#"{"encoder": "answerdotai/ModernBERT-large", "head_layers": 2,
|
||||
"max_len": 512, "head_max_len": 192,
|
||||
"temperature": [1.6369030475616455, 1.2514300346374512, 1.983399510383606],
|
||||
"temperature_by_options": {"choice:3-5": 1.7601518630981445, "choice:11+": 0.10058280825614929}}"#;
|
||||
|
||||
#[test]
|
||||
fn shipped_choice_11_bucket_is_clamped() {
|
||||
let c = DecisionConfig::from_json(ENGLISH).unwrap();
|
||||
assert_eq!(c.temperatures.for_question(0, 12), TEMP_MIN);
|
||||
assert_eq!(c.temperatures.for_question(0, 4), 1.7601518630981445);
|
||||
// No bucket: the per-type value.
|
||||
assert_eq!(c.temperatures.for_question(1, 7), 1.2514300346374512);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn buckets_split_where_the_reference_splits() {
|
||||
let b: Vec<_> = [1, 2, 3, 5, 6, 10, 11]
|
||||
.iter()
|
||||
.map(|&k| temp_bucket(2, k))
|
||||
.collect();
|
||||
assert_eq!(
|
||||
b,
|
||||
[
|
||||
"noul:2",
|
||||
"noul:2",
|
||||
"noul:3-5",
|
||||
"noul:3-5",
|
||||
"noul:6-10",
|
||||
"noul:6-10",
|
||||
"noul:11+"
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn language_overrides_match_on_primary_subtag() {
|
||||
let c = DecisionConfig::from_json(ENGLISH)
|
||||
.unwrap()
|
||||
.with_lang_temperatures(
|
||||
&Json::parse(br#"{"DE": {"temperature": [1.0, 1.4, 9.0]}}"#).unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
c.temperatures_for(Some("de-AT")).by_type,
|
||||
[1.0, 1.4, TEMP_MAX]
|
||||
);
|
||||
assert_eq!(c.temperatures_for(Some("fr")), &c.temperatures);
|
||||
assert_eq!(c.temperatures_for(None), &c.temperatures);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bad_temperatures_are_refused_or_neutralised() {
|
||||
assert!(DecisionConfig::from_json(r#"{"temperature": [1.0]}"#).is_err());
|
||||
let c = DecisionConfig::from_json(r#"{"temperature": [" 2.5 ", "x", null]}"#).unwrap();
|
||||
assert_eq!(c.temperatures.by_type, [2.5, 1.0, 1.0]);
|
||||
assert_eq!(
|
||||
(c.max_len, c.head_max_len),
|
||||
(DEFAULT_MAX_LEN, DEFAULT_HEAD_MAX_LEN)
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,249 @@
|
||||
//! `CandleHarness::systemone` end to end: request in, routed, validated,
|
||||
//! scored and shaped response out (#338).
|
||||
//!
|
||||
//! The tiny checkpoint (`testdata/laya/tiny/`) exercises the plumbing in
|
||||
//! CI. The released-weights test replays every reference case through
|
||||
//! the same entry point the HTTP handler calls and compares the whole
|
||||
//! response — answers, usage and routing — with what the reference
|
||||
//! server returned; it needs `LAYA_REFERENCE_DIR`.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
|
||||
use cortex_core::decisions::{Json, SystemOneRequest, limits};
|
||||
|
||||
use crate::harness::candle::{CandleHarness, SystemOneError};
|
||||
use crate::harness::decision_model::{CheckpointFiles, DecisionCheckpoint, LoadedDecisionModel};
|
||||
|
||||
fn testdata() -> PathBuf {
|
||||
Path::new(env!("CARGO_MANIFEST_DIR")).join("src/harness/testdata/laya")
|
||||
}
|
||||
|
||||
fn harness() -> Arc<CandleHarness> {
|
||||
CandleHarness::new(
|
||||
"http://localhost:13131".into(),
|
||||
&crate::config::CandleHarnessConfig::default(),
|
||||
)
|
||||
}
|
||||
|
||||
/// A decision model made of the given `(routing name, directory)`
|
||||
/// checkpoints, loaded on the CPU in f32.
|
||||
async fn decision_model(model_id: &str, checkpoints: &[(&str, PathBuf)]) -> LoadedDecisionModel {
|
||||
let mut loaded = Vec::new();
|
||||
for (name, dir) in checkpoints {
|
||||
let files = CheckpointFiles {
|
||||
tokenizer: if dir.join("tokenizer/tokenizer.json").exists() {
|
||||
dir.join("tokenizer/tokenizer.json")
|
||||
} else {
|
||||
dir.join("tokenizer.json")
|
||||
},
|
||||
dir: dir.clone(),
|
||||
};
|
||||
loaded.push(
|
||||
DecisionCheckpoint::load(name.to_string(), files, model_id, None)
|
||||
.await
|
||||
.unwrap_or_else(|e| panic!("load {name}: {e:#}")),
|
||||
);
|
||||
}
|
||||
LoadedDecisionModel {
|
||||
model_id: model_id.to_string(),
|
||||
spec: serde_json::from_value(serde_json::json!({
|
||||
"model_id": model_id,
|
||||
"harness": "candle",
|
||||
}))
|
||||
.expect("model spec"),
|
||||
devices: vec![0],
|
||||
checkpoints: loaded,
|
||||
poisoned: Arc::new(AtomicBool::new(false)),
|
||||
admission: crate::harness::admission::AdmissionController::new(
|
||||
&crate::config::AdmissionConfig::default(),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn request(body: &str) -> SystemOneRequest {
|
||||
SystemOneRequest::parse(body.as_bytes(), limits::DEFAULT_MAX_TOKEN_BUDGET)
|
||||
.expect("valid request")
|
||||
}
|
||||
|
||||
async fn tiny_harness() -> Arc<CandleHarness> {
|
||||
let h = harness();
|
||||
h.register_decision_model(
|
||||
decision_model("test/tiny", &[("english", testdata().join("tiny"))]).await,
|
||||
)
|
||||
.await;
|
||||
h
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn answers_every_question_in_order_with_exact_usage() {
|
||||
let h = tiny_harness().await;
|
||||
let req = request(
|
||||
r#"{"state": "a b a", "questions": {
|
||||
"z_last_alphabetically": {"type": "noul", "instructions": "a?"},
|
||||
"a_first": {"type": "choice", "instructions": "b?", "criteria": ["a", "b", "a b"]}
|
||||
}}"#,
|
||||
);
|
||||
let resp = h.systemone(None, req, None).await.expect("answered");
|
||||
let json = serde_json::to_value(&resp).unwrap();
|
||||
// Answers come back in request order, not sorted.
|
||||
let keys: Vec<&str> = resp.answers.0.iter().map(|(k, _)| k.as_str()).collect();
|
||||
assert_eq!(keys, ["z_last_alphabetically", "a_first"]);
|
||||
assert_eq!(json["answers"]["z_last_alphabetically"]["type"], "noul");
|
||||
let probs = json["answers"]["a_first"]["probabilities"]
|
||||
.as_object()
|
||||
.unwrap();
|
||||
let total: f64 = probs.values().map(|p| p.as_f64().unwrap()).sum();
|
||||
assert!((total - 1.0).abs() < 1e-3, "{probs:?}");
|
||||
// Usage is the exact number of tokens encoded, output tokens zero.
|
||||
assert!(json["usage"]["input_tokens"].as_u64().unwrap() > 0);
|
||||
assert_eq!(json["usage"]["output_tokens"], 0);
|
||||
assert_eq!(json["routing"]["model"], "english");
|
||||
assert_eq!(json["routing"]["repo"], "test/tiny");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pinning_a_checkpoint_the_model_lacks_is_404() {
|
||||
let h = tiny_harness().await;
|
||||
let req = request(
|
||||
r#"{"state": "a", "model": "multilingual",
|
||||
"questions": {"q": {"type": "noul", "instructions": "a?"}}}"#,
|
||||
);
|
||||
match h.systemone(None, req, None).await {
|
||||
Err(SystemOneError::Request(e)) => assert_eq!(e.status, 404, "{}", e.detail),
|
||||
other => panic!("expected a 404, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn routing_to_a_missing_checkpoint_falls_back_and_says_so() {
|
||||
// Devanagari routes to `multilingual`, which this model lacks.
|
||||
let h = tiny_harness().await;
|
||||
let req = request(
|
||||
r#"{"state": "मुझे पैसे वापस चाहिए",
|
||||
"questions": {"q": {"type": "noul", "instructions": "a?"}}}"#,
|
||||
);
|
||||
let resp = h.systemone(None, req, None).await.expect("answered");
|
||||
let routing = resp.routing.expect("routing");
|
||||
assert_eq!(routing.model, "english");
|
||||
assert!(
|
||||
routing
|
||||
.reason
|
||||
.contains("multilingual checkpoint is not loaded"),
|
||||
"{}",
|
||||
routing.reason
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn invalid_questions_are_rejected_before_admission() {
|
||||
let h = tiny_harness().await;
|
||||
let req =
|
||||
request(r#"{"state": "a", "questions": {"q": {"type": "maybe", "instructions": "a?"}}}"#);
|
||||
match h.systemone(None, req, None).await {
|
||||
Err(SystemOneError::Request(e)) => assert_eq!(e.status, 422, "{}", e.detail),
|
||||
other => panic!("expected a 422, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_gateway_routed_model_that_is_not_a_decision_model_is_refused() {
|
||||
let h = tiny_harness().await;
|
||||
let req = request(r#"{"state": "a", "questions": {}}"#);
|
||||
match h.systemone(Some("not/loaded"), req, None).await {
|
||||
Err(SystemOneError::Inference(crate::harness::candle::InferenceError::ModelNotLoaded(
|
||||
id,
|
||||
))) => {
|
||||
assert_eq!(id, "not/loaded")
|
||||
}
|
||||
other => panic!("expected ModelNotLoaded, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Every reference case through `systemone`, compared with the reference
|
||||
/// server's whole response. Routing and usage must match exactly;
|
||||
/// probabilities to within the engine's measured f32 distance from the
|
||||
/// reference (a last-digit wobble after 4-decimal rounding); the chosen
|
||||
/// answer exactly.
|
||||
#[tokio::test]
|
||||
#[ignore = "needs the released weights: set LAYA_REFERENCE_DIR"]
|
||||
async fn released_family_answers_like_the_reference_server() {
|
||||
let root = PathBuf::from(std::env::var("LAYA_REFERENCE_DIR").expect("LAYA_REFERENCE_DIR"));
|
||||
let h = harness();
|
||||
h.register_decision_model(
|
||||
decision_model(
|
||||
"convaiinnovations/laya",
|
||||
&[
|
||||
("english", root.clone()),
|
||||
("multilingual", root.join("multilingual")),
|
||||
("typed-decisions", root.join("typed-decisions")),
|
||||
],
|
||||
)
|
||||
.await,
|
||||
)
|
||||
.await;
|
||||
|
||||
let text = std::fs::read_to_string(testdata().join("reference.json")).unwrap();
|
||||
let reference = Json::parse(text.as_bytes()).unwrap();
|
||||
let Some(Json::Array(cases)) = reference.get("cases") else {
|
||||
panic!("no cases")
|
||||
};
|
||||
let mut max_dp = 0.0f64;
|
||||
for case in cases {
|
||||
let name = case.get("name").and_then(Json::as_str).unwrap();
|
||||
let body = serde_json::to_vec(case.get("request").unwrap()).unwrap();
|
||||
let req = SystemOneRequest::parse(&body, limits::DEFAULT_MAX_TOKEN_BUDGET).unwrap();
|
||||
let resp = h
|
||||
.systemone(Some("convaiinnovations/laya"), req, None)
|
||||
.await
|
||||
.unwrap_or_else(|e| panic!("{name}: {e:?}"));
|
||||
let want_json = case.get("response").unwrap();
|
||||
// Routing and usage compared as serialized text through `Json`:
|
||||
// exact, key order included. serde_json would re-round the
|
||||
// reference's floats on the way in.
|
||||
for key in ["routing", "usage", "model"] {
|
||||
let got = match key {
|
||||
"routing" => serde_json::to_string(&resp.routing).unwrap(),
|
||||
"usage" => serde_json::to_string(&resp.usage).unwrap(),
|
||||
_ => serde_json::to_string(&resp.model).unwrap(),
|
||||
};
|
||||
let want = serde_json::to_string(want_json.get(key).unwrap()).unwrap();
|
||||
assert_eq!(got, want, "{name}: {key}");
|
||||
}
|
||||
let got = serde_json::to_value(&resp).unwrap();
|
||||
let want: serde_json::Value =
|
||||
serde_json::from_str(&serde_json::to_string(want_json).unwrap()).unwrap();
|
||||
let (ga, wa) = (
|
||||
got["answers"].as_object().unwrap(),
|
||||
want["answers"].as_object().unwrap(),
|
||||
);
|
||||
assert_eq!(
|
||||
ga.keys().collect::<Vec<_>>(),
|
||||
wa.keys().collect::<Vec<_>>(),
|
||||
"{name}: answer ids"
|
||||
);
|
||||
for (qid, w) in wa {
|
||||
let g = &ga[qid];
|
||||
for key in ["choice", "type", "legend"] {
|
||||
assert_eq!(g.get(key), w.get(key), "{name}/{qid}: {key}");
|
||||
}
|
||||
if let Some(wp) = w["probabilities"].as_object() {
|
||||
for (opt, p) in wp {
|
||||
let d = (g["probabilities"][opt].as_f64().unwrap() - p.as_f64().unwrap()).abs();
|
||||
max_dp = max_dp.max(d);
|
||||
}
|
||||
}
|
||||
for key in ["noul", "score", "confidence", "answer_confidence"] {
|
||||
if let (Some(gv), Some(wv)) = (
|
||||
g.get(key).and_then(|v| v.as_f64()),
|
||||
w.get(key).and_then(|v| v.as_f64()),
|
||||
) {
|
||||
max_dp = max_dp.max((gv - wv).abs());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
eprintln!("released family: max |Δ| over probabilities and confidences {max_dp:.1e}");
|
||||
assert!(max_dp <= 5e-4, "max |Δ| {max_dp:e}");
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
//! Laya decision models: everything around the forward pass (#335).
|
||||
//!
|
||||
//! A Laya checkpoint answers typed questions over a *state* (text, a JSON
|
||||
//! document, a conversation) in one encoder forward pass: each question
|
||||
//! becomes one sequence,
|
||||
//!
|
||||
//! ```text
|
||||
//! [CLS] "<type> question: <instructions>" [SEP] [MASK] opt0 [MASK] opt1 … [SEP] state [SEP]
|
||||
//! ```
|
||||
//!
|
||||
//! and the decision head scores the hidden state at each option's `[MASK]`
|
||||
//! marker. This module is the part that touches no tensor:
|
||||
//!
|
||||
//! - [`question`] — validate and normalise the question definitions, and
|
||||
//! render each option's text;
|
||||
//! - [`sequence`] — assemble the token sequences under the model's token
|
||||
//! budgets ([`prepare`]);
|
||||
//! - [`shape`] — turn per-option logits into calibrated, typed answers
|
||||
//! ([`shape()`]);
|
||||
//! - [`config`] / [`tokenizer`] — the per-checkpoint pieces those need.
|
||||
//!
|
||||
//! The reference is the `laya` Python SDK (`laya/common.py`,
|
||||
//! `laya/agent.py`, `laya/serve.py`), whose behaviour this reproduces token
|
||||
//! for token: the parity tests replay fixtures recorded from it
|
||||
//! (`script/laya-reference.py`). The calibration arithmetic follows numpy's
|
||||
//! float32 evaluation order, so rounded probabilities match to the last
|
||||
//! published digit.
|
||||
//!
|
||||
//! The wire types (`/v1/systemone`) live in [`cortex_core::decisions`].
|
||||
|
||||
pub mod config;
|
||||
pub mod question;
|
||||
pub mod route;
|
||||
pub mod sequence;
|
||||
pub mod shape;
|
||||
pub mod tokenizer;
|
||||
|
||||
#[cfg(test)]
|
||||
mod handler_tests;
|
||||
|
||||
pub use config::DecisionConfig;
|
||||
pub use cortex_core::decisions::DecisionRequestError;
|
||||
pub use question::{QType, Question};
|
||||
pub use sequence::{Budget, OptionStats, Prepared, Row, prepare};
|
||||
pub use shape::{RowOutput, shape};
|
||||
pub use tokenizer::{DecisionTokenizer, SpecialTokens, TextEncoder};
|
||||
|
||||
/// The constant `model` field of every response: the decision head's name.
|
||||
/// Which checkpoint answered is reported in `routing`.
|
||||
pub const DECISION_MODEL_NAME: &str = "laya-rl-agent";
|
||||
|
||||
/// A request error in the reference's `ValueError` class: the server
|
||||
/// answers it with 422 and the message as `detail`.
|
||||
pub(crate) fn unprocessable(detail: impl Into<String>) -> DecisionRequestError {
|
||||
DecisionRequestError::new(422, detail)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod parity_tests;
|
||||
@@ -0,0 +1,209 @@
|
||||
//! Replays the reference fixtures (`testdata/laya/reference.json`, recorded
|
||||
//! by `script/laya-reference.py`) through [`prepare`] and [`shape`].
|
||||
//!
|
||||
//! Token ids and marker positions must match exactly. The answers and usage
|
||||
//! must serialize to exactly the same JSON — same keys, same order, same
|
||||
//! rounded values — when fed the reference's own raw logits.
|
||||
//!
|
||||
//! The English tokenizer is committed. The multilingual one is 34 MB, so the
|
||||
//! multilingual cases run only when `LAYA_ML_TOKENIZER_DIR` names a copy of
|
||||
//! the checkpoint's `multilingual/tokenizer/` directory.
|
||||
|
||||
use super::*;
|
||||
use cortex_core::decisions::{Json, SystemOneRequest, limits};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
fn testdata() -> PathBuf {
|
||||
Path::new(env!("CARGO_MANIFEST_DIR")).join("src/harness/testdata/laya")
|
||||
}
|
||||
|
||||
fn reference() -> Json {
|
||||
Json::parse(&std::fs::read(testdata().join("reference.json")).expect("fixtures")).expect("json")
|
||||
}
|
||||
|
||||
fn to_value(j: &Json) -> serde_json::Value {
|
||||
serde_json::to_value(j).expect("serializable")
|
||||
}
|
||||
|
||||
fn tokenizer_for(checkpoint: &str) -> Option<DecisionTokenizer> {
|
||||
let dir = match checkpoint {
|
||||
// typed-decisions ships the English tokenizer, byte for byte.
|
||||
"english" | "typed-decisions" => testdata().join("tokenizer"),
|
||||
"multilingual" => PathBuf::from(std::env::var_os("LAYA_ML_TOKENIZER_DIR")?),
|
||||
other => panic!("unknown checkpoint {other}"),
|
||||
};
|
||||
Some(DecisionTokenizer::from_dir(&dir).expect("tokenizer"))
|
||||
}
|
||||
|
||||
fn config_for(reference: &Json, checkpoint: &str) -> DecisionConfig {
|
||||
let cfg = reference
|
||||
.get("header")
|
||||
.and_then(|h| h.get("checkpoints"))
|
||||
.and_then(|c| c.get(checkpoint))
|
||||
.and_then(|c| c.get("cfg"))
|
||||
.expect("checkpoint cfg");
|
||||
DecisionConfig::from_json(&serde_json::to_string(cfg).unwrap()).unwrap()
|
||||
}
|
||||
|
||||
fn array(j: &Json) -> &[Json] {
|
||||
match j {
|
||||
Json::Array(a) => a,
|
||||
other => panic!("expected a list, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn num(j: &Json) -> f64 {
|
||||
match j {
|
||||
Json::Int(d) => d.parse().unwrap(),
|
||||
Json::Float(x) => *x,
|
||||
other => panic!("expected a number, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Replay every case whose tokenizer is available; return how many ran per
|
||||
/// checkpoint.
|
||||
fn replay(filter: impl Fn(&str) -> bool) -> Vec<(String, usize)> {
|
||||
let reference = reference();
|
||||
let mut ran: Vec<(String, usize)> = Vec::new();
|
||||
for case in array(reference.get("cases").unwrap()) {
|
||||
let name = case.get("name").and_then(Json::as_str).unwrap();
|
||||
if !filter(name) {
|
||||
continue;
|
||||
}
|
||||
let response = case.get("response").unwrap();
|
||||
let routing = response.get("routing").unwrap();
|
||||
let checkpoint = routing.get("model").and_then(Json::as_str).unwrap();
|
||||
let Some(tok) = tokenizer_for(checkpoint) else {
|
||||
continue;
|
||||
};
|
||||
let cfg = config_for(&reference, checkpoint);
|
||||
|
||||
let body = serde_json::to_vec(case.get("request").unwrap()).unwrap();
|
||||
let req = SystemOneRequest::parse(&body, limits::DEFAULT_MAX_TOKEN_BUDGET)
|
||||
.unwrap_or_else(|e| panic!("{name}: {e}"));
|
||||
let budget = Budget {
|
||||
max_len: req.max_len,
|
||||
head_max_len: req.head_max_len,
|
||||
};
|
||||
let prepared = prepare(&req.state, &req.questions, &cfg, &tok, budget)
|
||||
.unwrap_or_else(|e| panic!("{name}: {e}"));
|
||||
|
||||
let items = array(case.get("items").unwrap());
|
||||
assert_eq!(prepared.rows.len(), items.len(), "{name}: row count");
|
||||
let mut outputs = Vec::new();
|
||||
for (i, (row, item)) in prepared.rows.iter().zip(items).enumerate() {
|
||||
let ids: Vec<u32> = array(item.get("input_ids").unwrap())
|
||||
.iter()
|
||||
.map(|v| num(v) as u32)
|
||||
.collect();
|
||||
let markers: Vec<usize> = array(item.get("markers").unwrap())
|
||||
.iter()
|
||||
.map(|v| num(v) as usize)
|
||||
.collect();
|
||||
assert_eq!(row.input_ids, ids, "{name}: row {i} input_ids");
|
||||
assert_eq!(row.markers, markers, "{name}: row {i} markers");
|
||||
assert_eq!(row.qtype as usize, num(item.get("qtype").unwrap()) as usize);
|
||||
let opts = item.get("options").unwrap();
|
||||
let s = prepared.stats[i];
|
||||
assert_eq!(
|
||||
(s.options, s.options_distinct, s.tokens_per_option),
|
||||
(
|
||||
num(opts.get("options").unwrap()) as usize,
|
||||
num(opts.get("options_distinct").unwrap()) as usize,
|
||||
match opts.get("tokens_per_option").unwrap() {
|
||||
Json::Null => None,
|
||||
v => Some(num(v) as usize),
|
||||
}
|
||||
),
|
||||
"{name}: row {i} option stats"
|
||||
);
|
||||
outputs.push(RowOutput {
|
||||
logits: array(item.get("logits").unwrap())
|
||||
.iter()
|
||||
.map(|v| num(v) as f32)
|
||||
.collect(),
|
||||
act_probability: num(&array(item.get("act").unwrap())[0]) as f32,
|
||||
});
|
||||
}
|
||||
|
||||
let lang = routing
|
||||
.get("detection")
|
||||
.and_then(|d| d.get("language"))
|
||||
.and_then(Json::as_str);
|
||||
let (answers, usage) = shape(&prepared, &outputs, &cfg, lang);
|
||||
// Serialized strings, not values: key order is part of the contract.
|
||||
assert_eq!(
|
||||
serde_json::to_string(&answers).unwrap(),
|
||||
serde_json::to_string(response.get("answers").unwrap()).unwrap(),
|
||||
"{name}: answers"
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::to_value(&usage).unwrap(),
|
||||
to_value(response.get("usage").unwrap()),
|
||||
"{name}: usage"
|
||||
);
|
||||
match ran.iter_mut().find(|(c, _)| c == checkpoint) {
|
||||
Some((_, n)) => *n += 1,
|
||||
None => ran.push((checkpoint.to_string(), 1)),
|
||||
}
|
||||
}
|
||||
ran
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_english_fixture_matches_the_reference() {
|
||||
let ran = replay(|_| true);
|
||||
let count = |c: &str| ran.iter().find(|(n, _)| n == c).map_or(0, |(_, n)| *n);
|
||||
// The corpus routes 19 cases to English and 1 to typed-decisions; a
|
||||
// drop means a case silently stopped running.
|
||||
assert_eq!(count("english"), 19, "{ran:?}");
|
||||
assert_eq!(count("typed-decisions"), 1, "{ran:?}");
|
||||
if std::env::var_os("LAYA_ML_TOKENIZER_DIR").is_some() {
|
||||
assert_eq!(count("multilingual"), 8, "{ran:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_questions_need_no_forward() {
|
||||
let reference = reference();
|
||||
let tok = tokenizer_for("english").unwrap();
|
||||
let cfg = config_for(&reference, "english");
|
||||
let prepared = prepare(
|
||||
&Json::String("x".into()),
|
||||
&[],
|
||||
&cfg,
|
||||
&tok,
|
||||
Budget::default(),
|
||||
)
|
||||
.unwrap();
|
||||
assert!(prepared.rows.is_empty());
|
||||
let (answers, usage) = shape(&prepared, &[], &cfg, None);
|
||||
assert_eq!(serde_json::to_string(&answers).unwrap(), "{}");
|
||||
assert_eq!(
|
||||
serde_json::to_string(&usage).unwrap(),
|
||||
r#"{"input_tokens":0,"output_tokens":0}"#
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn options_that_cannot_all_fit_are_refused() {
|
||||
let reference = reference();
|
||||
let tok = tokenizer_for("english").unwrap();
|
||||
let cfg = config_for(&reference, "english");
|
||||
let questions = Json::parse(
|
||||
br#"{"q": {"type": "choice", "instructions": "i", "criteria": ["a", "b", "c", "d"]}}"#,
|
||||
)
|
||||
.unwrap();
|
||||
let Json::Object(questions) = questions else {
|
||||
unreachable!()
|
||||
};
|
||||
let budget = Budget {
|
||||
max_len: Some(10),
|
||||
head_max_len: None,
|
||||
};
|
||||
let err = prepare(&Json::String("x".into()), &questions, &cfg, &tok, budget).unwrap_err();
|
||||
assert_eq!(
|
||||
(err.status, err.detail.as_str()),
|
||||
(422, "question 'q' options exceed head_max_len=192")
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,458 @@
|
||||
//! Question validation, normalisation and option rendering.
|
||||
//!
|
||||
//! Every rejection here is a 422 whose message names the question and what
|
||||
//! to fix — the same text the reference returns, so a client written
|
||||
//! against it sees the same errors.
|
||||
|
||||
use super::{DecisionRequestError, unprocessable};
|
||||
use cortex_core::decisions::{HashKey, Json};
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// The three decision primitives. The discriminant is the index the model's
|
||||
/// type embedding and the calibration tables use.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum QType {
|
||||
/// Pick one of the caller's labels.
|
||||
Choice = 0,
|
||||
/// An ordinal level on the caller's rubric.
|
||||
Score = 1,
|
||||
/// Yes/no: the probability that a statement holds.
|
||||
Noul = 2,
|
||||
}
|
||||
|
||||
impl QType {
|
||||
pub const NAMES: [&'static str; 3] = ["choice", "score", "noul"];
|
||||
|
||||
pub fn name(self) -> &'static str {
|
||||
Self::NAMES[self as usize]
|
||||
}
|
||||
|
||||
fn parse(s: &str) -> Option<Self> {
|
||||
match s {
|
||||
"choice" => Some(Self::Choice),
|
||||
"score" => Some(Self::Score),
|
||||
"noul" => Some(Self::Noul),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A validated question, normalised the way the reference normalises it.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct Question {
|
||||
pub id: String,
|
||||
pub kind: QType,
|
||||
/// The instruction text: the caller's string, or the JSON rendering of
|
||||
/// a structured value.
|
||||
pub instructions: String,
|
||||
/// Each option's rendered text, in answer-index order. For `noul` this
|
||||
/// is always `[false, true]`, so index 1 is P(true).
|
||||
pub options: Vec<String>,
|
||||
/// `choice`: each option's label as the caller wrote it (the answer key).
|
||||
/// `score`: each level's description (the legend). `noul`: empty.
|
||||
pub labels: Vec<Json>,
|
||||
}
|
||||
|
||||
impl Question {
|
||||
/// Validate one question definition and normalise it.
|
||||
pub fn from_definition(id: &str, def: &Json) -> Result<Self, DecisionRequestError> {
|
||||
check_question(id, def)?;
|
||||
Ok(to_internal(id, def))
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate and normalise every question, in request order. The first
|
||||
/// invalid question is reported; nothing is normalised until all pass.
|
||||
pub fn questions_from_request(
|
||||
questions: &[(String, Json)],
|
||||
) -> Result<Vec<Question>, DecisionRequestError> {
|
||||
for (id, def) in questions {
|
||||
check_question(id, def)?;
|
||||
}
|
||||
Ok(questions
|
||||
.iter()
|
||||
.map(|(id, def)| to_internal(id, def))
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn q(id: &str) -> String {
|
||||
Json::String(id.to_string()).py_repr()
|
||||
}
|
||||
|
||||
fn check_question(id: &str, def: &Json) -> Result<(), DecisionRequestError> {
|
||||
let qid = q(id);
|
||||
if !matches!(def, Json::Object(_)) {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: definition must be a dict, got {}",
|
||||
def.py_type_name()
|
||||
)));
|
||||
}
|
||||
let t = def.get("type").unwrap_or(&Json::Null);
|
||||
let Some(kind) = t.as_str().and_then(QType::parse) else {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: unknown type {}; use one of ['choice', 'noul', 'score']",
|
||||
t.py_repr()
|
||||
)));
|
||||
};
|
||||
if def.get("instructions").is_none() {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: no 'instructions'; add the text the model should answer"
|
||||
)));
|
||||
}
|
||||
let crit = def.get("criteria");
|
||||
match kind {
|
||||
QType::Choice => check_choice(&qid, crit)?,
|
||||
QType::Score => {
|
||||
let Some(Json::Array(levels)) = crit else {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: a score question takes 'criteria' as a list of level \
|
||||
descriptions, index 0 first"
|
||||
)));
|
||||
};
|
||||
if levels.is_empty() {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: a score question needs at least one level"
|
||||
)));
|
||||
}
|
||||
if let Some(i) = levels.iter().position(|l| *l == Json::Null) {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: score level {i} is null; give every level a description, \
|
||||
index 0 first"
|
||||
)));
|
||||
}
|
||||
}
|
||||
QType::Noul => match crit {
|
||||
None | Some(Json::Null) => {}
|
||||
Some(Json::Object(m)) => {
|
||||
let mut keys: Vec<String> = m.iter().map(|(k, _)| k.to_lowercase()).collect();
|
||||
keys.sort();
|
||||
keys.dedup();
|
||||
if keys.iter().any(|k| k != "true" && k != "false") {
|
||||
let listed = Json::Array(keys.into_iter().map(Json::String).collect());
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: a noul question takes 'criteria' keyed only \
|
||||
'true'/'false' (either or both, and omitted is fine), got {}. Those \
|
||||
keys are the option texts the model reads; any other key was silently \
|
||||
dropped and replaced with the defaults. If you want the answer worded \
|
||||
differently, keep 'criteria' keyed 'true'/'false' and set 'labels' \
|
||||
instead.",
|
||||
listed.py_repr()
|
||||
)));
|
||||
}
|
||||
}
|
||||
Some(_) => {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: a noul question takes 'criteria' as a dict with optional \
|
||||
'true'/'false' descriptions, or omits it"
|
||||
)));
|
||||
}
|
||||
},
|
||||
}
|
||||
if let Some(labels) = def.get("labels") {
|
||||
if kind != QType::Noul {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: 'labels' is only supported for noul questions"
|
||||
)));
|
||||
}
|
||||
resolve_noul_labels(labels).map_err(|e| unprocessable(format!("question {qid}: {e}")))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn check_choice(qid: &str, crit: Option<&Json>) -> Result<(), DecisionRequestError> {
|
||||
let labels: Vec<Json> = match crit {
|
||||
Some(Json::Object(m)) => m.iter().map(|(k, _)| Json::String(k.clone())).collect(),
|
||||
Some(Json::Array(a)) => a.clone(),
|
||||
_ => {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: a choice question takes 'criteria' as a dict of label -> \
|
||||
description, or a list of labels"
|
||||
)));
|
||||
}
|
||||
};
|
||||
if labels.is_empty() {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: a choice question needs at least one criterion"
|
||||
)));
|
||||
}
|
||||
for (i, label) in labels.iter().enumerate() {
|
||||
match label {
|
||||
Json::Array(_) | Json::Object(_) => {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: choice label {i} is a {}; a label is rendered as option \
|
||||
text and used as the answer key, so it must be a scalar (a string, number \
|
||||
or bool), got {}",
|
||||
label.py_type_name(),
|
||||
label.py_repr()
|
||||
)));
|
||||
}
|
||||
Json::Null => {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: choice label {i} is null; a label is rendered as option \
|
||||
text and used as the answer key, so it must be a string, number or bool \
|
||||
-- a null label renders as the text \"None\" while its answer key is \
|
||||
\"null\""
|
||||
)));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
// Only the list form can repeat a key: an object's keys are unique.
|
||||
if matches!(crit, Some(Json::Array(_))) {
|
||||
let mut seen: HashMap<HashKey, usize> = HashMap::new();
|
||||
for (i, label) in labels.iter().enumerate() {
|
||||
let key = label.hash_key().expect("scalars are hashable");
|
||||
if key == HashKey::Unique {
|
||||
continue;
|
||||
}
|
||||
if let Some(first) = seen.get(&key) {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: choice label {i} ({}) repeats label {first}; the labels are \
|
||||
the answer keys, so every option needs its own (1, 1.0 and True are one key)",
|
||||
label.py_repr()
|
||||
)));
|
||||
}
|
||||
seen.insert(key, i);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
const DEFAULT_NOUL_FALSE: &str = "no, the statement does not hold";
|
||||
const DEFAULT_NOUL_TRUE: &str = "yes, the statement holds";
|
||||
|
||||
/// The words a noul question shows the model for its two outcomes.
|
||||
fn resolve_noul_labels(labels: &Json) -> Result<(String, String), &'static str> {
|
||||
const BAD: &str =
|
||||
"noul labels must map exactly 'false' and 'true' to distinct non-empty strings";
|
||||
let Json::Object(m) = labels else {
|
||||
return if *labels == Json::Null {
|
||||
Ok(("false".into(), "true".into()))
|
||||
} else {
|
||||
Err(BAD)
|
||||
};
|
||||
};
|
||||
let mut keys: Vec<&str> = m.iter().map(|(k, _)| k.as_str()).collect();
|
||||
keys.sort_unstable();
|
||||
if keys != ["false", "true"] {
|
||||
return Err(BAD);
|
||||
}
|
||||
let (Some(Json::String(f)), Some(Json::String(t))) = (labels.get("false"), labels.get("true"))
|
||||
else {
|
||||
return Err(BAD);
|
||||
};
|
||||
let (f, t) = (py_strip(f), py_strip(t));
|
||||
if f.is_empty() || t.is_empty() || f == t {
|
||||
return Err(BAD);
|
||||
}
|
||||
Ok((f.to_string(), t.to_string()))
|
||||
}
|
||||
|
||||
/// Python's `str.strip()` with no argument: Unicode whitespace.
|
||||
fn py_strip(s: &str) -> &str {
|
||||
s.trim_matches(|c: char| c.is_whitespace() || ('\u{1c}'..='\u{1f}').contains(&c))
|
||||
}
|
||||
|
||||
/// `render_criterion`: a string as itself, anything else as JSON.
|
||||
fn render_criterion(v: &Json) -> String {
|
||||
match v {
|
||||
Json::String(s) => s.clone(),
|
||||
other => other.dumps(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Only `None` and `""` mean "no description"; `0` and `false` are values.
|
||||
fn is_blank(v: &Json) -> bool {
|
||||
matches!(v, Json::Null) || matches!(v, Json::String(s) if s.is_empty())
|
||||
}
|
||||
|
||||
fn to_internal(id: &str, def: &Json) -> Question {
|
||||
let kind = def
|
||||
.get("type")
|
||||
.and_then(Json::as_str)
|
||||
.and_then(QType::parse)
|
||||
.expect("validated");
|
||||
let instructions = match def.get("instructions").expect("validated") {
|
||||
Json::String(s) => s.clone(),
|
||||
other => other.dumps(),
|
||||
};
|
||||
let crit = def.get("criteria");
|
||||
let (options, labels) = match kind {
|
||||
QType::Choice => {
|
||||
let pairs: Vec<(Json, Json)> = match crit {
|
||||
Some(Json::Object(m)) => m
|
||||
.iter()
|
||||
.map(|(k, v)| (Json::String(k.clone()), v.clone()))
|
||||
.collect(),
|
||||
Some(Json::Array(a)) => a.iter().map(|l| (l.clone(), Json::Null)).collect(),
|
||||
_ => unreachable!("validated"),
|
||||
};
|
||||
let options = pairs
|
||||
.iter()
|
||||
.map(|(k, v)| {
|
||||
if is_blank(v) {
|
||||
k.py_str()
|
||||
} else {
|
||||
format!("{}: {}", k.py_str(), render_criterion(v))
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
(options, pairs.into_iter().map(|(k, _)| k).collect())
|
||||
}
|
||||
QType::Score => {
|
||||
let Some(Json::Array(levels)) = crit else {
|
||||
unreachable!("validated")
|
||||
};
|
||||
let options = levels
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, c)| format!("level {i}: {}", render_criterion(c)))
|
||||
.collect();
|
||||
(options, levels.clone())
|
||||
}
|
||||
QType::Noul => {
|
||||
// Keys are matched case-insensitively; a repeat after
|
||||
// lower-casing keeps the later value, as a dict would.
|
||||
let mut desc: HashMap<String, Json> = HashMap::new();
|
||||
if let Some(Json::Object(m)) = crit {
|
||||
for (k, v) in m {
|
||||
desc.insert(k.to_lowercase(), v.clone());
|
||||
}
|
||||
}
|
||||
let (fl, tl) =
|
||||
resolve_noul_labels(def.get("labels").unwrap_or(&Json::Null)).expect("validated");
|
||||
let render = |key: &str, default: &str| match desc.get(key) {
|
||||
Some(v) if !is_blank(v) => render_criterion(v),
|
||||
_ => default.to_string(),
|
||||
};
|
||||
(
|
||||
vec![
|
||||
format!("{fl}: {}", render("false", DEFAULT_NOUL_FALSE)),
|
||||
format!("{tl}: {}", render("true", DEFAULT_NOUL_TRUE)),
|
||||
],
|
||||
Vec::new(),
|
||||
)
|
||||
}
|
||||
};
|
||||
Question {
|
||||
id: id.to_string(),
|
||||
kind,
|
||||
instructions,
|
||||
options,
|
||||
labels,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn def(s: &str) -> Json {
|
||||
Json::parse(s.as_bytes()).unwrap()
|
||||
}
|
||||
|
||||
fn err(s: &str) -> String {
|
||||
Question::from_definition("q", &def(s)).unwrap_err().detail
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn options_render_as_the_reference_renders_them() {
|
||||
let c = Question::from_definition(
|
||||
"q",
|
||||
&def(r#"{"type": "choice", "instructions": {"a": 1},
|
||||
"criteria": {"x": "", "y": null, "z": 0, "w": {"d": [1, 2.0]}, "v": "text"}}"#),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(c.instructions, r#"{"a": 1}"#);
|
||||
assert_eq!(
|
||||
c.options,
|
||||
["x", "y", "z: 0", r#"w: {"d": [1, 2.0]}"#, "v: text"]
|
||||
);
|
||||
|
||||
let l = Question::from_definition(
|
||||
"q",
|
||||
&def(r#"{"type": "choice", "instructions": "i", "criteria": [1, 2.5, false, "s"]}"#),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(l.options, ["1", "2.5", "False", "s"]);
|
||||
|
||||
let s = Question::from_definition(
|
||||
"q",
|
||||
&def(r#"{"type": "score", "instructions": "i", "criteria": ["calm", {"k": false}]}"#),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(s.options, ["level 0: calm", r#"level 1: {"k": false}"#]);
|
||||
|
||||
let n = Question::from_definition(
|
||||
"q",
|
||||
&def(
|
||||
r#"{"type": "noul", "instructions": "i", "criteria": {"TRUE": "held", "false": ""},
|
||||
"labels": {"false": " no ", "true": "yes"}}"#,
|
||||
),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
n.options,
|
||||
["no: no, the statement does not hold", "yes: held"]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_questions_name_the_question_and_the_fix() {
|
||||
assert_eq!(
|
||||
err("[]"),
|
||||
"question 'q': definition must be a dict, got list"
|
||||
);
|
||||
assert_eq!(
|
||||
err(r#"{"type": "pick"}"#),
|
||||
"question 'q': unknown type 'pick'; use one of ['choice', 'noul', 'score']"
|
||||
);
|
||||
assert_eq!(
|
||||
err(r#"{"type": "noul"}"#),
|
||||
"question 'q': no 'instructions'; add the text the model should answer"
|
||||
);
|
||||
assert!(
|
||||
err(r#"{"type": "choice", "instructions": "i", "criteria": []}"#)
|
||||
.ends_with("needs at least one criterion")
|
||||
);
|
||||
assert!(
|
||||
err(r#"{"type": "choice", "instructions": "i", "criteria": [[1]]}"#)
|
||||
.starts_with("question 'q': choice label 0 is a list;")
|
||||
);
|
||||
assert!(
|
||||
err(r#"{"type": "choice", "instructions": "i", "criteria": ["a", null]}"#)
|
||||
.starts_with("question 'q': choice label 1 is null;")
|
||||
);
|
||||
assert!(
|
||||
err(r#"{"type": "choice", "instructions": "i", "criteria": [1, "x", true]}"#)
|
||||
.starts_with("question 'q': choice label 2 (True) repeats label 0;")
|
||||
);
|
||||
assert_eq!(
|
||||
err(r#"{"type": "score", "instructions": "i", "criteria": ["a", null]}"#),
|
||||
"question 'q': score level 1 is null; give every level a description, index 0 first"
|
||||
);
|
||||
assert!(
|
||||
err(r#"{"type": "noul", "instructions": "i", "criteria": {"yes": "x"}}"#)
|
||||
.contains("got ['yes'].")
|
||||
);
|
||||
assert_eq!(
|
||||
err(r#"{"type": "score", "instructions": "i", "criteria": ["a"], "labels": {}}"#),
|
||||
"question 'q': 'labels' is only supported for noul questions"
|
||||
);
|
||||
assert_eq!(
|
||||
err(r#"{"type": "noul", "instructions": "i", "labels": {"false": "x", "true": " x"}}"#),
|
||||
"question 'q': noul labels must map exactly 'false' and 'true' to distinct non-empty strings"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn nan_labels_never_collide() {
|
||||
assert!(
|
||||
Question::from_definition(
|
||||
"q",
|
||||
&def(r#"{"type": "choice", "instructions": "i", "criteria": [NaN, NaN]}"#),
|
||||
)
|
||||
.is_ok()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,301 @@
|
||||
//! Choosing which checkpoint of a Laya family answers a request (#339).
|
||||
//!
|
||||
//! A released Laya repo carries three checkpoints: the English one at
|
||||
//! its root, `multilingual/` and `typed-decisions/`. The English
|
||||
//! checkpoint cannot read other scripts — upstream measured Khmer at
|
||||
//! 0.000 accuracy with 0.952 confidence — so an unpinned request is
|
||||
//! routed by the language of its state. This follows the reference's
|
||||
//! `Router._route` (`laya/router.py`) and its HTTP server's
|
||||
//! `_resolve_model` (`laya/serve.py`), including the `reason` strings.
|
||||
//!
|
||||
//! Precedence, as upstream: a checkpoint the caller names in `model`,
|
||||
//! then the detected script and language, then the default (English).
|
||||
//! Upstream's opt-in typed-decisions workflow routing and caller
|
||||
//! language hints are not exposed over HTTP by the reference, so they
|
||||
//! are not here either; the matched workflow is still reported.
|
||||
|
||||
use cortex_core::decisions::{DecisionRouting, Json};
|
||||
|
||||
use crate::harness::lang_detect;
|
||||
|
||||
/// The root checkpoint's routing name.
|
||||
pub const ENGLISH: &str = "english";
|
||||
pub const MULTILINGUAL: &str = "multilingual";
|
||||
pub const TYPED_DECISIONS: &str = "typed-decisions";
|
||||
|
||||
/// Where an unpinned request with no usable language signal goes.
|
||||
pub const DEFAULT_CHECKPOINT: &str = ENGLISH;
|
||||
|
||||
/// Every checkpoint a released family carries, root first.
|
||||
pub const FAMILY: [&str; 3] = [ENGLISH, MULTILINGUAL, TYPED_DECISIONS];
|
||||
|
||||
/// The reference's `_ALIASES`: names people are likely to type.
|
||||
const ALIASES: [(&str, &str); 10] = [
|
||||
("en", ENGLISH),
|
||||
("laya", ENGLISH),
|
||||
("default", ENGLISH),
|
||||
("multi", MULTILINGUAL),
|
||||
("ml", MULTILINGUAL),
|
||||
("laya-multilingual", MULTILINGUAL),
|
||||
("typed", TYPED_DECISIONS),
|
||||
("typed_decisions", TYPED_DECISIONS),
|
||||
("laya-typed-decisions", TYPED_DECISIONS),
|
||||
("decisions", TYPED_DECISIONS),
|
||||
];
|
||||
|
||||
/// The reference server's `_PUBLISHED_MODEL_IDS`: the Hugging Face ids
|
||||
/// the non-root checkpoints are published under.
|
||||
const PUBLISHED: [(&str, &str); 2] = [
|
||||
("convaiinnovations/laya-multilingual", MULTILINGUAL),
|
||||
("convaiinnovations/laya-typed-decisions", TYPED_DECISIONS),
|
||||
];
|
||||
|
||||
/// Question-id signatures of the four typed-decisions workflows.
|
||||
const WORKFLOWS: [(&str, [&str; 5]); 4] = [
|
||||
(
|
||||
"agent_trace_observability",
|
||||
["action", "needs_review", "outcome", "risk", "urgency"],
|
||||
),
|
||||
(
|
||||
"customer_service",
|
||||
["action", "category", "churn_risk", "needs_human", "urgency"],
|
||||
),
|
||||
(
|
||||
"invoice_processing",
|
||||
[
|
||||
"discrepancy_severity",
|
||||
"disposition",
|
||||
"duplicate",
|
||||
"matches_order",
|
||||
"urgency",
|
||||
],
|
||||
),
|
||||
(
|
||||
"security_incidents",
|
||||
[
|
||||
"credential_compromise",
|
||||
"disposition",
|
||||
"severity",
|
||||
"true_positive",
|
||||
"urgency",
|
||||
],
|
||||
),
|
||||
];
|
||||
|
||||
/// The checkpoint a request's `model` field pins, if it names one.
|
||||
///
|
||||
/// Anything else — a Jev service name like `jev-1`, the family's own
|
||||
/// repo id, an empty string — pins nothing and the request is routed.
|
||||
pub fn resolve_pin(model: Option<&str>) -> Option<&'static str> {
|
||||
let model = model?;
|
||||
if model.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let key = py_strip(model).to_lowercase();
|
||||
if let Some((_, ckpt)) = PUBLISHED.iter().find(|(id, _)| *id == key) {
|
||||
return Some(ckpt);
|
||||
}
|
||||
let key = ALIASES
|
||||
.iter()
|
||||
.find(|(alias, _)| *alias == key)
|
||||
.map(|(_, ckpt)| *ckpt)
|
||||
.unwrap_or(key.as_str());
|
||||
FAMILY.iter().copied().find(|c| *c == key)
|
||||
}
|
||||
|
||||
/// The typed-decisions workflow whose question ids these are exactly.
|
||||
pub fn match_workflow(questions: &[(String, Json)]) -> Option<&'static str> {
|
||||
let ids: std::collections::BTreeSet<&str> = questions.iter().map(|(q, _)| q.as_str()).collect();
|
||||
WORKFLOWS.iter().find_map(|(name, sig)| {
|
||||
let sig: std::collections::BTreeSet<&str> = sig.iter().copied().collect();
|
||||
(ids == sig).then_some(*name)
|
||||
})
|
||||
}
|
||||
|
||||
/// Decide which checkpoint answers, and say why.
|
||||
///
|
||||
/// `model_id` is the served family's repo id, used for the `repo` field.
|
||||
/// The result names a checkpoint whether or not this node has it
|
||||
/// loaded; the caller decides what to do about one it lacks.
|
||||
pub fn route(
|
||||
model_id: &str,
|
||||
state: &Json,
|
||||
questions: &[(String, Json)],
|
||||
pin: Option<&'static str>,
|
||||
) -> DecisionRouting {
|
||||
if let Some(key) = pin {
|
||||
return DecisionRouting {
|
||||
model: key.to_string(),
|
||||
repo: repo_of(model_id, key),
|
||||
reason: format!("explicit model={}", py_repr(key)),
|
||||
detection: None,
|
||||
workflow: None,
|
||||
};
|
||||
}
|
||||
let workflow = match_workflow(questions).map(str::to_string);
|
||||
let det = lang_detect::analyse(state);
|
||||
let (key, reason) = if det.script == "unknown" {
|
||||
(
|
||||
DEFAULT_CHECKPOINT,
|
||||
format!("no letters detected in state; using default ({DEFAULT_CHECKPOINT})"),
|
||||
)
|
||||
} else if det.script != "latin" {
|
||||
(
|
||||
MULTILINGUAL,
|
||||
format!(
|
||||
"non-Latin script ({}, {:.0}% of letters); the English checkpoint cannot read it",
|
||||
det.script,
|
||||
100.0 * det.non_latin_fraction
|
||||
),
|
||||
)
|
||||
} else if !det.is_english {
|
||||
let reason = if let Some(segment) = &det.mixed_segment {
|
||||
format!(
|
||||
"Latin script, mostly English, but a line or field reads as {} ({}); \
|
||||
the English checkpoint cannot read it",
|
||||
det.language.map(py_repr).unwrap_or_else(|| "None".into()),
|
||||
py_repr(&segment.chars().take(60).collect::<String>()),
|
||||
)
|
||||
} else if let Some(lang) = det.language {
|
||||
format!(
|
||||
"Latin script but language looks like {}, not English",
|
||||
py_repr(lang)
|
||||
)
|
||||
} else {
|
||||
format!(
|
||||
"Latin script, language not identified but {:.0}% non-English letters; \
|
||||
not safe for the English checkpoint",
|
||||
100.0 * det.diacritic_rate
|
||||
)
|
||||
};
|
||||
(MULTILINGUAL, reason)
|
||||
} else if det.language_undecided {
|
||||
(
|
||||
DEFAULT_CHECKPOINT,
|
||||
format!(
|
||||
"Latin script, language not identified and no non-English letters; \
|
||||
using default ({DEFAULT_CHECKPOINT})"
|
||||
),
|
||||
)
|
||||
} else {
|
||||
(ENGLISH, "English Latin text".to_string())
|
||||
};
|
||||
DecisionRouting {
|
||||
model: key.to_string(),
|
||||
repo: repo_of(model_id, key),
|
||||
reason,
|
||||
detection: Some(analysis_json(&det)),
|
||||
workflow,
|
||||
}
|
||||
}
|
||||
|
||||
/// The language per-language calibration keys on: the detected one.
|
||||
pub fn detected_language(routing: &DecisionRouting) -> Option<&str> {
|
||||
routing
|
||||
.detection
|
||||
.as_ref()
|
||||
.and_then(|d| d.get("language"))
|
||||
.and_then(Json::as_str)
|
||||
}
|
||||
|
||||
/// `<repo>` for the root checkpoint, `<repo>/<subfolder>` otherwise.
|
||||
pub fn repo_of(model_id: &str, checkpoint: &str) -> String {
|
||||
if checkpoint == ENGLISH {
|
||||
model_id.to_string()
|
||||
} else {
|
||||
format!("{model_id}/{checkpoint}")
|
||||
}
|
||||
}
|
||||
|
||||
/// The analysis as an order-preserving [`Json`], so the response keeps
|
||||
/// the reference's field order.
|
||||
fn analysis_json(det: &lang_detect::Analysis) -> Json {
|
||||
let text = serde_json::to_string(det).unwrap_or_else(|_| "null".into());
|
||||
Json::parse(text.as_bytes()).unwrap_or(Json::Null)
|
||||
}
|
||||
|
||||
fn py_repr(s: &str) -> String {
|
||||
Json::String(s.to_string()).py_repr()
|
||||
}
|
||||
|
||||
/// Python's `str.strip()`: Unicode whitespace, which includes the
|
||||
/// U+001C–U+001F separators Rust's `trim` leaves alone.
|
||||
fn py_strip(s: &str) -> &str {
|
||||
s.trim_matches(|c: char| c.is_whitespace() || ('\u{1c}'..='\u{1f}').contains(&c))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn pins_follow_the_reference_server() {
|
||||
assert_eq!(resolve_pin(Some("multilingual")), Some(MULTILINGUAL));
|
||||
assert_eq!(resolve_pin(Some(" ML ")), Some(MULTILINGUAL));
|
||||
assert_eq!(
|
||||
resolve_pin(Some("convaiinnovations/laya-typed-decisions")),
|
||||
Some(TYPED_DECISIONS)
|
||||
);
|
||||
assert_eq!(resolve_pin(Some("laya")), Some(ENGLISH));
|
||||
// A Jev service name, the family's own id, and nothing: no pin.
|
||||
assert_eq!(resolve_pin(Some("jev-1")), None);
|
||||
assert_eq!(resolve_pin(Some("convaiinnovations/laya")), None);
|
||||
assert_eq!(resolve_pin(Some("")), None);
|
||||
assert_eq!(resolve_pin(None), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workflow_needs_the_exact_id_set() {
|
||||
let qs = |ids: &[&str]| -> Vec<(String, Json)> {
|
||||
ids.iter().map(|i| (i.to_string(), Json::Null)).collect()
|
||||
};
|
||||
assert_eq!(
|
||||
match_workflow(&qs(&[
|
||||
"urgency",
|
||||
"action",
|
||||
"category",
|
||||
"churn_risk",
|
||||
"needs_human"
|
||||
])),
|
||||
Some("customer_service")
|
||||
);
|
||||
assert_eq!(match_workflow(&qs(&["urgency"])), None);
|
||||
}
|
||||
|
||||
/// Every case in the reference fixtures routes to the checkpoint, with
|
||||
/// the reason and detection, the reference reported.
|
||||
#[test]
|
||||
fn routes_every_reference_case_like_the_reference() {
|
||||
let text = std::fs::read_to_string(
|
||||
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
|
||||
.join("src/harness/testdata/laya/reference.json"),
|
||||
)
|
||||
.expect("read reference.json");
|
||||
let reference = Json::parse(text.as_bytes()).expect("parse reference.json");
|
||||
let Some(Json::Array(cases)) = reference.get("cases") else {
|
||||
panic!("no cases");
|
||||
};
|
||||
for case in cases {
|
||||
let name = case.get("name").and_then(Json::as_str).unwrap();
|
||||
let request = case.get("request").unwrap();
|
||||
let body = serde_json::to_vec(request).unwrap();
|
||||
let req = cortex_core::decisions::SystemOneRequest::parse(
|
||||
&body,
|
||||
cortex_core::decisions::limits::DEFAULT_MAX_TOKEN_BUDGET,
|
||||
)
|
||||
.unwrap();
|
||||
let routing = route(
|
||||
"convaiinnovations/laya",
|
||||
&req.state,
|
||||
&req.questions,
|
||||
resolve_pin(req.model.as_deref()),
|
||||
);
|
||||
let want = case.get("response").unwrap().get("routing").unwrap();
|
||||
assert_eq!(
|
||||
serde_json::to_string(&routing).unwrap(),
|
||||
serde_json::to_string(want).unwrap(),
|
||||
"{name}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,316 @@
|
||||
//! Sequence assembly: one token sequence per question, under the model's
|
||||
//! token budgets.
|
||||
//!
|
||||
//! ```text
|
||||
//! [CLS] "<type> question: <instructions>" [SEP] [MASK] opt0 [MASK] opt1 … [SEP] state [SEP]
|
||||
//! ```
|
||||
//!
|
||||
//! The budget rules, in the order they apply:
|
||||
//!
|
||||
//! 1. Each option is `[MASK]` plus at most 48 tokens of its text.
|
||||
//! 2. `head_max_len` bounds the instruction text and all options together.
|
||||
//! If fewer than 16 tokens would remain for the instructions, every
|
||||
//! option is cut to the same `max(4, (head_max_len - 16) / options)`
|
||||
//! tokens (marker included).
|
||||
//! 3. The instructions get what is left of `head_max_len`, but never fewer
|
||||
//! than 8 tokens.
|
||||
//! 4. The state fills the rest of `max_len`. A conversation (a JSON list)
|
||||
//! keeps its newest turns, so it is cut from the left; anything else is
|
||||
//! cut from the right.
|
||||
//! 5. The whole sequence is cut to `max_len`, dropping any marker that no
|
||||
//! longer fits — a question whose markers do not all fit is refused.
|
||||
|
||||
use super::question::{QType, Question, questions_from_request};
|
||||
use super::tokenizer::TextEncoder;
|
||||
use super::{DecisionConfig, DecisionRequestError, unprocessable};
|
||||
use cortex_core::decisions::Json;
|
||||
|
||||
/// Option text tokens kept before the head budget is considered.
|
||||
pub const MAX_OPTION_TOKENS: usize = 48;
|
||||
/// Head tokens the instructions are guaranteed before options are shrunk.
|
||||
const MIN_INSTRUCTION_ROOM: usize = 16;
|
||||
/// Per-option floor (marker included) when options are shrunk.
|
||||
const MIN_TOKENS_PER_OPTION: usize = 4;
|
||||
/// Instruction tokens kept however many the options take.
|
||||
const MIN_INSTRUCTION_TOKENS: usize = 8;
|
||||
|
||||
/// Per-request overrides of the checkpoint's token budgets.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct Budget {
|
||||
pub max_len: Option<usize>,
|
||||
pub head_max_len: Option<usize>,
|
||||
}
|
||||
|
||||
/// One model input row: one question over the state.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct Row {
|
||||
pub input_ids: Vec<u32>,
|
||||
/// Position of each option's `[MASK]` marker, in option order.
|
||||
pub markers: Vec<usize>,
|
||||
pub qtype: QType,
|
||||
}
|
||||
|
||||
/// What the head budget did to a question's options.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct OptionStats {
|
||||
pub options: usize,
|
||||
/// Options whose token span is still unique after the cut. Fewer than
|
||||
/// `options` means the model can no longer tell some of them apart.
|
||||
pub options_distinct: usize,
|
||||
/// The per-option cap applied, when the options had to be shrunk.
|
||||
pub tokens_per_option: Option<usize>,
|
||||
}
|
||||
|
||||
/// A request turned into model input: one [`Row`] per question, in request
|
||||
/// order. With no questions there are no rows and nothing to run.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct Prepared {
|
||||
pub questions: Vec<Question>,
|
||||
pub rows: Vec<Row>,
|
||||
pub stats: Vec<OptionStats>,
|
||||
/// Every token the forward pass encodes: the metered input.
|
||||
pub input_tokens: usize,
|
||||
/// The id batched rows are padded with.
|
||||
pub pad_id: u32,
|
||||
}
|
||||
|
||||
/// The text a state is tokenized from: a string as-is, anything else as
|
||||
/// `json.dumps(state, ensure_ascii=False)`.
|
||||
pub fn serialize_state(state: &Json) -> String {
|
||||
match state {
|
||||
Json::String(s) => s.clone(),
|
||||
other => other.dumps(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate a request's questions and build their model input.
|
||||
///
|
||||
/// `questions` are the request's definitions in request order; `cfg` and
|
||||
/// `tok` belong to the checkpoint the request was routed to.
|
||||
pub fn prepare(
|
||||
state: &Json,
|
||||
questions: &[(String, Json)],
|
||||
cfg: &DecisionConfig,
|
||||
tok: &dyn TextEncoder,
|
||||
budget: Budget,
|
||||
) -> Result<Prepared, DecisionRequestError> {
|
||||
let special = tok.special();
|
||||
if questions.is_empty() {
|
||||
return Ok(Prepared {
|
||||
questions: Vec::new(),
|
||||
rows: Vec::new(),
|
||||
stats: Vec::new(),
|
||||
input_tokens: 0,
|
||||
pad_id: special.pad_id,
|
||||
});
|
||||
}
|
||||
let questions = questions_from_request(questions)?;
|
||||
let max_len = budget.max_len.unwrap_or(cfg.max_len);
|
||||
let head_max_len = budget.head_max_len.unwrap_or(cfg.head_max_len);
|
||||
let encode = |text: &str| {
|
||||
tok.encode(text)
|
||||
.map_err(|e| DecisionRequestError::new(500, format!("tokenizer failed: {e}")))
|
||||
};
|
||||
// Tokenized once and sliced per question.
|
||||
let state_ids = encode(&serialize_state(state).replace(&special.mask_token, " "))?;
|
||||
let truncate_left = matches!(state, Json::Array(_));
|
||||
let frame = Frame {
|
||||
cls: special.cls_id,
|
||||
sep: special.sep_id,
|
||||
max_len,
|
||||
head_max_len,
|
||||
};
|
||||
|
||||
let mut rows = Vec::with_capacity(questions.len());
|
||||
let mut stats = Vec::with_capacity(questions.len());
|
||||
for question in &questions {
|
||||
let instructions = question.instructions.replace(&special.mask_token, " ");
|
||||
let head_ids = encode(&format!(
|
||||
"{} question: {instructions}",
|
||||
question.kind.name()
|
||||
))?;
|
||||
let mut option_ids = Vec::with_capacity(question.options.len());
|
||||
for option in &question.options {
|
||||
let mut ids = vec![special.mask_id];
|
||||
let text = encode(&format!(" {}", option.replace(&special.mask_token, " ")))?;
|
||||
ids.extend(text.into_iter().take(MAX_OPTION_TOKENS));
|
||||
option_ids.push(ids);
|
||||
}
|
||||
let (row, s) = frame.assemble(head_ids, option_ids, &state_ids, truncate_left);
|
||||
if row.markers.len() != question.options.len() {
|
||||
return Err(unprocessable(format!(
|
||||
"question {} options exceed head_max_len={head_max_len}",
|
||||
Json::String(question.id.clone()).py_repr()
|
||||
)));
|
||||
}
|
||||
rows.push(Row {
|
||||
qtype: question.kind,
|
||||
..row
|
||||
});
|
||||
stats.push(s);
|
||||
}
|
||||
let input_tokens = rows.iter().map(|r| r.input_ids.len()).sum();
|
||||
Ok(Prepared {
|
||||
questions,
|
||||
rows,
|
||||
stats,
|
||||
input_tokens,
|
||||
pad_id: special.pad_id,
|
||||
})
|
||||
}
|
||||
|
||||
/// The fixed parts of every sequence for one request.
|
||||
struct Frame {
|
||||
cls: u32,
|
||||
sep: u32,
|
||||
max_len: usize,
|
||||
head_max_len: usize,
|
||||
}
|
||||
|
||||
impl Frame {
|
||||
/// The pure budget arithmetic of one sequence, over already-tokenized parts.
|
||||
fn assemble(
|
||||
&self,
|
||||
mut head_ids: Vec<u32>,
|
||||
mut option_ids: Vec<Vec<u32>>,
|
||||
state_ids: &[u32],
|
||||
truncate_left: bool,
|
||||
) -> (Row, OptionStats) {
|
||||
let Frame {
|
||||
cls,
|
||||
sep,
|
||||
max_len,
|
||||
head_max_len,
|
||||
} = *self;
|
||||
let used = |o: &[Vec<u32>]| o.iter().map(Vec::len).sum::<usize>();
|
||||
// Signed: options can overrun the head budget.
|
||||
let mut opt_budget = head_max_len as isize - used(&option_ids) as isize;
|
||||
let mut tokens_per_option = None;
|
||||
if opt_budget < MIN_INSTRUCTION_ROOM as isize {
|
||||
let per = MIN_TOKENS_PER_OPTION
|
||||
.max(head_max_len.saturating_sub(MIN_INSTRUCTION_ROOM) / option_ids.len().max(1));
|
||||
tokens_per_option = Some(per);
|
||||
for o in &mut option_ids {
|
||||
o.truncate(per);
|
||||
}
|
||||
opt_budget = head_max_len as isize - used(&option_ids) as isize;
|
||||
}
|
||||
head_ids.truncate(opt_budget.max(MIN_INSTRUCTION_TOKENS as isize) as usize);
|
||||
|
||||
let mut ids = Vec::with_capacity(max_len);
|
||||
ids.push(cls);
|
||||
ids.extend(head_ids);
|
||||
ids.push(sep);
|
||||
let mut markers = Vec::with_capacity(option_ids.len());
|
||||
for o in &option_ids {
|
||||
markers.push(ids.len());
|
||||
ids.extend_from_slice(o);
|
||||
}
|
||||
ids.push(sep);
|
||||
let room = max_len.saturating_sub(ids.len() + 1);
|
||||
let state = if truncate_left {
|
||||
&state_ids[state_ids.len().saturating_sub(room)..]
|
||||
} else {
|
||||
&state_ids[..room.min(state_ids.len())]
|
||||
};
|
||||
ids.extend_from_slice(state);
|
||||
ids.push(sep);
|
||||
ids.truncate(max_len);
|
||||
markers.retain(|&m| m < max_len);
|
||||
|
||||
let mut distinct = option_ids.clone();
|
||||
distinct.sort();
|
||||
distinct.dedup();
|
||||
(
|
||||
Row {
|
||||
input_ids: ids,
|
||||
markers,
|
||||
qtype: QType::Choice,
|
||||
},
|
||||
OptionStats {
|
||||
options: option_ids.len(),
|
||||
options_distinct: distinct.len(),
|
||||
tokens_per_option,
|
||||
},
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
const CLS: u32 = 1;
|
||||
const SEP: u32 = 2;
|
||||
const M: u32 = 9;
|
||||
|
||||
fn frame(max_len: usize, head_max_len: usize) -> Frame {
|
||||
Frame {
|
||||
cls: CLS,
|
||||
sep: SEP,
|
||||
max_len,
|
||||
head_max_len,
|
||||
}
|
||||
}
|
||||
|
||||
fn opts(lens: &[usize]) -> Vec<Vec<u32>> {
|
||||
lens.iter()
|
||||
.enumerate()
|
||||
.map(|(i, &n)| {
|
||||
let mut o = vec![M];
|
||||
o.extend((0..n - 1).map(|j| 100 + (i * 50 + j) as u32));
|
||||
o
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn roomy_budget_keeps_everything() {
|
||||
let (row, s) = frame(64, 32).assemble(vec![10, 11], opts(&[2, 3]), &[50, 51], false);
|
||||
assert_eq!(
|
||||
row.input_ids,
|
||||
[CLS, 10, 11, SEP, M, 100, M, 150, 151, SEP, 50, 51, SEP]
|
||||
);
|
||||
assert_eq!(row.markers, [4, 6]);
|
||||
assert_eq!(
|
||||
(s.options, s.options_distinct, s.tokens_per_option),
|
||||
(2, 2, None)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn crowded_head_shrinks_options_evenly_and_keeps_eight_instruction_tokens() {
|
||||
// 3 options of 10 = 30 tokens against a 40-token head: 10 left < 16.
|
||||
let head: Vec<u32> = (10..40).collect();
|
||||
let (row, s) = frame(128, 40).assemble(head, opts(&[10, 10, 10]), &[], false);
|
||||
// per = max(4, (40 - 16) / 3) = 8; 40 - 24 = 16 instruction tokens.
|
||||
assert_eq!(s.tokens_per_option, Some(8));
|
||||
assert_eq!(row.markers, [18, 26, 34]);
|
||||
let (row, s) = frame(256, 40).assemble((10..40).collect(), opts(&[10; 12]), &[], false);
|
||||
// per = max(4, 24 / 12 = 2) = 4; the budget goes negative, 8 tokens survive.
|
||||
assert_eq!(s.tokens_per_option, Some(4));
|
||||
assert_eq!(row.markers[0], 1 + 8 + 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn state_is_cut_from_the_side_the_state_shape_asks_for() {
|
||||
let state: Vec<u32> = (50..60).collect();
|
||||
let (right, _) = frame(8, 32).assemble(vec![], opts(&[1]), &state, false);
|
||||
let (left, _) = frame(8, 32).assemble(vec![], opts(&[1]), &state, true);
|
||||
// [CLS] [SEP] [MASK] [SEP] = 4 tokens, 1 reserved for the last [SEP]: room 3.
|
||||
assert_eq!(right.input_ids, [CLS, SEP, M, SEP, 50, 51, 52, SEP]);
|
||||
assert_eq!(left.input_ids, [CLS, SEP, M, SEP, 57, 58, 59, SEP]);
|
||||
// No room at all: none of the state, not all of it.
|
||||
let (none, _) = frame(5, 32).assemble(vec![], opts(&[1]), &state, true);
|
||||
assert_eq!(none.input_ids, [CLS, SEP, M, SEP, SEP]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn markers_past_max_len_are_dropped_and_collisions_counted() {
|
||||
let same = vec![vec![M, 7, 7], vec![M, 7, 7]];
|
||||
let (row, s) = frame(5, 64).assemble(vec![], same, &[], false);
|
||||
// [CLS] [SEP] then markers at 2 and 5; only the first fits in 5.
|
||||
assert_eq!(row.markers, [2]);
|
||||
assert_eq!((s.options, s.options_distinct), (2, 1));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,318 @@
|
||||
//! Per-option logits → calibrated, typed answers.
|
||||
//!
|
||||
//! The reference computes this in numpy on the float32 logits the model
|
||||
//! returns, and publishes each probability rounded to 4 decimal places.
|
||||
//! Reproducing the rounded values exactly means reproducing the float32
|
||||
//! arithmetic, in numpy's order: division by a float32 temperature, a
|
||||
//! max-subtracted `exp`, and numpy's *pairwise* summation (not a running
|
||||
//! sum) for every reduction. The expected score is the exception: numpy
|
||||
//! promotes `arange(k) * p` to float64.
|
||||
|
||||
use super::question::QType;
|
||||
use super::sequence::Prepared;
|
||||
use super::{DECISION_MODEL_NAME, DecisionConfig};
|
||||
use cortex_core::decisions::{
|
||||
CollapsedOptions, DecisionAction, DecisionAnswer, DecisionUsage, OrderedMap, SystemOneResponse,
|
||||
};
|
||||
|
||||
/// The model's output for one [`super::Row`].
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct RowOutput {
|
||||
/// The scorer's raw logit at each option marker, before calibration.
|
||||
/// Entries past the row's marker count (batch padding) are ignored.
|
||||
pub logits: Vec<f32>,
|
||||
/// The act head's probability of acting: `softmax(act_logits)[0]`.
|
||||
pub act_probability: f32,
|
||||
}
|
||||
|
||||
/// Build the answers and usage for a prepared request from the model's
|
||||
/// per-row outputs (`outputs[i]` belongs to `prepared.rows[i]`).
|
||||
///
|
||||
/// `lang` selects per-language calibration where the checkpoint has it:
|
||||
/// the language the request was declared or detected as.
|
||||
pub fn shape(
|
||||
prepared: &Prepared,
|
||||
outputs: &[RowOutput],
|
||||
cfg: &DecisionConfig,
|
||||
lang: Option<&str>,
|
||||
) -> (OrderedMap<DecisionAnswer>, DecisionUsage) {
|
||||
assert_eq!(
|
||||
outputs.len(),
|
||||
prepared.rows.len(),
|
||||
"one model output per prepared row"
|
||||
);
|
||||
let temps = cfg.temperatures_for(lang);
|
||||
let mut answers = Vec::with_capacity(outputs.len());
|
||||
let mut collapsed = Vec::new();
|
||||
for ((question, row), (out, stats)) in prepared
|
||||
.questions
|
||||
.iter()
|
||||
.zip(&prepared.rows)
|
||||
.zip(outputs.iter().zip(&prepared.stats))
|
||||
{
|
||||
let k = row.markers.len();
|
||||
let p = calibrate(
|
||||
&out.logits[..k],
|
||||
temps.for_question(question.kind as usize, k),
|
||||
);
|
||||
let answer_confidence = round4(
|
||||
p.iter()
|
||||
.copied()
|
||||
.fold(f32::NEG_INFINITY, f32::max)
|
||||
.clamp(0.0, 1.0) as f64,
|
||||
);
|
||||
let action = DecisionAction {
|
||||
act_probability: round4(out.act_probability as f64),
|
||||
};
|
||||
let probabilities = |key: &dyn Fn(usize) -> String| {
|
||||
OrderedMap(
|
||||
p.iter()
|
||||
.enumerate()
|
||||
.map(|(i, &v)| (key(i), round4(v as f64)))
|
||||
.collect(),
|
||||
)
|
||||
};
|
||||
let answer = match question.kind {
|
||||
QType::Choice => DecisionAnswer::Choice {
|
||||
choice: question.labels[argmax(&p)].clone(),
|
||||
probabilities: probabilities(&|i| {
|
||||
question.labels[i]
|
||||
.json_key()
|
||||
.expect("choice labels are scalars")
|
||||
}),
|
||||
confidence: round4(entropy_confidence(&p) as f64),
|
||||
answer_confidence,
|
||||
action,
|
||||
},
|
||||
QType::Score => {
|
||||
// int64 * float32 promotes to float64 in numpy.
|
||||
let expected: Vec<f64> = p
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, &v)| i as f64 * v as f64)
|
||||
.collect();
|
||||
DecisionAnswer::Score {
|
||||
score: round4(pairwise_sum(&expected)),
|
||||
legend: OrderedMap(
|
||||
question
|
||||
.labels
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, c)| (i.to_string(), c.clone()))
|
||||
.collect(),
|
||||
),
|
||||
probabilities: probabilities(&|i| i.to_string()),
|
||||
confidence: round4(entropy_confidence(&p) as f64),
|
||||
answer_confidence,
|
||||
action,
|
||||
}
|
||||
}
|
||||
QType::Noul => {
|
||||
let yes = p[1] as f64;
|
||||
DecisionAnswer::Noul {
|
||||
noul: round4(yes),
|
||||
confidence: round4(yes.max(1.0 - yes)),
|
||||
answer_confidence,
|
||||
action,
|
||||
}
|
||||
}
|
||||
};
|
||||
answers.push((question.id.clone(), answer));
|
||||
if stats.options_distinct < stats.options {
|
||||
collapsed.push((
|
||||
question.id.clone(),
|
||||
CollapsedOptions {
|
||||
total: stats.options,
|
||||
distinct: stats.options_distinct,
|
||||
tokens_per_option: stats.tokens_per_option,
|
||||
},
|
||||
));
|
||||
}
|
||||
}
|
||||
let usage = DecisionUsage {
|
||||
input_tokens: prepared.input_tokens,
|
||||
output_tokens: 0,
|
||||
options: (!collapsed.is_empty()).then_some(OrderedMap(collapsed)),
|
||||
};
|
||||
(OrderedMap(answers), usage)
|
||||
}
|
||||
|
||||
/// Assemble the full response once routing is known.
|
||||
pub fn response(
|
||||
answers: OrderedMap<DecisionAnswer>,
|
||||
usage: DecisionUsage,
|
||||
routing: Option<cortex_core::decisions::DecisionRouting>,
|
||||
) -> SystemOneResponse {
|
||||
SystemOneResponse {
|
||||
model: DECISION_MODEL_NAME.to_string(),
|
||||
answers,
|
||||
usage,
|
||||
routing,
|
||||
}
|
||||
}
|
||||
|
||||
/// Temperature-scaled softmax, as the reference evaluates it in numpy:
|
||||
/// `z = logits / t`, `e = exp(z - z.max())`, `p = e / e.sum()`.
|
||||
///
|
||||
/// numpy's SIMD `expf` and libm's disagree by one ulp on about a third of
|
||||
/// inputs. That is invisible after the 4-decimal rounding on every
|
||||
/// reference fixture, and far below the engine's own distance from the
|
||||
/// reference's logits, so it is not reproduced; the steps either side of
|
||||
/// the `exp` are.
|
||||
fn calibrate(logits: &[f32], t: f64) -> Vec<f32> {
|
||||
let z = scale(logits, t);
|
||||
let max = z.iter().copied().fold(f32::NEG_INFINITY, f32::max);
|
||||
let e: Vec<f32> = z.iter().map(|&v| (v - max).exp()).collect();
|
||||
normalize(&e)
|
||||
}
|
||||
|
||||
/// `logits / t` in float32: numpy casts the Python-float temperature to the
|
||||
/// array's dtype rather than promoting the array.
|
||||
fn scale(logits: &[f32], t: f64) -> Vec<f32> {
|
||||
let t = t as f32;
|
||||
logits.iter().map(|&l| l / t).collect()
|
||||
}
|
||||
|
||||
/// `e / e.sum()` with numpy's pairwise float32 sum.
|
||||
fn normalize(e: &[f32]) -> Vec<f32> {
|
||||
let sum = pairwise_sum(e);
|
||||
e.iter().map(|&v| v / sum).collect()
|
||||
}
|
||||
|
||||
/// `1 - H(p) / log(k)`, clipped to [0, 1]: how concentrated the whole
|
||||
/// distribution is. Not calibrated; see `answer_confidence` for that.
|
||||
fn entropy_confidence(p: &[f32]) -> f32 {
|
||||
let k = p.len();
|
||||
if k < 2 {
|
||||
return 1.0;
|
||||
}
|
||||
let terms: Vec<f32> = p.iter().map(|&v| v * v.clamp(1e-12, 1.0).ln()).collect();
|
||||
let ent = -pairwise_sum(&terms);
|
||||
(1.0 - ent / (k as f64).ln() as f32).clamp(0.0, 1.0)
|
||||
}
|
||||
|
||||
/// The first index of the largest value, as `argmax` picks it.
|
||||
fn argmax(p: &[f32]) -> usize {
|
||||
let mut best = 0;
|
||||
for (i, &v) in p.iter().enumerate() {
|
||||
if v > p[best] {
|
||||
best = i;
|
||||
}
|
||||
}
|
||||
best
|
||||
}
|
||||
|
||||
/// Python's `round(x, 4)`: correctly rounded, ties to even on the exact
|
||||
/// binary value. Rust's fixed-precision formatting rounds the same way.
|
||||
fn round4(x: f64) -> f64 {
|
||||
format!("{x:.4}").parse().expect("formatted float parses")
|
||||
}
|
||||
|
||||
/// numpy's pairwise summation over a contiguous array, which the rounded
|
||||
/// outputs depend on: a plain loop below 8 elements, eight interleaved
|
||||
/// accumulators up to 128, and recursive halving (at a multiple of 8)
|
||||
/// beyond.
|
||||
fn pairwise_sum<T>(a: &[T]) -> T
|
||||
where
|
||||
T: Copy + std::ops::Add<Output = T> + Default,
|
||||
{
|
||||
const BLOCK: usize = 128;
|
||||
let n = a.len();
|
||||
if n < 8 {
|
||||
let mut res = T::default();
|
||||
for &v in a {
|
||||
res = res + v;
|
||||
}
|
||||
res
|
||||
} else if n <= BLOCK {
|
||||
let mut r = [a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7]];
|
||||
let mut i = 8;
|
||||
while i < n - n % 8 {
|
||||
for j in 0..8 {
|
||||
r[j] = r[j] + a[i + j];
|
||||
}
|
||||
i += 8;
|
||||
}
|
||||
let mut res = ((r[0] + r[1]) + (r[2] + r[3])) + ((r[4] + r[5]) + (r[6] + r[7]));
|
||||
while i < n {
|
||||
res = res + a[i];
|
||||
i += 1;
|
||||
}
|
||||
res
|
||||
} else {
|
||||
let mut n2 = n / 2;
|
||||
n2 -= n2 % 8;
|
||||
pairwise_sum(&a[..n2]) + pairwise_sum(&a[n2..])
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn pairwise_sum_differs_from_a_running_sum_where_numpy_does() {
|
||||
// 1e8 swamps every 1.0 added to it one at a time in f32; pairwise
|
||||
// grouping keeps them. `int(np.float32 array.sum())` of this array is 100000008.
|
||||
let mut a = vec![1.0f32; 16];
|
||||
a[0] = 1e8;
|
||||
let running = a.iter().fold(0.0f32, |s, &v| s + v);
|
||||
assert_eq!(running, 1e8);
|
||||
assert_eq!(pairwise_sum(&a), 100_000_008.0);
|
||||
}
|
||||
|
||||
/// Rows of `testdata/laya/numpy_softmax.txt`: the reference's
|
||||
/// calibration run in numpy on random float32 logits, with the bits of
|
||||
/// every intermediate. Bit equality pins the evaluation order — a
|
||||
/// running sum instead of numpy's pairwise one, or dividing by the
|
||||
/// temperature in float64, changes low bits that the 4-decimal rounding
|
||||
/// usually hides.
|
||||
fn numpy_rows() -> Vec<(f64, [Vec<f32>; 4])> {
|
||||
include_str!("../testdata/laya/numpy_softmax.txt")
|
||||
.lines()
|
||||
.filter(|l| !l.starts_with('#') && !l.trim().is_empty())
|
||||
.map(|l| {
|
||||
let mut f = l.split('|');
|
||||
let t: f64 = f.next().unwrap().parse().unwrap();
|
||||
let bits = |s: &str| -> Vec<f32> {
|
||||
s.split(',')
|
||||
.map(|h| f32::from_bits(u32::from_str_radix(&h[2..], 16).unwrap()))
|
||||
.collect()
|
||||
};
|
||||
let v: Vec<Vec<f32>> = f.map(bits).collect();
|
||||
(t, [v[0].clone(), v[1].clone(), v[2].clone(), v[3].clone()])
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn bits(v: &[f32]) -> Vec<u32> {
|
||||
v.iter().map(|x| x.to_bits()).collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn calibration_matches_numpy_either_side_of_exp() {
|
||||
let rows = numpy_rows();
|
||||
assert_eq!(rows.len(), 4);
|
||||
for (t, [logits, z, e, p]) in rows {
|
||||
assert_eq!(
|
||||
bits(&scale(&logits, t)),
|
||||
bits(&z),
|
||||
"k={} t={t}: z",
|
||||
logits.len()
|
||||
);
|
||||
assert_eq!(
|
||||
bits(&normalize(&e)),
|
||||
bits(&p),
|
||||
"k={} t={t}: p",
|
||||
logits.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn round4_ties_to_even_on_the_binary_value() {
|
||||
assert_eq!(round4(0.03125), 0.0312);
|
||||
assert_eq!(round4(1.00005), 1.0001); // 1.00005 is just above the tie in binary
|
||||
assert_eq!(round4(0.99995), 1.0);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
//! The tokenizer a decision checkpoint ships, and its special tokens.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
/// The special tokens the sequence layout uses, resolved to ids.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SpecialTokens {
|
||||
/// The option marker's text (`[MASK]`, `<mask>` for mmBERT). Literal
|
||||
/// occurrences in caller text are replaced with a space before
|
||||
/// tokenizing, so no caller can forge a marker.
|
||||
pub mask_token: String,
|
||||
pub cls_id: u32,
|
||||
pub sep_id: u32,
|
||||
pub mask_id: u32,
|
||||
/// Padding for batched rows; not part of any sequence.
|
||||
pub pad_id: u32,
|
||||
}
|
||||
|
||||
/// Text → token ids, with no special tokens added.
|
||||
pub trait TextEncoder {
|
||||
fn encode(&self, text: &str) -> Result<Vec<u32>, String>;
|
||||
fn special(&self) -> &SpecialTokens;
|
||||
}
|
||||
|
||||
/// A checkpoint's `tokenizer/` directory: `tokenizer.json` plus the special
|
||||
/// token names in `tokenizer_config.json`.
|
||||
pub struct DecisionTokenizer {
|
||||
tok: tokenizers::Tokenizer,
|
||||
special: SpecialTokens,
|
||||
}
|
||||
|
||||
impl DecisionTokenizer {
|
||||
pub fn from_dir(dir: &Path) -> anyhow::Result<Self> {
|
||||
let mut tok = tokenizers::Tokenizer::from_file(dir.join("tokenizer.json"))
|
||||
.map_err(|e| anyhow::anyhow!("{}: {e}", dir.join("tokenizer.json").display()))?;
|
||||
// The Python tokenizer call the reference makes pads and truncates
|
||||
// nothing unless asked; a tokenizer.json that carries either setting
|
||||
// would otherwise apply it here.
|
||||
tok.with_truncation(None)
|
||||
.map_err(|e| anyhow::anyhow!("disabling truncation: {e}"))?;
|
||||
tok.with_padding(None);
|
||||
let cfg_path = dir.join("tokenizer_config.json");
|
||||
let cfg: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(&cfg_path)?)
|
||||
.map_err(|e| anyhow::anyhow!("{}: {e}", cfg_path.display()))?;
|
||||
let name = |key: &str| -> anyhow::Result<String> {
|
||||
cfg.get(key)
|
||||
.and_then(|v| {
|
||||
v.as_str()
|
||||
.or_else(|| v.get("content").and_then(|c| c.as_str()))
|
||||
})
|
||||
.map(str::to_owned)
|
||||
.ok_or_else(|| anyhow::anyhow!("{}: no {key}", cfg_path.display()))
|
||||
};
|
||||
let id = |token: &str| -> anyhow::Result<u32> {
|
||||
tok.token_to_id(token)
|
||||
.ok_or_else(|| anyhow::anyhow!("special token {token:?} is not in the vocabulary"))
|
||||
};
|
||||
let mask_token = name("mask_token")?;
|
||||
let special = SpecialTokens {
|
||||
cls_id: id(&name("cls_token")?)?,
|
||||
sep_id: id(&name("sep_token")?)?,
|
||||
mask_id: id(&mask_token)?,
|
||||
pad_id: id(&name("pad_token")?)?,
|
||||
mask_token,
|
||||
};
|
||||
Ok(Self { tok, special })
|
||||
}
|
||||
}
|
||||
|
||||
impl TextEncoder for DecisionTokenizer {
|
||||
fn encode(&self, text: &str) -> Result<Vec<u32>, String> {
|
||||
self.tok
|
||||
.encode(text, false)
|
||||
.map(|e| e.get_ids().to_vec())
|
||||
.map_err(|e| e.to_string())
|
||||
}
|
||||
|
||||
fn special(&self) -> &SpecialTokens {
|
||||
&self.special
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,252 @@
|
||||
//! A loaded Laya decision model, as the rest of neuron sees it (#336).
|
||||
//!
|
||||
//! The weights live in a device worker's decision slab (CUDA) or in an
|
||||
//! `Arc` scored on the blocking pool (CPU); either way the caller holds
|
||||
//! only this struct and asks it to [`DecisionCheckpoint::decide`] a
|
||||
//! batch of already-assembled questions. Turning a `/v1/systemone`
|
||||
//! request into that batch — rendering options, budgets, tokenisation —
|
||||
//! and turning the scores back into calibrated answers happen outside
|
||||
//! this module.
|
||||
//!
|
||||
//! One Hugging Face repo can carry several checkpoints (the released
|
||||
//! `convaiinnovations/laya` has the English one at its root and
|
||||
//! `multilingual/` and `typed-decisions/` beside it), so a loaded model
|
||||
//! is a list of checkpoints, each named by its subfolder.
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use candle_core::DType;
|
||||
|
||||
use super::arch::laya::{CheckpointConfig, DecisionBatch, DecisionRow, LayaModel};
|
||||
use super::decision::{DecisionConfig, DecisionTokenizer};
|
||||
use super::device_worker::{DecisionHandle, DeviceWorkerHandle};
|
||||
use cortex_core::harness::ModelSpec;
|
||||
|
||||
/// Routing name of the checkpoint at the repo root. A released Laya
|
||||
/// family's root is its English checkpoint, and upstream routes to it
|
||||
/// under that name.
|
||||
pub const ROOT_CHECKPOINT: &str = super::decision::route::ENGLISH;
|
||||
|
||||
/// Where one checkpoint's weights live and how to reach them.
|
||||
enum Backend {
|
||||
/// Owned by a CUDA device worker; scored through its channel.
|
||||
Worker {
|
||||
worker: Arc<DeviceWorkerHandle>,
|
||||
handle: DecisionHandle,
|
||||
},
|
||||
/// CPU fallback. `LayaModel::forward` takes `&self`, so concurrent
|
||||
/// requests share it without a lock.
|
||||
Cpu(Arc<LayaModel>),
|
||||
}
|
||||
|
||||
/// The files one checkpoint needs, already on local disk.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CheckpointFiles {
|
||||
/// The checkpoint directory: holds `rl_agent_config.json`,
|
||||
/// `encoder/config.json` and `model.safetensors`.
|
||||
pub dir: PathBuf,
|
||||
pub tokenizer: PathBuf,
|
||||
}
|
||||
|
||||
pub struct DecisionCheckpoint {
|
||||
/// The checkpoint's routing name: [`ROOT_CHECKPOINT`] for the repo
|
||||
/// root, else its subfolder (`multilingual`, `typed-decisions`).
|
||||
pub name: String,
|
||||
pub config: CheckpointConfig,
|
||||
/// Budgets and fitted temperatures from `rl_agent_config.json`, as
|
||||
/// the prompt builder and calibration read them.
|
||||
pub calibration: DecisionConfig,
|
||||
pub tokenizer: DecisionTokenizer,
|
||||
/// Precision the weights were loaded in.
|
||||
pub dtype: DType,
|
||||
backend: Backend,
|
||||
}
|
||||
|
||||
impl DecisionCheckpoint {
|
||||
/// Load one checkpoint onto `worker`'s device, or onto the CPU when
|
||||
/// `worker` is `None`.
|
||||
pub async fn load(
|
||||
name: String,
|
||||
files: CheckpointFiles,
|
||||
model_id: &str,
|
||||
worker: Option<Arc<DeviceWorkerHandle>>,
|
||||
) -> Result<Self> {
|
||||
let config = CheckpointConfig::from_dir(&files.dir)?;
|
||||
let calibration_path = files.dir.join("rl_agent_config.json");
|
||||
let calibration = DecisionConfig::from_json(
|
||||
&std::fs::read_to_string(&calibration_path)
|
||||
.with_context(|| format!("read {}", calibration_path.display()))?,
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
let tokenizer_dir = files
|
||||
.tokenizer
|
||||
.parent()
|
||||
.context("tokenizer.json has no parent directory")?;
|
||||
let tokenizer = DecisionTokenizer::from_dir(tokenizer_dir)?;
|
||||
let (dtype, backend) = match worker {
|
||||
Some(worker) => {
|
||||
let wanted = device_dtype(config.laya.amp_dtype.as_deref());
|
||||
let (handle, dtype) = worker
|
||||
.load_decision(files.dir.clone(), model_id.to_string(), wanted)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("worker load_decision: {e:#}"))?;
|
||||
(dtype, Backend::Worker { worker, handle })
|
||||
}
|
||||
None => {
|
||||
let dir = files.dir.clone();
|
||||
let model = tokio::task::spawn_blocking(move || {
|
||||
LayaModel::load(&dir, &candle_core::Device::Cpu, DType::F32)
|
||||
})
|
||||
.await
|
||||
.context("decision load task")??;
|
||||
(DType::F32, Backend::Cpu(Arc::new(model)))
|
||||
}
|
||||
};
|
||||
Ok(Self {
|
||||
name,
|
||||
config,
|
||||
calibration,
|
||||
tokenizer,
|
||||
dtype,
|
||||
backend,
|
||||
})
|
||||
}
|
||||
|
||||
/// Score one batch: one [`DecisionRow`] per sequence, in order.
|
||||
pub async fn decide(&self, batch: DecisionBatch) -> Result<Vec<DecisionRow>> {
|
||||
match &self.backend {
|
||||
Backend::Worker { worker, handle } => worker
|
||||
.decide(*handle, batch)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{e:#}")),
|
||||
Backend::Cpu(model) => {
|
||||
let model = Arc::clone(model);
|
||||
tokio::task::spawn_blocking(move || model.forward(&batch))
|
||||
.await
|
||||
.context("decision forward task")?
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Release the weights. On a worker this drops them on the thread
|
||||
/// that owns the device context.
|
||||
pub async fn release(&self) {
|
||||
if let Backend::Worker { worker, handle } = &self.backend
|
||||
&& let Err(e) = worker.drop_decision(*handle).await
|
||||
{
|
||||
tracing::warn!(
|
||||
checkpoint = %self.name,
|
||||
error = %e,
|
||||
"decision unload: DropDecision RPC failed (weights may leak in worker slab)"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The precision to serve in on a GPU: the autocast dtype the reference
|
||||
/// itself runs there (`amp_dtype` in `rl_agent_config.json`), else bf16.
|
||||
fn device_dtype(amp_dtype: Option<&str>) -> DType {
|
||||
match amp_dtype {
|
||||
Some("fp16") | Some("float16") => DType::F16,
|
||||
Some("fp32") | Some("float32") => DType::F32,
|
||||
_ => DType::BF16,
|
||||
}
|
||||
}
|
||||
|
||||
pub struct LoadedDecisionModel {
|
||||
pub model_id: String,
|
||||
pub spec: ModelSpec,
|
||||
pub devices: Vec<u32>,
|
||||
/// At least one; the first is the default.
|
||||
pub checkpoints: Vec<DecisionCheckpoint>,
|
||||
/// Set when a forward returns a driver-level error, mirroring the
|
||||
/// other model kinds: `/models` reports it and requests fast-reject.
|
||||
pub poisoned: Arc<AtomicBool>,
|
||||
pub admission: super::admission::AdmissionController,
|
||||
}
|
||||
|
||||
impl LoadedDecisionModel {
|
||||
/// The checkpoint called `name`, or the default one for `None`.
|
||||
pub fn checkpoint(&self, name: Option<&str>) -> Option<&DecisionCheckpoint> {
|
||||
match name {
|
||||
None => self.checkpoints.first(),
|
||||
Some(n) => self.checkpoints.iter().find(|c| c.name == n),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn release(&self) {
|
||||
for c in &self.checkpoints {
|
||||
c.release().await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn tiny() -> CheckpointFiles {
|
||||
let dir =
|
||||
std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("src/harness/testdata/laya/tiny");
|
||||
CheckpointFiles {
|
||||
tokenizer: dir.join("tokenizer.json"),
|
||||
dir,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gpu_dtype_follows_the_reference_autocast() {
|
||||
assert_eq!(device_dtype(Some("bf16")), DType::BF16);
|
||||
assert_eq!(device_dtype(Some("fp16")), DType::F16);
|
||||
assert_eq!(device_dtype(None), DType::BF16);
|
||||
}
|
||||
|
||||
/// The whole async path on the CPU backend, and the device-worker
|
||||
/// path on a CPU build (the worker runs, with a CPU device), must
|
||||
/// give the same scores as the model called directly.
|
||||
#[tokio::test]
|
||||
async fn both_backends_score_like_the_model() {
|
||||
let files = tiny();
|
||||
let direct = LayaModel::load(&files.dir, &candle_core::Device::Cpu, DType::F32).unwrap();
|
||||
let batch = DecisionBatch {
|
||||
seqs: vec![super::super::arch::laya::DecisionSeq {
|
||||
input_ids: vec![5, 9, 11, 3, 7, 2],
|
||||
markers: vec![1, 3],
|
||||
qtype: 0,
|
||||
}],
|
||||
};
|
||||
let want = direct.forward(&batch).unwrap();
|
||||
|
||||
let cpu = DecisionCheckpoint::load("root".into(), files.clone(), "tiny", None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(cpu.decide(batch.clone()).await.unwrap(), want);
|
||||
|
||||
// A CPU build's worker holds a CPU device, so this exercises the
|
||||
// job round-trip (load, decide, drop) without a GPU — and the
|
||||
// worker must fall back to f32 there, reporting that it did.
|
||||
let worker = DeviceWorkerHandle::spawn(0).unwrap();
|
||||
let on_worker =
|
||||
DecisionCheckpoint::load("root".into(), files, "tiny", Some(worker.clone()))
|
||||
.await
|
||||
.unwrap();
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
{
|
||||
assert_eq!(on_worker.dtype, DType::F32);
|
||||
assert_eq!(on_worker.decide(batch).await.unwrap(), want);
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
{
|
||||
let got = on_worker.decide(batch).await.unwrap();
|
||||
for (g, w) in got[0].logits.iter().zip(&want[0].logits) {
|
||||
assert!((g - w).abs() < 5e-2, "{g} vs {w}");
|
||||
}
|
||||
}
|
||||
on_worker.release().await;
|
||||
assert!(on_worker.decide(DecisionBatch::default()).await.is_err());
|
||||
worker.shutdown().unwrap();
|
||||
}
|
||||
}
|
||||
@@ -13,12 +13,13 @@
|
||||
//! ARCH model state in this state slab will gain a companion
|
||||
//! `tp_models: HashMap<TpHandle, Box<TpLeaderModel>>`.
|
||||
|
||||
use crate::harness::arch::laya::LayaModel;
|
||||
use crate::harness::arch::snapshot::KvCacheSnapshot;
|
||||
use crate::harness::candle::ModelArch;
|
||||
#[cfg(feature = "cuda")]
|
||||
use crate::harness::device_worker::jobs::TpHandle;
|
||||
use crate::harness::device_worker::jobs::{
|
||||
ArchHandle, DenseLoad, ImageHandle, ImageInput, Job, KvSnapshotId,
|
||||
ArchHandle, DecisionHandle, DenseLoad, ImageHandle, ImageInput, Job, KvSnapshotId,
|
||||
};
|
||||
use crate::harness::image::ZImagePipeline;
|
||||
#[cfg(feature = "cuda")]
|
||||
@@ -65,6 +66,11 @@ struct DeviceWorkerState {
|
||||
image_models: HashMap<ImageHandle, Box<ZImagePipeline>>,
|
||||
/// Counter for minting fresh `ImageHandle`s.
|
||||
next_image_handle: u64,
|
||||
/// Laya decision-model slab (#336). Same lifecycle discipline as
|
||||
/// `image_models`, in its own namespace.
|
||||
decision_models: HashMap<DecisionHandle, Box<LayaModel>>,
|
||||
/// Counter for minting fresh `DecisionHandle`s.
|
||||
next_decision_handle: u64,
|
||||
/// Leader's NCCL state. Populated by `Job::NcclInit`; the
|
||||
/// underlying `Comm`'s libnccl handle lives bound to this thread
|
||||
/// for its entire lifetime. Subprocess workers maintain their own
|
||||
@@ -751,6 +757,62 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
|
||||
trim_device_pool(&state);
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
Job::LoadDecision {
|
||||
dir,
|
||||
model_id,
|
||||
dtype,
|
||||
reply,
|
||||
} => {
|
||||
let dtype = if state.device.is_cpu() {
|
||||
candle_core::DType::F32
|
||||
} else {
|
||||
dtype
|
||||
};
|
||||
let result = LayaModel::load(&dir, &state.device, dtype)
|
||||
.with_context(|| format!("load laya checkpoint {}", dir.display()))
|
||||
.map(|model| {
|
||||
let handle = DecisionHandle(state.next_decision_handle);
|
||||
state.next_decision_handle = state.next_decision_handle.wrapping_add(1);
|
||||
state.decision_models.insert(handle, Box::new(model));
|
||||
tracing::info!(
|
||||
device_index,
|
||||
model_id,
|
||||
checkpoint = %dir.display(),
|
||||
handle = handle.0,
|
||||
"device worker: decision model loaded"
|
||||
);
|
||||
(handle, dtype)
|
||||
});
|
||||
trim_device_pool(&state);
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
Job::DropDecision { handle, reply } => {
|
||||
let removed = state.decision_models.remove(&handle);
|
||||
let was_present = removed.is_some();
|
||||
// Explicit drop on this thread, as for DropArch/DropImage.
|
||||
drop(removed);
|
||||
if was_present {
|
||||
trim_device_pool(&state);
|
||||
}
|
||||
tracing::debug!(
|
||||
device_index,
|
||||
handle = handle.0,
|
||||
was_present,
|
||||
"device worker: decision model dropped"
|
||||
);
|
||||
let _ = reply.send(());
|
||||
}
|
||||
Job::Decide {
|
||||
handle,
|
||||
batch,
|
||||
reply,
|
||||
} => {
|
||||
let result = match state.decision_models.get(&handle) {
|
||||
Some(model) => model.forward(&batch),
|
||||
None => Err(anyhow::anyhow!("no decision model for handle {}", handle.0)),
|
||||
};
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
Job::Shutdown => unreachable!("Shutdown should break above"),
|
||||
}
|
||||
}
|
||||
@@ -823,6 +885,8 @@ fn init_state(device_index: u32) -> DeviceWorkerState {
|
||||
next_kv_snapshot_id: 1,
|
||||
image_models: HashMap::new(),
|
||||
next_image_handle: 1,
|
||||
decision_models: HashMap::new(),
|
||||
next_decision_handle: 1,
|
||||
nccl: NcclState::new(),
|
||||
tp_models: HashMap::new(),
|
||||
next_tp_handle: 1,
|
||||
@@ -841,6 +905,8 @@ fn init_state(device_index: u32) -> DeviceWorkerState {
|
||||
next_kv_snapshot_id: 1,
|
||||
image_models: HashMap::new(),
|
||||
next_image_handle: 1,
|
||||
decision_models: HashMap::new(),
|
||||
next_decision_handle: 1,
|
||||
nccl: NcclState::new(),
|
||||
}
|
||||
}
|
||||
@@ -1969,6 +2035,15 @@ fn drain_poisoned(job: Job, device_index: u32) {
|
||||
Job::GenerateImage { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
Job::LoadDecision { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
Job::DropDecision { reply, .. } => {
|
||||
let _ = reply.send(());
|
||||
}
|
||||
Job::Decide { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
Job::ForwardLogitsWithImages { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
|
||||
@@ -59,6 +59,15 @@ pub struct TpHandle(pub u64);
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub struct ImageHandle(pub u64);
|
||||
|
||||
/// Opaque handle to a Laya decision model ([`LayaModel`]) stored in the
|
||||
/// worker thread's decision slab (#336). Same discipline as
|
||||
/// [`ImageHandle`]: construction, every forward and `Drop` happen on the
|
||||
/// worker thread, and `Job::DropDecision` is the only way to release it.
|
||||
///
|
||||
/// [`LayaModel`]: crate::harness::arch::laya::LayaModel
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub struct DecisionHandle(pub u64);
|
||||
|
||||
/// Opaque handle to a prefix-cache snapshot (#11) stored worker-side
|
||||
/// next to the model slab. Scoped to the `ArchHandle` it was captured
|
||||
/// from — `Job::DropArch` drops every snapshot under its handle. The
|
||||
@@ -549,6 +558,33 @@ pub enum Job {
|
||||
max_dim: usize,
|
||||
reply: oneshot::Sender<Result<crate::harness::image::ImageGenResult>>,
|
||||
},
|
||||
/// Load one Laya checkpoint directory (#336) — the repo root or one
|
||||
/// of its subfolders, already on local disk — into the decision
|
||||
/// slab. File resolution (hf-hub) happens on the async side.
|
||||
///
|
||||
/// `dtype` is the precision wanted on a GPU; a worker whose device
|
||||
/// is the CPU loads in f32 instead (the CPU backend has no bf16
|
||||
/// matmul). The reply carries the dtype actually used.
|
||||
LoadDecision {
|
||||
dir: PathBuf,
|
||||
model_id: String,
|
||||
dtype: candle_core::DType,
|
||||
reply: oneshot::Sender<Result<(DecisionHandle, candle_core::DType)>>,
|
||||
},
|
||||
/// Remove a decision model from the slab and drop it on the worker
|
||||
/// thread. Mirrors `DropArch` / `DropImage`.
|
||||
DropDecision {
|
||||
handle: DecisionHandle,
|
||||
reply: oneshot::Sender<()>,
|
||||
},
|
||||
/// Score one batch of assembled questions (#336). One job per
|
||||
/// request: the whole batch is a single encoder forward, and the
|
||||
/// reply is CPU-side scores — no device tensor crosses back.
|
||||
Decide {
|
||||
handle: DecisionHandle,
|
||||
batch: crate::harness::arch::laya::DecisionBatch,
|
||||
reply: oneshot::Sender<Result<Vec<crate::harness::arch::laya::DecisionRow>>>,
|
||||
},
|
||||
/// Tell the worker to break its dispatch loop and exit. Any jobs
|
||||
/// queued after this in the channel reply `Err` to their oneshot
|
||||
/// senders (the senders are dropped on the worker's exit, which
|
||||
|
||||
@@ -51,7 +51,7 @@ use tokio::sync::oneshot;
|
||||
|
||||
#[cfg(feature = "cuda")]
|
||||
pub use jobs::TpHandle;
|
||||
pub use jobs::{ArchHandle, DenseLoad, ImageHandle, Job, KvSnapshotId};
|
||||
pub use jobs::{ArchHandle, DecisionHandle, DenseLoad, ImageHandle, Job, KvSnapshotId};
|
||||
|
||||
/// Errors returned by `DeviceWorkerHandle` submit methods.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
@@ -337,6 +337,85 @@ impl DeviceWorkerHandle {
|
||||
}
|
||||
}
|
||||
|
||||
/// Load one Laya checkpoint directory (#336) on the worker thread.
|
||||
/// Returns the handle and the dtype the weights were loaded in.
|
||||
pub async fn load_decision(
|
||||
&self,
|
||||
dir: std::path::PathBuf,
|
||||
model_id: String,
|
||||
dtype: candle_core::DType,
|
||||
) -> Result<(DecisionHandle, candle_core::DType), WorkerError> {
|
||||
if self.poisoned.load(Ordering::Acquire) {
|
||||
return Err(WorkerError::Poisoned {
|
||||
device_index: self.device_index,
|
||||
});
|
||||
}
|
||||
let (reply_tx, reply_rx) = oneshot::channel();
|
||||
self.tx
|
||||
.send(Job::LoadDecision {
|
||||
dir,
|
||||
model_id,
|
||||
dtype,
|
||||
reply: reply_tx,
|
||||
})
|
||||
.map_err(|_| WorkerError::Gone {
|
||||
device_index: self.device_index,
|
||||
})?;
|
||||
match reply_rx.await {
|
||||
Ok(result) => result.map_err(WorkerError::from),
|
||||
Err(_) => Err(WorkerError::Gone {
|
||||
device_index: self.device_index,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Drop the decision model for `handle` on the worker thread. Same
|
||||
/// poison-tolerant contract as [`Self::drop_arch`].
|
||||
pub async fn drop_decision(&self, handle: DecisionHandle) -> Result<(), WorkerError> {
|
||||
let (reply_tx, reply_rx) = oneshot::channel();
|
||||
self.tx
|
||||
.send(Job::DropDecision {
|
||||
handle,
|
||||
reply: reply_tx,
|
||||
})
|
||||
.map_err(|_| WorkerError::Gone {
|
||||
device_index: self.device_index,
|
||||
})?;
|
||||
reply_rx.await.map_err(|_| WorkerError::Gone {
|
||||
device_index: self.device_index,
|
||||
})
|
||||
}
|
||||
|
||||
/// Score one batch of assembled questions (#336). The reply is
|
||||
/// CPU-side scores; no device tensor crosses back.
|
||||
pub async fn decide(
|
||||
&self,
|
||||
handle: DecisionHandle,
|
||||
batch: crate::harness::arch::laya::DecisionBatch,
|
||||
) -> Result<Vec<crate::harness::arch::laya::DecisionRow>, WorkerError> {
|
||||
if self.poisoned.load(Ordering::Acquire) {
|
||||
return Err(WorkerError::Poisoned {
|
||||
device_index: self.device_index,
|
||||
});
|
||||
}
|
||||
let (reply_tx, reply_rx) = oneshot::channel();
|
||||
self.tx
|
||||
.send(Job::Decide {
|
||||
handle,
|
||||
batch,
|
||||
reply: reply_tx,
|
||||
})
|
||||
.map_err(|_| WorkerError::Gone {
|
||||
device_index: self.device_index,
|
||||
})?;
|
||||
match reply_rx.await {
|
||||
Ok(result) => result.map_err(WorkerError::from),
|
||||
Err(_) => Err(WorkerError::Gone {
|
||||
device_index: self.device_index,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Tell the worker to drop the `ModelArch` for `handle` on the
|
||||
/// worker thread (so CUDA tensors release on the right context).
|
||||
/// Returns `Ok(())` even if the handle wasn't in the slab — Drop
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -6,10 +6,13 @@ pub mod candle;
|
||||
pub mod capability;
|
||||
pub mod chat_template;
|
||||
pub mod context_limit;
|
||||
pub mod decision;
|
||||
pub mod decision_model;
|
||||
pub mod device_worker;
|
||||
pub mod engine;
|
||||
pub mod image;
|
||||
pub mod isq_cache;
|
||||
pub mod lang_detect;
|
||||
pub mod prefix_cache;
|
||||
pub mod preflight;
|
||||
pub mod preprocess;
|
||||
|
||||
@@ -54,6 +54,11 @@ pub enum SourceFormat {
|
||||
/// Takes precedence over the safetensors sniff because diffusers
|
||||
/// repos also contain `.safetensors` shards (in subdirs).
|
||||
Diffusers,
|
||||
/// A Laya decision-model checkpoint (#336): `rl_agent_config.json`
|
||||
/// plus `encoder/config.json` at the root, no top-level
|
||||
/// `config.json`. Served by the decision path — single device,
|
||||
/// unquantized. Takes precedence over the safetensors sniff.
|
||||
DecisionModel,
|
||||
}
|
||||
|
||||
/// Output of `preflight` for a load that can proceed. Carries the
|
||||
@@ -115,6 +120,12 @@ pub enum PreflightError {
|
||||
available: Vec<String>,
|
||||
nearest: Option<String>,
|
||||
},
|
||||
|
||||
/// The repo's format cannot be served with an option the spec
|
||||
/// asked for (e.g. `tensor_parallel` or `quant` on a decision
|
||||
/// model). `reason` says which and what to set instead.
|
||||
#[error("cannot load '{model_id}' as requested: {reason}")]
|
||||
Unsupported { model_id: String, reason: String },
|
||||
}
|
||||
|
||||
/// Run the placement check.
|
||||
@@ -230,6 +241,24 @@ pub async fn preflight(
|
||||
// The architecture compatibility check stays where it is —
|
||||
// `check_dense_config_supported` runs once `config.json` is
|
||||
// on disk, since it needs the parsed JSON.
|
||||
// Decision models (#336) load unquantized on one device.
|
||||
(SourceFormat::DecisionModel, tp, _) if tp > 1 => Err(PreflightError::Unsupported {
|
||||
model_id: source_id.to_string(),
|
||||
reason: format!(
|
||||
"decision models load on a single device; tensor_parallel={tp} is not supported"
|
||||
),
|
||||
}),
|
||||
(SourceFormat::DecisionModel, _, Some(q)) => Err(PreflightError::Unsupported {
|
||||
model_id: source_id.to_string(),
|
||||
reason: format!("decision models load unquantized; remove quant={q:?}"),
|
||||
}),
|
||||
(SourceFormat::DecisionModel, _, None) => Ok(PlacementPlan {
|
||||
model_id: source_id.to_string(),
|
||||
format: format.clone(),
|
||||
tp_size,
|
||||
picked_quant_file: None,
|
||||
}),
|
||||
|
||||
(SourceFormat::DenseSafetensors { .. } | SourceFormat::Mixed { .. }, _, _) => {
|
||||
Ok(PlacementPlan {
|
||||
model_id: source_id.to_string(),
|
||||
@@ -295,6 +324,9 @@ pub fn classify(filenames: &[&str]) -> SourceFormat {
|
||||
if crate::harness::image::is_diffusers_layout(filenames) {
|
||||
return SourceFormat::Diffusers;
|
||||
}
|
||||
if is_decision_layout(filenames) {
|
||||
return SourceFormat::DecisionModel;
|
||||
}
|
||||
let mut gguf_quants: Vec<String> = filenames
|
||||
.iter()
|
||||
.filter(|f| f.to_lowercase().ends_with(".gguf"))
|
||||
@@ -318,6 +350,16 @@ pub fn classify(filenames: &[&str]) -> SourceFormat {
|
||||
}
|
||||
}
|
||||
|
||||
/// A Laya decision checkpoint at the repo root: its decision config and
|
||||
/// its encoder config, with no top-level `config.json` (a causal LM
|
||||
/// that happened to ship a file with the same name still classifies as
|
||||
/// a text model).
|
||||
pub fn is_decision_layout(filenames: &[&str]) -> bool {
|
||||
filenames.contains(&"rl_agent_config.json")
|
||||
&& filenames.contains(&"encoder/config.json")
|
||||
&& !filenames.contains(&"config.json")
|
||||
}
|
||||
|
||||
/// Mirror of the quant-matching logic in `candle.rs::resolve_files` so
|
||||
/// preflight picks the same file the downstream loader would. Empty
|
||||
/// quant returns the first `.gguf` (any quant). Lowercased substring
|
||||
@@ -651,6 +693,20 @@ mod tests {
|
||||
tp_size,
|
||||
picked_quant_file: None,
|
||||
}),
|
||||
(SourceFormat::DecisionModel, tp, _) if tp > 1 => Err(PreflightError::Unsupported {
|
||||
model_id: source_id.to_string(),
|
||||
reason: format!("decision models load on a single device; tensor_parallel={tp}"),
|
||||
}),
|
||||
(SourceFormat::DecisionModel, _, Some(q)) => Err(PreflightError::Unsupported {
|
||||
model_id: source_id.to_string(),
|
||||
reason: format!("decision models load unquantized; remove quant={q:?}"),
|
||||
}),
|
||||
(SourceFormat::DecisionModel, _, None) => Ok(PlacementPlan {
|
||||
model_id: source_id.to_string(),
|
||||
format: format.clone(),
|
||||
tp_size,
|
||||
picked_quant_file: None,
|
||||
}),
|
||||
(SourceFormat::DenseSafetensors { .. } | SourceFormat::Mixed { .. }, _, _) => {
|
||||
Ok(PlacementPlan {
|
||||
model_id: source_id.to_string(),
|
||||
@@ -674,6 +730,66 @@ mod tests {
|
||||
assert_eq!(classify(&files), SourceFormat::Diffusers);
|
||||
}
|
||||
|
||||
/// The released `convaiinnovations/laya` listing: the English
|
||||
/// checkpoint at the root, two more in subfolders, and no root
|
||||
/// `config.json`.
|
||||
const LAYA_FILES: [&str; 12] = [
|
||||
"README.md",
|
||||
"rl_agent_config.json",
|
||||
"encoder/config.json",
|
||||
"model.safetensors",
|
||||
"tokenizer/tokenizer.json",
|
||||
"tokenizer/tokenizer_config.json",
|
||||
"multilingual/rl_agent_config.json",
|
||||
"multilingual/encoder/config.json",
|
||||
"multilingual/model.safetensors",
|
||||
"typed-decisions/rl_agent_config.json",
|
||||
"typed-decisions/encoder/config.json",
|
||||
"typed-decisions/model.safetensors",
|
||||
];
|
||||
|
||||
#[test]
|
||||
fn classify_laya_layout() {
|
||||
assert_eq!(classify(&LAYA_FILES), SourceFormat::DecisionModel);
|
||||
// A causal LM that also ships an rl config is still a text model.
|
||||
let lm = [
|
||||
"config.json",
|
||||
"rl_agent_config.json",
|
||||
"encoder/config.json",
|
||||
"model.safetensors",
|
||||
];
|
||||
assert!(matches!(
|
||||
classify(&lm),
|
||||
SourceFormat::DenseSafetensors { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn feasibility_decision_model() {
|
||||
let fmt = classify(&LAYA_FILES);
|
||||
let plan = decide(
|
||||
&spec("convaiinnovations/laya", None, None),
|
||||
&fmt,
|
||||
&LAYA_FILES,
|
||||
)
|
||||
.expect("single-device unquantized decision load should pass");
|
||||
assert_eq!(plan.format, SourceFormat::DecisionModel);
|
||||
let err = decide(
|
||||
&spec("convaiinnovations/laya", Some(2), None),
|
||||
&fmt,
|
||||
&LAYA_FILES,
|
||||
)
|
||||
.expect_err("tp>1 must be refused");
|
||||
assert!(matches!(err, PreflightError::Unsupported { .. }));
|
||||
let err = decide(
|
||||
&spec("convaiinnovations/laya", None, Some("q8_0")),
|
||||
&fmt,
|
||||
&LAYA_FILES,
|
||||
)
|
||||
.expect_err("quant must be refused");
|
||||
assert!(matches!(err, PreflightError::Unsupported { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn feasibility_diffusers_single_gpu_ok() {
|
||||
let files = ["model_index.json", "transformer/config.json"];
|
||||
|
||||
+46824
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,52 @@
|
||||
{
|
||||
"attention_bias": false,
|
||||
"attention_dropout": 0.0,
|
||||
"bos_token_id": 1,
|
||||
"classifier_activation": "gelu",
|
||||
"classifier_bias": false,
|
||||
"classifier_dropout": 0.0,
|
||||
"classifier_pooling": "cls",
|
||||
"cls_token_id": 1,
|
||||
"decoder_bias": true,
|
||||
"deterministic_flash_attn": false,
|
||||
"embedding_dropout": 0.0,
|
||||
"eos_token_id": 2,
|
||||
"global_attn_every_n_layers": 3,
|
||||
"hidden_activation": "gelu",
|
||||
"hidden_size": 32,
|
||||
"initializer_cutoff_factor": 2.0,
|
||||
"initializer_range": 0.3,
|
||||
"intermediate_size": 48,
|
||||
"layer_types": [
|
||||
"full_attention",
|
||||
"sliding_attention",
|
||||
"sliding_attention",
|
||||
"full_attention"
|
||||
],
|
||||
"local_attention": 8,
|
||||
"max_position_embeddings": 128,
|
||||
"mlp_bias": false,
|
||||
"mlp_dropout": 0.0,
|
||||
"model_type": "modernbert",
|
||||
"norm_bias": false,
|
||||
"norm_eps": 1e-05,
|
||||
"num_attention_heads": 2,
|
||||
"num_hidden_layers": 4,
|
||||
"pad_token_id": 0,
|
||||
"rope_parameters": {
|
||||
"full_attention": {
|
||||
"rope_theta": 1000.0,
|
||||
"rope_type": "default"
|
||||
},
|
||||
"sliding_attention": {
|
||||
"rope_theta": 300.0,
|
||||
"rope_type": "default"
|
||||
}
|
||||
},
|
||||
"sep_token_id": 2,
|
||||
"sparse_pred_ignore_index": -100,
|
||||
"sparse_prediction": false,
|
||||
"tie_word_embeddings": true,
|
||||
"transformers_version": "5.17.0",
|
||||
"vocab_size": 97
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
{
|
||||
"seed": 7,
|
||||
"seqs": [
|
||||
{
|
||||
"input_ids": [
|
||||
1,
|
||||
93,
|
||||
82,
|
||||
87,
|
||||
51,
|
||||
46,
|
||||
89,
|
||||
6,
|
||||
68,
|
||||
94,
|
||||
86,
|
||||
43,
|
||||
23,
|
||||
70,
|
||||
42,
|
||||
60,
|
||||
35,
|
||||
3,
|
||||
86,
|
||||
7,
|
||||
20,
|
||||
18,
|
||||
90,
|
||||
84,
|
||||
62,
|
||||
30,
|
||||
42,
|
||||
65,
|
||||
34,
|
||||
37,
|
||||
30,
|
||||
67,
|
||||
2,
|
||||
34,
|
||||
77,
|
||||
96,
|
||||
24,
|
||||
89,
|
||||
19,
|
||||
33
|
||||
],
|
||||
"markers": [
|
||||
3,
|
||||
9,
|
||||
17
|
||||
],
|
||||
"qtype": 0,
|
||||
"logits": [
|
||||
-0.4152399003505707,
|
||||
-0.5505563020706177,
|
||||
-0.600245475769043
|
||||
],
|
||||
"act": [
|
||||
0.4727683663368225,
|
||||
0.5272316336631775
|
||||
]
|
||||
},
|
||||
{
|
||||
"input_ids": [
|
||||
4,
|
||||
5,
|
||||
43,
|
||||
95,
|
||||
77,
|
||||
35,
|
||||
61,
|
||||
13,
|
||||
42,
|
||||
24,
|
||||
21,
|
||||
7,
|
||||
11,
|
||||
38,
|
||||
80,
|
||||
54,
|
||||
4,
|
||||
29,
|
||||
5,
|
||||
80,
|
||||
54,
|
||||
78,
|
||||
46
|
||||
],
|
||||
"markers": [
|
||||
4,
|
||||
11
|
||||
],
|
||||
"qtype": 1,
|
||||
"logits": [
|
||||
-0.1313231736421585,
|
||||
-0.3265465199947357
|
||||
],
|
||||
"act": [
|
||||
0.40568628907203674,
|
||||
0.5943136811256409
|
||||
]
|
||||
},
|
||||
{
|
||||
"input_ids": [
|
||||
67,
|
||||
40,
|
||||
83,
|
||||
41,
|
||||
77,
|
||||
96,
|
||||
58,
|
||||
39,
|
||||
76
|
||||
],
|
||||
"markers": [
|
||||
2
|
||||
],
|
||||
"qtype": 2,
|
||||
"logits": [
|
||||
0.534652590751648
|
||||
],
|
||||
"act": [
|
||||
0.4226877987384796,
|
||||
0.5773122310638428
|
||||
]
|
||||
},
|
||||
{
|
||||
"input_ids": [
|
||||
96,
|
||||
32,
|
||||
60,
|
||||
44,
|
||||
35,
|
||||
47,
|
||||
45,
|
||||
51,
|
||||
60,
|
||||
51,
|
||||
73,
|
||||
50,
|
||||
25,
|
||||
79,
|
||||
74,
|
||||
24,
|
||||
83,
|
||||
7,
|
||||
9,
|
||||
91,
|
||||
21,
|
||||
87,
|
||||
29,
|
||||
25,
|
||||
74,
|
||||
18,
|
||||
31,
|
||||
46,
|
||||
46,
|
||||
30,
|
||||
52
|
||||
],
|
||||
"markers": [
|
||||
5,
|
||||
7,
|
||||
12,
|
||||
20,
|
||||
26
|
||||
],
|
||||
"qtype": 0,
|
||||
"logits": [
|
||||
-0.4001395106315613,
|
||||
-0.5266168117523193,
|
||||
-0.44161877036094666,
|
||||
-0.43973028659820557,
|
||||
-0.5682165622711182
|
||||
],
|
||||
"act": [
|
||||
0.46282079815864563,
|
||||
0.537179172039032
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
Binary file not shown.
@@ -0,0 +1,16 @@
|
||||
{
|
||||
"encoder": "tiny-modernbert",
|
||||
"head_layers": 2,
|
||||
"max_len": 64,
|
||||
"head_max_len": 24,
|
||||
"act_costs": {
|
||||
"escalate": 0.5
|
||||
},
|
||||
"amp_dtype": "bf16",
|
||||
"temperature": [
|
||||
1.0,
|
||||
1.0,
|
||||
1.0
|
||||
],
|
||||
"temperature_by_options": {}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
{"version": "1.0", "truncation": null, "padding": null, "added_tokens": [], "normalizer": null, "pre_tokenizer": {"type": "Whitespace"}, "post_processor": null, "decoder": null, "model": {"type": "WordLevel", "vocab": {"[PAD]": 0, "[UNK]": 1, "[CLS]": 2, "[SEP]": 3, "[MASK]": 4, "a": 5, "b": 6}, "unk_token": "[UNK]"}}
|
||||
@@ -0,0 +1,7 @@
|
||||
{
|
||||
"cls_token": "[CLS]",
|
||||
"sep_token": "[SEP]",
|
||||
"mask_token": "[MASK]",
|
||||
"pad_token": "[PAD]",
|
||||
"unk_token": "[UNK]"
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,14 @@
|
||||
{
|
||||
"clean_up_tokenization_spaces": true,
|
||||
"cls_token": "[CLS]",
|
||||
"mask_token": "[MASK]",
|
||||
"model_input_names": [
|
||||
"input_ids",
|
||||
"attention_mask"
|
||||
],
|
||||
"model_max_length": 8192,
|
||||
"pad_token": "[PAD]",
|
||||
"sep_token": "[SEP]",
|
||||
"tokenizer_class": "PreTrainedTokenizerFast",
|
||||
"unk_token": "[UNK]"
|
||||
}
|
||||
@@ -158,6 +158,7 @@ fn preflight_kind(err: &PreflightError) -> &'static str {
|
||||
PreflightError::EmptyRepo { .. } => "empty_repo",
|
||||
PreflightError::TpRequiresSafetensors { .. } => "tp_requires_safetensors",
|
||||
PreflightError::QuantNotFound { .. } => "quant_not_found",
|
||||
PreflightError::Unsupported { .. } => "unsupported_load_option",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -820,3 +820,55 @@ async fn null_assistant_content_is_not_a_body_rejection() {
|
||||
assert_ne!(resp.status(), 422, "null content must not be a body error");
|
||||
assert_ne!(resp.status(), 400, "null content must not be a body error");
|
||||
}
|
||||
|
||||
// ── /v1/audio/generations (#382) ─────────────────────────────────
|
||||
//
|
||||
// The neuron half of the async audio job API: one synchronous render per
|
||||
// call. The envelope and error mapping are tested here; rendering needs an
|
||||
// audio engine (#383) and is validated on the fleet.
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_audio_no_candle_harness_503() {
|
||||
let url = spawn_neuron(fake_discovery()).await;
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("{url}/v1/audio/generations"))
|
||||
.json(&json!({"model": "m", "prompt": "lofi"}))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 503);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_audio_model_not_loaded_404() {
|
||||
let url = spawn_neuron_with_candle(fake_discovery()).await;
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("{url}/v1/audio/generations"))
|
||||
.json(&json!({"model": "ACE-Step/Ace-Step1.5", "prompt": "lofi"}))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 404);
|
||||
let body: serde_json::Value = resp.json().await.unwrap();
|
||||
assert_eq!(body["error"]["code"], "model_not_found");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_audio_invalid_params_400() {
|
||||
let url = spawn_neuron_with_candle(fake_discovery()).await;
|
||||
for bad in [
|
||||
json!({"model": "m", "prompt": ""}),
|
||||
json!({"model": "m", "prompt": "lofi", "duration": -3}),
|
||||
json!({"model": "m", "prompt": "lofi", "response_format": "ogg"}),
|
||||
] {
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("{url}/v1/audio/generations"))
|
||||
.json(&bad)
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 400, "{bad}");
|
||||
let body: serde_json::Value = resp.json().await.unwrap();
|
||||
assert_eq!(body["error"]["code"], "invalid_audio_params", "{bad}");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,6 +10,9 @@ Restart=on-failure
|
||||
RestartSec=5
|
||||
User=cortex
|
||||
Group=cortex
|
||||
# /var/lib/cortex, owned by the service user: async audio jobs (#382)
|
||||
# keep their records and rendered audio here.
|
||||
StateDirectory=cortex
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
|
||||
+4
-6
@@ -118,12 +118,10 @@ directly useful for P2 (#11 prefix KV caching, #23 chunked prefill,
|
||||
## Misc operational facts
|
||||
|
||||
- Fleet: beast 2×5090 (27B Q6K TP=2), benjy 4090 (8B), quadbrat 3060
|
||||
(1.7B); gateway on hanzalova:31313; all reachable by ssh as the
|
||||
operator (passwordless sudo) — used for sudoers sync, live probes,
|
||||
and killing wedged processes during P1.
|
||||
- The tea CLI token (`~/.config/tea/config.yml`) authenticates against
|
||||
git.lair.cafe for everything the gitea-mcp tools don't cover
|
||||
(milestones, labels, pin, merge).
|
||||
(1.7B); gateway on hanzalova:31313. Operator access to the hosts
|
||||
(used for sudoers sync, live probes, and killing wedged processes
|
||||
during P1) is described in `doc/infra/operator-access.md` in the
|
||||
private helexa/security repo.
|
||||
- Anthropic streaming is now production-real: Claude Code can be
|
||||
pointed at the gateway. Real-traffic feedback from that is useful
|
||||
input for P2 priorities (especially tool-use-heavy streams).
|
||||
|
||||
@@ -58,6 +58,6 @@ timeout_secs = 30
|
||||
# whether a mail also goes out. An unknown provider is a hard startup
|
||||
# failure rather than silently swallowed notifications.
|
||||
provider = "smtp"
|
||||
smtp_url = "smtps://user%40l4ir.net:CHANGEME@mail.l4ir.net:465"
|
||||
smtp_url = "smtps://mailer%40example.org:CHANGEME@smtp.example.org:465"
|
||||
from_addr = "helexa <no-reply@helexa.ai>"
|
||||
notify_to = "angels@helexa.ai"
|
||||
|
||||
+8
-1
@@ -26,8 +26,15 @@ region = "eu-central"
|
||||
|
||||
# Product tiers → real model ids (#166). The SPA sends the tier names.
|
||||
[aliases]
|
||||
"helexa/small" = "Qwen/Qwen3-8B"
|
||||
# helexa/small is the tier anonymous web chat uses. It points at the 1.7B on
|
||||
# quadbrat, the same target as the catalogue alias in models.toml: the 8B
|
||||
# needs benjy, and while benjy serves nothing an alias aimed there made
|
||||
# anonymous chat answer "no operator serves model 'Qwen/Qwen3-8B'".
|
||||
"helexa/small" = "Qwen/Qwen3-1.7B"
|
||||
"helexa/balanced" = "Qwen/Qwen3.8-27B"
|
||||
# System-One decisions (#333). Also the model /v1/systemone falls back
|
||||
# to when a request names none or names a service no operator serves.
|
||||
"helexa/one" = "convaiinnovations/laya"
|
||||
|
||||
# hanzalova operator cortex, plaintext over the WireGuard mesh (the mesh
|
||||
# is the trust boundary; no TLS pin needed for http://).
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user