[chore] do not build cuda extensions by default (#634)

CHANGELOG.md

@@ -5,6 +5,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## NEXT - TBD
+### Fixed
+- setup.py: hide CUDA extensions behind BUILD_CUDA_EXTENSIONS envvar
 ### Added
 - FSDP: better memory usage for reduce bucket ([#633](https://github.com/facebookresearch/fairscale/pull/633))
 

pyproject.toml

 [build-system]
 requires = [
     "setuptools >= 40.6.2",
-    "torch >= 1.4.0",
     "wheel >= 0.30.0"
 ]
 build-backend = "setuptools.build_meta:__legacy__"

setup.py

@@ -7,11 +7,8 @@
 
 import os
 import re
-import warnings
 
 import setuptools
-import torch
-from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension
 
 this_dir = os.path.dirname(os.path.abspath(__file__))
 
@@ -34,8 +31,9 @@ def find_version(version_file_path):
 extensions = []
 cmdclass = {}
 
-force_cuda = os.getenv("FORCE_CUDA", "0") == "1"
-if (torch.cuda.is_available() and CUDA_HOME is not None) or force_cuda:
+if os.getenv("BUILD_CUDA_EXTENSIONS", "0") == "1":
+    from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
     extensions.extend(
         [
             CUDAExtension(
@@ -51,8 +49,6 @@ if (torch.cuda.is_available() and CUDA_HOME is not None) or force_cuda:
     )
 
     cmdclass["build_ext"] = BuildExtension
-else:
-    warnings.warn("Cannot install FusedAdam cuda.")
 
 
 if __name__ == "__main__":