==> tests/inferd.c <==1/* A gpt2 worker on loopback answers the request meshd makes (single meshd: mono=0), its K_ERROR, and no worker. */
2#define _XOPEN_SOURCE 700 /* sexc compiles tests without the package's cflags (reported) */
3#include "mesh/check.h"
4#include "mesh/inferd.h"
5#include "mesh/log.h"
6
7#include <arpa/inet.h>
8#include <netinet/in.h>
9#include <pthread.h>
10#include <stdio.h>
11#include <stdlib.h>
12#include <string.h>
13#include <sys/socket.h>
14#include <sys/stat.h>
15#include <unistd.h>
16
17static int lfd;
18static void *worker(void *arg) {
19 int fail = *(int *)arg;
20 int c = accept(lfd, NULL, NULL);
21 uint8_t k;
22 mesh_buf b = MESH_BUF_INIT;
23 char *err = NULL;
24 if (mesh_infer_read_frame(c, &k, &b, NULL, &err) == 0) {
25 mesh_infer_msg q, a;
26 mesh_infer_msg_decode(b.data, b.len, &q, &err);
27 mesh_infer_msg_init(&a);
28 mesh_infer_msg_set(&a, "gen_tokens", mesh_infer_msg_get(&q, "max"));
29 mesh_infer_msg_set(&a, "mode", mesh_infer_msg_get(&q, "mode"));
30 mesh_infer_msg_set_body(&a, fail ? " out of memory \n" : q.body);
31 mesh_buf e = MESH_BUF_INIT;
32 mesh_infer_msg_encode(&a, &e);
33 mesh_infer_write_frame(c, fail ? MESH_INFER_K_ERROR : MESH_INFER_K_RESULT, e.data, e.len);
34 mesh_buf_free(&e);
35 mesh_infer_msg_free(&q);
36 mesh_infer_msg_free(&a);
37 }
38 mesh_buf_free(&b);
39 close(c);
40 return NULL;
41}
42
43int main(void) {
44 mesh_log_set_fd(-1);
45 /* A catalog with gpt2 at param 0: its worker listens on the control port's base. */
46 char root[64], p[128];
47 snprintf(root, sizeof root, "/tmp/mesh-inferd-%d", (int)getpid());
48 mkdir(root, 0700);
49 snprintf(p, sizeof p, "%s/gpt2", root);
50 mkdir(p, 0700);
51 snprintf(p, sizeof p, "%s/gpt2/program.toml", root);
52 FILE *f = fopen(p, "w");
53 fputs("name = \"gpt2\"\nabi = 6\nbackends = [\"native\"]\n", f);
54 fclose(f);
55 lfd = socket(AF_INET, SOCK_STREAM, 0);
56 struct sockaddr_in sa = {0};
57 sa.sin_family = AF_INET;
58 sa.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
59 bind(lfd, (struct sockaddr *)&sa, sizeof sa);
60 socklen_t sl = sizeof sa;
61 getsockname(lfd, (struct sockaddr *)&sa, &sl);
62 listen(lfd, 4);
63 char ctl[64];
64 snprintf(ctl, sizeof ctl, "127.0.0.1:%u", ntohs(sa.sin_port));
65 static _Atomic bool quit;
66 mesh_mailbox mb = mesh_mailbox_test();
67 mesh_shared *sh = mesh_shared_new(mb, root, false, 0, 1, &quit);
68 char *err = NULL;
69 CHECK_I(mesh_infer_svc_single(&sh->svc, ctl, &err), 0);
70 sh->has_svc = true;
71 mesh_infer_request r;
72 CHECK_I(mesh_infer_parse("exact max 4 mono=0 Hello there", &r, &err), 0);
73 int fail = 0;
74 pthread_t th;
75 pthread_create(&th, NULL, worker, &fail);
76 mesh_infer_msg res;
77 CHECK_I(mesh_inferd_run(sh, &r, &res, &err), 0);
78 pthread_join(th, NULL);
79 CHECK_S(res.body, "Hello there");
80 CHECK_S(mesh_infer_msg_get(&res, "gen_tokens"), "4");
81 CHECK_S(mesh_infer_msg_get(&res, "program"), "gpt2");
82 CHECK_S(mesh_infer_msg_get(&res, "coordinator"), "0");
83 CHECK_S(mesh_infer_msg_get(&res, "decode_sys"), "0");
84 CHECK(mesh_infer_msg_get(&res, "request_ms") != NULL);
85 mesh_infer_msg_free(&res);
86 fail = 1;
87 pthread_create(&th, NULL, worker, &fail);
88 CHECK_I(mesh_inferd_run(sh, &r, &res, &err), -1);
89 pthread_join(th, NULL);
90 CHECK_S(err, "gpt2 on member 0: out of memory");
91 free(err);
92 close(lfd);
93 CHECK_I(mesh_inferd_run(sh, &r, &res, &err), -1);
94 CHECK(strstr(err, "gpt2 on member 0 (127.0.0.1:") == err);
95 CHECK(strstr(err, "Connection refused (os error 111)") != NULL);
96 free(err);
97 /* Nothing placed (single, not loaded): no candidates. */
98 mesh_infer_request_free(&r);
99 CHECK_I(mesh_infer_parse("hi", &r, &err), 0);
100 CHECK_I(mesh_inferd_run(sh, &r, &res, &err), -1);
101 CHECK_S(err, "no gpt2, gpt2-prefill or gpt2-decode is placed on any member");
102 free(err);
103 mesh_infer_request_free(&r);
104 mesh_shared_free(sh);
105 free(mb.base);
106 unlink(p);
107 snprintf(p, sizeof p, "%s/gpt2", root);
108 rmdir(p);
109 rmdir(root);
110 return mesh_check_done();
111}