fix(build): cap ONNX Runtime build parallelism to avoid runner OOM (#348) (#349)

The generated Dockerfile passed a bare `--parallel` to ONNX Runtime's
build.sh, which expands to multiprocessing.cpu_count(). Combined with
per-nvcc multi-arch compilation, fully-uncached builds exhausted memory
on swap-less builders and the OOM killer terminated runner/system daemons
instead of the build.

Compute the parallelism in Python at Dockerfile-generation time and bake
literal `--parallel <jobs> --nvcc_threads 2` into the build step:

- usable_cpu_count() honors CPU affinity/cgroup pinning (Linux), falling
  back to cpu_count() elsewhere.
- available_memory_gb() reads MemAvailable from /proc/meminfo.
- build_parallelism() budgets ~2 GB of RAM per concurrent compiler slot
  and bounds the product `parallel_jobs * nvcc_threads` by that budget.
- The assumed budget is capped at MAX_BUILD_MEMORY_GB (64 GB): never plan
  for more than that, nor for more than is actually available; assume the
  cap when memory is unknown.

Baking literals keeps the build stage free of shell logic, so it is
portable across the Debian (apt) and RHEL (dnf) base-image paths.

Relates to TRI-1550.

Author: mc-nv

tools/gen_ort_dockerfile.py

@@ -26,6 +26,7 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 import argparse
+import multiprocessing
 import os
 import platform
 import re
@@ -88,6 +89,65 @@
 }
 
 
+def usable_cpu_count():
+    # Honor CPU affinity / cgroup pinning where available (Linux); fall back to
+    # the logical core count on platforms without sched_getaffinity.
+    try:
+        return len(os.sched_getaffinity(0))
+    except AttributeError:
+        return multiprocessing.cpu_count() or 1
+
+
+def available_memory_gb():
+    # Best-effort available RAM in GiB, or None if it cannot be determined.
+    # Reads MemAvailable from /proc/meminfo (Linux); returns None elsewhere so
+    # callers fall back to a CPU-only parallelism default.
+    try:
+        with open("/proc/meminfo") as f:
+            for line in f:
+                if line.startswith("MemAvailable:"):
+                    return int(line.split()[1]) / (1024 * 1024)
+    except (OSError, ValueError, IndexError):
+        pass
+    return None
+
+
+# Memory budget the build is assumed to have for compilation. Parallelism is
+# sized so the ONNX Runtime build's peak memory stays within this, leaving the
+# rest of the host free and avoiding OOM on swap-less builders (see TRI-1550).
+MAX_BUILD_MEMORY_GB = 64.0
+
+# Approximate RAM headroom reserved per concurrent compiler slot.
+MEM_GB_PER_SLOT = 2.0
+
+
+def build_parallelism(nvcc_threads):
+    # Cap ONNX Runtime build parallelism to avoid OOM on swap-less builders
+    # (see TRI-1550). Bare `--parallel` expands to cpu_count(), which combined
+    # with per-nvcc multi-arch compilation exhausts memory. Each concurrent
+    # compiler slot is budgeted ~MEM_GB_PER_SLOT of RAM, and the product
+    # `parallel_jobs * nvcc_threads` is bounded by that budget.
+    #
+    # The assumed memory budget is capped at MAX_BUILD_MEMORY_GB: we never plan
+    # for more than that (so a big-RAM host doesn't over-subscribe), nor for
+    # more than is actually available (so a small builder doesn't OOM). When the
+    # available memory is unknown, we fall back to the MAX_BUILD_MEMORY_GB
+    # assumption rather than to an unbounded core count.
+    #
+    # Computed at Dockerfile-generation time (same host that runs `docker
+    # build`), matching server/build.py; the resulting values are baked into
+    # the Dockerfile as literals so no shell logic runs in the build stage.
+    cpus = usable_cpu_count()
+
+    mem_gb = available_memory_gb()
+    budget_gb = (
+        min(mem_gb, MAX_BUILD_MEMORY_GB) if mem_gb is not None else MAX_BUILD_MEMORY_GB
+    )
+
+    mem_slots = max(1, int(budget_gb // MEM_GB_PER_SLOT))
+    return max(1, min(cpus, mem_slots // max(1, nvcc_threads)))
+
+
 def parse_cuda_arch_list(raw):
     """
     Parse CUDA_ARCH_LIST string to generate CUDA architecture codes:
@@ -391,16 +451,32 @@ def dockerfile_for_linux(output_file):
 
     df += """
 WORKDIR /workspace/onnxruntime
-ARG COMMON_BUILD_ARGS="--config ${{ONNXRUNTIME_BUILD_CONFIG}} --parallel --skip_submodule_sync --build_shared_lib \
+ARG COMMON_BUILD_ARGS="--config ${{ONNXRUNTIME_BUILD_CONFIG}} --skip_submodule_sync --build_shared_lib \
     --compile_no_warning_as_error --build_dir /workspace/build --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES='{}'  --cmake_extra_defines CMAKE_POLICY_VERSION_MINIMUM=3.5 --build_wheel"
 """.format(
         cuda_archs
     )
 
+    # Cap ONNX Runtime build parallelism to avoid OOM on swap-less builders
+    # (see TRI-1550). Values are computed in Python and baked in as literals so
+    # the build stage stays free of shell logic (portable across Debian/RHEL).
+    nvcc_threads = 2
+    ort_jobs = build_parallelism(nvcc_threads)
+    print(
+        "[INFO] ONNX Runtime build parallelism: --parallel {} --nvcc_threads {} "
+        "(usable cores {}, MemAvailable {})".format(
+            ort_jobs,
+            nvcc_threads,
+            usable_cpu_count(),
+            "{:.1f} GB".format(available_memory_gb())
+            if available_memory_gb() is not None
+            else "unknown",
+        )
+    )
     df += """
-RUN ./build.sh ${{COMMON_BUILD_ARGS}} --update --build {}
+RUN ./build.sh ${{COMMON_BUILD_ARGS}} --parallel {} --nvcc_threads {} --update --build {}
 """.format(
-        ep_flags
+        ort_jobs, nvcc_threads, ep_flags
     )
 
     df += """