Skip to content

Make model tests more scalable#154

Merged
kmaziarz merged 3 commits into
mainfrom
kmaziarz/test-models-in-separate-process
Jun 10, 2026
Merged

Make model tests more scalable#154
kmaziarz merged 3 commits into
mainfrom
kmaziarz/test-models-in-separate-process

Conversation

@kmaziarz

@kmaziarz kmaziarz commented Jun 9, 2026

Copy link
Copy Markdown
Contributor

With the upcoming integration of new models from RetroChimera, I found that the way we currently run single-step models tests does not scale, especially when larger models are included. First, we've been running tests for all models in a single process, which meant all models are loaded simultaneously; this is inherently bounded by total memory available on the CI runners. Second, downloading many heavy model checkpoints is slow, and has a large variance. This PR attempts to resolve both issues to enable scaling to more single-step models in the future.

To avoid keeping all models in memory simultaneously, one could try forcing garbage collection, but this does not clean up all state, especially for models that involve multiprocessing or call into non-Python-native libraries. This PR instead proposes to run the tests in separate processes, which ensures all state is cleaned up. Previously, in test_models.py we used to test every model twice, to also test that loading one model doesn't make another model unusable; in the new setup, this wouldn't make sense, as models are now completely separate. We do lose a bit of test coverage as we do not test interactions between the models, but doing so exhaustively would not scale anyway; also, for the majority of usecases, only one model is used at a time. Finally, to reduce the burden of downloading model checkpoints, I use the cache action to cache the checkpoint directory (keyed on the contents of default_checkpoint_ids.yml), so that CI runs that don't add new models or change existing ones can benefit from much faster checkpoint download.

@kmaziarz kmaziarz requested a review from AustinT June 9, 2026 12:03
@kmaziarz kmaziarz requested a review from jla-gardner June 9, 2026 15:34
@kmaziarz kmaziarz merged commit 5e26cdb into main Jun 10, 2026
37 checks passed
@kmaziarz kmaziarz deleted the kmaziarz/test-models-in-separate-process branch June 10, 2026 10:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants