Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[BE] Do not package caffe2 in wheel (#87986) #90433

Merged
merged 1 commit into from Dec 8, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
2 changes: 1 addition & 1 deletion .jenkins/pytorch/win-test-helpers/build_pytorch.bat
Expand Up @@ -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

Expand Down
32 changes: 23 additions & 9 deletions setup.py
Expand Up @@ -762,18 +762,22 @@ 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.

Returns:
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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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',
Expand All @@ -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',
Expand Down Expand Up @@ -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
Expand Down