diff --git a/README.md b/README.md index b21421f..58b7c7e 100644 --- a/README.md +++ b/README.md @@ -118,6 +118,9 @@ Optional env: `POPULACE_HF_REPO`, `POPULACE_HF_REVISION` to point at a different US dataset/revision; `POPULACE_UK_HF_REPO`, `POPULACE_UK_HF_REVISION` for the UK; and `POPULACE_BE_HF_REPO`, `POPULACE_BE_HF_REVISION` for Belgium. The Belgium repository defaults in code to `policyengine/populace-be-private`. Set `HF_TOKEN` -or `HUGGINGFACE_TOKEN` to read private datasets. US staging defaults to -`policyengine/populace-us-staging`; override it with `POPULACE_STAGING_HF_REPO` -and `POPULACE_STAGING_HF_REVISION`. +or `HUGGINGFACE_TOKEN` to read private datasets. Each country with the +`staging` capability declares its own staging repository in the country +registry; it never falls back to another country's repository. US staging +defaults to `policyengine/populace-us-staging`; override it with the +backward-compatible `POPULACE_STAGING_HF_REPO` and +`POPULACE_STAGING_HF_REVISION` variables. diff --git a/frontend/lib/microcosm/countries.test.ts b/frontend/lib/microcosm/countries.test.ts index b1a21e7..72063eb 100644 --- a/frontend/lib/microcosm/countries.test.ts +++ b/frontend/lib/microcosm/countries.test.ts @@ -35,6 +35,18 @@ test("every registration carries the full shape", () => { expect(registration.repo_env).toMatch(/^[A-Z][A-Z0-9_]*$/); expect(registration.revision_env).toMatch(/^[A-Z][A-Z0-9_]*$/); } + if (registration.staging) { + expect(registration.staging.repo).toMatch(/^[a-z0-9-]+\/[a-z0-9-]+$/); + expect(registration.staging.revision.length).toBeGreaterThan(0); + if ( + registration.staging.repo_env != null || + registration.staging.revision_env != null + ) { + expect(registration.staging.repo_env).toMatch(/^[A-Z][A-Z0-9_]*$/); + expect(registration.staging.revision_env).toMatch(/^[A-Z][A-Z0-9_]*$/); + } + } + expect(hasCapability(country, "staging")).toBe(registration.staging != null); } }); @@ -49,6 +61,11 @@ test("keeps the live registrations on their published repositories and labels", geography_id: "0100000US", visibility: "public", capabilities: COUNTRY_CAPABILITIES, + staging: { + repo: "policyengine/populace-us-staging", + repo_env: "POPULACE_STAGING_HF_REPO", + revision_env: "POPULACE_STAGING_HF_REVISION", + }, }); expect(countryRegistration("uk")).toMatchObject({ repo: "policyengine/populace-uk-private", @@ -71,12 +88,15 @@ test("keeps the live registrations on their published repositories and labels", test("selectable countries follow registry order and exclude fixtures", () => { expect(selectableCountries()).toEqual(["us", "uk", "be"]); expect(countryRegistration("zz").fixture).toBe(true); + expect(countryRegistration("am").fixture).toBe(true); expect(selectableCountries()).not.toContain("zz"); }); test("fixture registrations are valid countries without being selectable", () => { expect(isCountry("zz")).toBe(true); + expect(isCountry("am")).toBe(true); expect(parseCountry("zz")).toBe("zz"); + expect(parseCountry("am")).toBe("am"); }); test("country parsing is exact and defaults to the registry default", () => { @@ -96,6 +116,7 @@ test("country parsing is exact and defaults to the registry default", () => { test("capability gates read the registration", () => { expect(hasCapability("us", "staging")).toBe(true); + expect(hasCapability("am", "staging")).toBe(true); expect(hasCapability("us", "model_coverage")).toBe(true); expect(hasCapability("uk", "staging")).toBe(false); expect(hasCapability("be", "pipeline")).toBe(false); diff --git a/frontend/lib/microcosm/countries.ts b/frontend/lib/microcosm/countries.ts index 2b1de1e..1ee7c7d 100644 --- a/frontend/lib/microcosm/countries.ts +++ b/frontend/lib/microcosm/countries.ts @@ -38,6 +38,14 @@ export interface CountryRegistration { geography_id: string | null; visibility: RepositoryVisibility; capabilities: readonly CountryCapability[]; + // Optional staging telemetry repository. A staging-capable country must + // declare this instead of inheriting another country's repository. + staging?: { + repo: string; + revision: string; + repo_env?: string; + revision_env?: string; + }; // Other jurisdiction codes a cross-dataset bundle may use for this country. jurisdiction_aliases?: readonly string[]; // Conformance-only registration: a valid country that is never listed in @@ -68,6 +76,14 @@ export const COUNTRY_REGISTRY = { geography_id: "0100000US", visibility: "public", capabilities: ALL_CAPABILITIES, + staging: { + // Deprecated upstream identifier: the US staging publisher still uses + // the former Populace repository and deployment-variable names. + repo: "policyengine/populace-us-staging", + revision: "main", + repo_env: "POPULACE_STAGING_HF_REPO", + revision_env: "POPULACE_STAGING_HF_REVISION", + }, }, uk: { repo: "policyengine/populace-uk-private", @@ -106,6 +122,24 @@ export const COUNTRY_REGISTRY = { capabilities: ["calibration", "targets", "compare"], fixture: true, }, + // Synthetic non-US staging registration. It verifies that staging artifact + // resolution is country-scoped without exposing an unfinished country in + // selectors or treating its fixture repository as a published dataset. + am: { + repo: "policyengine/microcosm-am-fixture", + revision: "main", + label: "Armenia", + dataset_label: "Microcosm Armenia", + geography: "Armenia", + geography_id: null, + visibility: "private", + capabilities: ["calibration", "targets", "compare", "staging"], + staging: { + repo: "policyengine/microcosm-am-staging-fixture", + revision: "main", + }, + fixture: true, + }, } satisfies Record; export type MicrocosmCountry = keyof typeof COUNTRY_REGISTRY; diff --git a/frontend/lib/microcosm/staging-artifact.test.ts b/frontend/lib/microcosm/staging-artifact.test.ts index 7aa7053..1ca8327 100644 --- a/frontend/lib/microcosm/staging-artifact.test.ts +++ b/frontend/lib/microcosm/staging-artifact.test.ts @@ -5,10 +5,90 @@ import { loadStagingRun, loadStagingRuns, loadStagingTargetDiagnostics, + MICROCOSM_STAGING_HF_REPO, + MICROCOSM_STAGING_HF_REVISION, + stagingRepository, + stagingResolveUrl, stagingTargetChangeCacheTtlSeconds, stagingUnavailableReason, } from "./staging-artifact"; +test("resolves staging repositories from each country registration", () => { + expect(stagingRepository("us")).toEqual({ + repo: "policyengine/populace-us-staging", + revision: "main", + }); + expect(MICROCOSM_STAGING_HF_REPO).toBe("policyengine/populace-us-staging"); + expect(MICROCOSM_STAGING_HF_REVISION).toBe("main"); + + // Armenia is a fixture-only registration: this asserts country-specific + // repository selection without publishing a country or telemetry artifact. + expect(stagingRepository("am")).toEqual({ + repo: "policyengine/microcosm-am-staging-fixture", + revision: "main", + }); + expect(stagingResolveUrl("runs/am-fixture/progress.json", "am")).toBe( + "https://huggingface.co/datasets/policyengine/microcosm-am-staging-fixture/resolve/main/runs/am-fixture/progress.json", + ); + + const originalRepo = process.env.POPULACE_STAGING_HF_REPO; + const originalRevision = process.env.POPULACE_STAGING_HF_REVISION; + try { + process.env.POPULACE_STAGING_HF_REPO = "policyengine/us-staging-override"; + process.env.POPULACE_STAGING_HF_REVISION = "test-revision"; + expect(stagingRepository("us")).toEqual({ + repo: "policyengine/us-staging-override", + revision: "test-revision", + }); + expect(stagingRepository("am")).toEqual({ + repo: "policyengine/microcosm-am-staging-fixture", + revision: "main", + }); + } finally { + if (originalRepo === undefined) delete process.env.POPULACE_STAGING_HF_REPO; + else process.env.POPULACE_STAGING_HF_REPO = originalRepo; + if (originalRevision === undefined) delete process.env.POPULACE_STAGING_HF_REVISION; + else process.env.POPULACE_STAGING_HF_REVISION = originalRevision; + } +}); + +test("loads Armenia staging telemetry from Armenia's registered repository", async () => { + const originalFetch = globalThis.fetch; + const urls: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + urls.push(url); + if (url.endsWith("/runs.json")) { + return Response.json({ + runs: [ + { + run_id: "am-fixture", + candidate_release_id: "am-candidate", + status: "running", + }, + ], + }); + } + if (url.includes("/tree/main/runs?recursive=true")) return Response.json([]); + return new Response(null, { status: 404 }); + }) as typeof fetch; + + try { + await expect(loadStagingRuns(0, "am")).resolves.toMatchObject({ + available: true, + source_repo: "policyengine/microcosm-am-staging-fixture", + revision: "main", + runs: [{ run_id: "am-fixture", candidate_release_id: "am-candidate" }], + }); + expect(urls).toEqual([ + "https://huggingface.co/datasets/policyengine/microcosm-am-staging-fixture/resolve/main/runs.json", + "https://huggingface.co/api/datasets/policyengine/microcosm-am-staging-fixture/tree/main/runs?recursive=true", + ]); + } finally { + globalThis.fetch = originalFetch; + } +}); + test("names the country when staging is unavailable", () => { expect(stagingUnavailableReason("us")).toBeNull(); expect(stagingUnavailableReason("uk")).toBe( diff --git a/frontend/lib/microcosm/staging-artifact.ts b/frontend/lib/microcosm/staging-artifact.ts index b7628f6..9ae36f6 100644 --- a/frontend/lib/microcosm/staging-artifact.ts +++ b/frontend/lib/microcosm/staging-artifact.ts @@ -10,7 +10,7 @@ import { loadRelease, microcosmCountryGeography, } from "@/lib/microcosm/latest-artifact"; -import { hasCapability } from "@/lib/microcosm/countries"; +import { countryRegistration, hasCapability } from "@/lib/microcosm/countries"; import { type ReformValidation, buildReformValidation, @@ -34,21 +34,63 @@ const targetChangeCache = new Map(); export const MICROCOSM_STAGING_HF_REPO_ENV = "POPULACE_STAGING_HF_REPO"; export const MICROCOSM_STAGING_HF_REVISION_ENV = "POPULACE_STAGING_HF_REVISION"; -export const MICROCOSM_STAGING_HF_REPO = - // Deprecated upstream identifier: Microcosm staging still publishes under - // the former Populace repository and deployment-variable names. - process.env[MICROCOSM_STAGING_HF_REPO_ENV] ?? "policyengine/populace-us-staging"; -export const MICROCOSM_STAGING_HF_REVISION = - process.env[MICROCOSM_STAGING_HF_REVISION_ENV] ?? "main"; - -// Staging telemetry is served for countries registered with the `staging` -// capability; the single staging repository above is the one they read. +export interface StagingRepository { + repo: string; + revision: string; +} + +// Resolve staging telemetry from the country registry. The exported US values +// retain the legacy API and deployment-variable behavior for existing callers. +export function stagingRepository(country: MicrocosmCountry): StagingRepository | null { + const staging = countryRegistration(country).staging; + if (!hasCapability(country, "staging") || !staging) return null; + return { + repo: + (staging.repo_env ? process.env[staging.repo_env] : undefined) ?? + staging.repo, + revision: + (staging.revision_env ? process.env[staging.revision_env] : undefined) ?? + staging.revision, + }; +} + +export const MICROCOSM_STAGING_HF_REPO = stagingRepository("us")!.repo; +export const MICROCOSM_STAGING_HF_REVISION = stagingRepository("us")!.revision; + export function stagingUnavailableReason(country: MicrocosmCountry): string | null { - return hasCapability(country, "staging") + return stagingRepository(country) ? null : `${microcosmCountryGeography(country)} has no staging repository.`; } +function requiredStagingRepository(country: MicrocosmCountry): StagingRepository { + const repository = stagingRepository(country); + if (repository) return repository; + throw new Error( + stagingUnavailableReason(country) ?? "Staging repository is not configured.", + ); +} + +function stagingResolveUrlFor(repository: StagingRepository, path: string): string { + return `https://huggingface.co/datasets/${repository.repo}/resolve/${repository.revision}/${path}`; +} + +export function stagingResolveUrl(path: string, country: MicrocosmCountry = "us"): string { + return stagingResolveUrlFor(requiredStagingRepository(country), path); +} + +function stagingTreeUrl(repository: StagingRepository): string { + return `https://huggingface.co/api/datasets/${repository.repo}/tree/${repository.revision}/runs?recursive=true`; +} + +function stagingRepoUrl(repository: StagingRepository): string { + return `https://huggingface.co/api/datasets/${repository.repo}`; +} + +function stagingSource(country: MicrocosmCountry): StagingRepository { + return requiredStagingRepository(country); +} + function unavailableStaging(country: MicrocosmCountry) { const detail = stagingUnavailableReason(country); if (!detail) return null; @@ -64,15 +106,20 @@ class StagingFetchError extends Error { constructor( public readonly status: number, path: string, + repository: StagingRepository, ) { - super(stagingFetchMessage(status, path)); + super(stagingFetchMessage(status, path, repository)); } } -function stagingFetchMessage(status: number, path: string): string { +function stagingFetchMessage( + status: number, + path: string, + repository: StagingRepository, +): string { if (status === 401 || status === 403) { return ( - `Staging repo ${MICROCOSM_STAGING_HF_REPO} is not readable by this deployment ` + + `Staging repo ${repository.repo} is not readable by this deployment ` + `(${status} fetching ${path}). Set HF_TOKEN/HUGGINGFACE_TOKEN on the server, ` + "or publish staging telemetry to a public dataset repo." ); @@ -88,10 +135,6 @@ function hfHeaders(): HeadersInit | undefined { return token ? { Authorization: `Bearer ${token}` } : undefined; } -export function stagingResolveUrl(path: string): string { - return `https://huggingface.co/datasets/${MICROCOSM_STAGING_HF_REPO}/resolve/${MICROCOSM_STAGING_HF_REVISION}/${path}`; -} - function stagingFetchOptions(revalidate: number): RequestInit { return { headers: hfHeaders(), @@ -99,34 +142,58 @@ function stagingFetchOptions(revalidate: number): RequestInit { }; } -async function stagingJson(path: string, revalidate: number): Promise { - const res = await fetch(stagingResolveUrl(path), stagingFetchOptions(revalidate)); - if (!res.ok) throw new StagingFetchError(res.status, path); +async function stagingJson( + path: string, + revalidate: number, + country: MicrocosmCountry, +): Promise { + const repository = stagingSource(country); + const res = await fetch( + stagingResolveUrlFor(repository, path), + stagingFetchOptions(revalidate), + ); + if (!res.ok) throw new StagingFetchError(res.status, path, repository); return asObject(await res.json()); } -async function stagingJsonOrNull(path: string, revalidate: number): Promise { +async function stagingJsonOrNull( + path: string, + revalidate: number, + country: MicrocosmCountry, +): Promise { try { - return await stagingJson(path, revalidate); + return await stagingJson(path, revalidate, country); } catch (error) { if (error instanceof StagingFetchError && error.status !== 404) throw error; return null; } } -async function stagingTextOrNull(path: string, revalidate: number): Promise { - const res = await fetch(stagingResolveUrl(path), stagingFetchOptions(revalidate)); +async function stagingTextOrNull( + path: string, + revalidate: number, + country: MicrocosmCountry, +): Promise { + const repository = stagingSource(country); + const res = await fetch( + stagingResolveUrlFor(repository, path), + stagingFetchOptions(revalidate), + ); if (!res.ok) { if (res.status === 404) return null; - throw new StagingFetchError(res.status, path); + throw new StagingFetchError(res.status, path, repository); } return res.text(); } -async function stagingTree(revalidate: number): Promise { - const url = `https://huggingface.co/api/datasets/${MICROCOSM_STAGING_HF_REPO}/tree/${MICROCOSM_STAGING_HF_REVISION}/runs?recursive=true`; +async function stagingTree( + revalidate: number, + country: MicrocosmCountry, +): Promise { + const repository = stagingSource(country); + const url = stagingTreeUrl(repository); const res = await fetch(url, stagingFetchOptions(revalidate)); - if (!res.ok) throw new StagingFetchError(res.status, "runs tree"); + if (!res.ok) throw new StagingFetchError(res.status, "runs tree", repository); const tree = await res.json(); return Array.isArray(tree) ? tree.map(asObject) : []; } @@ -202,7 +269,8 @@ export async function loadStagingRuns( runs: [] as StagingRunSummary[], }; } - const index = await stagingJsonOrNull("runs.json", revalidate); + const repository = stagingSource(country); + const index = await stagingJsonOrNull("runs.json", revalidate, country); const indexedRuns = Array.isArray(index?.runs) ? (index.runs as JsonObject[]) .map((row) => { @@ -231,7 +299,7 @@ export async function loadStagingRuns( const runIds = new Set(indexedRuns.map((run) => run.run_id)); let treeMissing = false; try { - for (const entry of await stagingTree(revalidate)) { + for (const entry of await stagingTree(revalidate, country)) { if (typeof entry.path !== "string") continue; const match = /^runs\/([^/]+)\//.exec(entry.path); if (match) runIds.add(match[1]); @@ -248,16 +316,16 @@ export async function loadStagingRuns( // empty list — a silent empty state hides a broken token. if (index == null && treeMissing && runIds.size === 0) { const repoRes = await fetch( - `https://huggingface.co/api/datasets/${MICROCOSM_STAGING_HF_REPO}`, + stagingRepoUrl(repository), stagingFetchOptions(revalidate), ); if (!repoRes.ok) { return { available: false, - source_repo: MICROCOSM_STAGING_HF_REPO, - revision: MICROCOSM_STAGING_HF_REVISION, + source_repo: repository.repo, + revision: repository.revision, detail: - `Staging repo ${MICROCOSM_STAGING_HF_REPO} is not visible (HTTP ${repoRes.status}). ` + + `Staging repo ${repository.repo} is not visible (HTTP ${repoRes.status}). ` + "It is private — a missing or expired HF token reads as 404, not 401.", runs: [], }; @@ -274,7 +342,11 @@ export async function loadStagingRuns( const MAX_UNINDEXED_FETCH = 50; const fetched = await Promise.all( missing.slice(0, MAX_UNINDEXED_FETCH).map(async (runId) => { - const progress = await stagingJsonOrNull(`runs/${runId}/progress.json`, revalidate); + const progress = await stagingJsonOrNull( + `runs/${runId}/progress.json`, + revalidate, + country, + ); return summaryFromProgress(runId, progress); }), ); @@ -283,8 +355,8 @@ export async function loadStagingRuns( return { available: true, - source_repo: MICROCOSM_STAGING_HF_REPO, - revision: MICROCOSM_STAGING_HF_REVISION, + source_repo: repository.repo, + revision: repository.revision, truncated, runs: [...byId.values()].sort(sortRuns), }; @@ -313,16 +385,21 @@ export async function loadStagingCalibration( ): Promise { if (stagingUnavailableReason(country)) return null; assertSafeReleaseId(runId, "run"); - const progress = await stagingJsonOrNull(`runs/${runId}/progress.json`, revalidate); + const progress = await stagingJsonOrNull( + `runs/${runId}/progress.json`, + revalidate, + country, + ); const candidateReleaseId = stringValue(progress?.candidate_release_id) ?? runId; const diag = await stagingJsonOrNull( `runs/${runId}/calibration_diagnostics.json`, revalidate, + country, ); if (!diag) return null; const [buildManifest, releaseManifest] = await Promise.all([ - stagingJsonOrNull(`runs/${runId}/build_manifest.json`, revalidate), - stagingJsonOrNull(`runs/${runId}/release_manifest.json`, revalidate), + stagingJsonOrNull(`runs/${runId}/build_manifest.json`, revalidate, country), + stagingJsonOrNull(`runs/${runId}/release_manifest.json`, revalidate, country), ]); return buildCalibration( diag, @@ -346,6 +423,7 @@ export async function loadStagingTargetChangeDataset( const progress = await stagingJsonOrNull( `runs/${runId}/progress.json`, TARGET_CHANGE_MUTABLE_CACHE_SECONDS, + country, ); const ttlSeconds = stagingTargetChangeCacheTtlSeconds(progress?.status); const cacheKey = `${country}:${runId}:${releaseId}`; @@ -412,21 +490,22 @@ export async function loadStagingRun( cal, reformValidationRaw, ] = await Promise.all([ - stagingJsonOrNull(`runs/${runId}/progress.json`, revalidate), - stagingJsonOrNull(`runs/${runId}/run_manifest.json`, revalidate), - stagingJsonOrNull(`runs/${runId}/calibration_progress.json`, revalidate), - stagingTextOrNull(`runs/${runId}/events.ndjson`, revalidate), + stagingJsonOrNull(`runs/${runId}/progress.json`, revalidate, country), + stagingJsonOrNull(`runs/${runId}/run_manifest.json`, revalidate, country), + stagingJsonOrNull(`runs/${runId}/calibration_progress.json`, revalidate, country), + stagingTextOrNull(`runs/${runId}/events.ndjson`, revalidate, country), loadStagingCalibration(runId, revalidate, country), - stagingJsonOrNull(`runs/${runId}/reform_validation.json`, revalidate), + stagingJsonOrNull(`runs/${runId}/reform_validation.json`, revalidate, country), ]); const candidateReleaseId = stringValue(progress?.candidate_release_id) ?? stringValue(runManifest?.candidate_release_id) ?? runId; + const repository = stagingSource(country); return { available: true, - source_repo: MICROCOSM_STAGING_HF_REPO, - revision: MICROCOSM_STAGING_HF_REVISION, + source_repo: repository.repo, + revision: repository.revision, run_id: runId, candidate_release_id: candidateReleaseId, progress, @@ -454,7 +533,11 @@ export async function loadStagingReformValidationRaw( ): Promise { if (stagingUnavailableReason(country)) return null; assertSafeReleaseId(runId, "run"); - return stagingJsonOrNull(`runs/${runId}/reform_validation.json`, revalidate); + return stagingJsonOrNull( + `runs/${runId}/reform_validation.json`, + revalidate, + country, + ); } export async function loadStagingTargetDiagnostics(