fix(ci): keep the linear wire type shared, count tf32 warns at the source

Two CI failures on deepmodeling#5884:

- `test_linear_ener.py::test_pt_expt_consistent_with_ref` (all 4 cases, CPU
  + CUDA): the flat linear dict emitted `"type": "linear_ener"` while
  pt/tf write `"linear"`, so the cross-backend serialization comparison
  diverged. One wire format for all backends wins: emit `"linear"` again
  and keep `"linear_ener"` as an accepted atomic alias. The energy-specific
  name stays the config/model type, and it is now also registered as a
  MODEL alias in dpmodel/pt_expt so a serialized composition round-trips
  through `BaseModel.deserialize`.

- `test_get_model_dpa4.py::test_enable_tf32_warns_once[True]` (CPU +
  CUDA): counted records through caplog, which reads a ROOT handler, so
  the count moved with whatever global logging state earlier tests left
  behind (0 with `set_log_handles`' propagate=False, 2 in the CI shard).
  Count on the EMITTING logger with a dedicated handler instead -- exactly
  one record per `log.warning`, independent of propagation and of handlers
  attached elsewhere.

Also register test_zbl_bridging.py in `_AOTI_COMPILE_MODULES`: it freezes a
.pt2 and the drift guard flagged it, so its compile would otherwise run in
the CUDA lane.

Author: Han Wang

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -39,8 +39,8 @@
 )
 
 
-@BaseAtomicModel.register("linear_ener")
-@BaseAtomicModel.register("linear")  # legacy wire alias
+@BaseAtomicModel.register("linear")
+@BaseAtomicModel.register("linear_ener")  # accepted alias, never emitted
 class LinearEnergyAtomicModel(BaseAtomicModel):
     r"""Linear model makes linear combinations of several existing models.
 
@@ -472,9 +472,11 @@ def serialize(self) -> dict:
             {
                 "@class": "Model",
                 "@version": 3,
-                # energy-specific wire type: future linear dipole/polar models get
-                # their own, so the flat model dict dispatches unambiguously
-                "type": "linear_ener",
+                # ONE wire type across backends: pt/tf write "linear" here, so
+                # dpmodel must too or cross-backend conversion breaks.  The
+                # unambiguous energy-specific name lives in the config/model
+                # registry ("linear_ener"), which is also accepted here.
+                "type": "linear",
                 "models": [model.serialize() for model in self.models],
                 "type_map": self.type_map,
                 "weights": self.weights,

deepmd/dpmodel/model/dp_linear_model.py

@@ -26,7 +26,8 @@
 DPLinearModel_ = make_model(LinearEnergyAtomicModel, T_Bases=(NativeOP, BaseModel))
 
 
-@BaseModel.register("linear_ener")
+@BaseModel.register("linear_ener")  # config type
+@BaseModel.register("linear")  # wire type emitted by the flat serialize
 class LinearEnergyModel(DPModelCommon, DPLinearModel_):
     r"""Energy model over a linear combination of atomic models.
 

deepmd/pt_expt/model/dp_linear_model.py

@@ -33,7 +33,8 @@
 DPLinearModel_ = make_model(LinearEnergyAtomicModel, T_Bases=(BaseModel,))
 
 
-@BaseModel.register("linear_ener")
+@BaseModel.register("linear_ener")  # config type
+@BaseModel.register("linear")  # wire type emitted by the flat serialize
 class LinearEnergyModel(DPModelCommon, DPLinearModel_):
     # The graph .pt2 exportable is energy-contract machinery (public-key
     # translation over the CM's forward_common_lower_graph_exportable),

source/tests/common/dpmodel/test_zbl_bridging.py

@@ -132,7 +132,9 @@ def test_zbl_serialize_roundtrip_energy_identical():
     model = get_model(copy.deepcopy(ZBL_CONFIG))
     coord, atype, box = _close_pair_inputs()
     data = model.serialize()
