census-arms.ts
81 lines · scripts/census-arms.ts at 3f79d81
1#!/usr/bin/env -S bun --no-env-file --config=/dev/null
2/**
3 * Gate census arms for evals-periodic.yml (GSTA-49 TM2). Each arm is one host
4 * and model running the weekly gate census plan:
5 * - Monday cron (0 6 * * 1): claude-opus (the capture model, lib/eval-model.ts)
6 * and codex-sol (gpt-6.1-sol on the Codex host, adapter cases only);
7 * - monthly cron (0 6 1 * *): claude-fable and codex-astra, census jobs only;
8 * - dispatch: `census_arms` runs only the census with the named arms; without
9 * it a dispatch keeps the periodic lane plus the Monday arms when the census
10 * runs (main, or include_gate_census).
11 * Judges stay in the Claude lanes: every census plan skips them.
12 *
13 * CLI (GitHub step outputs and env lines):
14 * --select env SCHEDULE, CENSUS_ARMS, INCLUDE_GATE_CENSUS -> arms=<json>, periodic_lane=<bool>
15 * --env <arm> KEY=VALUE lines selecting the arm's host and model
16 * --matrix <dir> <arms> gate_matrix=<json> from <dir>/<arm>/manifest.json
17 */
18import * as fs from 'node:fs';
19import * as path from 'node:path';
20import type { EvalHost } from '../test/helpers/e2e-hosts';
22export interface CensusArm { host: EvalHost; model?: string }
24export const CENSUS_ARMS: Readonly<Record<string, CensusArm>> = Object.freeze({
25 // No model: the Claude runners' capture default (claude-opus-5-5 since TM1).
26 'claude-opus': { host: 'claude' },
27 'codex-sol': { host: 'codex', model: 'gpt-6.1-sol' },
28 'claude-fable': { host: 'claude', model: 'claude-fable-5-1' },
29 'codex-astra': { host: 'codex', model: 'gpt-6-astra' },
30});
32export const WEEKLY_CENSUS_SCHEDULE = '0 6 * * 1';
33export const MONTHLY_CENSUS_SCHEDULE = '0 6 1 * *';
34export const WEEKLY_CENSUS_ARMS: readonly string[] = ['claude-opus', 'codex-sol'];
35export const MONTHLY_CENSUS_ARMS: readonly string[] = ['claude-fable', 'codex-astra'];
37/** Which census arms run, and whether the periodic lane runs beside them. */
38export function selectCensus(opts: { schedule?: string; censusArms?: string; includeGateCensus: boolean }): { arms: string[]; periodicLane: boolean } {
39 if (opts.schedule === MONTHLY_CENSUS_SCHEDULE) return { arms: [...MONTHLY_CENSUS_ARMS], periodicLane: false };
40 const named = (opts.censusArms ?? '').split(/[\s,]+/).filter(Boolean);
41 const unknown = named.filter(arm => !Object.hasOwn(CENSUS_ARMS, arm));
42 if (unknown.length) throw new Error(`unknown census arm(s): ${unknown.join(', ')}; known: ${Object.keys(CENSUS_ARMS).join(', ')}`);
43 if (named.length) return { arms: [...new Set(named)], periodicLane: false };
44 return { arms: opts.includeGateCensus ? [...WEEKLY_CENSUS_ARMS] : [], periodicLane: true };
45}
47/** The env one arm's planner and executor run with. */
48export function censusArmEnv(arm: string): Record<string, string> {
49 const spec = CENSUS_ARMS[arm];
50 if (!spec) throw new Error(`unknown census arm: ${arm}`);
51 if (!spec.model) return { EVALS_HOST: spec.host };
52 return { EVALS_HOST: spec.host, [spec.host === 'codex' ? 'GSTACK_CODEX_MODEL' : 'EVALS_MODEL']: spec.model };
53}
55/** One executor per arm and planned slice, with that slice's planned CI ceiling. */
56export function censusMatrix(arms: readonly string[], planDir: string): Array<{ arm: string; slice: number; timeout: number }> {
57 return arms.flatMap(arm => {
58 const manifest = JSON.parse(fs.readFileSync(path.join(planDir, arm, 'manifest.json'), 'utf8')) as
59 { sliceCount: number; plan?: { ciTimeoutMinutes: number; sliceCiTimeoutMinutes?: number[] } };
60 if (!manifest.plan) throw new Error(`${arm}: census manifest has no budget plan`);
61 return Array.from({ length: manifest.sliceCount }, (_, i) => ({
62 arm, slice: i + 1, timeout: manifest.plan!.sliceCiTimeoutMinutes?.[i] ?? manifest.plan!.ciTimeoutMinutes,
63 }));
64 });
65}
67if (import.meta.main) {
68 const [mode, ...args] = process.argv.slice(2);
69 if (mode === '--select') {
70 const { arms, periodicLane } = selectCensus({ schedule: process.env.SCHEDULE, censusArms: process.env.CENSUS_ARMS,
71 includeGateCensus: process.env.INCLUDE_GATE_CENSUS === 'true' });
72 console.log(`arms=${JSON.stringify(arms)}\nperiodic_lane=${periodicLane}`);
73 } else if (mode === '--env' && args.length === 1) {
74 for (const [key, value] of Object.entries(censusArmEnv(args[0]!))) console.log(`${key}=${value}`);
75 } else if (mode === '--matrix' && args.length === 2) {
76 console.log(`gate_matrix=${JSON.stringify(censusMatrix(JSON.parse(args[1]!) as string[], args[0]!))}`);
77 } else {
78 console.error('usage: census-arms.ts --select | --env <arm> | --matrix <plan dir> <arms json>');
79 process.exit(2);
80 }
81}