Make solver resume work and keep the training history

[marisbasha]

Summary

Fixes #23. flyvis train-single ... resume=true failed in three places: reopening the network directory, the rest of recover() after the call from the issue, and the restored iteration. Continuing training also overwrote the stored loss history.

Changes

Reopening a network directory

MultiTaskSolver put delete_if_exists into the config it passes to NetworkDir. datamate stores the key when it creates the directory. When the directory already exists, it removes the key from the passed config before comparing. So opening the same directory again with the same config raised FileExistsError (incompatible config: passed=['-delete_if_exists'], stored=['+delete_if_exists: False']). This happens with datamate 0.2.7 and 1.0.

The flag is now passed through datamate's delete_if_exists context instead of the config. The context is imported from datamate, or from datamate.directory for datamate < 1.0. It is only entered when deletion is requested, so an enclosing delete_if_exists context still applies as before.

Directories created before this change still have the key stored, and datamate always drops it from the passed config of an existing directory. When the stored flag is the only difference to the passed config, MultiTaskSolver now opens such a directory by path. Any other difference still raises.

Rerunning train-single on a trained network

The FileExistsError above also stopped an accidental rerun of a trained network. Without it, the same command would open the directory and train from iteration 0 over the existing checkpoints. train_single.py now raises a FileExistsError when the directory already has checkpoints and resume is not set. The message names resume=true and delete_if_exists=true.

recover()

Besides the four-argument resolve_checkpoints call reported in #23, the body read checkpoint.index, checkpoints.index and checkpoints.path, none of which exist. So fixing only the call failed on the next line. recover() now resolves "best" with best_checkpoint_default_fn, and an int as a position into the sorted checkpoints. So -1 is the last one, as used by train-single.

Restored iteration

Checkpoints store the last completed iteration, self.iteration - 1, and recover() assigned it back unchanged. A resume repeated one iteration. Resuming from the initial checkpoint set the iteration to -1, where the scheduler picked the final learning rate. recover() now sets self.iteration to the stored value plus one, and to 0 for checkpoints without an iteration, such as the pretrained ones.

Checkpoint index after recovery

checkpoint() incremented _last_chkpt_ind and _curr_chkpt_ind separately. After recovering an older or the best checkpoint, the next checkpoint was saved under the last index but recorded the current index as the recovered one plus one. A later recover() of that index then returned early. The current index is now set to the newly saved checkpoint.

Training history

train() started the loss and activity lists empty on every call. It then overwrote dir.loss, dir.activity, dir.activity_min, dir.activity_max and dir.loss_<task> at each checkpoint. Continuing training with a larger n_iters, or resuming, kept only the iterations of the last call, and EnsembleView.training_loss showed partial curves. The lists now start from the stored history up to the current iteration.

Testing

Four new tests in tests/test_solver.py train a small solver on the mock Sintel data.

• The first trains 2 iterations, raises n_iters to 4, and trains again. It checks that dir.loss has 4 entries and that the first 2 are unchanged.
• The second trains 2 iterations and opens the same directory with a new solver. It recovers checkpoint -1, 0 and "best", and checks the iteration, checkpoint index and network parameters. It then continues to 4 iterations, and checks that delete_if_exists=True still clears the directory.
• The third recovers checkpoint 0 after training, makes a checkpoint, and checks that the current and last index are both the new one.
• The fourth creates a directory with delete_if_exists in its config, as earlier versions did. It checks that MultiTaskSolver opens it with the same config, and still raises for a different n_iters.

A new tests/test_train_single.py runs train_single.py as a subprocess on the mock data. It writes the initial checkpoint with train=false checkpoint_only=true. It checks that the same command without resume is refused, then that resume=true trains to n_iters with the full loss history.

All five fail on main. On the branch they pass with datamate 1.0.0, 0.2.7 and the current datamate main. With the CLI check removed, the subprocess test fails because the rerun trains over the checkpoints.

I also ran train-single by hand on the mock data, with each run stopped at its third checkpoint and then resumed with resume=true. Three cases: a new directory, a directory recreated with delete_if_exists=true, and a new directory created with delete_if_exists=true. On main every resume fails, with FileExistsError or the TypeError from #23. On the branch each resumes from the last checkpoint and finishes with the full loss history. That holds with datamate 1.0.0, 0.2.7 and the current datamate main. A run started with main's train-single, which stores delete_if_exists: false, also resumes on the branch with datamate 1.0.0 and 0.2.7.

Full suite (-m "not require_download and not require_large_download and not gpu", test_examples.py and test_sintel.py excluded): 196 passed, 12 skipped on main; 201 passed, 12 skipped on the branch, with datamate 1.0.0 and with 0.2.7. require_download tests against the pretrained models: 21 passed on both. tests/test_sintel.py, run offline with the mock data: 14 passed on both. ruff check and ruff format (0.5.5) are clean on the changed files.

[Copilot]

Copilot review overview

🟡 Changes recommended

Legacy directories remain unable to resume, and checkpoint bookkeeping becomes inconsistent after recovering a non-latest checkpoint.

Review effort: Balanced
Findings: 2

Open (2)

• Handle legacy metadata to resume existing training directories · New
• Synchronize current checkpoint index after recovery · New

What changed in this PR

Fixes solver resume behavior and preserves training history across continued runs.

Changes:

• Corrects checkpoint recovery and iteration restoration.
• Preserves prior loss and activity history.
• Adds CLI safeguards and resume regression tests.

File	Description
flyvis/​solver.py	Updates directory handling, history retention, and recovery.
flyvis_cli/​training/​train_single.py	Prevents accidental checkpoint overwrite.
tests/​test_solver.py	Tests continued training and recovery.
tests/​test_train_single.py	Tests CLI resume behavior.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Copilot review overview

🟢 Approval recommended

The resume failures and history loss are addressed coherently with focused regression coverage.

Review effort: Balanced
Findings: None

Resolved since last review (2)

• Synchronize current checkpoint index after recovery
• Handle legacy metadata to resume existing training directories