build: add memory-aware compile/link parallelism (#1718)

* build: add memory-aware compile/link parallelism Auto-detect available memory and CPU cores at configure time to calculate safe parallel job limits. With Ninja generator, creates separate job pools for compilation (~1.5GB/job) and linking (~4GB/job) so high-core machines can compile fast without OOM during linking. Changes: - New mooncake-common/limit_jobs.cmake module - Include from common.cmake for all build modes - Switch Dockerfile and CI workflows to Ninja - User can override via -DPARALLEL_COMPILE_JOBS / -DPARALLEL_LINK_JOBS Signed-off-by: staryxchen <staryxchen@tencent.com> * fix(ci): remove sudo from Ascend CI install step The Ascend CI job runs inside a Docker container as root, where sudo is not available. This caused exit code 127 after a successful build. Signed-off-by: staryxchen <staryxchen@tencent.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Signed-off-by: staryxchen <staryxchen@tencent.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-23 22:33:21 +08:00 · 2026-03-23 22:33:21 +08:00 · 79266ffb4d
parent cc2ecf6b1a
commit 79266ffb4d
7 changed files with 122 additions and 28 deletions
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@ -54,10 +54,10 @@ jobs:
        method: 'network'
        sub-packages: '["nvcc"]'

-    - name: Install coverage tools
+    - name: Install coverage tools and build utilities
      run: |
        sudo apt-get update
-        sudo apt-get install -y lcov gcovr
+        sudo apt-get install -y lcov gcovr ninja-build

    - name: Set up coverage compilation flags
      run: |
@ -87,14 +87,14 @@ jobs:
        sudo bash -x dependencies.sh -y
        mkdir build
        cd build
-        cmake .. -DUSE_HTTP=ON -DUSE_CXL=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DENABLE_ASAN=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Debug
+        cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CXL=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DENABLE_ASAN=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Debug
      shell: bash

    - name: Build project
      run: |
        cd build
-        make -j4
-        sudo make install
+        cmake --build .
+        sudo cmake --install .
      shell: bash

    - name: Build nvlink_allocator.so
@ -117,7 +117,7 @@ jobs:
        cd build
        export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
        ldconfig -v || echo "always continue"
-        MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 make test -j ARGS="-V"
+        MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -j --output-on-failure
      shell: bash

    - name: Generate coverage report
@ -209,18 +209,19 @@ jobs:
    - name: Configure project
      run: |
        apt update -y
+        apt install -y ninja-build
        bash -x dependencies.sh -y
        mkdir build
        cd build
-        cmake .. -DUSE_MUSA=ON -DUSE_MNNVL=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DUSE_CXL=ON -DUSE_TCP=ON -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF
+        cmake -G Ninja .. -DUSE_MUSA=ON -DUSE_MNNVL=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DUSE_CXL=ON -DUSE_TCP=ON -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF
      shell: bash

    - name: Build project
      run: |
        cd build
        source ~/.bashrc
-        make -j4
-        make install
+        cmake --build .
+        cmake --install .
      shell: bash

  test-wheel-ubuntu:
@ -411,6 +412,7 @@ jobs:
    - name: Install dependencies
      run: |
        sudo apt update -y
+        sudo apt install -y ninja-build
        sudo bash -x dependencies.sh -y
        df -h
      shell: bash
@ -422,9 +424,9 @@ jobs:
        cd build
        export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
        export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
-        cmake .. -DUSE_ETCD=OFF -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=OFF -DUSE_MNNVL=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
-        make -j4
-        sudo make install
+        cmake -G Ninja .. -DUSE_ETCD=OFF -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=OFF -DUSE_MNNVL=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
+        cmake --build .
+        sudo cmake --install .
        df -h
      shell: bash

@ -432,7 +434,7 @@ jobs:
      run: |
        mkdir build
        cd build
-        cmake .. -DUSE_ETCD=ON -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=ON -DUSE_MNNVL=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
+        cmake -G Ninja .. -DUSE_ETCD=ON -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=ON -DUSE_MNNVL=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
      shell: bash
      # TODO: lack USE_NVMEOF,USE_MNNVL

@ -441,15 +443,15 @@ jobs:
        export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
        export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
        cd build
-        make -j4
-        sudo make install
+        cmake --build .
+        sudo cmake --install .
        df -h
      shell: bash

    - name: Configure project with unit tests and examples
      run: |
        cd build
-        cmake .. -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON
+        cmake -G Ninja .. -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON
      shell: bash
      # TODO: lack WITH_RUST_EXAMPLE

@ -458,15 +460,15 @@ jobs:
        export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
        export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
        cd build
-        make -j4
-        sudo make install
+        cmake --build .
+        sudo cmake --install .
      shell: bash

    - name: Configure project
      run: |
        cd build
        rm -r */tests
-        cmake .. -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF -DUSE_HTTP=ON -DENABLE_SCCACHE=ON -DUSE_CXL=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0"
+        cmake -G Ninja .. -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF -DUSE_HTTP=ON -DENABLE_SCCACHE=ON -DUSE_CXL=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0"
      shell: bash

    - name: Build project
@ -474,8 +476,8 @@ jobs:
        export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
        export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
        cd build
-        make -j4
-        sudo make install
+        cmake --build .
+        sudo cmake --install .
      shell: bash

    - name: Build nvlink_allocator.so
--- a/.github/workflows/ci_ascend.yml
+++ b/.github/workflows/ci_ascend.yml
@ -68,8 +68,8 @@ jobs:
        source /usr/local/Ascend/cann-9.0.0/set_env.sh
        echo "Building..."
        cd build
-        make -j8
-        make install -j8
+        cmake --build . -j$(nproc)
+        cmake --install .
        echo "Mooncake installed successfully."

    - name: Run Hixl Mooncake Store Test
--- a/.github/workflows/ci_cu13.yml
+++ b/.github/workflows/ci_cu13.yml
@ -65,6 +65,7 @@ jobs:
    - name: Install dependencies
      run: |
        sudo apt update -y
+        sudo apt install -y ninja-build
        sudo bash -x dependencies.sh -y
        df -h
      shell: bash
@ -73,7 +74,7 @@ jobs:
      run: |
        mkdir build
        cd build
-        cmake .. \
+        cmake -G Ninja .. \
          -DUSE_ETCD=ON \
          -DUSE_REDIS=ON \
          -DUSE_HTTP=ON \
@ -96,8 +97,8 @@ jobs:
        export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
        export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
        cd build
-        make -j4
-        sudo make install
+        cmake --build .
+        sudo cmake --install .
        df -h
      shell: bash

--- a/dependencies.sh
+++ b/dependencies.sh
@ -103,6 +103,7 @@ echo -e "${YELLOW}This may take a few minutes...${NC}"

 SYSTEM_PACKAGES="build-essential \
                  cmake \
+                  ninja-build \
                  git \
                  wget \
                  unzip \
--- a/docker/mooncake.Dockerfile
+++ b/docker/mooncake.Dockerfile
@ -28,6 +28,7 @@ RUN apt-get update && \
        ca-certificates \
        curl \
        git \
+        ninja-build \
        python3 \
        python3-dev \
        python3-pip \
@ -44,7 +45,7 @@ RUN bash dependencies.sh -y
 # Configure & build Mooncake
 RUN mkdir -p build && \
    cd build && \
-    cmake .. \
+    cmake -G Ninja .. \
        -DBUILD_UNIT_TESTS=OFF \
        -DUSE_HTTP=ON \
        -DUSE_ETCD=ON \
@ -53,7 +54,7 @@ RUN mkdir -p build && \
        -DSTORE_USE_ETCD=ON \
        -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} && \
    export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH && \
