Harden lazy border_op registration and comm-artifact gating in pt_expt export

[Copilot]

This addresses review feedback on PR #5451 around pt_expt lazy comm-op initialization. The fixes make registration atomic/idempotent, prevent premature op-library loading for non-comm models, and align non-serialization comm call paths with explicit initialization.

• Thread-safe, idempotent comm registration

  • Added a module-level lock in deepmd/pt_expt/utils/comm.py.
  • ensure_comm_registered() now double-checks _registered inside the lock.
  • Duplicate-registration tolerance now also covers torch.library.register_autograd(...) runtime “already registered/has” cases.

• Correct ordering in with-comm export path

  • In deepmd/pt_expt/utils/serialization.py, _needs_with_comm_artifact(model) is evaluated before importing/calling ensure_comm_registered().
  • Non-comm models now fail fast with the existing ValueError instead of triggering unrelated op-load failures.

• Explicit initialization for non-serialization comm paths

  • Added explicit ensure_comm_registered() calls at descriptor comm_dict call sites:
    • deepmd/pt_expt/descriptor/repflows.py
    • deepmd/pt_expt/descriptor/repformers.py
  • Updated deepmd/pt_expt/utils/__init__.py and pt_expt test comments to reflect required lazy-init behavior.

• Focused regression coverage

  • Added tests for registration idempotency/duplicate autograd handling and for export-path gating-before-registration behavior:
    • source/tests/pt_expt/utils/test_comm_registration.py
    • source/tests/pt_expt/utils/test_serialization_with_comm_guard.py

# comm.py
if _registered:
    return
with _register_lock:
    if _registered:
        return
    _check_underlying_ops_loaded()
    ...
    try:
        torch.library.register_autograd(...)
    except RuntimeError as e:
        if "already has" not in str(e) and "already registered" not in str(e):
            raise
    _registered = True

Warning

Firewall rules blocked me from connecting to one or more addresses (expand for details)

I tried to connect to the following addresses, but was blocked by firewall rules:

• download-r2.pytorch.org
  • Triggering command: /opt/hostedtoolcache/uv/0.11.16/x86_64/uv uv pip install torch --index-url REDACTED (dns block)

If you need me to access, download, or install something from one of these locations, you can either:

• Configure Actions setup steps to set up my environment, which run before the firewall is enabled
• Add the appropriate URLs or hosts to the custom allowlist in this repository's Copilot coding agent settings (admins only)