[UCC][TORCH_UCC]Do integer driver version comparison for UCC (#1411)

eqy · web-flow · commit bf3c008e90d0 · 2022-07-05T13:32:03.000-07:00
* Integer driver number comparison

* packaging
diff --git a/apex/transformer/testing/distributed_test_base.py b/apex/transformer/testing/distributed_test_base.py
@@ -1,6 +1,7 @@
 import os
 import sys
 import unittest
+from packaging.version import Version, parse
 
 import torch
 from torch import distributed as dist
@@ -16,10 +17,10 @@
     HAS_TORCH_UCC = False
 
 # NOTE(mkozuki): Version guard for ucc. ref: https://github.com/openucx/ucc/issues/496
-_TORCH_UCC_COMPAT_NVIDIA_DRIVER_VERSION = "470.42.01"
+_TORCH_UCC_COMPAT_NVIDIA_DRIVER_VERSION = Version("470.42.01")
 _driver_version = None
 if torch.cuda.is_available():
-    _driver_version = collect_env.get_nvidia_driver_version(collect_env.run)
+    _driver_version = parse(collect_env.get_nvidia_driver_version(collect_env.run))
 HAS_TORCH_UCC_COMPAT_NVIDIA_DRIVER = _driver_version is not None and _driver_version >= _TORCH_UCC_COMPAT_NVIDIA_DRIVER_VERSION
 
 
diff --git a/requirements.txt b/requirements.txt
@@ -3,3 +3,4 @@ tqdm>=4.28.1
 numpy>=1.15.3
 PyYAML>=5.1
 pytest>=3.5.1
+packaging>=14.0