-    cmake --build . -j"$(nproc)"
+    cmake --build .

 # Build nvlink allocator to make wheel self-contained for CUDA paths
 RUN export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH && \
--- a/mooncake-common/common.cmake
+++ b/mooncake-common/common.cmake
@ -40,6 +40,9 @@ add_definitions(-DCONFIG_ERDMA)

 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)

+# Memory-aware build parallelism (compile vs. link job pools)
+include(${CMAKE_CURRENT_LIST_DIR}/limit_jobs.cmake)
+
 option(ENABLE_SCCACHE "Whether to open sccache" OFF)
 if (ENABLE_SCCACHE)
  find_program(SCCACHE sccache REQUIRED)
--- a/mooncake-common/limit_jobs.cmake
+++ b/mooncake-common/limit_jobs.cmake
@ -0,0 +1,86 @@
+# limit_jobs.cmake — Memory-aware build parallelism
+#
+# Auto-detects available memory and CPU count, calculates safe parallel job
+# limits for compilation and linking separately. With Ninja, creates job pools
+# so compilation uses many cores while memory-heavy linking is restricted.
+#
+# User overrides (cmake -D...):
+#   PARALLEL_COMPILE_JOBS  — override compile parallelism
+#   PARALLEL_LINK_JOBS     — override link parallelism
+#   MAX_COMPILER_MEMORY_MB — per-compile-job memory estimate (default: 1500)
+#   MAX_LINKER_MEMORY_MB   — per-link-job memory estimate (default: 4000)
+
+set(MAX_COMPILER_MEMORY_MB "1500" CACHE STRING
+    "Estimated peak memory per compile job in MB")
+set(MAX_LINKER_MEMORY_MB "4000" CACHE STRING
+    "Estimated peak memory per link job in MB")
+
+# Guard against invalid user input (division by zero)
+if(MAX_COMPILER_MEMORY_MB LESS_EQUAL 0)
+    message(WARNING "[limit_jobs] MAX_COMPILER_MEMORY_MB=${MAX_COMPILER_MEMORY_MB} "
+        "invalid, falling back to 1500")
+    set(MAX_COMPILER_MEMORY_MB 1500 CACHE STRING
+        "Estimated peak memory per compile job in MB" FORCE)
+endif()
+if(MAX_LINKER_MEMORY_MB LESS_EQUAL 0)
+    message(WARNING "[limit_jobs] MAX_LINKER_MEMORY_MB=${MAX_LINKER_MEMORY_MB} "
+        "invalid, falling back to 4000")
+    set(MAX_LINKER_MEMORY_MB 4000 CACHE STRING
+        "Estimated peak memory per link job in MB" FORCE)
+endif()
+
+# Detect system resources
+cmake_host_system_information(RESULT _available_mem_mb
+    QUERY AVAILABLE_PHYSICAL_MEMORY)
+cmake_host_system_information(RESULT _nproc
+    QUERY NUMBER_OF_LOGICAL_CORES)
+
+message(STATUS "[limit_jobs] Available memory: ${_available_mem_mb} MB, "
+    "CPU cores: ${_nproc}")
+
+# Calculate safe parallel jobs from memory
+math(EXPR _compile_jobs "${_available_mem_mb} / ${MAX_COMPILER_MEMORY_MB}")
+math(EXPR _link_jobs "${_available_mem_mb} / ${MAX_LINKER_MEMORY_MB}")
+
+# Clamp: [1, nproc]
+if(_compile_jobs LESS 1)
+    set(_compile_jobs 1)
+endif()
+if(_compile_jobs GREATER _nproc)
+    set(_compile_jobs ${_nproc})
+endif()
+if(_link_jobs LESS 1)
+    set(_link_jobs 1)
+endif()
+if(_link_jobs GREATER _nproc)
+    set(_link_jobs ${_nproc})
+endif()
+
+# Use auto-detected values unless user explicitly overrides with -D
+if(NOT DEFINED PARALLEL_COMPILE_JOBS)
+    set(PARALLEL_COMPILE_JOBS "${_compile_jobs}")
+endif()
+if(NOT DEFINED PARALLEL_LINK_JOBS)
+    set(PARALLEL_LINK_JOBS "${_link_jobs}")
+endif()
+
+message(STATUS "[limit_jobs] Compile jobs: ${PARALLEL_COMPILE_JOBS} "
+    "(${MAX_COMPILER_MEMORY_MB} MB/job), "
+    "Link jobs: ${PARALLEL_LINK_JOBS} (${MAX_LINKER_MEMORY_MB} MB/job)")
+
+# Apply to build system
+if(CMAKE_GENERATOR MATCHES "Ninja")
+    set_property(GLOBAL APPEND PROPERTY JOB_POOLS
+        compile_pool=${PARALLEL_COMPILE_JOBS}
+        link_pool=${PARALLEL_LINK_JOBS}
+    )
+    set(CMAKE_JOB_POOL_COMPILE "compile_pool" CACHE STRING "" FORCE)
+    set(CMAKE_JOB_POOL_LINK "link_pool" CACHE STRING "" FORCE)
+    message(STATUS "[limit_jobs] Ninja job pools: "
+        "compile=${PARALLEL_COMPILE_JOBS}, link=${PARALLEL_LINK_JOBS}")
+else()
+    message(STATUS "[limit_jobs] Hint: use -G Ninja for automatic "
+        "compile/link parallelism separation")
+    message(STATUS "[limit_jobs] With Make, recommend: "
+        "cmake --build . -j${PARALLEL_LINK_JOBS}")
+endif()