stages

Composable Puzzletron stage handlers.

The runner accepts an explicit handler map so stage implementation can migrate piece by piece without hard-wiring GPU-heavy imports into config preflight.

Classes

`ArtifactChoice`	One artifact source a stage can select from an upstream stage.
`StageSkipReason`	Reasons that may make a skipped stage an accepted terminal state.
`StageSpec`	Immutable scheduler-neutral metadata for one public Puzzletron stage.
`StageStatus`	Manifest terminal states accepted by Puzzletron stage consumers.
`StageTerminalState`	Typed terminal state shared by workers, manifests, resume, and orchestration.

Functions

`configured_parent_stage_ids`	Return enabled parents, contracting paths through disabled optional stages.
`configured_stage_ids`	Return the exact campaign DAG after replacing the legacy post-MIP tail.
`distributed_stage_ids`	Return stages whose existing runner launches distributed workers.
`enabled_stage_ids`	Return configured stages in deterministic topological order.
`required_stage_ids`	Return required public stages in deterministic registry order.
`semantic_stage_config`	Return configuration that can change the semantic result of one stage.
`selected_parent_stage_ids`	Return direct parents after selecting this stage's configured artifact inputs.
`stage_display_name`	Return the public stage label for its independently configured granularity.
`stage_ids`	Return public stage IDs in deterministic registry order.
`stage_is_enabled`	Return whether a stage is enabled by its required/default/conditional metadata.
`stage_spec`	Return the immutable specification for one public stage ID.
`stage_terminal_state`	Parse one manifest terminal state, returning `None` for invalid evidence.
`topological_mapping_items`	Return mapping items in dependency order, independent of key serialization order.
`topological_stage_ids`	Return a deterministic topological order or raise for invalid graph metadata.
`activation_stage`
`activation_diagnostic_stage`
`bypass_diagnostic_stage`	Compare sorted-prefix pruning against the same slice from a bypass checkpoint.
`aiperf_stage`
`build_library_stage`
`bypass_overfit_stage`	Run only the isolated same-batch nested-bypass acceptance probe.
`bypass_stage`
`convert_stage`	Stage 1: convert a local HF checkpoint into canonical Puzzletron AnyModel.
`distillation_stage`
`distillation_overfit_stage`	Replay one frozen minibatch while globally distilling selected solutions.
`depth_stage`
`evaluation_stage`
`post_distillation_evaluation_stage`	Evaluate globally distilled checkpoints through the shared evaluator.
`mip_stage`
`scoring_stage`
`sort_stage`
`sort_equivalence_stage`	Evaluate teacher and sorted teacher with the chunked AutoModel scorer.
`width_slice_equivalence_stage`	Compare content-identical physical/runtime slices on one canonical batch.
`vllm_stats_stage`	Collect runtime statistics from converted-teacher candidates before library assembly.

class ArtifactChoice

Bases: object

One artifact source a stage can select from an upstream stage.

__init__(artifact, parent, when=None, requires=(), fallback=False)

Parameters:

artifact (str)
parent (str)
when (tuple[str, bool] | None)
requires (tuple[tuple[str, bool], ...])
fallback (bool)

Return type:

None

artifact: str

fallback: bool = False

parent: str

requires: tuple[tuple[str, bool], ...] = ()

when: tuple[str, bool] | None = None

class StageSkipReason

Bases: str, Enum

Reasons that may make a skipped stage an accepted terminal state.

DISABLED = 'disabled'

__new__(value)

class StageSpec

Bases: object

Immutable scheduler-neutral metadata for one public Puzzletron stage.

__init__(stage_id, display_name, completion_artifacts=(), granularity_label=False, required=False, default_enabled=False, enabled_when=None, enabled_requires=(), parents=(), conditional_parents=(), artifact_choices=(), distributed=False, report_order=0, topology_order=0, semantic_config_sections=())

Parameters:

