==> src/inferd.c <==1/* inferd.rs: the pair, the call, the answer. */
2#include "mesh/inferd.h"
3#include "mesh/log.h"
4#include "mesh/manifest.h"
5#include "mesh/prog_table.h"
6#include "mesh/rsfmt.h"
7
8#include <arpa/inet.h>
9#include <errno.h>
10#include <fcntl.h>
11#include <netdb.h>
12#include <netinet/in.h>
13#include <netinet/tcp.h>
14#include <poll.h>
15#include <stdio.h>
16#include <stdlib.h>
17#include <string.h>
18#include <sys/socket.h>
19#include <sys/time.h>
20#include <unistd.h>
21
22size_t mesh_inferd_members_of(mesh_shared *sh, const char *name, uint32_t *out) {
23 if (!sh->peer) {
24 uint32_t s;
25 if (mesh_pt_find(sh->mb, name, &s)) { out[0] = 0; return 1; }
26 return 0;
27 }
28 char *text = mesh_shared_text(sh, &sh->fleet_text);
29 size_t n = 0;
30 char *save = NULL;
31 for (char *l = strtok_r(text, "\n", &save); l; l = strtok_r(NULL, "\n", &save)) {
32 char *w[64];
33 size_t nw = 0;
34 char *ws = NULL;
35 for (char *t = strtok_r(l, " \t\r", &ws); t && nw < 64; t = strtok_r(NULL, " \t\r", &ws)) w[nw++] = t;
36 if (nw > 2 && strcmp(w[0], "program") == 0 && strcmp(w[1], name) == 0) {
37 size_t i = 0;
38 while (i < nw && strcmp(w[i], "at") != 0) i++;
39 if (i >= nw || i + 1 >= nw) break;
40 if (strcmp(w[i + 1], "all") == 0) { for (uint32_t k = 0; k < sh->n; k++) out[n++] = k; break; }
41 char *cs = NULL;
42 for (char *x = strtok_r(w[i + 1], ",", &cs); x; x = strtok_r(NULL, ",", &cs)) {
43 uint32_t v;
44 if (mesh_rs_parse_u32(x, strlen(x), &v) == 0 && n < MESH_MAX_SYSTEMS) out[n++] = v;
45 }
46 break;
47 }
48 }
49 free(text);
50 return n;
51}
52
53static int param_of(mesh_shared *sh, const char *name, uint32_t *param, char **err) {
54 char *dir = mesh_buf_fmt("%s/%s", sh->programs, name);
55 mesh_manifest m;
56 char *ptx = NULL, *e = NULL;
57 bool ok = mesh_manifest_load(dir, &m, &ptx, &e);
58 free(dir);
59 free(ptx);
60 if (!ok) { *err = mesh_buf_fmt("%s: %s", name, e); free(e); return -1; }
61 *param = m.param;
62 mesh_manifest_free(&m);
63 return 0;
64}
65
66/* connect_timeout (3 s) to host:port; the fd or -1 with *err (std's words). */
67static int dial(const char *addr, int timeout_ms, char **err) {
68 const char *c = strrchr(addr, ':');
69 if (!c) { *err = mesh_buf_fmt("%s: invalid socket address syntax", addr); return -1; }
70 char host[256];
71 snprintf(host, sizeof host, "%.*s", (int)(c - addr), addr);
72 struct sockaddr_in sa;
73 memset(&sa, 0, sizeof sa);
74 sa.sin_family = AF_INET;
75 sa.sin_port = htons((uint16_t)atoi(c + 1));
76 if (inet_pton(AF_INET, host, &sa.sin_addr) != 1) { *err = mesh_buf_fmt("%s: invalid socket address syntax", addr); return -1; }
77 int fd = socket(AF_INET, SOCK_STREAM, 0);
78 if (fd < 0) { *err = mesh_buf_fmt("%s (os error %d)", strerror(errno), errno); return -1; }
79 int fl = fcntl(fd, F_GETFL);
80 fcntl(fd, F_SETFL, fl | O_NONBLOCK);
81 int r = connect(fd, (struct sockaddr *)&sa, sizeof sa);
82 if (r != 0 && errno == EINPROGRESS) {
83 struct pollfd p = {fd, POLLOUT, 0};
84 r = poll(&p, 1, timeout_ms);
85 if (r == 0) { close(fd); *err = mesh_xstrdup("connection timed out"); return -1; }
86 int so = 0;
87 socklen_t sl = sizeof so;
88 getsockopt(fd, SOL_SOCKET, SO_ERROR, &so, &sl);
89 if (so) { close(fd); *err = mesh_buf_fmt("%s (os error %d)", strerror(so), so); return -1; }
90 } else if (r != 0) {
91 int e = errno;
92 close(fd);
93 *err = mesh_buf_fmt("%s (os error %d)", strerror(e), e);
94 return -1;
95 }
96 fcntl(fd, F_SETFL, fl);
97 return fd;
98}
99
100int mesh_inferd_run(mesh_shared *sh, const mesh_infer_request *r, mesh_infer_msg *res, char **err) {
101 if (!sh->has_svc) { *err = mesh_xstrdup("this meshd has no service addresses"); return -1; }
102 const mesh_infer_svc *svc = &sh->svc;
103 uint32_t me = svc->me;
104 int64_t t0 = mesh_mono_ns();
105 uint8_t kind;
106 uint32_t target;
107 const char *prog;
108 mesh_infer_msg msg;
109 mesh_infer_msg_init(&msg);
110 char addr[300];
111 if (r->has_prefill && r->has_decode) {
112 uint32_t pp;
113 if (param_of(sh, MESH_INFER_PROGRAM_PREFILL, &pp, err)) { mesh_infer_msg_free(&msg); return -1; }
114 mesh_infer_svc_addr(svc, r->decode, r->prefill, pp, addr);
115 mesh_infer_msg_set(&msg, "prefill", addr);
116 mesh_infer_msg_setf(&msg, "prefill_sys", "%u", r->prefill);
117 mesh_infer_msg_setf(&msg, "decode_sys", "%u", r->decode);
118 kind = MESH_INFER_K_DECODE;
119 target = r->decode;
120 prog = MESH_INFER_PROGRAM_DECODE;
121 } else if (r->has_mono) {
122 kind = MESH_INFER_K_GEN;
123 target = r->mono;
124 prog = MESH_INFER_PROGRAM_MONO;
125 mesh_infer_msg_setf(&msg, "prefill_sys", "%u", r->mono);
126 mesh_infer_msg_setf(&msg, "decode_sys", "%u", r->mono);
127 } else {
128 /* The scheduler's placement: a prefill member, and a decode member other than it (this one if it decodes). */
129 uint32_t ps[MESH_MAX_SYSTEMS], ds[MESH_MAX_SYSTEMS];
130 size_t np = mesh_inferd_members_of(sh, MESH_INFER_PROGRAM_PREFILL, ps), nd = mesh_inferd_members_of(sh, MESH_INFER_PROGRAM_DECODE, ds);
131 bool has_d = false;
132 uint32_t dsel = 0;
133 if (np) {
134 for (size_t i = 0; i < nd && !has_d; i++) if (ds[i] == me && ds[i] != ps[0]) { dsel = ds[i]; has_d = true; }
135 for (size_t i = 0; i < nd && !has_d; i++) if (ds[i] != ps[0]) { dsel = ds[i]; has_d = true; }
136 }
137 if (np && has_d) {
138 uint32_t pp;
139 if (param_of(sh, MESH_INFER_PROGRAM_PREFILL, &pp, err)) { mesh_infer_msg_free(&msg); return -1; }
140 mesh_infer_svc_addr(svc, dsel, ps[0], pp, addr);
141 mesh_infer_msg_set(&msg, "prefill", addr);
142 mesh_infer_msg_setf(&msg, "prefill_sys", "%u", ps[0]);
143 mesh_infer_msg_setf(&msg, "decode_sys", "%u", dsel);
144 kind = MESH_INFER_K_DECODE;
145 target = dsel;
146 prog = MESH_INFER_PROGRAM_DECODE;
147 } else {
148 uint32_t mono[MESH_MAX_SYSTEMS];
149 size_t nm = mesh_inferd_members_of(sh, MESH_INFER_PROGRAM_MONO, mono);
150 bool here = false;
151 for (size_t i = 0; i < nm; i++) here |= mono[i] == me;
152 if (!here && !nm) { *err = mesh_xstrdup("no gpt2, gpt2-prefill or gpt2-decode is placed on any member"); mesh_infer_msg_free(&msg); return -1; }
153 target = here ? me : mono[0];
154 kind = MESH_INFER_K_GEN;
155 prog = MESH_INFER_PROGRAM_MONO;
156 mesh_infer_msg_setf(&msg, "prefill_sys", "%u", target);
157 mesh_infer_msg_setf(&msg, "decode_sys", "%u", target);
158 }
159 }
160 if (target >= svc->n) { *err = mesh_buf_fmt("member %u: the cluster has %u", target, svc->n); mesh_infer_msg_free(&msg); return -1; }
161 mesh_infer_msg_set(&msg, "mode", r->mode);
162 mesh_infer_msg_setf(&msg, "max", "%u", r->max);
163 if (r->has_len) mesh_infer_msg_setf(&msg, "len", "%u", r->len);
164 if (r->has_repeat) mesh_infer_msg_setf(&msg, "repeat", "%u", r->repeat);
165 mesh_infer_msg_set_body(&msg, r->prompt);
166 uint32_t param;
167 if (param_of(sh, prog, ¶m, err)) { mesh_infer_msg_free(&msg); return -1; }
168 mesh_infer_svc_addr(svc, me, target, param, addr);
169 char *e = NULL;
170 int fd = dial(addr, 3000, &e);
171 if (fd < 0) {
172 *err = mesh_buf_fmt("%s on member %u (%s): %s", prog, target, addr, e);
173 free(e);
174 mesh_infer_msg_free(&msg);
175 return -1;
176 }
177 int one = 1;
178 setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof one);
179 struct timeval tv = {MESH_INFER_REQUEST_WAIT_MS / 1000, 0};
180 setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof tv);
181 mesh_buf enc = MESH_BUF_INIT;
182 mesh_infer_msg_encode(&msg, &enc);
183 if (mesh_infer_write_frame(fd, kind, enc.data, enc.len) != 0) {
184 *err = mesh_buf_fmt("%s: %s (os error %d)", addr, strerror(errno), errno);
185 mesh_buf_free(&enc);
186 close(fd);
187 mesh_infer_msg_free(&msg);
188 return -1;
189 }
190 mesh_buf_free(&enc);
191 uint8_t k;
192 mesh_buf body = MESH_BUF_INIT;
193 if (mesh_infer_read_frame(fd, &k, &body, NULL, &e) != 0) {
194 *err = mesh_buf_fmt("%s on member %u (%s): %s", prog, target, addr, e);
195 free(e);
196 close(fd);
197 mesh_buf_free(&body);
198 mesh_infer_msg_free(&msg);
199 return -1;
200 }
201 close(fd);
202 if (mesh_infer_msg_decode(body.data ? body.data : "", body.len, res, err) != 0) { mesh_buf_free(&body); mesh_infer_msg_free(&msg); return -1; }
203 mesh_buf_free(&body);
204 if (k == MESH_INFER_K_ERROR) {
205 /* res.body.trim() */
206 const char *s = res->body;
207 size_t n = strlen(s);
208 while (n && (s[n - 1] == ' ' || (s[n - 1] >= 9 && s[n - 1] <= 13))) n--;
209 while (n && (*s == ' ' || (*s >= 9 && *s <= 13))) { s++; n--; }
210 *err = mesh_buf_fmt("%s on member %u: %.*s", prog, target, (int)n, s);
211 mesh_infer_msg_free(res);
212 mesh_infer_msg_free(&msg);
213 return -1;
214 }
215 if (k != MESH_INFER_K_RESULT) {
216 *err = mesh_buf_fmt("%s on member %u: answered a frame of kind %u", prog, target, k);
217 mesh_infer_msg_free(res);
218 mesh_infer_msg_free(&msg);
219 return -1;
220 }
221 mesh_infer_msg_setf(res, "coordinator", "%u", me);
222 const char *keys[2] = {"prefill_sys", "decode_sys"};
223 for (int i = 0; i < 2; i++) {
224 const char *v = mesh_infer_msg_get(&msg, keys[i]);
225 if (v) mesh_infer_msg_set(res, keys[i], v);
226 }
227 mesh_buf ms = MESH_BUF_INIT;
228 mesh_rs_fixed(&ms, (double)(mesh_mono_ns() - t0) / 1e6, 2);
229 mesh_infer_msg_set(res, "request_ms", mesh_buf_str(&ms));
230 mesh_buf_free(&ms);
231 mesh_infer_msg_set(res, "program", prog);
232 const char *g = mesh_infer_msg_get(res, "gen_tokens"), *tt = mesh_infer_msg_get(res, "ttft_ms"), *ts = mesh_infer_msg_get(res, "tok_s");
233 if (kind == MESH_INFER_K_DECODE) {
234 const char *ps = mesh_infer_msg_get(res, "prefill_sys"), *dsy = mesh_infer_msg_get(res, "decode_sys");
235 mesh_info("infer: %s prefill %s decode %s via %s on %u: %s tokens, ttft %s ms, %s tok/s", r->mode, ps ? ps : "?", dsy ? dsy : "?", prog, target,
236 g ? g : "?", tt ? tt : "?", ts ? ts : "?");
237 } else {
238 mesh_info("infer: %s mono %u via %s on %u: %s tokens, ttft %s ms, %s tok/s", r->mode, target, prog, target, g ? g : "?", tt ? tt : "?", ts ? ts : "?");
239 }
240 mesh_infer_msg_free(&msg);
241 return 0;
242}