aiofastnet/run_benchmark.sh at master · tarasko/aiofastnet · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
#!/usr/bin/env bash
set -euo pipefail

usage() {
    echo "Usage: $0 CPU [benchmark args...]"
    echo
    echo "Example:"
    echo "  $0 3"
    echo "  PYTHON=env314/bin/python $0 3 --help"
}

expand_cpu_list() {
    local list=${1//,/ }
    local item start end cpu

    for item in $list; do
        if [[ $item == *-* ]]; then
            start=${item%-*}
            end=${item#*-}
            for ((cpu = start; cpu <= end; cpu++)); do
                echo "$cpu"
            done
        else
            echo "$item"
        fi
    done
}

write_sysfs() {
    local value=$1
    local path=$2
    echo "$value" | sudo tee "$path" >/dev/null
}

restore() {
    local cpu dir governor min_freq max_freq

    if [[ ${#SAVED_CPUS[@]} -eq 0 ]]; then
        return
    fi

    echo "Restoring CPU frequency settings..."
    for cpu in "${SAVED_CPUS[@]}"; do
        dir="/sys/devices/system/cpu/cpu${cpu}/cpufreq"
        governor=${OLD_GOVERNOR[$cpu]}
        min_freq=${OLD_MIN_FREQ[$cpu]}
        max_freq=${OLD_MAX_FREQ[$cpu]}

        [[ -n $min_freq ]] && write_sysfs "$min_freq" "$dir/scaling_min_freq"
        [[ -n $max_freq ]] && write_sysfs "$max_freq" "$dir/scaling_max_freq"
        [[ -n $governor ]] && write_sysfs "$governor" "$dir/scaling_governor"
    done
}

if [[ $# -lt 1 || ${1:-} == "-h" || ${1:-} == "--help" ]]; then
    usage
    exit 1
fi

CPU=$1
shift
PYTHON=${PYTHON:-python}
CPUFREQ="/sys/devices/system/cpu/cpu${CPU}/cpufreq"

if [[ ! -d $CPUFREQ ]]; then
    echo "CPU $CPU does not expose cpufreq controls at $CPUFREQ" >&2
    exit 1
fi

declare -a POLICY_CPUS=()
declare -a SAVED_CPUS=()
declare -A SEEN=()
declare -A OLD_GOVERNOR=()
declare -A OLD_MIN_FREQ=()
declare -A OLD_MAX_FREQ=()

if [[ -r "$CPUFREQ/related_cpus" ]]; then
    while read -r cpu; do
        [[ -n ${SEEN[$cpu]:-} ]] && continue
        SEEN[$cpu]=1
        POLICY_CPUS+=("$cpu")
    done < <(expand_cpu_list "$(cat "$CPUFREQ/related_cpus")")
else
    POLICY_CPUS=("$CPU")
fi

trap restore EXIT INT TERM

echo "Selected logical CPU: $CPU"
if [[ -r "/sys/devices/system/cpu/cpu${CPU}/topology/thread_siblings_list" ]]; then
    echo "SMT siblings: $(cat "/sys/devices/system/cpu/cpu${CPU}/topology/thread_siblings_list")"
fi
echo "Frequency policy CPUs: ${POLICY_CPUS[*]}"

for cpu in "${POLICY_CPUS[@]}"; do
    dir="/sys/devices/system/cpu/cpu${cpu}/cpufreq"
    [[ -d $dir ]] || continue

    OLD_GOVERNOR[$cpu]=$(cat "$dir/scaling_governor" 2>/dev/null || true)
    OLD_MIN_FREQ[$cpu]=$(cat "$dir/scaling_min_freq" 2>/dev/null || true)
    OLD_MAX_FREQ[$cpu]=$(cat "$dir/scaling_max_freq" 2>/dev/null || true)
    SAVED_CPUS+=("$cpu")
done

for cpu in "${SAVED_CPUS[@]}"; do
    dir="/sys/devices/system/cpu/cpu${cpu}/cpufreq"
    max_freq=$(cat "$dir/cpuinfo_max_freq")

    echo "Setting CPU $cpu to performance governor and max frequency $max_freq"
    write_sysfs performance "$dir/scaling_governor"
    write_sysfs "$max_freq" "$dir/scaling_max_freq"
    write_sysfs "$max_freq" "$dir/scaling_min_freq"
done

echo "Running benchmark on logical CPU $CPU..."
taskset -c "$CPU" "$PYTHON" -m examples.benchmark "$@"