diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 754757ece..22787fab5 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -41,6 +41,8 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install system dependencies run: sudo apt-get update && sudo apt-get install -y ffmpeg libturbojpeg + - name: Install utils + run: pip install psutil - name: Build and install run: rm -rf .eggs && pip install -e . - name: Validate the installation @@ -73,6 +75,8 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install system dependencies run: sudo apt-get update && sudo apt-get install -y ffmpeg libturbojpeg + - name: Install utils + run: pip install psutil - name: Install PyTorch run: pip install torch==${{matrix.torch}}+cpu torchvision==${{matrix.torchvision}}+cpu -f https://download.pytorch.org/whl/torch_stable.html - name: Build and install @@ -113,6 +117,8 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install system dependencies run: sudo apt-get update && sudo apt-get install -y ffmpeg libturbojpeg + - name: Install utils + run: pip install psutil - name: Install Pillow run: pip install Pillow==6.2.2 if: ${{matrix.torchvision == '0.4.2'}} @@ -181,6 +187,8 @@ jobs: export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${CUDA_HOME}/include:${LD_LIBRARY_PATH} export PATH=${CUDA_HOME}/bin:${PATH} sudo apt-get install -y ninja-build + - name: Install utils + run: pip install psutil - name: Install Pillow run: pip install Pillow==6.2.2 if: ${{matrix.torchvision == '0.4.2'}} @@ -305,6 +313,8 @@ jobs: python-version: 3.7 - name: Install system dependencies run: brew install ffmpeg jpeg-turbo + - name: Install utils + run: pip install psutil - name: Install Pillow run: pip install Pillow==6.2.2 if: ${{matrix.torchvision == '0.4.2'}} diff --git a/setup.py b/setup.py index 086017725..52c287ca7 100644 --- a/setup.py +++ b/setup.py @@ -214,7 +214,14 @@ def get_extensions(): from torch.utils.cpp_extension import CppExtension, CUDAExtension # prevent ninja from using too many resources - os.environ.setdefault('MAX_JOBS', '4') + try: + import psutil + num_cpu = len(psutil.Process().cpu_affinity()) + cpu_use = max(4, num_cpu - 1) + except (ModuleNotFoundError, AttributeError): + cpu_use = 4 + + os.environ.setdefault('MAX_JOBS', str(cpu_use)) define_macros = [] extra_compile_args = {'cxx': []} include_dirs = []