stage_id (str)
display_name (str)
completion_artifacts (tuple[str, ...])
granularity_label (bool)
required (bool)
default_enabled (bool)
enabled_when (tuple[str, bool] | None)
enabled_requires (tuple[str, ...])
parents (tuple[str, ...])
conditional_parents (tuple[tuple[str, tuple[str, bool]], ...])
artifact_choices (tuple[ArtifactChoice, ...])
distributed (bool)
report_order (int)
topology_order (int)
semantic_config_sections (tuple[str, ...])

Return type:

None

artifact_choices: tuple[ArtifactChoice, ...] = ()

completion_artifacts: tuple[str, ...] = ()

conditional_parents: tuple[tuple[str, tuple[str, bool]], ...] = ()

default_enabled: bool = False

display_name: str

distributed: bool = False

enabled_requires: tuple[str, ...] = ()

enabled_when: tuple[str, bool] | None = None

granularity_label: bool = False

parents: tuple[str, ...] = ()

report_order: int = 0

required: bool = False

semantic_config_sections: tuple[str, ...] = ()

stage_id: str

topology_order: int = 0

class StageStatus

Bases: str, Enum

Manifest terminal states accepted by Puzzletron stage consumers.

IMPORTED = 'imported'

SKIPPED = 'skipped'

SUCCESS = 'success'

__new__(value)

class StageTerminalState

Bases: object

Typed terminal state shared by workers, manifests, resume, and orchestration.

__init__(status, skip_reason=None)

Parameters:

status (StageStatus)
skip_reason (StageSkipReason | None)

Return type:

None

allows_completion(stage_id, config)

Return whether graph and config semantics allow this terminal state.

Parameters:

stage_id (str)
config (Mapping[str, Any])

Return type:

bool

property produced_artifacts: bool: Return whether completion must be backed by stage artifacts.

skip_reason: StageSkipReason | None = None

status: StageStatus

activation_diagnostic_stage(config, manifest)

Parameters:

config (dict[str, Any])
manifest (StageManifest)

activation_stage(config, manifest)

Parameters:

config (dict[str, Any])
manifest (StageManifest)

aiperf_stage(config, manifest)

Parameters:

config (dict[str, Any])
manifest (StageManifest)

build_library_stage(config, manifest)

Parameters:

config (dict[str, Any])
manifest (StageManifest)

bypass_diagnostic_stage(config, manifest)

Compare sorted-prefix pruning against the same slice from a bypass checkpoint.

Parameters:

config (dict[str, Any])
manifest (StageManifest)

bypass_overfit_stage(config, manifest)

Run only the isolated same-batch nested-bypass acceptance probe.

Parameters:

config (dict[str, Any])
manifest (StageManifest)

bypass_stage(config, manifest)

Parameters:

config (dict[str, Any])
manifest (StageManifest)

configured_parent_stage_ids(stage_id, config)

Return enabled parents, contracting paths through disabled optional stages.

Parameters:

stage_id (str)
config (Mapping[str, Any])

Return type:

tuple[str, …]

configured_stage_ids(config, *, dynamic_post_mip_stage_ids=())

Return the exact campaign DAG after replacing the legacy post-MIP tail.

Parameters:

config (Mapping[str, Any])
dynamic_post_mip_stage_ids (Iterable[str])

Return type:

tuple[str, …]

convert_stage(config, manifest)

Stage 1: convert a local HF checkpoint into canonical Puzzletron AnyModel.

Parameters:

config (dict[str, Any])
manifest (StageManifest)

depth_stage(config, manifest)

Parameters:

config (dict[str, Any])
manifest (StageManifest)

distillation_overfit_stage(config, manifest)

Replay one frozen minibatch while globally distilling selected solutions.

Parameters:

config (dict[str, Any])
manifest (StageManifest)

distillation_stage(config, manifest)

Parameters:

config (dict[str, Any])
manifest (StageManifest)

distributed_stage_ids()

Return stages whose existing runner launches distributed workers.

Return type:: tuple[str, …]

enabled_stage_ids(config)

Return configured stages in deterministic topological order.

Parameters:: config (Mapping[str, Any])
Return type:: tuple[str, …]

evaluation_stage(config, manifest)

Parameters:

config (dict[str, Any])
manifest (StageManifest)

