TransformerEngine

1
# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
#
3
# See LICENSE for license information.
4

5
"""Transformer Engine version string."""
6
import os
7
from pathlib import Path
8
import subprocess
9

10
def te_version() -> str:
11
    """Transformer Engine version string
12

13
    Includes Git commit as local version, unless suppressed with
14
    NVTE_NO_LOCAL_VERSION environment variable.
15

16
    """
17
    root_path = Path(__file__).resolve().parent
18
    with open(root_path / "VERSION", "r") as f:
19
        version = f.readline().strip()
20
    if not int(os.getenv("NVTE_NO_LOCAL_VERSION", "0")):
21
        try:
22
            output = subprocess.run(
23
                ["git", "rev-parse" , "--short", "HEAD"],
24
                capture_output=True,
25
                cwd=root_path,
26
                check=True,
27
                universal_newlines=True,
28
            )
29
        except (subprocess.CalledProcessError, OSError):
30
            pass
31
        else:
32
            commit = output.stdout.strip()
33
            version += f"+{commit}"
34
    return version
35