Skip to content

Commit

Permalink
feat: Add VRAM and disk space matching (#471)
Browse files Browse the repository at this point in the history 
* test: Add VRAM match tests

* test: Add disk space match tests

* feat: Add vramMismatch match result

* feat: Add diskSpaceMismatch match result

* feat: Add match on VRAM

* feat: Add match on disk space

* test: Make error reporting more verbose
  • Loading branch information
@bgins
bgins authored Dec 18, 2024
1 parent 0050f2c commit 552d8e1
Show file tree
Hide file tree
Showing 2 changed files with 193 additions and 4 deletions.
82 changes: 82 additions & 0 deletions pkg/solver/matcher/match.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,50 @@ func (result ramMismatch) attributes() []attribute.KeyValue {
}
}

type vramMismatch struct {
resourceOffer data.ResourceOffer
jobOffer data.JobOffer
}

func (_ vramMismatch) matched() bool { return false }
func (_ vramMismatch) message() string { return "did not match VRAM" }
func (result vramMismatch) attributes() []attribute.KeyValue {
var resourceOfferVRAMs []int
for _, gpu := range result.resourceOffer.Spec.GPUs {
resourceOfferVRAMs = append(resourceOfferVRAMs, gpu.VRAM)
}

var jobOfferVRAMS []int
for _, gpu := range result.jobOffer.Spec.GPUs {
jobOfferVRAMS = append(jobOfferVRAMS, gpu.VRAM)
}

return []attribute.KeyValue{
attribute.String("match_result", fmt.Sprintf("%T", result)),
attribute.Bool("match_result.matched", result.matched()),
attribute.String("match_result.message", result.message()),
attribute.IntSlice("match_result.resource_offer.spec.gpus.vram", resourceOfferVRAMs),
attribute.IntSlice("match_result.job_offer.spec.gpus.vram", jobOfferVRAMS),
}
}

type diskSpaceMismatch struct {
resourceOffer data.ResourceOffer
jobOffer data.JobOffer
}

func (_ diskSpaceMismatch) matched() bool { return false }
func (_ diskSpaceMismatch) message() string { return "did not match disk space" }
func (result diskSpaceMismatch) attributes() []attribute.KeyValue {
return []attribute.KeyValue{
attribute.String("match_result", fmt.Sprintf("%T", result)),
attribute.Bool("match_result.matched", result.matched()),
attribute.String("match_result.message", result.message()),
attribute.Int("match_result.job_offer.spec.disk", result.jobOffer.Spec.Disk),
attribute.Int("match_result.resource_offer.spec.disk", result.resourceOffer.Spec.Disk),
}
}

type moduleIDError struct {
resourceOffer data.ResourceOffer
jobOffer data.JobOffer
Expand Down Expand Up @@ -227,6 +271,34 @@ func matchOffers(
}
}

// Skip VRAM check when job offer does not request VRAM
if len(jobOffer.Spec.GPUs) > 0 {
// Mismatch if job offer requests VRAM but resource provider has none
if len(resourceOffer.Spec.GPUs) == 0 {
return &vramMismatch{
jobOffer: jobOffer,
resourceOffer: resourceOffer,
}
}

// Mismatch if job offer largest VRAM is greater than resource offer largest VRAM
largestResourceOfferVRAM := getLargestVRAM(resourceOffer.Spec.GPUs)
largestJobOfferVRAM := getLargestVRAM(jobOffer.Spec.GPUs)
if largestResourceOfferVRAM < largestJobOfferVRAM {
return &vramMismatch{
jobOffer: jobOffer,
resourceOffer: resourceOffer,
}
}
}

if resourceOffer.Spec.Disk < jobOffer.Spec.Disk {
return &diskSpaceMismatch{
jobOffer: jobOffer,
resourceOffer: resourceOffer,
}
}

