-
Notifications
You must be signed in to change notification settings - Fork 80
Expand file tree
/
Copy paththread_default_unit.cpp
More file actions
170 lines (149 loc) · 6.51 KB
/
Copy paththread_default_unit.cpp
File metadata and controls
170 lines (149 loc) · 6.51 KB
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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
// thread_default_unit.cpp - white-box tests for the CPU thread-count helpers
// in transcribe-batch-util (default_n_threads / configure_sched_n_threads).
//
// The bug these pin: the per-arch default used std::thread::hardware_concurrency(),
// which reports the host's total cores and ignores the process CPU affinity
// mask. Under a constrained scheduler (taskset / cpuset / a CI runner pinned to
// N vCPUs) that over-counts, and the resulting oversubscription makes ggml's
// spin-wait barriers livelock. default_n_threads() counts the CPUs the process
// may actually run on instead. So the load-bearing assertion here is the
// affinity one (Linux): pin to K CPUs, expect K.
//
// configure_sched_n_threads() is exercised on the null-scheduler path, which
// isolates its resolution contract (requested>0 passes through; requested<=0
// falls back to default_n_threads()) without standing up a real backend graph —
// the real per-backend application is covered by the e2e/example runs.
#include "transcribe-batch-util.h"
#include <cstdio>
#if defined(__linux__)
# include <sched.h>
#elif defined(_WIN32)
# ifndef WIN32_LEAN_AND_MEAN
# define WIN32_LEAN_AND_MEAN
# endif
# ifndef NOMINMAX // else <windows.h>'s min/max macros clobber std::min/std::max
# define NOMINMAX
# endif
# include <windows.h>
#endif
namespace {
int g_failures = 0;
#define CHECK(cond) \
do { \
if (!(cond)) { \
std::fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, #cond); \
++g_failures; \
} \
} while (0)
// default_n_threads(): cap and range, independent of platform.
void test_default_range_and_cap() {
const int uncapped = transcribe::default_n_threads(/*cap=*/0);
CHECK(uncapped >= 1); // always at least one thread
const int capped8 = transcribe::default_n_threads(8);
CHECK(capped8 >= 1);
CHECK(capped8 <= 8); // cap is honored
CHECK(capped8 == (uncapped < 8 ? uncapped : 8)); // cap = min(uncapped, 8)
CHECK(transcribe::default_n_threads(1) == 1); // cap of 1 pins to 1
const int cap4 = transcribe::default_n_threads(4);
CHECK(cap4 >= 1 && cap4 <= 4);
}
// configure_sched_n_threads(nullptr, ...): resolution contract, no backends.
void test_configure_resolution_null_sched() {
// requested > 0 passes through verbatim.
CHECK(transcribe::configure_sched_n_threads(nullptr, 3) == 3);
CHECK(transcribe::configure_sched_n_threads(nullptr, 1) == 1);
// requested <= 0 falls back to default_n_threads().
const int def = transcribe::default_n_threads();
CHECK(transcribe::configure_sched_n_threads(nullptr, 0) == def);
CHECK(transcribe::configure_sched_n_threads(nullptr, -5) == def);
}
#if defined(__linux__)
// The actual bug: default_n_threads() must honor the process affinity mask,
// not the host core count. Pin to K CPUs and expect K. Skips gracefully if the
// environment forbids changing affinity (some sandboxes do).
void test_affinity_honored_linux() {
cpu_set_t original;
CPU_ZERO(&original);
if (sched_getaffinity(0, sizeof(original), &original) != 0) {
std::fprintf(stderr, "SKIP affinity test: sched_getaffinity failed\n");
return;
}
const int avail = CPU_COUNT(&original);
// Collect the CPU ids currently allowed, so we pin to real, permitted CPUs.
int cpus[CPU_SETSIZE];
int n_cpus = 0;
for (int c = 0; c < CPU_SETSIZE && n_cpus < avail; ++c) {
if (CPU_ISSET(c, &original)) {
cpus[n_cpus++] = c;
}
}
// Pin to exactly one CPU -> default_n_threads must be 1.
cpu_set_t one;
CPU_ZERO(&one);
CPU_SET(cpus[0], &one);
if (sched_setaffinity(0, sizeof(one), &one) != 0) {
std::fprintf(stderr, "SKIP affinity test: sched_setaffinity not permitted\n");
return;
}
CHECK(transcribe::default_n_threads(8) == 1);
CHECK(transcribe::default_n_threads(/*cap=*/0) == 1); // uncapped still sees 1 usable
// Pin to two CPUs -> expect 2 (only when the box actually has >= 2).
if (n_cpus >= 2) {
cpu_set_t two;
CPU_ZERO(&two);
CPU_SET(cpus[0], &two);
CPU_SET(cpus[1], &two);
if (sched_setaffinity(0, sizeof(two), &two) == 0) {
CHECK(transcribe::default_n_threads(8) == 2);
}
}
// Restore the process to its original affinity.
CHECK(sched_setaffinity(0, sizeof(original), &original) == 0);
}
#elif defined(_WIN32)
// Windows analog: GetProcessAffinityMask must drive the default, not the host
// core count. Pin the process to K CPUs via SetProcessAffinityMask and expect K.
// Skips gracefully if affinity changes aren't permitted.
void test_affinity_honored_windows() {
const HANDLE proc = GetCurrentProcess();
DWORD_PTR proc_mask = 0, sys_mask = 0;
if (!GetProcessAffinityMask(proc, &proc_mask, &sys_mask) || proc_mask == 0) {
std::fprintf(stderr, "SKIP affinity test: GetProcessAffinityMask failed\n");
return;
}
const DWORD_PTR original = proc_mask;
// Lowest set bit = one permitted CPU. Pin to it -> expect 1.
const DWORD_PTR one = original & (~original + 1);
if (!SetProcessAffinityMask(proc, one)) {
std::fprintf(stderr, "SKIP affinity test: SetProcessAffinityMask not permitted\n");
return;
}
CHECK(transcribe::default_n_threads(8) == 1);
CHECK(transcribe::default_n_threads(/*cap=*/0) == 1);
// Pin to two CPUs (lowest two set bits) -> expect 2, when available.
const DWORD_PTR second = (original & ~one) & (~(original & ~one) + 1);
if (second != 0) {
if (SetProcessAffinityMask(proc, one | second)) {
CHECK(transcribe::default_n_threads(8) == 2);
}
}
// Restore the process to its original affinity.
CHECK(SetProcessAffinityMask(proc, original) != 0);
}
#endif
} // namespace
int main() {
test_default_range_and_cap();
test_configure_resolution_null_sched();
#if defined(__linux__)
test_affinity_honored_linux();
#elif defined(_WIN32)
test_affinity_honored_windows();
#endif
if (g_failures != 0) {
std::fprintf(stderr, "thread_default_unit: %d failure(s)\n", g_failures);
return 1;
}
std::fprintf(stderr, "thread_default_unit: OK\n");
return 0;
}