Add `__version__` attribute to Python module (#675)

Add __version__ attribute to Python module
Signed-off-by: Tim Moon <tmoon@nvidia.com>

VERSION

-1.5.0dev
+1.5.0.dev0

setup.py

@@ -2,6 +2,8 @@
 #
 # See LICENSE for license information.
 
+"""Installation script."""
+
 import ctypes
 from functools import lru_cache
 import os
@@ -18,35 +20,11 @@ from typing import List, Optional, Tuple, Union
 import setuptools
 from setuptools.command.build_ext import build_ext
 
+from te_version import te_version
+
 # Project directory root
 root_path: Path = Path(__file__).resolve().parent
 
-@lru_cache(maxsize=1)
-def te_version() -> str:
-    """Transformer Engine version string
-
-    Includes Git commit as local version, unless suppressed with
-    NVTE_NO_LOCAL_VERSION environment variable.
-
-    """
-    with open(root_path / "VERSION", "r") as f:
-        version = f.readline().strip()
-    if not int(os.getenv("NVTE_NO_LOCAL_VERSION", "0")):
-        try:
-            output = subprocess.run(
-                ["git", "rev-parse" , "--short", "HEAD"],
-                capture_output=True,
-                cwd=root_path,
-                check=True,
-                universal_newlines=True,
-            )
-        except (CalledProcessError, OSError):
-            pass
-        else:
-            commit = output.stdout.strip()
-            version += f"+{commit}"
-    return version
-
 @lru_cache(maxsize=1)
 def with_debug_build() -> bool:
     """Whether to build with a debug configuration"""
@@ -266,7 +244,10 @@ def setup_requirements() -> Tuple[List[str], List[str], List[str]]:
 
     # Common requirements
     setup_reqs: List[str] = []
-    install_reqs: List[str] = ["pydantic"]
+    install_reqs: List[str] = [
+        "pydantic",
+        "importlib-metadata>=1.0; python_version<'3.8'",
+    ]
     test_reqs: List[str] = ["pytest"]
 
     def add_unique(l: List[str], vals: Union[str, List[str]]) -> None:

te_version.py

+# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Transformer Engine version string."""
+import os
+from pathlib import Path
+import subprocess
+
+def te_version() -> str:
+    """Transformer Engine version string
+
+    Includes Git commit as local version, unless suppressed with
+    NVTE_NO_LOCAL_VERSION environment variable.
+
+    """
+    root_path = Path(__file__).resolve().parent
+    with open(root_path / "VERSION", "r") as f:
+        version = f.readline().strip()
+    if not int(os.getenv("NVTE_NO_LOCAL_VERSION", "0")):
+        try:
+            output = subprocess.run(
+                ["git", "rev-parse" , "--short", "HEAD"],
+                capture_output=True,
+                cwd=root_path,
+                check=True,
+                universal_newlines=True,
+            )
+        except (subprocess.CalledProcessError, OSError):
+            pass
+        else:
+            commit = output.stdout.strip()
+            version += f"+{commit}"
+    return version

transformer_engine/__init__.py

@@ -3,6 +3,7 @@
 # See LICENSE for license information.
 
 """Top level package"""
+from ._version import __version__
 from . import common
 
 try:

transformer_engine/_version.py

+# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Version information"""
+import sys
+from packaging.version import Version
+
+if sys.version_info >= (3, 8):
+    from importlib import metadata
+else:
+    import importlib_metadata as metadata
+
+def _version_str() -> str:
+    """Transformer Engine version string"""
+
+    # Try getting version from package metadata
+    version_str = None
+    try:
+        version_str = metadata.version("transformer_engine")
+    except:
+        pass
+    if version_str:
+        return version_str
+
+    # Try getting version from Git root directory
+    try:
+        from te_version import te_version
+        version_str = te_version()
+    except:
+        pass
+    if version_str:
+        return version_str
+
+    # Could not deduce version
+    return "0.dev0+unknown"
+
+# Transformer Engine version
+__version__: Version = Version(_version_str())