Learn AI Create the Future — with Pranay
Main site

ML Code That Survives

Refactoring a 600-line training script

Pin the script's current answer to a number first, then move code in small steps that must keep producing that exact number — the test is what turns a terrifying rewrite into a boring afternoon.

Read these first

On this page 5
  1. Why it exists
  2. How it works
  3. A real example you have seen
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

Refactoring means rearranging code without changing what it does — and the safe way starts by writing down what it does.

A mechanic who is about to open an engine takes a reading first. Compression, idle speed, the number on the dial. Then he opens it, replaces the part, closes it, and takes the same reading again.

If the number matches, the repair did not break anything. Without that first reading, he has opinions and no evidence.

A 600-line training script is that engine. Everybody wants to tidy it. Almost nobody writes down the reading first, and that is why tidying it so often quietly breaks the model.

Why it exists

Long training scripts are not written; they accumulate. A file that started at 80 lines gains a feature here, a special case there, three commented-out experiments. Then a hyperparameter — a setting a person chooses rather than one the model learns — that somebody hardcoded during a bad week.

Nobody dares touch it, because the code and the results have become one thing in people's minds. "It works, do not look at it."

The problem is that unchangeable code stops the project. You cannot test it, reuse it, or hand it to a new joiner. You cannot answer "what if we changed the split?"

Refactoring is how the file becomes changeable again. Pinning the number first is how you do it without gambling.

How it works

  1. RECORD      run the script, write the final metric down
                 (0.925645 — that is now the truth)

  2. FENCE       write a test that runs the script and checks that number

  3. MOVE        cut ONE piece into a named function
                 |
                 v
  4. CHECK       run the test. Same number? Keep it. Different? Undo it.
                 |
                 +--> repeat from step 3, one piece at a time

The test in step 2 has a name: a characterisation test. It does not check that the code is right. It checks that the code is unchanged. Those are different jobs, and this is the one that makes refactoring safe.

A real example you have seen

Moving house. You do not tip everything into a lorry and hope. You photograph the shelf, pack one box, label it, and check the photo when you unpack.

Anyone who has moved without labelling boxes knows exactly what the alternative feels like — and what it costs in the weeks afterwards.

Remember this

  • Write down the current answer before you touch anything. No number, no refactor.
  • A characterisation test locks in current behaviour, not correct behaviour.
  • Move one piece at a time, and rerun the test after each move.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install scikit-learn

Verified with scikit-learn 1.7.2, numpy 1.26.4, Python 3.10, CPU. Runs in about two seconds. The dataset is generated in memory so the whole demonstration is self-contained.

The before, the after, and the mistake

A real 600-line script does not fit on a page, so here is its shape at 15 lines. Everything that matters about the procedure survives the shrinking.

refactor_demo.py
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# ---------- BEFORE: one function that does everything ----------
def train_everything(seed=0):
    rng = np.random.default_rng(seed)
    X = rng.normal(size=(400, 5)); X[:, 3] *= 50
    y = (X[:, 0] + 0.5*X[:, 1] - 0.02*X[:, 3] + rng.normal(0, 0.7, 400) > 0).astype(int)
    Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.3, random_state=seed)
    sc = StandardScaler().fit(Xtr)
    m = LogisticRegression(C=0.5, max_iter=1000).fit(sc.transform(Xtr), ytr)
    return roc_auc_score(yte, m.predict_proba(sc.transform(Xte))[:, 1])

# ---------- AFTER: four named steps, same behaviour ----------
def load_data(seed):
    rng = np.random.default_rng(seed)
    X = rng.normal(size=(400, 5)); X[:, 3] *= 50
    y = (X[:, 0] + 0.5*X[:, 1] - 0.02*X[:, 3] + rng.normal(0, 0.7, 400) > 0).astype(int)
    return X, y

def split(X, y, seed):
    return train_test_split(X, y, test_size=0.3, random_state=seed)

def fit_model(Xtr, ytr, C):
    scaler = StandardScaler().fit(Xtr)                 # fitted on train only
    return scaler, LogisticRegression(C=C, max_iter=1000).fit(scaler.transform(Xtr), ytr)

def evaluate(scaler, model, Xte, yte):
    return roc_auc_score(yte, model.predict_proba(scaler.transform(Xte))[:, 1])

