Move torch/version.py generation to cmake (pytorch#44577)

Summary:
Pull Request resolved: pytorch#44577

I would like to to move this to cmake so that I can depend on it
happening from other parts of the build.

This PR pulls out the logic for determining the version string and
writing the version file into its own module. `setup.py` still receives
the version string and uses it as before, but now the code for writing
out `torch/version.py` lives in a custom command in torch/CMakeLists.txt

I noticed a small inconsistency in how version info is populated.
`TORCH_BUILD_VERSION` is populated from `setup.py` at configuration
time, while `torch/version.py` is written at build time. So if, e.g. you
configured cmake on a certain git rev, then built it in on another, the
two versions would be inconsistent.

This does not appear to matter, so I opted to preserve the existing
behavior.

Test Plan: Imported from OSS

Reviewed By: bertmaher

Differential Revision: D23734781

Pulled By: suo

fbshipit-source-id: 4002c9ec8058503dc0550f8eece2256bc98c03a4

Author: suo

setup.py

@@ -192,7 +192,6 @@
 import distutils.command.clean
 import distutils.sysconfig
 import filecmp
-import subprocess
 import shutil
 import os
 import json
@@ -203,6 +202,7 @@
 from tools.setup_helpers.env import (IS_WINDOWS, IS_DARWIN, IS_LINUX,
                                      check_env_flag, build_type)
 from tools.setup_helpers.cmake import CMake
+from tools.generate_torch_version import get_torch_version
 
 ################################################################################
 # Parameters parsed from environment
@@ -276,22 +276,7 @@ def report(*args):
 # Version, create_version_file, and package_name
 ################################################################################
 package_name = os.getenv('TORCH_PACKAGE_NAME', 'torch')
-version = open('version.txt', 'r').read().strip()
-sha = 'Unknown'
-
-try:
-    sha = subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=cwd).decode('ascii').strip()
-except Exception:
-    pass
-
-if os.getenv('PYTORCH_BUILD_VERSION'):
-    assert os.getenv('PYTORCH_BUILD_NUMBER') is not None
-    build_number = int(os.getenv('PYTORCH_BUILD_NUMBER'))
-    version = os.getenv('PYTORCH_BUILD_VERSION')
-    if build_number > 1:
-        version += '.post' + str(build_number)
-elif sha != 'Unknown':
-    version += '+' + sha[:7]
+version = get_torch_version()
 report("Building wheel {}-{}".format(package_name, version))
 
 cmake = CMake()
@@ -330,18 +315,6 @@ def check_file(f):
                  cmake_only=CMAKE_ONLY,
                  cmake=cmake)
 
-    version_path = os.path.join(cwd, 'torch', 'version.py')
-    with open(version_path, 'w') as f:
-        f.write("__version__ = '{}'\n".format(version))
-        # NB: This is not 100% accurate, because you could have built the
-        # library code with DEBUG, but csrc without DEBUG (in which case
-        # this would claim to be a release build when it's not.)
-        f.write("debug = {}\n".format(repr(build_type.is_debug())))
-        cmake_cache_vars = defaultdict(lambda: None, cmake.get_cmake_cache_variables())
-        f.write("cuda = {}\n".format(repr(cmake_cache_vars['CUDA_VERSION'])))
-        f.write("git_version = {}\n".format(repr(sha)))
-        f.write("hip = {}\n".format(repr(cmake_cache_vars['HIP_VERSION'])))
-
     if CMAKE_ONLY:
         report('Finished running cmake. Run "ccmake build" or '
                '"cmake-gui build" to adjust build options and '

tools/generate_torch_version.py

@@ -0,0 +1,53 @@
+import argparse
+import os
+import subprocess
+from pathlib import Path
+
+def get_sha():
+    try:
+        return subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=pytorch_root).decode('ascii').strip()
+    except Exception:
+        return 'Unknown'
+
+def get_torch_version(sha=None):
+    pytorch_root = Path(__file__).parent.parent
+    version = open('version.txt', 'r').read().strip()
+    if sha is None:
+        sha = get_sha()
+
+    if os.getenv('PYTORCH_BUILD_VERSION'):
+        assert os.getenv('PYTORCH_BUILD_NUMBER') is not None
+        build_number = int(os.getenv('PYTORCH_BUILD_NUMBER'))
+        version = os.getenv('PYTORCH_BUILD_VERSION')
+        if build_number > 1:
+            version += '.post' + str(build_number)
+    elif sha != 'Unknown':
+        version += '+' + sha[:7]
+    return version
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Generate torch/version.py from build and environment metadata.")
+    parser.add_argument("--is_debug", type=bool, help="Whether this build is debug mode or not.")
+    parser.add_argument("--cuda_version", type=str)
+    parser.add_argument("--hip_version", type=str)
+
+    args = parser.parse_args()
+
+    assert args.is_debug is not None
+    args.cuda_version = None if args.cuda_version == '' else args.cuda_version
+    args.hip_version = None if args.hip_version == '' else args.hip_version
+
+    pytorch_root = Path(__file__).parent.parent
+    version_path = pytorch_root / "torch" / "version.py"
+    sha = get_sha()
+    version = get_torch_version(sha)
+
+    with open(version_path, 'w') as f:
+        f.write("__version__ = '{}'\n".format(version))
+        # NB: This is not 100% accurate, because you could have built the
+        # library code with DEBUG, but csrc without DEBUG (in which case
+        # this would claim to be a release build when it's not.)
+        f.write("debug = {}\n".format(repr(args.is_debug)))
+        f.write("cuda = {}\n".format(repr(args.cuda_version)))
+        f.write("git_version = {}\n".format(repr(sha)))
+        f.write("hip = {}\n".format(repr(args.hip_version)))

torch/CMakeLists.txt

@@ -284,3 +284,25 @@ if(NOT TORCH_PYTHON_LINK_FLAGS STREQUAL "")
 endif()
 
 install(TARGETS torch_python DESTINATION "${TORCH_INSTALL_LIB_DIR}")
+
+# Generate torch/version.py from the appropriate CMake cache variables.
+if(${CMAKE_BUILD_TYPE} STREQUAL "Debug")
+  set(TORCH_VERSION_DEBUG 1)
+else()
+  set(TORCH_VERSION_DEBUG 0)
+endif()
+
+add_custom_command(
+  OUTPUT ${TORCH_SRC_DIR}/version.py
+  COMMAND
+    "${PYTHON_EXECUTABLE}" ${TOOLS_PATH}/generate_torch_version.py
+      --is_debug=${TORCH_VERSION_DEBUG}
+      --cuda_version=${CUDA_VERSION}
+      --hip_version=${HIP_VERSION}
+  WORKING_DIRECTORY ${TORCH_ROOT}
+)
+add_custom_target(
+  gen_torch_version ALL
+  DEPENDS ${TORCH_SRC_DIR}/version.py
+)
+add_dependencies(torch_python gen_torch_version)