Pipeline runner

pstrain uses a small Python-native task runner in pstrain.lib.pipeline to orchestrate the training workflow. This document describes what it is, how it works, and why we built it instead of using Snakemake.

What it is

pstrain/lib/pipeline/
  runner.py     # Task, Pipeline, staleness, topo sort, execution
  context.py    # PipelineContext, config loading
  tasks.py      # Concrete tasks for the pstrain workflow

  • Task — an immutable dataclass: name, fn (callable), inputs: tuple[Path, ...], outputs: tuple[Path, ...], parallel_group: str, description: str, optional file dependencies in optional_inputs, and a lightweight configuration preflight callable.

  • Pipeline — registers tasks and resolves the DAG by matching one task’s outputs against another’s inputs. Plans, checks staleness, topologically sorts, and executes.

  • PipelineContext — per-run configuration (project dir, experiment, named config, derived feature/training params). Loaded from project/etc/configs.yaml.

The CLI entry points are:

  • pstrain build <target> — build a named target (e.g. cd-8g).

  • pstrain features — shortcut for pstrain build features.

  • pstrain step <name> — single-step debugging entry that delegates to the same pipeline.

How it works

Dependency resolution

Tasks declare file paths. The pipeline indexes outputs and uses inputs → outputs matching to walk the graph (the same model Snakemake uses).

Required external inputs must exist at plan time, including for cached targets. Missing generated inputs are scheduled through their producer. Optional inputs may be absent; their presence is recorded in the completion marker so adding or removing one invalidates its consumers. The training, model-definition, evaluation, and package tasks use this contract for shared/filler.dict.

Preflight callables validate configuration before dependencies execute. The training graph requires the native BW front end: features.ncep=13 and features.feat_type=1s_c_d_dd. Unsupported settings fail before extraction; the standalone features target retains the configurable extractor.

BW convergence keeps the signed per-frame likelihood delta and configured minimum number of passes, but requires a finite increase between zero and the threshold, inclusive. A likelihood decrease no longer counts as convergence. Training continues within the existing pass cap; it still saves each completed pass and does not roll back to a previous model automatically.

Staleness

A task is stale when any of:

  1. Any declared output is missing.

  2. The completion marker is missing, or an optional input’s presence differs from its completion record.

  3. The newest input mtime is greater than or equal to the oldest output mtime.

  4. Any upstream task is itself stale (transitively). The planner propagates staleness downstream because an upstream’s pending re-run will produce outputs newer than this task’s existing outputs.

--force marks every reachable task stale unconditionally.

Execution

Tasks run sequentially by default. Adjacent tasks sharing a parallel_group are batched together and dispatched to a ProcessPoolExecutor. This is how feature extraction fans out across the ~1000 audio files in train.fileids + test.fileids. Set -j N on the CLI to choose worker count.

The linear training chain (flat → ci-1g → ci-2g → … → cd-32g) runs in-process because each step depends on the previous one’s output.

Dry-run

--dry-run prints the topologically-sorted plan and never executes. The plan uses the same tab-separated shape as the run it predicts: two comment lines, a header row, then one row per stage.

# Plan for target: cd-1g
# 1263 task(s); 1263 stale
index	stage	tasks	status	description
1	provenance:split	1	unbuilt	Record effective split configuration
2	split	1	upstream:provenance:split	Partition all.transcription into train/test fileids + transcripts
3	provenance:features	1	unbuilt	Record effective features configuration
4-1135	features	1132	unbuilt
1136	provenance:training	1	unbuilt	Record effective training configuration
1137	flat	1	upstream:split	Initialize flat (uniform) acoustic model
1138	ci-1g	1	upstream:flat	Train CI-1g (1 Gaussian per state)
...
1263	cd-1g	1	upstream:cd-1g-init	Train tied CD-1g model

index is a position in the plan. A collapsed fan-out reports as one row carrying the range of positions it spans and the number of tasks in it, so a plan holding 1,132 per-utterance feature tasks still shows the shape of the build. --verbose lists every member instead.

status on an ordinary row is the reason the planner recorded, verbatim. Every status is a single word with no whitespace in it, because the plan is TSV: a value with a space in it survives cut -f but breaks anything that splits on whitespace. The reasons the planner produces are:

  • unbuilt — a declared output does not exist. This is the normal state of every stage in a fresh project, so it is not phrased as a fault.

  • stale — the outputs exist and have fallen behind their inputs. This is staleness in the strict sense, and no other status is.

  • incomplete — the outputs exist, but the private marker that says the task finished does not, so the outputs may be partial.

  • unconditional — the task declares no outputs, so there is nothing to compare and it runs every time.

  • unchanged — nothing to do.

  • forced — --force was passed, so nothing on disk was consulted.

  • upstream:<task> — this stage is due only because something it depends on is. The colon reads as a namespace separator, the way stage names such as provenance:split already use it.

