ADD __dcu_version_ attribute

nni/__init__.py

@@ -6,6 +6,11 @@ try:
 except ModuleNotFoundError:
     __version__ = '999.dev0'
 
+try:
+    from .version import __dcu_version__
+except ImportError:
+    pass
+
 from .runtime.log import _init_logger
 _init_logger()
 

nni/common/get_dcu_version.py

+import subprocess
+from pathlib import Path
+import os
+UNKNOWN = "Unknown"
+
+def sha_value(tf_root):
+    try:
+        return (
+            subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=tf_root)
+            .decode("ascii")
+            .strip()
+        )
+    except Exception:
+        return UNKNOWN
+
+def abi_value():
+    try:
+        return (
+            subprocess.check_output("echo '#include <string>' | gcc -x c++ -E -dM - | fgrep _GLIBCXX_USE_CXX11_ABI", shell=True)
+            .decode('ascii')
+            .strip()[-1]
+        )
+    except Exception:
+        return UNKNOWN
+
+def dtk_version_value():
+    try:
+       dtk_path=os.getenv('ROCM_PATH')
+       dtk_version_path = os.path.join(dtk_path, '.info', "version-dev")
+       with open(dtk_version_path, 'r',encoding='utf-8') as file:
+           lines = file.readlines()
+       dtk_version="dtk"+lines[0][:-2].replace(".", "")
+       return dtk_version
+    except Exception:
+       return UNKNOWN
+
+def nni_whl_name():
+    try:
+        tf_root = Path(__file__).parent.parent
+        sha = "git" + sha_value(tf_root)[0:7]
+        abi = "abi" + abi_value()
+        dtk_version = dtk_version_value()
+        whl_name = "+" + sha + "." + abi + "." + dtk_version
+        return whl_name
+    except Exception:
+        return UNKNOWN
+
+def dcu_version():
+    try:
+        release = os.environ.get('NNI_RELEASE')
+        nni_version = release or UNKNOWN
+        dcu_version = nni_version + nni_whl_name()
+        return dcu_version
+    except Exception:
+        return UNKNOWN

setup.py

@@ -67,7 +67,13 @@ from setuptools.command.develop import develop
 
 import setup_ts
 
+from nni.common import get_dcu_version
+
 release = os.environ.get('NNI_RELEASE')
+if (release):
+       release_nni = release + get_dcu_version.nni_whl_name()
+    else:
+       release_nni = release
 
 def _get_jupyter_lab_version():
     try:
@@ -92,7 +98,7 @@ def check_jupyter_lab_version():
 def _setup():
     setuptools.setup(
         name = 'nni',
-        version = release or '999.dev0',
+        version = release_nni or '999.dev0',
         description = 'Neural Network Intelligence project',
         long_description = open('README.md', encoding='utf-8').read(),
         long_description_content_type = 'text/markdown',
@@ -232,7 +238,11 @@ class Build(build):
 
         if os.path.islink('nni_node/main.js'):
             sys.exit('A development build already exists. Please uninstall NNI and run "python3 setup.py clean".')
-        open('nni/version.py', 'w').write(f"__version__ = '{release}'")
+        dcu_version = get_dcu_version.dcu_version()
+        version_path = "nni/version.py"
+        with open(version_path, "w") as f:
+            f.write(f"__version__ = '{release}'\n")
+            f.write(f"__dcu_version__ = '{dcu_version}'")
         super().run()
 
 class Develop(develop):