diff --git a/apps/provider-inventory/src/repositories/bid-screening/bid-screening.aggregator.spec.ts b/apps/provider-inventory/src/repositories/bid-screening/bid-screening.aggregator.spec.ts index 80c4cc05a..1a700986a 100644 --- a/apps/provider-inventory/src/repositories/bid-screening/bid-screening.aggregator.spec.ts +++ b/apps/provider-inventory/src/repositories/bid-screening/bid-screening.aggregator.spec.ts @@ -183,6 +183,123 @@ describe(aggregateCriteria.name, () => { const c = aggregateCriteria([makeUnit({ gpu: 0n })], makeRequirements()); expect(c.units[0].gpuTokens).toEqual([]); }); + + it("emits the single declared class for a unit with one persistent volume", () => { + const c = aggregateCriteria( + [ + makeUnit({ + storage: [ + { + name: "data", + quantity: 1000n, + attributes: [ + { key: "persistent", value: "true" }, + { key: "class", value: "beta2" } + ] + } + ] + }) + ], + makeRequirements() + ); + expect(c.units[0].persistentClasses).toEqual(["beta2"]); + }); + + it("emits every distinct class declared across a unit's persistent volumes", () => { + const c = aggregateCriteria( + [ + makeUnit({ + storage: [ + { + name: "data", + quantity: 1000n, + attributes: [ + { key: "persistent", value: "true" }, + { key: "class", value: "beta2" } + ] + }, + { + name: "logs", + quantity: 500n, + attributes: [ + { key: "persistent", value: "true" }, + { key: "class", value: "beta3" } + ] + } + ] + }) + ], + makeRequirements() + ); + expect(c.units[0].persistentClasses).toEqual(["beta2", "beta3"]); + }); + + it("dedupes repeated classes across persistent volumes within a unit", () => { + const c = aggregateCriteria( + [ + makeUnit({ + storage: [ + { + name: "data", + quantity: 1000n, + attributes: [ + { key: "persistent", value: "true" }, + { key: "class", value: "beta2" } + ] + }, + { + name: "extra", + quantity: 250n, + attributes: [ + { key: "persistent", value: "true" }, + { key: "class", value: "beta2" } + ] + } + ] + }) + ], + makeRequirements() + ); + expect(c.units[0].persistentClasses).toEqual(["beta2"]); + }); + + it("ignores ephemeral and ram volumes when collecting persistentClasses", () => { + const c = aggregateCriteria( + [ + makeUnit({ + storage: [ + { name: "scratch", quantity: 500n, attributes: [] }, + { name: "shm", quantity: 100n, attributes: [{ key: "class", value: "ram" }] } + ] + }) + ], + makeRequirements() + ); + expect(c.units[0].persistentClasses).toEqual([]); + }); + + it("emits per-unit persistentClasses independently across units in a mixed deployment", () => { + const c = aggregateCriteria( + [ + makeUnit({ + storage: [ + { + name: "data", + quantity: 1000n, + attributes: [ + { key: "persistent", value: "true" }, + { key: "class", value: "beta2" } + ] + } + ] + }), + makeUnit({ storage: [{ name: "scratch", quantity: 500n, attributes: [] }] }) + ], + makeRequirements() + ); + expect(c.units[0].persistentClasses).toEqual(["beta2"]); + expect(c.units[1].persistentClasses).toEqual([]); + }); }); }); diff --git a/apps/provider-inventory/src/repositories/bid-screening/bid-screening.aggregator.ts b/apps/provider-inventory/src/repositories/bid-screening/bid-screening.aggregator.ts index 6d30ab0af..82ae673e5 100644 --- a/apps/provider-inventory/src/repositories/bid-screening/bid-screening.aggregator.ts +++ b/apps/provider-inventory/src/repositories/bid-screening/bid-screening.aggregator.ts @@ -1,7 +1,7 @@ import { parseGPUAttributes } from "@src/lib/gpu-attribute-parser/gpu-attribute-parser"; import type { GroupSpecJSON } from "@src/lib/groupspec-mapper/groupspec-mapper"; import { parseStorageAttributes } from "@src/lib/storage-attribute-parser/storage-attribute-parser"; -import type { RequestedResourceUnit, ResourceAttribute } from "@src/types/inventory.types"; +import type { RequestedResourceUnit, RequestedStorage, ResourceAttribute } from "@src/types/inventory.types"; interface UnitFilters { gpuTokens: string[]; @@ -58,7 +58,10 @@ export function aggregateCriteria(resourceUnits: RequestedResourceUnit[], requir // ram volumes intentionally skipped — issue 4 will add them to totalMemory } - units.push({ gpuTokens: gpuTokensForUnit(unit.resources.gpu), persistentClasses: [] }); + units.push({ + gpuTokens: collectGpuTokens(unit.resources.gpu), + persistentClasses: collectPersistentStorageTokens(unit.resources.storage) + }); } const attributes: BidScreeningCriteria["attributes"] = []; @@ -98,7 +101,7 @@ function escapeRegex(input: string): string { return input.replace(/[\\.^$*+?()[\]{}|]/g, "\\$&"); } -function gpuTokensForUnit(gpu: { units: bigint; attributes: ResourceAttribute[] }): string[] { +function collectGpuTokens(gpu: { units: bigint; attributes: ResourceAttribute[] }): string[] { if (gpu.units === 0n) return []; const tokens: string[] = []; for (const parsed of parseGPUAttributes(gpu.attributes)) { @@ -107,3 +110,14 @@ function gpuTokensForUnit(gpu: { units: bigint; attributes: ResourceAttribute[] } return tokens; } + +function collectPersistentStorageTokens(storage: RequestedStorage[]): string[] { + const tokens: string[] = []; + for (const vol of storage) { + const parsed = parseStorageAttributes(vol.attributes); + if (parsed.classification === "persistent" && parsed.class && !tokens.includes(parsed.class)) { + tokens.push(parsed.class); + } + } + return tokens; +} diff --git a/apps/provider-inventory/src/repositories/bid-screening/bid-screening.repository.integration.ts b/apps/provider-inventory/src/repositories/bid-screening/bid-screening.repository.integration.ts index 216857a5d..6fd60dfc6 100644 --- a/apps/provider-inventory/src/repositories/bid-screening/bid-screening.repository.integration.ts +++ b/apps/provider-inventory/src/repositories/bid-screening/bid-screening.repository.integration.ts @@ -8,7 +8,7 @@ import type { GroupSpecJSON } from "@src/lib/groupspec-mapper/groupspec-mapper"; import { ResourcePair } from "@src/lib/resource-pair/resource-pair"; import { providerInventory } from "@src/model-schemas/provider-inventory/provider-inventory.schema"; import { DRIZZLE_DB } from "@src/providers/drizzle.provider"; -import type { RequestedResourceUnit, ResourceAttribute } from "@src/types/inventory.types"; +import type { RequestedResourceUnit, RequestedStorage, ResourceAttribute } from "@src/types/inventory.types"; import { AUDITOR, BidScreeningRepository } from "./bid-screening.repository"; describe(BidScreeningRepository.name, () => { @@ -216,6 +216,83 @@ describe(BidScreeningRepository.name, () => { }); }); + describe("storage_classes filter", () => { + it("requires the provider to declare the unit's persistent class via @>", async () => { + await seed({ + owner: "akash1beta2", + storageClasses: ["beta2"], + totalAvailablePersistent: 10_000n + }); + await seed({ + owner: "akash1beta3", + storageClasses: ["beta3"], + totalAvailablePersistent: 10_000n + }); + + const rows = await repository.findCandidates([unit({ storage: [persistentVolume("data", 1_000n, "beta2")] })], requirements()); + + expect(owners(rows)).toEqual(["akash1beta2"]); + }); + + it("requires the provider to declare every persistent class on the unit (containment, not overlap)", async () => { + await seed({ + owner: "akash1beta2Only", + storageClasses: ["beta2"], + totalAvailablePersistent: 10_000n + }); + await seed({ + owner: "akash1beta3Only", + storageClasses: ["beta3"], + totalAvailablePersistent: 10_000n + }); + await seed({ + owner: "akash1both", + storageClasses: ["beta2", "beta3"], + totalAvailablePersistent: 10_000n + }); + + const rows = await repository.findCandidates( + [ + unit({ + storage: [persistentVolume("data", 1_000n, "beta2"), persistentVolume("logs", 500n, "beta3")] + }) + ], + requirements() + ); + + expect(owners(rows)).toEqual(["akash1both"]); + }); + + it("omits the clause entirely for units without persistent storage, so providers with no storage classes stay in the result", async () => { + await seed({ owner: "akash1noStorage", storageClasses: [] }); + await seed({ owner: "akash1withStorage", storageClasses: ["beta2"] }); + + const rows = await repository.findCandidates([unit({ storage: [{ name: "scratch", quantity: 500n, attributes: [] }] })], requirements()); + + expect(owners(rows)).toEqual(["akash1noStorage", "akash1withStorage"]); + }); + + it("emits a clause only for units with persistent volumes in a mixed deployment", async () => { + await seed({ + owner: "akash1beta2", + storageClasses: ["beta2"], + totalAvailablePersistent: 10_000n + }); + await seed({ + owner: "akash1noStorage", + storageClasses: [], + totalAvailablePersistent: 0n + }); + + const rows = await repository.findCandidates( + [unit({ storage: [persistentVolume("data", 1_000n, "beta2")] }), unit({ storage: [{ name: "scratch", quantity: 500n, attributes: [] }] })], + requirements() + ); + + expect(owners(rows)).toEqual(["akash1beta2"]); + }); + }); + describe("online filter", () => { it("excludes rows where is_online is false", async () => { await seed({ owner: "akash1up" }); @@ -282,6 +359,7 @@ describe(BidScreeningRepository.name, () => { selfAttributes?: { key: string; value: string }[]; auditedBy?: string[]; gpuModels?: string[]; + storageClasses?: string[]; inventory?: unknown; } @@ -302,12 +380,20 @@ describe(BidScreeningRepository.name, () => { selfAttributes: input.selfAttributes ?? [], auditedBy: input.auditedBy ?? [], gpuModels: input.gpuModels ?? [], + storageClasses: input.storageClasses ?? [], inventory: input.inventory ?? { nodes: [], storage: {} } }); } }); -function unit(input: { cpu?: bigint; memory?: bigint; gpu?: bigint; count?: number; gpuAttributes?: ResourceAttribute[] }): RequestedResourceUnit { +function unit(input: { + cpu?: bigint; + memory?: bigint; + gpu?: bigint; + count?: number; + gpuAttributes?: ResourceAttribute[]; + storage?: RequestedStorage[]; +}): RequestedResourceUnit { return { id: 1, count: input.count ?? 1, @@ -315,11 +401,22 @@ function unit(input: { cpu?: bigint; memory?: bigint; gpu?: bigint; count?: numb cpu: { units: input.cpu ?? 0n, attributes: [] }, memory: { quantity: input.memory ?? 0n, attributes: [] }, gpu: { units: input.gpu ?? 0n, attributes: input.gpuAttributes ?? [] }, - storage: [] + storage: input.storage ?? [] } }; } +function persistentVolume(name: string, quantity: bigint, storageClass: string): RequestedStorage { + return { + name, + quantity, + attributes: [ + { key: "persistent", value: "true" }, + { key: "class", value: storageClass } + ] + }; +} + function requirements(input?: Partial): GroupSpecJSON["requirements"] { return { signedBy: input?.signedBy ?? { allOf: [], anyOf: [] }, diff --git a/apps/provider-inventory/src/repositories/bid-screening/bid-screening.repository.ts b/apps/provider-inventory/src/repositories/bid-screening/bid-screening.repository.ts index 37d2d941a..1e6707667 100644 --- a/apps/provider-inventory/src/repositories/bid-screening/bid-screening.repository.ts +++ b/apps/provider-inventory/src/repositories/bid-screening/bid-screening.repository.ts @@ -71,6 +71,10 @@ export class BidScreeningRepository { if (unit.gpuTokens.length > 0) { conditions.push(arrayOverlaps(providerInventory.gpuModels, unit.gpuTokens)); } + + if (unit.persistentClasses.length > 0) { + conditions.push(arrayContains(providerInventory.storageClasses, unit.persistentClasses)); + } } if (criteria.attributes.length > 0) {