Skip to content

Commit 7e0aa5e

Browse files
authored
fix(build): cap default build parallelism by available memory to prevent OOM (#8873) (#8875)
1 parent 601136a commit 7e0aa5e

1 file changed

Lines changed: 49 additions & 2 deletions

File tree

build.py

Lines changed: 49 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,51 @@ def log_verbose(msg):
106106
log(msg, force=True)
107107

108108

109+
def usable_cpu_count():
110+
# Honor CPU affinity / cgroup pinning where available (Linux); fall back to
111+
# the logical core count on platforms without sched_getaffinity.
112+
try:
113+
return len(os.sched_getaffinity(0))
114+
except AttributeError:
115+
return multiprocessing.cpu_count() or 1
116+
117+
118+
def available_memory_gb():
119+
# Best-effort available RAM in GiB, or None if it cannot be determined.
120+
# Reads MemAvailable from /proc/meminfo (Linux); returns None elsewhere so
121+
# callers fall back to a CPU-only parallelism default.
122+
try:
123+
with open("/proc/meminfo") as f:
124+
for line in f:
125+
if line.startswith("MemAvailable:"):
126+
return int(line.split()[1]) / (1024 * 1024)
127+
except (OSError, ValueError, IndexError):
128+
pass
129+
return None
130+
131+
132+
def default_build_parallel():
133+
# Triton is memory-heavy C++/CUDA: template-dense translation units routinely
134+
# peak at 1-2+ GB RAM per compiler process, so blindly using N (or 2*N) cores
135+
# can invoke the OOM killer. Cap parallelism so each job has ~2 GB headroom.
136+
mem_gb_per_job = 2.0
137+
jobs = usable_cpu_count()
138+
139+
mem_gb = available_memory_gb()
140+
if mem_gb is not None:
141+
mem_cap = max(1, int(mem_gb // mem_gb_per_job))
142+
if mem_cap < jobs:
143+
log(
144+
"Limiting build parallelism to {} (from {} usable cores) so each "
145+
"compile job has ~{:.0f} GB of RAM; {:.1f} GB available.".format(
146+
mem_cap, jobs, mem_gb_per_job, mem_gb
147+
)
148+
)
149+
jobs = min(jobs, mem_cap)
150+
151+
return max(1, jobs)
152+
153+
109154
def fail(msg):
110155
fail_if(True, msg)
111156

@@ -2340,7 +2385,9 @@ def enable_all():
23402385
type=int,
23412386
required=False,
23422387
default=None,
2343-
help="Build parallelism. Defaults to 2 * number-of-cores.",
2388+
help="Build parallelism. Defaults to the usable core count, "
2389+
"capped so each parallel compile job has ~2 GB of RAM headroom "
2390+
"to avoid out-of-memory build failures.",
23442391
)
23452392

23462393
parser.add_argument(
@@ -2663,7 +2710,7 @@ def enable_all():
26632710
os.environ.setdefault("TRITON_RELEASE_VERSION", FLAGS.version)
26642711

26652712
if FLAGS.build_parallel is None:
2666-
FLAGS.build_parallel = multiprocessing.cpu_count() * 2
2713+
FLAGS.build_parallel = default_build_parallel()
26672714

26682715
log("Building Triton Inference Server")
26692716
log("platform {}".format(target_platform()))

0 commit comments

Comments
 (0)