diff --git a/.jenkins/pytorch/win-test-helpers/build_pytorch.bat b/.jenkins/pytorch/win-test-helpers/build_pytorch.bat index 1ea791229113568..da28956cae9716d 100644 --- a/.jenkins/pytorch/win-test-helpers/build_pytorch.bat +++ b/.jenkins/pytorch/win-test-helpers/build_pytorch.bat @@ -145,7 +145,7 @@ python -c "import os, glob; os.system('python -mpip install ' + glob.glob('dist/ if "%BUILD_ENVIRONMENT%"=="" ( echo NOTE: To run `import torch`, please make sure to activate the conda environment by running `call %CONDA_PARENT_DIR%\Miniconda3\Scripts\activate.bat %CONDA_PARENT_DIR%\Miniconda3` in Command Prompt before running Git Bash. ) else ( - 7z a %TMP_DIR_WIN%\%IMAGE_COMMIT_TAG%.7z %CONDA_PARENT_DIR%\Miniconda3\Lib\site-packages\torch %CONDA_PARENT_DIR%\Miniconda3\Lib\site-packages\torchgen %CONDA_PARENT_DIR%\Miniconda3\Lib\site-packages\caffe2 %CONDA_PARENT_DIR%\Miniconda3\Lib\site-packages\functorch && copy /Y "%TMP_DIR_WIN%\%IMAGE_COMMIT_TAG%.7z" "%PYTORCH_FINAL_PACKAGE_DIR%\" + 7z a %TMP_DIR_WIN%\%IMAGE_COMMIT_TAG%.7z %CONDA_PARENT_DIR%\Miniconda3\Lib\site-packages\torch %CONDA_PARENT_DIR%\Miniconda3\Lib\site-packages\torchgen %CONDA_PARENT_DIR%\Miniconda3\Lib\site-packages\functorch && copy /Y "%TMP_DIR_WIN%\%IMAGE_COMMIT_TAG%.7z" "%PYTORCH_FINAL_PACKAGE_DIR%\" if errorlevel 1 exit /b if not errorlevel 0 exit /b diff --git a/setup.py b/setup.py index e3eb3ced6005b36..72a6bbae7b46022 100644 --- a/setup.py +++ b/setup.py @@ -762,6 +762,14 @@ def run(self): super().run() +def get_cmake_cache_vars(): + try: + return defaultdict(lambda: False, cmake.get_cmake_cache_variables()) + except FileNotFoundError: + # CMakeCache.txt does not exist. Probably running "python setup.py clean" over a clean directory. + return defaultdict(lambda: False) + + def configure_extension_build(): r"""Configures extension build options according to system environment and user's choice. @@ -769,11 +777,7 @@ def configure_extension_build(): The input to parameters ext_modules, cmdclass, packages, and entry_points as required in setuptools.setup. """ - try: - cmake_cache_vars = defaultdict(lambda: False, cmake.get_cmake_cache_variables()) - except FileNotFoundError: - # CMakeCache.txt does not exist. Probably running "python setup.py clean" over a clean directory. - cmake_cache_vars = defaultdict(lambda: False) + cmake_cache_vars = get_cmake_cache_vars() ################################################################################ # Configure compile flags @@ -877,7 +881,12 @@ def make_relative_rpath_args(path): ################################################################################ extensions = [] - packages = find_packages(exclude=('tools', 'tools.*')) + excludes = ['tools', 'tools.*'] + if not cmake_cache_vars['BUILD_CAFFE2']: + excludes.extend(['caffe2', 'caffe2.*']) + if not cmake_cache_vars['BUILD_FUNCTORCH']: + excludes.extend(['functorch', 'functorch.*']) + packages = find_packages(exclude=excludes) C = Extension("torch._C", libraries=main_libraries, sources=main_sources, @@ -1055,8 +1064,7 @@ def main(): 'include/ATen/native/quantized/*.h', 'include/ATen/native/quantized/cpu/*.h', 'include/ATen/quantized/*.h', - 'include/caffe2/utils/*.h', - 'include/caffe2/utils/**/*.h', + 'include/caffe2/serialize/*.h', 'include/c10/*.h', 'include/c10/macros/*.h', 'include/c10/core/*.h', @@ -1070,7 +1078,6 @@ def main(): 'include/c10/cuda/impl/*.h', 'include/c10/hip/*.h', 'include/c10/hip/impl/*.h', - 'include/caffe2/**/*.h', 'include/torch/*.h', 'include/torch/csrc/*.h', 'include/torch/csrc/api/include/torch/*.h', @@ -1158,6 +1165,13 @@ def main(): 'utils/model_dump/code.js', 'utils/model_dump/*.mjs', ] + + if get_cmake_cache_vars()['BUILD_CAFFE2']: + torch_package_data.extend([ + 'include/caffe2/**/*.h', + 'include/caffe2/utils/*.h', + 'include/caffe2/utils/**/*.h', + ]) torchgen_package_data = [ # Recursive glob doesn't work in setup.py, # https://github.com/pypa/setuptools/issues/1806