Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 49 additions & 2 deletions build.py
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,51 @@
log(msg, force=True)


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))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

neat trick!

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):

Check notice

Code scanning / CodeQL

Empty except Note

'except' clause does nothing but pass and there is no explanatory comment.
Comment thread
mc-nv marked this conversation as resolved.
Dismissed
pass
return None


def default_build_parallel():
# Triton is memory-heavy C++/CUDA: template-dense translation units routinely
# peak at 1-2+ GB RAM per compiler process, so blindly using N (or 2*N) cores
# can invoke the OOM killer. Cap parallelism so each job has ~2 GB headroom.
mem_gb_per_job = 2.0
jobs = usable_cpu_count()

mem_gb = available_memory_gb()
if mem_gb is not None:
mem_cap = max(1, int(mem_gb // mem_gb_per_job))
if mem_cap < jobs:
log(
"Limiting build parallelism to {} (from {} usable cores) so each "
"compile job has ~{:.0f} GB of RAM; {:.1f} GB available.".format(
mem_cap, jobs, mem_gb_per_job, mem_gb
)
)
jobs = min(jobs, mem_cap)

return max(1, jobs)


def fail(msg):
fail_if(True, msg)

Expand Down Expand Up @@ -2340,7 +2385,9 @@
type=int,
required=False,
default=None,
help="Build parallelism. Defaults to 2 * number-of-cores.",
help="Build parallelism. Defaults to the usable core count, "
"capped so each parallel compile job has ~2 GB of RAM headroom "
"to avoid out-of-memory build failures.",
)

parser.add_argument(
Expand Down Expand Up @@ -2663,7 +2710,7 @@
os.environ.setdefault("TRITON_RELEASE_VERSION", FLAGS.version)

if FLAGS.build_parallel is None:
FLAGS.build_parallel = multiprocessing.cpu_count() * 2
FLAGS.build_parallel = default_build_parallel()

log("Building Triton Inference Server")
log("platform {}".format(target_platform()))
Expand Down
Loading