Run a command bound to a specific socket's CPUs and memory

numactl-bind-socket.sh

#!/bin/bash
# numactl-bind-socket.sh — Run a command bound to a specific socket's CPUs and memory
# Usage:
#   ./numactl-bind-socket.sh --socket <id> --mode <physical|all> [--interleave <on|off>] <command> [args...]

set -euo pipefail

usage() {
    echo "Usage: $0 --socket <id> --mode <physical|all> [--interleave <on|off>] <command> [args...]"
    exit 1
}

if [ $# -lt 4 ]; then
    usage
fi

SOCKET=""
MODE=""
INTERLEAVE="off"

# --- Parse arguments ---
while [[ $# -gt 0 ]]; do
    case "$1" in
        --socket)
            SOCKET="$2"; shift 2 ;;
        --mode)
            MODE="$2"; shift 2 ;;
        --interleave)
            INTERLEAVE="$2"; shift 2 ;;
        --help|-h)
            usage ;;
        *)
            break ;;
    esac
done

if [[ -z "$SOCKET" || -z "$MODE" ]]; then
    usage
fi

# --- Determine NUMA nodes for this socket ---
NODES=$(lscpu --extended=CPU,SOCKET,NODE | awk -v s=$SOCKET '$2==s {print $3}' | sort -nu | paste -sd,)

if [ -z "$NODES" ]; then
    echo "Error: Could not find any NUMA nodes for socket $SOCKET"
    exit 1
fi

echo "Selected socket: $SOCKET"
echo "Associated NUMA nodes: $NODES"
echo "Mode: $MODE"
echo "Interleave: $INTERLEAVE"

# --- Determine memory binding policy ---
if [[ "$INTERLEAVE" == "on" ]]; then
    MEM_OP="--interleave=$NODES"
    echo "Memory policy: Interleave across nodes ($NODES)"
else
    MEM_OP="--membind=$NODES"
    echo "Memory policy: Bind to nodes ($NODES)"
fi

# --- Determine CPU binding based on mode ---
case "$MODE" in
    physical)
        CPUS=$(lscpu --parse=CPU,SOCKET,CORE | grep -v '^#' | \
               awk -F, -v s=$SOCKET '$2==s && !seen[$3]++ {print $1}' | paste -sd,)
        echo "Binding to physical cores only: $CPUS"
        exec numactl --physcpubind=$CPUS $MEM_OP "$@"
        ;;
    all)
        echo "Binding to all CPUs in NUMA node(s)"
        exec numactl --cpunodebind=$NODES $MEM_OP "$@"
        ;;
    *)
        echo "Invalid mode: $MODE (must be 'physical' or 'all')"
        exit 1
        ;;
esac