mip_stage(config, manifest)

Parameters:

config (dict[str, Any])
manifest (StageManifest)

post_distillation_evaluation_stage(config, manifest)

Evaluate globally distilled checkpoints through the shared evaluator.

Parameters:

config (dict[str, Any])
manifest (StageManifest)

required_stage_ids()

Return required public stages in deterministic registry order.

Return type:: tuple[str, …]

scoring_stage(config, manifest)

Parameters:

config (dict[str, Any])
manifest (StageManifest)

selected_parent_stage_ids(stage_id, config)

Return direct parents after selecting this stage’s configured artifact inputs.

Parameters:

stage_id (str)
config (Mapping[str, Any])

Return type:

tuple[str, …]

semantic_stage_config(config, stage_id, *, use_authored=True)

Return configuration that can change the semantic result of one stage.

Public stages declare their semantic sections alongside their other scheduler-neutral metadata. Dynamic stages that are not in the public registry retain the historical stage-ID section fallback.

Normalized worker configurations retain the independently loaded authored configuration under _runtime.authored_config. Semantic compatibility uses that authored view by default, while execution records may explicitly request the effective worker view.

Parameters:

config (Mapping[str, Any])
stage_id (str)
use_authored (bool)

Return type:

dict[str, Any]

sort_equivalence_stage(config, manifest)

Evaluate teacher and sorted teacher with the chunked AutoModel scorer.

Parameters:

config (dict[str, Any])
manifest (StageManifest)

sort_stage(config, manifest)

Parameters:

config (dict[str, Any])
manifest (StageManifest)

stage_display_name(stage_id, *, granularity=None)

Return the public stage label for its independently configured granularity.

Parameters:

stage_id (str)
granularity (str | None)

Return type:

str

stage_ids()

Return public stage IDs in deterministic registry order.

Return type:: tuple[str, …]

stage_is_enabled(stage_id, config)

Return whether a stage is enabled by its required/default/conditional metadata.

Parameters:

stage_id (str)
config (Mapping[str, Any])

Return type:

bool

stage_spec(stage_id)

Return the immutable specification for one public stage ID.

Parameters:: stage_id (str)
Return type:: StageSpec

stage_terminal_state(payload, *, expected_stage=None)

Parse one manifest terminal state, returning None for invalid evidence.

Parameters:

payload (Mapping[str, Any] | None)
expected_stage (str | None)

Return type:

StageTerminalState | None

topological_mapping_items(nodes, dependency_ids)

Return mapping items in dependency order, independent of key serialization order.

Parameters:

nodes (Mapping[str, Any])
dependency_ids (Callable[[str, Any], Iterable[str]])

Return type:

tuple[tuple[str, Any], …]

