Enable protov2 by default on UCX >= 1.18.0 (#384)

With protov1 not being actively developed anymore and protov2 now being more stable, it's time to switch permanently. See rapidsai/ucx-py#1121 for more details.

Closes rapidsai/ucx-py#1121 .

Authors:
  - Peter Andreas Entschev (https://github.com/pentschev)

Approvers:
  - Mads R. B. Kristensen (https://github.com/madsbk)

URL: #384

Author: pentschev

python/ucxx/ucxx/__init__.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: BSD-3-Clause
 
 """UCXX: Python bindings for the Unified Communication X library (UCX <www.openucx.org>)"""
@@ -37,6 +37,18 @@
 except ImportError:
     pynvml = None
 
+_ucx_version = get_ucx_version()
+__ucx_min_version__ = "1.15.0"
+__ucx_version__ = "%d.%d.%d" % _ucx_version
+
+if _ucx_version < tuple(int(i) for i in __ucx_min_version__.split(".")):
+    raise ImportError(
+        f"Support for UCX {__ucx_version__} has ended. Please upgrade to "
+        f"{__ucx_min_version__} or newer. If you believe the wrong version "
+        "is being loaded, please check the path from where UCX is loaded "
+        "by rerunning with the environment variable `UCX_LOG_LEVEL=debug`."
+    )
+
 # Setup UCX-Py logger
 logger = get_ucxpy_logger()
 
@@ -51,7 +63,7 @@
 if (
     pynvml is not None
     and "UCX_CUDA_COPY_MAX_REG_RATIO" not in os.environ
-    and get_ucx_version() >= (1, 12, 0)
+    and _ucx_version >= (1, 12, 0)
 ):
     try:
         pynvml.nvmlInit()
@@ -91,25 +103,14 @@ def _is_mig_device(handle):
     ):
         pass
 
-if "UCX_MAX_RNDV_RAILS" not in os.environ and get_ucx_version() >= (1, 12, 0):
+if "UCX_MAX_RNDV_RAILS" not in os.environ and _ucx_version >= (1, 12, 0):
     logger.info("Setting UCX_MAX_RNDV_RAILS=1")
     os.environ["UCX_MAX_RNDV_RAILS"] = "1"
 
-if "UCX_PROTO_ENABLE" not in os.environ:
+if "UCX_PROTO_ENABLE" not in os.environ and (1, 12, 0) <= _ucx_version < (1, 18, 0):
     # UCX protov2 still doesn't support CUDA async/managed memory
     logger.info("Setting UCX_PROTO_ENABLE=n")
     os.environ["UCX_PROTO_ENABLE"] = "n"
 
 
 from ._version import __git_commit__, __version__
-
-__ucx_min_version__ = "1.15.0"
-__ucx_version__ = "%d.%d.%d" % get_ucx_version()
-
-if get_ucx_version() < tuple(int(i) for i in __ucx_min_version__.split(".")):
-    raise ImportError(
-        f"Support for UCX {__ucx_version__} has ended. Please upgrade to "
-        f"{__ucx_min_version__} or newer. If you believe the wrong version "
-        "is being loaded, please check the path from where UCX is loaded "
-        "by rerunning with the environment variable `UCX_LOG_LEVEL=debug`."
-    )