diff --git a/services/README.md b/services/README.md index ee2880ca..f05a8a31 100644 --- a/services/README.md +++ b/services/README.md @@ -10,7 +10,7 @@ Go implementation of a distributed Gateway and pluggable Scheduler for AgentENV. - Gateway resolves `GET /nodes/{id}` via scheduler and proxies to the target node. - Gateway routes sandbox requests by existing sandbox-to-node binding. - Scheduler exposes gRPC API and supports pluggable strategy providers. -- Built-in strategies in v1: round_robin and random. +- Built-in strategies in v1: round_robin, random, and grouped_round_robin. - Scheduler supports both static node configuration and Kubernetes EndpointSlice discovery. - Scheduler sandbox binding store can be in-memory or Redis-backed. - Scheduler can run as a primary read/write service or as query-only replicas that serve only `LookupNode` from Redis. @@ -98,6 +98,9 @@ General config notes: - `SCHEDULER_REDIS_ADDR=` overrides `scheduler.redis_addr` from the environment. - `SCHEDULER_ARTIFACT_STORE_CAPACITY=` overrides `scheduler.artifact_store_capacity` from the environment. - `SCHEDULER_ARTIFACT_LOOKUP_NODE_LIMIT=` overrides `scheduler.artifact_lookup_node_limit` from the environment. +- `SCHEDULER_GROUPED_ROUND_ROBIN_MAX_SANDBOX_COUNT=` overrides `scheduler.grouped_round_robin.max_sandbox_count`. +- `SCHEDULER_GROUPED_ROUND_ROBIN_MAX_CPU_COUNT=` overrides `scheduler.grouped_round_robin.max_cpu_count`. +- `SCHEDULER_GROUPED_ROUND_ROBIN_MAX_MEMORY_MB=` overrides `scheduler.grouped_round_robin.max_memory_mb`. ### Scheduling strategy @@ -107,8 +110,28 @@ General config notes: |---|---| | `round_robin` (default) | Cycles through eligible nodes in stable order | | `random` | Picks a uniformly random eligible node | +| `grouped_round_robin` | Keeps requests for the same image/template in a resource-bounded group, then assigns the next group by global round-robin | -The strategy interface receives `RichNode` values that carry the node identity (ID + endpoint) together with the latest heartbeat `NodeSnapshot` (sandbox counts, CPU, memory, disk metrics). Current built-in strategies ignore the snapshot, but custom strategy implementations can use it for load-aware decisions. +The strategy interface receives `RichNode` values that carry the node identity (ID + endpoint) together with the latest heartbeat `NodeSnapshot` (sandbox counts, CPU, memory, disk metrics). + +The `grouped_round_robin` strategy uses the rootfs image reference from `POST /sandboxes-cold` or the exact template reference from `POST /sandboxes` as its grouping key. Requests for the same key stay on the current group's node until adding another request would exceed any configured group limit. A full group is closed permanently, and the next group is assigned to the next eligible node through one global round-robin cursor shared by all keys. If the open group's node is no longer ready or was removed by the node resource filter, the group closes early. + +Group accounting occurs atomically during scheduling, before the runtime finishes creating the sandbox, so bursts cannot overfill an open group while heartbeat metrics lag. Groups are advisory and in-memory; scheduler restart forgets them. The strategy requires a fresh READY heartbeat and skips nodes with missing or stale telemetry. + +`scheduler.grouped_round_robin.max_sandbox_count` is required and must be greater than zero when `scheduler.strategy` is `grouped_round_robin`. CPU and memory limits are optional additional bounds for cold-start requests: + +```json +"strategy": "grouped_round_robin", +"grouped_round_robin": { + "max_sandbox_count": 4, + "max_cpu_count": 8, + "max_memory_mb": 16384 +} +``` + +When a cold-start request omits `cpuCount` or `memoryMB` while the corresponding group limit is enabled, the scheduler charges that dimension at the full configured group limit. Runtime defaults are node-local and are not available in the gateway hint, so this conservative behavior closes that request's group instead of allowing an unknown default to exceed the budget. + +Template-based create requests do not carry CPU or memory values, so their groups are bounded by `max_sandbox_count` only. Image tags and template aliases are used exactly as supplied and may be mutable; a stale grouping hint can reduce cache affinity, but cannot bypass group or node resource limits. ### Node resource limit diff --git a/services/api/proto/scheduler.pb.go b/services/api/proto/scheduler.pb.go index 73047601..b01db82b 100644 --- a/services/api/proto/scheduler.pb.go +++ b/services/api/proto/scheduler.pb.go @@ -359,7 +359,10 @@ func (x *NewColdSandboxHint) GetMetadata() map[string]string { type NewSandboxHint struct { state protoimpl.MessageState `protogen:"open.v1"` // Sandbox metadata key/value pairs parsed from the request body. - Metadata map[string]string `protobuf:"bytes,1,rep,name=metadata,proto3" json:"metadata,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + Metadata map[string]string `protobuf:"bytes,1,rep,name=metadata,proto3" json:"metadata,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + // Template or snapshot reference supplied by POST /sandboxes. The grouped + // round-robin strategy treats the exact reference as a workload identity. + TemplateId string `protobuf:"bytes,2,opt,name=template_id,json=templateId,proto3" json:"template_id,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -401,6 +404,13 @@ func (x *NewSandboxHint) GetMetadata() map[string]string { return nil } +func (x *NewSandboxHint) GetTemplateId() string { + if x != nil { + return x.TemplateId + } + return "" +} + type ScheduleRequest struct { state protoimpl.MessageState `protogen:"open.v1"` Hint *ScheduleRequestHint `protobuf:"bytes,2,opt,name=hint,proto3" json:"hint,omitempty"` @@ -2327,9 +2337,11 @@ const file_api_proto_scheduler_proto_rawDesc = "" + "\bmetadata\x18\x04 \x03(\v2..scheduler.v1.NewColdSandboxHint.MetadataEntryR\bmetadata\x1a;\n" + "\rMetadataEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + - "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\x95\x01\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\xb6\x01\n" + "\x0eNewSandboxHint\x12F\n" + - "\bmetadata\x18\x01 \x03(\v2*.scheduler.v1.NewSandboxHint.MetadataEntryR\bmetadata\x1a;\n" + + "\bmetadata\x18\x01 \x03(\v2*.scheduler.v1.NewSandboxHint.MetadataEntryR\bmetadata\x12\x1f\n" + + "\vtemplate_id\x18\x02 \x01(\tR\n" + + "templateId\x1a;\n" + "\rMetadataEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"N\n" + diff --git a/services/api/proto/scheduler.proto b/services/api/proto/scheduler.proto index 169459f4..8f330f75 100644 --- a/services/api/proto/scheduler.proto +++ b/services/api/proto/scheduler.proto @@ -51,6 +51,9 @@ message NewColdSandboxHint { message NewSandboxHint { // Sandbox metadata key/value pairs parsed from the request body. map metadata = 1; + // Template or snapshot reference supplied by POST /sandboxes. The grouped + // round-robin strategy treats the exact reference as a workload identity. + string template_id = 2; } message ScheduleRequest { diff --git a/services/gateway/internal/schedule_hint.go b/services/gateway/internal/schedule_hint.go index 44284a27..7eb4b137 100644 --- a/services/gateway/internal/schedule_hint.go +++ b/services/gateway/internal/schedule_hint.go @@ -134,7 +134,8 @@ func parseNewColdSandboxHint(body []byte) *schedulerv1.NewColdSandboxHint { // newSandboxBody mirrors the subset of NewSandbox (src/api/openapi.yml) that is // relevant for scheduling. type newSandboxBody struct { - Metadata map[string]string `json:"metadata"` + TemplateID string `json:"templateID"` + Metadata map[string]string `json:"metadata"` } // parseNewSandboxHint extracts the structured sandbox hint from the request @@ -149,6 +150,7 @@ func parseNewSandboxHint(body []byte) *schedulerv1.NewSandboxHint { if err := json.Unmarshal(body, &parsed); err != nil { return hint } + hint.TemplateId = parsed.TemplateID hint.Metadata = parsed.Metadata return hint } diff --git a/services/gateway/internal/schedule_hint_test.go b/services/gateway/internal/schedule_hint_test.go index c5839286..b0ead1fe 100644 --- a/services/gateway/internal/schedule_hint_test.go +++ b/services/gateway/internal/schedule_hint_test.go @@ -35,6 +35,9 @@ func TestBuildScheduleHintNewSandbox(t *testing.T) { if hint.GetNewColdSandbox() != nil { t.Fatalf("did not expect cold sandbox hint") } + if got := hint.GetNewSandbox().GetTemplateId(); got != "tmpl" { + t.Fatalf("template_id = %q, want tmpl", got) + } // Body must remain available for the upstream request. body, err := io.ReadAll(r.Body) @@ -158,6 +161,25 @@ func TestParseNewColdSandboxHint(t *testing.T) { }) } +func TestParseNewSandboxHint(t *testing.T) { + t.Run("template and metadata", func(t *testing.T) { + hint := parseNewSandboxHint([]byte(`{"templateID":"tmpl","metadata":{"team":"infra"}}`)) + if got := hint.GetTemplateId(); got != "tmpl" { + t.Fatalf("template_id = %q, want tmpl", got) + } + if got := hint.GetMetadata()["team"]; got != "infra" { + t.Fatalf("metadata team = %q, want infra", got) + } + }) + + t.Run("malformed json", func(t *testing.T) { + hint := parseNewSandboxHint([]byte("{not json")) + if hint.GetTemplateId() != "" || len(hint.GetMetadata()) != 0 { + t.Fatalf("expected empty best-effort hint, got %v", hint) + } + }) +} + func TestCaptureRequestBodyNil(t *testing.T) { r := newHintRequest(t, http.MethodPost, "/sandboxes-cold", "") body, err := captureRequestBody(r) diff --git a/services/scheduler/cmd/main.go b/services/scheduler/cmd/main.go index 3e8574a5..2a672170 100644 --- a/services/scheduler/cmd/main.go +++ b/services/scheduler/cmd/main.go @@ -69,7 +69,14 @@ func main() { svc := scheduler.NewService( logger, registry, - scheduler.NewStrategy(cfg.Scheduler.Strategy), + scheduler.NewStrategy( + cfg.Scheduler.Strategy, + scheduler.WithGroupedRoundRobinLimits(scheduler.GroupedRoundRobinLimits{ + MaxSandboxCount: cfg.Scheduler.GroupedRoundRobin.MaxSandboxCount, + MaxCPUCount: cfg.Scheduler.GroupedRoundRobin.MaxCPUCount, + MaxMemoryMB: cfg.Scheduler.GroupedRoundRobin.MaxMemoryMB, + }), + ), store, scheduler.WithArtifactStore(scheduler.NewInMemoryArtifactStore( cfg.Scheduler.ArtifactStoreCapacity, diff --git a/services/scheduler/internal/grouped_round_robin_strategy_test.go b/services/scheduler/internal/grouped_round_robin_strategy_test.go new file mode 100644 index 00000000..b56d2bc7 --- /dev/null +++ b/services/scheduler/internal/grouped_round_robin_strategy_test.go @@ -0,0 +1,299 @@ +package scheduler + +import ( + "fmt" + "sync" + "testing" + + schedulerv1 "agentenv/services/api/proto" +) + +func TestGroupedRoundRobinInterleavesImageGroups(t *testing.T) { + strategy := NewGroupedRoundRobinStrategy(GroupedRoundRobinLimits{MaxSandboxCount: 2}) + nodes := readyNodes("a", "b", "c") + + got := []string{ + selectNodeID(t, strategy, nodes, coldHint("ubuntu", 0, 0)), + selectNodeID(t, strategy, nodes, coldHint("python", 0, 0)), + selectNodeID(t, strategy, nodes, coldHint("ubuntu", 0, 0)), + selectNodeID(t, strategy, nodes, coldHint("python", 0, 0)), + selectNodeID(t, strategy, nodes, coldHint("ubuntu", 0, 0)), + selectNodeID(t, strategy, nodes, coldHint("python", 0, 0)), + } + want := []string{"a", "b", "a", "b", "c", "a"} + if !equalNodeIDs(got, want) { + t.Fatalf("placements = %v, want %v", got, want) + } +} + +func TestGroupedRoundRobinClosesGroupAtCPULimit(t *testing.T) { + strategy := NewGroupedRoundRobinStrategy(GroupedRoundRobinLimits{ + MaxSandboxCount: 10, + MaxCPUCount: 4, + }) + nodes := readyNodes("a", "b") + + got := []string{ + selectNodeID(t, strategy, nodes, coldHint("ubuntu", 3, 0)), + selectNodeID(t, strategy, nodes, coldHint("ubuntu", 2, 0)), + selectNodeID(t, strategy, nodes, coldHint("ubuntu", 2, 0)), + selectNodeID(t, strategy, nodes, coldHint("ubuntu", 1, 0)), + } + want := []string{"a", "b", "b", "a"} + if !equalNodeIDs(got, want) { + t.Fatalf("placements = %v, want %v", got, want) + } +} + +func TestGroupedRoundRobinClosesGroupAtMemoryLimit(t *testing.T) { + strategy := NewGroupedRoundRobinStrategy(GroupedRoundRobinLimits{ + MaxSandboxCount: 10, + MaxMemoryMB: 1024, + }) + nodes := readyNodes("a", "b") + + got := []string{ + selectNodeID(t, strategy, nodes, coldHint("ubuntu", 0, 768)), + selectNodeID(t, strategy, nodes, coldHint("ubuntu", 0, 512)), + selectNodeID(t, strategy, nodes, coldHint("ubuntu", 0, 512)), + } + want := []string{"a", "b", "b"} + if !equalNodeIDs(got, want) { + t.Fatalf("placements = %v, want %v", got, want) + } +} + +func TestGroupedRoundRobinTreatsOmittedColdResourcesConservatively(t *testing.T) { + strategy := NewGroupedRoundRobinStrategy(GroupedRoundRobinLimits{ + MaxSandboxCount: 10, + MaxCPUCount: 4, + MaxMemoryMB: 1024, + }) + nodes := readyNodes("a", "b") + + got := []string{ + selectNodeID(t, strategy, nodes, coldHint("ubuntu", 0, 0)), + selectNodeID(t, strategy, nodes, coldHint("ubuntu", 0, 0)), + } + want := []string{"a", "b"} + if !equalNodeIDs(got, want) { + t.Fatalf("placements for omitted resources = %v, want %v", got, want) + } +} + +func TestGroupedRoundRobinOversizedRequestDoesNotLeaveOpenGroup(t *testing.T) { + strategy := NewGroupedRoundRobinStrategy(GroupedRoundRobinLimits{ + MaxSandboxCount: 10, + MaxCPUCount: 2, + }) + nodes := readyNodes("a", "b") + + first := selectNodeID(t, strategy, nodes, coldHint("ubuntu", 4, 0)) + second := selectNodeID(t, strategy, nodes, coldHint("ubuntu", 4, 0)) + if first != "a" || second != "b" { + t.Fatalf("oversized placements = %s %s, want a b", first, second) + } + if len(strategy.groups) != 0 { + t.Fatalf("oversized requests left %d open groups, want 0", len(strategy.groups)) + } +} + +func TestGroupedRoundRobinClosesGroupWhenNodeBecomesIneligible(t *testing.T) { + strategy := NewGroupedRoundRobinStrategy(GroupedRoundRobinLimits{MaxSandboxCount: 3}) + nodes := readyNodes("a", "b") + + if got := selectNodeID(t, strategy, nodes, coldHint("ubuntu", 0, 0)); got != "a" { + t.Fatalf("first placement = %s, want a", got) + } + + // The service removes resource-constrained nodes before calling Select. + if got := selectNodeID(t, strategy, readyNodes("b", "c"), coldHint("ubuntu", 0, 0)); got != "b" { + t.Fatalf("replacement placement = %s, want b", got) + } +} + +func TestGroupedRoundRobinSkipsNodesWithoutReadyHeartbeat(t *testing.T) { + strategy := NewGroupedRoundRobinStrategy(GroupedRoundRobinLimits{MaxSandboxCount: 2}) + nodes := []RichNode{ + {Node: Node{ID: "missing"}}, + { + Node: Node{ID: "unhealthy"}, + Snapshot: &schedulerv1.NodeSnapshot{ + Status: schedulerv1.NodeStatus_NODE_STATUS_UNHEALTHY, + }, + }, + readyNodes("ready")[0], + } + + if got := selectNodeID(t, strategy, nodes, coldHint("ubuntu", 0, 0)); got != "ready" { + t.Fatalf("placement = %s, want ready", got) + } +} + +func TestGroupedRoundRobinGroupsTemplatesByExactReference(t *testing.T) { + strategy := NewGroupedRoundRobinStrategy(GroupedRoundRobinLimits{MaxSandboxCount: 2}) + nodes := readyNodes("a", "b") + + got := []string{ + selectNodeID(t, strategy, nodes, templateHint("base")), + selectNodeID(t, strategy, nodes, templateHint("alias")), + selectNodeID(t, strategy, nodes, templateHint("base")), + selectNodeID(t, strategy, nodes, templateHint("alias")), + selectNodeID(t, strategy, nodes, templateHint("base")), + } + want := []string{"a", "b", "a", "b", "a"} + if !equalNodeIDs(got, want) { + t.Fatalf("placements = %v, want %v", got, want) + } +} + +func TestGroupedRoundRobinFallsBackToGlobalRoundRobin(t *testing.T) { + strategy := NewGroupedRoundRobinStrategy(GroupedRoundRobinLimits{MaxSandboxCount: 2}) + nodes := readyNodes("a", "b", "c") + + tooLong := coldHint("x", 0, 0) + tooLong.GetNewColdSandbox().Images[0] = string(make([]byte, maxGroupedRoundRobinKeyBytes+1)) + got := []string{ + selectNodeID(t, strategy, nodes, nil), + selectNodeID(t, strategy, nodes, coldHint("", 0, 0)), + selectNodeID(t, strategy, nodes, tooLong), + } + want := []string{"a", "b", "c"} + if !equalNodeIDs(got, want) { + t.Fatalf("fallback placements = %v, want %v", got, want) + } +} + +func TestGroupedRoundRobinFallbackDoesNotAdvanceGroupCursor(t *testing.T) { + strategy := NewGroupedRoundRobinStrategy(GroupedRoundRobinLimits{MaxSandboxCount: 2}) + nodes := readyNodes("a", "b") + + if got := selectNodeID(t, strategy, nodes, nil); got != "a" { + t.Fatalf("fallback placement = %s, want a", got) + } + if got := selectNodeID(t, strategy, nodes, coldHint("ubuntu", 0, 0)); got != "a" { + t.Fatalf("first group placement = %s, want a", got) + } +} + +func TestGroupedRoundRobinBoundsOpenGroupState(t *testing.T) { + strategy := NewGroupedRoundRobinStrategy(GroupedRoundRobinLimits{MaxSandboxCount: 2}) + nodes := readyNodes("a") + + for i := 0; i <= maxOpenGroupedRoundRobinGroups; i++ { + selectNodeID(t, strategy, nodes, coldHint(fmt.Sprintf("image-%d", i), 0, 0)) + } + if got := len(strategy.groups); got != maxOpenGroupedRoundRobinGroups { + t.Fatalf("open group count = %d, want %d", got, maxOpenGroupedRoundRobinGroups) + } + if _, ok := strategy.groups["image:image-0"]; ok { + t.Fatal("oldest open group was not evicted") + } +} + +func TestGroupedRoundRobinCountsConcurrentPlacementsAtomically(t *testing.T) { + strategy := NewGroupedRoundRobinStrategy(GroupedRoundRobinLimits{MaxSandboxCount: 10}) + nodes := readyNodes("a", "b") + + const requests = 200 + type result struct { + nodeID string + err error + } + results := make(chan result, requests) + var wg sync.WaitGroup + wg.Add(requests) + for i := 0; i < requests; i++ { + go func() { + defer wg.Done() + node, err := strategy.Select(nodes, coldHint("ubuntu", 0, 0)) + results <- result{nodeID: node.ID, err: err} + }() + } + wg.Wait() + close(results) + + counts := map[string]int{} + for result := range results { + if result.err != nil { + t.Fatalf("Select returned error: %v", result.err) + } + counts[result.nodeID]++ + } + if counts["a"] != requests/2 || counts["b"] != requests/2 { + t.Fatalf("concurrent placement counts = %v, want equal distribution", counts) + } +} + +func TestNewStrategySelectsGroupedRoundRobinCaseInsensitively(t *testing.T) { + strategy := NewStrategy( + " GROUPED_ROUND_ROBIN ", + WithGroupedRoundRobinLimits(GroupedRoundRobinLimits{MaxSandboxCount: 2}), + ) + if strategy.Name() != "grouped_round_robin" { + t.Fatalf("strategy name = %q, want grouped_round_robin", strategy.Name()) + } +} + +func readyNodes(ids ...string) []RichNode { + nodes := make([]RichNode, 0, len(ids)) + for _, id := range ids { + nodes = append(nodes, RichNode{ + Node: Node{ID: id, Endpoint: "http://" + id}, + Snapshot: &schedulerv1.NodeSnapshot{ + Status: schedulerv1.NodeStatus_NODE_STATUS_READY, + }, + }) + } + return nodes +} + +func coldHint(image string, cpuCount uint32, memoryMB uint64) *schedulerv1.ScheduleRequestHint { + images := []string(nil) + if image != "" { + images = []string{image} + } + return &schedulerv1.ScheduleRequestHint{ + Kind: &schedulerv1.ScheduleRequestHint_NewColdSandbox{ + NewColdSandbox: &schedulerv1.NewColdSandboxHint{ + Images: images, + CpuCount: cpuCount, + MemoryMb: memoryMB, + }, + }, + } +} + +func templateHint(templateID string) *schedulerv1.ScheduleRequestHint { + return &schedulerv1.ScheduleRequestHint{ + Kind: &schedulerv1.ScheduleRequestHint_NewSandbox{ + NewSandbox: &schedulerv1.NewSandboxHint{TemplateId: templateID}, + }, + } +} + +func selectNodeID( + t *testing.T, + strategy Strategy, + nodes []RichNode, + hint *schedulerv1.ScheduleRequestHint, +) string { + t.Helper() + node, err := strategy.Select(nodes, hint) + if err != nil { + t.Fatalf("Select returned error: %v", err) + } + return node.ID +} + +func equalNodeIDs(got, want []string) bool { + if len(got) != len(want) { + return false + } + for i := range got { + if got[i] != want[i] { + return false + } + } + return true +} diff --git a/services/scheduler/internal/node_registry.go b/services/scheduler/internal/node_registry.go index 5596d298..9542b926 100644 --- a/services/scheduler/internal/node_registry.go +++ b/services/scheduler/internal/node_registry.go @@ -21,11 +21,9 @@ type NodeRegistry interface { ListP2pPeers(clusterID string, backend string, excludeNodeID string, now time.Time) []*schedulerv1.P2PPeer FilterP2pPeers(clusterID string, backend string, nodeIDs []string, excludeNodeID string, now time.Time) []*schedulerv1.P2PPeer GetObserved(nodeID string, clusterID string, now time.Time) (*schedulerv1.ObservedNode, bool) - // PeekObserved returns the latest heartbeat-reported NodeSnapshot for a node. - // Unlike GetObserved, it does not derive status from discovery state or TTL, - // and returns only the raw snapshot suitable for scheduling decisions. - // Returns nil if the node has never sent a heartbeat. - PeekObserved(nodeID string) *schedulerv1.NodeSnapshot + // SchedulingSnapshot returns the latest NodeSnapshot with status derived + // from current discovery state and heartbeat TTL. + SchedulingSnapshot(nodeID string, now time.Time) *schedulerv1.NodeSnapshot UnregisterObserved(nodeID string, serviceInstanceID string) error } @@ -331,7 +329,9 @@ func (r *AtomicNodeRegistry) GetObserved(nodeID string, clusterID string, now ti return r.deriveObservedNodeViewLocked(record, nowMs), true } -func (r *AtomicNodeRegistry) PeekObserved(nodeID string) *schedulerv1.NodeSnapshot { +func (r *AtomicNodeRegistry) SchedulingSnapshot(nodeID string, now time.Time) *schedulerv1.NodeSnapshot { + nowMs := now.UTC().UnixMilli() + r.mu.RLock() defer r.mu.RUnlock() record, ok := r.observed[nodeID] @@ -342,7 +342,9 @@ func (r *AtomicNodeRegistry) PeekObserved(nodeID string) *schedulerv1.NodeSnapsh if snapshot == nil { return nil } - return cloneSnapshot(snapshot) + out := cloneSnapshot(snapshot) + out.Status = r.derivedStatusLocked(record, nowMs) + return out } func (r *AtomicNodeRegistry) UnregisterObserved(nodeID string, serviceInstanceID string) error { @@ -374,33 +376,38 @@ func (r *AtomicNodeRegistry) deriveObservedNodeViewLocked(record observedNodeRec } nodeID := out.GetNodeId() - knownNode, inDiscovery := r.nodesByID[nodeID] - isLingering := r.lingeringIDs[nodeID] if inDiscovery && strings.TrimSpace(knownNode.Endpoint) != "" { out.Endpoint = knownNode.Endpoint } + out.Snapshot.Status = r.derivedStatusLocked(record, nowMs) + return out +} + +// derivedStatusLocked computes the effective scheduling status for a heartbeat +// record. r.mu must be held by the caller. +func (r *AtomicNodeRegistry) derivedStatusLocked(record observedNodeRecord, nowMs int64) schedulerv1.NodeStatus { ttl := record.reportTTL if ttl <= 0 { ttl = defaultObservedReportTTL } - if out.GetLastSeenUnixMs() > 0 && nowMs-out.GetLastSeenUnixMs() > ttl.Milliseconds() { - out.Snapshot.Status = schedulerv1.NodeStatus_NODE_STATUS_UNHEALTHY - } else if !inDiscovery { - out.Snapshot.Status = schedulerv1.NodeStatus_NODE_STATUS_CONNECTING - } else if isLingering { - out.Snapshot.Status = schedulerv1.NodeStatus_NODE_STATUS_LINGERING - } else { - // Active — keep the status reported by the node. - if out.Snapshot.GetStatus() == schedulerv1.NodeStatus_NODE_STATUS_UNSPECIFIED { - out.Snapshot.Status = schedulerv1.NodeStatus_NODE_STATUS_CONNECTING - } + if lastSeen := record.node.GetLastSeenUnixMs(); lastSeen > 0 && nowMs-lastSeen > ttl.Milliseconds() { + return schedulerv1.NodeStatus_NODE_STATUS_UNHEALTHY } - - return out + nodeID := record.node.GetNodeId() + if _, inDiscovery := r.nodesByID[nodeID]; !inDiscovery { + return schedulerv1.NodeStatus_NODE_STATUS_CONNECTING + } + if r.lingeringIDs[nodeID] { + return schedulerv1.NodeStatus_NODE_STATUS_LINGERING + } + if status := record.node.GetSnapshot().GetStatus(); status != schedulerv1.NodeStatus_NODE_STATUS_UNSPECIFIED { + return status + } + return schedulerv1.NodeStatus_NODE_STATUS_CONNECTING } func cloneObservedNode(node *schedulerv1.ObservedNode) *schedulerv1.ObservedNode { diff --git a/services/scheduler/internal/node_registry_test.go b/services/scheduler/internal/node_registry_test.go index 66d656a1..a774c295 100644 --- a/services/scheduler/internal/node_registry_test.go +++ b/services/scheduler/internal/node_registry_test.go @@ -290,6 +290,115 @@ func TestLingeringNodeBecomesUnhealthyAfterTTL(t *testing.T) { } } +func TestSchedulingSnapshotDerivesCurrentStatus(t *testing.T) { + start := time.Unix(100, 0) + tests := []struct { + name string + lingering bool + queryOffset time.Duration + want schedulerv1.NodeStatus + }{ + { + name: "fresh ready heartbeat", + want: schedulerv1.NodeStatus_NODE_STATUS_READY, + }, + { + name: "expired heartbeat", + queryOffset: 2 * time.Second, + want: schedulerv1.NodeStatus_NODE_STATUS_UNHEALTHY, + }, + { + name: "lingering overrides reported ready", + lingering: true, + want: schedulerv1.NodeStatus_NODE_STATUS_LINGERING, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + registry := NewAtomicNodeRegistry(nil, time.Second) + node := Node{ID: "node-a", Endpoint: "http://node-a"} + if tt.lingering { + registry.Set(nil, []Node{node}) + } else { + registry.Set([]Node{node}, nil) + } + if _, _, err := registry.Heartbeat(&schedulerv1.HeartbeatRequest{ + NodeId: node.ID, + ClusterId: "cluster-a", + ServiceInstanceId: "svc-a", + Snapshot: &schedulerv1.NodeSnapshot{Status: schedulerv1.NodeStatus_NODE_STATUS_READY}, + }, start); err != nil { + t.Fatalf("heartbeat: %v", err) + } + + snapshot := registry.SchedulingSnapshot(node.ID, start.Add(tt.queryOffset)) + if snapshot == nil { + t.Fatal("expected scheduling snapshot") + } + if got := snapshot.GetStatus(); got != tt.want { + t.Fatalf("status = %v, want %v", got, tt.want) + } + }) + } +} + +func TestSchedulingSnapshotReturnsNilWithoutHeartbeat(t *testing.T) { + registry := NewAtomicNodeRegistry( + []Node{{ID: "node-a", Endpoint: "http://node-a"}}, + time.Second, + ) + + if snapshot := registry.SchedulingSnapshot("node-a", time.Unix(100, 0)); snapshot != nil { + t.Fatalf("expected nil scheduling snapshot, got %+v", snapshot) + } +} + +func TestSchedulingSnapshotReturnsDeepClone(t *testing.T) { + registry := NewAtomicNodeRegistry( + []Node{{ID: "node-a", Endpoint: "http://node-a"}}, + time.Second, + ) + now := time.Unix(100, 0) + if _, _, err := registry.Heartbeat(&schedulerv1.HeartbeatRequest{ + NodeId: "node-a", + ClusterId: "cluster-a", + ServiceInstanceId: "svc-a", + Snapshot: &schedulerv1.NodeSnapshot{ + Status: schedulerv1.NodeStatus_NODE_STATUS_READY, + AllocatedCpu: 2, + Disks: []*schedulerv1.DiskMetric{{ + MountPoint: "/", + UsedBytes: 10, + }}, + }, + }, now); err != nil { + t.Fatalf("heartbeat: %v", err) + } + + first := registry.SchedulingSnapshot("node-a", now) + if first == nil || len(first.GetDisks()) != 1 { + t.Fatalf("unexpected first scheduling snapshot: %+v", first) + } + first.Status = schedulerv1.NodeStatus_NODE_STATUS_UNHEALTHY + first.AllocatedCpu = 99 + first.Disks[0].UsedBytes = 99 + + second := registry.SchedulingSnapshot("node-a", now) + if second == nil { + t.Fatal("expected second scheduling snapshot") + } + if got := second.GetStatus(); got != schedulerv1.NodeStatus_NODE_STATUS_READY { + t.Fatalf("stored status changed through returned snapshot: %v", got) + } + if got := second.GetAllocatedCpu(); got != 2 { + t.Fatalf("stored allocated CPU changed through returned snapshot: %d", got) + } + if got := second.GetDisks()[0].GetUsedBytes(); got != 10 { + t.Fatalf("stored disk metric changed through returned snapshot: %d", got) + } +} + func heartbeatWithConfig(t *testing.T, registry *AtomicNodeRegistry, nodeID, clusterID, svcID, cpuJSON string) string { t.Helper() var mi *schedulerv1.MachineInfo diff --git a/services/scheduler/internal/service.go b/services/scheduler/internal/service.go index a284dfa8..d9a89027 100644 --- a/services/scheduler/internal/service.go +++ b/services/scheduler/internal/service.go @@ -89,7 +89,7 @@ func (s *Service) Schedule(_ context.Context, req *schedulerv1.ScheduleRequest) for _, n := range discovered { rich = append(rich, RichNode{ Node: n, - Snapshot: s.nodes.PeekObserved(n.ID), + Snapshot: s.nodes.SchedulingSnapshot(n.ID, start), }) } @@ -130,7 +130,7 @@ func summarizeScheduleHint(hint *schedulerv1.ScheduleRequestHint) string { c := k.NewColdSandbox return fmt.Sprintf("new_cold_sandbox cpu=%d memory_mb=%d images=%v", c.GetCpuCount(), c.GetMemoryMb(), c.GetImages()) case *schedulerv1.ScheduleRequestHint_NewSandbox: - return "new_sandbox" + return fmt.Sprintf("new_sandbox template=%q", k.NewSandbox.GetTemplateId()) default: return "none" } diff --git a/services/scheduler/internal/service_test.go b/services/scheduler/internal/service_test.go index 4f326b89..f9979ba0 100644 --- a/services/scheduler/internal/service_test.go +++ b/services/scheduler/internal/service_test.go @@ -232,6 +232,52 @@ func TestScheduleReturnsUnavailableWhenRegistryIsEmpty(t *testing.T) { } } +func TestGroupedRoundRobinScheduleSkipsStaleHeartbeat(t *testing.T) { + registry := NewAtomicNodeRegistry( + []Node{ + {ID: "node-a", Endpoint: "http://node-a"}, + {ID: "node-b", Endpoint: "http://node-b"}, + }, + time.Second, + ) + now := time.Now() + for _, heartbeat := range []struct { + nodeID string + at time.Time + }{ + {nodeID: "node-a", at: now.Add(-time.Minute)}, + {nodeID: "node-b", at: now}, + } { + _, _, err := registry.Heartbeat(&schedulerv1.HeartbeatRequest{ + NodeId: heartbeat.nodeID, + ClusterId: "cluster-1", + ServiceInstanceId: "service-" + heartbeat.nodeID, + Snapshot: &schedulerv1.NodeSnapshot{ + Status: schedulerv1.NodeStatus_NODE_STATUS_READY, + }, + }, heartbeat.at) + if err != nil { + t.Fatalf("heartbeat %s failed: %v", heartbeat.nodeID, err) + } + } + + service := NewService( + zap.NewNop(), + registry, + NewGroupedRoundRobinStrategy(GroupedRoundRobinLimits{MaxSandboxCount: 2}), + NewInMemoryBindingStore(defaultObservedReportTTL), + ) + response, err := service.Schedule(context.Background(), &schedulerv1.ScheduleRequest{ + Hint: coldHint("ubuntu", 0, 0), + }) + if err != nil { + t.Fatalf("Schedule returned error: %v", err) + } + if got := response.GetNode().GetNodeId(); got != "node-b" { + t.Fatalf("scheduled node = %q, want node-b", got) + } +} + func TestScheduleOnlyConsidersReadyNodes(t *testing.T) { registry := NewAtomicNodeRegistry(nil, defaultObservedReportTTL) // node-a: active, node-b: lingering diff --git a/services/scheduler/internal/strategy.go b/services/scheduler/internal/strategy.go index 165588bb..21e362e7 100644 --- a/services/scheduler/internal/strategy.go +++ b/services/scheduler/internal/strategy.go @@ -1,8 +1,12 @@ package scheduler import ( + "container/list" "errors" "math/rand" + "sort" + "strings" + "sync" "sync/atomic" schedulerv1 "agentenv/services/api/proto" @@ -48,10 +52,249 @@ func (s *RandomStrategy) Name() string { return "random" } -func NewStrategy(name string) Strategy { - switch name { +// GroupedRoundRobinLimits bounds a same-workload placement group. Zero CPU and +// memory limits disable those checks; MaxSandboxCount is always enforced. +type GroupedRoundRobinLimits struct { + MaxSandboxCount uint32 + MaxCPUCount uint32 + MaxMemoryMB uint64 +} + +type groupedRoundRobinRequest struct { + key string + cpuCount uint32 + memoryMB uint64 +} + +type groupedRoundRobinGroup struct { + nodeID string + sandboxCount uint32 + cpuCount uint64 + memoryMB uint64 +} + +type groupedRoundRobinGroupEntry struct { + key string + group groupedRoundRobinGroup +} + +const ( + maxGroupedRoundRobinKeyBytes = 1024 + maxOpenGroupedRoundRobinGroups = 10_000 +) + +// GroupedRoundRobinStrategy keeps same-workload requests on one node until the +// current group reaches a configured budget. New groups share a global +// round-robin cursor so popular workloads spread progressively across nodes. +type GroupedRoundRobinStrategy struct { + mu sync.Mutex + lastGroupNodeID string + lastFallbackNodeID string + limits GroupedRoundRobinLimits + groups map[string]*list.Element + lru list.List +} + +func NewGroupedRoundRobinStrategy(limits GroupedRoundRobinLimits) *GroupedRoundRobinStrategy { + if limits.MaxSandboxCount == 0 { + // Production config rejects this, but keep direct construction bounded. + limits.MaxSandboxCount = 1 + } + return &GroupedRoundRobinStrategy{ + limits: limits, + groups: make(map[string]*list.Element), + } +} + +func (s *GroupedRoundRobinStrategy) Select(nodes []RichNode, hint *schedulerv1.ScheduleRequestHint) (RichNode, error) { + ready := readyGroupedRoundRobinNodes(nodes) + if len(ready) == 0 { + return RichNode{}, ErrNoNodes + } + + request, grouped := groupedRoundRobinRequestFromHint(hint, s.limits) + + s.mu.Lock() + defer s.mu.Unlock() + + if !grouped { + return selectNext(ready, &s.lastFallbackNodeID), nil + } + + if element, ok := s.groups[request.key]; ok { + entry := element.Value.(*groupedRoundRobinGroupEntry) + if node, eligible := findNode(ready, entry.group.nodeID); eligible && + groupCanFit(entry.group, request, s.limits) { + addToGroup(&entry.group, request) + if groupIsFull(entry.group, s.limits) { + s.removeGroup(element) + } else { + s.lru.MoveToFront(element) + } + return node, nil + } + s.removeGroup(element) + } + + node := selectNext(ready, &s.lastGroupNodeID) + group := groupedRoundRobinGroup{nodeID: node.ID} + addToGroup(&group, request) + if !groupIsFull(group, s.limits) { + s.putGroup(request.key, group) + } + return node, nil +} + +func (s *GroupedRoundRobinStrategy) Name() string { + return "grouped_round_robin" +} + +func selectNext(nodes []RichNode, lastNodeID *string) RichNode { + index := 0 + if *lastNodeID != "" { + index = sort.Search(len(nodes), func(i int) bool { + return nodes[i].ID > *lastNodeID + }) + if index == len(nodes) { + index = 0 + } + } + node := nodes[index] + *lastNodeID = node.ID + return node +} + +func (s *GroupedRoundRobinStrategy) putGroup(key string, group groupedRoundRobinGroup) { + element := s.lru.PushFront(&groupedRoundRobinGroupEntry{key: key, group: group}) + s.groups[key] = element + if len(s.groups) <= maxOpenGroupedRoundRobinGroups { + return + } + s.removeGroup(s.lru.Back()) +} + +func (s *GroupedRoundRobinStrategy) removeGroup(element *list.Element) { + if element == nil { + return + } + entry := element.Value.(*groupedRoundRobinGroupEntry) + delete(s.groups, entry.key) + s.lru.Remove(element) +} + +func readyGroupedRoundRobinNodes(nodes []RichNode) []RichNode { + ready := make([]RichNode, 0, len(nodes)) + for _, node := range nodes { + if node.Snapshot == nil || + node.Snapshot.GetStatus() != schedulerv1.NodeStatus_NODE_STATUS_READY { + continue + } + ready = append(ready, node) + } + sort.Slice(ready, func(i, j int) bool { + return ready[i].ID < ready[j].ID + }) + return ready +} + +func groupedRoundRobinRequestFromHint( + hint *schedulerv1.ScheduleRequestHint, + limits GroupedRoundRobinLimits, +) (groupedRoundRobinRequest, bool) { + var request groupedRoundRobinRequest + switch kind := hint.GetKind().(type) { + case *schedulerv1.ScheduleRequestHint_NewColdSandbox: + images := kind.NewColdSandbox.GetImages() + if len(images) == 0 { + return groupedRoundRobinRequest{}, false + } + request = groupedRoundRobinRequest{ + key: "image:" + strings.TrimSpace(images[0]), + cpuCount: kind.NewColdSandbox.GetCpuCount(), + memoryMB: kind.NewColdSandbox.GetMemoryMb(), + } + // The runtime fills omitted cold-start resources from node-local + // machine defaults. The scheduler cannot know those defaults for every + // candidate, so charge an unknown dimension at the full configured + // group limit. This conservatively prevents an omitted value from + // allowing a group to exceed its resource budget. + if request.cpuCount == 0 && limits.MaxCPUCount > 0 { + request.cpuCount = limits.MaxCPUCount + } + if request.memoryMB == 0 && limits.MaxMemoryMB > 0 { + request.memoryMB = limits.MaxMemoryMB + } + case *schedulerv1.ScheduleRequestHint_NewSandbox: + request.key = "template:" + strings.TrimSpace(kind.NewSandbox.GetTemplateId()) + default: + return groupedRoundRobinRequest{}, false + } + if strings.HasSuffix(request.key, ":") || len(request.key) > maxGroupedRoundRobinKeyBytes { + return groupedRoundRobinRequest{}, false + } + return request, true +} + +func findNode(nodes []RichNode, nodeID string) (RichNode, bool) { + for _, node := range nodes { + if node.ID == nodeID { + return node, true + } + } + return RichNode{}, false +} + +func groupCanFit(group groupedRoundRobinGroup, request groupedRoundRobinRequest, limits GroupedRoundRobinLimits) bool { + if group.sandboxCount >= limits.MaxSandboxCount { + return false + } + if exceedsLimit(group.cpuCount, uint64(request.cpuCount), uint64(limits.MaxCPUCount)) { + return false + } + return !exceedsLimit(group.memoryMB, request.memoryMB, limits.MaxMemoryMB) +} + +func exceedsLimit(current, added, limit uint64) bool { + return limit > 0 && (added > limit || current > limit-added) +} + +func addToGroup(group *groupedRoundRobinGroup, request groupedRoundRobinRequest) { + group.sandboxCount++ + group.cpuCount += uint64(request.cpuCount) + group.memoryMB += request.memoryMB +} + +func groupIsFull(group groupedRoundRobinGroup, limits GroupedRoundRobinLimits) bool { + return group.sandboxCount >= limits.MaxSandboxCount || + (limits.MaxCPUCount > 0 && group.cpuCount >= uint64(limits.MaxCPUCount)) || + (limits.MaxMemoryMB > 0 && group.memoryMB >= limits.MaxMemoryMB) +} + +type strategyOptions struct { + groupedRoundRobinLimits GroupedRoundRobinLimits +} + +type StrategyOption func(*strategyOptions) + +func WithGroupedRoundRobinLimits(limits GroupedRoundRobinLimits) StrategyOption { + return func(options *strategyOptions) { + options.groupedRoundRobinLimits = limits + } +} + +func NewStrategy(name string, opts ...StrategyOption) Strategy { + options := strategyOptions{ + groupedRoundRobinLimits: GroupedRoundRobinLimits{MaxSandboxCount: 1}, + } + for _, opt := range opts { + opt(&options) + } + + switch strings.ToLower(strings.TrimSpace(name)) { case "random": return NewRandomStrategy() + case "grouped_round_robin": + return NewGroupedRoundRobinStrategy(options.groupedRoundRobinLimits) case "round_robin": fallthrough default: diff --git a/services/shared/config/config.go b/services/shared/config/config.go index 7977084e..18610eb2 100644 --- a/services/shared/config/config.go +++ b/services/shared/config/config.go @@ -57,6 +57,14 @@ type NodeResourceLimit struct { MaxAllocatedMemoryBytesIncludingPaused *uint64 `json:"max_allocated_memory_bytes_including_paused"` } +// GroupedRoundRobinConfig bounds each same-workload placement group. A group is +// closed before adding a sandbox that would exceed any enabled limit. +type GroupedRoundRobinConfig struct { + MaxSandboxCount uint32 `json:"max_sandbox_count"` + MaxCPUCount uint32 `json:"max_cpu_count"` + MaxMemoryMB uint64 `json:"max_memory_mb"` +} + type SchedulerConfig struct { GRPCListenAddr string `json:"grpc_listen_addr"` MetricsListenAddr string `json:"metrics_listen_addr"` @@ -69,6 +77,7 @@ type SchedulerConfig struct { Nodes []Node `json:"nodes"` Discovery SchedulerDiscoveryConfig `json:"discovery"` NodeResourceLimit *NodeResourceLimit `json:"node_resource_limit"` + GroupedRoundRobin GroupedRoundRobinConfig `json:"grouped_round_robin"` } func (s *SchedulerConfig) UnmarshalJSON(data []byte) error { @@ -84,6 +93,7 @@ func (s *SchedulerConfig) UnmarshalJSON(data []byte) error { Nodes *[]Node `json:"nodes"` Discovery *SchedulerDiscoveryConfig `json:"discovery"` NodeResourceLimit *NodeResourceLimit `json:"node_resource_limit"` + GroupedRoundRobin *GroupedRoundRobinConfig `json:"grouped_round_robin"` } parsed := wire{} @@ -109,6 +119,9 @@ func (s *SchedulerConfig) UnmarshalJSON(data []byte) error { if parsed.NodeResourceLimit != nil { s.NodeResourceLimit = parsed.NodeResourceLimit } + if parsed.GroupedRoundRobin != nil { + s.GroupedRoundRobin = *parsed.GroupedRoundRobin + } if parsed.RedisAddr != nil { s.RedisAddr = *parsed.RedisAddr } @@ -353,6 +366,30 @@ func overrideWithEnv(cfg *Config) error { cfg.Scheduler.ArtifactLookupNodeLimit = limit } + if v := strings.TrimSpace(os.Getenv("SCHEDULER_GROUPED_ROUND_ROBIN_MAX_SANDBOX_COUNT")); v != "" { + limit, err := strconv.ParseUint(v, 10, 32) + if err != nil { + return fmt.Errorf("invalid SCHEDULER_GROUPED_ROUND_ROBIN_MAX_SANDBOX_COUNT %q: %w", v, err) + } + cfg.Scheduler.GroupedRoundRobin.MaxSandboxCount = uint32(limit) + } + + if v := strings.TrimSpace(os.Getenv("SCHEDULER_GROUPED_ROUND_ROBIN_MAX_CPU_COUNT")); v != "" { + limit, err := strconv.ParseUint(v, 10, 32) + if err != nil { + return fmt.Errorf("invalid SCHEDULER_GROUPED_ROUND_ROBIN_MAX_CPU_COUNT %q: %w", v, err) + } + cfg.Scheduler.GroupedRoundRobin.MaxCPUCount = uint32(limit) + } + + if v := strings.TrimSpace(os.Getenv("SCHEDULER_GROUPED_ROUND_ROBIN_MAX_MEMORY_MB")); v != "" { + limit, err := strconv.ParseUint(v, 10, 64) + if err != nil { + return fmt.Errorf("invalid SCHEDULER_GROUPED_ROUND_ROBIN_MAX_MEMORY_MB %q: %w", v, err) + } + cfg.Scheduler.GroupedRoundRobin.MaxMemoryMB = limit + } + if v := strings.TrimSpace(os.Getenv("GATEWAY_REQUEST_TIMEOUT")); v != "" { d, err := time.ParseDuration(v) if err != nil { @@ -443,6 +480,10 @@ func (c Config) validate(schedulerQueryOnly bool) error { } return nil } + if strings.EqualFold(strings.TrimSpace(c.Scheduler.Strategy), "grouped_round_robin") && + c.Scheduler.GroupedRoundRobin.MaxSandboxCount == 0 { + return errors.New("scheduler.grouped_round_robin.max_sandbox_count must be greater than zero for grouped_round_robin strategy") + } if c.Scheduler.ArtifactStoreCapacity <= 0 { return errors.New("scheduler.artifact_store_capacity must be greater than zero") } diff --git a/services/shared/config/config_test.go b/services/shared/config/config_test.go index dbe47f0f..73a58dae 100644 --- a/services/shared/config/config_test.go +++ b/services/shared/config/config_test.go @@ -38,6 +38,88 @@ func TestDefaultSchedulerDiscoveryModeIsStatic(t *testing.T) { if got := cfg.Scheduler.ArtifactLookupNodeLimit; got != 0 { t.Fatalf("expected scheduler artifact lookup node limit 0, got %d", got) } + if got := cfg.Scheduler.GroupedRoundRobin; got != (GroupedRoundRobinConfig{}) { + t.Fatalf("expected zero-value grouped round-robin config, got %+v", got) + } +} + +func TestLoadParsesSchedulerGroupedRoundRobin(t *testing.T) { + tmpDir := t.TempDir() + path := filepath.Join(tmpDir, "config.json") + content := `{ + "scheduler": { + "strategy": "grouped_round_robin", + "grouped_round_robin": { + "max_sandbox_count": 4, + "max_cpu_count": 8, + "max_memory_mb": 16384 + } + } + }` + if err := os.WriteFile(path, []byte(content), 0o644); err != nil { + t.Fatalf("write config file failed: %v", err) + } + + cfg, err := Load(path, "scheduler") + if err != nil { + t.Fatalf("load config failed: %v", err) + } + want := GroupedRoundRobinConfig{ + MaxSandboxCount: 4, + MaxCPUCount: 8, + MaxMemoryMB: 16384, + } + if got := cfg.Scheduler.GroupedRoundRobin; got != want { + t.Fatalf("grouped round-robin config = %+v, want %+v", got, want) + } +} + +func TestLoadRejectsGroupedRoundRobinWithoutSandboxGroupLimit(t *testing.T) { + tmpDir := t.TempDir() + path := filepath.Join(tmpDir, "config.json") + content := `{ + "scheduler": { + "strategy": "grouped_round_robin", + "grouped_round_robin": { + "max_cpu_count": 8, + "max_memory_mb": 16384 + } + } + }` + if err := os.WriteFile(path, []byte(content), 0o644); err != nil { + t.Fatalf("write config file failed: %v", err) + } + + if _, err := Load(path, "scheduler"); err == nil { + t.Fatal("expected grouped_round_robin strategy without max_sandbox_count to fail") + } +} + +func TestLoadAppliesSchedulerGroupedRoundRobinEnv(t *testing.T) { + t.Setenv("SCHEDULER_STRATEGY", "grouped_round_robin") + t.Setenv("SCHEDULER_GROUPED_ROUND_ROBIN_MAX_SANDBOX_COUNT", "3") + t.Setenv("SCHEDULER_GROUPED_ROUND_ROBIN_MAX_CPU_COUNT", "6") + t.Setenv("SCHEDULER_GROUPED_ROUND_ROBIN_MAX_MEMORY_MB", "12288") + + cfg, err := Load("", "scheduler") + if err != nil { + t.Fatalf("load config failed: %v", err) + } + if got := cfg.Scheduler.GroupedRoundRobin; got != (GroupedRoundRobinConfig{ + MaxSandboxCount: 3, + MaxCPUCount: 6, + MaxMemoryMB: 12288, + }) { + t.Fatalf("unexpected grouped round-robin config from env: %+v", got) + } +} + +func TestLoadRejectsInvalidSchedulerGroupedRoundRobinEnv(t *testing.T) { + t.Setenv("SCHEDULER_GROUPED_ROUND_ROBIN_MAX_SANDBOX_COUNT", "-1") + + if _, err := Load("", "scheduler"); err == nil { + t.Fatal("expected invalid grouped round-robin env to fail") + } } func TestLoadSchedulerAllowsQueryOnlyWithRedisWithoutNodes(t *testing.T) {