diff --git a/src/fairchem/core/calculate/__init__.py b/src/fairchem/core/calculate/__init__.py index 75fead76ae..7cc9868327 100644 --- a/src/fairchem/core/calculate/__init__.py +++ b/src/fairchem/core/calculate/__init__.py @@ -17,6 +17,7 @@ FormationEnergyCalculator, ) from fairchem.core.components.batch_server import ( + ModelSpec, setup_batch_predict_server, setup_multiplexed_batch_predict_server, ) @@ -27,6 +28,7 @@ "FormationEnergyCalculator", "InferenceBatcher", "InferenceSettings", + "ModelSpec", "get_local_inference_raycluster", "get_slurm_inference_raycluster", "setup_batch_predict_server", diff --git a/src/fairchem/core/components/batch_server.py b/src/fairchem/core/components/batch_server.py index 992820a1fe..6faa5b07cc 100644 --- a/src/fairchem/core/components/batch_server.py +++ b/src/fairchem/core/components/batch_server.py @@ -7,14 +7,19 @@ from __future__ import annotations +import copy +import hashlib +import io import json import logging import os +import re import time -from collections import defaultdict, deque +from collections import OrderedDict, defaultdict, deque from dataclasses import asdict, dataclass, field +from functools import cached_property from multiprocessing import cpu_count -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal import ray import torch @@ -22,10 +27,11 @@ from ray.serve.schema import ApplicationStatus from fairchem.core.datasets.atomic_data import atomicdata_list_to_batch +from fairchem.core.units.mlip_unit.api.inference import guess_inference_settings if TYPE_CHECKING: from fairchem.core.datasets.atomic_data import AtomicData - from fairchem.core.units.mlip_unit import MLIPPredictUnit + from fairchem.core.units.mlip_unit import InferenceSettings, MLIPPredictUnit def _to_cpu(obj: Any) -> Any: @@ -35,7 +41,6 @@ def _to_cpu(obj: Any) -> Any: handles arbitrary object graphs containing tensors, ``nn.Module`` instances, OmegaConf containers, etc., without needing to walk and mutate the structure. """ - import io buf = io.BytesIO() torch.save(obj, buf) @@ -47,6 +52,107 @@ def _to_cpu(obj: Any) -> Any: # __init__s) so the two setup helpers can't drift apart silently. DEFAULT_MAX_BATCH_SIZE = 512 DEFAULT_BATCH_WAIT_TIMEOUT_S = 0.1 +MAX_NUM_MODELS_PER_REPLICA = 3 +MODEL_SPEC_CACHE_CAPACITY = MAX_NUM_MODELS_PER_REPLICA * 4 + + +def _canonicalize_model_spec_value(value: Any) -> Any: + """Convert nested model configuration values to stable JSON-compatible data.""" + if isinstance(value, dict): + return { + str(key): _canonicalize_model_spec_value(item) + for key, item in sorted(value.items(), key=lambda pair: str(pair[0])) + } + if isinstance(value, (set, frozenset)): + normalized = [_canonicalize_model_spec_value(item) for item in value] + return sorted( + normalized, + key=lambda item: json.dumps(item, sort_keys=True, separators=(",", ":")), + ) + if isinstance(value, (list, tuple)): + return [_canonicalize_model_spec_value(item) for item in value] + if isinstance(value, torch.dtype): + return str(value).removeprefix("torch.") + return value + + +@dataclass(frozen=True) +class ModelSpec: + """Typed configuration and deterministic identity for a multiplexed model.""" + + checkpoint: str + inference_settings: InferenceSettings | str = "default" + device: str | None = None + overrides: dict | None = None + source: Literal["auto", "path", "registry"] = "auto" + _canonical_config: dict[str, Any] = field(init=False, repr=False, compare=False) + _loader_settings: InferenceSettings = field(init=False, repr=False, compare=False) + _loader_overrides: dict | None = field(init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + if not isinstance(self.checkpoint, str) or not self.checkpoint: + raise ValueError("checkpoint must be a non-empty string") + if self.source not in ("auto", "path", "registry"): + raise ValueError( + "source must be one of 'auto', 'path', or 'registry', " + f"got {self.source!r}" + ) + + settings = copy.deepcopy(guess_inference_settings(self.inference_settings)) + overrides = ( + copy.deepcopy(self.overrides) if self.overrides is not None else None + ) + object.__setattr__(self, "inference_settings", copy.deepcopy(settings)) + object.__setattr__(self, "overrides", copy.deepcopy(overrides)) + object.__setattr__(self, "_loader_settings", settings) + object.__setattr__(self, "_loader_overrides", overrides) + + settings_config = settings.to_omegaconf() + settings_config.pop("_target_", None) + object.__setattr__( + self, + "_canonical_config", + { + "checkpoint": self.checkpoint, + "inference_settings": _canonicalize_model_spec_value(settings_config), + "device": self.device, + "overrides": _canonicalize_model_spec_value(overrides or {}), + "source": self.source, + }, + ) + + def canonical_dict(self) -> dict[str, Any]: + """Return a copy of the stable configuration used to derive ``model_id``.""" + return copy.deepcopy(self._canonical_config) + + @cached_property + def model_id(self) -> str: + """Return a readable deterministic identity token for Ray Serve routing.""" + canonical_json = json.dumps( + self.canonical_dict(), sort_keys=True, separators=(",", ":") + ) + digest = hashlib.sha256(canonical_json.encode("utf-8")).hexdigest()[:12] + short_name = re.sub( + r"[^A-Za-z0-9._-]+", "-", os.path.basename(self.checkpoint) + ).strip("-._") + short_name = (short_name or "model")[:48] + return f"{short_name}-{digest}" + + def resolve_device(self) -> str: + """Resolve the device on the replica when the client leaves it unspecified.""" + return self.device or ("cuda" if torch.cuda.is_available() else "cpu") + + def loader_settings(self) -> InferenceSettings: + """Return typed inference settings matching this spec's identity snapshot.""" + return copy.deepcopy(self._loader_settings) + + def loader_overrides(self) -> dict | None: + """Return model overrides matching this spec's identity snapshot.""" + return copy.deepcopy(self._loader_overrides) + + +class ModelSpecNotRegisteredError(KeyError): + """Raised when a routed identity has no configuration on the replica.""" @dataclass @@ -128,7 +234,7 @@ def get_predict_unit_attribute(self, attribute_name: str, **kwargs) -> Any: # server's device). return _to_cpu(getattr(self.predict_unit, attribute_name)) - def validate_atoms_data(self, atoms_info: dict, task_name: str) -> dict: + def validate_atoms_data(self, atoms_info: dict, task_name: str, **kwargs) -> dict: """ Run the predict unit's validation and return the (possibly mutated) atoms.info. @@ -211,7 +317,10 @@ def _run_batched_inference( return results async def __call__( - self, data: AtomicData, undo_element_references: bool = True + self, + data: AtomicData, + undo_element_references: bool = True, + **kwargs, ) -> dict: """ Main entry point for inference requests. @@ -335,9 +444,9 @@ class MultiplexedBatchPredictServer(BatchPredictServerMixin): Unlike ``BatchPredictServer`` which serves a single pre-loaded model, this deployment loads models on demand using ``@serve.multiplexed``. - Different clients can request different models by passing a ``model_id`` - and an LRU cache keeps up to ``max_num_models_per_replica`` models - resident on each replica. + Different clients request models with a :class:`ModelSpec`; its deterministic + ``model_id`` is used only for Serve routing and LRU cache identity. The spec + travels with each request and supplies the loader configuration. **Batching with per-model routing.** ``@serve.batch`` collects requests from concurrent ``__call__`` invocations. Because @@ -369,9 +478,52 @@ def __init__( Defaults to False. """ self.split_oom_batch = split_oom_batch + self._specs: OrderedDict[str, ModelSpec] = OrderedDict() + self._active_spec_counts: dict[str, int] = defaultdict(int) + self._spec_capacity = MODEL_SPEC_CACHE_CAPACITY self.configure_batching(max_batch_size, batch_wait_timeout_s) logging.info("MultiplexedBatchPredictServer initialized") + def _register_spec(self, spec: ModelSpec, *, pin: bool = False) -> str: + """Record a model configuration long enough for the multiplexed loader.""" + if not isinstance(spec, ModelSpec): + raise TypeError(f"spec must be a ModelSpec, got {type(spec).__name__}") + + model_id = spec.model_id + existing = self._specs.get(model_id) + if existing is not None and existing.canonical_dict() != spec.canonical_dict(): + raise ValueError(f"ModelSpec hash collision for model_id={model_id!r}") + self._specs[model_id] = spec + self._specs.move_to_end(model_id) + if pin: + self._active_spec_counts[model_id] += 1 + self._evict_specs() + return model_id + + def _release_spec(self, model_id: str) -> None: + """Unpin an in-flight spec and restore the steady-state cache bound.""" + remaining = self._active_spec_counts[model_id] - 1 + if remaining > 0: + self._active_spec_counts[model_id] = remaining + else: + self._active_spec_counts.pop(model_id, None) + self._evict_specs() + + def _evict_specs(self) -> None: + """Evict least-recent specs that are not needed by active requests.""" + while len(self._specs) > self._spec_capacity: + evictable_id = next( + ( + candidate_id + for candidate_id in self._specs + if candidate_id not in self._active_spec_counts + ), + None, + ) + if evictable_id is None: + return + del self._specs[evictable_id] + async def is_multiplexed(self) -> bool: return True @@ -382,6 +534,7 @@ async def predict( self, data_list: list[AtomicData], model_id_list: list[str], + spec_list: list[ModelSpec], undo_element_references: bool = True, ) -> list[dict]: """ @@ -398,6 +551,8 @@ async def predict( data_list: List of AtomicData objects (automatically batched by Ray Serve). model_id_list: Corresponding model IDs, one per request. + spec_list: Corresponding model specs, used to keep in-flight requests + registered even when the bounded spec table evicts older entries. undo_element_references: Whether to undo element references. Ray Serve batches this into a list; the first value is used. @@ -411,17 +566,26 @@ async def predict( else undo_element_references ) - # Group (original_index, data) pairs by model_id - groups: dict[str, list[tuple[int, AtomicData]]] = defaultdict(list) - for i, (data, model_id) in enumerate(zip(data_list, model_id_list)): - groups[model_id].append((i, data)) + # Group (original_index, data, spec) tuples by model_id. + groups: dict[str, list[tuple[int, AtomicData, ModelSpec]]] = defaultdict(list) + for i, (data, model_id, spec) in enumerate( + zip(data_list, model_id_list, spec_list) + ): + groups[model_id].append((i, data, spec)) results: list[dict | None] = [None] * len(data_list) for model_id, indexed_items in groups.items(): + group_spec = indexed_items[0][2] + registered_model_id = self._register_spec(group_spec) + if registered_model_id != model_id: + raise ValueError( + f"Batched ModelSpec identity mismatch: received {model_id!r}, " + f"derived {registered_model_id!r}." + ) predict_unit = await self.get_model(model_id) # LRU cache hit - indices, group_data = zip(*indexed_items) + indices, group_data, _ = zip(*indexed_items) group_results = self._run_batched_inference( list(group_data), predict_unit, undo_refs ) @@ -430,103 +594,92 @@ async def predict( return results - @serve.multiplexed(max_num_models_per_replica=3) + @serve.multiplexed(max_num_models_per_replica=MAX_NUM_MODELS_PER_REPLICA) async def get_model(self, model_id: str): - """ - Load (or retrieve from cache) a predict unit by model key. - - The ``@serve.multiplexed`` decorator caches the *return value* of - this method in an LRU cache (keyed by ``model_id``). Returning the - ``predict_unit`` directly means the cached object is the unit itself, - so subsequent calls for the same ``model_id`` skip the loading code - and return the cached unit without touching any instance state. - - Args: - model_id: Key in the format ``"checkpoint_name_or_path:settings"`` - where ``settings`` is one of the recognized inference setting - names (e.g. ``"default"``, ``"batch"``, ``"turbo"``) or an - empty string for the default settings. - - Returns: - The loaded ``MLIPPredictUnit`` for this model_id. - """ - parts = model_id.split(":", 1) - checkpoint = parts[0] - settings_name = parts[1] if len(parts) > 1 and parts[1] else "default" - - device = "cuda" if torch.cuda.is_available() else "cpu" + """Load or retrieve the predict unit identified by a registered spec.""" + try: + spec = self._specs[model_id] + except KeyError as err: + raise ModelSpecNotRegisteredError( + f"No ModelSpec is registered for model_id={model_id!r} on this " + "replica. Send the ModelSpec with the request before loading it." + ) from err + self._specs.move_to_end(model_id) + + loader_kwargs = { + "inference_settings": spec.loader_settings(), + "device": spec.resolve_device(), + } + overrides = spec.loader_overrides() + if overrides: + loader_kwargs["overrides"] = overrides - if os.path.isfile(checkpoint): + use_path = spec.source == "path" or ( + spec.source == "auto" and os.path.isfile(spec.checkpoint) + ) + if use_path: from fairchem.core.units.mlip_unit import load_predict_unit - predict_unit = load_predict_unit( - checkpoint, - inference_settings=settings_name, - device=device, - ) + predict_unit = load_predict_unit(spec.checkpoint, **loader_kwargs) else: from fairchem.core.calculate import pretrained_mlip predict_unit = pretrained_mlip.get_predict_unit( - checkpoint, - inference_settings=settings_name, - device=device, + spec.checkpoint, **loader_kwargs ) - logging.info(f"MultiplexedBatchPredictServer loaded model_id={model_id!r}") + logging.info( + "MultiplexedBatchPredictServer loaded model_id=%r spec=%s", + model_id, + json.dumps(spec.canonical_dict(), sort_keys=True), + ) return predict_unit async def get_predict_unit_attribute( - self, attribute_name: str, model_id: str | None = None + self, attribute_name: str, spec: ModelSpec ) -> Any: - """ - Get an attribute from a loaded predict unit. + """Get an attribute after registering and loading the requested model.""" + model_id = self._register_spec(spec, pin=True) + try: + predict_unit = await self.get_model(model_id) + return _to_cpu(getattr(predict_unit, attribute_name)) + finally: + self._release_spec(model_id) - Uses the ``multiplexed_model_id`` set on the request by the caller - to resolve the correct model first. - """ - model_id = model_id or serve.get_multiplexed_model_id() - predict_unit = await self.get_model(model_id) - attr = getattr(predict_unit, attribute_name) - # Move any CUDA tensors to CPU before returning so callers (which - # may be CPU-only Ray workers) can deserialize the result without - # requiring CUDA. - return _to_cpu(attr) - - async def validate_atoms_data(self, atoms_info: dict, task_name: str) -> dict: - """ - Run model-specific validation after loading the correct model. - """ + async def validate_atoms_data( + self, atoms_info: dict, task_name: str, spec: ModelSpec + ) -> dict: + """Run model-specific validation after registering the requested model.""" from ase import Atoms - model_id = serve.get_multiplexed_model_id() - predict_unit = await self.get_model(model_id) - stub = Atoms() - stub.info = atoms_info - predict_unit.validate_atoms_data(stub, task_name) - return stub.info + model_id = self._register_spec(spec, pin=True) + try: + predict_unit = await self.get_model(model_id) + stub = Atoms() + stub.info = atoms_info + predict_unit.validate_atoms_data(stub, task_name) + return stub.info + finally: + self._release_spec(model_id) async def __call__( - self, data: AtomicData, undo_element_references: bool = True + self, + data: AtomicData, + spec: ModelSpec, + undo_element_references: bool = True, ) -> dict: - """ - Main entry point for multiplexed inference requests. - - ``serve.get_multiplexed_model_id()`` is called here (per-request - context) and forwarded explicitly to ``predict()``. Inside the - ``@serve.batch`` function only one request context is active, so - the model ID cannot be reliably read there. - - Args: - data: Single AtomicData object. - undo_element_references: Whether to undo element references. - - Returns: - Prediction dictionary for this system. - """ - model_id = serve.get_multiplexed_model_id() - predictions = await self.predict(data, model_id, undo_element_references) - return predictions + """Register the request's spec and forward its identity into the batch.""" + model_id = self._register_spec(spec, pin=True) + try: + routed_model_id = serve.get_multiplexed_model_id() + if routed_model_id != model_id: + raise ValueError( + "Ray Serve multiplexed_model_id does not match the request's " + f"ModelSpec: routed={routed_model_id!r}, expected={model_id!r}." + ) + return await self.predict(data, model_id, spec, undo_element_references) + finally: + self._release_spec(model_id) def _init_ray_and_serve( diff --git a/src/fairchem/core/units/mlip_unit/predict.py b/src/fairchem/core/units/mlip_unit/predict.py index d9a4dbbad1..dc48ff8a03 100644 --- a/src/fairchem/core/units/mlip_unit/predict.py +++ b/src/fairchem/core/units/mlip_unit/predict.py @@ -36,7 +36,7 @@ get_device_for_local_rank, setup_env_local_multi_gpu, ) -from fairchem.core.components.batch_server import get_app_handle_with_retry +from fairchem.core.components.batch_server import ModelSpec, get_app_handle_with_retry from fairchem.core.datasets.atomic_data import AtomicData, warn_if_upcasting from fairchem.core.models.uma.nn.execution_backends import ( ExecutionMode, @@ -900,9 +900,10 @@ class BatchServerPredictUnit(MLIPPredictUnitProtocol): Works with both ``BatchPredictServer`` (single model) and ``MultiplexedBatchPredictServer`` (on-demand model loading). For - multiplexed deployments, pass ``multiplexed_model_id`` to ``from_deployment_connection_info`` - which binds the Ray Serve ``multiplexed_model_id`` to the handle so - that all requests are transparently routed to the correct model. + multiplexed deployments, pass a ``ModelSpec`` to + ``from_deployment_connection_info``. Its deterministic ``model_id`` is + bound to the Ray Serve handle for replica-affinity routing, while the + complete spec travels with every request. Can be constructed directly with a server handle, or via ``from_deployment_connection_info`` to connect to an already-running deployment. @@ -921,32 +922,32 @@ class BatchServerPredictUnit(MLIPPredictUnitProtocol): def __init__( self, server_handle: DeploymentHandle, - multiplexed_model_id: str | None = None, + model_spec: ModelSpec | None = None, ): """ Args: server_handle: Ray Serve deployment handle for a ``BatchPredictServer`` or ``MultiplexedBatchPredictServer``. - multiplexed_model_id: Optional model identifier for multiplexed - deployments in the format - ``"checkpoint_name_or_path:settings"``. When provided, the - handle is configured with Ray Serve's - ``multiplexed_model_id`` so that all calls are routed to - the correct model on the server. + model_spec: Typed model configuration for a multiplexed deployment. + Its derived ``model_id`` is bound to the handle for routing, and + the full spec is sent with each remote call. """ - if multiplexed_model_id is not None: + if model_spec is not None: + if not isinstance(model_spec, ModelSpec): + raise TypeError( + f"model_spec must be a ModelSpec, got {type(model_spec).__name__}" + ) if not server_handle.is_multiplexed.remote().result(): raise ValueError( - f"multiplexed_model_id={multiplexed_model_id!r} was " - "provided but the deployment is not a multiplexed " - "server. Use MultiplexedBatchPredictServer or remove " - "the multiplexed_model_id argument." + f"model_spec={model_spec!r} was provided but the deployment " + "is not a multiplexed server. Use " + "MultiplexedBatchPredictServer or remove the model_spec argument." ) server_handle = server_handle.options( - multiplexed_model_id=multiplexed_model_id + multiplexed_model_id=model_spec.model_id ) self.server_handle = server_handle - self._multiplexed_model_id = multiplexed_model_id + self._model_spec = model_spec # Identity-based cache for ``validate_atoms_data``. # ``ase_calculator.calculate()`` calls validate on every # optimizer step with the same ``atoms.info`` dict object; @@ -966,14 +967,14 @@ def __init__( self._request_timeout_s = _resolve_batch_server_timeout() @property - def multiplexed_model_id(self) -> str | None: - """ - The multiplexed model ID bound to this unit's server handle. + def model_spec(self) -> ModelSpec | None: + """The model configuration sent to a multiplexed deployment.""" + return self._model_spec - Read-only — changing this after construction would have no effect - on the already-configured handle. - """ - return self._multiplexed_model_id + @property + def multiplexed_model_id(self) -> str | None: + """The identity token derived from ``model_spec`` for Serve routing.""" + return self._model_spec.model_id if self._model_spec is not None else None @classmethod def from_deployment_connection_info( @@ -981,7 +982,7 @@ def from_deployment_connection_info( deployment_name: str = "predict-server", ray_address: str | None = None, namespace: str | None = None, - multiplexed_model_id: str | None = None, + model_spec: ModelSpec | None = None, ) -> BatchServerPredictUnit: """ Connect to an already-running server by deployment name. @@ -993,9 +994,7 @@ def from_deployment_connection_info( is unset, assumes Ray is already initialised locally. namespace: Ray namespace. Falls back to ``RAY_NAMESPACE_SERVE_FAIRCHEM`` env var. - multiplexed_model_id: Optional model identifier for multiplexed - deployments in the format - ``"checkpoint_name_or_path:settings"``. + model_spec: Typed model configuration for a multiplexed deployment. Returns: A ``BatchServerPredictUnit`` connected to the remote deployment. @@ -1015,7 +1014,7 @@ def from_deployment_connection_info( return cls( cls._handle_cache[cache_key], - multiplexed_model_id=multiplexed_model_id, + model_spec=model_spec, ) def predict(self, data: AtomicData, undo_element_references: bool = True) -> dict: @@ -1027,9 +1026,11 @@ def predict(self, data: AtomicData, undo_element_references: bool = True) -> dic Returns: Prediction dictionary """ - result = self.server_handle.remote(data, undo_element_references).result( - timeout_s=self._request_timeout_s - ) + result = self.server_handle.remote( + data, + spec=self._model_spec, + undo_element_references=undo_element_references, + ).result(timeout_s=self._request_timeout_s) return result def validate_atoms_data(self, atoms: Atoms, task_name: str) -> None: @@ -1053,7 +1054,7 @@ def validate_atoms_data(self, atoms: Atoms, task_name: str) -> None: if key in self._validated_info_keys: return updated_info = self.server_handle.validate_atoms_data.remote( - dict(atoms.info), task_name + dict(atoms.info), task_name, spec=self._model_spec ).result(timeout_s=self._request_timeout_s) atoms.info.update(updated_info) self._validated_info_keys.add(key) @@ -1061,23 +1062,23 @@ def validate_atoms_data(self, atoms: Atoms, task_name: str) -> None: @cached_property def dataset_to_tasks(self) -> dict: return self.server_handle.get_predict_unit_attribute.remote( - "dataset_to_tasks" + "dataset_to_tasks", spec=self._model_spec ).result(timeout_s=self._request_timeout_s) @cached_property def atom_refs(self) -> dict | None: - return self.server_handle.get_predict_unit_attribute.remote("atom_refs").result( - timeout_s=self._request_timeout_s - ) + return self.server_handle.get_predict_unit_attribute.remote( + "atom_refs", spec=self._model_spec + ).result(timeout_s=self._request_timeout_s) @cached_property def inference_settings(self) -> InferenceSettings: return self.server_handle.get_predict_unit_attribute.remote( - "inference_settings" + "inference_settings", spec=self._model_spec ).result(timeout_s=self._request_timeout_s) @cached_property def form_elem_refs(self) -> dict: return self.server_handle.get_predict_unit_attribute.remote( - "form_elem_refs" + "form_elem_refs", spec=self._model_spec ).result(timeout_s=self._request_timeout_s) diff --git a/tests/core/components/test_batch_server.py b/tests/core/components/test_batch_server.py index f5439b24d8..0374a46e1b 100644 --- a/tests/core/components/test_batch_server.py +++ b/tests/core/components/test_batch_server.py @@ -10,9 +10,9 @@ live deployment). 3. Multiplexed server (on-demand model loading via MultiplexedBatchPredictServer). -Models: uma-s-1p1, uma-s-1p2 (module-level pytestmark). Locked to +Models: uma-s-1p1, uma-s-1p2 (integration-test markers). Locked to UMA-S only because the base GPU runner OOMs with uma-m-1p1's - Ray Serve replicas. + Ray Serve replicas. Pure ModelSpec tests remain CPU-safe. CI: test_gpu_sweep (models shard). """ @@ -20,6 +20,7 @@ import json import uuid +from collections import OrderedDict, defaultdict from contextlib import suppress from pathlib import Path @@ -33,6 +34,9 @@ from fairchem.core import FAIRChemCalculator from fairchem.core.components.batch_server import ( + MODEL_SPEC_CACHE_CAPACITY, + ModelSpec, + MultiplexedBatchPredictServer, get_ray_connection_info, setup_batch_predict_server, setup_multiplexed_batch_predict_server, @@ -40,6 +44,7 @@ ) from fairchem.core.datasets.atomic_data import AtomicData from fairchem.core.launchers.cluster.ray_cluster import find_free_port +from fairchem.core.units.mlip_unit.api.inference import InferenceSettings from fairchem.core.units.mlip_unit.predict import BatchServerPredictUnit from tests.conftest import sweep_model, uma_models @@ -48,7 +53,170 @@ MULTIPLEXED_DEPLOYMENT_NAME = "multiplexed-predict-server" NAMESPACE = "fairchem_inference_test" -pytestmark = [pytest.mark.gpu, pytest.mark.pretrained("uma-s-1p1", "uma-s-1p2")] + +def test_model_spec_default_preset_has_canonical_identity(): + implicit = ModelSpec("x") + explicit = ModelSpec("x", "default") + + assert implicit.inference_settings == explicit.inference_settings + assert implicit.model_id == explicit.model_id + + +def test_model_spec_preserves_colon_bearing_checkpoint(): + spec = ModelSpec("s3://bucket/uma.pt", source="path") + + assert spec.checkpoint == "s3://bucket/uma.pt" + assert spec.canonical_dict()["checkpoint"] == "s3://bucket/uma.pt" + assert spec.model_id.startswith("uma.pt-") + + +def test_model_spec_identity_covers_settings_and_canonicalizes_sets(): + stress_a = InferenceSettings( + execution_mode="general", + predict_untrained_stress={"omat", "omol"}, + ) + stress_b = InferenceSettings( + execution_mode="general", + predict_untrained_stress={"omol", "omat"}, + ) + fast = InferenceSettings( + execution_mode="umas_fast_pytorch", + predict_untrained_stress={"omat", "omol"}, + ) + + assert ModelSpec("x", stress_a).model_id == ModelSpec("x", stress_b).model_id + assert ModelSpec("x", stress_a).model_id != ModelSpec("x", fast).model_id + assert ( + ModelSpec("x", stress_a).model_id + != ModelSpec("x", InferenceSettings(execution_mode="general")).model_id + ) + assert ( + ModelSpec("x", stress_a, device="cpu").model_id + != ModelSpec("x", stress_a, device="cuda").model_id + ) + assert ( + ModelSpec("x", stress_a, overrides={"backbone": {"max_neighbors": 64}}).model_id + != ModelSpec( + "x", stress_a, overrides={"backbone": {"max_neighbors": 128}} + ).model_id + ) + + +def test_model_spec_rejects_unknown_preset_at_construction(): + with pytest.raises(AssertionError, match="inference setting name"): + ModelSpec(checkpoint="x", inference_settings="nonsense") + + +def test_model_spec_model_id_is_stable(): + assert ModelSpec("x").model_id == "x-050e09e7dbf7" + + +def test_model_spec_identity_is_immutable_after_construction(): + settings = InferenceSettings(predict_untrained_stress={"omat"}) + overrides = {"backbone": {"max_neighbors": 64}} + spec = ModelSpec("x", settings, overrides=overrides) + original_id = spec.model_id + + settings.predict_untrained_stress.add("omol") + overrides["backbone"]["max_neighbors"] = 128 + spec.inference_settings.predict_untrained_stress.add("oc20") + spec.overrides["backbone"]["max_neighbors"] = 256 + + assert spec.model_id == original_id + assert spec.canonical_dict()["inference_settings"]["predict_untrained_stress"] == [ + "omat" + ] + assert spec.canonical_dict()["overrides"] == {"backbone": {"max_neighbors": 64}} + assert spec.loader_settings().predict_untrained_stress == {"omat"} + assert spec.loader_overrides() == {"backbone": {"max_neighbors": 64}} + + +def test_batch_server_predict_unit_binds_and_sends_model_spec(): + class FakeResponse: + def __init__(self, value): + self.value = value + + def result(self, timeout_s=None): + return self.value + + class FakeRemoteMethod: + def __init__(self, value): + self.value = value + + def remote(self, *args, **kwargs): + return FakeResponse(self.value) + + class FakeHandle: + def __init__(self): + self.is_multiplexed = FakeRemoteMethod(True) + self.bound_model_id = None + self.calls = [] + + def options(self, *, multiplexed_model_id): + self.bound_model_id = multiplexed_model_id + return self + + def remote(self, *args, **kwargs): + self.calls.append((args, kwargs)) + return FakeResponse({"energy": torch.tensor([1.0])}) + + handle = FakeHandle() + spec = ModelSpec("x", InferenceSettings(execution_mode="general")) + unit = BatchServerPredictUnit(handle, model_spec=spec) + data = object() + + result = unit.predict(data, undo_element_references=False) + + assert handle.bound_model_id == spec.model_id + assert unit.model_spec is spec + assert unit.multiplexed_model_id == spec.model_id + assert handle.calls == [ + ( + (data,), + {"spec": spec, "undo_element_references": False}, + ) + ] + assert result["energy"].item() == 1.0 + + +def test_model_spec_registry_evicts_least_recently_used_spec(): + server_class = MultiplexedBatchPredictServer.func_or_class + server = object.__new__(server_class) + server._specs = OrderedDict() + server._active_spec_counts = defaultdict(int) + server._spec_capacity = MODEL_SPEC_CACHE_CAPACITY + specs = [ + ModelSpec(f"model-{index}") for index in range(MODEL_SPEC_CACHE_CAPACITY + 1) + ] + + for spec in specs[:MODEL_SPEC_CACHE_CAPACITY]: + server._register_spec(spec) + server._register_spec(specs[0]) + server._register_spec(specs[-1]) + + assert len(server._specs) == MODEL_SPEC_CACHE_CAPACITY + assert specs[0].model_id in server._specs + assert specs[1].model_id not in server._specs + assert specs[-1].model_id in server._specs + + +def test_model_spec_registry_does_not_evict_in_flight_specs(): + server_class = MultiplexedBatchPredictServer.func_or_class + server = object.__new__(server_class) + server._specs = OrderedDict() + server._active_spec_counts = defaultdict(int) + server._spec_capacity = 2 + specs = [ModelSpec(f"model-{index}") for index in range(3)] + + for spec in specs: + server._register_spec(spec, pin=True) + + assert list(server._specs) == [spec.model_id for spec in specs] + + server._release_spec(specs[0].model_id) + + assert list(server._specs) == [spec.model_id for spec in specs[1:]] + assert len(server._specs) == server._spec_capacity @pytest.fixture() @@ -138,6 +306,8 @@ def local_ray_cluster_with_head_file(local_ray_cluster_with_inference, dashboard # --------------------------------------------------------------------------- +@pytest.mark.gpu() +@pytest.mark.pretrained("uma-s-1p1", "uma-s-1p2") def test_rayserve_remote_task_multiple_concurrent(local_ray_cluster_with_inference): """Test multiple concurrent Ray remote tasks hitting the inference server.""" @@ -173,6 +343,8 @@ def compute_predictions(dep_name: str, atoms_dict: dict): # --------------------------------------------------------------------------- +@pytest.mark.gpu() +@pytest.mark.pretrained("uma-s-1p1", "uma-s-1p2") def test_rayserve_external_multiple_systems(local_ray_cluster_with_head_file): """Test BatchServerPredictUnit from outside Ray with multiple systems.""" conn_info = get_ray_connection_info(local_ray_cluster_with_head_file) @@ -207,6 +379,8 @@ def test_rayserve_external_multiple_systems(local_ray_cluster_with_head_file): ), f"Stress shape mismatch for {atoms.get_chemical_formula()}" +@pytest.mark.gpu() +@pytest.mark.pretrained("uma-s-1p1", "uma-s-1p2") def test_rayserve_external_model_metadata(local_ray_cluster_with_inference): """Test that BatchServerPredictUnit correctly fetches model metadata.""" @@ -223,6 +397,8 @@ def test_rayserve_external_model_metadata(local_ray_cluster_with_inference): ), f"Expected 'omat' in tasks, got: {list(dataset_to_tasks.keys())}" +@pytest.mark.gpu() +@pytest.mark.pretrained("uma-s-1p1", "uma-s-1p2") def test_rayserve_external_vs_local_comparison( local_ray_cluster_with_inference, uma_predict_unit ): @@ -275,14 +451,13 @@ def test_rayserve_external_vs_local_comparison( @pytest.fixture() -def uma_multiplexed_model_id(request): +def uma_model_spec(request): """ - Multiplexed model ID for the sweep model, or first available UMA model. + Model spec for the sweep model, or first available UMA model. Honors ``--sweep-model`` so per-model sweep CI jobs target the requested checkpoint. Skips when the sweep value is a filesystem - path — the multiplexed server is keyed by registered model name, - so paths cannot be exercised here. + path because this fixture exercises registry-backed model loading. """ available_uma = uma_models() if not available_uma: @@ -297,7 +472,7 @@ def uma_multiplexed_model_id(request): model = sweep else: model = available_uma[0] - return f"{model}:default" + return ModelSpec(model, source="registry") @pytest.fixture() @@ -337,12 +512,14 @@ def local_multiplexed_cluster(): ray.shutdown() +@pytest.mark.gpu() +@pytest.mark.pretrained("uma-s-1p1", "uma-s-1p2") def test_multiplexed_single_model( - local_multiplexed_cluster, uma_multiplexed_model_id, uma_predict_unit + local_multiplexed_cluster, uma_model_spec, uma_predict_unit ): """Test loading a single model via the multiplexed server.""" unit = BatchServerPredictUnit.from_deployment_connection_info( - multiplexed_model_id=uma_multiplexed_model_id, + model_spec=uma_model_spec, deployment_name=MULTIPLEXED_DEPLOYMENT_NAME, ) @@ -364,28 +541,30 @@ def test_multiplexed_single_model( npt.assert_allclose(stress_mux, stress_local, atol=ATOL) -def test_multiplexed_switch_models(local_multiplexed_cluster, uma_multiplexed_model_id): +@pytest.mark.gpu() +@pytest.mark.pretrained("uma-s-1p1", "uma-s-1p2") +def test_multiplexed_switch_models(local_multiplexed_cluster, uma_model_spec): """Test switching between two different model keys.""" available_uma = uma_models() if len(available_uma) < 2: pytest.skip("Need at least 2 UMA models to test switching") - # uma_multiplexed_model_id already encodes the sweep target (or first - # UMA model). Pick any other UMA model as the second key. - primary = uma_multiplexed_model_id.split(":")[0] + # uma_model_spec already identifies the sweep target (or first UMA model). + # Pick any other UMA model as the second spec. + primary = uma_model_spec.checkpoint other_candidates = [m for m in available_uma if m != primary] if not other_candidates: pytest.skip("No second UMA model available that differs from the primary") - key_a = uma_multiplexed_model_id - key_b = f"{other_candidates[0]}:default" + spec_a = uma_model_spec + spec_b = ModelSpec(other_candidates[0]) unit_a = BatchServerPredictUnit.from_deployment_connection_info( - multiplexed_model_id=key_a, + model_spec=spec_a, deployment_name=MULTIPLEXED_DEPLOYMENT_NAME, ) unit_b = BatchServerPredictUnit.from_deployment_connection_info( - multiplexed_model_id=key_b, + model_spec=spec_b, deployment_name=MULTIPLEXED_DEPLOYMENT_NAME, ) @@ -402,20 +581,18 @@ def test_multiplexed_switch_models(local_multiplexed_cluster, uma_multiplexed_mo ), "Different models should produce different energies" -def test_multiplexed_concurrent_requests( - local_multiplexed_cluster, uma_multiplexed_model_id -): +@pytest.mark.gpu() +@pytest.mark.pretrained("uma-s-1p1", "uma-s-1p2") +def test_multiplexed_concurrent_requests(local_multiplexed_cluster, uma_model_spec): """Test concurrent requests to the multiplexed server.""" @ray.remote - def compute_predictions_mux( - dep_name: str, multiplexed_model_id: str, atoms_dict: dict - ): + def compute_predictions_mux(dep_name: str, model_spec: ModelSpec, atoms_dict: dict): """Ray remote task using BatchServerPredictUnit directly.""" atoms = Atoms.fromdict(atoms_dict) atomic_data = AtomicData.from_ase(atoms, task_name="omat") unit = BatchServerPredictUnit.from_deployment_connection_info( - multiplexed_model_id=multiplexed_model_id, + model_spec=model_spec, deployment_name=dep_name, ) return unit.predict(atomic_data, undo_element_references=True) @@ -424,9 +601,7 @@ def compute_predictions_mux( atoms_dicts = [a.todict() for a in systems] futures = [ - compute_predictions_mux.remote( - MULTIPLEXED_DEPLOYMENT_NAME, uma_multiplexed_model_id, d - ) + compute_predictions_mux.remote(MULTIPLEXED_DEPLOYMENT_NAME, uma_model_spec, d) for d in atoms_dicts ] results = ray.get(futures)