diff --git a/deepmd/entrypoints/convert_backend.py b/deepmd/entrypoints/convert_backend.py index 43cb901449..817d3996a4 100644 --- a/deepmd/entrypoints/convert_backend.py +++ b/deepmd/entrypoints/convert_backend.py @@ -30,6 +30,14 @@ def convert_backend( If True, export .pt2/.pte models with per-atom virial correction. This adds ~2.5x inference cost. Default False. Silently ignored (with a warning) for backends that don't support the flag. + + Notes + ----- + Backend conversion preserves an explicit ``lower_input_kind`` reported by + the source serializer. Sources without this metadata retain the target's + automatic lower selection for backward compatibility. A target backend + that cannot represent an explicit non-dense lower is rejected rather than + silently changing the model function. """ inp_backend: Backend = Backend.detect_backend_by_model(INPUT)() out_backend: Backend = Backend.detect_backend_by_model(OUTPUT)() @@ -40,8 +48,18 @@ def convert_backend( sig = inspect.signature(out_hook) hook_kwargs: dict[str, Any] = {} + lower_input_kind = data.get("lower_input_kind") if "lower_kind" in sig.parameters: - hook_kwargs["lower_kind"] = "auto" + hook_kwargs["lower_kind"] = ( + lower_input_kind if lower_input_kind is not None else "auto" + ) + elif lower_input_kind not in (None, "nlist"): + raise ValueError( + f"Cannot preserve lower_input_kind {lower_input_kind!r} when " + f"converting to output backend {out_backend.name!r}: its " + "deserializer does not accept a lower_kind. Retrain or freeze the " + "model with that backend instead of converting this artifact." + ) if "do_atomic_virial" in sig.parameters: hook_kwargs["do_atomic_virial"] = atomic_virial elif atomic_virial: diff --git a/deepmd/jax/utils/serialization.py b/deepmd/jax/utils/serialization.py index 62a6851160..33801e4149 100644 --- a/deepmd/jax/utils/serialization.py +++ b/deepmd/jax/utils/serialization.py @@ -427,6 +427,7 @@ def restore_model(model_params: dict, model_state: dict) -> BaseModel: "jax_version": jax.__version__, "model": model_dict, "model_def_script": model_def_script, + "lower_input_kind": "nlist", "@variables": {}, } if min_nbor_dist is not None: @@ -436,6 +437,7 @@ def restore_model(model_params: dict, model_state: dict) -> BaseModel: data = load_dp_model(model_file) data.pop("constants") data["@variables"].pop("stablehlo") + data["lower_input_kind"] = "nlist" return data elif model_file.endswith(".savedmodel"): raise ValueError( diff --git a/deepmd/pt/model/model/model.py b/deepmd/pt/model/model/model.py index 5f89ff50ad..02ecedbce7 100644 --- a/deepmd/pt/model/model/model.py +++ b/deepmd/pt/model/model/model.py @@ -26,6 +26,17 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: "min_nbor_dist", torch.tensor(-1.0, dtype=torch.float64, device=env.DEVICE) ) + def export_lower_input_kind(self) -> str: + """Return the lower-input ABI that preserves this model's semantics. + + Returns + ------- + str + ``"nlist"`` for the standard PyTorch model contract. Models with + a graph-native deployment ABI override this method. + """ + return "nlist" + def compute_or_load_stat( self, sampled_func: Any, diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index c0adf618c8..029f6c1dd0 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -454,6 +454,17 @@ def has_spin(self) -> bool: """Returns whether it has spin input and output.""" return True + def export_lower_input_kind(self) -> str: + """Return the dense ABI used by the virtual-atom spin model. + + Returns + ------- + str + ``"nlist"``, because virtual atoms are expanded inside the + bounded neighbor-list contract. + """ + return "nlist" + @torch.jit.export def has_message_passing(self) -> bool: """Returns whether the model has message passing.""" diff --git a/deepmd/pt/utils/serialization.py b/deepmd/pt/utils/serialization.py index db23eef4dc..6cdd65712f 100644 --- a/deepmd/pt/utils/serialization.py +++ b/deepmd/pt/utils/serialization.py @@ -53,6 +53,7 @@ def serialize_from_file(model_file: str) -> dict: "pt_version": str(torch.__version__), "model": model_dict, "model_def_script": model_def_script, + "lower_input_kind": model.export_lower_input_kind(), "@variables": {}, } if model.get_min_nbor_dist() is not None: diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 2ad89a22c4..eb0bcf6081 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -52,6 +52,10 @@ # --------------------------------------------------------------------------- PT2_EXTRA_PREFIX = "model/extra/" +_SUPPORTED_LOWER_INPUT_KINDS = frozenset( + {"nlist", "graph", "dpa1_canonical", "dpa4c_canonical"} +) + def _strip_shape_assertions(graph_module: torch.nn.Module) -> None: """Neutralise deferred shape-guard assertion nodes in an exported graph. @@ -1242,7 +1246,8 @@ def serialize_from_file(model_file: str) -> dict: dict The serialized model data. If the archive contains ``model_def_script.json`` (training config), it is included - under the ``"model_def_script"`` key. + under the ``"model_def_script"`` key. ``lower_input_kind`` records + the concrete lower ABI from the artifact metadata. """ if model_file.endswith(".pt2"): return _serialize_from_file_pt2(model_file) @@ -1252,10 +1257,20 @@ def serialize_from_file(model_file: str) -> dict: def _serialize_from_file_pte(model_file: str) -> dict: """Serialize a .pte model file to a dictionary.""" - extra_files = {"model.json": "", "model_def_script.json": ""} + extra_files = { + "model.json": "", + "model_def_script.json": "", + "metadata.json": "", + } torch.export.load(model_file, extra_files=extra_files) model_dict = json.loads(extra_files["model.json"]) model_dict = _json_to_numpy(model_dict) + metadata = ( + json.loads(extra_files["metadata.json"]) if extra_files["metadata.json"] else {} + ) + model_dict["lower_input_kind"] = metadata.get( + "lower_input_kind", model_dict.get("lower_input_kind", "nlist") + ) if extra_files["model_def_script.json"]: model_dict["model_def_script"] = json.loads( extra_files["model_def_script.json"] @@ -1273,6 +1288,7 @@ def _serialize_from_file_pt2(model_file: str) -> dict: model_json_entry = PT2_EXTRA_PREFIX + "model.json" model_def_script_entry = PT2_EXTRA_PREFIX + "model_def_script.json" + metadata_entry = PT2_EXTRA_PREFIX + "metadata.json" with zipfile.ZipFile(model_file, "r") as zf: names = zf.namelist() if model_json_entry not in names: @@ -1283,8 +1299,15 @@ def _serialize_from_file_pt2(model_file: str) -> dict: model_def_script_json = "" if model_def_script_entry in names: model_def_script_json = zf.read(model_def_script_entry).decode("utf-8") + metadata_json = "" + if metadata_entry in names: + metadata_json = zf.read(metadata_entry).decode("utf-8") model_dict = json.loads(model_json) model_dict = _json_to_numpy(model_dict) + metadata = json.loads(metadata_json) if metadata_json else {} + model_dict["lower_input_kind"] = metadata.get( + "lower_input_kind", model_dict.get("lower_input_kind", "nlist") + ) if model_def_script_json: model_dict["model_def_script"] = json.loads(model_def_script_json) return model_dict @@ -1393,14 +1416,21 @@ def deserialize_to_file( (``atype``/``n_node``/``edge_index``/``edge_vec``/``edge_mask`` and the destination/source CSR views) with a DYNAMIC edge axis ``E`` (``Dim("nedge", min=2)``), so the artifact accepts any system size. - ``"auto"`` (used by ``convert-backend``) resolves to ``"graph"`` for an - exportable graph-lower ``.pt2`` and ``"nlist"`` otherwise (see - :func:`_resolve_lower_kind`). A graph lower always preserves the fused - inference operators (``DP_CUDA_INFER >= 2``) and the per-atom virial. + ``"auto"`` resolves to ``"graph"`` for an exportable graph-lower + ``.pt2`` and ``"nlist"`` otherwise (see :func:`_resolve_lower_kind`). + Backend conversion passes the source artifact's concrete lower kind + instead, preserving its execution semantics. A graph lower always + preserves the fused inference operators (``DP_CUDA_INFER >= 2``) and + the per-atom virial. The selected schema is recorded as ``lower_input_kind`` in ``metadata.json``. """ lower_kind = _resolve_lower_kind(model_file, data, lower_kind) + if lower_kind not in _SUPPORTED_LOWER_INPUT_KINDS: + raise ValueError( + f"Unsupported lower_kind {lower_kind!r}; expected one of " + f"{sorted(_SUPPORTED_LOWER_INPUT_KINDS)}." + ) if data["model"].get("type") == "native_spin" and lower_kind not in ( "graph", "dpa4c_canonical", diff --git a/deepmd/tf/utils/serialization.py b/deepmd/tf/utils/serialization.py index 1d2f1b597f..81d46c1e2d 100644 --- a/deepmd/tf/utils/serialization.py +++ b/deepmd/tf/utils/serialization.py @@ -50,6 +50,7 @@ def serialize_from_file(model_file: str) -> dict: "tf_version": tf.__version__, "model": model_dict, "model_def_script": jdata["model"], + "lower_input_kind": "nlist", } # neighbor stat information try: diff --git a/deepmd/tf2/utils/serialization.py b/deepmd/tf2/utils/serialization.py index 63bbdf6583..b859d16b10 100644 --- a/deepmd/tf2/utils/serialization.py +++ b/deepmd/tf2/utils/serialization.py @@ -611,6 +611,7 @@ def serialize_from_file(model_file: str) -> dict: "backend": "TensorFlow2", "model": model_payload, "model_def_script": model_def_script, + "lower_input_kind": "nlist", "shared_links": state.get("shared_links"), "@variables": { "current_step": int(state.get("current_step", 0)), diff --git a/doc/backend.md b/doc/backend.md index e09f230eea..d85a0ee8a2 100644 --- a/doc/backend.md +++ b/doc/backend.md @@ -154,3 +154,12 @@ then selects its `edge_vec`, dense `nlist`, NeighborGraph, or compact ## Convert model files between backends If a model is supported by two backends, one can use [`dp convert-backend`](./cli.rst) to convert the model file between these two backends. + +Backend conversion preserves the concrete `lower_input_kind` reported by the +source serializer. Dense TensorFlow, TensorFlow 2, JAX, and standard PyTorch +models therefore remain dense `nlist` models when converted to `.pt2`; model +families with a graph-native deployment ABI report their corresponding kind. +Compiled `.pt2` and `.pte` artifacts retain the kind recorded in their metadata. +A conversion is rejected when the target backend cannot represent an explicit +source kind. Legacy model files without lower metadata retain target-specific +automatic selection for backward compatibility. diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index 2a34e2bbe5..913aca4d51 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -275,6 +275,7 @@ def setUp(self) -> None: "model": model.serialize(), "backend": "test", "model_def_script": model_def_script, + "lower_input_kind": "nlist", } def tearDown(self) -> None: @@ -319,6 +320,7 @@ def setUp(self) -> None: "model": model.serialize(), "backend": "test", "model_def_script": model_def_script, + "lower_input_kind": "nlist", } def tearDown(self) -> None: @@ -375,6 +377,7 @@ def setUp(self) -> None: "model": model.serialize(), "backend": "test", "model_def_script": model_def_script, + "lower_input_kind": "nlist", } def tearDown(self) -> None: @@ -422,6 +425,7 @@ def setUp(self) -> None: "model": model.serialize(), "backend": "test", "model_def_script": model_def_script, + "lower_input_kind": "nlist", } def tearDown(self) -> None: diff --git a/source/tests/jax/test_hlo.py b/source/tests/jax/test_hlo.py index 34488f4629..2e923205bd 100644 --- a/source/tests/jax/test_hlo.py +++ b/source/tests/jax/test_hlo.py @@ -1,9 +1,21 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Regression tests for metadata exposed by serialized JAX HLO models.""" +from types import ( + SimpleNamespace, +) + +import pytest +from typing_extensions import ( + Self, +) + from deepmd.jax.model.hlo import ( HLO, ) +from deepmd.jax.utils import ( + serialization, +) def test_hlo_get_nnei_uses_stored_selection() -> None: @@ -17,3 +29,47 @@ def test_hlo_get_nnei_uses_stored_selection() -> None: model.sel = [6, 12, 1] assert model.get_nnei() == sum(model.sel) + + +def test_hlo_serialization_declares_dense_lower( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A JAX HLO artifact exposes its dense source execution semantics.""" + stored_data = { + "model": {}, + "constants": {}, + "@variables": {"stablehlo": b"module"}, + } + monkeypatch.setattr(serialization, "load_dp_model", lambda _path: stored_data) + + data = serialization.serialize_from_file("model.hlo") + + assert data["lower_input_kind"] == "nlist" + + +def test_checkpoint_serialization_declares_dense_lower( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A JAX training checkpoint exposes its dense source semantics.""" + + class Checkpointer: + def __init__(self, _handler: object) -> None: + pass + + def __enter__(self) -> Self: + return self + + def __exit__(self, *_args: object) -> None: + pass + + def restore(self, *_args: object, **_kwargs: object) -> SimpleNamespace: + return SimpleNamespace( + state={}, + model_def_script={"model_dict": {}}, + ) + + monkeypatch.setattr(serialization.ocp, "Checkpointer", Checkpointer) + + data = serialization.serialize_from_file("model.jax") + + assert data["lower_input_kind"] == "nlist" diff --git a/source/tests/pt/model/test_ener_spin_model.py b/source/tests/pt/model/test_ener_spin_model.py index 49864375e2..80199934d2 100644 --- a/source/tests/pt/model/test_ener_spin_model.py +++ b/source/tests/pt/model/test_ener_spin_model.py @@ -111,6 +111,10 @@ def test_output_shape( torch.testing.assert_close(result["force"].shape, [nframes, nloc, 3]) torch.testing.assert_close(result["force_mag"].shape, [nframes, nloc, 3]) + def test_export_lower_input_kind(self) -> None: + """Virtual-atom spin models retain the dense export ABI.""" + self.assertEqual(self.model.export_lower_input_kind(), "nlist") + def test_input_output_process(self) -> None: nframes, nloc = self.coord.shape[:2] self.real_ntypes = self.model.spin.get_ntypes_real() diff --git a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py index 7ee61a1526..f086d24c2c 100644 --- a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py +++ b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py @@ -14,9 +14,13 @@ import tempfile import zipfile +import numpy as np import pytest import torch +from deepmd.entrypoints.convert_backend import ( + convert_backend, +) from deepmd.pt_expt.model.graph_lower import ( graph_edge_dtype, ) @@ -24,6 +28,7 @@ _needs_with_comm_artifact, _supports_graph_export, deserialize_to_file, + serialize_from_file, ) # dpa1 with attn_layer == 0 — the energy model exercised by the graph path. @@ -86,11 +91,159 @@ def _read_metadata(pt2_path: str) -> dict: return json.loads(raw) +@pytest.mark.parametrize( + "lower_input_kind", + ["nlist", "graph", "dpa1_canonical", "dpa4c_canonical", "edge_vec"], +) +def test_pt2_serialization_preserves_lower_input_kind( + tmp_path, lower_input_kind: str +) -> None: + """The interchange dictionary exposes the artifact's lower semantics.""" + model_file = tmp_path / "model.pt2" + with zipfile.ZipFile(model_file, "w") as zf: + zf.writestr("model/extra/model.json", json.dumps({"model": {}})) + zf.writestr( + "model/extra/metadata.json", + json.dumps({"lower_input_kind": lower_input_kind}), + ) + + data = serialize_from_file(str(model_file)) + + assert data["lower_input_kind"] == lower_input_kind + + +@pytest.mark.parametrize( + "lower_input_kind", + ["nlist", "graph", "dpa1_canonical", "dpa4c_canonical", "edge_vec"], +) +def test_pte_serialization_preserves_lower_input_kind( + tmp_path, monkeypatch: pytest.MonkeyPatch, lower_input_kind: str +) -> None: + """PTE extra metadata has the same interchange contract as PT2.""" + + def load_exported_program(_model_file: str, *, extra_files: dict[str, str]) -> None: + extra_files["model.json"] = json.dumps({"model": {}}) + extra_files["model_def_script.json"] = "" + extra_files["metadata.json"] = json.dumps( + {"lower_input_kind": lower_input_kind} + ) + + monkeypatch.setattr(torch.export, "load", load_exported_program) + + data = serialize_from_file(str(tmp_path / "model.pte")) + + assert data["lower_input_kind"] == lower_input_kind + + +@pytest.mark.parametrize( + ("embedded_lower_input_kind", "expected"), + [("graph", "graph"), (None, "nlist")], +) +def test_pt2_serialization_legacy_lower_input_kind_fallback( + tmp_path, + embedded_lower_input_kind: str | None, + expected: str, +) -> None: + """Legacy PT2 archives use embedded metadata, then dense fallback.""" + model_file = tmp_path / "model.pt2" + model_data: dict[str, object] = {"model": {}} + if embedded_lower_input_kind is not None: + model_data["lower_input_kind"] = embedded_lower_input_kind + with zipfile.ZipFile(model_file, "w") as zf: + zf.writestr("model/extra/model.json", json.dumps(model_data)) + + data = serialize_from_file(str(model_file)) + + assert data["lower_input_kind"] == expected + + +@pytest.mark.parametrize( + ("embedded_lower_input_kind", "expected"), + [("graph", "graph"), (None, "nlist")], +) +def test_pte_serialization_legacy_lower_input_kind_fallback( + tmp_path, + monkeypatch: pytest.MonkeyPatch, + embedded_lower_input_kind: str | None, + expected: str, +) -> None: + """Legacy PTE archives use embedded metadata, then dense fallback.""" + model_data: dict[str, object] = {"model": {}} + if embedded_lower_input_kind is not None: + model_data["lower_input_kind"] = embedded_lower_input_kind + + def load_exported_program(_model_file: str, *, extra_files: dict[str, str]) -> None: + extra_files["model.json"] = json.dumps(model_data) + extra_files["model_def_script.json"] = "" + extra_files["metadata.json"] = "" + + monkeypatch.setattr(torch.export, "load", load_exported_program) + + data = serialize_from_file(str(tmp_path / "model.pte")) + + assert data["lower_input_kind"] == expected + + @pytest.fixture(scope="module") def dpa1_dpmodel_data() -> dict: return _build_dpa1_data() +def test_convert_regular_pt_dpa1_preserves_dense_semantics(tmp_path) -> None: + """A nonzero-davg PT artifact remains numerically dense after conversion.""" + from deepmd.infer import ( + DeepPot, + ) + from deepmd.pt.utils.serialization import deserialize_to_file as deserialize_to_pt + from deepmd.pt.utils.serialization import serialize_from_file as serialize_from_pt + + data = _build_dpa1_data() + descriptor_variables = data["model"]["descriptor"]["@variables"] + descriptor_variables["davg"] = np.full_like(descriptor_variables["davg"], 0.01) + source_model = tmp_path / "model.pth" + converted_model = tmp_path / "model.pt2" + deserialize_to_pt(str(source_model), copy.deepcopy(data)) + + source_data = serialize_from_pt(str(source_model)) + assert source_data["lower_input_kind"] == "nlist" + + convert_backend(INPUT=str(source_model), OUTPUT=str(converted_model)) + assert _read_metadata(str(converted_model))["lower_input_kind"] == "nlist" + + coord = np.array( + [ + [0.0, 0.0, 0.0], + [1.1, 0.2, 0.1], + [0.3, 1.4, 0.2], + [1.2, 1.1, 0.8], + ], + dtype=np.float64, + )[None, ...] + atype = np.array([[0, 1, 0, 1]], dtype=np.int32) + source_result = DeepPot(str(source_model), auto_batch_size=False).eval( + coord, None, atype + ) + converted_result = DeepPot(str(converted_model), auto_batch_size=False).eval( + coord, None, atype + ) + np.testing.assert_allclose( + converted_result[0], source_result[0], rtol=1e-10, atol=1e-10 + ) + np.testing.assert_allclose( + converted_result[1], source_result[1], rtol=1e-10, atol=1e-10 + ) + + +def test_deserialize_rejects_unknown_lower_kind(dpa1_dpmodel_data, tmp_path) -> None: + """The target serializer owns validation of its supported lower ABIs.""" + with pytest.raises(ValueError, match="Unsupported lower_kind 'unknown'"): + deserialize_to_file( + str(tmp_path / "model.pt2"), + copy.deepcopy(dpa1_dpmodel_data), + lower_kind="unknown", + ) + + def test_graph_pt2_has_lower_input_kind_graph(dpa1_dpmodel_data) -> None: """``lower_kind="graph"`` -> metadata ``lower_input_kind == "graph"``.""" with tempfile.TemporaryDirectory() as d: diff --git a/source/tests/test_convert_backend.py b/source/tests/test_convert_backend.py index 063caec868..873daabe7c 100644 --- a/source/tests/test_convert_backend.py +++ b/source/tests/test_convert_backend.py @@ -9,7 +9,7 @@ ) -def test_convert_backend_automatically_selects_lower_kind( +def test_convert_backend_uses_auto_for_unannotated_source( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, object] = {} @@ -48,3 +48,102 @@ def detect_backend(path: str) -> type[InputBackend] | type[OutputBackend]: assert captured["lower_kind"] == "auto" assert captured["do_atomic_virial"] is False + + +@pytest.mark.parametrize( + "lower_input_kind", + ["nlist", "graph", "dpa1_canonical", "dpa4c_canonical", "edge_vec"], +) +def test_convert_backend_preserves_explicit_lower_kind( + monkeypatch: pytest.MonkeyPatch, + lower_input_kind: str, +) -> None: + captured: dict[str, object] = {} + + class InputBackend: + name = "input" + + @staticmethod + def serialize_hook(path: str) -> dict[str, str]: + return {"path": path, "lower_input_kind": lower_input_kind} + + class OutputBackend: + name = "output" + + @staticmethod + def deserialize_hook( + path: str, + data: dict[str, str], + *, + lower_kind: str = "nlist", + ) -> None: + captured.update(path=path, data=data, lower_kind=lower_kind) + + def detect_backend(path: str) -> type[InputBackend] | type[OutputBackend]: + return InputBackend if path.endswith(".input") else OutputBackend + + monkeypatch.setattr(Backend, "detect_backend_by_model", detect_backend) + + convert_backend(INPUT="model.input", OUTPUT="model.output") + + assert captured["lower_kind"] == lower_input_kind + + +@pytest.mark.parametrize("lower_input_kind", [None, "nlist"]) +def test_convert_backend_allows_dense_compatible_source_for_plain_output( + monkeypatch: pytest.MonkeyPatch, + lower_input_kind: str | None, +) -> None: + captured: dict[str, object] = {} + + class InputBackend: + name = "input" + + @staticmethod + def serialize_hook(path: str) -> dict[str, str]: + data = {"path": path} + if lower_input_kind is not None: + data["lower_input_kind"] = lower_input_kind + return data + + class OutputBackend: + name = "output" + + @staticmethod + def deserialize_hook(path: str, data: dict[str, str]) -> None: + captured.update(path=path, data=data) + + def detect_backend(path: str) -> type[InputBackend] | type[OutputBackend]: + return InputBackend if path.endswith(".input") else OutputBackend + + monkeypatch.setattr(Backend, "detect_backend_by_model", detect_backend) + + convert_backend(INPUT="model.input", OUTPUT="model.output") + + assert captured["path"] == "model.output" + + +def test_convert_backend_rejects_graph_for_dense_only_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class InputBackend: + name = "input" + + @staticmethod + def serialize_hook(path: str) -> dict[str, str]: + return {"path": path, "lower_input_kind": "graph"} + + class OutputBackend: + name = "output" + + @staticmethod + def deserialize_hook(path: str, data: dict[str, str]) -> None: + raise AssertionError("dense-only output hook must not be called") + + def detect_backend(path: str) -> type[InputBackend] | type[OutputBackend]: + return InputBackend if path.endswith(".input") else OutputBackend + + monkeypatch.setattr(Backend, "detect_backend_by_model", detect_backend) + + with pytest.raises(ValueError, match="Cannot preserve lower_input_kind 'graph'"): + convert_backend(INPUT="model.input", OUTPUT="model.output") diff --git a/source/tests/tf2/test_serialization.py b/source/tests/tf2/test_serialization.py new file mode 100644 index 0000000000..a8ac2af15b --- /dev/null +++ b/source/tests/tf2/test_serialization.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""TensorFlow 2 backend serialization contracts.""" + +import os + +import pytest + +if os.environ.get("DP_TEST_TF2_ONLY") != "1": + pytest.skip( + "TF2 tests require DP_TEST_TF2_ONLY=1", + allow_module_level=True, + ) + +from deepmd.tf2.utils import ( + serialization, +) + + +def test_checkpoint_serialization_declares_dense_lower( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A TF2 training checkpoint exposes its dense source semantics.""" + state = { + "backend": "TensorFlow2", + "model_def_script": {}, + "current_step": 0, + } + monkeypatch.setattr( + serialization, + "_load_checkpoint_state", + lambda _path: ("checkpoint", state), + ) + monkeypatch.setattr( + serialization, + "_restore_models_from_checkpoint", + lambda _checkpoint, _script, _state: {}, + ) + monkeypatch.setattr( + serialization, + "_serialize_models", + lambda _models, _script: {}, + ) + + data = serialization.serialize_from_file("model.tf2") + + assert data["lower_input_kind"] == "nlist"