fix(pt_expt): store .pt2 metadata under model/extra/ to match PyTorch 2.11 layout

``aoti_compile_and_package`` writes every entry of the compiled ``.pt2``
archive under a top-level ``model/`` directory; deepmd-kit then appended
its own metadata JSON blobs (``metadata.json``, ``model.json``,
``model_def_script.json``) at the root-level ``extra/`` path via
``zipfile``.  Starting with PyTorch 2.11, the strict single-model
loader ``torch.export.pt2_archive._package.load_pt2`` refuses archives
that carry files outside ``model/``:
    RuntimeError: [enforce fail at inline_container.cc:340] .
    file in archive is not in a subdirectory model/: extra/metadata.json
``torch._inductor.package.package.load_package`` catches this error and
falls back to the legacy C++ loader, but prints the misleading warning
``Loading outdated pt2 file. Please regenerate your package.`` every
time the archive is opened -- even though the archive version itself
(``archive_version == '0'``) is already current.
Move the deepmd-kit metadata blobs into ``model/extra/`` so that the
fast path through ``load_pt2`` accepts the archive cleanly and the
misleading warning disappears.  A module-level constant
``PT2_EXTRA_PREFIX`` in ``deepmd.pt_expt.utils.serialization`` is the
single source of truth for the prefix; both the writer
(``_deserialize_to_file_pt2``) and the readers
(``_serialize_from_file_pt2``, ``DeepEval._load_pt2``) derive their
entry names from it.
The C++ reader in ``source/api_cc/src/commonPTExpt.h::read_zip_entry``
needs no changes: it already matches ``entry_name`` as a
``/``-delimited suffix, so ``"extra/metadata.json"`` resolves against
both the old root-level and the new ``model/extra/`` location
transparently.  The ``test_pt2_has_metadata`` assertion in
``source/tests/pt_expt/infer/test_deep_eval.py`` is updated to expect
the new paths.

Author: OutisLi

deepmd/pt_expt/infer/deep_eval.py

@@ -166,9 +166,9 @@ def _init_from_model_json(self, model_json_str: str) -> None:
             self._model_output_def = ModelOutputDef(self._dpmodel.atomic_output_def())
 
     def _init_from_metadata(self) -> None:
-        """Initialize DeepEval from ``extra/metadata.json`` alone.
+        """Initialize DeepEval from ``metadata.json`` alone.
 
-        Used when the ``.pt2`` / ``.pte`` archive ships no ``extra/model.json``
+        Used when the ``.pt2`` / ``.pte`` archive ships no ``model.json``
         (e.g. for backends that do not travel through the dpmodel round-trip).
         The metadata contract is the same one the C++ ``DeepPotPTExpt``
         reader consumes, so anything that validates against the C++ side
@@ -241,33 +241,44 @@ def _load_pte(self, model_file: str) -> None:
     def _load_pt2(self, model_file: str) -> None:
         """Load a .pt2 (AOTInductor) model file.
 
-        ``extra/model.json`` is optional — it only enables the dpmodel
+        ``model.json`` is optional — it only enables the dpmodel
         round-trip (used by ``eval_descriptor``, ``eval_typeebd``, etc.).
         Pure AOTI inference (``DeepPot.eval`` / ``dp test`` / ASE
-        calculator) only needs ``extra/metadata.json``, matching the
-        contract the C++ ``DeepPotPTExpt`` reader enforces.  Backends that
-        cannot produce ``model.json``.
+        calculator) only needs ``metadata.json``, matching the contract
+        the C++ ``DeepPotPTExpt`` reader enforces.
+
+        Archive entries are located under ``model/extra/`` so that the
+        PyTorch 2.11 ``load_pt2`` loader accepts the archive without the
+        "outdated pt2 file" fallback warning.
         """
         import zipfile
 
         from torch._inductor import (
             aoti_load_package,
         )
 
