diff --git a/README.md b/README.md index e32c43b..8bea1dd 100644 --- a/README.md +++ b/README.md @@ -557,6 +557,29 @@ cap concurrency, request rate, or both: arrival instead of failing (NGINX `limit_req` without `nodelay`). Waits share the 256-request queue bound and the five-minute maximum; a request whose wait would exceed it still fails fast with `429` and `retry-after`. +- `requestClasses`: optional non-empty array of `standard`, `interactive`, or + `foreground`. Only those classes use this limiter's rate budget, concurrency + slots, and waiting room; omitted means all classes. Selected classes share + one bucket. Other applicable model/alias limits and normal runtime capacity + still apply. + +For example, bind a cloud model's cost throttle to background work while chat +keeps its normal priority and capacity: + +```json +"rateLimit": { + "rate": "4.5/m", + "burst": 0, + "queue": true, + "requestClasses": ["standard"] +} +``` + +Set this on the concrete cloud model to cover every alias that resolves to it. +The limit follows the resolved model when routes change. A caller marked +`standard`, including asynchronous inbox work, remains paced. Foreground and +interactive requests bypass this particular limit without spending its budget. +The rate counts model requests, not tokens or dollars. The limiter is the standard token-bucket used by NGINX `limit_req` and Envoy's local rate limiter: tokens refill continuously at one per `rateMs`, requests @@ -571,7 +594,7 @@ same limited model draw from the same budget, and a direct model ID request is limited identically. Rate-limit rejections never poison target failure backoff or trigger model failover; the 429 surfaces directly with `retry-after`. -Live state (active, queued, next-allowed time per limiter) is exposed as +Live state (selected request classes, active, queued, next-allowed time per limiter) is exposed as `rateLimits` in `GET /gateway/routing`. Config reloads keep existing buckets for unchanged ids and prune removed ones. diff --git a/docs/request-scheduling.md b/docs/request-scheduling.md index aa5f854..0b4d545 100644 --- a/docs/request-scheduling.md +++ b/docs/request-scheduling.md @@ -106,6 +106,16 @@ the next token. Completed requests still consume their budget. Set an individual is exhausted. Limits and scheduling settings can be hot-reloaded without canceling active requests. +An optional `rateLimit.requestClasses` array scopes a model or alias limit to +specific existing priorities. For example, `requestClasses: ["standard"]` with +`rate: "4.5/m"`, `burst: 0`, and `queue: true` spaces background calls without +making foreground or interactive callers wait for that budget. The selected +classes share the limiter; omitted means all three classes. Exempt requests +neither reserve its concurrency slots nor spend its rate tokens. Every other +applicable limit and runtime-capacity rule still applies, and priority within +each compute group is unchanged. Config reloads preserve the spent bucket while +updating which classes use it, including already waiting requests. + Adaptive requests also wait through a needed cold start instead of receiving the earlier startup retry response. Model-loading admission checks priority before changing residency and yields before its next mutation when higher diff --git a/src/rate-limit.mjs b/src/rate-limit.mjs index 070e52f..7be9436 100644 --- a/src/rate-limit.mjs +++ b/src/rate-limit.mjs @@ -21,6 +21,8 @@ // the next conforming arrival instead (nginx `limit_req` with a burst queue // and no `nodelay`), bounded by the same maximum wait. +import { REQUEST_CLASSES, normalizeSchedulingClass } from './request-class.mjs'; + const MAX_WAIT_MS = 5 * 60 * 1000; export const RATE_LIMIT_QUEUE_CAPACITY = 256; @@ -54,6 +56,17 @@ export function normalizeRateLimit(value) { throw new Error('rateLimit must be a string like "20/m" or an object'); } const settings = value.rateLimit ?? value; + const selection = {}; + if (Object.hasOwn(settings, 'requestClasses')) { + if ( + !Array.isArray(settings.requestClasses) || + !settings.requestClasses.length || + settings.requestClasses.some((requestClass) => !REQUEST_CLASSES.includes(requestClass)) + ) { + throw new Error('rateLimit.requestClasses must be a non-empty array of standard, interactive, or foreground'); + } + selection.requestClasses = REQUEST_CLASSES.filter((requestClass) => settings.requestClasses.includes(requestClass)); + } if (settings.queue != null && typeof settings.queue !== 'boolean') { throw new Error(`rateLimit.queue must be a boolean: ${JSON.stringify(settings.queue)}`); } @@ -77,7 +90,7 @@ export function normalizeRateLimit(value) { throw new Error('rateLimit.burst must be a non-negative integer'); } burst = settings.burst ?? 0; - return { maxConcurrent, burst, rateMs, ...queue }; + return { maxConcurrent, burst, rateMs, ...queue, ...selection }; } const rate = settings.rate ?? settings.requestsPerMinute ?? settings.requests; let rateCount = null; @@ -132,7 +145,8 @@ export function normalizeRateLimit(value) { // An explicit burst overrides this. burst: rateMs == null ? 0 : (burst ?? Math.max(0, Math.floor(rateCount) - 1)), rateMs, - ...queue + ...queue, + ...selection }; } @@ -308,6 +322,13 @@ export function createGcra({ rateMs, burst }) { */ export function createRateLimitRegistry(limits = {}) { const limiters = new Map(); + const matchingLimiter = (id, requestClass) => { + const entry = limiters.get(id); + return entry && + (!entry.settings.requestClasses || entry.settings.requestClasses.includes(normalizeSchedulingClass(requestClass))) + ? entry + : null; + }; const listeners = new Set(); let wakeTimer = null; let wakeAt = Infinity; @@ -370,7 +391,8 @@ export function createRateLimitRegistry(limits = {}) { entry.settings.maxConcurrent === settings.maxConcurrent && entry.settings.rateMs === settings.rateMs && entry.settings.burst === settings.burst && - entry.settings.queue === settings.queue + entry.settings.queue === settings.queue && + JSON.stringify(entry.settings.requestClasses) === JSON.stringify(settings.requestClasses) ) { continue; } @@ -402,8 +424,10 @@ export function createRateLimitRegistry(limits = {}) { clearTimeout(wakeTimer); listeners.clear(); }, - tryAcquireChain(ids) { - const entries = [...new Set(ids)].map((id) => [id, limiters.get(id)]).filter(([, entry]) => entry); + tryAcquireChain(ids, { requestClass = 'standard' } = {}) { + const entries = [...new Set(ids)] + .map((id) => [id, matchingLimiter(id, requestClass)]) + .filter(([, entry]) => entry); let waitMs = 0; for (const [id, entry] of entries) { const delay = entry.gcra?.nextAllowedInMs() ?? 0; @@ -435,7 +459,7 @@ export function createRateLimitRegistry(limits = {}) { releases.push(release); } try { - consumeRateBudget({ limiter: (id) => limiters.get(id) }, ids); + consumeRateBudget({ limiter: (id) => matchingLimiter(id, requestClass) }, ids, { requestClass }); } catch (error) { releases.forEach((fn) => fn()); if (error.name !== 'RateBudgetExhaustedError') throw error; @@ -450,8 +474,8 @@ export function createRateLimitRegistry(limits = {}) { notify(); }; }, - limiter(id) { - return limiters.get(id) ?? null; + limiter(id, { requestClass } = {}) { + return requestClass === undefined ? (limiters.get(id) ?? null) : matchingLimiter(id, requestClass); }, get size() { return limiters.size; @@ -463,6 +487,7 @@ export function createRateLimitRegistry(limits = {}) { rateMs: entry.settings.rateMs ?? null, burst: entry.settings.burst ?? null, queue: entry.settings.queue === true, + requestClasses: entry.settings.requestClasses ?? [...REQUEST_CLASSES], active: entry.semaphore?.active ?? 0, queued: (entry.semaphore?.queued ?? 0) + entry.rateQueued, rateLimited: entry.gcra ? entry.gcra.nextAllowedInMs(now) > 0 : false, @@ -476,13 +501,37 @@ export function createRateLimitRegistry(limits = {}) { export async function acquireRateLimitSlot( registry, id, - { signal = null, timeoutMs = MAX_WAIT_MS, rateBudget = true } = {} + { signal = null, timeoutMs = MAX_WAIT_MS, rateBudget = true, requestClass = 'standard' } = {} ) { - const entry = registry.limiter(id); + const entry = registry.limiter(id, { requestClass }); if (!entry) return () => {}; // Concurrency waits in a bounded FIFO queue, exactly like runtime slots. signal?.throwIfAborted?.(); - const releaseSemaphore = entry.semaphore ? await entry.semaphore.acquire(signal, { timeoutMs }) : () => {}; + let releaseSemaphore = () => {}; + if (entry.semaphore) { + // A class-selection reload can exempt a request already in the legacy + // semaphore queue. Remove that waiter without interrupting admitted work. + const exemption = new AbortController(); + const unsubscribe = registry.subscribe?.(() => { + if (!registry.limiter(id, { requestClass })) exemption.abort(); + }); + try { + releaseSemaphore = await entry.semaphore.acquire( + signal ? AbortSignal.any([signal, exemption.signal]) : exemption.signal, + { timeoutMs } + ); + } catch (error) { + signal?.throwIfAborted?.(); + if (!exemption.signal.aborted) throw error; + return () => {}; + } finally { + unsubscribe?.(); + } + if (!registry.limiter(id, { requestClass })) { + releaseSemaphore(); + return () => {}; + } + } try { signal?.throwIfAborted?.(); if (rateBudget && entry.gcra) { @@ -516,10 +565,10 @@ export async function acquireRateLimitSlot( * atomic: if any limiter rejects, every token already taken is refunded, so a * request that never reached a model costs no budget on the scopes it passed. */ -export function consumeRateBudget(registry, ids) { +export function consumeRateBudget(registry, ids, { requestClass = 'standard' } = {}) { const consumed = []; for (const id of ids) { - const gcra = registry.limiter(id)?.gcra; + const gcra = registry.limiter(id, { requestClass })?.gcra; if (!gcra) continue; const decision = gcra.tryAcquire(); if (decision.allowed) { @@ -546,15 +595,16 @@ export function consumeRateBudget(registry, ids) { export async function awaitRateBudget( registry, ids, - { signal = null, started = Date.now(), capacity = RATE_LIMIT_QUEUE_CAPACITY } = {} + { signal = null, started = Date.now(), capacity = RATE_LIMIT_QUEUE_CAPACITY, requestClass = 'standard' } = {} ) { const maxQueued = Number.isSafeInteger(capacity) && capacity >= 0 ? capacity : RATE_LIMIT_QUEUE_CAPACITY; for (;;) { signal?.throwIfAborted?.(); try { - return consumeRateBudget(registry, ids); + return consumeRateBudget(registry, ids, { requestClass }); } catch (error) { - const entry = error?.name === 'RateBudgetExhaustedError' ? registry.limiter(error.limiterId) : null; + const entry = + error?.name === 'RateBudgetExhaustedError' ? registry.limiter(error.limiterId, { requestClass }) : null; if (!entry?.settings.queue || Date.now() + error.retryAfterMs - started > MAX_WAIT_MS) throw error; // Reserve before yielding: a concurrent caller can have passed the // server's precheck while this caller was acquiring semaphore slots. @@ -563,7 +613,7 @@ export async function awaitRateBudget( } entry.rateQueued += 1; try { - await delay(error.retryAfterMs, signal); + await delay(error.retryAfterMs, signal, registry); } finally { entry.rateQueued -= 1; } @@ -571,16 +621,25 @@ export async function awaitRateBudget( } } -function delay(ms, signal) { +function delay(ms, signal, registry) { return new Promise((resolve, reject) => { - const onAbort = () => { + let timer; + let unsubscribe; + const cleanup = () => { clearTimeout(timer); + unsubscribe?.(); + signal?.removeEventListener?.('abort', onAbort); + }; + const onAbort = () => { + cleanup(); reject(signal.reason ?? abortError()); }; - const timer = setTimeout(() => { - signal?.removeEventListener?.('abort', onAbort); + const finish = () => { + cleanup(); resolve(); - }, ms); + }; + timer = setTimeout(finish, ms); + unsubscribe = registry?.subscribe?.(finish); signal?.addEventListener?.('abort', onAbort, { once: true }); }); } diff --git a/src/request-class.mjs b/src/request-class.mjs new file mode 100644 index 0000000..051d1b3 --- /dev/null +++ b/src/request-class.mjs @@ -0,0 +1,5 @@ +export const REQUEST_CLASSES = ['standard', 'interactive', 'foreground']; + +export function normalizeSchedulingClass(value) { + return typeof value === 'string' && REQUEST_CLASSES.includes(value) ? value : 'standard'; +} diff --git a/src/request-scheduler.mjs b/src/request-scheduler.mjs index d715b63..42f2737 100644 --- a/src/request-scheduler.mjs +++ b/src/request-scheduler.mjs @@ -1,4 +1,7 @@ import { currentNodeId } from './cluster.mjs'; +import { REQUEST_CLASSES, normalizeSchedulingClass } from './request-class.mjs'; + +export { REQUEST_CLASSES, normalizeSchedulingClass } from './request-class.mjs'; // Priority scheduling for connected gateway inference requests. // @@ -9,12 +12,6 @@ import { currentNodeId } from './cluster.mjs'; // back without a backlog burst. Nothing here pauses or resumes a GPU: an // accepted request keeps its stream and simply waits its turn. -export const REQUEST_CLASSES = ['standard', 'interactive', 'foreground']; - -export function normalizeSchedulingClass(value) { - return typeof value === 'string' && REQUEST_CLASSES.includes(value) ? value : 'standard'; -} - function classRank(requestClass) { return REQUEST_CLASSES.indexOf(normalizeSchedulingClass(requestClass)); } @@ -271,14 +268,16 @@ export class RequestScheduler { return standard.indexOf(lease) < this.capacity(state); } - higherDemandForScopes(lease, scopeIds) { + higherDemandForScopes(lease, scopeIds, appliesToClass = () => true) { if (this.mode !== 'adaptive' || !scopeIds.length) return false; const scopes = new Set(scopeIds); const rank = classRank(lease.requestClass); for (const state of this.groups.values()) { if ( [...state.leases, ...state.waiters].some( - (entry) => classRank(entry.requestClass) > rank && entry.scopeIds?.some((id) => scopes.has(id)) + (entry) => + classRank(entry.requestClass) > rank && + entry.scopeIds?.some((id) => scopes.has(id) && appliesToClass(id, entry.requestClass)) ) ) return true; diff --git a/src/server.mjs b/src/server.mjs index 12ff2f8..ae0c18f 100644 --- a/src/server.mjs +++ b/src/server.mjs @@ -577,6 +577,8 @@ const QUEUE_BACKPRESSURE_CODES = new Set([ 'RUNTIME_QUEUE_TIMEOUT', 'RUNTIME_QUEUE_FULL', 'INTERACTIVE_PRIORITY', + 'MODEL_RATE_LIMITED', + 'MODEL_RATE_LIMIT_QUEUE_FULL', 'REQUEST_QUEUE_FULL', 'REQUEST_SCHEDULING_TIMEOUT', 'REQUEST_SCHEDULER_CLOSED' @@ -604,9 +606,14 @@ function errorStatusCode(error) { export function shouldFailoverModelRequest(error, res = null) { if ( - ['INTERACTIVE_PRIORITY', 'REQUEST_QUEUE_FULL', 'REQUEST_SCHEDULING_TIMEOUT', 'REQUEST_SCHEDULER_CLOSED'].some( - (code) => code === error?.code || code === error?.upstreamCode - ) + [ + 'INTERACTIVE_PRIORITY', + 'MODEL_RATE_LIMITED', + 'MODEL_RATE_LIMIT_QUEUE_FULL', + 'REQUEST_QUEUE_FULL', + 'REQUEST_SCHEDULING_TIMEOUT', + 'REQUEST_SCHEDULER_CLOSED' + ].some((code) => code === error?.code || code === error?.upstreamCode) ) return false; if (res?.headersSent || res?.writableEnded || res?.destroyed) return false; @@ -3068,9 +3075,13 @@ export function createLloomServer( * and rate budget. A single limiter per scope id is shared across aliases so * a model reachable through two aliases contributes to one queue. */ - async function admitRateLimited(resolved, scopeIds, { signal = null, started = Date.now() } = {}) { + async function admitRateLimited( + resolved, + scopeIds, + { signal = null, started = Date.now(), requestClass = 'standard' } = {} + ) { if (!scopeIds.length) return () => {}; - const limited = scopeIds.filter((id) => rateLimitRegistry.limiter(id)); + const limited = scopeIds.filter((id) => rateLimitRegistry.limiter(id, { requestClass })); if (!limited.length) return () => {}; // FIFO fairness across scopes: acquire semaphore slots first (queued in // call order), then rate budget. Failover between candidates happens @@ -3081,11 +3092,12 @@ export function createLloomServer( }; try { for (const id of limited) { - const entry = rateLimitRegistry.limiter(id); + const entry = rateLimitRegistry.limiter(id, { requestClass }); + if (!entry) continue; if ((entry.semaphore?.queued ?? 0) + entry.rateQueued >= RATE_LIMIT_QUEUE_CAPACITY) { throw new ModelRateLimitError(resolved.requestedId, 5, { queueFull: true }); } - releases.push(await acquireRateLimitSlot(rateLimitRegistry, id, { signal, rateBudget: false })); + releases.push(await acquireRateLimitSlot(rateLimitRegistry, id, { signal, rateBudget: false, requestClass })); } // Consume rate budget only once every concurrency slot is held, and // atomically, so an inner rejection cannot burn the budget of the outer @@ -3093,7 +3105,8 @@ export function createLloomServer( await awaitRateBudget(rateLimitRegistry, limited, { signal, started, - capacity: RATE_LIMIT_QUEUE_CAPACITY + capacity: RATE_LIMIT_QUEUE_CAPACITY, + requestClass }); return releaseAll; } catch (error) { @@ -3171,13 +3184,14 @@ export function createLloomServer( requestScheduler.readyFor(schedulingLease) && !requestScheduler.higherDemandForScopes( schedulingLease, - rateLimitScopeIds.filter((id) => rateLimitRegistry.limiter(id)) + rateLimitScopeIds.filter((id) => rateLimitRegistry.limiter(id, { requestClass })), + (id, otherClass) => Boolean(rateLimitRegistry.limiter(id, { requestClass: otherClass })) ); const canAdmit = () => { if (!readyForAdmission()) return false; if (rateAdmitted) return true; try { - const release = rateLimitRegistry.tryAcquireChain(rateLimitScopeIds); + const release = rateLimitRegistry.tryAcquireChain(rateLimitScopeIds, { requestClass }); if (!release) return false; rateLimitRelease = release; rateAdmitted = true; @@ -3201,7 +3215,8 @@ export function createLloomServer( else rateLimitRelease = await admitRateLimited(resolved, rateLimitScopeIds, { signal: client.signal, - started + started, + requestClass }); } catch (error) { // A rejected admission never reaches the request lifecycle that would diff --git a/test/adaptive-priority.test.mjs b/test/adaptive-priority.test.mjs index 63e3305..2489027 100644 --- a/test/adaptive-priority.test.mjs +++ b/test/adaptive-priority.test.mjs @@ -309,6 +309,28 @@ test('foreground on local compute does not pause independent cloud work', { time assert.equal((await foreground).status, 200); }); +test( + 'standard-only pacing leaves a connected waiter queued while interactive and foreground bypass its budget', + { timeout: 10000 }, + async (t) => { + const f = await fixture(t, { + foregroundIdleMs: 0, + backgroundMaxConcurrency: 4, + modelRateLimit: { rate: '1/s', burst: 0, queue: true, requestClasses: ['standard'] } + }); + assert.equal((await f.request('initial')).status, 200); + const waiting = f.request('paced-standard'); + await new Promise((resolve) => setTimeout(resolve, 30)); + const chatStarted = Date.now(); + assert.equal((await f.request('interactive', { requestClass: 'interactive' })).status, 200); + assert.equal((await f.request('foreground', { requestClass: 'foreground' })).status, 200); + assert.ok(Date.now() - chatStarted < 500, 'chat does not wait for the one-second standard rate interval'); + assert.deepEqual(f.hits, ['initial', 'interactive', 'foreground']); + assert.equal((await waiting).status, 200); + assert.deepEqual(f.hits, ['initial', 'interactive', 'foreground', 'paced-standard']); + } +); + test('cancelled queued background never reaches upstream', { timeout: 10000 }, async (t) => { const f = await fixture(t); const foreground = f.request('hold-foreground', { requestClass: 'foreground' }); diff --git a/test/federated-priority.test.mjs b/test/federated-priority.test.mjs index 9464559..bc2f541 100644 --- a/test/federated-priority.test.mjs +++ b/test/federated-priority.test.mjs @@ -18,7 +18,10 @@ async function listen(server) { return `http://127.0.0.1:${server.address().port}`; } -async function fixture(t, { holdStream = false, providerStatus = 200, rawFallback = false, mode = 'interrupt' } = {}) { +async function fixture( + t, + { holdStream = false, providerStatus = 200, rawFallback = false, mode = 'interrupt', modelRateLimit } = {} +) { const hits = []; const fallbackHits = []; const streamResponses = []; @@ -91,7 +94,13 @@ async function fixture(t, { holdStream = false, providerStatus = 200, rawFallbac { id: 'chat', kind: 'chat' }, { id: 'speech', kind: 'audio_speech' }, { id: 'transcription', kind: 'audio_transcription' } - ].map((model) => ({ ...model, backend: 'provider', runtime: 'limited', upstreamModel: model.id })), + ].map((model) => ({ + ...model, + backend: 'provider', + runtime: 'limited', + upstreamModel: model.id, + ...(model.id === 'chat' && modelRateLimit ? { rateLimit: modelRateLimit } : {}) + })), runtimes: { limited: { enabled: true, maxConcurrency: 2, interactiveReservedSlots: 1 } } }; const manager = new RuntimeManager(config, { logger }); @@ -186,6 +195,48 @@ async function fixture(t, { holdStream = false, providerStatus = 200, rawFallbac }; } +test( + 'serving-node standard-only model limits preserve the forwarded class and keep raw providers private', + { timeout: 10000 }, + async (t) => { + const f = await fixture(t, { + mode: 'adaptive', + modelRateLimit: { + rate: '1/h', + burst: 0, + queue: false, + requestClasses: ['standard'] + } + }); + const chat = (requestClass, stream = false) => + f.request({ + route: '/v1/chat/completions', + model: 'chat', + requestClass, + body: { stream, messages: [{ role: 'user', content: 'Synthetic request' }] } + }).response; + assert.equal((await chat('standard')).status, 200); + assert.equal((await chat('standard')).status, 429); + const routing = await fetch(`http://127.0.0.1:${f.front.server.address().port}/gateway/routing`, { + headers: { authorization: 'Bearer front-fixture-key' } + }).then((response) => response.json()); + assert.deepEqual(routing.targetBackoffs, [], 'a remote class limit does not mark the model target unhealthy'); + assert.equal((await chat('interactive')).status, 200); + const foreground = await chat('foreground', true); + assert.equal(foreground.status, 200); + assert.match(await foreground.text(), /\[DONE\]/); + assert.equal(f.hits.length, 3); + assert.ok(f.hits.every((hit) => hit.headers['x-lloom-request-class'] === undefined)); + assert.deepEqual( + f.serving.metrics + .snapshot() + .recent.filter((entry) => entry.ok) + .map((entry) => entry.requestClass), + ['standard', 'interactive', 'foreground'] + ); + } +); + const pathways = [ ['embeddings', {}], [ diff --git a/test/rate-limit.test.mjs b/test/rate-limit.test.mjs index f93eecd..dd6c55e 100644 --- a/test/rate-limit.test.mjs +++ b/test/rate-limit.test.mjs @@ -62,6 +62,12 @@ async function chatGateway(fixture) { }); return { chat, + reload: async () => + fetch(`http://127.0.0.1:${port}/gateway/runtimes/keep-warm`, { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ reloadConfig: true }) + }), routing: () => fetch(`http://127.0.0.1:${port}/gateway/routing`).then((r) => r.json()), async close() { upstream.closeAllConnections(); @@ -155,6 +161,200 @@ test('the token bucket admits exactly n requests back-to-back then spaces by int assert.ok(count >= ceiling - 1, `2/m must keep admitting at the steady rate over 120s, saw ${count}`); }); +test('class-scoped limits validate and retain their selection after normalization', () => { + const settings = normalizeRateLimit({ rate: '4.5/m', burst: 0, queue: true, requestClasses: ['standard'] }); + assert.deepEqual(settings.requestClasses, ['standard']); + assert.equal(settings.rateMs, 60_000 / 4.5); + assert.deepEqual(normalizeRateLimit(settings), settings); + assert.deepEqual( + normalizeRateLimit({ maxConcurrent: 1, requestClasses: ['foreground', 'standard', 'foreground'] }).requestClasses, + ['standard', 'foreground'] + ); + for (const requestClasses of [null, [], 'standard', ['background'], ['standard', 1]]) { + assert.throws(() => normalizeRateLimit({ rate: '1/m', requestClasses }), /requestClasses/); + } +}); + +test('exempt classes neither wait for nor consume standard-only slots and rate budget', async () => { + const registry = createRateLimitRegistry({ + cloud: { maxConcurrent: 1, rate: '1/h', burst: 0, requestClasses: ['standard'] } + }); + const standard = registry.tryAcquireChain(['cloud'], { requestClass: 'standard' }); + assert.equal(typeof standard, 'function'); + assert.equal(registry.tryAcquireChain(['cloud'], { requestClass: 'standard' }), null); + const deadline = registry.limiter('cloud').gcra.nextAllowedInMs(); + for (const requestClass of ['foreground', 'interactive']) { + const release = registry.tryAcquireChain(['cloud'], { requestClass }); + assert.equal(typeof release, 'function'); + release(); + const legacy = await acquireRateLimitSlot(registry, 'cloud', { requestClass }); + legacy(); + await awaitRateBudget(registry, ['cloud'], { requestClass }); + } + assert.equal(registry.status()[0].active, 1); + assert.ok(registry.limiter('cloud').gcra.nextAllowedInMs() <= deadline); + assert.deepEqual(registry.status()[0].requestClasses, ['standard']); + standard(); + registry.close(); +}); + +test('class selection reload preserves spent budget and still composes with global limits', () => { + const registry = createRateLimitRegistry({ + outer: { rate: '1/h', burst: 0 }, + cloud: { rate: '1/h', burst: 0, requestClasses: ['standard'] } + }); + registry.tryAcquireChain(['cloud'], { requestClass: 'standard' })(); + const bucket = registry.limiter('cloud').gcra; + registry.sync({ + cloud: { rate: '1/h', burst: 0, requestClasses: ['interactive'] }, + outer: { rate: '1/h', burst: 0 } + }); + assert.equal(registry.limiter('cloud').gcra, bucket); + registry.tryAcquireChain(['outer', 'cloud'], { requestClass: 'foreground' })(); + assert.equal(registry.tryAcquireChain(['outer', 'cloud'], { requestClass: 'foreground' }), null); + assert.equal(registry.tryAcquireChain(['cloud'], { requestClass: 'interactive' }), null); + registry.tryAcquireChain(['cloud'], { requestClass: 'standard' })(); + registry.close(); +}); + +test( + 'reloading a class selection promptly releases legacy rate and semaphore waiters that become exempt', + { timeout: 3000 }, + async () => { + const registry = createRateLimitRegistry({ + rate: { rate: '1/m', burst: 0, queue: true, requestClasses: ['standard'] }, + slot: { maxConcurrent: 1, requestClasses: ['standard'] } + }); + const abort = new AbortController(); + const occupied = await acquireRateLimitSlot(registry, 'slot'); + const first = await acquireRateLimitSlot(registry, 'rate'); + first(); + const rateWait = awaitRateBudget(registry, ['rate'], { signal: abort.signal }); + const slotWait = acquireRateLimitSlot(registry, 'slot', { signal: abort.signal }); + assert.equal(registry.limiter('rate').rateQueued, 1); + assert.equal(registry.limiter('slot').semaphore.queued, 1); + try { + registry.sync({ + rate: { rate: '1/m', burst: 0, queue: true, requestClasses: ['foreground'] }, + slot: { maxConcurrent: 1, requestClasses: ['foreground'] } + }); + await rateWait; + (await slotWait)(); + assert.equal(registry.limiter('rate').rateQueued, 0); + assert.equal(registry.limiter('slot').semaphore.active, 1, 'the originally admitted request is preserved'); + assert.equal(registry.limiter('slot').semaphore.queued, 0); + } finally { + abort.abort(); + occupied(); + await Promise.allSettled([rateWait, slotWait]); + registry.close(); + } + } +); + +test( + 'adaptive config reload releases a paced standard request without resetting its model bucket', + { timeout: 5000 }, + async () => { + const directory = await fs.mkdtemp(path.join(os.tmpdir(), 'lloom-class-reload-')); + const fixture = await chatFixture( + { alias: { members: ['model-a'] } }, + { + models: [ + { + id: 'model-a', + backend: 'fixture', + upstreamModel: 'fixture', + rateLimit: { + rate: '1/h', + burst: 0, + queue: true, + requestClasses: ['standard'] + } + } + ] + } + ); + fixture.config.server.requestScheduling.mode = 'adaptive'; + fixture.config.server.foregroundIdleMs = 0; + fixture.config.sourcePath = path.join(directory, 'config.json'); + await fs.writeFile(fixture.config.sourcePath, JSON.stringify(fixture.config)); + const gateway = await chatGateway(fixture); + const abort = new AbortController(); + let waiting; + try { + assert.equal((await gateway.chat('alias')).status, 200); + let settled = false; + waiting = gateway.chat('alias', { signal: abort.signal }).then((response) => { + settled = true; + return response; + }); + await new Promise((resolve) => setTimeout(resolve, 30)); + assert.equal(settled, false); + fixture.config.models[0].rateLimit.requestClasses = ['foreground']; + await fs.writeFile(fixture.config.sourcePath, JSON.stringify(fixture.config)); + assert.equal((await gateway.reload()).status, 200); + assert.equal((await waiting).status, 200); + assert.deepEqual((await gateway.routing()).rateLimits[0].requestClasses, ['foreground']); + assert.equal( + (await gateway.routing()).rateLimits[0].rateLimited, + true, + 'the original standard request spent this bucket' + ); + } finally { + abort.abort(); + if (waiting) await Promise.allSettled([waiting]); + await gateway.close(); + await fs.rm(directory, { recursive: true, force: true }); + } + } +); + +for (const mode of ['off', 'interrupt', 'adaptive']) { + test( + `${mode} gateway enforces a standard-only model budget across aliases and bypasses it for chat`, + { timeout: 5000 }, + async () => { + const fixture = await chatFixture( + { first: { members: ['model-a'] }, second: { members: ['model-a'] } }, + { + models: [ + { + id: 'model-a', + backend: 'fixture', + upstreamModel: 'fixture', + rateLimit: { + rate: '1/h', + burst: 0, + queue: false, + requestClasses: ['standard'] + } + } + ] + } + ); + fixture.config.server.requestScheduling.mode = mode; + fixture.config.server.foregroundIdleMs = 0; + const gateway = await chatGateway(fixture); + const options = (requestClass) => ({ + headers: { 'content-type': 'application/json', 'x-lloom-request-class': requestClass } + }); + try { + assert.equal((await gateway.chat('first', options('standard'))).status, 200); + assert.equal((await gateway.chat('second', options('standard'))).status, 429); + assert.equal((await gateway.chat('model-a', options('standard'))).status, 429); + for (const requestClass of ['interactive', 'foreground']) { + assert.equal((await gateway.chat('second', options(requestClass))).status, 200); + assert.equal((await gateway.chat('model-a', options(requestClass))).status, 200); + } + assert.deepEqual((await gateway.routing()).rateLimits[0].requestClasses, ['standard']); + } finally { + await gateway.close(); + } + } + ); +} + test('the concurrency semaphore queues FIFO, respects aborts, and releases idempotently', async () => { const semaphore = createSemaphore(1); const first = await semaphore.acquire(); @@ -221,6 +421,9 @@ test('config validation rejects malformed rateLimit metadata and normalizes vali const loaded = await loadConfig(configPath); assert.deepEqual(loaded.models[0].rateLimit, { maxConcurrent: 2, burst: 0, rateMs: null }); assert.deepEqual(loaded.aliases.fast.rateLimit, { maxConcurrent: null, burst: 4, rateMs: 12_000 }); + base.models[0].rateLimit.requestClasses = ['standard']; + await fs.writeFile(configPath, JSON.stringify(base)); + assert.deepEqual((await loadConfig(configPath)).models[0].rateLimit.requestClasses, ['standard']); for (const mutation of [ () => (base.aliases.fast.rateLimit = '5/x'), diff --git a/test/request-scheduler.test.mjs b/test/request-scheduler.test.mjs index ee4afc8..92ebea2 100644 --- a/test/request-scheduler.test.mjs +++ b/test/request-scheduler.test.mjs @@ -271,6 +271,11 @@ test('shared limiter scopes prioritize higher demand even across separate comput const higher = await gate.admit({ groupId: 'device-b', scopeIds: ['shared-model'], requestClass: 'foreground' }); assert.equal(gate.readyFor(lower), true, 'its independent compute group is available'); assert.equal(gate.higherDemandForScopes(lower, ['shared-model']), true, 'its shared budget belongs to higher demand'); + assert.equal( + gate.higherDemandForScopes(lower, ['shared-model'], (_id, requestClass) => requestClass === 'standard'), + false, + 'a standard-only budget does not belong to an exempt higher class in another compute group' + ); assert.equal(gate.higherDemandForScopes(lower, ['independent-model']), false); assert.equal(gate.higherDemandForScopes(higher, ['shared-model']), false); higher.release();