43e6446587
Add new scheduler component for distributed ML workload orchestration: - Hub-based coordination for multi-worker clusters - Pacing controller for rate limiting job submissions - Priority queue with preemption support - Port allocator for dynamic service discovery - Protocol handlers for worker-scheduler communication - Service manager with OS-specific implementations - Connection management and state persistence - Template system for service deployment Includes comprehensive test suite: - Unit tests for all core components - Integration tests for distributed scenarios - Benchmark tests for performance validation - Mock fixtures for isolated testing Refs: scheduler-architecture.md
91 lines
2.2 KiB
Go
91 lines
2.2 KiB
Go
package scheduler_test
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/jfraeys/fetch_ml/internal/scheduler"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestPortAllocator_BasicAllocation(t *testing.T) {
|
|
pa := scheduler.NewPortAllocator(10000, 10010)
|
|
|
|
// Allocate first port
|
|
port1, err := pa.Allocate("service-1")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 10000, port1)
|
|
|
|
// Allocate second port
|
|
port2, err := pa.Allocate("service-2")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 10001, port2)
|
|
}
|
|
|
|
func TestPortAllocator_Exhaustion(t *testing.T) {
|
|
pa := scheduler.NewPortAllocator(10000, 10002) // Only 3 ports available
|
|
|
|
// Allocate all ports
|
|
port1, _ := pa.Allocate("service-1")
|
|
assert.Equal(t, 10000, port1)
|
|
|
|
port2, _ := pa.Allocate("service-2")
|
|
assert.Equal(t, 10001, port2)
|
|
|
|
port3, _ := pa.Allocate("service-3")
|
|
assert.Equal(t, 10002, port3)
|
|
|
|
// Should fail - no ports left
|
|
_, err := pa.Allocate("service-4")
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestPortAllocator_Release(t *testing.T) {
|
|
pa := scheduler.NewPortAllocator(10000, 10005)
|
|
|
|
// Allocate and release
|
|
port, _ := pa.Allocate("service-1")
|
|
assert.Equal(t, 10000, port)
|
|
|
|
pa.Release(port)
|
|
|
|
// Should be able to allocate again
|
|
port2, err := pa.Allocate("service-2")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 10000, port2) // Reuses released port
|
|
}
|
|
|
|
func TestPortAllocator_DuplicateServiceID(t *testing.T) {
|
|
pa := scheduler.NewPortAllocator(10000, 10010)
|
|
|
|
// Allocate for service
|
|
port1, err := pa.Allocate("service-1")
|
|
require.NoError(t, err)
|
|
|
|
// Allocate again with same service ID - gets new port (current behavior)
|
|
port2, err := pa.Allocate("service-1")
|
|
require.NoError(t, err)
|
|
assert.NotEqual(t, port1, port2) // Each call returns new port
|
|
}
|
|
|
|
func TestPortAllocator_ConcurrentAccess(t *testing.T) {
|
|
pa := scheduler.NewPortAllocator(10000, 10100) // 101 ports
|
|
done := make(chan bool, 10)
|
|
|
|
// Concurrent allocations
|
|
for i := 0; i < 10; i++ {
|
|
go func(id int) {
|
|
for j := 0; j < 10; j++ {
|
|
pa.Allocate("service-")
|
|
}
|
|
done <- true
|
|
}(i)
|
|
}
|
|
|
|
for i := 0; i < 10; i++ {
|
|
<-done
|
|
}
|
|
|
|
// All 100 ports should be allocated
|
|
// (10 goroutines * 10 allocations, but only 100 unique service IDs)
|
|
}
|