+        from deepmd.pt_expt.utils.serialization import (
+            PT2_EXTRA_PREFIX,
+        )
+
+        md_entry = PT2_EXTRA_PREFIX + "metadata.json"
+        model_json_entry = PT2_EXTRA_PREFIX + "model.json"
+        mds_entry = PT2_EXTRA_PREFIX + "model_def_script.json"
+
         # Read metadata from the .pt2 ZIP archive
         with zipfile.ZipFile(model_file, "r") as zf:
             names = zf.namelist()
-            if "extra/metadata.json" not in names:
+            if md_entry not in names:
                 raise ValueError(
-                    f"Invalid .pt2 file '{model_file}': missing 'extra/metadata.json'"
+                    f"Invalid .pt2 file '{model_file}': missing '{md_entry}'"
                 )
-            md = zf.read("extra/metadata.json").decode("utf-8")
+            md = zf.read(md_entry).decode("utf-8")
             model_json_str = ""
-            if "extra/model.json" in names:
-                model_json_str = zf.read("extra/model.json").decode("utf-8")
+            if model_json_entry in names:
+                model_json_str = zf.read(model_json_entry).decode("utf-8")
             mds = ""
-            if "extra/model_def_script.json" in names:
-                mds = zf.read("extra/model_def_script.json").decode("utf-8")
+            if mds_entry in names:
+                mds = zf.read(mds_entry).decode("utf-8")
 
         self.metadata = json.loads(md)
         self._model_def_script = json.loads(mds) if mds else {}
@@ -1118,15 +1129,15 @@ def _require_dpmodel(self, feature: str) -> None:
 
         ``eval_descriptor`` / ``eval_typeebd`` / ``eval_fitting_last_layer``
         all introspect the dpmodel's internal sub-modules, which requires
-        ``extra/model.json`` to have been present at load time.  Archives
+        ``model.json`` to have been present at load time.  Archives
         shipped without ``model.json`` (metadata-only mode) can still run
         the main ``eval`` inference path but cannot expose these hooks.
         """
         if self._dpmodel is None:
             raise NotImplementedError(
                 f"{feature} requires the dpmodel instance, which is only "
                 "available when the .pt2 / .pte archive contains "
-                "'extra/model.json'. The loaded archive is metadata-only; "
+                "'model.json'. The loaded archive is metadata-only; "
                 "re-export with the full dpmodel serialisation to enable "
                 "this feature."
             )

deepmd/pt_expt/utils/serialization.py

@@ -15,6 +15,25 @@
     traverse_model_dict,
 )
 
+# ---------------------------------------------------------------------------
+# AOTInductor ``.pt2`` archive layout.
+#
+# PyTorch 2.11 tightened the single-model ``.pt2`` convention so that every
+# entry in the ZIP archive must live under the top-level ``model/`` directory.
+# Any stray root-level file makes
+# ``torch.export.pt2_archive._package.load_pt2`` raise ``RuntimeError`` at
+# load time; the upper-level ``torch._inductor.package.package.load_package``
+# then emits a misleading ``Loading outdated pt2 file. Please regenerate
+# your package.`` warning and falls back to the legacy C++ loader.
+#
+# deepmd-kit therefore stores its metadata JSON blobs under ``model/extra/``
+# so that the strict ``load_pt2`` loader accepts the archive without
+# complaint.  The C++ reader (``commonPTExpt.h::read_zip_entry``) resolves
+# this layout transparently because it matches ``entry_name`` as a
+# ``/``-delimited suffix.
+# ---------------------------------------------------------------------------
+PT2_EXTRA_PREFIX = "model/extra/"
+
 
 def _strip_shape_assertions(graph_module: torch.nn.Module) -> None:
     """Remove shape-guard assertion nodes from a spin model's exported graph.
@@ -316,21 +335,23 @@ def _serialize_from_file_pte(model_file: str) -> dict:
 def _serialize_from_file_pt2(model_file: str) -> dict:
     """Serialize a .pt2 model file to a dictionary.
 
-    Reads the model dict stored in the extra/ directory of the .pt2 ZIP archive.
+    Reads the model dict stored in the ``model/extra/`` directory of the
+    ``.pt2`` ZIP archive.
     """
     import zipfile
 
