==> src/worker.c <==1/* mps_worker.rs. */
2#include "mesh/mps_worker.h"
3#include "mesh/cuda.h"
4#include "mesh/log.h"
5#include "mesh/mps_wire.h"
6#include "mesh/text.h"
7
8#include "mesh/abi/base.h"
9#include "mesh/abi/control_block.h"
10
11#include <errno.h>
12#include <fcntl.h>
13#include <pthread.h>
14#include <stdatomic.h>
15#include <stdio.h>
16#include <stdlib.h>
17#include <string.h>
18#include <sys/mman.h>
19#include <sys/stat.h>
20#include <time.h>
21#include <unistd.h>
22
23#define MPS_PIPE_ENV "CUDA_MPS_PIPE_DIRECTORY"
24/* Stream polling while a kernel runs: bounds how late meshd sees a halt. */
25#define POLL_US 20
26
27static pthread_mutex_t out_mu = PTHREAD_MUTEX_INITIALIZER;
28
29static void send_frame(uint8_t op, const uint8_t *b, size_t n) {
30 pthread_mutex_lock(&out_mu);
31 mesh_mps_write_frame(1, op, b, n); /* meshd gone: nothing to tell; EOF on stdin ends us */
32 pthread_mutex_unlock(&out_mu);
33}
34
35/* u32 ok, then the body or the error's text. */
36static void result(uint8_t op, bool ok, const uint8_t *body, size_t n, const char *err) {
37 uint8_t *b = NULL;
38 size_t bn = 0;
39 mesh_mps_put_u32(&b, &bn, ok ? 1 : 0);
40 if (ok) mesh_mps_put_bytes(&b, &bn, body, n);
41 else mesh_mps_put_bytes(&b, &bn, err, strlen(err));
42 send_frame(op, b, bn);
43 free(b);
44}
45
46static _Noreturn void fatal(pid_t pid, const char *e) {
47 mesh_error("worker %d: %s", (int)pid, e);
48 mesh_log_exit(1);
49}
50
51typedef struct {
52 CUcontext ctx;
53 CUstream stream;
54 uint32_t gen;
55} watch;
56
57static void sleep_us(unsigned us) {
58 struct timespec t = {0, (long)us * 1000L};
59 nanosleep(&t, NULL);
60}
61
62/* W_DONE once the stream of launch gen finished (or failed). */
63static void *monitor(void *a) {
64 watch *w = a;
65 const mesh_cuda_fns *cu = mesh_cuda();
66 cu->cuCtxSetCurrent(w->ctx);
67 char *err = NULL;
68 for (;;) {
69 CUresult r = cu->cuStreamQuery(w->stream);
70 if (r == CUDA_SUCCESS) break;
71 if (r == CUDA_ERROR_NOT_READY) {
72 sleep_us(POLL_US);
73 continue;
74 }
75 char *d = mesh_cu_result_debug(r);
76 err = mesh_fmt("stream: %s", d);
77 free(d);
78 break;
79 }
80 if (err) mesh_error("worker %d: %s", (int)getpid(), err);
81 uint8_t *b = NULL;
82 size_t bn = 0;
83 mesh_mps_put_u32(&b, &bn, w->gen);
84 mesh_mps_put_u32(&b, &bn, err ? 0 : 1);
85 if (err) mesh_mps_put_bytes(&b, &bn, err, strlen(err));
86 send_frame(MESH_W_DONE, b, bn);
87 free(b);
88 free(err);
89 free(w);
90 return NULL;
91}
92
93/* Context (green with --sms), the mailbox file registered for the device, a stream. */
94static bool setup(const char *path, uint32_t sms, mesh_cuda_gpu *g, mesh_cuda_stream *st, uint64_t *dev, uint8_t **mb, char **err) {
95 if (!mesh_cuda_gpu_open(sms, g, err)) return false;
96 int fd = open(path, O_RDWR);
97 if (fd < 0) {
98 int e = errno;
99 return mesh_errf(err, "mailbox %s: %s (os error %d)", path, strerror(e), e);
100 }
101 struct stat s;
102 if (fstat(fd, &s) != 0) {
103 int e = errno;
104 close(fd);
105 return mesh_errf(err, "%s (os error %d)", strerror(e), e);
106 }
107 if ((size_t)s.st_size != MESH_MAILBOX_BYTES) {
108 close(fd);
109 return mesh_errf(err, "mailbox %s: %lld bytes, not %d", path, (long long)s.st_size, MESH_MAILBOX_BYTES);
110 }
111 void *p = mmap(NULL, MESH_MAILBOX_BYTES, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
112 int e = errno;
113 close(fd);
114 if (p == MAP_FAILED) return mesh_errf(err, "mmap %s: %s (os error %d)", path, strerror(e), e);
115 const mesh_cuda_fns *cu = mesh_cuda();
116 CUdeviceptr d = 0;
117 if (!mesh_cu_ck_debug(cu->cuMemHostRegister_v2(p, MESH_MAILBOX_BYTES, CU_MEMHOSTREGISTER_PORTABLE | CU_MEMHOSTREGISTER_DEVICEMAP),
118 "cuMemHostRegister", err) ||
119 !mesh_cu_ck_debug(cu->cuMemHostGetDevicePointer_v2(&d, p, 0), "cuMemHostGetDevicePointer", err))
120 return false;
121 if (!mesh_cuda_new_stream(g, st, err)) return false;
122 *dev = d;
123 *mb = p;
124 return true;
125}
126
127/* The module and the loader's per-image facts; the context's stack sized for it now. */
128static bool load(mesh_cuda_gpu *g, const uint8_t *cubin, uint32_t threads, uint32_t smem, uint32_t scratch, mesh_cuda_kernel *k,
129 uint32_t facts[5], char **err) {
130 if (!mesh_cuda_load_cubin(g, cubin, MESH_PROGRAM_ENTRY, k, err)) return false;
131 size_t limit;
132 bool ok = mesh_cuda_kernel_set_shape(k, threads, smem, err) && mesh_cuda_kernel_alloc_scratch(k, scratch, err) &&
133 mesh_cuda_kernel_global_u32(k, MESH_RT_ABI_SYMBOL, &facts[0], err) && mesh_cuda_kernel_regs(k, &facts[1], err) &&
134 mesh_cuda_kernel_max_blocks_per_sm(k, &facts[2], err) && mesh_cuda_kernel_local_bytes(k, &facts[3], err) &&
135 mesh_cuda_stack_limit(g, &limit, err) && (facts[3] <= limit || mesh_cuda_set_stack_limit(g, facts[3], err));
136 facts[4] = g->sms;
137 if (!ok) mesh_cuda_unload(g, k, NULL);
138 return ok;
139}
140
141_Noreturn void mesh_mps_worker_main(int argc, char **argv) {
142 char *err = NULL;
143 if (!mesh_cuda_force_eager(&err)) {
144 mesh_error("worker: %s", err);
145 mesh_log_exit(2);
146 }
147 const char *path = NULL, *sms_s = NULL;
148 for (int i = 0; i + 1 < argc; i++) {
149 if (strcmp(argv[i], "--mailbox") == 0) path = argv[i + 1];
150 if (strcmp(argv[i], "--sms") == 0) sms_s = argv[i + 1];
151 }
152 if (!path) {
153 fputs("usage: meshd worker --mailbox FILE [--sms N] (started by meshd --isolation mps)\n", stderr);
154 mesh_log_exit(2);
155 }
156 pid_t pid = getpid();
157 uint32_t sms = 0;
158 if (sms_s && !mesh_parse_u32(sms_s, strlen(sms_s), 10, &sms)) {
159 char *q = mesh_debug_str(sms_s);
160 mesh_error("worker: bad --sms %s", q);
161 mesh_log_exit(2);
162 }
163 if (!getenv(MPS_PIPE_ENV)) {
164 mesh_error("worker %d: %s is not set; refusing to run outside meshd's MPS daemon", (int)pid, MPS_PIPE_ENV);
165 mesh_log_exit(2);
166 }
167 mesh_cuda_gpu g;
168 mesh_cuda_stream st;
169 uint64_t dev;
170 uint8_t *mb;
171 if (!setup(path, sms, &g, &st, &dev, &mb, &err)) fatal(pid, err);
172 bool launched = false;
173 uint32_t slot_of_launch = 0;
174 bool loaded = false;
175 mesh_cuda_kernel k;
176 char *kname = NULL;
177 uint8_t op, *b;
178 size_t n;
179 while (mesh_mps_read_frame(0, &op, &b, &n)) {
180 if (op == MESH_W_LOAD && !loaded) {
181 uint32_t threads = mesh_mps_u32_at(b, n, 0), smem = mesh_mps_u32_at(b, n, 4), scratch = mesh_mps_u32_at(b, n, 8);
182 size_t nl = mesh_mps_u32_at(b, n, 12);
183 if (nl > (n > 16 ? n - 16 : 0)) nl = n > 16 ? n - 16 : 0;
184 char *name = mesh_strndup((const char *)b + 16, nl);
185 uint32_t facts[5];
186 char *e = NULL;
187 if (load(&g, b + 16 + nl, threads, smem, scratch, &k, facts, &e)) {
188 uint8_t *v = NULL;
189 size_t vn = 0;
190 for (int i = 0; i < 5; i++) mesh_mps_put_u32(&v, &vn, facts[i]);
191 mesh_mps_put_u64(&v, &vn, dev);
192 mesh_mps_put_u64(&v, &vn, k.scratch);
193 result(MESH_W_LOADED, true, v, vn, NULL);
194 free(v);
195 loaded = true;
196 free(kname);
197 kname = name;
198 } else {
199 char *m = mesh_fmt("%s: %s", name, e);
200 result(MESH_W_LOADED, false, NULL, 0, m);
201 free(m);
202 free(e);
203 free(name);
204 }
205 } else if (op == MESH_W_LOAD) {
206 result(MESH_W_LOADED, false, NULL, 0, "a worker holds one image at a time");
207 } else if (op == MESH_W_UNLOAD) {
208 if (loaded) {
209 /* meshd unloads only a halted image: nothing of this context runs. */
210 bool done = false;
211 char *e = NULL;
212 if (!mesh_cuda_query(&g, &st, &done, &e) || !done || !mesh_cuda_unload(&g, &k, &e)) {
213 char *m = mesh_fmt("unload %s: %s", kname, e ? e : "its kernel still runs");
214 fatal(pid, m);
215 }
216 loaded = false;
217 }
218 } else if (op == MESH_W_LAUNCH) {
219 uint32_t gen = mesh_mps_u32_at(b, n, 0), slot = mesh_mps_u32_at(b, n, 4), nodes = mesh_mps_u32_at(b, n, 8);
220 launched = true;
221 slot_of_launch = slot;
222 char *e = NULL;
223 bool ok;
224 if (!loaded) ok = mesh_errf(&e, "launch before load");
225 else {
226 bool done = false;
227 ok = mesh_cuda_query(&g, &st, &done, &e);
228 if (ok && !done) ok = mesh_errf(&e, "its previous kernel is still on the stream");
229 if (ok) {
230 uint64_t d = dev;
231 uint32_t s = slot;
232 void *params[] = {&d, &s};
233 ok = mesh_cuda_launch_on(&g, &st, &k, nodes, params, &e);
234 }
235 }
236 result(MESH_W_LAUNCHED, ok, NULL, 0, e);
237 free(e);
238 if (ok) {
239 watch *w = mesh_alloc(sizeof *w);
240 *w = (watch){g.ctx, st.s, gen};
241 pthread_t t;
242 if (pthread_create(&t, NULL, monitor, w) == 0) pthread_detach(t);
243 }
244 } else if (op == MESH_W_READY) {
245 result(MESH_W_READY, true, NULL, 0, NULL);
246 } else if (op == MESH_W_QUIT) {
247 free(b);
248 break;
249 } else {
250 char *m = mesh_fmt("unknown request %u", op);
251 fatal(pid, m);
252 }
253 free(b);
254 }
255 /* W_QUIT comes only once meshd halted us; EOF may mean meshd died with our
256 * kernel resident, and nobody else will stop it: P_SHUTDOWN ourselves. */
257 bool done = true;
258 char *e = NULL;
259 if (launched && mesh_cuda_query(&g, &st, &done, &e) && !done) {
260 size_t at = mesh_prog_ctl(slot_of_launch) + MESH_P_SHUTDOWN;
261 if (at < MESH_MAILBOX_BYTES) atomic_store_explicit((_Atomic uint32_t *)(void *)(mb + at), 1u, memory_order_release);
262 struct timespec t0, t1;
263 clock_gettime(CLOCK_MONOTONIC, &t0);
264 for (;;) {
265 done = true;
266 if (!mesh_cuda_query(&g, &st, &done, NULL) || done) break;
267 clock_gettime(CLOCK_MONOTONIC, &t1);
268 if (t1.tv_sec - t0.tv_sec >= 2) break;
269 sleep_us(POLL_US);
270 }
271 }
272 free(e);
273 /* Still running: exit without the unload, which would wait for it forever (the exit frees the context). */
274 done = false;
275 if (loaded && mesh_cuda_query(&g, &st, &done, NULL) && done) mesh_cuda_unload(&g, &k, NULL);
276 free(kname);
277 mesh_cuda_gpu_free(&g);
278 mesh_log_exit(0);
279}