Gracefully handle optree less than minimum version (pytorch#150977)

Gracefully handle optree less than minimum version (pytorch#150956)

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 pytorch#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: pytorch#150956
Approved by: https://github.com/albanD, https://github.com/atalman, https://github.com/XuehaiPan

(cherry picked from commit 061832b)

Co-authored-by: rzou <[email protected]>

Author: pytorchbot

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

torch/utils/_pytree.py

@@ -163,21 +163,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] = []