From a95b7b88d1240a2f57731345d0ff71a2490bb3d3 Mon Sep 17 00:00:00 2001 From: Pol Date: Wed, 21 Jan 2026 12:47:39 +0100 Subject: [PATCH 1/7] Run trainings lazily in tests --- .gitignore | 3 ++ tests/cli/__init__.py | 21 -------- tests/cli/test_eval_model.py | 6 +-- tests/cli/test_export_model.py | 2 +- tests/cli/test_train_model.py | 32 +++++++----- tests/conftest.py | 58 ++++++++++++++++++++++ tests/resources/generate-outputs.sh | 30 ----------- tests/resources/run_trainings.sh | 77 +++++++++++++++++++++++++++++ tests/utils/__init__.py | 7 --- tests/utils/test_augmentation.py | 2 +- tests/utils/test_evaluate_model.py | 2 +- tests/utils/test_io.py | 71 +++++++++++++++----------- tests/utils/test_long_range.py | 2 +- tests/utils/test_output_gradient.py | 2 +- tests/utils/test_scaler.py | 4 +- tox.ini | 2 - 16 files changed, 209 insertions(+), 112 deletions(-) delete mode 100755 tests/resources/generate-outputs.sh create mode 100755 tests/resources/run_trainings.sh diff --git a/.gitignore b/.gitignore index df94cece1f..dbb99fd266 100644 --- a/.gitignore +++ b/.gitignore @@ -174,6 +174,9 @@ outputs/ examples/basic_usage/*.xyz extensions/ +# lock files for model training +tests/resources/*.trainlock + # sphinx gallery docs/src/generated_examples/ *execution_times* diff --git a/tests/cli/__init__.py b/tests/cli/__init__.py index a3dc25c3d0..e69de29bb2 100644 --- a/tests/cli/__init__.py +++ b/tests/cli/__init__.py @@ -1,21 +0,0 @@ -from pathlib import Path - -from metatrain.utils.architectures import get_default_hypers - - -MODEL_HYPERS = get_default_hypers("soap_bpnn")["model"] - -RESOURCES_PATH = Path(__file__).parents[1] / "resources" - -DATASET_PATH_QM9 = RESOURCES_PATH / "qm9_reduced_100.xyz" -DATASET_PATH_ETHANOL = RESOURCES_PATH / "ethanol_reduced_100.xyz" -DATASET_PATH_CARBON = RESOURCES_PATH / "carbon_reduced_100.xyz" -DATASET_PATH_QM7X = RESOURCES_PATH / "qm7x_reduced_100.xyz" -DATASET_PATH_DOS = RESOURCES_PATH / "dos_100.xyz" -EVAL_OPTIONS_PATH = RESOURCES_PATH / "eval.yaml" -MODEL_PATH = RESOURCES_PATH / "model-32-bit.pt" -MODEL_PATH_64_BIT = RESOURCES_PATH / "model-64-bit.ckpt" -MODEL_PATH_PET = RESOURCES_PATH / "model-pet.ckpt" -OPTIONS_PATH = RESOURCES_PATH / "options.yaml" -OPTIONS_PET_PATH = RESOURCES_PATH / "options-pet.yaml" -OPTIONS_EXTRA_DATA_PATH = RESOURCES_PATH / "options-extra-data.yaml" diff --git a/tests/cli/test_eval_model.py b/tests/cli/test_eval_model.py index ddf9e9e028..122cb0799a 100644 --- a/tests/cli/test_eval_model.py +++ b/tests/cli/test_eval_model.py @@ -18,11 +18,11 @@ from metatrain.utils.data.writers import DiskDatasetWriter from metatrain.utils.neighbor_lists import get_system_with_neighbor_lists -from . import EVAL_OPTIONS_PATH, MODEL_HYPERS, MODEL_PATH, RESOURCES_PATH +from ..conftest import EVAL_OPTIONS_PATH, MODEL_HYPERS, RESOURCES_PATH @pytest.fixture -def model(): +def model(MODEL_PATH): return torch.jit.load(MODEL_PATH) @@ -31,7 +31,7 @@ def options(): return OmegaConf.load(EVAL_OPTIONS_PATH) -def test_eval_cli(monkeypatch, tmp_path): +def test_eval_cli(monkeypatch, tmp_path, MODEL_PATH): """Test succesful run of the eval script via the CLI with default arguments""" monkeypatch.chdir(tmp_path) shutil.copy(RESOURCES_PATH / "qm9_reduced_100.xyz", "qm9_reduced_100.xyz") diff --git a/tests/cli/test_export_model.py b/tests/cli/test_export_model.py index e6f7bd368e..ffcdaed75e 100644 --- a/tests/cli/test_export_model.py +++ b/tests/cli/test_export_model.py @@ -19,7 +19,7 @@ from metatrain.utils.architectures import find_all_architectures from metatrain.utils.io import load_model -from . import RESOURCES_PATH +from ..conftest import RESOURCES_PATH @pytest.mark.parametrize("path", [Path("exported.pt"), "exported.pt"]) diff --git a/tests/cli/test_train_model.py b/tests/cli/test_train_model.py index 2c98687053..5d6aff8374 100644 --- a/tests/cli/test_train_model.py +++ b/tests/cli/test_train_model.py @@ -29,14 +29,12 @@ from metatrain.utils.pydantic import MetatrainValidationError from metatrain.utils.testing._utils import WANDB_AVAILABLE -from . import ( +from ..conftest import ( DATASET_PATH_CARBON, DATASET_PATH_DOS, DATASET_PATH_ETHANOL, DATASET_PATH_QM7X, DATASET_PATH_QM9, - MODEL_PATH_64_BIT, - MODEL_PATH_PET, OPTIONS_EXTRA_DATA_PATH, OPTIONS_PATH, OPTIONS_PET_PATH, @@ -667,7 +665,7 @@ def test_same_name_targets_extra_data( train_model(options_extra) -def test_restart(options, monkeypatch, tmp_path): +def test_restart(options, monkeypatch, tmp_path, MODEL_PATH_64_BIT): """Test that continuing training from a checkpoint runs without an error raise.""" monkeypatch.chdir(tmp_path) shutil.copy(DATASET_PATH_QM9, "qm9_reduced_100.xyz") @@ -675,7 +673,7 @@ def test_restart(options, monkeypatch, tmp_path): train_model(options, restart_from=MODEL_PATH_64_BIT) -def test_finetune(options_pet, caplog, monkeypatch, tmp_path): +def test_finetune(options_pet, caplog, monkeypatch, tmp_path, MODEL_PATH_PET): monkeypatch.chdir(tmp_path) options_pet["architecture"]["training"]["finetune"] = { @@ -695,7 +693,7 @@ def test_finetune(options_pet, caplog, monkeypatch, tmp_path): assert f"Starting finetuning from '{MODEL_PATH_PET}'" in caplog.text -def test_transfer_learn(options_pet, caplog, monkeypatch, tmp_path): +def test_transfer_learn(options_pet, caplog, monkeypatch, tmp_path, MODEL_PATH_PET): monkeypatch.chdir(tmp_path) options_pet_transfer_learn = copy.deepcopy(options_pet) @@ -719,7 +717,9 @@ def test_transfer_learn(options_pet, caplog, monkeypatch, tmp_path): assert f"Starting finetuning from '{MODEL_PATH_PET}'" in caplog.text -def test_transfer_learn_with_forces(options_pet, caplog, monkeypatch, tmp_path): +def test_transfer_learn_with_forces( + options_pet, caplog, monkeypatch, tmp_path, MODEL_PATH_PET +): monkeypatch.chdir(tmp_path) options_pet_transfer_learn = copy.deepcopy(options_pet) @@ -752,7 +752,9 @@ def test_transfer_learn_with_forces(options_pet, caplog, monkeypatch, tmp_path): assert f"Starting finetuning from '{MODEL_PATH_PET}'" in caplog.text -def test_transfer_learn_variant(options_pet, caplog, monkeypatch, tmp_path): +def test_transfer_learn_variant( + options_pet, caplog, monkeypatch, tmp_path, MODEL_PATH_PET +): monkeypatch.chdir(tmp_path) options_pet_transfer_learn = copy.deepcopy(options_pet) @@ -782,7 +784,9 @@ def test_transfer_learn_variant(options_pet, caplog, monkeypatch, tmp_path): assert f"Starting finetuning from '{MODEL_PATH_PET}'" in caplog.text -def test_transfer_learn_inherit_heads(options_pet, caplog, monkeypatch, tmp_path): +def test_transfer_learn_inherit_heads( + options_pet, caplog, monkeypatch, tmp_path, MODEL_PATH_PET +): monkeypatch.chdir(tmp_path) options_pet_transfer_learn = copy.deepcopy(options_pet) @@ -808,7 +812,7 @@ def test_transfer_learn_inherit_heads(options_pet, caplog, monkeypatch, tmp_path def test_transfer_learn_inherit_heads_invalid_source( - options_pet, caplog, monkeypatch, tmp_path + options_pet, caplog, monkeypatch, tmp_path, MODEL_PATH_PET ): monkeypatch.chdir(tmp_path) @@ -838,7 +842,7 @@ def test_transfer_learn_inherit_heads_invalid_source( def test_transfer_learn_inherit_heads_invalid_destination( - options_pet, caplog, monkeypatch, tmp_path + options_pet, caplog, monkeypatch, tmp_path, MODEL_PATH_PET ): monkeypatch.chdir(tmp_path) @@ -863,7 +867,9 @@ def test_transfer_learn_inherit_heads_invalid_destination( @pytest.mark.parametrize("move_folder", [True, False]) -def test_restart_auto(options, caplog, monkeypatch, tmp_path, move_folder): +def test_restart_auto( + options, caplog, monkeypatch, tmp_path, move_folder, MODEL_PATH_64_BIT +): """Test that continuing with the `auto` keyword results in a continuation from the most recent checkpoint.""" monkeypatch.chdir(tmp_path) @@ -912,7 +918,7 @@ def test_restart_auto_no_outputs(options, caplog, monkeypatch, tmp_path): assert "Restart training from" not in caplog.text -def test_restart_different_dataset(options, monkeypatch, tmp_path): +def test_restart_different_dataset(options, monkeypatch, tmp_path, MODEL_PATH_64_BIT): """Test that continuing training from a checkpoint runs without an error raise with a different dataset than the original.""" monkeypatch.chdir(tmp_path) diff --git a/tests/conftest.py b/tests/conftest.py index 922cb36ccd..121b654916 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,8 +1,66 @@ import math import os +import subprocess +from pathlib import Path + +import pytest + +from metatrain.utils.architectures import get_default_hypers def pytest_xdist_auto_num_workers(): """Limit the number of workers used by pytest""" n_processes = os.cpu_count() or 1 return min(12, math.ceil(n_processes * 0.8)) + +MODEL_HYPERS = get_default_hypers("soap_bpnn")["model"] + +# ------------------------------- +# PATHS TO RESOURCES +# ------------------------------- +RESOURCES_PATH = Path(__file__).parent / "resources" + +DATASET_PATH_QM9 = RESOURCES_PATH / "qm9_reduced_100.xyz" +DATASET_PATH_ETHANOL = RESOURCES_PATH / "ethanol_reduced_100.xyz" +DATASET_PATH_CARBON = RESOURCES_PATH / "carbon_reduced_100.xyz" +DATASET_PATH_QM7X = RESOURCES_PATH / "qm7x_reduced_100.xyz" +DATASET_PATH_DOS = RESOURCES_PATH / "dos_100.xyz" +EVAL_OPTIONS_PATH = RESOURCES_PATH / "eval.yaml" +OPTIONS_PATH = RESOURCES_PATH / "options.yaml" +OPTIONS_PET_PATH = RESOURCES_PATH / "options-pet.yaml" +OPTIONS_EXTRA_DATA_PATH = RESOURCES_PATH / "options-extra-data.yaml" + +# ------------------------------- +# PATHS TO TRAINED MODELS +# ------------------------------- +# These files are generated from training on each test run +# and take some time to generate. Therefore, we make them a +# fixture so that they are generated lazily only if there +# is a test that requires them. Since we support parallel +# test runs, we use the unique identifier for the test run to +# so that the multiple workers do not replicate work and wait +# for each other. +def ensure_path(mode: str, uid: str) -> Path: + """Checks if the path for a model exists, and if not + runs the training script to generate it.""" + path = RESOURCES_PATH / f"model-{mode}-{uid}.pt" + if not path.exists(): + subprocess.run( + ["bash", str(RESOURCES_PATH / "run_trainings.sh"), mode, uid], check=True + ) + return path + + +@pytest.fixture(scope="session") +def MODEL_PATH(testrun_uid) -> Path: + return ensure_path("32-bit", testrun_uid) + + +@pytest.fixture(scope="session") +def MODEL_PATH_64_BIT(testrun_uid) -> Path: + return ensure_path("64-bit", testrun_uid) + + +@pytest.fixture(scope="session") +def MODEL_PATH_PET(testrun_uid) -> Path: + return ensure_path("pet", testrun_uid) diff --git a/tests/resources/generate-outputs.sh b/tests/resources/generate-outputs.sh deleted file mode 100755 index 921895d1a6..0000000000 --- a/tests/resources/generate-outputs.sh +++ /dev/null @@ -1,30 +0,0 @@ -#!/bin/bash -set -eux - -echo "Generating data for testing..." - -ROOT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) - -cd "$ROOT_DIR" - -mtt train options.yaml -o model-32-bit.pt -r base_precision=32 -mtt train options.yaml -o model-64-bit.pt -r base_precision=64 -mtt train options-pet.yaml -o model-pet.pt - -set +x # disable command echoing for sensitive private token check -TOKEN_PRESENT=false -if [[ -n "${HUGGINGFACE_TOKEN_METATRAIN:-}" ]]; then - TOKEN_PRESENT=true -fi -set -x - -if [ $TOKEN_PRESENT = true ]; then - hf upload \ - "metatensor/metatrain-test" \ - "model-32-bit.ckpt" \ - "model.ckpt" \ - --commit-message="Overwrite test model with new version" \ - --token="$HUGGINGFACE_TOKEN_METATRAIN" -else - echo "HUGGINGFACE_TOKEN_METATRAIN is not set, skipping upload." -fi diff --git a/tests/resources/run_trainings.sh b/tests/resources/run_trainings.sh new file mode 100755 index 0000000000..e5062ba2bd --- /dev/null +++ b/tests/resources/run_trainings.sh @@ -0,0 +1,77 @@ +#!/bin/bash +set -eux + +ROOT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) + +MODE=$1 +TRAIN_ID=$2 + +if [ -z "$MODE" ]; then + echo "Error: First argument of the script is the mode" + echo " There is no first argument, please set it to " + echo " '32-bit', '64-bit' or 'pet'" + exit 1 +fi + +# If there is a TRAIN_ID, implement a lock file +if [ -n "$TRAIN_ID" ]; then + echo "Creating lockfile" + LOCKFILE="$ROOT_DIR/$MODE-$TRAIN_ID.trainlock" + if [ -f $LOCKFILE ]; then + # Wait until the lock file is removed + while [ -f $LOCKFILE ]; do + sleep 5 + done + exit 0 + else + touch $LOCKFILE + fi +fi + +echo "Clearing previous generated files..." +# Clean previous generated files +rm $ROOT_DIR/model-$MODE-*.pt $ROOT_DIR/model-$MODE-*.ckpt || true +rm $ROOT_DIR/$MODE-*.lock || true + +echo "Generating data for testing..." + +cd "$ROOT_DIR" +# The generated files are uniquely identified by the TRAIN_ID passed as second argument, +# in this way a test run can know if the files have already been generated. +if [ "$MODE" == "32-bit" ]; then + mtt train options.yaml -o model-32-bit-$TRAIN_ID.pt -r base_precision=32 +elif [ "$MODE" == "64-bit" ]; then + mtt train options.yaml -o model-64-bit-$TRAIN_ID.pt -r base_precision=64 +elif [ "$MODE" == "pet" ]; then + mtt train options-pet.yaml -o model-pet-$TRAIN_ID.ckpt -r base_precision=pet +else + echo "Error: Unknown training mode (first argument): '$MODE'" + echo " Please set it to '32-bit', '64-bit' or 'pet'" + exit 1 +fi + +rm $LOCKFILE || true + +# If the mode is 32-bit, we will try to upload the model to Hugging Face, +# otherwise we are done here +if [ "$MODE" != "32-bit" ]; then + exit 0 +fi + +set +x # disable command echoing for sensitive private token check +TOKEN_PRESENT=false +if [[ -n "${HUGGINGFACE_TOKEN_METATRAIN:-}" ]]; then + TOKEN_PRESENT=true +fi +set -x + +if [ $TOKEN_PRESENT = true ]; then + hf upload \ + "metatensor/metatrain-test" \ + "model-32-bit-$TRAIN_ID.ckpt" \ + "model.ckpt" \ + --commit-message="Overwrite test model with new version" \ + --token="$HUGGINGFACE_TOKEN_METATRAIN" +else + echo "HUGGINGFACE_TOKEN_METATRAIN is not set, skipping upload." +fi diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py index bfe83e9403..8b13789179 100644 --- a/tests/utils/__init__.py +++ b/tests/utils/__init__.py @@ -1,8 +1 @@ -from pathlib import Path -from metatrain.utils.architectures import get_default_hypers - - -MODEL_HYPERS = get_default_hypers("soap_bpnn")["model"] - -RESOURCES_PATH = Path(__file__).parents[1] / "resources" diff --git a/tests/utils/test_augmentation.py b/tests/utils/test_augmentation.py index faabad5bec..fea0651016 100644 --- a/tests/utils/test_augmentation.py +++ b/tests/utils/test_augmentation.py @@ -11,7 +11,7 @@ from metatrain.utils.data import DatasetInfo, DiskDataset, TargetInfo from metatrain.utils.data.target_info import get_generic_target_info -from . import RESOURCES_PATH +from ..conftest import RESOURCES_PATH @pytest.fixture diff --git a/tests/utils/test_evaluate_model.py b/tests/utils/test_evaluate_model.py index 9f86c74c47..5d00f3faf6 100644 --- a/tests/utils/test_evaluate_model.py +++ b/tests/utils/test_evaluate_model.py @@ -10,7 +10,7 @@ get_system_with_neighbor_lists, ) -from . import MODEL_HYPERS, RESOURCES_PATH +from ..conftest import MODEL_HYPERS, RESOURCES_PATH @pytest.mark.parametrize("training", [True, False]) diff --git a/tests/utils/test_io.py b/tests/utils/test_io.py index de190aff90..5f37ed3821 100644 --- a/tests/utils/test_io.py +++ b/tests/utils/test_io.py @@ -14,7 +14,36 @@ trainer_from_checkpoint, ) -from . import RESOURCES_PATH + +@pytest.fixture(scope="module", params=["pathlib", "str", "file_url"]) +def load_path(request, MODEL_PATH): + """Fixture that provides the model checkpoint path in different formats. + + This is meant to be used in the `test_load_model_exported` test to + check that loading works with all of these formats. + """ + if request.param == "pathlib": + return MODEL_PATH + elif request.param == "str": + return str(MODEL_PATH) + elif request.param == "file_url": + return f"file:{str(MODEL_PATH)}" + + +@pytest.fixture(scope="module", params=["pathlib", "str", "file_url"]) +def load_path_ckpt(MODEL_PATH, request): + """Fixture that provides the model checkpoint path in different formats. + + This is meant to be used in the `test_load_model_checkpoint` test to + check that loading works with all of these formats. + """ + ckpt_path = MODEL_PATH.with_suffix(".ckpt") + if request.param == "pathlib": + return ckpt_path + elif request.param == "str": + return str(ckpt_path) + elif request.param == "file_url": + return f"file:{str(ckpt_path)}" @pytest.mark.parametrize("filename", ["example.txt", Path("example.txt")]) @@ -35,21 +64,13 @@ def test_warning_on_missing_suffix(filename): assert isinstance(result, type(filename)) -def test_is_exported_file(): - assert is_exported_file(RESOURCES_PATH / "model-32-bit.pt") - assert not is_exported_file(RESOURCES_PATH / "model-32-bit.ckpt") +def test_is_exported_file(MODEL_PATH): + assert is_exported_file(MODEL_PATH) + assert not is_exported_file(MODEL_PATH.with_suffix(".ckpt")) -@pytest.mark.parametrize( - "path", - [ - RESOURCES_PATH / "model-32-bit.ckpt", - str(RESOURCES_PATH / "model-32-bit.ckpt"), - f"file:{str(RESOURCES_PATH / 'model-32-bit.ckpt')}", - ], -) -def test_load_model_checkpoint(path): - model = load_model(path) +def test_load_model_checkpoint(load_path_ckpt): + model = load_model(load_path_ckpt) assert type(model) is SoapBpnn # TODO: test that weights are the expected if loading with `context == 'export'`. @@ -57,9 +78,9 @@ def test_load_model_checkpoint(path): # currently weights of the `"export"` and the `"restart"` context are the same... -def test_load_model_checkpoint_wrong_version(monkeypatch, tmp_path): +def test_load_model_checkpoint_wrong_version(monkeypatch, tmp_path, MODEL_PATH_64_BIT): monkeypatch.chdir(tmp_path) - path = RESOURCES_PATH / "model-64-bit.ckpt" + path = MODEL_PATH_64_BIT.with_suffix(".ckpt") model = torch.load(path, weights_only=False, map_location="cpu") model["model_ckpt_version"] = 5000000 @@ -76,9 +97,11 @@ def test_load_model_checkpoint_wrong_version(monkeypatch, tmp_path): model_from_checkpoint(checkpoint, context="restart") -def test_load_trainer_checkpoint_wrong_version(monkeypatch, tmp_path): +def test_load_trainer_checkpoint_wrong_version( + monkeypatch, tmp_path, MODEL_PATH_64_BIT +): monkeypatch.chdir(tmp_path) - path = RESOURCES_PATH / "model-64-bit.ckpt" + path = MODEL_PATH_64_BIT.with_suffix(".ckpt") model = torch.load(path, weights_only=False, map_location="cpu") model["trainer_ckpt_version"] = 5000000 @@ -96,16 +119,8 @@ def test_load_trainer_checkpoint_wrong_version(monkeypatch, tmp_path): trainer_from_checkpoint(checkpoint, context="restart", hypers={}) -@pytest.mark.parametrize( - "path", - [ - RESOURCES_PATH / "model-32-bit.pt", - str(RESOURCES_PATH / "model-32-bit.pt"), - f"file:{str(RESOURCES_PATH / 'model-32-bit.pt')}", - ], -) -def test_load_model_exported(path): - model = load_model(path) +def test_load_model_exported(load_path): + model = load_model(load_path) assert type(model) is AtomisticModel diff --git a/tests/utils/test_long_range.py b/tests/utils/test_long_range.py index a8c8e85b10..3b6b24fb0f 100644 --- a/tests/utils/test_long_range.py +++ b/tests/utils/test_long_range.py @@ -11,7 +11,7 @@ get_system_with_neighbor_lists, ) -from . import RESOURCES_PATH +from ..conftest import RESOURCES_PATH @pytest.mark.parametrize("periodicity", [True, False]) diff --git a/tests/utils/test_output_gradient.py b/tests/utils/test_output_gradient.py index 3799173e94..89365627e3 100644 --- a/tests/utils/test_output_gradient.py +++ b/tests/utils/test_output_gradient.py @@ -11,7 +11,7 @@ ) from metatrain.utils.output_gradient import compute_gradient -from . import MODEL_HYPERS, RESOURCES_PATH +from ..conftest import MODEL_HYPERS, RESOURCES_PATH @pytest.mark.parametrize("is_training", [True, False]) diff --git a/tests/utils/test_scaler.py b/tests/utils/test_scaler.py index 1405594857..70c3bb7888 100644 --- a/tests/utils/test_scaler.py +++ b/tests/utils/test_scaler.py @@ -1,6 +1,5 @@ import copy import math -from pathlib import Path import metatensor.torch import metatensor.torch as mts @@ -19,8 +18,7 @@ ) from metatrain.utils.scaler import Scaler, remove_scale - -RESOURCES_PATH = Path(__file__).parents[1] / "resources" +from ..conftest import RESOURCES_PATH @pytest.mark.parametrize("batch_size", [1, 2]) diff --git a/tox.ini b/tox.ini index 7b39199e67..a6b3077690 100644 --- a/tox.ini +++ b/tox.ini @@ -103,8 +103,6 @@ setenv = COVERAGE_RUN = 1 allowlist_externals = bash -commands_pre = - bash {toxinidir}/tests/resources/generate-outputs.sh commands = pytest \ --numprocesses=auto \ From cf6729a3c822ed214b46a3dcb25268b3fdf99ed6 Mon Sep 17 00:00:00 2001 From: Pol Date: Wed, 21 Jan 2026 13:02:11 +0100 Subject: [PATCH 2/7] lint --- tests/conftest.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 121b654916..df17ac16cd 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -13,6 +13,7 @@ def pytest_xdist_auto_num_workers(): n_processes = os.cpu_count() or 1 return min(12, math.ceil(n_processes * 0.8)) + MODEL_HYPERS = get_default_hypers("soap_bpnn")["model"] # ------------------------------- @@ -30,10 +31,11 @@ def pytest_xdist_auto_num_workers(): OPTIONS_PET_PATH = RESOURCES_PATH / "options-pet.yaml" OPTIONS_EXTRA_DATA_PATH = RESOURCES_PATH / "options-extra-data.yaml" + # ------------------------------- # PATHS TO TRAINED MODELS # ------------------------------- -# These files are generated from training on each test run +# These files are generated from training on each test run # and take some time to generate. Therefore, we make them a # fixture so that they are generated lazily only if there # is a test that requires them. Since we support parallel From 548a207892e7b9a16cfa9759dbc4a9c933da02a2 Mon Sep 17 00:00:00 2001 From: Pol Date: Wed, 21 Jan 2026 13:08:20 +0100 Subject: [PATCH 3/7] Add __init__.py --- tests/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/__init__.py diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 From d0e20fbd1c615e0cea37f1875dc98adfbefa4869 Mon Sep 17 00:00:00 2001 From: Pol Date: Wed, 21 Jan 2026 19:10:34 +0100 Subject: [PATCH 4/7] Fix export tests --- tests/cli/test_export_model.py | 40 ++++++++++++++++++-------------- tests/resources/run_trainings.sh | 4 ++-- 2 files changed, 25 insertions(+), 19 deletions(-) diff --git a/tests/cli/test_export_model.py b/tests/cli/test_export_model.py index ffcdaed75e..c7b22a5843 100644 --- a/tests/cli/test_export_model.py +++ b/tests/cli/test_export_model.py @@ -19,16 +19,14 @@ from metatrain.utils.architectures import find_all_architectures from metatrain.utils.io import load_model -from ..conftest import RESOURCES_PATH - @pytest.mark.parametrize("path", [Path("exported.pt"), "exported.pt"]) -def test_export(monkeypatch, tmp_path, path, caplog): +def test_export(monkeypatch, tmp_path, path, caplog, MODEL_PATH_64_BIT): """Tests the export_model function.""" monkeypatch.chdir(tmp_path) caplog.set_level(logging.INFO) - checkpoint_path = RESOURCES_PATH / "model-64-bit.ckpt" + checkpoint_path = MODEL_PATH_64_BIT.with_suffix(".ckpt") export_model(checkpoint_path, path) # Test if extensions are saved @@ -43,20 +41,28 @@ def test_export(monkeypatch, tmp_path, path, caplog): @pytest.mark.parametrize("output", [None, "exported.pt"]) @pytest.mark.parametrize("model_type", ["32-bit", "64-bit", "pet"]) -def test_export_cli(monkeypatch, tmp_path, output, model_type): +def test_export_cli(request, monkeypatch, tmp_path, output, model_type): """Test that the export cli runs without an error raise.""" monkeypatch.chdir(tmp_path) + fixture_name = { + "32-bit": "MODEL_PATH", + "64-bit": "MODEL_PATH_64_BIT", + "pet": "MODEL_PATH_PET", + }.get(model_type) + + model_path = request.getfixturevalue(fixture_name) + command = [ "mtt", "export", - str(RESOURCES_PATH / f"model-{model_type}.ckpt"), + str(model_path.with_suffix(".ckpt")), ] if output is not None: command += ["-o", output] else: - output = f"model-{model_type}.pt" + output = model_path.name subprocess.check_call(command) assert Path(output).is_file() @@ -83,21 +89,21 @@ def test_export_cli(monkeypatch, tmp_path, output, model_type): assert next(model.parameters()).device.type == "cpu" -def test_export_with_env(monkeypatch, tmp_path): +def test_export_with_env(monkeypatch, tmp_path, MODEL_PATH): """Test that export with env variable works for local file.""" monkeypatch.chdir(tmp_path) command = [ "mtt", "export", - str(RESOURCES_PATH / "model-32-bit.ckpt"), + str(MODEL_PATH.with_suffix(".ckpt")), ] env = os.environ.copy() env["HF_TOKEN"] = "1234" subprocess.check_call(command, env=env) - assert Path("model-32-bit.pt").is_file() + assert Path(MODEL_PATH.name).is_file() def test_export_cli_unknown_architecture(tmpdir): @@ -113,11 +119,11 @@ def test_export_cli_unknown_architecture(tmpdir): assert architecture_name in stdout -def test_reexport(monkeypatch, tmp_path): +def test_reexport(monkeypatch, tmp_path, MODEL_PATH_64_BIT): """Test that an already exported model can be loaded and again exported.""" monkeypatch.chdir(tmp_path) - checkpoint_path = RESOURCES_PATH / "model-64-bit.ckpt" + checkpoint_path = MODEL_PATH_64_BIT.with_suffix(".ckpt") export_model(checkpoint_path, "exported.pt") export_model("exported.pt", "exported_new.pt") @@ -200,7 +206,7 @@ def test_token_env_error(): subprocess.check_call(command, env=env) -def test_metadata(monkeypatch, tmp_path): +def test_metadata(monkeypatch, tmp_path, MODEL_PATH): """Test that the export cli does inject metadata.""" monkeypatch.chdir(tmp_path) @@ -211,17 +217,17 @@ def test_metadata(monkeypatch, tmp_path): command = [ "mtt", "export", - str(RESOURCES_PATH / "model-32-bit.ckpt"), + str(MODEL_PATH.with_suffix(".ckpt")), "--metadata=metadata.yaml", ] subprocess.check_call(command) - model = load_model("model-32-bit.pt", extensions_directory="extensions/") + model = load_model(MODEL_PATH.name, extensions_directory="extensions/") assert f"This is the {model_name} model" in str(model.metadata()) -def test_export_checkpoint_with_metadata(monkeypatch, tmp_path): +def test_export_checkpoint_with_metadata(monkeypatch, tmp_path, MODEL_PATH): """Tests that the metadata is correctly assigned to the exported model if the checkpoint has the metadata inside.""" @@ -234,7 +240,7 @@ def test_export_checkpoint_with_metadata(monkeypatch, tmp_path): command = [ "mtt", "export", - str(RESOURCES_PATH / "model-32-bit.ckpt"), + str(MODEL_PATH.with_suffix(".ckpt")), "-o=model-32-bit-with-metadata.ckpt", "--metadata=metadata.yaml", ] diff --git a/tests/resources/run_trainings.sh b/tests/resources/run_trainings.sh index e5062ba2bd..828c44e8c8 100755 --- a/tests/resources/run_trainings.sh +++ b/tests/resources/run_trainings.sh @@ -31,7 +31,7 @@ fi echo "Clearing previous generated files..." # Clean previous generated files rm $ROOT_DIR/model-$MODE-*.pt $ROOT_DIR/model-$MODE-*.ckpt || true -rm $ROOT_DIR/$MODE-*.lock || true +rm $ROOT_DIR/$MODE-*.trainlock || true echo "Generating data for testing..." @@ -43,7 +43,7 @@ if [ "$MODE" == "32-bit" ]; then elif [ "$MODE" == "64-bit" ]; then mtt train options.yaml -o model-64-bit-$TRAIN_ID.pt -r base_precision=64 elif [ "$MODE" == "pet" ]; then - mtt train options-pet.yaml -o model-pet-$TRAIN_ID.ckpt -r base_precision=pet + mtt train options-pet.yaml -o model-pet-$TRAIN_ID.pt else echo "Error: Unknown training mode (first argument): '$MODE'" echo " Please set it to '32-bit', '64-bit' or 'pet'" From f63f6d9a7fe04f4a6b7f705019949d1cdb65ab49 Mon Sep 17 00:00:00 2001 From: Pol Date: Wed, 21 Jan 2026 19:27:46 +0100 Subject: [PATCH 5/7] Fix eval and abc tests --- tests/cli/test_eval_model.py | 20 ++++++++++++-------- tests/utils/test_abc.py | 4 ++-- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/tests/cli/test_eval_model.py b/tests/cli/test_eval_model.py index 122cb0799a..bcb56dc4ff 100644 --- a/tests/cli/test_eval_model.py +++ b/tests/cli/test_eval_model.py @@ -54,15 +54,17 @@ def test_eval_cli(monkeypatch, tmp_path, MODEL_PATH): assert Path("output.xyz").is_file() -@pytest.mark.parametrize("model_name", ["model-32-bit.pt", "model-64-bit.pt"]) -def test_eval(monkeypatch, tmp_path, caplog, model_name, options): +@pytest.mark.parametrize("model_type", ["32-bit", "64-bit"]) +def test_eval(request, monkeypatch, tmp_path, caplog, model_type, options): """Test that eval via python API runs without an error raise.""" monkeypatch.chdir(tmp_path) caplog.set_level(logging.INFO) shutil.copy(RESOURCES_PATH / "qm9_reduced_100.xyz", "qm9_reduced_100.xyz") - model = torch.jit.load(RESOURCES_PATH / model_name) + model_path = request.getfixturevalue(model_type) + + model = torch.jit.load(model_path) eval_model( model=model, @@ -84,15 +86,17 @@ def test_eval(monkeypatch, tmp_path, caplog, model_name, options): frames[0].info["energy"] -@pytest.mark.parametrize("model_name", ["model-32-bit.pt", "model-64-bit.pt"]) -def test_eval_batch_size(monkeypatch, tmp_path, caplog, model_name, options): +@pytest.mark.parametrize("model_type", ["32-bit", "64-bit"]) +def test_eval_batch_size(request, monkeypatch, tmp_path, caplog, model_type, options): """Test that eval via python API runs without an error raise.""" monkeypatch.chdir(tmp_path) caplog.set_level(logging.DEBUG) shutil.copy(RESOURCES_PATH / "qm9_reduced_100.xyz", "qm9_reduced_100.xyz") - model = torch.jit.load(RESOURCES_PATH / model_name) + model_path = request.getfixturevalue(model_type) + + model = torch.jit.load(model_path) eval_model( model=model, @@ -180,14 +184,14 @@ def test_eval_no_targets(monkeypatch, tmp_path, model, options): @pytest.mark.parametrize("suffix", [".zip", ".mts"]) -def test_eval_disk_dataset(monkeypatch, tmp_path, caplog, suffix): +def test_eval_disk_dataset(monkeypatch, tmp_path, caplog, suffix, MODEL_PATH): """Test that eval via python API runs without an error raise.""" monkeypatch.chdir(tmp_path) caplog.set_level(logging.INFO) shutil.copy(RESOURCES_PATH / "qm9_reduced_100.xyz", "qm9_reduced_100.xyz") - model = torch.jit.load(RESOURCES_PATH / "model-32-bit.pt") + model = torch.jit.load(MODEL_PATH) options = OmegaConf.create( { diff --git a/tests/utils/test_abc.py b/tests/utils/test_abc.py index a96ee16995..fef2924319 100644 --- a/tests/utils/test_abc.py +++ b/tests/utils/test_abc.py @@ -42,7 +42,7 @@ def load_checkpoint( def test_trainer_interface(): message = ( "missing '__checkpoint_version__' class attribute for " - "'utils.test_abc.MyTrainer'" + "'tests.utils.test_abc.MyTrainer'" ) with pytest.raises(TypeError, match=message): _ = MyTrainer({}) @@ -105,7 +105,7 @@ def test_model_interface(): ] for attr in EXPECTED_ATTRS: - message = f"missing '{attr}' class attribute for 'utils.test_abc.MyModel'" + message = f"missing '{attr}' class attribute for 'tests.utils.test_abc.MyModel'" with pytest.raises(TypeError, match=message): _ = MyModel({}, DatasetInfo("", [], {}), ModelMetadata()) From a062dbca6ddd9a9fd77f241a176e8458a8db4f75 Mon Sep 17 00:00:00 2001 From: Pol Date: Wed, 21 Jan 2026 19:32:17 +0100 Subject: [PATCH 6/7] Fix transfer learning tests --- tests/cli/test_train_model.py | 32 ++++++++++++++++++++++---------- 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/tests/cli/test_train_model.py b/tests/cli/test_train_model.py index 5d6aff8374..6e10a04d0d 100644 --- a/tests/cli/test_train_model.py +++ b/tests/cli/test_train_model.py @@ -670,7 +670,7 @@ def test_restart(options, monkeypatch, tmp_path, MODEL_PATH_64_BIT): monkeypatch.chdir(tmp_path) shutil.copy(DATASET_PATH_QM9, "qm9_reduced_100.xyz") - train_model(options, restart_from=MODEL_PATH_64_BIT) + train_model(options, restart_from=MODEL_PATH_64_BIT.with_suffix("ckpt")) def test_finetune(options_pet, caplog, monkeypatch, tmp_path, MODEL_PATH_PET): @@ -696,10 +696,12 @@ def test_finetune(options_pet, caplog, monkeypatch, tmp_path, MODEL_PATH_PET): def test_transfer_learn(options_pet, caplog, monkeypatch, tmp_path, MODEL_PATH_PET): monkeypatch.chdir(tmp_path) + ckpt_path = MODEL_PATH_PET.with_suffix(".ckpt") + options_pet_transfer_learn = copy.deepcopy(options_pet) options_pet_transfer_learn["architecture"]["training"]["finetune"] = { "method": "heads", - "read_from": str(MODEL_PATH_PET), + "read_from": str(ckpt_path), "config": { "head_modules": ["node_heads", "edge_heads"], "last_layer_modules": ["node_last_layers", "edge_last_layers"], @@ -714,7 +716,7 @@ def test_transfer_learn(options_pet, caplog, monkeypatch, tmp_path, MODEL_PATH_P caplog.set_level(logging.INFO) train_model(options_pet_transfer_learn) - assert f"Starting finetuning from '{MODEL_PATH_PET}'" in caplog.text + assert f"Starting finetuning from '{ckpt_path}'" in caplog.text def test_transfer_learn_with_forces( @@ -722,10 +724,12 @@ def test_transfer_learn_with_forces( ): monkeypatch.chdir(tmp_path) + ckpt_path = MODEL_PATH_PET.with_suffix(".ckpt") + options_pet_transfer_learn = copy.deepcopy(options_pet) options_pet_transfer_learn["architecture"]["training"]["finetune"] = { "method": "heads", - "read_from": str(MODEL_PATH_PET), + "read_from": str(ckpt_path), "config": { "head_modules": ["node_heads", "edge_heads"], "last_layer_modules": ["node_last_layers", "edge_last_layers"], @@ -749,7 +753,7 @@ def test_transfer_learn_with_forces( caplog.set_level(logging.INFO) train_model(options_pet_transfer_learn) - assert f"Starting finetuning from '{MODEL_PATH_PET}'" in caplog.text + assert f"Starting finetuning from '{ckpt_path}'" in caplog.text def test_transfer_learn_variant( @@ -757,10 +761,12 @@ def test_transfer_learn_variant( ): monkeypatch.chdir(tmp_path) + ckpt_path = MODEL_PATH_PET.with_suffix(".ckpt") + options_pet_transfer_learn = copy.deepcopy(options_pet) options_pet_transfer_learn["architecture"]["training"]["finetune"] = { "method": "full", - "read_from": str(MODEL_PATH_PET), + "read_from": str(ckpt_path), } options_pet_transfer_learn["training_set"]["systems"]["read_from"] = ( "ethanol_reduced_100.xyz" @@ -781,7 +787,7 @@ def test_transfer_learn_variant( caplog.set_level(logging.INFO) train_model(options_pet_transfer_learn) - assert f"Starting finetuning from '{MODEL_PATH_PET}'" in caplog.text + assert f"Starting finetuning from '{ckpt_path}'" in caplog.text def test_transfer_learn_inherit_heads( @@ -789,10 +795,12 @@ def test_transfer_learn_inherit_heads( ): monkeypatch.chdir(tmp_path) + ckpt_path = MODEL_PATH_PET.with_suffix(".ckpt") + options_pet_transfer_learn = copy.deepcopy(options_pet) options_pet_transfer_learn["architecture"]["training"]["finetune"] = { "method": "full", - "read_from": str(MODEL_PATH_PET), + "read_from": str(ckpt_path), "config": {}, "inherit_heads": { "mtt::energy": "energy", @@ -816,12 +824,14 @@ def test_transfer_learn_inherit_heads_invalid_source( ): monkeypatch.chdir(tmp_path) + ckpt_path = MODEL_PATH_PET.with_suffix(".ckpt") + options_pet_transfer_learn_invalid_source = copy.deepcopy(options_pet) options_pet_transfer_learn_invalid_source["architecture"]["training"][ "finetune" ] = { "method": "full", - "read_from": str(MODEL_PATH_PET), + "read_from": str(ckpt_path), "config": {}, "inherit_heads": { "mtt::energy": "foo", @@ -846,10 +856,12 @@ def test_transfer_learn_inherit_heads_invalid_destination( ): monkeypatch.chdir(tmp_path) + ckpt_path = MODEL_PATH_PET.with_suffix(".ckpt") + options_pet_transfer_learn_invalid_dest = copy.deepcopy(options_pet) options_pet_transfer_learn_invalid_dest["architecture"]["training"]["finetune"] = { "method": "full", - "read_from": str(MODEL_PATH_PET), + "read_from": str(ckpt_path), "inherit_heads": { "mtt::foo": "energy", }, From 2d1ae9494781d86b8d5d1c181b1f8a650c1816b9 Mon Sep 17 00:00:00 2001 From: Pol Date: Wed, 21 Jan 2026 20:19:07 +0100 Subject: [PATCH 7/7] Final test fixes --- tests/cli/test_eval_model.py | 14 ++++++++++++-- tests/cli/test_train_model.py | 14 +++++++++----- 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/tests/cli/test_eval_model.py b/tests/cli/test_eval_model.py index bcb56dc4ff..5bae7aa7e7 100644 --- a/tests/cli/test_eval_model.py +++ b/tests/cli/test_eval_model.py @@ -62,7 +62,12 @@ def test_eval(request, monkeypatch, tmp_path, caplog, model_type, options): shutil.copy(RESOURCES_PATH / "qm9_reduced_100.xyz", "qm9_reduced_100.xyz") - model_path = request.getfixturevalue(model_type) + fixture_name = { + "32-bit": "MODEL_PATH", + "64-bit": "MODEL_PATH_64_BIT", + }.get(model_type) + + model_path = request.getfixturevalue(fixture_name) model = torch.jit.load(model_path) @@ -94,7 +99,12 @@ def test_eval_batch_size(request, monkeypatch, tmp_path, caplog, model_type, opt shutil.copy(RESOURCES_PATH / "qm9_reduced_100.xyz", "qm9_reduced_100.xyz") - model_path = request.getfixturevalue(model_type) + fixture_name = { + "32-bit": "MODEL_PATH", + "64-bit": "MODEL_PATH_64_BIT", + }.get(model_type) + + model_path = request.getfixturevalue(fixture_name) model = torch.jit.load(model_path) diff --git a/tests/cli/test_train_model.py b/tests/cli/test_train_model.py index 6e10a04d0d..cf5753601c 100644 --- a/tests/cli/test_train_model.py +++ b/tests/cli/test_train_model.py @@ -670,15 +670,17 @@ def test_restart(options, monkeypatch, tmp_path, MODEL_PATH_64_BIT): monkeypatch.chdir(tmp_path) shutil.copy(DATASET_PATH_QM9, "qm9_reduced_100.xyz") - train_model(options, restart_from=MODEL_PATH_64_BIT.with_suffix("ckpt")) + train_model(options, restart_from=MODEL_PATH_64_BIT.with_suffix(".ckpt")) def test_finetune(options_pet, caplog, monkeypatch, tmp_path, MODEL_PATH_PET): monkeypatch.chdir(tmp_path) + ckpt_path = MODEL_PATH_PET.with_suffix(".ckpt") + options_pet["architecture"]["training"]["finetune"] = { "method": "heads", - "read_from": str(MODEL_PATH_PET), + "read_from": str(ckpt_path), "config": { "head_modules": ["node_heads", "edge_heads"], "last_layer_modules": ["node_last_layers", "edge_last_layers"], @@ -690,7 +692,7 @@ def test_finetune(options_pet, caplog, monkeypatch, tmp_path, MODEL_PATH_PET): caplog.set_level(logging.INFO) train_model(options_pet) - assert f"Starting finetuning from '{MODEL_PATH_PET}'" in caplog.text + assert f"Starting finetuning from '{ckpt_path}'" in caplog.text def test_transfer_learn(options_pet, caplog, monkeypatch, tmp_path, MODEL_PATH_PET): @@ -888,6 +890,8 @@ def test_restart_auto( shutil.copy(DATASET_PATH_QM9, "qm9_reduced_100.xyz") caplog.set_level(logging.INFO) + ckpt_path = MODEL_PATH_64_BIT.with_suffix(".ckpt") + # Make up an output directory with some checkpoints true_checkpoint_dir = Path("outputs/2021-09-02/00-10-05") # as well as some lower-priority checkpoints @@ -905,7 +909,7 @@ def test_restart_auto( for checkpoint_dir in fake_checkpoints_dirs + [true_checkpoint_dir]: time.sleep(0.1) checkpoint_dir.mkdir(parents=True, exist_ok=True) - shutil.copy(MODEL_PATH_64_BIT, checkpoint_dir / checkpoint_name) + shutil.copy(ckpt_path, checkpoint_dir / checkpoint_name) # also check that the timestamp-based implementation works with moved folders if move_folder: @@ -939,7 +943,7 @@ def test_restart_different_dataset(options, monkeypatch, tmp_path, MODEL_PATH_64 options["training_set"]["systems"]["read_from"] = "ethanol_reduced_100.xyz" options["training_set"]["targets"]["energy"]["key"] = "energy" - train_model(options, restart_from=MODEL_PATH_64_BIT) + train_model(options, restart_from=MODEL_PATH_64_BIT.with_suffix(".ckpt")) @pytest.mark.parametrize("seed", [None, 1234])