Spaces:

automindy
/

fokemal

Sleeping

fokemal / server /src /__tests__ /lib /fallback-loop-model-bench.test.ts

Sync to upstream v0.13.7, keep HF Space adaptations

dcedb25 22 days ago

5.94 kB

import { describe, it, expect, beforeAll, beforeEach, vi } from 'vitest';

// Model-level failure benching: a model that keeps failing upstream must sink
// out of routing on EVERY key that can serve it, not only on the key that
// happened to fail last — the failure window counts across keys, so a per-key
// bench would leave the siblings feeding the same sick model.

vi.mock('../../services/health.js', () => ({
  checkKeyHealth: vi.fn(),
  // recordUpstreamSuccess touches this on the streak-clearing path.
  markKeyHealthyFromRequest: vi.fn(),
}));

import { initDb, getDb } from '../../db/index.js';
import { encrypt } from '../../lib/crypto.js';
import {
  newFallbackState,
  recordRetryableFailure,
  recordUpstreamSuccess,
  MODEL_FAILURE_COOLDOWN_MS,
  MODEL_FAILURE_THRESHOLD,
  MODEL_FAILURE_WINDOW_MS,
  resetModelFailureWindows,
} from '../../lib/fallback-loop.js';
import {
  getActiveCooldownsForKeys,
  isOnCooldown,
  resetKeyLocalityCache,
} from '../../services/ratelimit.js';
import type { RouteResult } from '../../services/router.js';

const PLATFORM = 'groq';
let keyA = 0;
let keyB = 0;
let modelA: { id: number; model_id: string };
let modelB: { id: number; model_id: string };

function insertKey(label: string): number {
  const { encrypted, iv, authTag } = encrypt(`test-${label}`);
  const info = getDb().prepare(`
    INSERT INTO api_keys (platform, label, encrypted_key, iv, auth_tag, status, enabled)
    VALUES (?, ?, ?, ?, ?, 'healthy', 1)
  `).run(PLATFORM, label, encrypted, iv, authTag);
  return Number(info.lastInsertRowid);
}

function routeFor(model: { id: number; model_id: string }, keyId: number): RouteResult {
  return {
    provider: {} as any,
    modelId: model.model_id,
    modelDbId: model.id,
    apiKey: 'k',
    keyId,
    platform: PLATFORM,
    displayName: model.model_id,
    rpdLimit: null,
    tpdLimit: null,
  };
}

// A plain 5xx: retryable, no quota signal, so the per-key bench is the short
// transient one (90s) — far below the model bench, which is what makes the
// assertions below able to tell the two apart.
const upstream500 = () => Object.assign(new Error('Groq API error 500: upstream'), { status: 500 });

function failOnce(model: { id: number; model_id: string }, keyId: number, now?: number) {
  // Fresh state per call: each failure stands for a separate request.
  recordRetryableFailure(routeFor(model, keyId), upstream500(), newFallbackState(), now);
}

function cooldownFor(model: { id: number; model_id: string }, keyId: number) {
  return getActiveCooldownsForKeys([keyId])
    .get(keyId)
    ?.find(c => c.platform === PLATFORM && c.modelId === model.model_id);
}

beforeAll(() => {
  process.env.ENCRYPTION_KEY = '0'.repeat(64);
  initDb(':memory:');
  const db = getDb();
  db.prepare('DELETE FROM api_keys').run();
  keyA = insertKey('key-a');
  keyB = insertKey('key-b');
  resetKeyLocalityCache();
  const models = db.prepare(
    'SELECT id, model_id FROM models WHERE platform = ? ORDER BY id LIMIT 2'
  ).all(PLATFORM) as { id: number; model_id: string }[];
  expect(models.length).toBe(2);
  [modelA, modelB] = models;
});

beforeEach(() => {
  getDb().prepare('DELETE FROM rate_limit_cooldowns').run();
  // The failure windows are module state and outlive the DB wipe above, so a
  // case would otherwise start partway to the threshold.
  resetModelFailureWindows();
});

describe('model-level failure benching covers every key of the model', () => {
  it('benches BOTH keys once the window trips, even though only one key failed last', () => {
    // Three failures inside the window, spread across the two keys.
    failOnce(modelA, keyA);
    failOnce(modelA, keyB);
    expect(MODEL_FAILURE_THRESHOLD).toBe(3);
    failOnce(modelA, keyA);

    for (const keyId of [keyA, keyB]) {
      expect(isOnCooldown(PLATFORM, modelA.model_id, keyId)).toBe(true);
      const cd = cooldownFor(modelA, keyId);
      expect(cd).toBeDefined();
      // The per-key transient bench is 90s; only the model bench reaches here.
      expect(cd!.remainingMs).toBeGreaterThan(MODEL_FAILURE_COOLDOWN_MS / 2);
    }
  });

  it('starts the streak over after a served request, so the next failure benches nothing extra', () => {
    failOnce(modelB, keyA);
    failOnce(modelB, keyB);
    // A served request is counter-evidence: the two failures above no longer count.
    recordUpstreamSuccess(routeFor(modelB, keyA), 100);
    failOnce(modelB, keyA);

    // Only the short per-key benches, no model-wide one.
    const benchedB = cooldownFor(modelB, keyB);
    expect(benchedB).toBeDefined();
    expect(benchedB!.remainingMs).toBeLessThan(MODEL_FAILURE_COOLDOWN_MS / 2);
    const benchedA = cooldownFor(modelB, keyA);
    expect(benchedA).toBeDefined();
    expect(benchedA!.remainingMs).toBeLessThan(MODEL_FAILURE_COOLDOWN_MS / 2);
  });

  // The window is a SLIDING one, which is only observable if the clock can be
  // moved. Without an injectable `now` this case would have to sleep for 15
  // real minutes, so the behaviour was never pinned.
  it('does not trip when the failures are spread beyond the sliding window', () => {
    const t0 = Date.now();
    failOnce(modelA, keyA, t0);
    failOnce(modelA, keyA, t0 + MODEL_FAILURE_WINDOW_MS + 1);
    failOnce(modelA, keyA, t0 + 2 * (MODEL_FAILURE_WINDOW_MS + 1));

    // Three failures, but never three inside one window: no model-wide bench.
    const benched = cooldownFor(modelA, keyB);
    expect(benched).toBeUndefined();
  });

  it('trips when the same three failures fall inside one window', () => {
    const t0 = Date.now();
    failOnce(modelA, keyA, t0);
    failOnce(modelA, keyA, t0 + 1000);
    failOnce(modelA, keyA, t0 + 2000);

    // keyB never failed, but the model-wide bench covers it.
    const benched = cooldownFor(modelA, keyB);
    expect(benched).toBeDefined();
    expect(benched!.remainingMs).toBeGreaterThan(MODEL_FAILURE_COOLDOWN_MS / 2);
  });
});

	import { describe, it, expect, beforeAll, beforeEach, vi } from 'vitest';

	// Model-level failure benching: a model that keeps failing upstream must sink
	// out of routing on EVERY key that can serve it, not only on the key that
	// happened to fail last — the failure window counts across keys, so a per-key
	// bench would leave the siblings feeding the same sick model.

	vi.mock('../../services/health.js', () => ({
	checkKeyHealth: vi.fn(),
	// recordUpstreamSuccess touches this on the streak-clearing path.
	markKeyHealthyFromRequest: vi.fn(),
	}));

	import { initDb, getDb } from '../../db/index.js';
	import { encrypt } from '../../lib/crypto.js';
	import {
	newFallbackState,
	recordRetryableFailure,
	recordUpstreamSuccess,
	MODEL_FAILURE_COOLDOWN_MS,
	MODEL_FAILURE_THRESHOLD,
	MODEL_FAILURE_WINDOW_MS,
	resetModelFailureWindows,
	} from '../../lib/fallback-loop.js';
	import {
	getActiveCooldownsForKeys,
	isOnCooldown,
	resetKeyLocalityCache,
	} from '../../services/ratelimit.js';
	import type { RouteResult } from '../../services/router.js';

	const PLATFORM = 'groq';
	let keyA = 0;
	let keyB = 0;
	let modelA: { id: number; model_id: string };
	let modelB: { id: number; model_id: string };

	function insertKey(label: string): number {
	const { encrypted, iv, authTag } = encrypt(`test-${label}`);
	const info = getDb().prepare(`
	INSERT INTO api_keys (platform, label, encrypted_key, iv, auth_tag, status, enabled)
	VALUES (?, ?, ?, ?, ?, 'healthy', 1)
	`).run(PLATFORM, label, encrypted, iv, authTag);
	return Number(info.lastInsertRowid);
	}

	function routeFor(model: { id: number; model_id: string }, keyId: number): RouteResult {
	return {
	provider: {} as any,
	modelId: model.model_id,
	modelDbId: model.id,
	apiKey: 'k',
	keyId,
	platform: PLATFORM,
	displayName: model.model_id,
	rpdLimit: null,
	tpdLimit: null,
	};
	}

	// A plain 5xx: retryable, no quota signal, so the per-key bench is the short
	// transient one (90s) — far below the model bench, which is what makes the
	// assertions below able to tell the two apart.
	const upstream500 = () => Object.assign(new Error('Groq API error 500: upstream'), { status: 500 });

	function failOnce(model: { id: number; model_id: string }, keyId: number, now?: number) {
	// Fresh state per call: each failure stands for a separate request.
	recordRetryableFailure(routeFor(model, keyId), upstream500(), newFallbackState(), now);
	}

	function cooldownFor(model: { id: number; model_id: string }, keyId: number) {
	return getActiveCooldownsForKeys([keyId])
	.get(keyId)
	?.find(c => c.platform === PLATFORM && c.modelId === model.model_id);
	}

	beforeAll(() => {
	process.env.ENCRYPTION_KEY = '0'.repeat(64);
	initDb(':memory:');
	const db = getDb();
	db.prepare('DELETE FROM api_keys').run();
	keyA = insertKey('key-a');
	keyB = insertKey('key-b');
	resetKeyLocalityCache();
	const models = db.prepare(
	'SELECT id, model_id FROM models WHERE platform = ? ORDER BY id LIMIT 2'
	).all(PLATFORM) as { id: number; model_id: string }[];
	expect(models.length).toBe(2);
	[modelA, modelB] = models;
	});

	beforeEach(() => {
	getDb().prepare('DELETE FROM rate_limit_cooldowns').run();
	// The failure windows are module state and outlive the DB wipe above, so a
	// case would otherwise start partway to the threshold.
	resetModelFailureWindows();
	});

	describe('model-level failure benching covers every key of the model', () => {
	it('benches BOTH keys once the window trips, even though only one key failed last', () => {
	// Three failures inside the window, spread across the two keys.
	failOnce(modelA, keyA);
	failOnce(modelA, keyB);
	expect(MODEL_FAILURE_THRESHOLD).toBe(3);
	failOnce(modelA, keyA);

	for (const keyId of [keyA, keyB]) {
	expect(isOnCooldown(PLATFORM, modelA.model_id, keyId)).toBe(true);
	const cd = cooldownFor(modelA, keyId);
	expect(cd).toBeDefined();
	// The per-key transient bench is 90s; only the model bench reaches here.
	expect(cd!.remainingMs).toBeGreaterThan(MODEL_FAILURE_COOLDOWN_MS / 2);
	}
	});

	it('starts the streak over after a served request, so the next failure benches nothing extra', () => {
	failOnce(modelB, keyA);
	failOnce(modelB, keyB);
	// A served request is counter-evidence: the two failures above no longer count.
	recordUpstreamSuccess(routeFor(modelB, keyA), 100);
	failOnce(modelB, keyA);

	// Only the short per-key benches, no model-wide one.
	const benchedB = cooldownFor(modelB, keyB);
	expect(benchedB).toBeDefined();
	expect(benchedB!.remainingMs).toBeLessThan(MODEL_FAILURE_COOLDOWN_MS / 2);
	const benchedA = cooldownFor(modelB, keyA);
	expect(benchedA).toBeDefined();
	expect(benchedA!.remainingMs).toBeLessThan(MODEL_FAILURE_COOLDOWN_MS / 2);
	});

	// The window is a SLIDING one, which is only observable if the clock can be
	// moved. Without an injectable `now` this case would have to sleep for 15
	// real minutes, so the behaviour was never pinned.
	it('does not trip when the failures are spread beyond the sliding window', () => {
	const t0 = Date.now();
	failOnce(modelA, keyA, t0);
	failOnce(modelA, keyA, t0 + MODEL_FAILURE_WINDOW_MS + 1);
	failOnce(modelA, keyA, t0 + 2 * (MODEL_FAILURE_WINDOW_MS + 1));

	// Three failures, but never three inside one window: no model-wide bench.
	const benched = cooldownFor(modelA, keyB);
	expect(benched).toBeUndefined();
	});

	it('trips when the same three failures fall inside one window', () => {
	const t0 = Date.now();
	failOnce(modelA, keyA, t0);
	failOnce(modelA, keyA, t0 + 1000);
	failOnce(modelA, keyA, t0 + 2000);

	// keyB never failed, but the model-wide bench covers it.
	const benched = cooldownFor(modelA, keyB);
	expect(benched).toBeDefined();
	expect(benched!.remainingMs).toBeGreaterThan(MODEL_FAILURE_COOLDOWN_MS / 2);
	});
	});