Gracefully handle optree less than minimum version (#150956)

zou3519 · pytorchmergebot · commit 061832bc7a67 · 2025-04-10T01:22:50.000Z
Summary: - We are saying the minimum version of pytree that PyTorch can use is 0.13.0 - If a user imports torch.utils._cxx_pytree, it will raise an ImportError if optree doesn't exist or exists and is less than the minimum version. Fixes #150889. There are actually two parts to that issue: 1. dtensor imports torch.utils._cxx_pytree, but the optree installed in the environment might be too old. Instead, raising ImportError in torch.utils._cxx_pytree solves the issue. 2. We emit an "optree too low version" warning. I've deleted the warning in favor of the more explicit ImportError. Test Plan: - code reading Pull Request resolved: #150956 Approved by: https://github.com/albanD, https://github.com/atalman, https://github.com/XuehaiPan
diff --git a/torch/utils/_cxx_pytree.py b/torch/utils/_cxx_pytree.py
@@ -20,6 +20,19 @@
 from typing_extensions import deprecated, TypeIs
 
 import optree
+
+from torch._vendor.packaging.version import Version
+
+
+if Version(optree.__version__) < Version("0.13.0"):  # type: ignore[attr-defined]
+    raise ImportError(
+        "torch.utils._cxx_pytree depends on optree, which is an optional dependency "
+        "of PyTorch. To use it, please upgrade your optree package to >= 0.13.0"
+    )
+
+del Version
+
+
 from optree import PyTreeSpec as TreeSpec  # direct import for type annotations
 
 import torch.utils._pytree as python_pytree
diff --git a/torch/utils/_pytree.py b/torch/utils/_pytree.py
@@ -173,21 +173,12 @@ class _SerializeNodeDef(NamedTuple):
 try:
     _optree_version = importlib.metadata.version("optree")
 except importlib.metadata.PackageNotFoundError:
+    # optree was not imported
     _cxx_pytree_dynamo_traceable = _cxx_pytree_exists = False
 else:
+    # optree was imported
     _cxx_pytree_exists = True
-    from torch._vendor.packaging.version import Version
-
-    _cxx_pytree_dynamo_traceable = Version(_optree_version) >= Version("0.13.0")
-    if not _cxx_pytree_dynamo_traceable:
-        warnings.warn(
-            "optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. "
-            "C++ pytree support is disabled. "
-            "Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.",
-            FutureWarning,
-        )
-
-    del Version
+    _cxx_pytree_dynamo_traceable = True
 
 _cxx_pytree_imported = False
 _cxx_pytree_pending_imports: list[Any] = []
