diff --git a/CHANGELOG.md b/CHANGELOG.md index 68961dbc..b36ff139 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,16 @@ # Changelog +## Unreleased + +- Add an `inference_speed` benchmark that measures, per system size, both **model + throughput** (the raw forward pass — pure network for mlip models, forced ASE + recompute for external models) and **MD throughput** (end-to-end ns/day). It reuses + the `scaling` dataset and produces a (hardware-relative) speed score from the per-atom + model forward time, contributing to the overall model score. The GUI switches between + metrics on log–log axes with power-law fits and variance, and a summary table + including the model's graph cutoff. +- Benchmarks can now reuse another benchmark's dataset via the `dataset_name` attribute. + ## Release 0.1.4 - Add an Apache-2.0 `LICENSE` file and declare the license in `pyproject.toml`. diff --git a/docs/source/api_reference/general/inference_speed.rst b/docs/source/api_reference/general/inference_speed.rst new file mode 100644 index 00000000..3f6aac4a --- /dev/null +++ b/docs/source/api_reference/general/inference_speed.rst @@ -0,0 +1,20 @@ +.. _inference_speed_api: + +Inference Speed +=============== + +.. module:: mlipaudit.benchmarks.inference_speed.inference_speed + +.. autoclass:: InferenceSpeedBenchmark + + .. automethod:: __init__ + + .. automethod:: run_model + + .. automethod:: analyze + +.. autoclass:: InferenceSpeedResult + +.. autoclass:: InferenceSpeedStructureResult + +.. autoclass:: InferenceSpeedModelOutput diff --git a/docs/source/api_reference/index.rst b/docs/source/api_reference/index.rst index 34655795..ab9add98 100644 --- a/docs/source/api_reference/index.rst +++ b/docs/source/api_reference/index.rst @@ -40,3 +40,4 @@ Benchmark implementations biomolecules/sampling general/stability general/scaling + general/inference_speed diff --git a/docs/source/benchmarks/general/index.rst b/docs/source/benchmarks/general/index.rst index c7cbb1ac..6d37bbc1 100644 --- a/docs/source/benchmarks/general/index.rst +++ b/docs/source/benchmarks/general/index.rst @@ -11,3 +11,4 @@ applicable across molecular systems. Stability Scaling + Inference Speed diff --git a/docs/source/benchmarks/general/inference_speed.rst b/docs/source/benchmarks/general/inference_speed.rst new file mode 100644 index 00000000..d5c2939f --- /dev/null +++ b/docs/source/benchmarks/general/inference_speed.rst @@ -0,0 +1,46 @@ +.. _inference_speed: + +Inference Speed +=============== + +Purpose +------- + +This benchmark measures how fast a machine-learned interatomic potential (**MLIP**) runs +and how that speed scales with system size. It reports two complementary metrics so +models can be compared on performance as well as accuracy. + +Description +----------- + +The benchmark reuses the ``scaling`` dataset (a size-stratified set of protein chains). +For each system it measures, with warm-up and outlier trimming: + +* **Model throughput** — the raw model forward pass (energy + forces), independent of + the simulation engine. For mlip models this is the pure network forward on a + pre-built graph (mirroring mlip-jax's ``scripts/time_inference.py``); for external + ASE calculators it is a forced recomputation on the pre-built atoms (which also + includes the calculator's neighbour-list construction). Reported as **atoms/s**. +* **MD throughput** — an end-to-end short **NVT** **MD** simulation at **300 K**, + timed per episode (discarding the first to ignore compilation). Reported as + **ns/day**. + +The gap between the two reflects simulation overhead (neighbour lists, the integrator +and engine). The GUI lets you switch metrics, plots them on log–log axes with +power-law fits and per-system variance, and shows a per-model summary including the +model's **graph cutoff** (which drives neighbour count and therefore speed). + +Dataset +------- + +This benchmark reuses the ``scaling`` dataset — see :ref:`scaling` for details of the +size-stratified protein chains. + +Interpretation +-------------- + +The benchmark produces a score in ``[0, 1]`` based on the per-atom **model forward +time** ``t`` via a Hill function ``1 / (1 + (t / t₀)ᵏ)`` averaged over systems, so +faster models score higher. The forward time (rather than the MD step time) is scored +because it is engine-independent. Because ``t`` is wall-clock time, this score is +hardware-dependent and is only comparable across models run on the same GPU. diff --git a/src/mlipaudit/benchmark.py b/src/mlipaudit/benchmark.py index ed47c41d..36261925 100644 --- a/src/mlipaudit/benchmark.py +++ b/src/mlipaudit/benchmark.py @@ -104,6 +104,11 @@ class Benchmark(ABC): reusable_output_id: tuple[str, ...] | None = None + #: Name of the dataset (subdirectory and HuggingFace archive) this benchmark reads + #: its inputs from. Defaults to ``name``; set it to another benchmark's name to + #: reuse that benchmark's dataset instead of shipping a duplicate. + dataset_name: str | None = None + def __init__( self, force_field: ForceField | ASECalculator, @@ -224,17 +229,24 @@ def check_can_run_model(cls, force_field: ForceField) -> bool: return True + @property + def _dataset_name(self) -> str: + """The dataset directory/archive name, defaulting to the benchmark name.""" + return self.dataset_name or self.name + def _download_data(self) -> None: """Download the data from the data input directory if not already exists.""" - already_exists = (self.data_input_dir / self.name).exists() + already_exists = (self.data_input_dir / self._dataset_name).exists() if not already_exists: hf_hub_download( repo_id="InstaDeepAI/MLIPAudit-data", - filename=f"{self.name}.zip", + filename=f"{self._dataset_name}.zip", local_dir=self.data_input_dir, repo_type="dataset", ) - with zipfile.ZipFile(self.data_input_dir / f"{self.name}.zip", "r") as z: + with zipfile.ZipFile( + self.data_input_dir / f"{self._dataset_name}.zip", "r" + ) as z: z.extractall(self.data_input_dir) @abstractmethod diff --git a/src/mlipaudit/benchmarks/__init__.py b/src/mlipaudit/benchmarks/__init__.py index 5e899c38..3845b5fc 100644 --- a/src/mlipaudit/benchmarks/__init__.py +++ b/src/mlipaudit/benchmarks/__init__.py @@ -34,6 +34,12 @@ FoldingStabilityModelOutput, FoldingStabilityResult, ) +from mlipaudit.benchmarks.inference_speed.inference_speed import ( + InferenceSpeedBenchmark, + InferenceSpeedModelOutput, + InferenceSpeedResult, + InferenceSpeedStructureResult, +) from mlipaudit.benchmarks.noncovalent_interactions.noncovalent_interactions import ( NoncovalentInteractionsBenchmark, NoncovalentInteractionsModelOutput, diff --git a/src/mlipaudit/benchmarks/inference_speed/inference_speed.py b/src/mlipaudit/benchmarks/inference_speed/inference_speed.py new file mode 100644 index 00000000..030ef1cd --- /dev/null +++ b/src/mlipaudit/benchmarks/inference_speed/inference_speed.py @@ -0,0 +1,466 @@ +# Copyright 2025 InstaDeep Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Inference-speed benchmark. + +Measures how fast a model runs molecular dynamics and how that speed scales with +system size, and turns it into a single throughput score. It reuses the ``scaling`` +benchmark's size-stratified protein dataset and its timing helpers, but additionally +records the per-episode timing spread and the MD timestep, and produces a +hardware-relative speed score. +""" + +import functools +import logging +import os +import time +from pathlib import Path +from typing import Any + +from ase.io import read as ase_read +from mlip.simulation import SimulationState +from pydantic import BaseModel, ConfigDict, NonNegativeFloat, PositiveInt + +from mlipaudit.benchmark import ( + DEFAULT_CHARGE, + DEFAULT_SPIN, + Benchmark, + BenchmarkResult, + ModelOutput, +) +from mlipaudit.benchmarks.scaling.scaling import ( + NUM_DEV_SYSTEMS, + SIMULATION_CONFIG, + SIMULATION_CONFIG_DEV, + Timer, + get_molecule_size_from_name, +) +from mlipaudit.run_mode import RunMode +from mlipaudit.scoring import compute_speed_score +from mlipaudit.utils.simulation import get_simulation_engine + +#: Number of forward passes to discard (JIT compilation / lazy init) and to time when +#: measuring model throughput. After timing, the slowest `FORWARD_TRIM_FRACTION` of +#: passes are dropped (garbage-collection / scheduling spikes) before averaging. +NUM_FORWARD_WARMUP = 2 +NUM_FORWARD_TIMED = 25 +NUM_FORWARD_TIMED_DEV = 3 +FORWARD_TRIM_FRACTION = 0.2 + +#: Score parameters for the Hill function ``score = 1 / (1 + (t / t0) ** k)``, where +#: ``t`` is the per-atom model forward-pass time in seconds (the scored, engine- +#: independent metric). ``SCORE_PER_ATOM_FORWARD_TIME_MIDPOINT`` (``t0``) is the +#: per-atom forward time that scores 0.5 and ``SCORE_SHARPNESS`` (``k``) controls how +#: sharply models separate. These are calibrated for our H100 reference hardware; the +#: score is only comparable across models run on the same hardware. +#: TODO: tune ``t0`` against real model runs once available. +SCORE_PER_ATOM_FORWARD_TIME_MIDPOINT = 1.0e-6 +SCORE_SHARPNESS = 1.0 + +logger = logging.getLogger("mlipaudit") + + +class InferenceSpeedModelOutput(ModelOutput): + """Model output for the inference-speed benchmark. + + Attributes: + structure_names: The names of the structures used. + simulation_states: A list of final simulation states for each corresponding + structure. None if the simulation failed. + average_episode_times: A list of average episode times for each corresponding + structure, excluding the first episode to ignore the compilation time. + None if the simulation failed. + episode_times: A list, per structure, of the individual episode durations + (excluding the first episode) used to quantify timing variance. Empty for + structures whose simulation failed. + forward_times: A list, per structure, of the individual timed model + forward-pass durations (excluding warm-up). Empty for structures whose + forward pass failed. + """ + + structure_names: list[str] + simulation_states: list[SimulationState | None] + average_episode_times: list[float | None] + episode_times: list[list[float]] = [] + forward_times: list[list[float]] = [] + + model_config = ConfigDict(arbitrary_types_allowed=True) + + +class InferenceSpeedStructureResult(BaseModel): + """Result object for a single structure. + + Attributes: + structure_name: The structure name. + num_atoms: The number of atoms in the structure. + num_steps: The number of steps in the simulation. + num_episodes: The number of episodes in the simulation. + average_episode_time: The average episode time of the simulation, excluding + the first episode to ignore the compilation time. + average_step_time: The average step time of the simulation, excluding the + first episode to ignore the compilation time. None if the MD run failed. + timestep_fs: The MD timestep in femtoseconds, used to convert step times into + a throughput (ns/day). + episode_times: The individual episode durations (excluding the first episode), + used to quantify timing variance. Empty if unavailable. + average_forward_time: The average wall-clock time of a single model forward + pass (energy + forces), excluding warm-up. This is the engine-independent + model-throughput metric. None if the forward pass failed. + forward_times: The individual timed forward-pass durations, used to quantify + variance. Empty if unavailable. + failed: Whether both the MD run and the forward pass failed for this structure. + """ + + structure_name: str + num_atoms: PositiveInt + num_steps: PositiveInt + num_episodes: PositiveInt + average_episode_time: NonNegativeFloat | None = None + average_step_time: NonNegativeFloat | None = None + timestep_fs: float | None = None + episode_times: list[float] = [] + average_forward_time: NonNegativeFloat | None = None + forward_times: list[float] = [] + + failed: bool = False + + +class InferenceSpeedResult(BenchmarkResult): + """Result object for the inference-speed benchmark. + + Attributes: + structure_names: The names of the structures. + structures: List of per structure results. + graph_cutoff_angstrom: The model's graph (neighbour-list) cutoff in Angstrom, + which influences neighbour count and therefore speed. None if unavailable + (e.g. some external calculators do not expose it). + """ + + structure_names: list[str] + structures: list[InferenceSpeedStructureResult] + graph_cutoff_angstrom: float | None = None + + +class InferenceSpeedBenchmark(Benchmark): + """Benchmark measuring model and MD throughput and how they scale with size. + + For each structure (reusing the ``scaling`` dataset) it measures two complementary + speeds: the **model forward-pass** time (energy + forces, engine-independent) and + the **MD step** time (end-to-end, including neighbour lists, the integrator and the + simulation engine). The gap between them reflects simulation overhead. The model is + scored with a Hill function on its per-atom forward-pass time so that faster models + score higher; the score is wall-clock based and only comparable across models run on + the same hardware. + + Attributes: + name: The unique benchmark name (``inference_speed``). + category: The benchmark category, used for grouping in the UI. + dataset_name: Set to ``scaling`` so this benchmark reuses the ``scaling`` + dataset rather than shipping a duplicate. + result_class: The `InferenceSpeedResult` type returned by `analyze`. + model_output_class: The `InferenceSpeedModelOutput` type. + required_elements: The element types present in the input files. + """ + + name = "inference_speed" + category = "General" + dataset_name = "scaling" + result_class = InferenceSpeedResult + model_output_class = InferenceSpeedModelOutput + + required_elements = {"N", "H", "O", "S", "C"} + + def run_model(self) -> None: + """For each structure, time the model forward pass (model throughput) and run a + short MD simulation (MD throughput). The two measurements fail independently. + """ + simulation_states: list[SimulationState | None] = [] + average_episode_times: list[float | None] = [] + episode_times: list[list[float]] = [] + forward_times: list[list[float]] = [] + for structure_name in self._structure_names: + try: + atoms = ase_read( + self.data_input_dir / self._dataset_name / f"{structure_name}.xyz" + ) + atoms.info["charge"] = DEFAULT_CHARGE + atoms.info["spin"] = DEFAULT_SPIN + except Exception as e: + logger.info("Error reading system %s: %s", structure_name, str(e)) + simulation_states.append(None) + average_episode_times.append(None) + episode_times.append([]) + forward_times.append([]) + continue + + # Model throughput: timed forward passes on a copy to avoid side effects. + forward_times.append(self._measure_model_throughput(atoms.copy())) + + # MD throughput: short MD simulation timed per episode. + try: + timer = Timer() + md_engine = get_simulation_engine( + atoms=atoms, + force_field=self.force_field, + **self._md_kwargs, + ) + md_engine.attach_logger(timer.log) + md_engine.run() + + simulation_states.append(md_engine.state) + average_episode_times.append(timer.average_episode_time) + episode_times.append(timer.episode_times) + + except Exception as e: + logger.info( + "Error running simulation on system %s: %s", str(atoms), str(e) + ) + simulation_states.append(None) + average_episode_times.append(None) + episode_times.append([]) + + self.model_output = InferenceSpeedModelOutput( + structure_names=self._structure_names, + simulation_states=simulation_states, + average_episode_times=average_episode_times, + episode_times=episode_times, + forward_times=forward_times, + ) + + def _measure_model_throughput(self, atoms: Any) -> list[float]: + """Time single-structure model forward passes (energy + forces). + + For mlip ``ForceField`` models this times the pure network forward on a + pre-built graph (excluding neighbour-list construction), mirroring mlip-jax's + ``scripts/time_inference.py``. For external ASE calculators it times a forced + recomputation on the pre-built atoms (which includes the calculator's own + neighbour-list build, as there is no jittable forward to isolate). Warm-up + passes absorb compilation, the result is read to force device synchronisation, + and the slowest `FORWARD_TRIM_FRACTION` of passes are dropped before the caller + averages. + + Args: + atoms: The structure to run inference on. + + Returns: + The kept per-pass durations in seconds, or an empty list if it failed. + """ + try: + forward = self._build_forward_fn(atoms) + + for _ in range(NUM_FORWARD_WARMUP): + forward() + + times = [] + for _ in range(self._num_forward_timed): + start = time.perf_counter() + forward() + times.append(time.perf_counter() - start) + + # Drop the slowest passes (garbage-collection / scheduling spikes). + times.sort() + keep = max(1, round((1.0 - FORWARD_TRIM_FRACTION) * len(times))) + return times[:keep] + + except Exception as e: + logger.info( + "Error measuring model throughput on system %s: %s", str(atoms), str(e) + ) + return [] + + def _build_forward_fn(self, atoms: Any) -> Any: + """Build a zero-argument closure that runs (and blocks on) one forward pass. + + Args: + atoms: The structure to run inference on. + + Returns: + A callable taking no arguments that performs one forward pass. + """ + from mlip.models import ForceField # noqa: PLC0415 + + if isinstance(self.force_field, ForceField): + # mlip model: time the pure network forward on a pre-built graph. + import jax # noqa: PLC0415 + from mlip.data.chemical_system import ChemicalSystem # noqa: PLC0415 + from mlip.graph import Graph # noqa: PLC0415 + + force_field = self.force_field + chem_system = ChemicalSystem.from_ase_atoms( + atoms, get_property_fields=False + ) + graph = Graph.from_chemical_system( + chem_system, + force_field.cutoff_distance, + long_range_cutoff_angstrom=force_field.long_range_cutoff_distance, + ) + jitted = jax.jit(force_field) + + def mlip_forward() -> None: + prediction = jitted(graph) + jax.block_until_ready(prediction.forces) + + return mlip_forward + + # External ASE calculator: force a full recompute each call (ASE caches results + # for an unchanged system, so we invalidate via system_changes=all_changes). + from ase.calculators.calculator import all_changes # noqa: PLC0415 + + calculator = self.force_field + + def ase_forward() -> None: + calculator.calculate(atoms, ["energy", "forces"], all_changes) + _ = calculator.results["forces"] # read to force device synchronisation + + return ase_forward + + def analyze(self) -> InferenceSpeedResult: + """Aggregate the timings and compute the throughput score. + + Returns: + An `InferenceSpeedResult` object. + + Raises: + RuntimeError: If called before `run_model()`. + """ + if self.model_output is None: + raise RuntimeError("Must call run_model() first.") + + timestep_fs = float(self._md_kwargs["timestep_fs"]) + num_steps_per_episode = ( + self._md_kwargs["num_steps"] // self._md_kwargs["num_episodes"] + ) + + structure_results = [] + for i, structure_name in enumerate(self._structure_names): + episode_times = ( + self.model_output.episode_times[i] + if i < len(self.model_output.episode_times) + else [] + ) + forward_times = ( + self.model_output.forward_times[i] + if i < len(self.model_output.forward_times) + else [] + ) + + average_episode_time = self.model_output.average_episode_times[i] + average_step_time = ( + average_episode_time / num_steps_per_episode + if average_episode_time is not None + else None + ) + average_forward_time = ( + sum(forward_times) / len(forward_times) if forward_times else None + ) + + structure_results.append( + InferenceSpeedStructureResult( + structure_name=structure_name, + num_atoms=get_molecule_size_from_name(structure_name), + num_steps=self._md_kwargs["num_steps"], + num_episodes=self._md_kwargs["num_episodes"], + average_episode_time=average_episode_time, + average_step_time=average_step_time, + timestep_fs=timestep_fs, + episode_times=episode_times, + average_forward_time=average_forward_time, + forward_times=forward_times, + failed=average_step_time is None and average_forward_time is None, + ) + ) + + if len(self.model_output.simulation_states) == 0: + return InferenceSpeedResult( + structure_names=self._structure_names, failed=True + ) + + return InferenceSpeedResult( + structure_names=self._structure_names, + structures=structure_results, + score=self._compute_score(structure_results), + graph_cutoff_angstrom=self._graph_cutoff_angstrom(), + ) + + def _graph_cutoff_angstrom(self) -> float | None: + """Best-effort retrieval of the model's graph cutoff in Angstrom. + + mlip ``ForceField`` models expose ``cutoff_distance``; external ASE + calculators vary, so we probe a few common attribute names and fall back to + None when none are present. + + Returns: + The graph cutoff in Angstrom, or None if it cannot be determined. + """ + for attr in ("cutoff_distance", "cutoff", "r_max"): + value = getattr(self.force_field, attr, None) + if isinstance(value, (int, float)): + return float(value) + return None + + @staticmethod + def _compute_score( + structure_results: list[InferenceSpeedStructureResult], + ) -> float: + """Score speed via a Hill function on the per-atom model forward time. + + Each structure contributes ``1 / (1 + (t / t0) ** k)`` where ``t`` is its + per-atom forward-pass time (size-normalised so a single midpoint is meaningful + across system sizes); structures with no successful forward pass score 0. The + benchmark score is the mean. The forward pass is used (rather than the MD step) + because it is engine-independent. + + Args: + structure_results: The per-structure results. + + Returns: + The mean speed score in [0, 1]. + """ + per_atom_forward_times = [ + r.average_forward_time / r.num_atoms + if r.average_forward_time is not None + else None + for r in structure_results + ] + scores = compute_speed_score( + per_atom_forward_times, + midpoint=SCORE_PER_ATOM_FORWARD_TIME_MIDPOINT, + sharpness=SCORE_SHARPNESS, + ) + return float(scores.mean()) + + @functools.cached_property + def _structure_filenames(self) -> list[str]: + structure_names = sorted( + os.listdir(self.data_input_dir / self._dataset_name), + key=get_molecule_size_from_name, + ) + if self.run_mode == RunMode.DEV: + return structure_names[:NUM_DEV_SYSTEMS] + return structure_names + + @functools.cached_property + def _structure_names(self) -> list[str]: + return [Path(filename).stem for filename in self._structure_filenames] + + @functools.cached_property + def _md_kwargs(self) -> dict[str, Any]: + return ( + SIMULATION_CONFIG_DEV if self.run_mode == RunMode.DEV else SIMULATION_CONFIG + ) + + @functools.cached_property + def _num_forward_timed(self) -> int: + return ( + NUM_FORWARD_TIMED_DEV if self.run_mode == RunMode.DEV else NUM_FORWARD_TIMED + ) diff --git a/src/mlipaudit/scoring.py b/src/mlipaudit/scoring.py index ddf37aa2..ccbb0d6a 100644 --- a/src/mlipaudit/scoring.py +++ b/src/mlipaudit/scoring.py @@ -58,6 +58,41 @@ def compute_metric_score( return scores +def compute_speed_score( + values: list[float | None], midpoint: float, sharpness: float +) -> np.ndarray: + """Compute a monotonically decreasing score in (0, 1] for timing values. + + Uses a Hill function ``score = 1 / (1 + (value / midpoint) ** sharpness)``, so + smaller (faster) values score closer to 1, ``value == midpoint`` scores 0.5, and + larger (slower) values decay towards 0. Unlike `compute_metric_score`, there is no + flat region: being faster always improves the score, which suits speed metrics. + + Args: + values: Timing values (e.g. per-atom step times). ``None`` entries score 0. + midpoint: The value that scores 0.5. Must be positive. + sharpness: Controls how sharply scores separate around the midpoint. Must be + positive. + + Returns: + A NumPy array of scores in [0, 1]. + + Raises: + ValueError: If ``midpoint`` or ``sharpness`` is not positive. + """ + if midpoint <= 0: + raise ValueError("midpoint must be a positive number.") + if sharpness <= 0: + raise ValueError("sharpness must be a positive number.") + + numeric_values = np.array( + [v if v is not None else np.nan for v in values], dtype=float + ) + scores = 1.0 / (1.0 + np.power(numeric_values / midpoint, sharpness)) + scores[np.isnan(numeric_values)] = 0.0 + return scores + + def compute_benchmark_score( errors: list[list[float | None]], thresholds: list[float], diff --git a/src/mlipaudit/ui/__init__.py b/src/mlipaudit/ui/__init__.py index b50c93e6..c864384a 100644 --- a/src/mlipaudit/ui/__init__.py +++ b/src/mlipaudit/ui/__init__.py @@ -16,6 +16,7 @@ from mlipaudit.ui.conformer_selection import conformer_selection_page from mlipaudit.ui.dihedral_scan import dihedral_scan_page from mlipaudit.ui.folding_stability import folding_stability_page +from mlipaudit.ui.inference_speed import inference_speed_page from mlipaudit.ui.leaderboard import leaderboard_page from mlipaudit.ui.noncovalent_interactions import noncovalent_interactions_page from mlipaudit.ui.nudged_elastic_band import nudged_elastic_band_page diff --git a/src/mlipaudit/ui/inference_speed.py b/src/mlipaudit/ui/inference_speed.py new file mode 100644 index 00000000..bae4bec4 --- /dev/null +++ b/src/mlipaudit/ui/inference_speed.py @@ -0,0 +1,387 @@ +# Copyright 2025 InstaDeep Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Callable, TypeAlias + +import altair as alt +import numpy as np +import pandas as pd +import streamlit as st + +from mlipaudit.benchmarks import InferenceSpeedBenchmark, InferenceSpeedResult +from mlipaudit.ui.page_wrapper import UIPageWrapper +from mlipaudit.ui.utils import ( + display_failed_models, + fetch_selected_models, + filter_failed_results, + get_failed_models, +) + +ModelName: TypeAlias = str +BenchmarkResultForMultipleModels: TypeAlias = dict[ModelName, InferenceSpeedResult] + +#: ns/day = timestep_fs * NS_PER_DAY_FACTOR / step_time_seconds. +NS_PER_DAY_FACTOR = 0.0864 +#: Used when older results do not store the MD timestep (all past runs used 1 fs). +DEFAULT_TIMESTEP_FS = 1.0 + + +def _atoms_per_s(time_s: float, num_atoms: int, timestep_fs: float) -> float: + return num_atoms / time_s + + +def _ns_per_day(time_s: float, num_atoms: int, timestep_fs: float) -> float: + return timestep_fs * NS_PER_DAY_FACTOR / time_s + + +#: Selectable y-axis metrics. ``family`` selects which per-structure time drives the +#: metric: "model" uses the model forward-pass time (engine-independent), "md" uses the +#: MD step time (end-to-end). ``value`` maps (time_s, num_atoms, timestep_fs) -> value. +METRICS: dict[str, dict] = { + "Model throughput (atoms/s)": { + "family": "model", + "value": _atoms_per_s, + "format": ".0f", + }, + "Model forward passes/s": { + "family": "model", + "value": lambda time_s, num_atoms, timestep_fs: 1.0 / time_s, + "format": ".1f", + }, + "Model forward time (s)": { + "family": "model", + "value": lambda time_s, num_atoms, timestep_fs: time_s, + "format": ".4f", + }, + "MD throughput (ns/day)": { + "family": "md", + "value": _ns_per_day, + "format": ".1f", + }, + "MD steps/s": { + "family": "md", + "value": lambda time_s, num_atoms, timestep_fs: 1.0 / time_s, + "format": ".1f", + }, + "MD step time (s)": { + "family": "md", + "value": lambda time_s, num_atoms, timestep_fs: time_s, + "format": ".4f", + }, +} + + +def _structure_time_and_samples(structure, family: str) -> tuple: + """Return ``(central_time_s, [sample_times_s])`` for the metric family. + + The samples are used for error bars. Returns ``(None, [])`` if the structure has + no measurement for that family. + """ + if family == "model": + return structure.average_forward_time, list(structure.forward_times) + steps_per_episode = structure.num_steps / structure.num_episodes + samples = [e / steps_per_episode for e in structure.episode_times if e > 0] + return structure.average_step_time, samples + + +def _process_data_into_dataframe( + data: BenchmarkResultForMultipleModels, + selected_models: list[str], + metric_name: str, +) -> pd.DataFrame: + """Build a per-structure dataframe with the selected metric and its variance. + + Args: + data: The benchmark results per model. + selected_models: The models to include. + metric_name: The metric to compute (a key of `METRICS`). + + Returns: + A dataframe with one row per (model, structure). + """ + spec = METRICS[metric_name] + value_fn, family = spec["value"], spec["family"] + df_data = [] + for model_name, result in data.items(): + if model_name not in selected_models: + continue + cutoff = getattr(result, "graph_cutoff_angstrom", None) + for structure in result.structures: + time_s, samples = _structure_time_and_samples(structure, family) + if time_s is None or time_s <= 0: + continue + + timestep_fs = structure.timestep_fs or DEFAULT_TIMESTEP_FS + num_atoms = structure.num_atoms + metric_value = value_fn(time_s, num_atoms, timestep_fs) + + # Per-sample metric values give the spread shown as error bars. Results + # without per-sample times simply have no error bar. + sample_values = [ + value_fn(s, num_atoms, timestep_fs) for s in samples if s > 0 + ] + std = float(np.std(sample_values)) if sample_values else 0.0 + + df_data.append({ + "Model name": model_name, + "Structure": structure.structure_name, + "Num atoms": num_atoms, + "Time (s)": time_s, + "Graph cutoff (Å)": cutoff, + metric_name: metric_value, + "Metric low": max(metric_value - std, metric_value * 1e-3), + "Metric high": metric_value + std, + }) + return pd.DataFrame(df_data) + + +def _fit_scaling(num_atoms: np.ndarray, step_times: np.ndarray) -> tuple[float, float]: + """Fit a power law ``step_time = a * N^k`` in log-log space. + + Args: + num_atoms: System sizes. + step_times: Average step times (s). + + Returns: + A tuple ``(exponent_k, r_squared)``. + """ + log_n, log_t = np.log(num_atoms), np.log(step_times) + k, log_a = np.polyfit(log_n, log_t, 1) + predicted = log_a + k * log_n + ss_res = float(np.sum((log_t - predicted) ** 2)) + ss_tot = float(np.sum((log_t - log_t.mean()) ** 2)) + r_squared = 1.0 - ss_res / ss_tot if ss_tot > 0 else 1.0 + return float(k), r_squared + + +def _build_fit_lines(df: pd.DataFrame, metric_name: str) -> pd.DataFrame: + """Build smooth fitted scaling curves (in the selected metric) per model. + + Args: + df: The per-structure dataframe. + metric_name: The metric being plotted. + + Returns: + A dataframe of fitted curve points, empty if no model has >= 2 sizes. + """ + value_fn = METRICS[metric_name]["value"] + rows = [] + for model_name, group in df.groupby("Model name"): + if group["Num atoms"].nunique() < 2: + continue + num_atoms = group["Num atoms"].to_numpy(dtype=float) + times = group["Time (s)"].to_numpy(dtype=float) + k, log_a = np.polyfit(np.log(num_atoms), np.log(times), 1) + grid = np.geomspace(num_atoms.min(), num_atoms.max(), num=50) + fitted_times = np.exp(log_a) * grid**k + for n, fitted_time in zip(grid, fitted_times): + rows.append({ + "Model name": model_name, + "Num atoms": n, + metric_name: value_fn(fitted_time, n, DEFAULT_TIMESTEP_FS), + }) + return pd.DataFrame(rows) + + +def _summary_table(df: pd.DataFrame, metric_name: str) -> pd.DataFrame: + """Build the per-model summary (scaling exponent, fit quality, throughput). + + Args: + df: The per-structure dataframe. + metric_name: The metric being plotted (used for the headline value column). + + Returns: + A summary dataframe indexed by model name. + """ + rows = [] + for model_name, group in df.groupby("Model name"): + largest = group.loc[group["Num atoms"].idxmax()] + num_atoms = group["Num atoms"].to_numpy(dtype=float) + times = group["Time (s)"].to_numpy(dtype=float) + if group["Num atoms"].nunique() >= 2: + exponent, r_squared = _fit_scaling(num_atoms, times) + exponent_str = f"{exponent:.2f}" + r_squared_str = f"{r_squared:.3f}" + else: + exponent_str, r_squared_str = "N/A", "N/A" + + rows.append({ + "Model name": model_name, + "Graph cutoff (Å)": group["Graph cutoff (Å)"].iloc[0], + "Scaling exponent (time ∝ Nᵏ)": exponent_str, + "R²": r_squared_str, + f"{metric_name} @ largest system": largest[metric_name], + "Largest system (atoms)": int(largest["Num atoms"]), + }) + return pd.DataFrame(rows).set_index("Model name") + + +def plot_all_models_performance( + df: pd.DataFrame, metric_name: str, log_scale: bool +) -> alt.Chart: + """Plot the inference-speed curves for all models together. + + Args: + df: The per-structure dataframe. + metric_name: The metric to plot on the y-axis (a key of `METRICS`). + log_scale: Whether to use log-log axes. + + Returns: + The Altair chart. + """ + scale_type = "log" if log_scale else "linear" + value_format = METRICS[metric_name]["format"] + color = alt.Color("Model name:N", title="Model", legend=alt.Legend(title="Model")) + x = alt.X( + "Num atoms:Q", + title="System size (number of atoms)", + scale=alt.Scale(type=scale_type, zero=False), + ) + + points = ( + alt.Chart(df) + .mark_point(size=70, filled=True, opacity=0.85) + .encode( + x=x, + y=alt.Y( + f"{metric_name}:Q", + title=metric_name, + scale=alt.Scale(type=scale_type, zero=False), + ), + color=color, + tooltip=[ + alt.Tooltip("Model name:N", title="Model"), + alt.Tooltip("Structure:N", title="Structure"), + alt.Tooltip("Num atoms:Q", title="Number of atoms"), + alt.Tooltip(f"{metric_name}:Q", title=metric_name, format=value_format), + ], + ) + ) + + error_bars = ( + alt.Chart(df) + .mark_rule(opacity=0.6) + .encode( + x=x, + y=alt.Y("Metric low:Q", title=metric_name), + y2="Metric high:Q", + color=color, + ) + ) + + layers = [error_bars, points] + + fit_df = _build_fit_lines(df, metric_name) + if not fit_df.empty: + lines = ( + alt.Chart(fit_df) + .mark_line(strokeWidth=2) + .encode(x=x, y=alt.Y(f"{metric_name}:Q", title=metric_name), color=color) + ) + layers.append(lines) + + chart = alt.layer(*layers).properties(width=800, height=500).interactive() + st.altair_chart(chart, use_container_width=True) + return chart + + +def inference_speed_page( + data_func: Callable[[], BenchmarkResultForMultipleModels], +) -> None: + """Page for the visualization app for the inference-speed page. + + Args: + data_func: A data function that delivers the results on request. It does + not take any arguments and returns a dictionary with model names as + keys and the benchmark results objects as values. + """ + st.markdown("# Inference speed") + + st.markdown( + "This module assesses how fast MLIPs run, across several systems of varying " + "size. Two complementary speeds are reported: **model throughput** (the raw " + "forward pass, engine-independent) and **MD throughput** (end-to-end, with the " + "simulation engine). Switch between them with the metric selector; the gap " + "between the two reflects simulation overhead." + ) + + st.markdown( + "For more information, see the " + "[docs](https://instadeepai.github.io/mlipaudit" + "/benchmarks/general/inference_speed.html)." + ) + + if "inference_speed_data" not in st.session_state: + st.session_state.inference_speed_data = data_func() + + data = st.session_state.inference_speed_data + + if not data: + st.markdown("**No results to display**.") + return + + failed_models = get_failed_models(data) + display_failed_models(failed_models) + data = filter_failed_results(data) + + st.markdown("## Inference speed: throughput vs system size") + + selected_models = fetch_selected_models(available_models=list(data.keys())) + + if not selected_models: + st.markdown("**No results to display**.") + return + + col_metric, col_scale = st.columns([3, 1]) + with col_metric: + metric_name = st.selectbox("Metric", options=list(METRICS.keys()), index=0) + with col_scale: + log_scale = st.checkbox("Log–log axes", value=True) + + df = _process_data_into_dataframe(data, selected_models, metric_name) + + if df.empty: + st.markdown("**No results to display**.") + return + + plot_all_models_performance(df, metric_name, log_scale) + + st.caption( + "Points are per-system measurements (error bars show the spread across " + "repeats); lines are power-law fits. For external (ASE) models the model " + "metric includes neighbour-list construction, whereas for mlip models it is " + "the pure network forward. All times are wall-clock and hardware-relative — " + "only compare models run on the same GPU." + ) + + st.markdown("### Summary") + metric_column = f"{metric_name} @ largest system" + metric_format = "{:" + METRICS[metric_name]["format"] + "}" + st.dataframe( + _summary_table(df, metric_name).style.format({metric_column: metric_format}) + ) + + +class InferenceSpeedPageWrapper(UIPageWrapper): + """Page wrapper for the inference-speed benchmark.""" + + @classmethod + def get_page_func( # noqa: D102 + cls, + ) -> Callable[[Callable[[], BenchmarkResultForMultipleModels]], None]: + return inference_speed_page + + @classmethod + def get_benchmark_class(cls) -> type[InferenceSpeedBenchmark]: # noqa: D102 + return InferenceSpeedBenchmark diff --git a/tests/inference_speed/test_inference_speed.py b/tests/inference_speed/test_inference_speed.py new file mode 100644 index 00000000..0d6335ef --- /dev/null +++ b/tests/inference_speed/test_inference_speed.py @@ -0,0 +1,125 @@ +# Copyright 2025 InstaDeep Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re +from pathlib import Path + +import numpy as np +import pytest +from mlip.simulation import SimulationState + +from mlipaudit.benchmarks import ( + InferenceSpeedBenchmark, + InferenceSpeedModelOutput, + InferenceSpeedResult, +) +from mlipaudit.run_mode import RunMode + +INPUT_DATA_DIR = Path(__file__).parent.parent / "data" + + +@pytest.fixture +def inference_speed_benchmark( + request, + mocked_benchmark_init, # Use the generic init mock + mock_force_field, # Use the generic force field mock +) -> InferenceSpeedBenchmark: + """Assembles a fully configured and isolated InferenceSpeed instance. + This fixture is parameterized to handle the `run_mode` flag. + + Returns: + An initialized InferenceSpeed instance. + """ + is_fast_run = getattr(request, "param", False) + run_mode = RunMode.DEV if is_fast_run else RunMode.STANDARD + + return InferenceSpeedBenchmark( + force_field=mock_force_field, + data_input_dir=INPUT_DATA_DIR, + run_mode=run_mode, + ) + + +def test_reuses_scaling_dataset(inference_speed_benchmark): + """The benchmark reads its inputs from the shared ``scaling`` dataset.""" + assert inference_speed_benchmark._dataset_name == "scaling" + + +@pytest.mark.parametrize("inference_speed_benchmark", [True], indirect=True) +def test_full_run_with_mocked_engine(inference_speed_benchmark): + """Integration test testing the analysis of a full run of the benchmark.""" + benchmark = inference_speed_benchmark + + num_frames = 10 + positions_2jof = np.tile(np.ones((284, 3)), reps=(num_frames, 1, 1)) + positions_1r0r = np.tile(np.ones((748, 3)), reps=(num_frames, 1, 1)) + + benchmark.model_output = InferenceSpeedModelOutput( + structure_names=["284_2jof_A", "748_1r0r_I"], + simulation_states=[ + SimulationState(positions=positions_2jof), + SimulationState(positions=positions_1r0r), + ], + average_episode_times=[0.05, 0.1], + episode_times=[[0.04, 0.06], [0.09, 0.11]], + forward_times=[[0.002, 0.003], [0.004, 0.006]], + ) + + result = benchmark.analyze() + assert type(result) is InferenceSpeedResult + + assert len(result.structures) == 2 + s0 = result.structures[0] + assert s0.structure_name == "284_2jof_A" + assert s0.num_atoms == 284 + # MD throughput metric. + assert s0.average_step_time == 0.05 + assert s0.timestep_fs == 1 + assert s0.episode_times == [0.04, 0.06] + # Model throughput metric (mean of the timed forward passes). + assert s0.average_forward_time == 0.0025 + assert s0.forward_times == [0.002, 0.003] + assert not s0.failed + + # The benchmark produces a speed score in [0, 1] from the forward-pass times. + assert result.score is not None + assert 0.0 <= result.score <= 1.0 + + +@pytest.mark.parametrize("inference_speed_benchmark", [True], indirect=True) +def test_structure_fails_only_when_both_measurements_fail(inference_speed_benchmark): + """A structure is `failed` only if both the forward pass and MD produced nothing.""" + benchmark = inference_speed_benchmark + benchmark.model_output = InferenceSpeedModelOutput( + structure_names=["284_2jof_A", "748_1r0r_I"], + simulation_states=[None, None], + # First: MD failed but forward succeeded -> not failed. + # Second: both failed -> failed. + average_episode_times=[None, None], + episode_times=[[], []], + forward_times=[[0.002, 0.003], []], + ) + + result = benchmark.analyze() + assert result.structures[0].average_step_time is None + assert result.structures[0].average_forward_time == 0.0025 + assert not result.structures[0].failed + assert result.structures[1].failed + + +def test_analyze_raises_error_if_run_first(inference_speed_benchmark): + """Verifies the RuntimeError when analyze is called before run_model.""" + expected_message = "Must call run_model() first." + with pytest.raises(RuntimeError, match=re.escape(expected_message)): + inference_speed_benchmark.analyze() diff --git a/tests/test_scoring.py b/tests/test_scoring.py index 985a8355..21835c53 100644 --- a/tests/test_scoring.py +++ b/tests/test_scoring.py @@ -17,7 +17,29 @@ from numpy.testing import assert_allclose from mlipaudit.benchmarks import ConformerSelectionResult -from mlipaudit.scoring import compute_benchmark_score, compute_metric_score +from mlipaudit.scoring import ( + compute_benchmark_score, + compute_metric_score, + compute_speed_score, +) + + +def test_compute_speed_score(): + """Test the Hill-function speed score.""" + midpoint, sharpness = 2.0, 1.0 + + with pytest.raises(ValueError): + compute_speed_score([1.0], midpoint=0.0, sharpness=sharpness) + with pytest.raises(ValueError): + compute_speed_score([1.0], midpoint=midpoint, sharpness=0.0) + + # value == midpoint -> 0.5; faster -> higher; None -> 0. + scores = compute_speed_score([0.0, midpoint, None], midpoint, sharpness) + assert_allclose(scores, np.array([1.0, 0.5, 0.0])) + + # Monotonically decreasing in the value. + decreasing = compute_speed_score([1.0, 2.0, 4.0], midpoint, sharpness) + assert decreasing[0] > decreasing[1] > decreasing[2] def test_compute_metric_score():