1/* sched_tests.rs needs_round_trip_and_terms, and the error messages. */
2#include "mesh/check.h"
3#include "mesh/sched_needs.h"
4
5#include <stdlib.h>
6
7static msm_member member(uint32_t id, const char *gpu, const char *host, const char *arch) {
8    msm_member m = msm_member_default();
9    m.id = id;
10    MSM_SET(m.gpu, gpu);
11    MSM_SET(m.host, host);
12    MSM_SET(m.arch, arch);
13    return m;
14}
15
16static mesh_term term(const char *s) {
17    mesh_term t;
18    char *err = NULL;
19    if (!mesh_term_parse(s, &t, &err)) {
20        fprintf(stderr, "%s\n", err);
21        free(err), err = NULL;
22        abort();
23    }
24    return t;
25}
26
27static void round_trip(const mesh_needs *n) {
28    char *e = msn_encode(n);
29    CHECK(strchr(e, ' ') == NULL);
30    mesh_needs d;
31    char *err = NULL;
32    CHECK(msn_decode(e, &d, &err));
33    CHECK(mesh_needs_eq(&d, n));
34    mesh_needs_free(&d);
35    free(e);
36}
37
38int main(void) {
39    mesh_needs n = msn_default();
40    mesh_strs_free(n.backends, n.nbackends), n.backends = NULL, n.nbackends = 0;
41    msn_strs_push(&n.backends, &n.nbackends, "cuda");
42    msn_strs_push(&n.backends, &n.nbackends, "wasm");
43    n.arch = mt_strdup("aarch64");
44    n.min_mem_mb = 4096;
45    msn_terms_push(&n.prefer, &n.nprefer, term("cuda"));
46    msn_terms_push(&n.prefer, &n.nprefer, term("host:dgx-00"));
47    msn_terms_push(&n.avoid, &n.navoid, term("system:3"));
48    n.replicas = (mesh_replicas){false, 2};
49    n.pack = true;
50    n.colocate_with = mt_strdup("pipehash");
51    msn_strs_push(&n.feeds, &n.nfeeds, "a");
52    n.group = mt_strdup("g");
53    n.group_size = 3;
54    msn_terms_push(&n.group_members, &n.ngroup_members, term("arch:aarch64"));
55    n.priority = -2;
56    n.cpu = 0.5f;
57    n.mem_mb = 64;
58    n.migrate = MESH_MIGRATE_RESTART;
59    n.lockstep = true;
60    n.view_bytes = 0x10000;
61    n.warps = 8;
62    n.shared = 0x3000;
63    round_trip(&n);
64    char *e = msn_encode(&n);
65    CHECK_S(e, "b=cuda,wasm;ar=aarch64;mm=4096;pf=cuda,host:dgx-00;av=system:3;r=2;sp=pack;co=pipehash;fd=a;g=g;gs=3;gm=arch:aarch64;pr=-2;cpu=0.5;mem=64;mg=restart;ls=1;vb=65536;w=8;sh=12288");
66    free(e);
67    char *s = msn_summary(&n);
68    CHECK_S(s, "replicas 2 pack backends cuda,wasm arch aarch64 min_mem_mb 4096 prefer cuda,host:dgx-00 avoid system:3 colocate_with pipehash group g size 3 members arch:aarch64 priority -2 cpu 0.50 mem_mb 64 view 64 KiB node 8 warps, 12288 shared bytes migrate restart");
69    free(s);
70    mesh_needs_free(&n);
71    n = msn_default();
72    round_trip(&n);
73    e = msn_encode(&n);
74    CHECK_S(e, "b=cuda;ar=-;mm=0;pf=-;av=-;r=all;sp=spread;co=-;fd=-;g=-;gs=0;gm=-;pr=0;cpu=0;mem=0;mg=state;ls=0;vb=0;w=1;sh=0");
75    free(e);
76    s = msn_summary(&n);
77    CHECK_S(s, "replicas all backends cuda view 0 KiB migrate state");
78    free(s);
79    mesh_needs_free(&n);
80    char *err = NULL;
81    mesh_needs d;
82    CHECK(!msn_decode("cpu=nan", &d, &err));
83    CHECK_S(err, "needs: cpu=\"nan\"");
84    free(err), err = NULL;
85    CHECK(!msn_decode("zz=1", &d, &err));
86    CHECK_S(err, "needs: unknown key \"zz\"");
87    free(err), err = NULL;
88    CHECK(!msn_decode("w=0", &d, &err));
89    CHECK_S(err, "needs: w=\"0\"");
90    free(err), err = NULL;
91    CHECK(!msn_decode("b", &d, &err));
92    CHECK_S(err, "needs: \"b\"");
93    free(err), err = NULL;
94    mesh_term t;
95    CHECK(!mesh_term_parse("host:", &t, &err));
96    CHECK_S(err, "\"host:\": a backend (cuda, metal, wasm, native), \"host:NAME\", \"system:K\" or \"arch:A\"");
97    free(err), err = NULL;
98    CHECK(!mesh_term_parse("gpu", &t, &err));
99    free(err), err = NULL;
100    CHECK(!mesh_term_parse("system:9", &t, &err));
101    free(err), err = NULL;
102    msm_member mac = member(3, "metal", "monumentals-mac-studio.local", "aarch64");
103    msm_member dgx = member(0, "cuda", "dgx-00", "aarch64");
104    msm_member being = member(2, "cuda", "being", "x86_64");
105    t = term("host:monumentals-mac-studio");
106    CHECK(msn_term_matches(&t, &mac));
107    free(t.s);
108    t = term("metal");
109    CHECK(msn_term_matches(&t, &mac) && !msn_term_matches(&t, &dgx));
110    free(t.s);
111    t = term("arch:x86_64");
112    CHECK(msn_term_matches(&t, &being));
113    free(t.s);
114    t = term("system:2");
115    CHECK(msn_term_matches(&t, &being) && !msn_term_matches(&t, &mac));
116    char *w = mesh_term_word(&t);
117    CHECK_S(w, "system:2");
118    free(w);
119    return mesh_check_done();
120}