moduleID, err := data.GetModuleID(jobOffer.Module)
if err != nil {
return &moduleIDError{
Expand Down Expand Up @@ -295,6 +367,16 @@ func matchOffers(
}
}

func getLargestVRAM(gpus []data.GPUSpec) int {
largestVRAM := 0
for _, gpu := range gpus {
if gpu.VRAM > largestVRAM {
largestVRAM = gpu.VRAM
}
}
return largestVRAM
}

func logMatch(result matchResult) {
switch r := result.(type) {
case offersMatched:
Expand Down
115 changes: 111 additions & 4 deletions pkg/solver/matcher/matcher_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,14 @@ func TestMatchOffers(t *testing.T) {
CPU: 1000,
GPU: 1000,
RAM: 1024,
GPUs: []data.GPUSpec{
{
Name: "NVIDIA RTX 4090",
Vendor: "NVIDIA",
VRAM: 24576, // MB
},
},
Disk: 2924295844659, // Bytes
},
DefaultPricing: data.DealPricing{
InstructionPrice: 10,
Expand All @@ -29,9 +37,11 @@ func TestMatchOffers(t *testing.T) {

basicJobOffer := data.JobOffer{
Spec: data.MachineSpec{
CPU: 1000,
GPU: 1000,
RAM: 1024,
CPU: 1000,
GPU: 1000,
RAM: 1024,
GPUs: []data.GPUSpec{},
Disk: 0, // Bytes
},
Mode: data.MarketPrice,
Services: services,
Expand All @@ -51,6 +61,13 @@ func TestMatchOffers(t *testing.T) {
Path: "/lilypad_module.json.tmpl",
}

sdxlModuleConfig := data.ModuleConfig{
Name: "",
Repo: "https://github.com/Lilypad-Tech/lilypad-module-sdxl",
Hash: "v0.9-lilypad1",
Path: "/lilypad_module.json.tmpl",
}

testCases := []struct {
name string
resourceOffer func(offer data.ResourceOffer) data.ResourceOffer
Expand Down Expand Up @@ -100,6 +117,96 @@ func TestMatchOffers(t *testing.T) {
},
shouldMatch: false,
},
{
name: "VRAM match when job creator specifies VRAM",
resourceOffer: func(offer data.ResourceOffer) data.ResourceOffer {
offer.Spec.GPUs = []data.GPUSpec{{VRAM: 24576}}
return offer
},
jobOffer: func(offer data.JobOffer) data.JobOffer {
offer.Module = sdxlModuleConfig
offer.Spec.GPUs = []data.GPUSpec{{VRAM: 20000}}
return offer
},
shouldMatch: true,
},
{
name: "VRAM match when job creator does not specify VRAM",
resourceOffer: func(offer data.ResourceOffer) data.ResourceOffer {
offer.Spec.GPUs = []data.GPUSpec{{VRAM: 24576}}
return offer
},
jobOffer: func(offer data.JobOffer) data.JobOffer {
offer.Module = sdxlModuleConfig
offer.Spec.GPUs = []data.GPUSpec{}
return offer
},
shouldMatch: true,
},
{
name: "VRAM requested is more than resource offer VRAM",
resourceOffer: func(offer data.ResourceOffer) data.ResourceOffer {
offer.Spec.GPUs = []data.GPUSpec{{VRAM: 24576}}
return offer
},
jobOffer: func(offer data.JobOffer) data.JobOffer {
offer.Module = sdxlModuleConfig
offer.Spec.GPUs = []data.GPUSpec{{VRAM: 40000}}
return offer
},
shouldMatch: false,
},
{
name: "VRAM requested but resource offer has none",
resourceOffer: func(offer data.ResourceOffer) data.ResourceOffer {
offer.Spec.GPUs = []data.GPUSpec{}
return offer
},
jobOffer: func(offer data.JobOffer) data.JobOffer {
offer.Module = sdxlModuleConfig
offer.Spec.GPUs = []data.GPUSpec{{VRAM: 49152}}
return offer
},
shouldMatch: false,
},
{
name: "Disk space match when job creator specifies disk space",
resourceOffer: func(offer data.ResourceOffer) data.ResourceOffer {
offer.Spec.Disk = 2924295844659 // Bytes
return offer
},
jobOffer: func(offer data.JobOffer) data.JobOffer {
offer.Module = sdxlModuleConfig
offer.Spec.Disk = 1000000000000
return offer
},
shouldMatch: true,
},
{
name: "Disk space match when job creator does not specify disk space",
resourceOffer: func(offer data.ResourceOffer) data.ResourceOffer {
offer.Spec.Disk = 2924295844659 // Bytes
return offer
},
jobOffer: func(offer data.JobOffer) data.JobOffer {
offer.Module = sdxlModuleConfig
offer.Spec.Disk = 0 // zero-value
return offer
},
shouldMatch: true,
},
{
name: "Disk space requested is more than resource offer disk space",
resourceOffer: func(offer data.ResourceOffer) data.ResourceOffer {
offer.Spec.Disk = 2924295844659 // Bytes
return offer
},
jobOffer: func(offer data.JobOffer) data.JobOffer {
offer.Spec.Disk = 4000000000000
return offer
},
shouldMatch: false,
},
{
name: "Resource provider supports module",
resourceOffer: func(offer data.ResourceOffer) data.ResourceOffer {
Expand Down Expand Up @@ -212,7 +319,7 @@ func TestMatchOffers(t *testing.T) {
t.Run(tc.name, func(t *testing.T) {
result := matchOffers(tc.resourceOffer(basicResourceOffer), tc.jobOffer(basicJobOffer))
if result.matched() != tc.shouldMatch {
t.Errorf("Expected match to be %v, but got %v", tc.shouldMatch, result)
t.Errorf("Expected match to be %v, but got %+v", tc.shouldMatch, result)
}
})
}
Expand Down

0 comments on commit 552d8e1

Please sign in to comment.