From f1759e1e8dc6812ac3515f284ade6762072686e8 Mon Sep 17 00:00:00 2001 From: Arn Date: Mon, 6 Jul 2026 17:36:42 +0100 Subject: [PATCH] Isolate completion marker changes --- conftest.py | 2 + ml_peg/analysis/conftest.py | 111 +++ .../iron_properties/calc_iron_properties.py | 11 +- ml_peg/calcs/conftest.py | 295 +++++++- .../nebs/li_diffusion/calc_li_diffusion.py | 23 +- .../physicality/diatomics/calc_diatomics.py | 12 +- .../physicality/locality/calc_locality.py | 15 +- .../oxidation_states/calc_oxidation_states.py | 10 +- .../surfaces/CMRAds200/calc_CMRAds200.py | 2 +- ml_peg/calcs/utils/completion.py | 318 ++++++++ ml_peg/calcs/utils/utils.py | 122 +++- ml_peg/cli/cli.py | 12 +- scripts/migrate_completion_markers.py | 171 +++++ scripts/mlpeg_job_lock.py | 146 ++++ tests/test_completion.py | 677 ++++++++++++++++++ tests/test_download_utils.py | 175 +++++ tests/test_migrate_completion_markers.py | 183 +++++ 17 files changed, 2230 insertions(+), 55 deletions(-) create mode 100644 ml_peg/analysis/conftest.py create mode 100644 ml_peg/calcs/utils/completion.py create mode 100644 scripts/migrate_completion_markers.py create mode 100644 scripts/mlpeg_job_lock.py create mode 100644 tests/test_completion.py create mode 100644 tests/test_download_utils.py create mode 100644 tests/test_migrate_completion_markers.py diff --git a/conftest.py b/conftest.py index 0b19685a5..20598401e 100644 --- a/conftest.py +++ b/conftest.py @@ -10,6 +10,8 @@ from ml_peg import models +pytest_plugins = ["pytester"] + def pytest_addoption(parser): """Add flag to run tests for extra MLIPs.""" diff --git a/ml_peg/analysis/conftest.py b/ml_peg/analysis/conftest.py new file mode 100644 index 000000000..053c49caa --- /dev/null +++ b/ml_peg/analysis/conftest.py @@ -0,0 +1,111 @@ +"""Configure pytest for analysis.""" + +from __future__ import annotations + +import pytest +from pytest import CallInfo, Item, Parser + +from ml_peg.calcs.utils import completion + + +def pytest_addoption(parser: Parser) -> None: + """ + Add custom CLI inputs to pytest. + + Parameters + ---------- + parser + Pytest parser object. + """ + parser.addoption( + "--force-analysis", + action="store_true", + default=False, + help="Run analysis even if it previously completed", + ) + + +def pytest_runtest_setup(item: Item) -> None: + """ + Skip analysis that previously completed with identical inputs. + + Unlike calculations, analysis aggregates all models into shared outputs, + so completion is tracked per test rather than per model. Modules missing + any of MODELS, CALC_PATH or OUT_PATH always run and keep no marker. + + Parameters + ---------- + item + Pytest test item. + """ + module = getattr(item, "module", None) + models = getattr(module, "MODELS", None) + calc_path = getattr(module, "CALC_PATH", None) + out_path = getattr(module, "OUT_PATH", None) + if models is None or calc_path is None or out_path is None: + return + + completion.clear_data_files() + item._analysis_fingerprint = completion.analysis_fingerprint( + item.path.parent, models, calc_path + ) + item._analysis_out_path = out_path + + # Empty model name: analysis markers live directly in OUT_PATH, not per model + if not item.config.getoption("--force-analysis") and completion.is_complete( + out_path, "", item.name, item._analysis_fingerprint + ): + pytest.skip("Analysis previously completed. Use --force-analysis to re-run.") + + # The analysis will run: drop any previous marker entry so a failed or + # interrupted run cannot leave outputs masked as complete + completion.unmark_complete(out_path, "", item.name) + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_makereport(item: Item, call: CallInfo): + """ + Record whether the test call phase passed. + + Needed as pytest does not expose the test outcome to + pytest_runtest_teardown, which must only mark analysis as completed if + the test passed. + + Parameters + ---------- + item + Pytest test item. + call + Result of the test phase that just ran. + + Yields + ------ + Result + Hook wrapper result holding the test report. + """ + outcome = yield + report = outcome.get_result() + if report.when == "call" and report.passed: + item._analysis_passed = True + + +def pytest_runtest_teardown(item: Item) -> None: + """ + Mark completed analysis. + + Parameters + ---------- + item + Pytest test item. + """ + fingerprint = getattr(item, "_analysis_fingerprint", None) + if fingerprint is None or not getattr(item, "_analysis_passed", False): + return + + completion.mark_complete( + item._analysis_out_path, + "", + item.name, + fingerprint, + completion.used_data_files(), + ) diff --git a/ml_peg/calcs/bulk_crystal/iron_properties/calc_iron_properties.py b/ml_peg/calcs/bulk_crystal/iron_properties/calc_iron_properties.py index e80e1dc28..07b16ffbd 100644 --- a/ml_peg/calcs/bulk_crystal/iron_properties/calc_iron_properties.py +++ b/ml_peg/calcs/bulk_crystal/iron_properties/calc_iron_properties.py @@ -973,8 +973,8 @@ def run_iron_properties(model_name: str, model: Any) -> None: @pytest.mark.slow -@pytest.mark.parametrize("model_name", MODELS) -def test_iron_properties(model_name: str) -> None: +@pytest.mark.parametrize("mlip", MODELS.items()) +def test_iron_properties(mlip: tuple[str, Any]) -> None: """ Run iron properties benchmark for each registered model. @@ -983,7 +983,8 @@ def test_iron_properties(model_name: str) -> None: Parameters ---------- - model_name - Name of the model to evaluate. + mlip + Name of the model to evaluate and the model itself. """ - run_iron_properties(model_name, MODELS[model_name]) + model_name, model = mlip + run_iron_properties(model_name, model) diff --git a/ml_peg/calcs/conftest.py b/ml_peg/calcs/conftest.py index 7b2855061..32ff627f6 100644 --- a/ml_peg/calcs/conftest.py +++ b/ml_peg/calcs/conftest.py @@ -2,9 +2,16 @@ from __future__ import annotations -from pytest import Config, Parser +from pathlib import Path +from typing import Any + +import pytest +from pytest import CallInfo, Config, Item, Parser, TerminalReporter from ml_peg import models +from ml_peg.calcs.utils import completion + +DRY_RUN_STATUS = pytest.StashKey[list]() def pytest_addoption(parser: Parser) -> None: @@ -28,6 +35,18 @@ def pytest_addoption(parser: Parser) -> None: default=False, help="Only run mock model, ignoring other models", ) + parser.addoption( + "--force-calcs", + action="store_true", + default=False, + help="Run calculations even if they previously completed", + ) + parser.addoption( + "--dry-run", + action="store_true", + default=False, + help="Report which calculations would run, without running any", + ) def pytest_configure(config: Config) -> None: @@ -42,3 +61,277 @@ def pytest_configure(config: Config) -> None: # Set current models from CLI input models.run_mock = config.getoption("--run-mock") models.mock_only = config.getoption("--mock-only") + + +def _item_mlip(item: Item) -> tuple[str, Any] | None: + """ + Get the model a test item is parametrized with, if any. + + Parameters + ---------- + item + Pytest test item. + + Returns + ------- + tuple[str, Any] | None + The item's (model_name, model) "mlip" parameter, or None if the test + is not parametrized over models. + """ + callspec = getattr(item, "callspec", None) + mlip = callspec.params.get("mlip") if callspec is not None else None + if isinstance(mlip, tuple) and len(mlip) == 2 and isinstance(mlip[0], str): + return mlip + return None + + +def _completion_test_name(item: Item) -> str: + """ + Get the test name to key completion markers on for a model ("mlip") test. + + Uses the unparametrized test name, extended with any parameters other + than the model itself (e.g. a case index), so each parametrized case is + tracked separately per model. + + Parameters + ---------- + item + Pytest test item. + + Returns + ------- + str + Test name to key completion markers on. + """ + callspec = getattr(item, "callspec", None) + extra = { + name: value + for name, value in (callspec.params.items() if callspec is not None else ()) + if name != "mlip" + } + if not extra: + return item.originalname + suffix = "-".join(f"{name}={value}" for name, value in sorted(extra.items())) + return f"{item.originalname}[{suffix}]" + + +def _module_models(item: Item) -> dict[str, Any] | None: + """ + Get the module-level MODELS dict for a test item, if defined. + + Parameters + ---------- + item + Pytest test item. + + Returns + ------- + dict[str, Any] | None + The test module's MODELS dict, or None if not defined. + """ + module_models = getattr(getattr(item, "module", None), "MODELS", None) + return module_models if isinstance(module_models, dict) else None + + +def _model_statuses(item: Item) -> dict[str, bool] | None: + """ + Get completion status for each model a test item would run. + + Parameters + ---------- + item + Pytest test item. + + Returns + ------- + dict[str, bool] | None + Whether each model's calculation previously completed with identical + inputs, or None if the test does not run models. + """ + mlip = _item_mlip(item) + if mlip is not None: + names = [mlip[0]] + test_name = _completion_test_name(item) + else: + module_models = _module_models(item) + if module_models is None: + return None + names = list(module_models) + test_name = item.name + + calc_dir = item.path.parent + return { + name: completion.is_complete( + calc_dir / "outputs", + name, + test_name, + completion.calc_fingerprint(calc_dir, name), + ) + for name in names + } + + +def pytest_collection_modifyitems(config: Config, items: list[Item]) -> None: + """ + Report which calculations would run in dry run or collect-only mode. + + In dry run mode, all calculations are also skipped. + + Parameters + ---------- + config + Pytest configuration object. + items + Collected test items. + """ + dry_run = config.getoption("--dry-run", default=False) + if not dry_run and not config.getoption("collectonly", default=False): + return + + calcs_dir = Path(__file__).parent + status_lines = config.stash.setdefault(DRY_RUN_STATUS, []) + skip_dry_run = pytest.mark.skip(reason="dry run") + for item in items: + if not item.path.is_relative_to(calcs_dir): + continue + for name, done in (_model_statuses(item) or {}).items(): + status = "up to date" if done else "would run" + status_lines.append(f"{status}: {item.nodeid} - {name}") + if dry_run: + item.add_marker(skip_dry_run) + + +def pytest_terminal_summary( + terminalreporter: TerminalReporter, exitstatus: int, config: Config +) -> None: + """ + Print the calculation statuses gathered in dry run mode. + + Parameters + ---------- + terminalreporter + Pytest terminal reporter. + exitstatus + Exit status that will be reported to the operating system. + config + Pytest configuration object. + """ + status_lines = config.stash.get(DRY_RUN_STATUS, None) + if not status_lines: + return + + terminalreporter.section("calculations dry run") + for line in sorted(status_lines): + terminalreporter.write_line(line) + + pending = sum(line.startswith("would run") for line in status_lines) + terminalreporter.write_line( + f"{pending} calculation(s) to run, {len(status_lines) - pending} up to date" + ) + + +def pytest_runtest_setup(item: Item) -> None: + """ + Skip calculations that previously completed with identical inputs. + + Tests parametrized over models ("mlip") are skipped per model. For tests + that loop over the module's MODELS dict instead, completed models are + removed from the dict, so the loop only runs the remaining models, and + the test is skipped entirely if none remain. Pruned models are restored + in pytest_runtest_teardown. + + Parameters + ---------- + item + Pytest test item. + """ + mlip = _item_mlip(item) + module_models = _module_models(item) + if mlip is None and module_models is None: + return + + completion.clear_data_files() + if mlip is None: + item._pruned_models = {} + if item.config.getoption("--force-calcs"): + return + + statuses = _model_statuses(item) + + if mlip is not None: + if statuses[mlip[0]]: + pytest.skip( + f"'{mlip[0]}' previously completed. Use --force-calcs to re-run." + ) + return + + for name, done in statuses.items(): + if done: + print(f"[skip] {item.name}: '{name}' previously completed") + item._pruned_models[name] = module_models.pop(name) + + if item._pruned_models and not module_models: + pytest.skip("All models previously completed. Use --force-calcs to re-run.") + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_makereport(item: Item, call: CallInfo): + """ + Record whether the test call phase passed. + + Needed as pytest does not expose the test outcome to + pytest_runtest_teardown, which must only mark calculations as completed + if the test passed. + + Parameters + ---------- + item + Pytest test item. + call + Result of the test phase that just ran. + + Yields + ------ + Result + Hook wrapper result holding the test report. + """ + outcome = yield + report = outcome.get_result() + if report.when == "call" and report.passed: + item._calcs_passed = True + + +def pytest_runtest_teardown(item: Item) -> None: + """ + Mark completed calculations and restore pruned models. + + Parameters + ---------- + item + Pytest test item. + """ + mlip = _item_mlip(item) + pruned = getattr(item, "_pruned_models", None) + module_models = _module_models(item) + + if getattr(item, "_calcs_passed", False): + if mlip is not None: + test_name, names = _completion_test_name(item), [mlip[0]] + elif pruned is not None: + test_name, names = item.name, list(module_models) + else: + test_name, names = item.name, [] + + calc_dir = item.path.parent + data_files = completion.used_data_files() + for name in names: + completion.mark_complete( + calc_dir / "outputs", + name, + test_name, + completion.calc_fingerprint(calc_dir, name), + data_files, + ) + + if pruned is not None and module_models is not None: + module_models.update(pruned) diff --git a/ml_peg/calcs/nebs/li_diffusion/calc_li_diffusion.py b/ml_peg/calcs/nebs/li_diffusion/calc_li_diffusion.py index bded9ec8b..c64c90158 100644 --- a/ml_peg/calcs/nebs/li_diffusion/calc_li_diffusion.py +++ b/ml_peg/calcs/nebs/li_diffusion/calc_li_diffusion.py @@ -3,6 +3,7 @@ from __future__ import annotations from pathlib import Path +from typing import Any from warnings import warn from ase import Atoms @@ -62,8 +63,10 @@ def relaxed_structs() -> dict[str, Atoms]: @pytest.mark.slow -@pytest.mark.parametrize("model_name", MODELS) -def test_li_diffusion_b(relaxed_structs: dict[str, Atoms], model_name: str) -> None: +@pytest.mark.parametrize("mlip", MODELS.items()) +def test_li_diffusion_b( + relaxed_structs: dict[str, Atoms], mlip: tuple[str, Any] +) -> None: """ Run calculations required for lithium diffusion along path B. @@ -71,9 +74,10 @@ def test_li_diffusion_b(relaxed_structs: dict[str, Atoms], model_name: str) -> N ---------- relaxed_structs Relaxed input structures, indexed by structure name and model name. - model_name - Name of model to use. + mlip + Name of model to use and model to get calculator. """ + model_name, _ = mlip init_struct = relaxed_structs[f"{model_name}-LiFePO4_start_bc.cif"] final_struct = relaxed_structs[f"{model_name}-LiFePO4_end_b.cif"] neb = NEB( @@ -109,8 +113,10 @@ def test_li_diffusion_b(relaxed_structs: dict[str, Atoms], model_name: str) -> N @pytest.mark.slow -@pytest.mark.parametrize("model_name", MODELS) -def test_li_diffusion_c(relaxed_structs: dict[str, Atoms], model_name: str) -> None: +@pytest.mark.parametrize("mlip", MODELS.items()) +def test_li_diffusion_c( + relaxed_structs: dict[str, Atoms], mlip: tuple[str, Any] +) -> None: """ Run calculations required for lithium diffusion along path C. @@ -118,9 +124,10 @@ def test_li_diffusion_c(relaxed_structs: dict[str, Atoms], model_name: str) -> N ---------- relaxed_structs Relaxed input structures, indexed by structure name and model name. - model_name - Name of model to use. + mlip + Name of model to use and model to get calculator. """ + model_name, _ = mlip init_struct = relaxed_structs[f"{model_name}-LiFePO4_start_bc.cif"] final_struct = relaxed_structs[f"{model_name}-LiFePO4_end_c.cif"] neb = NEB( diff --git a/ml_peg/calcs/physicality/diatomics/calc_diatomics.py b/ml_peg/calcs/physicality/diatomics/calc_diatomics.py index be8f6c9e0..841bdc93a 100644 --- a/ml_peg/calcs/physicality/diatomics/calc_diatomics.py +++ b/ml_peg/calcs/physicality/diatomics/calc_diatomics.py @@ -6,6 +6,7 @@ import itertools import json from pathlib import Path +from typing import Any from warnings import warn from ase import Atoms @@ -247,14 +248,15 @@ def run_diatomics(model_name: str, model) -> None: @pytest.mark.slow -@pytest.mark.parametrize("model_name", MODELS) -def test_diatomics(model_name: str) -> None: +@pytest.mark.parametrize("mlip", MODELS.items()) +def test_diatomics(mlip: tuple[str, Any]) -> None: """ Run diatomics benchmark for each registered model. Parameters ---------- - model_name - Name of the model to evaluate. + mlip + Name of the model to evaluate and the model itself. """ - run_diatomics(model_name, MODELS[model_name]) + model_name, model = mlip + run_diatomics(model_name, model) diff --git a/ml_peg/calcs/physicality/locality/calc_locality.py b/ml_peg/calcs/physicality/locality/calc_locality.py index a89159bda..2768bc2b4 100644 --- a/ml_peg/calcs/physicality/locality/calc_locality.py +++ b/ml_peg/calcs/physicality/locality/calc_locality.py @@ -4,6 +4,7 @@ from copy import copy from pathlib import Path +from typing import Any from warnings import warn from ase import Atoms @@ -163,8 +164,8 @@ def prepared_solute() -> dict[str, Atoms]: return solutes -@pytest.mark.parametrize("model_name", MODELS) -def test_ghost_atoms(prepared_solute: dict[str, Atoms], model_name: str) -> None: +@pytest.mark.parametrize("mlip", MODELS.items()) +def test_ghost_atoms(prepared_solute: dict[str, Atoms], mlip: tuple[str, Any]) -> None: """ Run ghost atom tests. @@ -172,9 +173,10 @@ def test_ghost_atoms(prepared_solute: dict[str, Atoms], model_name: str) -> None ---------- prepared_solute Solute structure to add ghost atoms to for each model. - model_name + mlip Name of model use and model to get calculator. """ + model_name, _ = mlip solute = prepared_solute[model_name] ghost_num = 20 # number of ghost atoms @@ -195,8 +197,8 @@ def test_ghost_atoms(prepared_solute: dict[str, Atoms], model_name: str) -> None write(write_dir / "system_ghost.xyz", [solute, system_ghost]) -@pytest.mark.parametrize("model_name", MODELS) -def test_rand_h(prepared_solute: dict[str, Atoms], model_name: str) -> None: +@pytest.mark.parametrize("mlip", MODELS.items()) +def test_rand_h(prepared_solute: dict[str, Atoms], mlip: tuple[str, Any]) -> None: """ Run test adding random hydrogens to solute. @@ -204,9 +206,10 @@ def test_rand_h(prepared_solute: dict[str, Atoms], model_name: str) -> None: ---------- prepared_solute Solute structure to add hydrogens to for each model. - model_name + mlip Name of model use and model to get calculator. """ + model_name, _ = mlip solute = prepared_solute[model_name] rand_trials = 30 diff --git a/ml_peg/calcs/physicality/oxidation_states/calc_oxidation_states.py b/ml_peg/calcs/physicality/oxidation_states/calc_oxidation_states.py index 4a825ab18..fddd43ba1 100644 --- a/ml_peg/calcs/physicality/oxidation_states/calc_oxidation_states.py +++ b/ml_peg/calcs/physicality/oxidation_states/calc_oxidation_states.py @@ -74,17 +74,17 @@ def test_iron_oxidation_state_md(mlip: tuple[str, Any]) -> None: warn(f"Error during MD for {salt}: {exc}", stacklevel=2) -@pytest.mark.parametrize("model_name", MODELS) -def test_iron_oxygen_rdfs(model_name: str) -> None: +@pytest.mark.parametrize("mlip", MODELS.items()) +def test_iron_oxygen_rdfs(mlip: tuple[str, Any]) -> None: """ Calculate Fe-O RDFs from NVT MLMD for oxidation states tests. Parameters ---------- - model_name - Name of MLIP. + mlip + Name of MLIP and the model itself. """ - model_name = model_name + model_name, _ = mlip out_dir = OUT_PATH / model_name rmax = 6.0 diff --git a/ml_peg/calcs/surfaces/CMRAds200/calc_CMRAds200.py b/ml_peg/calcs/surfaces/CMRAds200/calc_CMRAds200.py index b65d84283..f7ee97fae 100644 --- a/ml_peg/calcs/surfaces/CMRAds200/calc_CMRAds200.py +++ b/ml_peg/calcs/surfaces/CMRAds200/calc_CMRAds200.py @@ -59,7 +59,7 @@ def test_cmrads200(mlip: tuple[str, Any]) -> None: data_dir = download_github_data( filename="CMR.zip", github_uri="https://github.com/benshi97/ml-peg-transition-metal-data/raw/main", - force=True, + force=False, ) cmr_structs_list = read(data_dir / "CMRads_structures.extxyz", index=":") diff --git a/ml_peg/calcs/utils/completion.py b/ml_peg/calcs/utils/completion.py new file mode 100644 index 000000000..65d5f8e7e --- /dev/null +++ b/ml_peg/calcs/utils/completion.py @@ -0,0 +1,318 @@ +"""Track completed benchmark calculations so they can be skipped on re-runs.""" + +from __future__ import annotations + +from collections.abc import Iterable +import hashlib +import json +from pathlib import Path +from typing import Any + +MARKER_FILENAME = ".completed.json" + +# Benchmark data files recorded since the current test started +_data_files: set[str] = set() + +# Content hashes of local files (e.g. model checkpoints), keyed by path, size +# and mtime so each file is only re-hashed after it changes +_file_hashes: dict[tuple[str, int, int], str] = {} + + +def record_data_file(path: Path | str) -> None: + """ + Record a benchmark data file used by the current test. + + Called by the download utilities so that completion markers can store + which cached data files a calculation depends on. + + Parameters + ---------- + path + Path to the locally cached data file. + """ + _data_files.add(str(path)) + + +def clear_data_files() -> None: + """Reset the record of benchmark data files used by the current test.""" + _data_files.clear() + + +def used_data_files() -> list[str]: + """ + Get benchmark data files recorded since the last reset. + + Returns + ------- + list[str] + Sorted paths of recorded data files. + """ + return sorted(_data_files) + + +def _model_config(model_name: str) -> dict[str, Any]: + """ + Get a model's configuration from the models file. + + Parameters + ---------- + model_name + Name of the model to look up. + + Returns + ------- + dict[str, Any] + Model configuration, or an empty dict if not defined (e.g. "mock"). + """ + from ml_peg import models + from ml_peg.models.get_models import _load_models_yaml + + config = _load_models_yaml(models.models_file).get(model_name) + return config if isinstance(config, dict) else {} + + +def _local_files(config: Any) -> list[Path]: + """ + Find existing local files referenced in a model configuration. + + Parameters + ---------- + config + Model configuration, or a nested value within it. + + Returns + ------- + list[Path] + Paths of existing local files, such as model checkpoints. + """ + if isinstance(config, dict): + return [file for value in config.values() for file in _local_files(value)] + if isinstance(config, (list, tuple)): + return [file for value in config for file in _local_files(value)] + if isinstance(config, str): + try: + path = Path(config) + if path.is_file(): + return [path] + except (OSError, ValueError): + pass + return [] + + +def _file_sha256(path: Path) -> str: + """ + Hash a file's contents, cached per process. + + Parameters + ---------- + path + Path of the file to hash. + + Returns + ------- + str + Hex digest of the file's contents. + """ + stat = path.stat() + key = (str(path), stat.st_size, stat.st_mtime_ns) + if key not in _file_hashes: + sha = hashlib.sha256() + with open(path, "rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + sha.update(chunk) + _file_hashes[key] = sha.hexdigest() + return _file_hashes[key] + + +def calc_fingerprint( + calc_dir: Path, model_name: str, config: dict[str, Any] | None = None +) -> str: + """ + Fingerprint a calculation's inputs for one model. + + Hashes the Python source files in the benchmark directory, the model's + configuration, and the contents of any local files the configuration + references, such as model checkpoints. + + Parameters + ---------- + calc_dir + Directory containing the benchmark's calculation script(s). + model_name + Name of the model the calculation runs with. + config + Model configuration. Default is the model's entry in the models file. + + Returns + ------- + str + Hex digest identifying the calculation's inputs. + """ + config = _model_config(model_name) if config is None else config + + sha = hashlib.sha256() + for source in sorted(Path(calc_dir).glob("*.py")): + sha.update(source.name.encode()) + sha.update(source.read_bytes()) + sha.update(json.dumps(config, sort_keys=True, default=str).encode()) + for path in _local_files(config): + sha.update(f"{path}:{_file_sha256(path)}".encode()) + return sha.hexdigest() + + +def analysis_fingerprint( + analysis_dir: Path, + model_names: Iterable[str], + calc_path: Path, + configs: dict[str, dict[str, Any]] | None = None, +) -> str: + """ + Fingerprint an analysis benchmark's inputs across all models. + + Hashes the Python and YAML sources in the analysis directory, each model's + configuration, and the raw bytes of each model's calculation completion + marker. Re-running a calculation rewrites its marker, so this invalidates + the analysis transitively without hashing the calculation outputs + themselves. Local files referenced by model configurations are not hashed + either, as the calculation markers already reflect them. + + Parameters + ---------- + analysis_dir + Directory containing the benchmark's analysis script(s). + model_names + Names of the models the analysis aggregates. + calc_path + Benchmark calculation outputs directory, containing per-model + completion markers. + configs + Model configurations by name. Default is each model's entry in the + models file. + + Returns + ------- + str + Hex digest identifying the analysis' inputs. + """ + sha = hashlib.sha256() + for source in sorted( + [*Path(analysis_dir).glob("*.py"), *Path(analysis_dir).glob("*.yml")] + ): + sha.update(source.name.encode()) + sha.update(source.read_bytes()) + for name in sorted(model_names): + config = _model_config(name) if configs is None else configs.get(name, {}) + sha.update(name.encode()) + sha.update(json.dumps(config, sort_keys=True, default=str).encode()) + marker = Path(calc_path) / name / MARKER_FILENAME + if marker.is_file(): + sha.update(marker.read_bytes()) + else: + sha.update(f"absent:{name}".encode()) + return sha.hexdigest() + + +def _read_marker(marker: Path) -> dict[str, Any]: + """ + Read a completion marker file. + + Parameters + ---------- + marker + Path to the marker file. + + Returns + ------- + dict[str, Any] + Marker contents, or an empty dict if missing or unreadable. + """ + try: + content = json.loads(marker.read_text(encoding="utf8")) + except (OSError, json.JSONDecodeError): + return {} + return content if isinstance(content, dict) else {} + + +def is_complete( + out_path: Path, model_name: str, test_name: str, fingerprint: str +) -> bool: + """ + Check whether a calculation previously completed with identical inputs. + + Data files are considered unchanged while their cached copies exist, + matching the criterion the download utilities use to reuse the cache. + + Parameters + ---------- + out_path + Benchmark outputs directory. + model_name + Name of the model the calculation runs with. + test_name + Name of the test running the calculation. + fingerprint + Current fingerprint of the calculation's inputs. + + Returns + ------- + bool + Whether the calculation previously completed with identical inputs. + """ + entry = _read_marker(out_path / model_name / MARKER_FILENAME).get(test_name) + if not isinstance(entry, dict) or entry.get("fingerprint") != fingerprint: + return False + return all(Path(path).exists() for path in entry.get("data_files", ())) + + +def mark_complete( + out_path: Path, + model_name: str, + test_name: str, + fingerprint: str, + data_files: list[str], +) -> None: + """ + Mark a calculation as completed. + + Parameters + ---------- + out_path + Benchmark outputs directory. + model_name + Name of the model the calculation ran with. + test_name + Name of the test that ran the calculation. + fingerprint + Fingerprint of the calculation's inputs. + data_files + Paths of cached data files the calculation used. + """ + marker = out_path / model_name / MARKER_FILENAME + marker.parent.mkdir(parents=True, exist_ok=True) + content = _read_marker(marker) + content[test_name] = {"fingerprint": fingerprint, "data_files": list(data_files)} + marker.write_text(json.dumps(content, indent=2, sort_keys=True), encoding="utf8") + + +def unmark_complete(out_path: Path, model_name: str, test_name: str) -> None: + """ + Remove a test's completion marker entry, if present. + + Called before a test re-runs, so that a failed or interrupted run cannot + leave outputs masked as complete by an earlier success. + + Parameters + ---------- + out_path + Benchmark outputs directory. + model_name + Name of the model the calculation ran with. + test_name + Name of the test to remove the marker entry for. + """ + marker = out_path / model_name / MARKER_FILENAME + content = _read_marker(marker) + if test_name not in content: + return + del content[test_name] + marker.write_text(json.dumps(content, indent=2, sort_keys=True), encoding="utf8") diff --git a/ml_peg/calcs/utils/utils.py b/ml_peg/calcs/utils/utils.py index 3f01fbb02..1ce24bda9 100644 --- a/ml_peg/calcs/utils/utils.py +++ b/ml_peg/calcs/utils/utils.py @@ -3,6 +3,7 @@ from __future__ import annotations import contextlib +import fcntl import os import pathlib from pathlib import Path @@ -10,12 +11,57 @@ import requests +from ml_peg.calcs.utils.completion import record_data_file from ml_peg.data.data import download # Local cache directory BENCHMARK_DATA_DIR = pathlib.Path.home() / ".cache" / "ml_peg" +@contextlib.contextmanager +def cache_lock(path: Path): + """ + Hold an exclusive inter-process lock while downloading/extracting a file. + + Blocks until any other process holding the lock for the same path releases + it. The lock is tied to an open file descriptor, so it is released + automatically if the holding process dies. The `.lock` file exists + whether or not anyone holds the lock, so its presence is not informative. + But deletion is not necessarily safe: fcntl locks live on the inode, so a + process locking a recreated file is not excluded by holders of the old one. + + Parameters + ---------- + path + Path to the cached file to lock. + """ + lock_path = path.parent / (path.name + ".lock") + lock_path.parent.mkdir(parents=True, exist_ok=True) + with open(lock_path, "w") as lock_file: + fcntl.lockf(lock_file, fcntl.LOCK_EX) + try: + yield + finally: + fcntl.lockf(lock_file, fcntl.LOCK_UN) + + +def _extracted_marker(path: Path) -> Path: + """ + Get the marker file recording that a cached file was fully extracted. + + Parameters + ---------- + path + Path to the cached file. + + Returns + ------- + Path + Path to the marker file. + """ + return path.parent / (path.name + ".extracted") + + def download_s3_data( key: str, filename: str | Path, @@ -45,16 +91,30 @@ def download_s3_data( Path to directory containing extracted data. """ local_path = Path(BENCHMARK_DATA_DIR) / filename + marker = _extracted_marker(local_path) - # Download file if not already cached or if force is True - if force or not local_path.exists(): - print(f"[download] Downloading {endpoint}/{bucket}/{key}") - download(key=key, filename=local_path, bucket=bucket, endpoint=endpoint) - else: - print(f"[cache] Found cached file: {local_path.name}") + if force: + marker.unlink(missing_ok=True) - # Extract contents if necessary and return path - return extract_zip(local_path) + if marker.exists(): + print(f"[cache] Found cached file: {local_path.name}") + else: + # Only one process downloads and extracts; the rest block on the lock, + # then find the marker and skip + with cache_lock(local_path): + if not marker.exists(): + if force or not local_path.exists(): + print(f"[download] Downloading {endpoint}/{bucket}/{key}") + download( + key=key, filename=local_path, bucket=bucket, endpoint=endpoint + ) + else: + print(f"[cache] Found cached file: {local_path.name}") + extract_zip(local_path) + marker.touch() + + record_data_file(local_path) + return local_path.parent def download_github_data(filename: str, github_uri: str, force: bool = False) -> Path: @@ -79,23 +139,38 @@ def download_github_data(filename: str, github_uri: str, force: bool = False) -> """ uri = f"{github_uri}/{filename}" local_path = Path(BENCHMARK_DATA_DIR) / filename + marker = _extracted_marker(local_path) - # Download file if not already cached or if force is True - if force or not local_path.exists(): - print(f"[download] Downloading {filename} from {uri}") - - response = requests.get(uri) - response.raise_for_status() - local_path.parent.mkdir(parents=True, exist_ok=True) - - with open(local_path, "wb") as f_out: - f_out.write(response.content) + if force: + marker.unlink(missing_ok=True) - else: + if marker.exists(): print(f"[cache] Found cached file: {local_path.name}") - - # Extract contents if necessary and return path - return extract_zip(local_path) + else: + # Only one process downloads and extracts; the rest block on the lock, + # then find the marker and skip + with cache_lock(local_path): + if not marker.exists(): + if force or not local_path.exists(): + print(f"[download] Downloading {filename} from {uri}") + + response = requests.get(uri) + response.raise_for_status() + local_path.parent.mkdir(parents=True, exist_ok=True) + + # Write via a temporary file so a crash mid-download cannot + # leave a plausible-looking partial file behind + part_path = local_path.parent / (local_path.name + ".part") + with open(part_path, "wb") as f_out: + f_out.write(response.content) + part_path.replace(local_path) + else: + print(f"[cache] Found cached file: {local_path.name}") + extract_zip(local_path) + marker.touch() + + record_data_file(local_path) + return local_path.parent def extract_zip(filename: Path) -> Path: @@ -112,9 +187,10 @@ def extract_zip(filename: Path) -> Path: Path Parent directory of unziped file. """ + extract_dir = filename.parent + # If it's a zip, try to extract it if filename.suffix == ".zip": - extract_dir = filename.parent try: with zipfile.ZipFile(filename, "r") as zip_ref: zip_ref.extractall(extract_dir) diff --git a/ml_peg/cli/cli.py b/ml_peg/cli/cli.py index 60f710ab4..6cba44fe7 100644 --- a/ml_peg/cli/cli.py +++ b/ml_peg/cli/cli.py @@ -267,8 +267,13 @@ def run_calcs( pytest.main(options) -@app.command(name="analyse", help="Run analysis") +@app.command( + name="analyse", + help="Run analysis", + context_settings={"allow_extra_args": True, "ignore_unknown_options": True}, +) def run_analysis( + ctx: Context, models: Annotated[ str | None, Option( @@ -303,6 +308,8 @@ def run_analysis( Parameters ---------- + ctx + Typer Context. Automatically set. models Models to run analysis for, in comma-separated list. Default is `None`, corresponding to all available models. @@ -336,6 +343,9 @@ def run_analysis( if models_file: options.extend(["--models-file", models_file]) + # Parse any custom options to pytest + options.extend(ctx.args) + pytest.main(options) diff --git a/scripts/migrate_completion_markers.py b/scripts/migrate_completion_markers.py new file mode 100644 index 000000000..d0ff4f742 --- /dev/null +++ b/scripts/migrate_completion_markers.py @@ -0,0 +1,171 @@ +""" +Migrate completion markers to content-based fingerprints. + +`calc_fingerprint` used to hash local files referenced by a model's +configuration (e.g. checkpoints) by size and modification time; it now hashes +their contents. Every marker written before the change therefore holds an +outdated fingerprint, and all calculations would rerun. + +Run this script once per machine after pulling the change and before the next +pytest run, passing the same --models-file the pytest runs use (fingerprints +include the model's configuration, so markers only match when the same model +definitions are read). Every marker entry whose fingerprint is valid under the +old scheme +is rewritten to the new scheme; entries that match neither scheme are +genuinely stale and are left untouched so they rerun. Before a marker is +modified, its original content is backed up to `.completed_old.json` alongside +it (an existing backup is never overwritten). Safe to run repeatedly. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +from ml_peg import models +from ml_peg.calcs.utils.completion import ( + MARKER_FILENAME, + _local_files, + _model_config, + calc_fingerprint, +) + +CALCS_DIR = Path(__file__).resolve().parent.parent / "ml_peg" / "calcs" +BACKUP_FILENAME = ".completed_old.json" + + +def legacy_fingerprint(calc_dir: Path, model_name: str) -> str: + """ + Fingerprint a calculation's inputs as `calc_fingerprint` did before. + + Identical to the current `calc_fingerprint` except that local files are + hashed by size and modification time instead of by content. + + Parameters + ---------- + calc_dir + Directory containing the benchmark's calculation script(s). + model_name + Name of the model the calculation runs with. + + Returns + ------- + str + Hex digest identifying the calculation's inputs. + """ + config = _model_config(model_name) + + sha = hashlib.sha256() + for source in sorted(Path(calc_dir).glob("*.py")): + sha.update(source.name.encode()) + sha.update(source.read_bytes()) + sha.update(json.dumps(config, sort_keys=True, default=str).encode()) + for path in _local_files(config): + stat = path.stat() + sha.update(f"{path}:{stat.st_size}:{stat.st_mtime_ns}".encode()) + return sha.hexdigest() + + +def migrate(calcs_dir: Path | None = None) -> dict[str, int]: + """ + Migrate all completion markers under a calculations directory. + + Parameters + ---------- + calcs_dir + Directory searched recursively for `outputs//.completed.json`. + Default is the repository's calculations directory. + + Returns + ------- + dict[str, int] + Marker files scanned and unreadable, and entries migrated, already + current, and left stale. + """ + calcs_dir = CALCS_DIR if calcs_dir is None else calcs_dir + counts = dict.fromkeys(("markers", "migrated", "current", "stale", "unreadable"), 0) + # Fingerprints per (calc_dir, model): markers hold one entry per test + fingerprints: dict[tuple[Path, str], tuple[str, str]] = {} + + for marker in sorted(calcs_dir.glob(f"**/outputs/*/{MARKER_FILENAME}")): + counts["markers"] += 1 + try: + original = marker.read_text(encoding="utf8") + content = json.loads(original) + except (OSError, json.JSONDecodeError) as err: + counts["unreadable"] += 1 + print(f"[unreadable] {marker}: {err}") + continue + if not isinstance(content, dict): + counts["unreadable"] += 1 + print(f"[unreadable] {marker}: not a JSON object") + continue + + model_name = marker.parent.name + calc_dir = marker.parent.parent.parent + key = (calc_dir, model_name) + if key not in fingerprints: + fingerprints[key] = ( + calc_fingerprint(calc_dir, model_name), + legacy_fingerprint(calc_dir, model_name), + ) + new, legacy = fingerprints[key] + + changed = False + for entry in content.values(): + stored = entry.get("fingerprint") if isinstance(entry, dict) else None + if stored == new: + counts["current"] += 1 + elif stored == legacy: + entry["fingerprint"] = new + counts["migrated"] += 1 + changed = True + else: + counts["stale"] += 1 + + if changed: + backup = marker.parent / BACKUP_FILENAME + if not backup.exists(): + backup.write_text(original, encoding="utf8") + marker.write_text( + json.dumps(content, indent=2, sort_keys=True), encoding="utf8" + ) + + return counts + + +def main(argv: list[str] | None = None) -> None: + """ + Run the migration and print a summary. + + Parameters + ---------- + argv + Command line arguments. Default is `sys.argv[1:]`. + """ + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--models-file", + default=None, + help="Model definitions the pytest runs use (pytest's --models-file). " + "Default is models.yml in the models directory.", + ) + args = parser.parse_args(argv) + if args.models_file: + models.models_file = args.models_file + + print(f"Models file: {models.models_file}") + counts = migrate() + print( + f"Scanned {counts['markers']} marker file(s): " + f"{counts['migrated']} entr(y/ies) migrated, " + f"{counts['current']} already current, " + f"{counts['stale']} stale (left to rerun), " + f"{counts['unreadable']} unreadable." + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/mlpeg_job_lock.py b/scripts/mlpeg_job_lock.py new file mode 100644 index 000000000..a59c74f7d --- /dev/null +++ b/scripts/mlpeg_job_lock.py @@ -0,0 +1,146 @@ +""" +Pytest plugin: cross-job lock files so concurrent SLURM jobs share the work. + +Multiple independent jobs can run the same pytest invocation against the same +checkout. Before a test runs, the job atomically claims a lock file under +``MLPEG_LOCK_DIR`` (``O_CREAT | O_EXCL`` on shared storage); if another job +already holds the claim, the test is skipped and the job moves on to the next +one. Together with ml-peg's completion markers this makes the jobs share the +(model, benchmark) pool without ever computing the same pair twice. + +Locks are released on test teardown whether the test passed or failed, so +only jobs killed mid-test (e.g. by the walltime) leave locks behind; locks +older than ``MLPEG_LOCK_STALE_HOURS`` (default 25, just over the 24 h +walltime) are treated as such leftovers and reclaimed. + +The plugin is a no-op unless ``MLPEG_LOCK_DIR`` is set. Load it with +``pytest -p mlpeg_job_lock`` (the scripts directory must be on PYTHONPATH). +""" + +from __future__ import annotations + +import hashlib +import os +from pathlib import Path +import re +import time + +import pytest +from pytest import Item + +# Locks claimed by this process, released in pytest_runtest_teardown +_held_locks: set[Path] = set() + + +def _lock_dir() -> Path | None: + """ + Get the shared lock directory, if configured. + + Returns + ------- + Path | None + Directory from the MLPEG_LOCK_DIR environment variable, or None. + """ + value = os.environ.get("MLPEG_LOCK_DIR") + return Path(value) if value else None + + +def _lock_path(lock_dir: Path, nodeid: str) -> Path: + """ + Map a test node ID to its lock file. + + Parameters + ---------- + lock_dir + Shared lock directory. + nodeid + Pytest test node ID, e.g. "path/to/calc_x.py::test_x[model]". + + Returns + ------- + Path + Lock file path: a readable slug plus a hash of the full node ID. + """ + slug = re.sub(r"[^A-Za-z0-9._-]+", "_", nodeid)[-80:] + digest = hashlib.sha256(nodeid.encode()).hexdigest()[:16] + return lock_dir / f"{slug}.{digest}.lock" + + +def _try_claim(lock: Path) -> bool: + """ + Atomically create a lock file claiming a test for this job. + + Parameters + ---------- + lock + Lock file path. + + Returns + ------- + bool + Whether the claim succeeded (False if the lock already exists). + """ + try: + fd = os.open(lock, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + except FileExistsError: + return False + with os.fdopen(fd, "w", encoding="utf8") as file: + file.write(f"{os.environ.get('SLURM_JOB_ID', os.getpid())}\n") + return True + + +def pytest_runtest_setup(item: Item) -> None: + """ + Claim the test's lock, or skip the test if another job holds it. + + Parameters + ---------- + item + Pytest test item. + """ + lock_dir = _lock_dir() + if lock_dir is None: + return + lock_dir.mkdir(parents=True, exist_ok=True) + lock = _lock_path(lock_dir, item.nodeid) + + if not _try_claim(lock): + try: + age = time.time() - lock.stat().st_mtime + except OSError: + # Lock released between the claim attempt and stat + age = None + stale = float(os.environ.get("MLPEG_LOCK_STALE_HOURS", "25")) * 3600 + if age is not None and age < stale: + pytest.skip(f"claimed by another job ({lock.name})") + # Stale or vanished lock: reclaim it. Unlink + create is not atomic; + # the worst case is two jobs redoing one test, which is harmless. + try: + lock.unlink() + except OSError: + pass + if not _try_claim(lock): + pytest.skip(f"claimed by another job ({lock.name})") + + _held_locks.add(lock) + + +def pytest_runtest_teardown(item: Item) -> None: + """ + Release the test's lock if this job claimed it. + + Parameters + ---------- + item + Pytest test item. + """ + lock_dir = _lock_dir() + if lock_dir is None: + return + lock = _lock_path(lock_dir, item.nodeid) + if lock in _held_locks: + _held_locks.discard(lock) + try: + lock.unlink() + except OSError: + pass diff --git a/tests/test_completion.py b/tests/test_completion.py new file mode 100644 index 000000000..a7a4f8667 --- /dev/null +++ b/tests/test_completion.py @@ -0,0 +1,677 @@ +"""Tests for skipping previously completed calculations.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +from pytest import Pytester + +import ml_peg.analysis +import ml_peg.calcs +from ml_peg.calcs.utils import completion + +CALCS_CONFTEST = Path(ml_peg.calcs.__file__).parent / "conftest.py" +ANALYSIS_CONFTEST = Path(ml_peg.analysis.__file__).parent / "conftest.py" + +CALC_FILE = ''' +"""Fake benchmark writing one line per model run.""" + +from pathlib import Path + +OUT_PATH = Path(__file__).parent / "outputs" + +MODELS = {"model-a": None, "model-b": None} + + +def test_fake(): + """Pretend to run calculations for each model.""" + for model_name, _model in MODELS.items(): + out_dir = OUT_PATH / model_name + out_dir.mkdir(parents=True, exist_ok=True) + with open(out_dir / "runs.txt", "a") as file: + file.write("run\\n") +''' + +FAILING_CALC_FILE = ( + CALC_FILE + + """ + if model_name == "model-b": + raise ValueError("Calculation failed") +""" +) + +PARAM_CALC_FILE = ''' +"""Fake parametrized benchmark writing one line per model run.""" + +from pathlib import Path +from typing import Any + +import pytest + +OUT_PATH = Path(__file__).parent / "outputs" + +MODELS = {"model-a": None, "model-b": None} + + +@pytest.mark.parametrize("mlip", MODELS.items()) +def test_fake_param(mlip: tuple[str, Any]): + """Pretend to run a calculation for one model.""" + model_name, _model = mlip + out_dir = OUT_PATH / model_name + out_dir.mkdir(parents=True, exist_ok=True) + with open(out_dir / "runs.txt", "a") as file: + file.write("run\\n") + if model_name == "model-b": + raise ValueError("Calculation failed") +''' + + +STACKED_CALC_FILE = ''' +"""Fake benchmark parametrized over models and cases, logging each model run.""" + +from pathlib import Path +from typing import Any + +import pytest + +OUT_PATH = Path(__file__).parent / "outputs" + +MODELS = {"model-a": None, "model-b": None} + + +@pytest.mark.parametrize("mlip", MODELS.items()) +@pytest.mark.parametrize("case_idx", range(2)) +def test_fake_stacked(mlip: tuple[str, Any], case_idx: int): + """Pretend to run a calculation for one model and case.""" + model_name, _model = mlip + out_dir = OUT_PATH / model_name + out_dir.mkdir(parents=True, exist_ok=True) + with open(out_dir / "runs.txt", "a") as file: + file.write("run\\n") +''' + + +YAML_CALC_FILE = ''' +"""Fake benchmark selecting models named in models_list.txt from models.yml.""" + +from pathlib import Path +from typing import Any + +import pytest + +from ml_peg import models + +HERE = Path(__file__).parent +models.models_file = HERE / "models.yml" + +OUT_PATH = HERE / "outputs" + +MODELS = {name: None for name in (HERE / "models_list.txt").read_text().split()} + + +@pytest.mark.parametrize("mlip", MODELS.items()) +def test_fake_yaml(mlip: tuple[str, Any]): + """Pretend to run a calculation for one model.""" + model_name, _model = mlip + out_dir = OUT_PATH / model_name + out_dir.mkdir(parents=True, exist_ok=True) + with open(out_dir / "runs.txt", "a") as file: + file.write("run\\n") +''' + + +ANALYSIS_FILE = ''' +"""Fake analysis benchmark logging each run.""" + +from pathlib import Path + +HERE = Path(__file__).parent +CALC_PATH = HERE / "calc_outputs" +OUT_PATH = HERE / "app_data" + +MODELS = {"model-a": None, "model-b": None} + + +def test_analysis(): + """Pretend to analyse all models.""" + OUT_PATH.mkdir(parents=True, exist_ok=True) + with open(OUT_PATH / "runs.txt", "a") as file: + file.write("run\\n") +''' + +FAILING_ANALYSIS_FILE = ( + ANALYSIS_FILE + + """ + raise ValueError("Analysis failed") +""" +) + +FLAKY_ANALYSIS_FILE = ( + ANALYSIS_FILE + + """ + if (HERE / "fail.flag").exists(): + raise ValueError("Analysis failed") +""" +) + +DATA_ANALYSIS_FILE = ANALYSIS_FILE.replace( + " OUT_PATH.mkdir(parents=True, exist_ok=True)", + """ from ml_peg.calcs.utils import completion + + data_file = HERE / "cache.txt" + data_file.write_text("data") + completion.record_data_file(data_file) + OUT_PATH.mkdir(parents=True, exist_ok=True)""", +) + +NO_CALC_PATH_ANALYSIS_FILE = ANALYSIS_FILE.replace( + 'CALC_PATH = HERE / "calc_outputs"\n', "" +) + + +def _runs(pytester: Pytester, model_name: str) -> int: + """ + Count runs recorded by the fake benchmark for a model. + + Parameters + ---------- + pytester + Pytester fixture the fake benchmark ran under. + model_name + Name of the model to count runs for. + + Returns + ------- + int + Number of recorded runs. + """ + runs_file = pytester.path / "outputs" / model_name / "runs.txt" + return runs_file.read_text().count("run") if runs_file.exists() else 0 + + +def test_fingerprint_tracks_sources_and_config(tmp_path): + """Test fingerprint changes with calc scripts and model configuration.""" + calc_file = tmp_path / "calc_fake.py" + calc_file.write_text("A = 1\n") + config = {"class_name": "mace"} + + fingerprint = completion.calc_fingerprint(tmp_path, "model", config=config) + assert fingerprint == completion.calc_fingerprint(tmp_path, "model", config=config) + + calc_file.write_text("A = 2\n") + new_fingerprint = completion.calc_fingerprint(tmp_path, "model", config=config) + assert new_fingerprint != fingerprint + + assert ( + completion.calc_fingerprint(tmp_path, "model", config={"class_name": "orb"}) + != new_fingerprint + ) + + +def test_fingerprint_tracks_local_config_files(tmp_path): + """Test fingerprint changes when local files in the configuration change.""" + (tmp_path / "calc_fake.py").write_text("A = 1\n") + checkpoint = tmp_path / "model.ckpt" + checkpoint.write_text("weights") + config = {"kwargs": {"model_path": str(checkpoint)}} + + fingerprint = completion.calc_fingerprint(tmp_path, "model", config=config) + checkpoint.write_text("retrained weights") + assert completion.calc_fingerprint(tmp_path, "model", config=config) != fingerprint + + +def test_fingerprint_ignores_checkpoint_mtime(tmp_path): + """Test fingerprint is unchanged when a checkpoint is re-copied unmodified.""" + (tmp_path / "calc_fake.py").write_text("A = 1\n") + checkpoint = tmp_path / "model.ckpt" + checkpoint.write_text("weights") + config = {"kwargs": {"model_path": str(checkpoint)}} + + fingerprint = completion.calc_fingerprint(tmp_path, "model", config=config) + + # Bump mtime without touching the content, as re-staging with cp does + os.utime(checkpoint, ns=(1, 1)) + assert completion.calc_fingerprint(tmp_path, "model", config=config) == fingerprint + + # Same size, new content: the file is re-hashed and the fingerprint changes + checkpoint.write_text("weightz") + assert completion.calc_fingerprint(tmp_path, "model", config=config) != fingerprint + + +def test_marker_roundtrip(tmp_path): + """Test completion markers record fingerprints and data files.""" + out_path = tmp_path / "outputs" + data_file = tmp_path / "data.zip" + data_file.write_text("data") + + assert not completion.is_complete(out_path, "model", "test_x", "abc") + + completion.mark_complete(out_path, "model", "test_x", "abc", [str(data_file)]) + assert completion.is_complete(out_path, "model", "test_x", "abc") + assert not completion.is_complete(out_path, "model", "test_x", "other") + assert not completion.is_complete(out_path, "model", "test_y", "abc") + + # Data is treated as changed once the cached file is gone + data_file.unlink() + assert not completion.is_complete(out_path, "model", "test_x", "abc") + + # Removing an entry invalidates it; absent entries are a no-op + completion.mark_complete(out_path, "model", "test_y", "abc", []) + completion.unmark_complete(out_path, "model", "test_y") + assert not completion.is_complete(out_path, "model", "test_y", "abc") + completion.unmark_complete(out_path, "model", "missing") + + +def test_completed_calcs_skipped(pytester: Pytester): + """Test completed calculations are skipped until their inputs change.""" + pytester.makeconftest(CALCS_CONFTEST.read_text()) + pytester.makepyfile(calc_fake=CALC_FILE) + + result = pytester.runpytest_subprocess("calc_fake.py") + result.assert_outcomes(passed=1) + assert _runs(pytester, "model-a") == 1 + assert _runs(pytester, "model-b") == 1 + + marker_file = pytester.path / "outputs" / "model-a" / completion.MARKER_FILENAME + assert "test_fake" in json.loads(marker_file.read_text()) + + # Unchanged inputs: test is skipped and calculations do not re-run + result = pytester.runpytest_subprocess("calc_fake.py") + result.assert_outcomes(skipped=1) + assert _runs(pytester, "model-a") == 1 + + # Missing outputs for one model: only that model re-runs + marker_file_b = pytester.path / "outputs" / "model-b" / completion.MARKER_FILENAME + marker_file_b.unlink() + result = pytester.runpytest_subprocess("calc_fake.py") + result.assert_outcomes(passed=1) + assert _runs(pytester, "model-a") == 1 + assert _runs(pytester, "model-b") == 2 + + # --force-calcs: everything re-runs + result = pytester.runpytest_subprocess("calc_fake.py", "--force-calcs") + result.assert_outcomes(passed=1) + assert _runs(pytester, "model-a") == 2 + + # Changed benchmark code: everything re-runs + (pytester.path / "calc_fake.py").write_text(CALC_FILE + "\n# changed\n") + result = pytester.runpytest_subprocess("calc_fake.py") + result.assert_outcomes(passed=1) + assert _runs(pytester, "model-a") == 3 + + +def test_parametrized_calcs_skipped_per_model(pytester: Pytester): + """Test model-parametrized calculations are skipped and marked per model.""" + pytester.makeconftest(CALCS_CONFTEST.read_text()) + pytester.makepyfile(calc_param=PARAM_CALC_FILE) + + # model-a passes, model-b fails + result = pytester.runpytest_subprocess("calc_param.py") + result.assert_outcomes(passed=1, failed=1) + assert _runs(pytester, "model-a") == 1 + assert _runs(pytester, "model-b") == 1 + + # Only the passed model is skipped; the failed model re-runs + result = pytester.runpytest_subprocess("calc_param.py") + result.assert_outcomes(skipped=1, failed=1) + assert _runs(pytester, "model-a") == 1 + assert _runs(pytester, "model-b") == 2 + + # Changed benchmark code: everything re-runs + fixed_calc = PARAM_CALC_FILE.replace( + ' if model_name == "model-b":\n raise ValueError' + '("Calculation failed")\n', + "", + ) + (pytester.path / "calc_param.py").write_text(fixed_calc) + result = pytester.runpytest_subprocess("calc_param.py") + result.assert_outcomes(passed=2) + assert _runs(pytester, "model-a") == 2 + assert _runs(pytester, "model-b") == 3 + + # --force-calcs: everything re-runs + result = pytester.runpytest_subprocess("calc_param.py", "--force-calcs") + result.assert_outcomes(passed=2) + assert _runs(pytester, "model-a") == 3 + + +def test_stacked_parametrized_calcs_skipped_per_case(pytester: Pytester): + """Test calculations stacking extra parameters on the model track each case.""" + pytester.makeconftest(CALCS_CONFTEST.read_text()) + pytester.makepyfile(calc_stacked=STACKED_CALC_FILE) + + # Every model and case combination runs + result = pytester.runpytest_subprocess("calc_stacked.py") + result.assert_outcomes(passed=4) + assert _runs(pytester, "model-a") == 2 + assert _runs(pytester, "model-b") == 2 + + # Completion is marked per case, not once per model + marker_file = pytester.path / "outputs" / "model-a" / completion.MARKER_FILENAME + marker = json.loads(marker_file.read_text()) + assert set(marker) == { + "test_fake_stacked[case_idx=0]", + "test_fake_stacked[case_idx=1]", + } + + # Unchanged inputs: nothing re-runs + result = pytester.runpytest_subprocess("calc_stacked.py") + result.assert_outcomes(skipped=4) + assert _runs(pytester, "model-a") == 2 + + # One case's marker missing for one model: only that case re-runs + del marker["test_fake_stacked[case_idx=1]"] + marker_file.write_text(json.dumps(marker)) + result = pytester.runpytest_subprocess("calc_stacked.py") + result.assert_outcomes(passed=1, skipped=3) + assert _runs(pytester, "model-a") == 3 + assert _runs(pytester, "model-b") == 2 + + +def test_yaml_changes_invalidate_only_changed_models(pytester: Pytester): + """Test only models with unchanged configuration are skipped across runs.""" + pytester.makeconftest(CALCS_CONFTEST.read_text()) + pytester.makepyfile(calc_yaml=YAML_CALC_FILE) + models_yml = pytester.path / "models.yml" + models_list = pytester.path / "models_list.txt" + + models_yml.write_text( + "model-a:\n version: 1\nmodel-b:\n version: 1\nmodel-c:\n version: 1\n" + ) + models_list.write_text("model-a model-b") + + result = pytester.runpytest_subprocess("calc_yaml.py") + result.assert_outcomes(passed=2) + + # Change model-b's configuration only, and select model-c as well + models_yml.write_text( + "model-a:\n version: 1\nmodel-b:\n version: 2\nmodel-c:\n version: 1\n" + ) + models_list.write_text("model-a model-b model-c") + + result = pytester.runpytest_subprocess("calc_yaml.py") + result.assert_outcomes(skipped=1, passed=2) + assert _runs(pytester, "model-a") == 1 # overlapping and unchanged: skipped + assert _runs(pytester, "model-b") == 2 # overlapping but changed: re-ran + assert _runs(pytester, "model-c") == 1 # newly selected: ran + + # Nothing changed: all models are skipped + result = pytester.runpytest_subprocess("calc_yaml.py") + result.assert_outcomes(skipped=3) + + +def test_dry_run_reports_pending_calcs(pytester: Pytester): + """Test --dry-run reports pending calculations without running any.""" + pytester.makeconftest(CALCS_CONFTEST.read_text()) + pytester.makepyfile(calc_param=PARAM_CALC_FILE) + + # Nothing has run yet: all models pending, nothing executes + result = pytester.runpytest_subprocess("calc_param.py", "--dry-run") + result.assert_outcomes(skipped=2) + assert _runs(pytester, "model-a") == 0 + result.stdout.fnmatch_lines( + [ + "*would run:*model-a*", + "*would run:*model-b*", + "*2 calculation(s) to run, 0 up to date*", + ] + ) + + # model-a passes, model-b fails + pytester.runpytest_subprocess("calc_param.py") + + # Only the failed model is still pending; dry run executes nothing + result = pytester.runpytest_subprocess("calc_param.py", "--dry-run") + result.assert_outcomes(skipped=2) + assert _runs(pytester, "model-a") == 1 + assert _runs(pytester, "model-b") == 1 + result.stdout.fnmatch_lines( + [ + "*up to date:*model-a*", + "*would run:*model-b*", + "*1 calculation(s) to run, 1 up to date*", + ] + ) + + # --collect-only also reports calculation statuses + result = pytester.runpytest_subprocess("calc_param.py", "--collect-only") + assert _runs(pytester, "model-a") == 1 + result.stdout.fnmatch_lines( + [ + "*up to date:*model-a*", + "*would run:*model-b*", + "*1 calculation(s) to run, 1 up to date*", + ] + ) + + +def test_failed_calcs_not_marked(pytester: Pytester): + """Test failed calculations are not marked as completed.""" + pytester.makeconftest(CALCS_CONFTEST.read_text()) + pytester.makepyfile(calc_fake=FAILING_CALC_FILE) + + result = pytester.runpytest_subprocess("calc_fake.py") + result.assert_outcomes(failed=1) + marker_file = pytester.path / "outputs" / "model-a" / completion.MARKER_FILENAME + assert not marker_file.exists() + + # Both models re-run after a failure + result = pytester.runpytest_subprocess("calc_fake.py") + result.assert_outcomes(failed=1) + assert _runs(pytester, "model-a") == 2 + assert _runs(pytester, "model-b") == 2 + + +def test_analysis_fingerprint_tracks_inputs(tmp_path): + """Test analysis fingerprint tracks sources, configs and calc markers.""" + analysis_dir = tmp_path / "analysis" + analysis_dir.mkdir() + calc_path = tmp_path / "calc_outputs" + (analysis_dir / "analyse_fake.py").write_text("A = 1\n") + (analysis_dir / "metrics.yml").write_text("metric: 1\n") + configs = {"model-a": {"class_name": "mace"}, "model-b": {}} + + def fingerprint(models=configs, config_map=None): + names = list(models) + return completion.analysis_fingerprint( + analysis_dir, names, calc_path, configs=config_map or models + ) + + reference = fingerprint() + assert fingerprint() == reference + + # Analysis source changed + (analysis_dir / "analyse_fake.py").write_text("A = 2\n") + changed_source = fingerprint() + assert changed_source != reference + + # metrics.yml changed + (analysis_dir / "metrics.yml").write_text("metric: 2\n") + changed_yaml = fingerprint() + assert changed_yaml != changed_source + + # A model's configuration changed + changed_config = fingerprint( + config_map={"model-a": {"class_name": "orb"}, "model-b": {}} + ) + assert changed_config != changed_yaml + + # A model added + assert fingerprint({**configs, "model-c": {}}) != changed_yaml + + # A calc marker appearing, then changing, each change the fingerprint + marker = calc_path / "model-a" / completion.MARKER_FILENAME + marker.parent.mkdir(parents=True) + marker.write_text("{}") + marker_written = fingerprint() + assert marker_written != changed_yaml + marker.write_text('{"test": {}}') + assert fingerprint() != marker_written + + # MODELS as a dict and as a list of names fingerprint identically + assert fingerprint(list(configs), config_map=configs) == fingerprint() + + +def test_analysis_fingerprint_absent_markers_distinct(tmp_path): + """Test absent calc markers for different models do not collide.""" + calc_path = tmp_path / "calc_outputs" + + fingerprint_a = completion.analysis_fingerprint( + tmp_path, ["model-a"], calc_path, configs={"model-a": {}} + ) + fingerprint_b = completion.analysis_fingerprint( + tmp_path, ["model-b"], calc_path, configs={"model-b": {}} + ) + assert fingerprint_a != fingerprint_b + + +def _analysis_runs(pytester: Pytester) -> int: + """ + Count runs recorded by the fake analysis benchmark. + + Parameters + ---------- + pytester + Pytester fixture the fake benchmark ran under. + + Returns + ------- + int + Number of recorded runs. + """ + runs_file = pytester.path / "app_data" / "runs.txt" + return runs_file.read_text().count("run") if runs_file.exists() else 0 + + +def test_completed_analysis_skipped(pytester: Pytester): + """Test completed analysis is skipped until its inputs change.""" + pytester.makeconftest(ANALYSIS_CONFTEST.read_text()) + pytester.makepyfile(analyse_fake=ANALYSIS_FILE) + calc_marker = ( + pytester.path / "calc_outputs" / "model-a" / completion.MARKER_FILENAME + ) + calc_marker.parent.mkdir(parents=True) + calc_marker.write_text("{}") + + result = pytester.runpytest_subprocess("analyse_fake.py") + result.assert_outcomes(passed=1) + assert _analysis_runs(pytester) == 1 + + marker_file = pytester.path / "app_data" / completion.MARKER_FILENAME + assert "test_analysis" in json.loads(marker_file.read_text()) + + # Unchanged inputs: analysis is skipped + result = pytester.runpytest_subprocess("analyse_fake.py") + result.assert_outcomes(skipped=1) + assert _analysis_runs(pytester) == 1 + + # Skipping preserves the marker: a still-unchanged third run skips again + result = pytester.runpytest_subprocess("analyse_fake.py") + result.assert_outcomes(skipped=1) + assert _analysis_runs(pytester) == 1 + + # Changed analysis code: re-runs + (pytester.path / "analyse_fake.py").write_text(ANALYSIS_FILE + "\n# changed\n") + result = pytester.runpytest_subprocess("analyse_fake.py") + result.assert_outcomes(passed=1) + assert _analysis_runs(pytester) == 2 + + # Changed calc marker (calculation re-ran): analysis re-runs + calc_marker.write_text('{"test_calc": {}}') + result = pytester.runpytest_subprocess("analyse_fake.py") + result.assert_outcomes(passed=1) + assert _analysis_runs(pytester) == 3 + + # New metrics.yml next to the analysis script: re-runs, then skips again + metrics_file = pytester.path / "metrics.yml" + metrics_file.write_text("metric: 1\n") + result = pytester.runpytest_subprocess("analyse_fake.py") + result.assert_outcomes(passed=1) + assert _analysis_runs(pytester) == 4 + result = pytester.runpytest_subprocess("analyse_fake.py") + result.assert_outcomes(skipped=1) + + # Changed metrics.yml: re-runs + metrics_file.write_text("metric: 2\n") + result = pytester.runpytest_subprocess("analyse_fake.py") + result.assert_outcomes(passed=1) + assert _analysis_runs(pytester) == 5 + + # --force-analysis: re-runs despite a valid marker, and rewrites it; the + # stale fingerprint proves the skip below relies on the rewritten marker + marker = json.loads(marker_file.read_text()) + marker["test_analysis"]["fingerprint"] = "stale" + marker_file.write_text(json.dumps(marker)) + result = pytester.runpytest_subprocess("analyse_fake.py", "--force-analysis") + result.assert_outcomes(passed=1) + assert _analysis_runs(pytester) == 6 + result = pytester.runpytest_subprocess("analyse_fake.py") + result.assert_outcomes(skipped=1) + + +def test_failed_analysis_not_marked(pytester: Pytester): + """Test failed analysis is not marked as completed.""" + pytester.makeconftest(ANALYSIS_CONFTEST.read_text()) + pytester.makepyfile(analyse_fake=FAILING_ANALYSIS_FILE) + + result = pytester.runpytest_subprocess("analyse_fake.py") + result.assert_outcomes(failed=1) + marker_file = pytester.path / "app_data" / completion.MARKER_FILENAME + assert not marker_file.exists() + + # A later passing run writes the marker + (pytester.path / "analyse_fake.py").write_text(ANALYSIS_FILE) + result = pytester.runpytest_subprocess("analyse_fake.py") + result.assert_outcomes(passed=1) + assert "test_analysis" in json.loads(marker_file.read_text()) + + +def test_failed_rerun_invalidates_marker(pytester: Pytester): + """Test a failed re-run with unchanged inputs does not skip afterwards.""" + pytester.makeconftest(ANALYSIS_CONFTEST.read_text()) + pytester.makepyfile(analyse_fake=FLAKY_ANALYSIS_FILE) + + result = pytester.runpytest_subprocess("analyse_fake.py") + result.assert_outcomes(passed=1) + + # Force a re-run, with fingerprinted inputs unchanged, that fails partway + flag = pytester.path / "fail.flag" + flag.write_text("") + result = pytester.runpytest_subprocess("analyse_fake.py", "--force-analysis") + result.assert_outcomes(failed=1) + + # The earlier success must not mask the failure: the analysis re-runs + flag.unlink() + result = pytester.runpytest_subprocess("analyse_fake.py") + result.assert_outcomes(passed=1) + assert _analysis_runs(pytester) == 3 + + +def test_analysis_data_files_tracked(pytester: Pytester): + """Test analysis re-runs once a recorded data file is gone from the cache.""" + pytester.makeconftest(ANALYSIS_CONFTEST.read_text()) + pytester.makepyfile(analyse_fake=DATA_ANALYSIS_FILE) + + result = pytester.runpytest_subprocess("analyse_fake.py") + result.assert_outcomes(passed=1) + result = pytester.runpytest_subprocess("analyse_fake.py") + result.assert_outcomes(skipped=1) + + (pytester.path / "cache.txt").unlink() + result = pytester.runpytest_subprocess("analyse_fake.py") + result.assert_outcomes(passed=1) + assert _analysis_runs(pytester) == 2 + + +def test_analysis_without_convention_always_runs(pytester: Pytester): + """Test analysis modules missing CALC_PATH always run and keep no marker.""" + pytester.makeconftest(ANALYSIS_CONFTEST.read_text()) + pytester.makepyfile(analyse_fake=NO_CALC_PATH_ANALYSIS_FILE) + + for _ in range(2): + result = pytester.runpytest_subprocess("analyse_fake.py") + result.assert_outcomes(passed=1) + assert _analysis_runs(pytester) == 2 + assert not (pytester.path / "app_data" / completion.MARKER_FILENAME).exists() diff --git a/tests/test_download_utils.py b/tests/test_download_utils.py new file mode 100644 index 000000000..65342f2c4 --- /dev/null +++ b/tests/test_download_utils.py @@ -0,0 +1,175 @@ +"""Tests for the cached download and extraction utilities.""" + +from __future__ import annotations + +import multiprocessing +from pathlib import Path +import zipfile + +import pytest + +import ml_peg.calcs.utils.utils as utils +from ml_peg.calcs.utils.utils import ( + cache_lock, + download_github_data, + extract_zip, +) + +MEMBER_NAME = "data/payload.txt" +MEMBER_CONTENT = b"0123456789abcdef" * 65536 # 1 MiB + + +def make_zip(path: Path) -> None: + """ + Create a zip file containing a single known member. + + Parameters + ---------- + path + Path of the zip file to create. + """ + path.parent.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(path, "w") as zip_ref: + zip_ref.writestr(MEMBER_NAME, MEMBER_CONTENT) + + +@pytest.fixture +def cache_dir(tmp_path, monkeypatch): + """Point the benchmark data cache at a temporary directory.""" + monkeypatch.setattr(utils, "BENCHMARK_DATA_DIR", tmp_path) + return tmp_path + + +def test_extract_zip_non_zip_returns_parent(tmp_path): + """Non-zip files are returned untouched without raising.""" + path = tmp_path / "data.txt" + path.write_text("not a zip") + assert extract_zip(path) == tmp_path + + +def test_cached_zip_extracted_once(cache_dir): + """A cached zip is extracted on first use only, then never rewritten.""" + make_zip(cache_dir / "data.zip") + + data_dir = download_github_data("data.zip", "https://unused.invalid") + extracted = data_dir / MEMBER_NAME + assert extracted.read_bytes() == MEMBER_CONTENT + assert (cache_dir / "data.zip.extracted").exists() + + # A second call must not re-extract over the cache: corrupt the extracted + # file and check it is left alone + extracted.write_bytes(b"corrupted") + download_github_data("data.zip", "https://unused.invalid") + assert extracted.read_bytes() == b"corrupted" + + +class FakeResponse: + """Minimal stand-in for requests.Response serving fixed bytes.""" + + def __init__(self, content: bytes): + self.content = content + + def raise_for_status(self) -> None: + """Match the requests.Response interface.""" + + +def test_force_re_downloads_and_re_extracts(cache_dir, tmp_path, monkeypatch): + """force=True re-downloads, removes the marker and repairs the data.""" + make_zip(cache_dir / "data.zip") + + data_dir = download_github_data("data.zip", "https://unused.invalid") + extracted = data_dir / MEMBER_NAME + extracted.write_bytes(b"corrupted") + + zip_bytes = (cache_dir / "data.zip").read_bytes() + monkeypatch.setattr(utils.requests, "get", lambda uri: FakeResponse(zip_bytes)) + download_github_data("data.zip", "https://unused.invalid", force=True) + + assert extracted.read_bytes() == MEMBER_CONTENT + + +def _hammer(cache_dir_str: str, barrier, errors) -> None: + """ + Repeatedly resolve the cached zip while asserting its contents are intact. + + Parameters + ---------- + cache_dir_str + Benchmark cache directory to use. + barrier + Barrier synchronising all workers to maximise contention. + errors + Queue collecting assertion failures from workers. + """ + try: + utils.BENCHMARK_DATA_DIR = Path(cache_dir_str) + barrier.wait(timeout=30) + for _ in range(50): + data_dir = download_github_data("data.zip", "https://unused.invalid") + content = (data_dir / MEMBER_NAME).read_bytes() + if content != MEMBER_CONTENT: + errors.put(f"read {len(content)} bytes, expected full member") + return + except Exception as err: + errors.put(repr(err)) + + +def test_concurrent_workers_see_complete_files(cache_dir): + """Concurrent workers never observe empty or partially extracted files.""" + make_zip(cache_dir / "data.zip") + + ctx = multiprocessing.get_context("fork") + barrier = ctx.Barrier(4) + errors = ctx.Queue() + workers = [ + ctx.Process(target=_hammer, args=(str(cache_dir), barrier, errors)) + for _ in range(4) + ] + for worker in workers: + worker.start() + for worker in workers: + worker.join(timeout=120) + + assert errors.empty(), errors.get() + + +def _try_lock(lock_path_str: str, result) -> None: + """ + Attempt a non-blocking lock from a child process. + + Parameters + ---------- + lock_path_str + Path of the lock file to try. + result + Queue receiving True if the lock was acquired, False if it was busy. + """ + import fcntl + + with open(lock_path_str, "w") as lock_file: + try: + fcntl.lockf(lock_file, fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError: + result.put(False) + else: + result.put(True) + + +def test_cache_lock_is_exclusive(cache_dir): + """Another process cannot take the lock while it is held.""" + path = cache_dir / "data.zip" + lock_path = cache_dir / "data.zip.lock" + ctx = multiprocessing.get_context("fork") + result = ctx.Queue() + + with cache_lock(path): + proc = ctx.Process(target=_try_lock, args=(str(lock_path), result)) + proc.start() + proc.join(timeout=30) + assert result.get(timeout=10) is False + + # And it is acquirable again once released + proc = ctx.Process(target=_try_lock, args=(str(lock_path), result)) + proc.start() + proc.join(timeout=30) + assert result.get(timeout=10) is True diff --git a/tests/test_migrate_completion_markers.py b/tests/test_migrate_completion_markers.py new file mode 100644 index 000000000..5b4cc9cb4 --- /dev/null +++ b/tests/test_migrate_completion_markers.py @@ -0,0 +1,183 @@ +"""Tests for the completion marker migration script.""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from ml_peg import models +from ml_peg.calcs.utils import completion + +SCRIPT_PATH = Path(__file__).parent.parent / "scripts" / "migrate_completion_markers.py" +_spec = importlib.util.spec_from_file_location( + "migrate_completion_markers", SCRIPT_PATH +) +migrate_markers = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(migrate_markers) + + +@pytest.fixture +def scene(tmp_path, monkeypatch): + """Fake calc directory, checkpoint, models file and cached data file.""" + calc_dir = tmp_path / "calc_fake" + calc_dir.mkdir() + (calc_dir / "calc_fake.py").write_text("A = 1\n") + + checkpoint = tmp_path / "model.ckpt" + checkpoint.write_text("weights") + + data_file = tmp_path / "data.zip" + data_file.write_text("data") + + models_yml = tmp_path / "models.yml" + models_yml.write_text( + f"model-x:\n kwargs:\n model_path: {checkpoint}\n" + "model-plain:\n class_name: mace\n" + ) + monkeypatch.setattr(models, "models_file", models_yml) + + return SimpleNamespace( + calcs_dir=tmp_path, + calc_dir=calc_dir, + out_path=calc_dir / "outputs", + checkpoint=checkpoint, + data_file=data_file, + models_yml=models_yml, + ) + + +def write_marker( + calc_dir: Path, model_name: str, fingerprint: str, data_files: list[str] +) -> Path: + """ + Write a completion marker with a single test entry. + + Parameters + ---------- + calc_dir + Fake benchmark directory to write the marker under. + model_name + Name of the model the marker belongs to. + fingerprint + Fingerprint to store in the entry. + data_files + Data file paths to store in the entry. + + Returns + ------- + Path + Path of the written marker file. + """ + marker = calc_dir / "outputs" / model_name / completion.MARKER_FILENAME + marker.parent.mkdir(parents=True, exist_ok=True) + content = {"test_fake": {"fingerprint": fingerprint, "data_files": data_files}} + marker.write_text(json.dumps(content, indent=2, sort_keys=True), encoding="utf8") + return marker + + +def test_old_marker_migrated_with_backup(scene): + """Test old-style markers are migrated and backed up exactly.""" + legacy = migrate_markers.legacy_fingerprint(scene.calc_dir, "model-x") + marker = write_marker(scene.calc_dir, "model-x", legacy, [str(scene.data_file)]) + original = marker.read_text() + + new = completion.calc_fingerprint(scene.calc_dir, "model-x") + assert not completion.is_complete(scene.out_path, "model-x", "test_fake", new) + + counts = migrate_markers.migrate(scene.calcs_dir) + assert counts["migrated"] == 1 + assert counts["stale"] == 0 + + assert completion.is_complete(scene.out_path, "model-x", "test_fake", new) + content = json.loads(marker.read_text()) + assert content["test_fake"]["data_files"] == [str(scene.data_file)] + backup = marker.parent / migrate_markers.BACKUP_FILENAME + assert backup.read_text() == original + + +def test_second_run_is_noop_and_keeps_backup(scene): + """Test a second run changes nothing and never clobbers the backup.""" + legacy = migrate_markers.legacy_fingerprint(scene.calc_dir, "model-x") + marker = write_marker(scene.calc_dir, "model-x", legacy, []) + backup = marker.parent / migrate_markers.BACKUP_FILENAME + backup.write_text("pre-existing backup") + + migrate_markers.migrate(scene.calcs_dir) + assert backup.read_text() == "pre-existing backup" + + migrated = marker.read_text() + counts = migrate_markers.migrate(scene.calcs_dir) + assert counts["migrated"] == 0 + assert counts["current"] == 1 + assert marker.read_text() == migrated + assert backup.read_text() == "pre-existing backup" + + +def test_stale_entry_left_untouched(scene): + """Test entries matching neither scheme are left stale to rerun.""" + marker = write_marker(scene.calc_dir, "model-x", "bogus", []) + original = marker.read_text() + + counts = migrate_markers.migrate(scene.calcs_dir) + assert counts["stale"] == 1 + assert counts["migrated"] == 0 + assert marker.read_text() == original + assert not (marker.parent / migrate_markers.BACKUP_FILENAME).exists() + + new = completion.calc_fingerprint(scene.calc_dir, "model-x") + assert not completion.is_complete(scene.out_path, "model-x", "test_fake", new) + + +def test_model_without_local_files_is_noop(scene): + """Test markers of models with no local files are already current.""" + fingerprint = completion.calc_fingerprint(scene.calc_dir, "model-plain") + assert fingerprint == migrate_markers.legacy_fingerprint( + scene.calc_dir, "model-plain" + ) + marker = write_marker(scene.calc_dir, "model-plain", fingerprint, []) + original = marker.read_text() + + counts = migrate_markers.migrate(scene.calcs_dir) + assert counts["current"] == 1 + assert counts["migrated"] == 0 + assert marker.read_text() == original + assert not (marker.parent / migrate_markers.BACKUP_FILENAME).exists() + + +def test_main_models_file_option(scene, monkeypatch, capsys): + """Test --models-file selects the model definitions like pytest's option.""" + legacy = migrate_markers.legacy_fingerprint(scene.calc_dir, "model-x") + marker = write_marker(scene.calc_dir, "model-x", legacy, []) + + # Without the right models file, model-x's config is empty and the marker + # cannot migrate + monkeypatch.setattr(models, "models_file", scene.calcs_dir / "missing.yml") + monkeypatch.setattr(migrate_markers, "CALCS_DIR", scene.calcs_dir) + + migrate_markers.main(["--models-file", str(scene.models_yml)]) + output = capsys.readouterr().out + assert str(scene.models_yml) in output + assert "1 entr(y/ies) migrated" in output + assert json.loads(marker.read_text())["test_fake"][ + "fingerprint" + ] == completion.calc_fingerprint(scene.calc_dir, "model-x") + + +def test_corrupt_marker_reported_not_fatal(scene, capsys): + """Test corrupt markers are reported while other markers still migrate.""" + broken = scene.calc_dir / "outputs" / "model-broken" / completion.MARKER_FILENAME + broken.parent.mkdir(parents=True) + broken.write_text("not json") + + legacy = migrate_markers.legacy_fingerprint(scene.calc_dir, "model-x") + write_marker(scene.calc_dir, "model-x", legacy, []) + + counts = migrate_markers.migrate(scene.calcs_dir) + assert counts["unreadable"] == 1 + assert counts["migrated"] == 1 + assert broken.read_text() == "not json" + assert "model-broken" in capsys.readouterr().out