A collapsed fan-out takes its status from what its members actually reported, so it preserves the same distinctions a per-member listing would. When every member agrees the group borrows their word, as the features row above does. When they disagree the group says it is mixed and carries the count in each state, as in mixed: 3 unbuilt, 37 unchanged. The known states are listed in a fixed order — unbuilt, unconditional, incomplete, forced, stale, unchanged — and anything else follows them sorted, so the same plan always prints the same row.

There are no bullet markers and no continuation line for the description, so every row carries the same columns and a plan pastes as a TSV beside the progress rows it foretells.

Why we built our own

A previous iteration used Snakemake. The workflow we actually have is small enough that Snakemake’s pull-ins didn’t pay off:

  • The DAG has ~15 logical nodes plus a fan-out over fileids. Not the large, branching, multi-sample DAG Snakemake is designed for.

  • Every Snakefile rule’s run: block just called into pstrain.lib.steps.run_* Python functions. Snakemake was a thin shim, not actually orchestrating shell commands or managing envs.

  • Snakemake pulls ~25 transitive dependencies (gitpython, jinja2, pulp, nbformat, …) for what amounts to “if output is older than input, re-run.”

  • The DSL is not real Python. Hard to test, hard to type-check, hard to debug. Inputs/outputs were duplicated between Snakefile rules and Step classes.

The runner replaces ~1100 lines of Snakefile + features.smk + targets.yaml with ~400 lines of Python in pstrain/lib/pipeline/. Adds zero runtime dependencies. Everything is one process, importable and debuggable.

What we explicitly don’t support

  • Cluster execution (Slurm, Kubernetes, etc.). If you need to run training on a cluster, Snakemake or Dagster would be a better fit.

  • Per-task conda envs. pstrain has one Python environment.

  • Content-hash staleness. Mtime parity with Snakemake is enough; layer a content-hash check on top if a real need shows up.

Multi-pronunciation training

Baum-Welch training defaults to multi-pronunciation mode: each word with k variants in the dictionary contributes k parallel phone paths to the per-utterance training graph, and forward-backward sums posteriors across them. Variant arc weights are initialized uniformly (1/k) so dictionary row order doesn’t pick the acoustic targets.

Opt out per-config by setting training.multipron_training: false in etc/configs.yaml; that config’s runs fall through to the legacy linear path (bit-identical to pstrain’s pre-multipron behavior).

See multi-pron-training.md for the full design and the as-built layout.

Stage-specific Baum-Welch control

Training schedules are configured independently because the upstream recipes do not apply one variance history and endpoint to every stage:

training:
  ci: {max_iterations: 10, min_iterations: 1, convergence_ratio: 0.001}
  untied: {max_iterations: 10, min_iterations: 1, convergence_ratio: 0.001}
  tied: {max_iterations: 10, min_iterations: 1, convergence_ratio: 0.001}

All three retain the SphinxTrain signed likelihood-delta arithmetic, but only finite, nonnegative changes within the threshold can indicate convergence. They may stop before their ten-pass cap after min_iterations; upstream stage 30 is a converge-with-cap loop, not a fixed-count loop. The separately frozen Arctic benchmark pin runs the same ten-pass untied cap. All stages use the A7c-matched 0.001 decision threshold by default.

Variance accumulation is deliberately code-defined by stage. CI and each newly split tied stage use one-pass variance on their first iteration and centered two-pass variance thereafter. CD-untied uses centered two-pass variance from its first iteration, matching the unconditional -2passvar yes in scripts/30.cd_hmm_untied/baum_welch.pl.

Adding a new pipeline node

  1. In pstrain/lib/pipeline/tasks.py, write a builder that closes over ctx and returns a Task:

    def _make_my_task(ctx: PipelineContext) -> Task:
    src = ctx.model_dir("ci-8g")
    out = ctx.model_dir("my-thing")

    def run() -> None:
        from pstrain.lib.something import do_thing
        do_thing(src=src, out=out)

    return Task(
        name="my-thing",
        fn=run,
        inputs=tuple(ctx.model_files("ci-8g")),
        outputs=tuple(ctx.model_files("my-thing")),
        description="Do my thing",
    )

  2. Register it in build_pipeline() and add it to TARGETS if it should be a named build target.

  3. If it’s a fan-out (one task per fileid, etc.), make fn a functools.partial over a top-level worker function so it pickles for ProcessPoolExecutor, and set parallel_group="some-name".