-    assert data["type"] == "linear_ener"
+    # the flat wire type is "linear" -- the SAME string pt/tf write, so a
+    # composition round-trips across backends
+    assert data["type"] == "linear"
     m2 = BaseModel.deserialize(data)
     assert type(m2) is LinearEnergyModel
     e1 = model.call_common(coord, atype, box=box, neighbor_graph_method="dense")[

source/tests/conftest.py

@@ -36,6 +36,7 @@
         "pt_expt/model/test_export_with_comm.py",
         "pt_expt/model/test_dpa1_graph_lower.py",
         "pt_expt/model/test_graph_export.py",
+        "pt_expt/model/test_zbl_bridging.py",
         "pt_expt/model/test_graph_export_with_comm.py",
         "pt_expt/utils/test_graph_pt2_metadata.py",
         "pt_expt/infer/test_deep_eval_metadata_only.py",

source/tests/pt_expt/model/test_get_model_dpa4.py

@@ -274,7 +274,7 @@ def test_default_unsupported_values_pass(self) -> None:
 # `enable_tf32` toggles TF32 matmul precision in pt but is ignored by pt_expt
 # (always "highest" precision); a truthy value must emit a warn-once message.
 @pytest.mark.parametrize("enable_tf32", [True, False])  # truthy warns, falsy silent
-def test_enable_tf32_warns_once(enable_tf32, caplog, monkeypatch) -> None:
+def test_enable_tf32_warns_once(enable_tf32, monkeypatch) -> None:
     import importlib
 
     # the package __init__ rebinds the name ``get_model`` to the function, so
@@ -285,27 +285,37 @@ def test_enable_tf32_warns_once(enable_tf32, caplog, monkeypatch) -> None:
     # test ordering (other get_sezm_model calls may have already warned)
     monkeypatch.setattr(gm_mod, "_WARNED_ONCE", set())
 
-    # caplog captures via a ROOT-logger handler and relies on propagation,
-    # but any earlier test that ran main() (e.g. test_dp_freeze's dispatcher
-    # tests) leaves set_log_handles' global ``deepmd``-logger
-    # propagate=False behind (deepmd/loggers/loggers.py), silently emptying
-    # caplog.records. monkeypatch restores the attribute afterwards.
-    monkeypatch.setattr(logging.getLogger("deepmd"), "propagate", True)
-
-    raw = _make_raw_model_config(enable_tf32=enable_tf32)
-
-    with caplog.at_level(logging.WARNING, logger=gm_mod.log.name):
-        gm_mod.get_sezm_model(raw)
-    matches = [r for r in caplog.records if "enable_tf32" in r.getMessage()]
-    if enable_tf32:
-        assert len(matches) == 1, caplog.text
-        # a second call must NOT warn again (warn-once per process)
-        caplog.clear()
-        with caplog.at_level(logging.WARNING, logger=gm_mod.log.name):
+    # Count emissions on the EMITTING logger with our own handler rather than
+    # through caplog: caplog reads a root handler, so whatever global logging
+    # state earlier tests left behind (set_log_handles flips the ``deepmd``
+    # logger's propagate off and installs its own handlers) changes how many
+    # records reach it -- zero when propagation is off, more than one when the
+    # record is seen through several attached handlers.  A handler on the
+    # emitting logger sees exactly one record per ``log.warning`` call.
+    records: list[logging.LogRecord] = []
+
+    class _Collect(logging.Handler):
+        def emit(self, record: logging.LogRecord) -> None:
+            records.append(record)
+
+    handler = _Collect(level=logging.WARNING)
+    old_level = gm_mod.log.level
+    gm_mod.log.setLevel(logging.WARNING)
+    gm_mod.log.addHandler(handler)
+    try:
+        gm_mod.get_sezm_model(_make_raw_model_config(enable_tf32=enable_tf32))
+        matches = [r for r in records if "enable_tf32" in r.getMessage()]
+        if enable_tf32:
+            assert len(matches) == 1, [r.getMessage() for r in records]
+            # a second call must NOT warn again (warn-once per process)
+            records.clear()
             gm_mod.get_sezm_model(_make_raw_model_config(enable_tf32=enable_tf32))
-        assert not [r for r in caplog.records if "enable_tf32" in r.getMessage()]
-    else:
-        assert not matches, caplog.text
+            assert not [r for r in records if "enable_tf32" in r.getMessage()]
+        else:
+            assert not matches, [r.getMessage() for r in records]
+    finally:
+        gm_mod.log.removeHandler(handler)
+        gm_mod.log.setLevel(old_level)
 
 
 class TestNativeSpinErrorTranslation(unittest.TestCase):

source/tests/pt_expt/model/test_zbl_bridging.py

@@ -209,7 +209,9 @@ def test_serialize_roundtrip(self) -> None:
         )
 
         data = self.pt_expt_model.serialize()
-        assert data["type"] == "linear_ener"
+        # the flat wire type is "linear" -- the SAME string pt/tf write, so a
+        # composition round-trips across backends
+        assert data["type"] == "linear"
         m2 = BaseModel.deserialize(data).to(torch.device("cpu")).eval()
         assert type(m2) is LinearEnergyModel
         out = self.pt_expt_model.forward(self.coord, self.atype, box=self.box)