mirror of
https://github.com/nestriness/nestri.git
synced 2026-11-03 22:22:39 +02:00
Wanjohi
5a399686ca
Create a room under a per-user advisory lock (limit, card and insert in one transaction) and book cards under per-machine locks. A free plan with no model is placed on an entry card only. Boxes with no card hold their machine's first card; the new migration backfills single-GPU fleet boxes and adds box.gpu_slot. Machine.Gpu accepts slot and renderNode; session.start carries gpuSlot and gpuPciId. A stopped room re-checks its card when a run is requested.
361 lines
14 KiB
TypeScript
361 lines
14 KiB
TypeScript
import { afterAll, describe, expect, test } from 'bun:test';
|
||
|
||
import { AccessToken } from '@nestri/core/access-token/index';
|
||
import { Box } from '@nestri/core/box/index';
|
||
import { Fixtures } from '@nestri/core/db/fixtures';
|
||
import { testDb } from '@nestri/core/db/test';
|
||
import { Env } from '@nestri/core/env';
|
||
import { Identifier } from '@nestri/core/id';
|
||
import { Machine } from '@nestri/core/machine/index';
|
||
import { Team } from '@nestri/core/team/index';
|
||
|
||
import { app } from '../app/index';
|
||
import { TEST_FRONTEND_URL } from './setup';
|
||
|
||
const sql = testDb();
|
||
const createdUserIds: string[] = [];
|
||
const createdOrgIds: string[] = [];
|
||
|
||
const V710 = { model: 'AMD Radeon Pro V710', vramMib: 28672, pciId: '1002:7461' };
|
||
const BIG = { model: 'NVIDIA H100', vramMib: 81920, pciId: '10de:2330' };
|
||
|
||
function useFleet(organisationId: string, entry?: string) {
|
||
Env.init({
|
||
NODE_ENV: 'test',
|
||
FRONTEND_URL: TEST_FRONTEND_URL,
|
||
FLEET_ORGANISATION_ID: organisationId,
|
||
...(entry !== undefined ? { FLEET_ENTRY_GPUS: entry } : {})
|
||
});
|
||
}
|
||
|
||
async function person(label: string, plan?: 'paid') {
|
||
const owner = await Fixtures.owner(label);
|
||
createdUserIds.push(owner.userId);
|
||
if (plan) {
|
||
await Team.setPlan({ id: owner.teamId, plan, subscriptionStatus: 'active' });
|
||
}
|
||
const pat = await AccessToken.create({
|
||
id: Identifier.ascending('accessToken'),
|
||
ownerUserId: owner.userId,
|
||
teamId: null,
|
||
name: label
|
||
});
|
||
return {
|
||
owner,
|
||
headers: {
|
||
authorization: `Bearer ${pat.token}`,
|
||
'content-type': 'application/json'
|
||
} as Record<string, string>
|
||
};
|
||
}
|
||
|
||
async function fleet(gpusPerMachine: (typeof V710)[][]) {
|
||
const operator = await Fixtures.owner('operator');
|
||
createdUserIds.push(operator.userId);
|
||
const f = await Fixtures.fleet(operator, gpusPerMachine.length);
|
||
createdOrgIds.push(f.organisationId);
|
||
for (const [i, gpus] of gpusPerMachine.entries()) {
|
||
await Machine.setGpus({ id: f.machines[i]!.id, gpus });
|
||
}
|
||
useFleet(f.organisationId);
|
||
return f;
|
||
}
|
||
|
||
const post = (headers: Record<string, string>, body: object) =>
|
||
app.request('/box', { method: 'POST', headers, body: JSON.stringify(body) });
|
||
const json = async (res: Response) => (await res.json()) as any;
|
||
|
||
afterAll(async () => {
|
||
await sql`delete from session where box_id in (select id from box where user_id = any(${createdUserIds}))`;
|
||
await sql`delete from box where user_id = any(${createdUserIds})`;
|
||
await sql`delete from install_token where created_by_user_id = any(${createdUserIds})`;
|
||
await sql`delete from machine where owner_user_id = any(${createdUserIds})`;
|
||
await sql`delete from organisation where id = any(${createdOrgIds})`;
|
||
await sql`delete from "user" where id = any(${createdUserIds})`;
|
||
});
|
||
|
||
describe('GET /location', () => {
|
||
test('lists your own GPUs one by one and the fleet grouped by model, per your plan', async () => {
|
||
await fleet([[V710], [V710], [BIG]]);
|
||
const free = await person('loc-free');
|
||
const host = await Fixtures.machine(free.owner, 'loc-host');
|
||
await Machine.setGpus({
|
||
id: host,
|
||
gpus: [
|
||
{ model: 'NVIDIA GeForce RTX 3060', vramMib: 12288, pciId: '10de:2503' },
|
||
{ model: 'AMD Radeon RX 9060 XT', vramMib: 16304, pciId: '1002:7590' }
|
||
]
|
||
});
|
||
|
||
const res = await app.request('/location', { headers: free.headers });
|
||
expect(res.status).toBe(200);
|
||
const { data } = await json(res);
|
||
expect(data.own.map((o: any) => [o.machineId, o.pciId, o.model, o.vramMib, o.online])).toEqual([
|
||
[host, '10de:2503', 'NVIDIA GeForce RTX 3060', 12288, false],
|
||
[host, '1002:7590', 'AMD Radeon RX 9060 XT', 16304, false]
|
||
]);
|
||
expect(data.fleet).toEqual([
|
||
{ model: V710.model, vramMib: 28672, available: true, entry: true, allowed: true },
|
||
{ model: BIG.model, vramMib: 81920, available: true, entry: false, allowed: false }
|
||
]);
|
||
|
||
const pro = await person('loc-pro', 'paid');
|
||
const proFleet = (await json(await app.request('/location', { headers: pro.headers }))).data
|
||
.fleet;
|
||
expect(proFleet.every((f: any) => f.allowed)).toBe(true);
|
||
expect(proFleet.find((f: any) => f.model === BIG.model).entry).toBe(false);
|
||
});
|
||
|
||
test('a model with every card taken is listed as not available', async () => {
|
||
await fleet([[V710]]);
|
||
const p = await person('loc-taken');
|
||
expect((await post(p.headers, { on: 'fleet', model: V710.model })).status).toBe(201);
|
||
const other = await person('loc-other', 'paid');
|
||
const { data } = await json(await app.request('/location', { headers: other.headers }));
|
||
expect(data.fleet[0].available).toBe(false);
|
||
});
|
||
});
|
||
|
||
describe('POST /box on a chosen GPU', () => {
|
||
test('an own room is pinned to the chosen GPU; several own hosts are fine', async () => {
|
||
const p = await person('own-gpu');
|
||
const a = await Fixtures.machine(p.owner, 'host-a');
|
||
const b = await Fixtures.machine(p.owner, 'host-b');
|
||
await Machine.setGpus({ id: a, gpus: [V710] });
|
||
await Machine.setGpus({ id: b, gpus: [BIG, V710] });
|
||
|
||
const res = await post(p.headers, { on: 'own', machineId: b, pciId: BIG.pciId });
|
||
expect(res.status).toBe(201);
|
||
const box = (await json(res)).data;
|
||
expect(box.machineId).toBe(b);
|
||
expect(box.gpuPciId).toBe(BIG.pciId);
|
||
expect((await Box.fromID(box.id))!.gpuPciId).toBe(BIG.pciId);
|
||
|
||
// Own rooms are never counted against a plan.
|
||
for (let i = 0; i < 4; i++) {
|
||
expect((await post(p.headers, { on: 'own', machineId: a, pciId: V710.pciId })).status).toBe(
|
||
201
|
||
);
|
||
}
|
||
});
|
||
|
||
test('an own room needs a GPU the host reports, on a host that is yours', async () => {
|
||
const p = await person('own-bad');
|
||
const stranger = await person('own-stranger');
|
||
const mine = await Fixtures.machine(p.owner, 'mine');
|
||
const theirs = await Fixtures.machine(stranger.owner, 'theirs');
|
||
await Machine.setGpus({ id: mine, gpus: [V710] });
|
||
await Machine.setGpus({ id: theirs, gpus: [V710] });
|
||
|
||
expect((await post(p.headers, { on: 'own', machineId: mine, pciId: BIG.pciId })).status).toBe(
|
||
404
|
||
);
|
||
expect((await post(p.headers, { on: 'own', machineId: theirs, pciId: V710.pciId })).status).toBe(
|
||
404
|
||
);
|
||
expect((await post(p.headers, { on: 'own', machineId: mine })).status).toBe(400);
|
||
});
|
||
|
||
test('the old body still works: one host, no GPU named', async () => {
|
||
const p = await person('old-body');
|
||
const host = await Fixtures.machine(p.owner, 'old');
|
||
const res = await post(p.headers, { label: 'den' });
|
||
expect(res.status).toBe(201);
|
||
const box = (await json(res)).data;
|
||
expect(box.machineId).toBe(host);
|
||
expect(box.gpuPciId).toBeNull();
|
||
});
|
||
|
||
test('a fleet room is placed by model on the least-loaded machine with one free', async () => {
|
||
const f = await fleet([[V710], [V710]]);
|
||
const a = await person('fleet-model-a');
|
||
const b = await person('fleet-model-b');
|
||
const boxA = (await json(await post(a.headers, { on: 'fleet', model: V710.model }))).data;
|
||
const boxB = (await json(await post(b.headers, { on: 'fleet', model: V710.model }))).data;
|
||
expect(boxA.gpuPciId).toBe(V710.pciId);
|
||
expect(new Set([boxA.machineId, boxB.machineId])).toEqual(new Set(f.machines.map((m) => m.id)));
|
||
|
||
// Both cards are now taken.
|
||
const c = await person('fleet-model-c');
|
||
expect((await post(c.headers, { on: 'fleet', model: V710.model })).status).toBe(429);
|
||
});
|
||
|
||
test('free: one cloud room, entry models only', async () => {
|
||
await fleet([[V710], [BIG], [V710]]);
|
||
const p = await person('free-limits');
|
||
|
||
const big = await post(p.headers, { on: 'fleet', model: BIG.model });
|
||
expect(big.status).toBe(403);
|
||
const refusal = await json(big);
|
||
expect(refusal.code).toBe('gpu_needs_pro');
|
||
expect(refusal.message).toBe(`Upgrade to create a room on ${BIG.model}`);
|
||
|
||
expect((await post(p.headers, { on: 'fleet', model: V710.model })).status).toBe(201);
|
||
const second = await post(p.headers, { on: 'fleet', model: V710.model });
|
||
expect(second.status).toBe(403);
|
||
const limit = await json(second);
|
||
expect(limit.code).toBe('room_limit');
|
||
expect(limit.message).toBe(`Upgrade to create a room on ${V710.model}`);
|
||
|
||
// Your own GPU is still open at the limit.
|
||
const host = await Fixtures.machine(p.owner, 'free-own');
|
||
await Machine.setGpus({ id: host, gpus: [BIG] });
|
||
expect((await post(p.headers, { on: 'own', machineId: host, pciId: BIG.pciId })).status).toBe(
|
||
201
|
||
);
|
||
});
|
||
|
||
test('pro: three cloud rooms on any model, then room_limit', async () => {
|
||
await fleet([[BIG], [BIG], [BIG], [BIG]]);
|
||
const p = await person('pro-limits', 'paid');
|
||
for (let i = 0; i < 3; i++) {
|
||
const res = await post(p.headers, { on: 'fleet', model: BIG.model });
|
||
expect(res.status).toBe(201);
|
||
}
|
||
const over = await post(p.headers, { on: 'fleet', model: BIG.model });
|
||
expect(over.status).toBe(403);
|
||
expect((await json(over)).code).toBe('room_limit');
|
||
|
||
const billing = await json(await app.request('/billing', { headers: p.headers }));
|
||
expect(billing.data.plan).toBe('paid');
|
||
expect(billing.data.rooms).toEqual({ used: 3, allowed: 3 });
|
||
});
|
||
|
||
test('the entry list is configuration', async () => {
|
||
const f = await fleet([[BIG]]);
|
||
useFleet(f.organisationId, 'h100');
|
||
const p = await person('entry-config');
|
||
expect((await post(p.headers, { on: 'fleet', model: BIG.model })).status).toBe(201);
|
||
});
|
||
});
|
||
|
||
describe('GET /billing', () => {
|
||
test('answers plan and rooms, with no windows', async () => {
|
||
await fleet([[V710]]);
|
||
const p = await person('billing-shape');
|
||
const { data } = await json(await app.request('/billing', { headers: p.headers }));
|
||
expect(Object.keys(data).sort()).toEqual(['plan', 'rooms', 'teamId']);
|
||
expect(data.rooms).toEqual({ used: 0, allowed: 1 });
|
||
});
|
||
});
|
||
|
||
describe('the idle limit', () => {
|
||
test('a cloud room with nothing live for 3 hours is stopped, and its record stays', async () => {
|
||
const f = await fleet([[V710]]);
|
||
const p = await person('idle');
|
||
const box = (await json(await post(p.headers, { on: 'fleet', model: V710.model }))).data;
|
||
await Box.setState({ id: box.id, state: 'running' });
|
||
|
||
// Fresh: left alone.
|
||
expect(await Box.stopIdle(f.machines[0]!.id)).toEqual([]);
|
||
|
||
await sql`update box set time_created = now() - interval '4 hours' where id = ${box.id}`;
|
||
expect(await Box.stopIdle(f.machines[0]!.id)).toEqual([box.id]);
|
||
const after = (await Box.fromID(box.id))!;
|
||
expect(after.state).toBe('stopped');
|
||
expect(after.stopReason).toBe(Box.IDLE_REASON);
|
||
});
|
||
|
||
test('your own hardware has no idle limit', async () => {
|
||
const p = await person('idle-own');
|
||
const host = await Fixtures.machine(p.owner, 'idle-own-host');
|
||
const box = (await json(await post(p.headers, { label: 'x' }))).data;
|
||
await Box.setState({ id: box.id, state: 'running' });
|
||
await sql`update box set time_created = now() - interval '9 hours' where id = ${box.id}`;
|
||
expect(await Box.stopIdle(host)).toEqual([]);
|
||
expect((await Box.fromID(box.id))!.state).toBe('running');
|
||
});
|
||
});
|
||
|
||
describe('review blockers on rooms', () => {
|
||
test('free with no model is placed on an entry card only', async () => {
|
||
await fleet([[BIG], [V710]]);
|
||
const p = await person('free-any');
|
||
const res = await post(p.headers, { on: 'fleet' });
|
||
expect(res.status).toBe(201);
|
||
expect((await json(res)).data.gpuPciId).toBe(V710.pciId);
|
||
|
||
const only = await fleet([[BIG]]);
|
||
const q = await person('free-none');
|
||
const refused = await post(q.headers, { on: 'fleet' });
|
||
expect(refused.status).toBe(429);
|
||
expect(only.machines.length).toBe(1);
|
||
});
|
||
|
||
test('two concurrent POST /box for a free user: exactly one room', async () => {
|
||
await fleet([[V710], [V710], [V710]]);
|
||
const p = await person('race-limit');
|
||
const res = await Promise.all([
|
||
post(p.headers, { on: 'fleet', model: V710.model }),
|
||
post(p.headers, { on: 'fleet', model: V710.model })
|
||
]);
|
||
expect(res.map((r) => r.status).sort()).toEqual([201, 403]);
|
||
expect(await Box.onFleet({ userId: p.owner.userId, organisationId: Env.get().FLEET_ORGANISATION_ID! }))
|
||
.toHaveLength(1);
|
||
});
|
||
|
||
test('two people racing for the last card: one gets it', async () => {
|
||
await fleet([[V710]]);
|
||
const a = await person('race-a');
|
||
const b = await person('race-b');
|
||
const res = await Promise.all([
|
||
post(a.headers, { on: 'fleet', model: V710.model }),
|
||
post(b.headers, { on: 'fleet', model: V710.model })
|
||
]);
|
||
expect(res.map((r) => r.status).sort()).toEqual([201, 429]);
|
||
});
|
||
|
||
test('a box with no card holds its machine’s first card; slots tell identical cards apart', async () => {
|
||
const f = await fleet([[V710]]);
|
||
const legacy = await person('legacy');
|
||
await Box.create({
|
||
id: Identifier.ascending('box'),
|
||
userId: legacy.owner.userId,
|
||
machineId: f.machines[0]!.id,
|
||
gpuPciId: null,
|
||
label: 'old',
|
||
tier: 'md'
|
||
});
|
||
const other = await person('legacy-other', 'paid');
|
||
expect((await post(other.headers, { on: 'fleet', model: V710.model })).status).toBe(429);
|
||
const loc = (await json(await app.request('/location', { headers: other.headers }))).data;
|
||
expect(loc.fleet[0].available).toBe(false);
|
||
|
||
const g = await fleet([[
|
||
{ ...V710, slot: '0000:45:00.0' },
|
||
{ ...V710, slot: '0000:46:00.0' }
|
||
]]);
|
||
const p = await person('slots', 'paid');
|
||
const one = (await json(await post(p.headers, { on: 'fleet', model: V710.model }))).data;
|
||
const two = (await json(await post(p.headers, { on: 'fleet', model: V710.model }))).data;
|
||
expect([one.gpuSlot, two.gpuSlot].sort()).toEqual(['0000:45:00.0', '0000:46:00.0']);
|
||
expect(g.machines.length).toBe(1);
|
||
const jobs = (await import('@nestri/core/session/index')).Session;
|
||
expect(jobs.StartJob.shape.gpuSlot).toBeDefined();
|
||
});
|
||
|
||
test('a stopped room whose card was taken is re-placed on the same model, or refused', async () => {
|
||
const f = await fleet([[V710], [V710]]);
|
||
const a = await person('reclaim-a');
|
||
const room = (await json(await post(a.headers, { on: 'fleet', model: V710.model }))).data;
|
||
await Box.setState({ id: room.id, state: 'stopped', stopReason: 'idle for 3 hours', stopClean: true });
|
||
const b = await person('reclaim-b');
|
||
const bBox = (await json(await post(b.headers, { on: 'fleet', model: V710.model }))).data;
|
||
const c = await person('reclaim-c', 'paid');
|
||
const cBox = (await json(await post(c.headers, { on: 'fleet', model: V710.model }))).data;
|
||
expect(cBox.machineId).toBe(room.machineId);
|
||
expect(bBox.machineId).not.toBe(room.machineId);
|
||
|
||
const { Placement } = await import('@nestri/core/box/placement');
|
||
await expect(Placement.reclaim(room.id)).rejects.toThrow(/free right now/);
|
||
expect((await Box.fromID(room.id))!.state).toBe('stopped');
|
||
|
||
await Box.setState({ id: bBox.id, state: 'stopped', stopReason: 'x', stopClean: true });
|
||
await Placement.reclaim(room.id);
|
||
const moved = (await Box.fromID(room.id))!;
|
||
expect(moved.machineId).toBe(bBox.machineId);
|
||
expect(moved.state).toBe('created');
|
||
expect(f.machines.length).toBe(2);
|
||
});
|
||
});
|