Testing

  • tests/test_pipeline_runner.py — runner behavior in isolation: topo sort, staleness, propagation, dry-run, parallel fan-out, cycles, failures.

  • tests/test_pipeline_tasks.py — task graph validation: every registered target has a producer, every declared target is registered, the cd-8g plan includes the full chain in dependency order.

  • tests/test_pipeline_integration.py — end-to-end training runs against a real audio corpus (CMU Arctic via PSTRAIN_TEST_PROJECT).

Experimental split variance regularization

training.split_variance_floor_fraction defaults to 0.0, which disables regularization and preserves the existing training path. Existing version-1 complete profiles that omit this newly optional field retain its schema default. A user-selected finite fraction greater than zero and at most one enables a lower bound on every saved variance coordinate in split training stages. This is an experimental coefficient, not a universal tuning recommendation. CI splits always use the CI one-Gaussian model as their reference; CD splits always use the CD one-Gaussian model. Later splits retain that same reference instead of ratcheting against their parent.

For each codebook and feature coordinate, the lower bound is the selected fraction of its matching one-Gaussian reference variance, broadcast across all densities. Reference zeros contribute a zero bound; the mechanism does not invent observations or guarantee a positive variance for unobserved coordinates. Reference values must be finite and nonnegative. Candidate values must be finite, but may be negative: single-precision normalization can leave small negative second-moment residuals in sparsely observed split densities, and the bound replaces them. Where the reference is zero, such a residual is stored as zero. Stored variances are unfloored normalization output either way; training and HMM.load apply their own evaluation floor before scoring, so a stored zero is never evaluated as zero. The current native front end supports one stream of 39 features; incompatible shapes are rejected rather than reshaped.

The reference is loaded once per split training run. Its variance file is an explicit task input when enabled, and the fraction participates in the existing training configuration fingerprint. As with other model inputs, file staleness uses modification times; changes that preserve those times are not independently detected by the runner. Each successful pass records the reference path, SHA-256, fraction, zero-bound count, and number of clamped coordinates in telemetry.

Density counts are staged before normalization. Normalized candidate parameters are then regularized and published with those counts, with checked rollback on write failures or interruptions. Checkpoints and subsequent passes consume the published bounded variances, so the saved model carries the constraint without decoder support. This publication is not an atomic directory snapshot for concurrent readers and does not promise recovery after process death.

The direct run_bw_training API also accepts variance_floor_reference and variance_floor_fraction. Callers must supply a reference with the same codebook mapping; matching dimensions alone cannot establish state identity. The initial input is validated but not rewritten before the first scoring pass: the floor applies to completed updates. Other pipeline training stages remain unchanged.

Explicit checkpoint recovery

When iteration checkpoints are enabled, pstrain checkpoints MODEL_DIR lists the retained updates. Update N contains the parameters produced by pass N; those parameters are evaluated by pass N+1. The last saved update normally remains unevaluated. A convergence decision or a successful save alone is not evidence that the saved update passed an alignment health check.

New pass telemetry binds the evaluated native scoring files (mdef, means, variances, mixture weights, and transitions) to their existing BW fingerprint. A separate per-file snapshot includes density counts, which identify the retained generation but are not part of native scoring identity. The report marks a checkpoint evaluated-healthy only when that snapshot matches and its evaluation processed utterances and frames, had finite likelihood statistics, and stayed within the configured skip limit. This is an alignment/statistics check, not a recognition-quality guarantee or a monotonic likelihood rule. Older telemetry without these bindings is reported-unverified; missing evaluation is unevaluated. Modified checkpoint bytes produce an evidence-mismatch. A failed save may follow a healthy input evaluation, so the stop decision is not used as a health proxy.

Use pstrain checkpoints MODEL_DIR --restore N --dry-run to inspect a selected restoration, then omit --dry-run to perform it. --json supplies structured output. Selection is always explicit, including for unevaluated or unverified snapshots; the tool never picks an older model automatically. Stop concurrent training and readers before restoration. The selected checkpoint must have the same state mapping as the destination. Its raw parameters and counts are copied without normalization; absent checkpoint counts remove destination counts.

The current model, counts, telemetry, and success metadata are first retained in MODEL_DIR/recovery-history/. Completion markers, model provenance, and any cached sendump are invalidated before publishing the selected parameters. Caught write failures and interruptions restore the previous parameter/count set, but leave success metadata invalidated. This keeps an interrupted recovery from presenting a model as a completed pipeline stage. Original checkpoints and training diagnostics are retained. A restored model is an explicit recovery artifact, not newly successful training; a pipeline build must revalidate its stage. Publication does not support concurrent readers or promise recovery from process death.

Training still retries the scheduled utterances on each pass. This tool adds no automatic rollback, final scoring pass, exclusion policy, or model selection rule.