def run(seed=0, C=0.5):
    X, y = load_data(seed)
    Xtr, Xte, ytr, yte = split(X, y, seed)
    scaler, model = fit_model(Xtr, ytr, C)
    return evaluate(scaler, model, Xte, yte)

# ---------- AFTER, with one line lost in the move ----------
def split_buggy(X, y, seed):
    return train_test_split(X, y, test_size=0.3)       # random_state quietly dropped

def run_buggy(seed=0, C=0.5):
    X, y = load_data(seed)
    Xtr, Xte, ytr, yte = split_buggy(X, y, seed)
    scaler, model = fit_model(Xtr, ytr, C)
    return evaluate(scaler, model, Xte, yte)

BASELINE = 0.925645                                    # recorded from the old script

def matches(value):
    return abs(value - BASELINE) < 1e-6

print(f"before   AUC {train_everything():.6f}")
print(f"after    AUC {run():.6f}   matches baseline: {matches(run())}")

buggy = [run_buggy() for _ in range(3)]
for i, v in enumerate(buggy, 1):
    print(f"buggy #{i}                    matches baseline: {matches(v)}")
print(f"the three buggy runs even agree with each other: {len(set(buggy)) == 1}")
Output
before   AUC 0.925645
after    AUC 0.925645   matches baseline: True
buggy #1                    matches baseline: False
buggy #2                    matches baseline: False
buggy #3                    matches baseline: False
the three buggy runs even agree with each other: False

The buggy version's actual AUC values are deliberately not printed: they differ on every run and on every machine, which is the entire symptom.

The walkthrough

The bug is one missing keyword argument. split_buggy dropped random_state=seed while the line was being moved into its own function. Nothing raises. Nothing looks wrong in review. The model still trains, and the score is in a plausible range every single time.

Without the baseline, that bug ships. Over two thousand runs of the buggy version the AUC averages 0.937, sits between 0.91 and 0.96 about ninety percent of the time, and occasionally strays as far as 0.87 or 0.99 — all depending on the split it happened to draw. Anyone eyeballing "about 0.93" accepts it, and every experiment afterwards is comparing against a moving target.

With the baseline, the bug is caught in one second. matches baseline: False on the first run after the move, pointing at exactly the change that caused it, because you moved one thing.

The last line finds the deeper problem. The three buggy runs do not even agree with each other. A refactor that introduces non-determinism is worse than one that introduces a wrong constant, because there is nothing stable left to compare against — see random seeds and reproducibility.

Turn the baseline into a real test. Save the script above as refactor_demo.py, put this beside it, and run python -m pytest -q test_refactor.py:

test_refactor.py
from refactor_demo import run          # the script above, saved beside this file

def test_pipeline_unchanged():
    assert abs(run(seed=0, C=0.5) - 0.925645) < 1e-6
Output
.                                                                        [100%]
1 passed in 2.03s