+    model_json_entry = PT2_EXTRA_PREFIX + "model.json"
+    model_def_script_entry = PT2_EXTRA_PREFIX + "model_def_script.json"
     with zipfile.ZipFile(model_file, "r") as zf:
-        if "extra/model.json" not in zf.namelist():
+        names = zf.namelist()
+        if model_json_entry not in names:
             raise ValueError(
-                f"Invalid .pt2 file '{model_file}': missing 'extra/model.json'"
+                f"Invalid .pt2 file '{model_file}': missing '{model_json_entry}'"
             )
-        model_json = zf.read("extra/model.json").decode("utf-8")
+        model_json = zf.read(model_json_entry).decode("utf-8")
         model_def_script_json = ""
-        if "extra/model_def_script.json" in zf.namelist():
-            model_def_script_json = zf.read("extra/model_def_script.json").decode(
-                "utf-8"
-            )
+        if model_def_script_entry in names:
+            model_def_script_json = zf.read(model_def_script_entry).decode("utf-8")
     model_dict = json.loads(model_json)
     model_dict = _json_to_numpy(model_dict)
     if model_def_script_json:
@@ -563,13 +584,20 @@ def _deserialize_to_file_pt2(
     # Compile via AOTInductor into a .pt2 package
     aoti_compile_and_package(exported, package_path=model_file)
 
-    # Embed metadata into the .pt2 ZIP archive
+    # Embed metadata into the .pt2 ZIP archive.  Entries are placed under
+    # ``model/extra/`` so the strict PyTorch 2.11 ``load_pt2`` loader
+    # accepts the archive without emitting the "outdated pt2 file"
+    # fallback warning.  See the module-level comment on
+    # ``PT2_EXTRA_PREFIX`` for the rationale.
     model_def_script = data.get("model_def_script") or {}
     metadata["output_keys"] = output_keys
     with zipfile.ZipFile(model_file, "a") as zf:
-        zf.writestr("extra/metadata.json", json.dumps(metadata))
-        zf.writestr("extra/model_def_script.json", json.dumps(model_def_script))
+        zf.writestr(PT2_EXTRA_PREFIX + "metadata.json", json.dumps(metadata))
+        zf.writestr(
+            PT2_EXTRA_PREFIX + "model_def_script.json",
+            json.dumps(model_def_script),
+        )
         zf.writestr(
-            "extra/model.json",
+            PT2_EXTRA_PREFIX + "model.json",
             json.dumps(data_for_json, separators=(",", ":")),
         )

source/tests/pt_expt/infer/test_deep_eval.py

@@ -638,14 +638,14 @@ def test_pt2_file_is_zip(self) -> None:
         self.assertTrue(zipfile.is_zipfile(self.tmpfile.name))
 
     def test_pt2_has_metadata(self) -> None:
-        """The .pt2 ZIP should contain metadata entries."""
+        """The .pt2 ZIP should contain metadata entries under ``model/extra/``."""
         with zipfile.ZipFile(self.tmpfile.name, "r") as zf:
             names = zf.namelist()
-            self.assertIn("extra/metadata.json", names)
-            self.assertIn("extra/model_def_script.json", names)
-            self.assertIn("extra/model.json", names)
-            self.assertNotIn("extra/output_keys.json", names)
-            self.assertNotIn("extra/model_params.json", names)
+            self.assertIn("model/extra/metadata.json", names)
+            self.assertIn("model/extra/model_def_script.json", names)
+            self.assertIn("model/extra/model.json", names)
+            self.assertNotIn("model/extra/output_keys.json", names)
+            self.assertNotIn("model/extra/model_params.json", names)
 
     def test_eval_consistency(self) -> None:
         """Test that DeepPot.eval gives same results as direct model forward."""