topological_stage_ids(specs=(StageSpec(stage_id='convert', display_name='Convert Checkpoint', completion_artifacts=('ckpts/teacher/config.json',), granularity_label=False, required=True, default_enabled=False, enabled_when=None, enabled_requires=(), parents=(), conditional_parents=(), artifact_choices=(), distributed=False, report_order=0, topology_order=0, semantic_config_sections=('convert',)), StageSpec(stage_id='tokenize_data', display_name='Tokenize Data', completion_artifacts=('dataset_cache/*.tokens',), granularity_label=False, required=False, default_enabled=False, enabled_when=None, enabled_requires=(), parents=('convert',), conditional_parents=(), artifact_choices=(), distributed=False, report_order=1, topology_order=1, semantic_config_sections=('tokenize_data', 'convert', 'data', 'dataset_path', 'train_token_cache_path', 'validation_token_cache_path', 'pruning', 'replacement_scoring', 'depth_importance', 'sort_sanity', 'width_sanity')), StageSpec(stage_id='vllm_stats', display_name='{unit} vLLM Stats', completion_artifacts=('artifacts/vllm_stats/summary.json',), granularity_label=True, required=False, default_enabled=False, enabled_when=None, enabled_requires=(), parents=('convert',), conditional_parents=(), artifact_choices=(), distributed=False, report_order=2, topology_order=2, semantic_config_sections=('vllm_stats', 'build_library', 'library')), StageSpec(stage_id='depth_importance', display_name='Depth Importance Estimation', completion_artifacts=('depth/iterative/trajectory.json',), granularity_label=False, required=False, default_enabled=False, enabled_when=None, enabled_requires=(), parents=('tokenize_data',), conditional_parents=(), artifact_choices=(), distributed=True, report_order=3, topology_order=3, semantic_config_sections=('depth_importance', 'pruning', 'replacement_scoring')), StageSpec(stage_id='width_importance', display_name='Width Importance Estimation', completion_artifacts=('pruning/pruning_scores/automodel/*/activation_passes_manifest.json',), granularity_label=False, required=True, default_enabled=False, enabled_when=None, enabled_requires=(), parents=('tokenize_data',), conditional_parents=(), artifact_choices=(), distributed=True, report_order=4, topology_order=4, semantic_config_sections=('width_importance', 'pruning')), StageSpec(stage_id='sort', display_name='Sort Checkpoint', completion_artifacts=('ckpts/sorted_teacher/config.json',), granularity_label=False, required=True, default_enabled=False, enabled_when=None, enabled_requires=(), parents=('width_importance',), conditional_parents=(), artifact_choices=(), distributed=True, report_order=5, topology_order=5, semantic_config_sections=('sort', 'pruning')), StageSpec(stage_id='sort_sanity', display_name='Sort Sanity Check', completion_artifacts=('artifacts/sort_sanity/summary.json',), granularity_label=False, required=False, default_enabled=False, enabled_when=None, enabled_requires=(), parents=('sort',), conditional_parents=(), artifact_choices=(), distributed=True, report_order=6, topology_order=6, semantic_config_sections=('sort_sanity', 'sanity', 'sort', 'pruning', 'replacement_scoring')), StageSpec(stage_id='width_sanity', display_name='Width Sanity Check', completion_artifacts=('artifacts/width_sanity/summary.json',), granularity_label=False, required=False, default_enabled=False, enabled_when=None, enabled_requires=('sort_sanity',), parents=('sort_sanity',), conditional_parents=(), artifact_choices=(), distributed=True, report_order=7, topology_order=7, semantic_config_sections=('width_sanity', 'sanity', 'pruning', 'replacement_scoring')), StageSpec(stage_id='slicing_sanity', display_name='Slicing Sanity Check', completion_artifacts=('artifacts/slicing_sanity/summary.json',), granularity_label=False, required=False, default_enabled=False, enabled_when=None, enabled_requires=('width_sanity',), parents=('width_sanity',), conditional_parents=(), artifact_choices=(), distributed=False, report_order=8, topology_order=8, semantic_config_sections=('slicing_sanity', 'sanity', 'sort', 'pruning', 'replacement_scoring')), StageSpec(stage_id='bypass_sanity', display_name='Bypass Sanity Check', completion_artifacts=('artifacts/bypass_sanity/summary.json',), granularity_label=False, required=False, default_enabled=False, enabled_when=None, enabled_requires=(), parents=('sort',), conditional_parents=(), artifact_choices=(), distributed=True, report_order=9, topology_order=9, semantic_config_sections=('bypass_sanity', 'sanity', 'bypass', 'pruning')), StageSpec(stage_id='bypass', display_name='{unit} Bypass', completion_artifacts=('artifacts/bypass/dp_observations.jsonl', 'artifacts/bypass/local_kd_loss_history.json'), granularity_label=True, required=False, default_enabled=False, enabled_when=None, enabled_requires=(), parents=('bypass_sanity',), conditional_parents=(), artifact_choices=(), distributed=True, report_order=10, topology_order=10, semantic_config_sections=('bypass', 'pruning')), StageSpec(stage_id='build_library', display_name='Build Block Library', completion_artifacts=('replacement_library.json', 'candidate_library.json'), granularity_label=False, required=True, default_enabled=False, enabled_when=None, enabled_requires=(), parents=('bypass',), conditional_parents=(('vllm_stats', ('vllm_stats.enabled', True)),), artifact_choices=(), distributed=True, report_order=11, topology_order=11, semantic_config_sections=('build_library', 'vllm_stats', 'library', 'bypass')), StageSpec(stage_id='replacement_scoring', display_name='Replace-one-{unit_lower} Scoring', completion_artifacts=('artifacts/replacement_scoring/summary.json',), granularity_label=True, required=True, default_enabled=False, enabled_when=None, enabled_requires=(), parents=('build_library',), conditional_parents=(), artifact_choices=(), distributed=True, report_order=12, topology_order=12, semantic_config_sections=('replacement_scoring', 'build_library', 'library', 'pruning')), StageSpec(stage_id='mip', display_name='MIP Search', completion_artifacts=('mip/profiles/*/mip_grid.json',), granularity_label=False, required=True, default_enabled=False, enabled_when=None, enabled_requires=(), parents=('vllm_stats', 'depth_importance', 'replacement_scoring'), conditional_parents=(), artifact_choices=(), distributed=False, report_order=13, topology_order=13, semantic_config_sections=('mip', 'realize_model', 'replacement_scoring', 'vllm_stats', 'library', 'bypass')), StageSpec(stage_id='zero_shot_evaluation', display_name='Zero-shot Evaluation', completion_artifacts=('artifacts/zero_shot_evaluation',), granularity_label=False, required=False, default_enabled=False, enabled_when=None, enabled_requires=(), parents=('mip',), conditional_parents=(), artifact_choices=(), distributed=True, report_order=14, topology_order=14, semantic_config_sections=('zero_shot_evaluation', 'convert', 'replacement_scoring')), StageSpec(stage_id='aiperf', display_name='AIPerf', completion_artifacts=('artifacts/aiperf/**/aiperf_results.json',), granularity_label=False, required=False, default_enabled=False, enabled_when=None, enabled_requires=(), parents=('mip',), conditional_parents=(), artifact_choices=(), distributed=False, report_order=15, topology_order=15, semantic_config_sections=('aiperf', 'zero_shot_evaluation')), StageSpec(stage_id='global_distillation_sanity', display_name='Global Distillation Sanity Check', completion_artifacts=('artifacts/global_distillation_sanity/**/global_distillation_sanity_summary.json',), granularity_label=False, required=False, default_enabled=False, enabled_when=None, enabled_requires=(), parents=('mip',), conditional_parents=(), artifact_choices=(), distributed=True, report_order=16, topology_order=16, semantic_config_sections=('global_distillation_sanity', 'sanity', 'global_distillation', 'replacement_scoring', 'calibration')), StageSpec(stage_id='global_distillation', display_name='Global Distillation', completion_artifacts=('artifacts/global_distillation/**/global_distillation_summary.json',), granularity_label=False, required=False, default_enabled=False, enabled_when=None, enabled_requires=(), parents=('global_distillation_sanity',), conditional_parents=(), artifact_choices=(), distributed=True, report_order=17, topology_order=17, semantic_config_sections=('global_distillation', 'zero_shot_evaluation', 'replacement_scoring')), StageSpec(stage_id='post_distillation_evaluation', display_name='Post Distillation Evaluation', completion_artifacts=('artifacts/post_distillation_evaluation/**/evaluation_summary.json',), granularity_label=False, required=False, default_enabled=False, enabled_when=None, enabled_requires=(), parents=('global_distillation',), conditional_parents=(), artifact_choices=(), distributed=True, report_order=18, topology_order=18, semantic_config_sections=('post_distillation_evaluation', 'global_distillation', 'zero_shot_evaluation', 'replacement_scoring'))))

Return a deterministic topological order or raise for invalid graph metadata.

Parameters:: specs (Iterable[StageSpec])
Return type:: tuple[str, …]

vllm_stats_stage(config, manifest)

Collect runtime statistics from converted-teacher candidates before library assembly.

Parameters:

config (dict[str, Any])
manifest (StageManifest)

width_slice_equivalence_stage(config, manifest)

Compare content-identical physical/runtime slices on one canonical batch.

Parameters:

config (dict[str, Any])
manifest (StageManifest)