(The elapsed time is your machine's, not a number to match.) Then run it after every single move. Add a second one at a different seed, so a change that only holds for seed=0 cannot hide.

The order to pull things apart

Long scripts come apart in a reliable order, easiest and safest first.

  1. Pull out constants. Magic numbers to named module-level values, then to a config file. Zero behaviour change, immediate readability gain.
  2. Pull out pure functions — pieces that take inputs and return outputs with no file reads or global state. evaluate above is the model. These are also the pieces that become testable.
  3. Pull out I/O behind a boundary. One load_data function, one save_model function. Once they are isolated, the rest of the pipeline can be tested with tiny in-memory fixtures.
  4. Pull out the training loop last. It has the most hidden state, and it is where an accidental change does the most damage.
  5. Delete dead code, and delete it separately. Commented-out experiments go in a commit of their own, so a bisect can tell "we deleted things" from "we moved things".

Common mistakes

Refactoring and improving in the same commit. If the number changes, you cannot tell whether the improvement worked or the refactor broke something. Behaviour-preserving changes and behaviour-changing changes are two commits, always.

A baseline on one seed only. Some bugs are invisible at seed=0. Record two or three seeds, and treat the set of them as the fence.

A tolerance that is too loose. abs(new - old) < 0.01 passes for the buggy version above about a third of the time — a test that fails two runs in three is a test nobody trusts, and it will be deleted rather than believed. Make the tolerance as tight as your pipeline's genuine determinism allows — 1e-6 for CPU scikit-learn, looser for GPU floating point, and document why.

Trusting a green test on a slow script. If the characterisation test takes forty minutes, nobody runs it between moves, so it protects nothing. Build the fence on a tiny subset that finishes in seconds, and run the full version once at the end.

Splitting into functions that still share globals. Four functions passing state through module-level variables is one function wearing four hats. Every seam must pass its data as arguments — getting out of the notebook makes the same point about the same mistake.

Try it yourself

Introduce a second realistic bug: in fit_model, drop the C argument so LogisticRegression() takes its default. Run the check. On this dataset the AUC barely moves — so extend your fence to also assert on the model's coefficients, and watch the same bug get caught instantly. That exercise teaches which quantity to pin, which is the hardest judgement in this lesson.

What to learn next

Researcher — Mathematics and papers.

Characterisation testing as a formal device

The technique is Feathers (2004), Working Effectively with Legacy Code, where a characterisation test is defined as a test that documents actual behaviour rather than intended behaviour, written specifically to make a system changeable. The distinction matters: a correctness test can fail because the code is wrong, whereas a characterisation test can only fail because the code changed, which is exactly the invariant a refactor must preserve.

Fowler (1999; 2nd ed. 2018), Refactoring, supplies the complementary discipline: behaviour-preserving transformations applied one at a time under a green test suite, with the catalogue of named moves (Extract Function, Introduce Parameter Object, Replace Temp with Query) that the "order to pull things apart" above is an ML-specific specialisation of. The theoretical guarantee is weak — behaviour preservation is undecidable in general, and refactoring tools verify only syntactic preconditions — so the empirical fence does the real work.

ML code weakens the guarantee further in a specific way: many pipelines are stochastic, so "same behaviour" is only checkable modulo a seed. The pinning technique therefore has an implicit precondition that the pipeline is deterministic given its seed, which is why the buggy example's failure to reproduce itself is the more serious finding. Pineau et al. (2021), Improving Reproducibility in Machine Learning Research (JMLR), and the NeurIPS reproducibility programme make determinism-given-seed the entry requirement for any claim of a controlled change.

Why ML code decays faster than ordinary code

Sculley et al. (2015), Hidden Technical Debt in Machine Learning Systems (NeurIPS), catalogues the mechanisms, and several bear directly on the 600-line script: glue code (the majority of a mature ML system is plumbing around a small modelling core), pipeline jungles (accreted data preparation with no clean seams), configuration debt (hyperparameters spread across code, flags and defaults), and entanglement, summarised as CACE — changing anything changes everything, because a model's parameters depend jointly on every preprocessing decision.

CACE is precisely why the pin-the-number method is more necessary here than in ordinary software. In a conventional module, moving a function cannot change the output of an unrelated function. In a training pipeline, changing the order of two preprocessing operations changes fitted parameters, which changes every downstream metric, with no type error and no exception anywhere.

Breck et al. (2017), The ML Test Score (IEEE Big Data), operationalise the response as a rubric of 28 tests across data, model, infrastructure and monitoring. Two items map onto this lesson exactly: "model specs are unit tested" and "the full ML pipeline is integration tested", both of which the characterisation test is the cheapest available approximation to when starting from an untested script.

Choosing what to pin

The pinned quantity determines which bugs the fence can catch, and the exercise at the end of the developer block is a small demonstration of a general result: a scalar metric is a lossy hash of the pipeline's behaviour. A summary statistic like AUC is invariant to many real changes (it is rank-based, so any monotone transform of the scores leaves it fixed), so a fence built on it is blind to calibration changes (reliability diagrams) and to hyperparameter drift whose effect happens to be small on this sample.

Stronger fences, in increasing sensitivity and cost:

  • Per-example predictions — hash the full prediction vector. Catches anything that changes any output, including changes with zero effect on the aggregate metric. This is the ML analogue of approval testing / golden files.
  • Fitted parameters — coefficients, or a hash of the state dict. Catches optimisation-path changes that the predictions round away.
  • Intermediate artefacts — the feature matrix after preprocessing, checked by shape, dtype and column-wise summary statistics, which localises a failure to a stage instead of reporting it at the end.

The trade-off is brittleness: parameter-level fences fail on legitimate library upgrades and hardware changes, generating maintenance noise. The standard resolution is tiered — a fast, tolerant metric fence run on every move, and a strict artefact fence run before merge and updated deliberately with a recorded justification, which is the same pattern data validation uses for schemas.

What to learn next