# 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`](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:

```yaml
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`:

   ```python
   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.