Skip to content

fix(consumers): metadata consumers parse only the config subset they read#952

Open
ocg-goodfire wants to merge 1 commit into
feature/jaxfrom
bridge/task-consumer-config-schema-drift
Open

fix(consumers): metadata consumers parse only the config subset they read#952
ocg-goodfire wants to merge 1 commit into
feature/jaxfrom
bridge/task-consumer-config-schema-drift

Conversation

@ocg-goodfire

Copy link
Copy Markdown
Collaborator

Why

Stored runs' pinned launch_config.yaml refuse to parse under the current schema, so consumers can't open them. At feature/jax tip only 7 of 104 stored runs parse — the #939 anneal fields are the dominant cohort, but the sweep found ~10 distinct drift kinds: removed data/runtime/optimizer fields, deleted loss-type union tags, and fields from branches ahead of tip (runs launched today from an unmerged branch use data.train_batch_replay, unknown at tip). Every future schema change recreates this, in both directions.

The root cause for the metadata consumers is accidental coupling: run_metadata / PDAdapter (feeding autointerp/clustering) re-validated the full LMExperimentConfig (extra="forbid") when all they read is target.spec, pd.decomposition_targets, and four data fields.

What

  • ConsumerRunConfig (+ load_consumer_config): parses only the consumer-read subset, extra="ignore" elsewhere, strict on the read fields. The full schema remains the train/resume contract.
  • resolve_target: target resolution factored out of _resolve_target(LMExperimentConfig) so it works from the subset (the LMExperimentConfig path delegates to it).
  • run_metadata and PDAdapter.cfg switched onto the consumer parse. The run_metadata toy catch-all is gone — resolve_target returns the closed LM union, so the match is exhaustive by type.
  • Test: a drifted config (unknown fields at every level, deleted loss tag) parses via ConsumerRunConfig and refuses via LMExperimentConfig; missing/invalid read-fields still refuse.

Swept against the real runs dir: run_metadata now opens 103/103 stored runs (was 7/104 via the full parse; one run dir gained a config mid-sweep, hence 103 vs 104).

What this does NOT fix

open_jax_run (harvest/clustering workers) keeps the full-schema load_run_dir_config: it must rebuild the exact TrainState pytree (optimizers + per-loss-term persistent state, e.g. PPGD sources) to restore orbax, so it genuinely depends on the full config. Unblocking it for existing runs is the one-shot #939 migration (script attached to that PR, never run — coordinated separately in the bridge thread); shrinking it structurally is gated on the checkpoint-split work (decomposition item = components + ci_fn only).

Validation

  • make check clean
  • pytest param_decomp_lab -m 'not slow': 230 passed
  • Read-only sweep of run_metadata over all stored runs on cluster: 103/103 OK

bridge thread: consumer-config-schema-drift

🤖 Generated with Claude Code

https://claude.ai/code/session_01W9D6bgyT239nJ3E9gF3S3s

…read

run_metadata / PDAdapter re-validated a run's pinned launch_config.yaml
against the full LMExperimentConfig (extra=forbid), so every trainer-schema
change locked consumers out of stored runs: at tip only 7/104 stored runs
parsed (post-#939 anneal fields, removed data/runtime/optimizer fields,
deleted loss-type union tags, and fields from branches ahead of tip).

ConsumerRunConfig parses only what those consumers read (target.spec,
pd.decomposition_targets, four data fields) with extra=ignore, strict on the
read fields. Target resolution is factored into resolve_target so it no
longer needs the full schema. run_metadata now opens 103/103 stored runs.

open_jax_run keeps the full-schema load_run_dir_config: it must rebuild the
exact TrainState pytree (optimizers + per-loss-term persistent state) to
restore orbax. Shrinking it the same way is gated on the checkpoint-split
work.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01W9D6bgyT239nJ3E9gF3S3s
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant