Respect user-set CMAKE_PREFIX_PATH (#61904)

Summary:
Fixes the case where the `CMAKE_PREFIX_PATH` variable gets silently overwritten by a user specified environment variable.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/61904

Reviewed By: walterddr, malfet

Differential Revision: D29792014

Pulled By: cbalioglu

fbshipit-source-id: babacc8d5a1490bff1e14247850cc00c6ba9e6be

diff --git a/tools/setup_helpers/cmake.py b/tools/setup_helpers/cmake.py
index 4541f97..5200972 100644 (file)
--- a/tools/setup_helpers/cmake.py
+++ b/tools/setup_helpers/cmake.py
@@ -9,7 +9,7 @@ from subprocess import check_call, check_output, CalledProcessError
 import sys
 import sysconfig
 from setuptools import distutils  # type: ignore[import]
-from typing import IO, Any, Dict, List, Optional, Union
+from typing import IO, Any, Dict, List, Optional, Union, cast
 
 from . import which
 from .env import (BUILD_DIR, IS_64BIT, IS_DARWIN, IS_WINDOWS, check_negative_env_flag)
@@ -224,10 +224,8 @@ class CMake:
         _mkdir_p(self.build_dir)
 
         # Store build options that are directly stored in environment variables
-        build_options: Dict[str, CMakeValue] = {
-            # The default value cannot be easily obtained in CMakeLists.txt. We set it here.
-            'CMAKE_PREFIX_PATH': sysconfig.get_path('purelib')
-        }
+        build_options: Dict[str, CMakeValue] = {}
+
         # Build options that do not start with "BUILD_", "USE_", or "CMAKE_" and are directly controlled by env vars.
         # This is a dict that maps environment variables to the corresponding variable name in CMake.
         additional_options = {
@@ -281,6 +279,16 @@ class CMake:
             elif var.startswith(('BUILD_', 'USE_', 'CMAKE_')) or var.endswith(('EXITCODE', 'EXITCODE__TRYRUN_OUTPUT')):
                 build_options[var] = val
 
+        # The default value cannot be easily obtained in CMakeLists.txt. We set it here.
+        py_lib_path = sysconfig.get_path('purelib')
+        cmake_prefix_path = build_options.get('CMAKE_PREFIX_PATH', None)
+        if cmake_prefix_path:
+            build_options["CMAKE_PREFIX_PATH"] = (
+                cast(str, py_lib_path) + ";" + cast(str, cmake_prefix_path)
+            )
+        else:
+            build_options['CMAKE_PREFIX_PATH'] = py_lib_path
+
         # Some options must be post-processed. Ideally, this list will be shrunk to only one or two options in the
         # future, as CMake can detect many of these libraries pretty comfortably. We have them here for now before CMake
         # integration is completed. They appear here not in the CMake.defines call below because they start with either