1/* @ajax/mesh-infer-wire: the wire of the inference services (programs/gpt2;
2 * docs/plan-disagg-inference.md).
3 *
4 * Every gpt2 worker listens on TCP; frames are a u32 length (LE, of what
5 * follows), a u8 kind, a payload. A request (K_GEN, K_DECODE, K_PREFILL)
6 * carries a Msg: `key=value` lines (sorted by key), a blank line, a body. KV
7 * travels as K_LAYER frames (layer, rows, a checksum, the bytes); K_FIRST
8 * ends a prefill; K_RESULT and K_ERROR answer. */
9#ifndef MESH_INFER_WIRE_H
10#define MESH_INFER_WIRE_H
11
12#include "mesh/buf.h"
13
14#include <stddef.h>
15#include <stdint.h>
16
17#define MESH_INFER_K_GEN 1
18#define MESH_INFER_K_DECODE 2
19#define MESH_INFER_K_PREFILL 3
20#define MESH_INFER_K_LAYER 10
21#define MESH_INFER_K_FIRST 11
22#define MESH_INFER_K_RESULT 20
23#define MESH_INFER_K_ERROR 21
24/* A frame's payload, at most (a 1024-token f32 layer is 6.3 MB). */
25#define MESH_INFER_MAX_FRAME (64u << 20)
26/* How long meshd waits for a request's answer (exact-cpu on a long prompt is slow): 600 s. */
27#define MESH_INFER_REQUEST_WAIT_MS 600000
28
29#define MESH_INFER_PROGRAM_MONO "gpt2"
30#define MESH_INFER_PROGRAM_PREFILL "gpt2-prefill"
31#define MESH_INFER_PROGRAM_DECODE "gpt2-decode"
32
33/* Writes one frame to fd (whole; EINTR retried): 0, or -1 with errno set. */
34int mesh_infer_write_frame(int fd, uint8_t kind, const void *body, size_t len);
35/* Reads one frame: 0 with *kind and body (appended), and (when wire_us) the
36 * microseconds the payload took after its header; -1 with *err ("read: ..."). */
37int mesh_infer_read_frame(int fd, uint8_t *kind, mesh_buf *body, int64_t *wire_us, char **err);
38
39/* key=value lines (sorted, one per key), a blank line, a body. */
40typedef struct {
41    char **k, **v;
42    size_t n, cap;
43    char *body;
44} mesh_infer_msg;
45
46void mesh_infer_msg_init(mesh_infer_msg *m);
47void mesh_infer_msg_free(mesh_infer_msg *m);
48void mesh_infer_msg_set(mesh_infer_msg *m, const char *k, const char *v);
49void mesh_infer_msg_setf(mesh_infer_msg *m, const char *k, const char *fmt, ...) __attribute__((format(printf, 3, 4)));
50void mesh_infer_msg_set_body(mesh_infer_msg *m, const char *body);
51const char *mesh_infer_msg_get(const mesh_infer_msg *m, const char *k);
52/* Missing / not-a-number errors as the Rust: "missing k", "k: not a number". */
53int mesh_infer_msg_num(const mesh_infer_msg *m, const char *k, uint64_t *out, char **err);
54void mesh_infer_msg_encode(const mesh_infer_msg *m, mesh_buf *out);
55/* 0, or -1 with *err: "a message that is not UTF-8", "a header line without '=': \"..\"". */
56int mesh_infer_msg_decode(const void *b, size_t n, mesh_infer_msg *out, char **err);
57
58/* A layer's checksum: the wrapping sum of its little-endian u64 words (a short tail zero-padded). */
59uint64_t mesh_infer_layer_sum(const void *b, size_t n);
60/* A K_LAYER payload: layer, rows, checksum, the bytes (appended to out). */
61void mesh_infer_layer_frame(mesh_buf *out, uint32_t layer, uint32_t rows, const void *bytes, size_t n);
62/* A K_LAYER payload's layer, rows and bytes (pointing into b), its checksum verified. */
63int mesh_infer_parse_layer(const void *b, size_t n, uint32_t *layer, uint32_t *rows, const uint8_t **bytes, size_t *len,
64                           char **err);
65
66#endif