1/* jobs.rs's `select_is_the_schedulers_filter_and_score`: select with the fleet scheduler's model. */
2#include "mesh/jobs/sched.h"
3#include "mesh/check.h"
4#include "mesh/strbuf.h"
5#include "mesh/strmap.h"
6
7#include <stdlib.h>
8#include <string.h>
9
10/* The scheduler's model of n CUDA members running "p". */
11static mjs_cx cx(uint32_t n) {
12    mjs_cx c = {0};
13    c.members = mxalloc(n * sizeof *c.members);
14    c.nmembers = n;
15    for (uint32_t k = 0; k < n; k++) {
16        msm_member m = msm_member_default();
17        m.id = k;
18        m.up = true;
19        MSM_SET(m.gpu, "cuda");
20        MSM_SET(m.arch, "aarch64");
21        m.cpus = 8;
22        m.ram_mb = 1000;
23        m.mem_avail_mb = 800;
24        m.gpu_share = mof32_some(0.9f);
25        m.load_known = true;
26        c.members[k] = m;
27    }
28    c.need_names = mxalloc(sizeof *c.need_names);
29    c.need_names[0] = mstrdup("p");
30    c.needs = mxalloc(sizeof *c.needs);
31    c.needs[0] = msn_default();
32    c.nneeds = 1;
33    c.w = msw_default();
34    return c;
35}
36
37static mj_member_progs progs4(void) {
38    mj_member_progs p = {0};
39    for (uint32_t k = 0; k < 4; k++) {
40        mj_prog *x = mxalloc(sizeof *x);
41        *x = (mj_prog){mstrdup(k == 3 ? "q" : "p"), 0, 8};
42        mj_member_progs_set(&p, k, x, 1);
43    }
44    return p;
45}
46
47static void oks(const mj_picks *v, const char *want) {
48    char got[64] = "";
49    for (size_t i = 0; i < v->n; i++)
50        if (v->v[i].ok) {
51            char b[8];
52            snprintf(b, sizeof b, "%s%u", got[0] ? "," : "", v->v[i].member);
53            strcat(got, b);
54        }
55    CHECK_S(got, want);
56}
57
58static const char *why(const mj_picks *v, uint32_t k) {
59    for (size_t i = 0; i < v->n; i++)
60        if (v->v[i].member == k) return v->v[i].why ? v->v[i].why : "";
61    return "";
62}
63
64static mj_spec spec(uint32_t members) {
65    mj_setup su = {6, 0};
66    return (mj_spec){.id = 1, .program = mstrdup("p"), .op = 254, .items = mj_ranges_of(0, 100), .attempts = 3, .setup = mj_setup_clone(&su, 1),
67                     .nsetup = 1, .submitted_ms = 5, .members = members};
68}
69
70static void perf(mjs_cx *c, uint32_t k, float task_us) {
71    msm_perf p = {0};
72    strcpy(p.exec, "cuda");
73    p.run = true;
74    p.task_us = mof32_some(task_us);
75    msm_perfs_set(&c->perf, "p", k, &p);
76}
77
78int main(void) {
79    mj_member_progs progs = progs4();
80    mj_spec job = spec(0);
81    /* No fleet (a single meshd): alive and running it. */
82    mj_picks v = mj_select(&job, 0xf, &progs, NULL);
83    oks(&v, "0,1,2");
84    mj_picks_free(&v);
85    /* Cordoned, busy (an idle member runs it), not running it, down. */
86    mjs_cx c = cx(5);
87    c.members[1].cordoned = true;
88    c.members[2].busy = mstrdup("cpu psi 0.40");
89    mj_sched *s = mjs_sched(c);
90    v = mj_select(&job, 0x0f, &progs, s);
91    oks(&v, "0");
92    CHECK_S(why(&v, 1), "cordoned");
93    CHECK(strncmp(why(&v, 2), "busy (owner policy: cpu psi 0.40)", 33) == 0);
94    CHECK(strstr(why(&v, 3), "does not run") != NULL);
95    CHECK_S(why(&v, 4), "down");
96    char *l = mj_pick_line(&v.v[0]);
97    CHECK(strncmp(l, "pick system 0 eligible ", 23) == 0);
98    free(l);
99    mj_picks_free(&v);
100    /* The busy member alone: eligible, floor-sized leases. */
101    ((mjs_cx *)s->ctx)->members[0].cordoned = true;
102    v = mj_select(&job, 0xf, &progs, s);
103    oks(&v, "2");
104    l = mj_pick_line(&v.v[0]);
105    CHECK(v.v[0].busy && strstr(l, "floor-sized"));
106    free(l);
107    mj_picks_free(&v);
108    s->release(s->ctx);
109    free(s);
110    /* Backends, a failure there, the job's members. */
111    c = cx(3);
112    MSM_SET(c.members[1].gpu, "metal");
113    ms_set(&c.members[2].fails, "p", "trap", true);
114    s = mjs_sched(c);
115    v = mj_select(&job, 0x7, &progs, s);
116    oks(&v, "0");
117    CHECK(strncmp(why(&v, 1), "backends cuda", 13) == 0);
118    CHECK(strstr(why(&v, 2), "failed there: trap") != NULL);
119    mj_picks_free(&v);
120    s->release(s->ctx);
121    free(s);
122    mj_spec j2 = spec(0x2);
123    s = mjs_sched(cx(3));
124    v = mj_select(&j2, 0x7, &progs, s);
125    oks(&v, "1");
126    mj_picks_free(&v);
127    s->release(s->ctx);
128    free(s);
129    mj_spec_free(&j2);
130    /* Measured performance and share order them: the faster member first, ties to the lower id. */
131    c = cx(3);
132    perf(&c, 0, 400.0f);
133    perf(&c, 1, 200.0f);
134    perf(&c, 2, 100.0f);
135    s = mjs_sched(c);
136    v = mj_select(&job, 0x7, &progs, s);
137    oks(&v, "2,1,0");
138    mj_picks_free(&v);
139    ((mjs_cx *)s->ctx)->members[2].gpu_share = mof32_some(0.1f);
140    v = mj_select(&job, 0x7, &progs, s);
141    CHECK(v.v[0].ok && v.v[0].member == 1);
142    mj_picks_free(&v);
143    s->release(s->ctx);
144    free(s);
145    mj_spec_free(&job);
146    mj_member_progs_free(&progs);
147    return mesh_check_done();
148}