diff --git a/.jenkins/pytorch/win-test-helpers/build_pytorch.bat b/.jenkins/pytorch/win-test-helpers/build_pytorch.bat index b85dad0616cd71c..114cd725182ab90 100644 --- a/.jenkins/pytorch/win-test-helpers/build_pytorch.bat +++ b/.jenkins/pytorch/win-test-helpers/build_pytorch.bat @@ -144,7 +144,7 @@ python setup.py bdist_wheel && sccache --show-stats && python -c "import os, glo 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 40defe2424befbf..2ef8b7f29dbb8d7 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, @@ -1047,8 +1056,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', @@ -1062,7 +1070,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', @@ -1147,6 +1154,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