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
67 lines
1.8 KiB
Go
67 lines
1.8 KiB
Go
package scheduler_test
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/jfraeys/fetch_ml/internal/scheduler"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestStateStore_BasicOperations(t *testing.T) {
|
|
dir := t.TempDir()
|
|
ss, err := scheduler.NewStateStore(dir + "/state.json")
|
|
require.NoError(t, err)
|
|
|
|
// Append some events
|
|
events := []scheduler.StateEvent{
|
|
{Type: scheduler.EventJobEnqueued, TaskID: "task-1", Timestamp: time.Now()},
|
|
{Type: scheduler.EventJobAssigned, TaskID: "task-1", WorkerID: "worker-1", Timestamp: time.Now()},
|
|
{Type: scheduler.EventJobCompleted, TaskID: "task-1", WorkerID: "worker-1", Timestamp: time.Now()},
|
|
}
|
|
|
|
for _, e := range events {
|
|
err := ss.Append(e)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
// Replay events
|
|
replayed, err := ss.Replay()
|
|
require.NoError(t, err)
|
|
assert.Len(t, replayed, 3)
|
|
assert.Equal(t, "task-1", replayed[0].TaskID)
|
|
}
|
|
|
|
func TestStateStore_Persistence(t *testing.T) {
|
|
dir := t.TempDir()
|
|
|
|
// Create store and append events
|
|
ss1, err := scheduler.NewStateStore(dir + "/state.json")
|
|
require.NoError(t, err)
|
|
event := scheduler.StateEvent{
|
|
Type: scheduler.EventJobEnqueued,
|
|
TaskID: "persistent-task",
|
|
Timestamp: time.Now(),
|
|
}
|
|
err = ss1.Append(event)
|
|
require.NoError(t, err)
|
|
|
|
// Create new store instance pointing to same directory
|
|
ss2, err := scheduler.NewStateStore(dir + "/state.json")
|
|
require.NoError(t, err)
|
|
replayed, err := ss2.Replay()
|
|
require.NoError(t, err)
|
|
assert.Len(t, replayed, 1)
|
|
assert.Equal(t, "persistent-task", replayed[0].TaskID)
|
|
}
|
|
|
|
func TestStateStore_ReplayEmpty(t *testing.T) {
|
|
dir := t.TempDir()
|
|
ss, err := scheduler.NewStateStore(dir + "/state.json")
|
|
require.NoError(t, err)
|
|
|
|
replayed, err := ss.Replay()
|
|
require.NoError(t, err)
|
|
assert.Empty(t, replayed)
|
|
}
|