Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions protocol/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -323,6 +323,29 @@ type ResourceLimits struct {
OutputBytes int64 `json:"output_bytes,omitempty"`
}

type ResourceCapacity struct {
CPUCount int `json:"cpu_count,omitempty"`
MemoryBytes int64 `json:"memory_bytes,omitempty"`
DiskBytes int64 `json:"disk_bytes,omitempty"`
}

func ValidateResourceLimitsAgainstCapacity(limits ResourceLimits, capacity ResourceCapacity) error {
var errs []error
if limits.CPUPercent > 0 && capacity.CPUCount > 0 {
cpuCapacity := capacity.CPUCount * 100
if limits.CPUPercent > cpuCapacity {
errs = append(errs, fmt.Errorf("resource_limits.cpu_percent %d exceeds worker CPU capacity %d", limits.CPUPercent, cpuCapacity))
}
}
if limits.MemoryBytes > 0 && capacity.MemoryBytes > 0 && limits.MemoryBytes > capacity.MemoryBytes {
errs = append(errs, fmt.Errorf("resource_limits.memory_bytes %d exceeds worker memory_bytes %d", limits.MemoryBytes, capacity.MemoryBytes))
}
if limits.WorkspaceBytes > 0 && capacity.DiskBytes > 0 && limits.WorkspaceBytes > capacity.DiskBytes {
errs = append(errs, fmt.Errorf("resource_limits.workspace_bytes %d exceeds worker disk_bytes %d", limits.WorkspaceBytes, capacity.DiskBytes))
}
return errors.Join(errs...)
}

func (l ResourceLimits) Validate() error {
var errs []error
if l.CPUPercent < 0 {
Expand Down
39 changes: 39 additions & 0 deletions protocol/types_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -433,6 +433,45 @@ func TestResourceLimitsRejectNegativeValues(t *testing.T) {
}
}

func TestValidateResourceLimitsAgainstCapacity(t *testing.T) {
limits := protocol.ResourceLimits{
CPUPercent: 1_100,
MemoryBytes: 64 << 30,
WorkspaceBytes: 1 << 40,
}
capacity := protocol.ResourceCapacity{
CPUCount: 4,
MemoryBytes: 8 << 30,
DiskBytes: 20 << 30,
}

err := protocol.ValidateResourceLimitsAgainstCapacity(limits, capacity)
if err == nil {
t.Fatal("expected oversized resource limits to fail")
}
for _, want := range []string{
"resource_limits.cpu_percent 1100 exceeds worker CPU capacity 400",
"resource_limits.memory_bytes",
"resource_limits.workspace_bytes",
} {
if !strings.Contains(err.Error(), want) {
t.Fatalf("expected resource capacity error to contain %q, got %v", want, err)
}
}

if err := protocol.ValidateResourceLimitsAgainstCapacity(protocol.ResourceLimits{
CPUPercent: 400,
MemoryBytes: 8 << 30,
WorkspaceBytes: 20 << 30,
}, capacity); err != nil {
t.Fatalf("limits at capacity should pass: %v", err)
}

if err := protocol.ValidateResourceLimitsAgainstCapacity(limits, protocol.ResourceCapacity{}); err != nil {
t.Fatalf("unknown capacity should not reject limits: %v", err)
}
}

func TestProviderContractRejectsMalformedConfigSchemaDigest(t *testing.T) {
contract := validBatchProviderContract()
contract.ConfigSchemaDigest = "sha256:not-hex"
Expand Down
Loading