mirror of
https://github.com/mudler/LocalAI.git
synced 2026-08-02 03:20:12 -04:00
878a0d00a1
* fix(distributed): stop the probe reaper from orphaning busy backends
The reconciler's liveness probe is a 1s gRPC HealthCheck, and a single
failed probe deleted the model's node_models row. A backend that is
merely busy cannot answer it: single-threaded Python backends (video and
avatar generation) block for minutes inside one request, so the reaper
was deleting registry rows for backends that were alive and mid-request.
The model then vanished from the nodes page while it was still
generating, and because the row was gone the in-flight decrement had
nothing to decrement ("DecrementInFlight: no matching row or already
zero"). Every subsequent request re-routed and re-staged the full model
from scratch.
Two guards:
- Replicas with in-flight requests are excluded in SQL. A row that is
actively serving is proof of life, and the running request is
exactly what stops the backend from answering the probe.
- Idle replicas must miss three CONSECUTIVE probes before removal, so
a transient blip cannot orphan a live replica. A successful probe
resets the streak.
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Assisted-by: Claude:claude-opus-5 [Claude Code]
* fix(distributed): drop the local model stub when its last replica goes
In distributed mode every routed model leaves an in-process stub in the
frontend's ModelLoader, and DistributedModelStore.Range reports local
stubs UNION the registry rows. Every registry removal path deletes only
the DB row, so the stub outlived the replica and the model was reported
as loaded forever.
That is the "loaded on the home page, absent from every node" ghost:
/system reads the union and still sees the stub, while /api/nodes/models
reads the registry and correctly sees nothing. It never self-healed,
and both frontend replicas showed it independently.
The replica-removed chokepoint could not fix this as it stood, because
it held a SINGLE hook that the prefix cache already owned, and it was
registered only when the prefix cache was enabled. Registering a second
listener would have silently displaced the first.
- Turn replicaRemovedHook into a list (AddReplicaRemovedHook), so
independent subsystems can each register without displacing others.
- Add NewLocalStubInvalidator, which drops the local stub once no
healthy replica of the model remains anywhere in the cluster, and
wire it unconditionally in startup.
The stub is kept while another node still serves the model: the
frontend is right to consider it loaded, and each request re-routes
through SmartRouter to pick a live replica anyway.
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Assisted-by: Claude:claude-opus-5 [Claude Code]
* fix(distributed): stop staging checksum sidecars back to workers
The file transfer server writes a "<file>.sha256" sidecar next to every
file it accepts. The sender walked the model directory with no filter,
so it staged those sidecars too, and the receiver duly wrote a sidecar
for each sidecar. Every staging pass multiplied the tree:
config.json -> config.json.sha256 -> config.json.sha256.sha256 -> ...
One LongCat snapshot had grown to 498 files, 466 of them chained, up to
29 levels deep, and the staged file count climbed on every pass. This
inflates each transfer and grows disk without bound on both ends.
Skip hash sidecars in stageDirectory, and mirror the skip in
countStageableFiles so the progress bar still reaches 100%. The check is
"a sidecar sitting next to a real file" rather than a blanket suffix
ban, so a model that genuinely ships a .sha256 payload with no
corresponding base file is still transferred.
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Assisted-by: Claude:claude-opus-5 [Claude Code]
* fix(distributed): classify the liveness probe instead of gating on in_flight
The previous commit excluded replicas with in-flight requests from the probe
reaper. That was the wrong guard, and could invert the bug it fixed.
in_flight has no decrement guarantee: track() balances its increment with a
defer, but a frontend killed mid-request never runs it, and the load-time
reservation is released only when the first inference completes. Nothing
resets a leaked counter. Gating the reaper on it therefore meant a leaked
counter would shield a genuinely dead replica from ever being reaped.
Nor was patience alone a fix: three misses at the default interval is ~90s of
silence, while the generation that triggered this blocks for 15+ minutes.
The real conflation was in the probe itself. A gRPC HealthCheck against the
backend's serving port measures "is it idle enough to answer", not "does the
process exist", and probeLoadedModels discarded the error that tells them
apart. Because the gRPC client is lazy, the status code is decisive:
- DeadlineExceeded: transport fine, nothing serviced the RPC. Busy.
- Unavailable: nothing is listening. Gone.
ModelProber now returns a ProbeOutcome, and only ProbeUnreachable counts
toward the reap threshold. ProbeBusy clears the streak: it is evidence of
life. A blackholed network reads as busy too, deliberately, since whole-node
failure is the health monitor's job.
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Assisted-by: Claude:claude-opus-5 [Claude Code]
* feat(distributed): reconcile replicas against worker-reported processes
Probing a backend's own serving port cannot distinguish "busy" from "gone"
without inferring it from an error code. The worker can answer directly: it
spawned the process, holds the handle, and its reply is not blocked by
whatever that backend is doing.
Adds a models.running request-reply subject. The worker answers out of its
in-memory process table, reporting each live process as (modelID,
replicaIndex, address) — the supervisor's process keys are `modelID#replica`,
which is isomorphic to a NodeModel row, so the reconciler can diff the two
directly.
reconcileNodeProcesses runs before the port probe and reaps rows for models
the worker is not running. Models the worker vouches for get updated_at
bumped, which takes them out of the port prober's stale set entirely: that is
what keeps a backend deep in a long generation away from the probe in the
first place, rather than relying on classifying its silence after the fact.
A worker that does not answer is skipped, not assumed empty. A messaging
failure says nothing about the processes, and assuming the worst would delete
a node's rows on a transient NATS blip; the port probe stays as the fallback
for those nodes. Rows younger than probeStaleAfter are ignored so a freshly
created row is never judged against a process table that has not caught up.
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Assisted-by: Claude:claude-opus-5 [Claude Code]
* fix(distributed): stop in_flight leaking and pin replicas against eviction
A leaked in_flight counter is not cosmetic. FindLRUModel,
FindGlobalLRUModelWithZeroInFlight and the router's eviction query all require
in_flight = 0, so a replica whose counter never came back is pinned and its
VRAM is unreclaimable for the lifetime of the process.
Two halves.
The source: routing reserves in_flight = 1 at load time so a freshly loaded
replica is not evicted out from under the request that caused the load. That
reservation was released ONLY by the first inference completing, so a route
torn down before any inference ran (client disconnect, handler error, failure
between load and the backend call) stranded it. newRouteResult now wires the
reservation to a sync.Once fired by whichever comes first, the first inference
or route teardown, and replaces three copies of the old wiring.
The backstop: a sweeper for counters leaked by paths that cannot run a defer
at all, such as a frontend killed mid-request.
Identifying a leak by elapsed time alone is unsafe. IncrementInFlight stamps
last_used at request START and nothing moves it while the request runs, so a
long generation is indistinguishable from a leak by age, and resetting there
would expose a serving model to eviction. The probe supplies the missing bit:
a backend that answers a health check promptly is not inside a request,
because that is precisely what a busy one cannot do. Requiring the row to also
be idle for 30 minutes covers backends that serve in parallel and can answer
while working, since those keep last_used fresh through each new increment.
Two existing tests asserted the old behaviour ("No decrement on Release").
That assertion was the leak, so both now pin the release instead.
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Assisted-by: Claude:claude-opus-5 [Claude Code]
---------
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
1580 lines
71 KiB
Go
1580 lines
71 KiB
Go
package nodes
|
|
|
|
import (
|
|
"context"
|
|
"runtime"
|
|
"time"
|
|
|
|
. "github.com/onsi/ginkgo/v2"
|
|
. "github.com/onsi/gomega"
|
|
|
|
"github.com/mudler/LocalAI/core/services/testutil"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
var _ = Describe("NodeRegistry", func() {
|
|
var (
|
|
db *gorm.DB
|
|
registry *NodeRegistry
|
|
)
|
|
|
|
BeforeEach(func() {
|
|
if runtime.GOOS == "darwin" {
|
|
Skip("testcontainers requires Docker, not available on macOS CI")
|
|
}
|
|
db = testutil.SetupTestDB()
|
|
var err error
|
|
registry, err = NewNodeRegistry(db)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
})
|
|
|
|
// Helper to build a minimal BackendNode.
|
|
makeNode := func(name, address string, vram uint64) *BackendNode {
|
|
return &BackendNode{
|
|
Name: name,
|
|
NodeType: NodeTypeBackend,
|
|
Address: address,
|
|
TotalVRAM: vram,
|
|
AvailableVRAM: vram,
|
|
}
|
|
}
|
|
|
|
Describe("Register", func() {
|
|
It("sets StatusPending when autoApprove is false", func() {
|
|
node := makeNode("worker-1", "10.0.0.1:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, false)).To(Succeed())
|
|
Expect(node.Status).To(Equal(StatusPending))
|
|
|
|
fetched, err := registry.GetByName(context.Background(), "worker-1")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(fetched.Status).To(Equal(StatusPending))
|
|
})
|
|
|
|
It("sets StatusHealthy when autoApprove is true", func() {
|
|
node := makeNode("worker-2", "10.0.0.2:50051", 4_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(node.Status).To(Equal(StatusHealthy))
|
|
})
|
|
})
|
|
|
|
Describe("Re-registration", func() {
|
|
It("keeps a pending node pending on re-register with autoApprove=false", func() {
|
|
node := makeNode("re-pending", "10.0.0.3:50051", 4_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, false)).To(Succeed())
|
|
Expect(node.Status).To(Equal(StatusPending))
|
|
|
|
// Re-register same name, still no auto-approve
|
|
node2 := makeNode("re-pending", "10.0.0.3:50052", 4_000_000_000)
|
|
Expect(registry.Register(context.Background(), node2, false)).To(Succeed())
|
|
Expect(node2.Status).To(Equal(StatusPending))
|
|
|
|
// ID is preserved from original registration
|
|
Expect(node2.ID).To(Equal(node.ID))
|
|
})
|
|
|
|
It("restores a previously approved node to healthy on re-register with autoApprove=false", func() {
|
|
node := makeNode("re-approved", "10.0.0.4:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(node.Status).To(Equal(StatusHealthy))
|
|
|
|
// Simulate the node becoming unhealthy
|
|
Expect(registry.MarkUnhealthy(context.Background(), node.ID)).To(Succeed())
|
|
fetched, err := registry.GetByName(context.Background(), "re-approved")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(fetched.Status).To(Equal(StatusUnhealthy))
|
|
|
|
// Re-register with autoApprove=false — should restore to healthy
|
|
// because the node was previously approved (status != pending)
|
|
node2 := makeNode("re-approved", "10.0.0.4:50052", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), node2, false)).To(Succeed())
|
|
Expect(node2.Status).To(Equal(StatusHealthy))
|
|
})
|
|
})
|
|
|
|
Describe("ApproveNode", func() {
|
|
It("transitions a pending node to healthy", func() {
|
|
node := makeNode("approve-me", "10.0.0.5:50051", 4_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, false)).To(Succeed())
|
|
Expect(node.Status).To(Equal(StatusPending))
|
|
|
|
Expect(registry.ApproveNode(context.Background(), node.ID)).To(Succeed())
|
|
|
|
fetched, err := registry.Get(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(fetched.Status).To(Equal(StatusHealthy))
|
|
})
|
|
|
|
It("returns error for non-existent node ID", func() {
|
|
err := registry.ApproveNode(context.Background(), "non-existent-id")
|
|
Expect(err).To(HaveOccurred())
|
|
Expect(err.Error()).To(ContainSubstring("not found or not in pending status"))
|
|
})
|
|
|
|
It("returns error for an already-healthy node", func() {
|
|
node := makeNode("already-healthy", "10.0.0.6:50051", 4_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
err := registry.ApproveNode(context.Background(), node.ID)
|
|
Expect(err).To(HaveOccurred())
|
|
Expect(err.Error()).To(ContainSubstring("not found or not in pending status"))
|
|
})
|
|
})
|
|
|
|
Describe("MarkOffline", func() {
|
|
It("sets status to offline and clears model records", func() {
|
|
node := makeNode("offline-test", "10.0.0.7:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
// Load a model on the node
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "llama-7b", 0, "loaded", "10.0.0.7:50052", 0)).To(Succeed())
|
|
models, err := registry.GetNodeModels(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(models).To(HaveLen(1))
|
|
|
|
// Mark offline
|
|
Expect(registry.MarkOffline(context.Background(), node.ID)).To(Succeed())
|
|
|
|
fetched, err := registry.Get(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(fetched.Status).To(Equal(StatusOffline))
|
|
|
|
// Model records should be cleared
|
|
models, err = registry.GetNodeModels(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(models).To(BeEmpty())
|
|
})
|
|
|
|
It("returns error for non-existent node", func() {
|
|
err := registry.MarkOffline(context.Background(), "does-not-exist")
|
|
Expect(err).To(HaveOccurred())
|
|
Expect(err.Error()).To(ContainSubstring("not found"))
|
|
})
|
|
})
|
|
|
|
Describe("SetNodeModel ID stability", func() {
|
|
It("preserves the ID when called twice for the same node+model", func() {
|
|
node := makeNode("stable-id-node", "10.0.0.99:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "my-model", 0, "loaded", "10.0.0.99:50052", 0)).To(Succeed())
|
|
nm1, err := registry.GetNodeModel(context.Background(), node.ID, "my-model", 0)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
|
|
// Call again with different state/address
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "my-model", 0, "loaded", "10.0.0.99:50053", 0)).To(Succeed())
|
|
nm2, err := registry.GetNodeModel(context.Background(), node.ID, "my-model", 0)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
|
|
Expect(nm2.ID).To(Equal(nm1.ID), "ID should remain stable across SetNodeModel calls")
|
|
Expect(nm2.Address).To(Equal("10.0.0.99:50053"), "Address should be updated")
|
|
})
|
|
})
|
|
|
|
Describe("FindNodeWithVRAM", func() {
|
|
It("selects the node with sufficient VRAM", func() {
|
|
small := makeNode("small-gpu", "10.0.0.10:50051", 4_000_000_000)
|
|
big := makeNode("big-gpu", "10.0.0.11:50051", 16_000_000_000)
|
|
Expect(registry.Register(context.Background(), small, true)).To(Succeed())
|
|
Expect(registry.Register(context.Background(), big, true)).To(Succeed())
|
|
|
|
// Request 8 GB — only big-gpu qualifies
|
|
found, err := registry.FindNodeWithVRAM(context.Background(), 8_000_000_000)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(found.Name).To(Equal("big-gpu"))
|
|
})
|
|
|
|
It("returns error when no node has enough VRAM", func() {
|
|
small := makeNode("tiny-gpu", "10.0.0.12:50051", 2_000_000_000)
|
|
Expect(registry.Register(context.Background(), small, true)).To(Succeed())
|
|
|
|
_, err := registry.FindNodeWithVRAM(context.Background(), 32_000_000_000)
|
|
Expect(err).To(HaveOccurred())
|
|
})
|
|
})
|
|
|
|
Describe("FindIdleNode", func() {
|
|
It("returns the node with no loaded models", func() {
|
|
busy := makeNode("busy-node", "10.0.0.20:50051", 8_000_000_000)
|
|
idle := makeNode("idle-node", "10.0.0.21:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), busy, true)).To(Succeed())
|
|
Expect(registry.Register(context.Background(), idle, true)).To(Succeed())
|
|
|
|
// Load a model on the busy node
|
|
Expect(registry.SetNodeModel(context.Background(), busy.ID, "model-a", 0, "loaded", "", 0)).To(Succeed())
|
|
|
|
found, err := registry.FindIdleNode(context.Background())
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(found.Name).To(Equal("idle-node"))
|
|
})
|
|
|
|
It("returns error when all nodes have models loaded", func() {
|
|
n := makeNode("all-busy", "10.0.0.22:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), n, true)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), n.ID, "model-x", 0, "loaded", "", 0)).To(Succeed())
|
|
|
|
_, err := registry.FindIdleNode(context.Background())
|
|
Expect(err).To(HaveOccurred())
|
|
})
|
|
})
|
|
|
|
Describe("FindLeastLoadedNode", func() {
|
|
It("returns the node with fewer in-flight requests", func() {
|
|
heavy := makeNode("heavy-node", "10.0.0.30:50051", 8_000_000_000)
|
|
light := makeNode("light-node", "10.0.0.31:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), heavy, true)).To(Succeed())
|
|
Expect(registry.Register(context.Background(), light, true)).To(Succeed())
|
|
|
|
// Set up models with different in-flight counts
|
|
Expect(registry.SetNodeModel(context.Background(), heavy.ID, "model-a", 0, "loaded", "", 0)).To(Succeed())
|
|
Expect(registry.IncrementInFlight(context.Background(), heavy.ID, "model-a", 0)).To(Succeed())
|
|
Expect(registry.IncrementInFlight(context.Background(), heavy.ID, "model-a", 0)).To(Succeed())
|
|
Expect(registry.IncrementInFlight(context.Background(), heavy.ID, "model-a", 0)).To(Succeed())
|
|
|
|
Expect(registry.SetNodeModel(context.Background(), light.ID, "model-b", 0, "loaded", "", 0)).To(Succeed())
|
|
Expect(registry.IncrementInFlight(context.Background(), light.ID, "model-b", 0)).To(Succeed())
|
|
|
|
found, err := registry.FindLeastLoadedNode(context.Background())
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(found.Name).To(Equal("light-node"))
|
|
})
|
|
})
|
|
|
|
Describe("FindAndLockNodeWithModel", func() {
|
|
It("returns the correct node and increments in-flight", func() {
|
|
node := makeNode("lock-node", "10.0.0.40:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "my-model", 0, "loaded", "10.0.0.40:50052", 0)).To(Succeed())
|
|
|
|
foundNode, foundNM, err := registry.FindAndLockNodeWithModel(context.Background(), "my-model", nil, nil)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(foundNode.ID).To(Equal(node.ID))
|
|
Expect(foundNM.ModelName).To(Equal("my-model"))
|
|
|
|
// Verify in-flight was incremented
|
|
nm, err := registry.GetNodeModel(context.Background(), node.ID, "my-model", 0)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(nm.InFlight).To(Equal(1))
|
|
})
|
|
|
|
It("returns error when model is not loaded anywhere", func() {
|
|
_, _, err := registry.FindAndLockNodeWithModel(context.Background(), "nonexistent-model", nil, nil)
|
|
Expect(err).To(HaveOccurred())
|
|
})
|
|
|
|
It("selects the node with fewer in-flight when multiple exist", func() {
|
|
n1 := makeNode("lock-heavy", "10.0.0.41:50051", 8_000_000_000)
|
|
n2 := makeNode("lock-light", "10.0.0.42:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), n1, true)).To(Succeed())
|
|
Expect(registry.Register(context.Background(), n2, true)).To(Succeed())
|
|
|
|
Expect(registry.SetNodeModel(context.Background(), n1.ID, "shared-model", 0, "loaded", "", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), n2.ID, "shared-model", 0, "loaded", "", 0)).To(Succeed())
|
|
|
|
// Add in-flight to n1
|
|
Expect(registry.IncrementInFlight(context.Background(), n1.ID, "shared-model", 0)).To(Succeed())
|
|
Expect(registry.IncrementInFlight(context.Background(), n1.ID, "shared-model", 0)).To(Succeed())
|
|
|
|
foundNode, _, err := registry.FindAndLockNodeWithModel(context.Background(), "shared-model", nil, nil)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(foundNode.Name).To(Equal("lock-light"))
|
|
})
|
|
|
|
It("filters by candidateNodeIDs even when an excluded node has lower in_flight", func() {
|
|
// Reproduces the selector-mismatch loop: a model loaded on a node
|
|
// the selector now excludes (excluded) and on a node it includes
|
|
// (included). Without the filter the excluded node wins on
|
|
// in_flight ASC; with the filter the included node is returned
|
|
// directly so Route() can serve from its existing replica.
|
|
excluded := makeNode("excluded-node", "10.0.0.43:50051", 8_000_000_000)
|
|
included := makeNode("included-node", "10.0.0.44:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), excluded, true)).To(Succeed())
|
|
Expect(registry.Register(context.Background(), included, true)).To(Succeed())
|
|
|
|
Expect(registry.SetNodeModel(context.Background(), excluded.ID, "filtered-model", 0, "loaded", "", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), included.ID, "filtered-model", 0, "loaded", "", 0)).To(Succeed())
|
|
|
|
// Make `included` strictly busier than `excluded` so the unfiltered
|
|
// query would prefer the excluded one — proving the filter is
|
|
// what's steering the result, not the in_flight ordering.
|
|
Expect(registry.IncrementInFlight(context.Background(), included.ID, "filtered-model", 0)).To(Succeed())
|
|
Expect(registry.IncrementInFlight(context.Background(), included.ID, "filtered-model", 0)).To(Succeed())
|
|
|
|
foundNode, foundNM, err := registry.FindAndLockNodeWithModel(context.Background(), "filtered-model", []string{included.ID}, nil)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(foundNode.ID).To(Equal(included.ID))
|
|
Expect(foundNM.NodeID).To(Equal(included.ID))
|
|
})
|
|
|
|
It("round-robins between replicas when in_flight ties (last_used tiebreaker)", func() {
|
|
// Three replicas of the same model on three nodes, all with in_flight=0.
|
|
// Without the last_used tiebreaker, the node with the largest available_vram
|
|
// would win every pick and one node would take ~all the load. With it,
|
|
// each successful pick refreshes last_used so the next pick rotates to
|
|
// the oldest-used replica.
|
|
fat := makeNode("rr-fat", "10.0.0.50:50051", 24_000_000_000)
|
|
mid := makeNode("rr-mid", "10.0.0.51:50051", 16_000_000_000)
|
|
small := makeNode("rr-small", "10.0.0.52:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), fat, true)).To(Succeed())
|
|
Expect(registry.Register(context.Background(), mid, true)).To(Succeed())
|
|
Expect(registry.Register(context.Background(), small, true)).To(Succeed())
|
|
|
|
Expect(registry.SetNodeModel(context.Background(), fat.ID, "rr-model", 0, "loaded", "", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), mid.ID, "rr-model", 0, "loaded", "", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), small.ID, "rr-model", 0, "loaded", "", 0)).To(Succeed())
|
|
|
|
// Decrement back to 0 after each pick so the next call sees a tie.
|
|
// (FindAndLockNodeWithModel atomically increments to lock the row.)
|
|
picks := make([]string, 0, 9)
|
|
for i := 0; i < 9; i++ {
|
|
n, nm, err := registry.FindAndLockNodeWithModel(context.Background(), "rr-model", nil, nil)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
picks = append(picks, n.Name)
|
|
Expect(registry.DecrementInFlight(context.Background(), n.ID, "rr-model", nm.ReplicaIndex)).To(Succeed())
|
|
}
|
|
|
|
// Each replica should have been picked at least twice across 9 ties —
|
|
// proves we're rotating, not pinning to the largest-VRAM node.
|
|
counts := map[string]int{}
|
|
for _, p := range picks {
|
|
counts[p]++
|
|
}
|
|
Expect(counts["rr-fat"]).To(BeNumerically(">=", 2), "fat node was picked %d times across 9 ties: %v", counts["rr-fat"], picks)
|
|
Expect(counts["rr-mid"]).To(BeNumerically(">=", 2), "mid node was picked %d times across 9 ties: %v", counts["rr-mid"], picks)
|
|
Expect(counts["rr-small"]).To(BeNumerically(">=", 2), "small node was picked %d times across 9 ties: %v", counts["rr-small"], picks)
|
|
})
|
|
|
|
It("returns not-found when the model is loaded only on excluded nodes", func() {
|
|
loadedExcluded := makeNode("excl-only-node", "10.0.0.45:50051", 8_000_000_000)
|
|
emptyIncluded := makeNode("empty-included-node", "10.0.0.46:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), loadedExcluded, true)).To(Succeed())
|
|
Expect(registry.Register(context.Background(), emptyIncluded, true)).To(Succeed())
|
|
|
|
Expect(registry.SetNodeModel(context.Background(), loadedExcluded.ID, "no-match-model", 0, "loaded", "", 0)).To(Succeed())
|
|
|
|
// Filter restricts to a node that does not have the model — the
|
|
// query must return an error so Route() falls through to schedule
|
|
// a fresh load on a matching node instead of reusing the excluded
|
|
// replica.
|
|
_, _, err := registry.FindAndLockNodeWithModel(context.Background(), "no-match-model", []string{emptyIncluded.ID}, nil)
|
|
Expect(err).To(HaveOccurred())
|
|
})
|
|
|
|
It("agrees with PickBestReplica on a seeded dataset (policy mirror)", func() {
|
|
// Guard against drift between the SQL ORDER BY in
|
|
// FindAndLockNodeWithModel and the canonical Go implementation in
|
|
// PickBestReplica. The two layers will eventually diverge in
|
|
// caller (DB-backed atomic pick vs in-memory snapshot pick for the
|
|
// per-frontend rotating cache), but the policy itself must stay
|
|
// the single source of truth. If this test fails, update *both*
|
|
// sides — never just one.
|
|
//
|
|
// Scenario exercises all three tiers:
|
|
// - "loser-busy" has the most VRAM but in_flight=2 — loses tier 1.
|
|
// - "loser-recent" ties at in_flight=0 but its last_used is the
|
|
// newest of the in_flight=0 group — loses tier 2.
|
|
// - "winner-mid" and "winner-fat" both tie at in_flight=0 and
|
|
// share the oldest last_used — tier 3 decides: fattest wins.
|
|
loserBusy := makeNode("mirror-loser-busy", "10.0.0.70:50051", 32_000_000_000)
|
|
loserRecent := makeNode("mirror-loser-recent", "10.0.0.71:50051", 8_000_000_000)
|
|
winnerMid := makeNode("mirror-winner-mid", "10.0.0.72:50051", 16_000_000_000)
|
|
winnerFat := makeNode("mirror-winner-fat", "10.0.0.73:50051", 24_000_000_000)
|
|
for _, n := range []*BackendNode{loserBusy, loserRecent, winnerMid, winnerFat} {
|
|
Expect(registry.Register(context.Background(), n, true)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), n.ID, "mirror-model", 0, "loaded", "", 0)).To(Succeed())
|
|
}
|
|
|
|
// Force in_flight=2 on the "busy" node so tier 1 disqualifies it.
|
|
Expect(registry.IncrementInFlight(context.Background(), loserBusy.ID, "mirror-model", 0)).To(Succeed())
|
|
Expect(registry.IncrementInFlight(context.Background(), loserBusy.ID, "mirror-model", 0)).To(Succeed())
|
|
|
|
// Slam last_used to known values so the test is deterministic
|
|
// regardless of clock resolution between the helpers above.
|
|
base := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
|
|
set := func(id string, t time.Time) {
|
|
Expect(db.Model(&NodeModel{}).
|
|
Where("node_id = ? AND model_name = ?", id, "mirror-model").
|
|
Update("last_used", t).Error).To(Succeed())
|
|
}
|
|
set(loserBusy.ID, base) // newest doesn't matter — already disqualified by tier 1
|
|
set(loserRecent.ID, base.Add(time.Hour))
|
|
set(winnerMid.ID, base)
|
|
set(winnerFat.ID, base)
|
|
|
|
// Pull the same dataset both pickers will operate on. The Go
|
|
// picker is a faithful representation of the policy; the SQL is
|
|
// the production path.
|
|
var rows []NodeModel
|
|
Expect(db.Where("model_name = ? AND state = ?", "mirror-model", "loaded").
|
|
Find(&rows).Error).To(Succeed())
|
|
candidates := make([]ReplicaCandidate, 0, len(rows))
|
|
for _, nm := range rows {
|
|
var bn BackendNode
|
|
Expect(db.First(&bn, "id = ? AND status = ?", nm.NodeID, StatusHealthy).Error).To(Succeed())
|
|
candidates = append(candidates, ReplicaCandidate{
|
|
NodeID: nm.NodeID,
|
|
Address: bn.Address,
|
|
ReplicaIndex: nm.ReplicaIndex,
|
|
InFlight: nm.InFlight,
|
|
LastUsed: nm.LastUsed,
|
|
AvailableVRAM: bn.AvailableVRAM,
|
|
})
|
|
}
|
|
goPick := PickBestReplica(candidates)
|
|
Expect(goPick).ToNot(BeNil())
|
|
|
|
sqlNode, _, err := registry.FindAndLockNodeWithModel(context.Background(), "mirror-model", nil, nil)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
|
|
Expect(sqlNode.ID).To(Equal(goPick.NodeID),
|
|
"SQL ORDER BY picked %s; PickBestReplica picked %s — policy has drifted",
|
|
sqlNode.ID, goPick.NodeID)
|
|
// Sanity check: the policy says winner-fat wins on tier 3.
|
|
Expect(goPick.NodeID).To(Equal(winnerFat.ID))
|
|
})
|
|
})
|
|
|
|
Describe("FindAndLockNodeWithModel preference", func() {
|
|
var nodeA, nodeB *BackendNode
|
|
|
|
BeforeEach(func() {
|
|
nodeA = makeNode("pref-a", "10.0.0.70:50051", 8_000_000_000)
|
|
nodeB = makeNode("pref-b", "10.0.0.71:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), nodeA, true)).To(Succeed())
|
|
Expect(registry.Register(context.Background(), nodeB, true)).To(Succeed())
|
|
// Both loaded+healthy for model "pref-model", in_flight 0.
|
|
Expect(registry.SetNodeModel(context.Background(), nodeA.ID, "pref-model", 0, "loaded", "", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), nodeB.ID, "pref-model", 0, "loaded", "", 0)).To(Succeed())
|
|
})
|
|
|
|
It("locks the preferred node when eligible", func() {
|
|
node, nm, err := registry.FindAndLockNodeWithModel(context.Background(), "pref-model", nil, &RoutePreference{PreferredNodeID: nodeB.ID})
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(node.ID).To(Equal(nodeB.ID))
|
|
Expect(nm.NodeID).To(Equal(nodeB.ID))
|
|
|
|
// in_flight is incremented atomically via gorm.Expr, so verify the
|
|
// persisted value through a re-fetch (the returned struct mirrors
|
|
// the pre-increment read, like the default-pick path).
|
|
persisted, err := registry.GetNodeModel(context.Background(), nodeB.ID, "pref-model", 0)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(persisted.InFlight).To(Equal(1))
|
|
})
|
|
|
|
It("falls back to default order when preferred not loaded", func() {
|
|
node, _, err := registry.FindAndLockNodeWithModel(context.Background(), "pref-model", nil, &RoutePreference{PreferredNodeID: "ZZZ"})
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(node.ID).To(BeElementOf(nodeA.ID, nodeB.ID))
|
|
})
|
|
|
|
It("nil preference behaves like before", func() {
|
|
node, _, err := registry.FindAndLockNodeWithModel(context.Background(), "pref-model", nil, nil)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(node).ToNot(BeNil())
|
|
})
|
|
|
|
It("locks the EXACT preferred replica when the node hosts two replicas", func() {
|
|
// A single node hosts replica 0 and replica 1 of a model, both
|
|
// loaded+healthy. The preference must lock the SPECIFIC replica
|
|
// requested, not the least-loaded replica on the node.
|
|
node := makeNode("pref-multi", "10.0.0.72:50051", 16_000_000_000)
|
|
node.MaxReplicasPerModel = 2
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "multi-model", 0, "loaded", "addr0", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "multi-model", 1, "loaded", "addr1", 0)).To(Succeed())
|
|
|
|
// pref={node, 1} must lock replica 1 specifically.
|
|
gotNode, nm1, err := registry.FindAndLockNodeWithModel(context.Background(), "multi-model", nil,
|
|
&RoutePreference{PreferredNodeID: node.ID, PreferredReplica: 1})
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(gotNode.ID).To(Equal(node.ID))
|
|
Expect(nm1.ReplicaIndex).To(Equal(1))
|
|
|
|
// pref={node, 0} must lock replica 0 specifically.
|
|
_, nm0, err := registry.FindAndLockNodeWithModel(context.Background(), "multi-model", nil,
|
|
&RoutePreference{PreferredNodeID: node.ID, PreferredReplica: 0})
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(nm0.ReplicaIndex).To(Equal(0))
|
|
})
|
|
})
|
|
|
|
Describe("LoadedReplicaStats", func() {
|
|
var n1, n2, n3 *BackendNode
|
|
|
|
BeforeEach(func() {
|
|
n1 = makeNode("stats-1", "10.0.0.80:50051", 8_000_000_000)
|
|
n2 = makeNode("stats-2", "10.0.0.81:50051", 8_000_000_000)
|
|
n3 = makeNode("stats-3", "10.0.0.82:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), n1, true)).To(Succeed())
|
|
Expect(registry.Register(context.Background(), n2, true)).To(Succeed())
|
|
Expect(registry.Register(context.Background(), n3, true)).To(Succeed())
|
|
// n1 loaded+busy, n2 loaded+idle, n3 has a different model only.
|
|
Expect(registry.SetNodeModel(context.Background(), n1.ID, "stats-model", 0, "loaded", "10.0.0.80:6000", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), n2.ID, "stats-model", 0, "loaded", "10.0.0.81:6000", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), n3.ID, "other-model", 0, "loaded", "", 0)).To(Succeed())
|
|
Expect(registry.IncrementInFlight(context.Background(), n1.ID, "stats-model", 0)).To(Succeed())
|
|
Expect(registry.IncrementInFlight(context.Background(), n1.ID, "stats-model", 0)).To(Succeed())
|
|
})
|
|
|
|
It("returns loaded healthy replicas with in-flight counts", func() {
|
|
stats, err := registry.LoadedReplicaStats(context.Background(), "stats-model", nil)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(stats).To(HaveLen(2))
|
|
byNode := map[string]ReplicaCandidate{}
|
|
for _, s := range stats {
|
|
byNode[s.NodeID] = s
|
|
}
|
|
Expect(byNode).To(HaveKey(n1.ID))
|
|
Expect(byNode).To(HaveKey(n2.ID))
|
|
Expect(byNode[n1.ID].InFlight).To(Equal(2))
|
|
Expect(byNode[n2.ID].InFlight).To(Equal(0))
|
|
})
|
|
|
|
It("filters to the candidate node set when provided", func() {
|
|
stats, err := registry.LoadedReplicaStats(context.Background(), "stats-model", []string{n2.ID})
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(stats).To(HaveLen(1))
|
|
Expect(stats[0].NodeID).To(Equal(n2.ID))
|
|
})
|
|
|
|
It("excludes unhealthy nodes", func() {
|
|
Expect(registry.MarkUnhealthy(context.Background(), n1.ID)).To(Succeed())
|
|
stats, err := registry.LoadedReplicaStats(context.Background(), "stats-model", nil)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(stats).To(HaveLen(1))
|
|
Expect(stats[0].NodeID).To(Equal(n2.ID))
|
|
})
|
|
|
|
It("returns empty for a model with no loaded replicas", func() {
|
|
stats, err := registry.LoadedReplicaStats(context.Background(), "no-such-model", nil)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(stats).To(BeEmpty())
|
|
})
|
|
})
|
|
|
|
Describe("MarkHealthy and MarkUnhealthy round-trip", func() {
|
|
It("transitions healthy -> unhealthy -> healthy", func() {
|
|
node := makeNode("roundtrip-node", "10.0.0.60:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(node.Status).To(Equal(StatusHealthy))
|
|
|
|
// Mark unhealthy
|
|
Expect(registry.MarkUnhealthy(context.Background(), node.ID)).To(Succeed())
|
|
fetched, err := registry.Get(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(fetched.Status).To(Equal(StatusUnhealthy))
|
|
|
|
// Mark healthy again
|
|
Expect(registry.MarkHealthy(context.Background(), node.ID)).To(Succeed())
|
|
fetched, err = registry.Get(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(fetched.Status).To(Equal(StatusHealthy))
|
|
})
|
|
|
|
It("returns error for non-existent node", func() {
|
|
err := registry.MarkHealthy(context.Background(), "does-not-exist")
|
|
Expect(err).To(HaveOccurred())
|
|
Expect(err.Error()).To(ContainSubstring("not found"))
|
|
})
|
|
})
|
|
|
|
Describe("NodeLabel CRUD", func() {
|
|
It("sets and retrieves labels for a node", func() {
|
|
node := makeNode("label-node", "10.0.0.70:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
Expect(registry.SetNodeLabel(context.Background(), node.ID, "env", "prod")).To(Succeed())
|
|
Expect(registry.SetNodeLabel(context.Background(), node.ID, "region", "us-east")).To(Succeed())
|
|
|
|
labels, err := registry.GetNodeLabels(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(labels).To(HaveLen(2))
|
|
|
|
labelMap := make(map[string]string)
|
|
for _, l := range labels {
|
|
labelMap[l.Key] = l.Value
|
|
}
|
|
Expect(labelMap["env"]).To(Equal("prod"))
|
|
Expect(labelMap["region"]).To(Equal("us-east"))
|
|
})
|
|
|
|
It("overwrites existing label with same key", func() {
|
|
node := makeNode("label-overwrite", "10.0.0.71:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
Expect(registry.SetNodeLabel(context.Background(), node.ID, "env", "dev")).To(Succeed())
|
|
Expect(registry.SetNodeLabel(context.Background(), node.ID, "env", "prod")).To(Succeed())
|
|
|
|
labels, err := registry.GetNodeLabels(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(labels).To(HaveLen(1))
|
|
Expect(labels[0].Value).To(Equal("prod"))
|
|
})
|
|
|
|
It("removes a single label by key", func() {
|
|
node := makeNode("label-remove", "10.0.0.72:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
Expect(registry.SetNodeLabel(context.Background(), node.ID, "env", "prod")).To(Succeed())
|
|
Expect(registry.SetNodeLabel(context.Background(), node.ID, "region", "us-east")).To(Succeed())
|
|
|
|
Expect(registry.RemoveNodeLabel(context.Background(), node.ID, "env")).To(Succeed())
|
|
|
|
labels, err := registry.GetNodeLabels(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(labels).To(HaveLen(1))
|
|
Expect(labels[0].Key).To(Equal("region"))
|
|
})
|
|
|
|
It("SetNodeLabels replaces all labels", func() {
|
|
node := makeNode("label-replace", "10.0.0.73:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
Expect(registry.SetNodeLabel(context.Background(), node.ID, "old-key", "old-val")).To(Succeed())
|
|
|
|
newLabels := map[string]string{"new-a": "val-a", "new-b": "val-b"}
|
|
Expect(registry.SetNodeLabels(context.Background(), node.ID, newLabels)).To(Succeed())
|
|
|
|
labels, err := registry.GetNodeLabels(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(labels).To(HaveLen(2))
|
|
|
|
labelMap := make(map[string]string)
|
|
for _, l := range labels {
|
|
labelMap[l.Key] = l.Value
|
|
}
|
|
Expect(labelMap).To(Equal(newLabels))
|
|
})
|
|
})
|
|
|
|
Describe("GetWithExtras", func() {
|
|
It("returns the node enriched with its labels map", func() {
|
|
node := makeNode("extras-node", "10.0.0.80:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.SetNodeLabel(context.Background(), node.ID, "env", "prod")).To(Succeed())
|
|
Expect(registry.SetNodeLabel(context.Background(), node.ID, "region", "us-east")).To(Succeed())
|
|
|
|
got, err := registry.GetWithExtras(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(got).ToNot(BeNil())
|
|
Expect(got.ID).To(Equal(node.ID))
|
|
Expect(got.Name).To(Equal("extras-node"))
|
|
Expect(got.Labels).To(Equal(map[string]string{"env": "prod", "region": "us-east"}))
|
|
})
|
|
|
|
It("returns an empty (non-nil) labels map when the node has none", func() {
|
|
node := makeNode("extras-no-labels", "10.0.0.81:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
got, err := registry.GetWithExtras(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(got).ToNot(BeNil())
|
|
Expect(got.Labels).ToNot(BeNil())
|
|
Expect(got.Labels).To(BeEmpty())
|
|
})
|
|
|
|
It("returns an error for an unknown node", func() {
|
|
_, err := registry.GetWithExtras(context.Background(), "does-not-exist")
|
|
Expect(err).To(HaveOccurred())
|
|
})
|
|
})
|
|
|
|
Describe("FindNodesBySelector", func() {
|
|
It("returns nodes matching all labels in selector", func() {
|
|
n1 := makeNode("sel-match", "10.0.0.80:50051", 8_000_000_000)
|
|
n2 := makeNode("sel-nomatch", "10.0.0.81:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), n1, true)).To(Succeed())
|
|
Expect(registry.Register(context.Background(), n2, true)).To(Succeed())
|
|
|
|
Expect(registry.SetNodeLabel(context.Background(), n1.ID, "env", "prod")).To(Succeed())
|
|
Expect(registry.SetNodeLabel(context.Background(), n1.ID, "region", "us-east")).To(Succeed())
|
|
Expect(registry.SetNodeLabel(context.Background(), n2.ID, "env", "dev")).To(Succeed())
|
|
|
|
nodes, err := registry.FindNodesBySelector(context.Background(), map[string]string{"env": "prod", "region": "us-east"})
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(nodes).To(HaveLen(1))
|
|
Expect(nodes[0].Name).To(Equal("sel-match"))
|
|
})
|
|
|
|
It("returns empty when no nodes match", func() {
|
|
n := makeNode("sel-empty", "10.0.0.82:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), n, true)).To(Succeed())
|
|
Expect(registry.SetNodeLabel(context.Background(), n.ID, "env", "dev")).To(Succeed())
|
|
|
|
nodes, err := registry.FindNodesBySelector(context.Background(), map[string]string{"env": "prod"})
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(nodes).To(BeEmpty())
|
|
})
|
|
|
|
It("ignores unhealthy nodes", func() {
|
|
n := makeNode("sel-unhealthy", "10.0.0.83:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), n, true)).To(Succeed())
|
|
Expect(registry.SetNodeLabel(context.Background(), n.ID, "env", "prod")).To(Succeed())
|
|
Expect(registry.MarkUnhealthy(context.Background(), n.ID)).To(Succeed())
|
|
|
|
nodes, err := registry.FindNodesBySelector(context.Background(), map[string]string{"env": "prod"})
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(nodes).To(BeEmpty())
|
|
})
|
|
|
|
It("matches nodes with more labels than selector requires", func() {
|
|
n := makeNode("sel-superset", "10.0.0.84:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), n, true)).To(Succeed())
|
|
Expect(registry.SetNodeLabel(context.Background(), n.ID, "env", "prod")).To(Succeed())
|
|
Expect(registry.SetNodeLabel(context.Background(), n.ID, "region", "us-east")).To(Succeed())
|
|
Expect(registry.SetNodeLabel(context.Background(), n.ID, "tier", "gpu")).To(Succeed())
|
|
|
|
nodes, err := registry.FindNodesBySelector(context.Background(), map[string]string{"env": "prod"})
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(nodes).To(HaveLen(1))
|
|
Expect(nodes[0].Name).To(Equal("sel-superset"))
|
|
})
|
|
|
|
It("returns all healthy nodes for empty selector", func() {
|
|
n1 := makeNode("sel-all-1", "10.0.0.85:50051", 8_000_000_000)
|
|
n2 := makeNode("sel-all-2", "10.0.0.86:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), n1, true)).To(Succeed())
|
|
Expect(registry.Register(context.Background(), n2, true)).To(Succeed())
|
|
|
|
nodes, err := registry.FindNodesBySelector(context.Background(), map[string]string{})
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(len(nodes)).To(BeNumerically(">=", 2))
|
|
})
|
|
})
|
|
|
|
Describe("ModelSchedulingConfig CRUD", func() {
|
|
It("creates and retrieves a scheduling config", func() {
|
|
config := &ModelSchedulingConfig{
|
|
ModelName: "llama-7b",
|
|
NodeSelector: `{"gpu.vendor":"nvidia"}`,
|
|
MinReplicas: 1,
|
|
MaxReplicas: 3,
|
|
}
|
|
Expect(registry.SetModelScheduling(context.Background(), config)).To(Succeed())
|
|
Expect(config.ID).ToNot(BeEmpty())
|
|
|
|
fetched, err := registry.GetModelScheduling(context.Background(), "llama-7b")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(fetched).ToNot(BeNil())
|
|
Expect(fetched.ModelName).To(Equal("llama-7b"))
|
|
Expect(fetched.NodeSelector).To(Equal(`{"gpu.vendor":"nvidia"}`))
|
|
Expect(fetched.MinReplicas).To(Equal(1))
|
|
Expect(fetched.MaxReplicas).To(Equal(3))
|
|
})
|
|
|
|
It("updates existing config via SetModelScheduling", func() {
|
|
config := &ModelSchedulingConfig{
|
|
ModelName: "update-model",
|
|
MinReplicas: 1,
|
|
MaxReplicas: 2,
|
|
}
|
|
Expect(registry.SetModelScheduling(context.Background(), config)).To(Succeed())
|
|
|
|
config2 := &ModelSchedulingConfig{
|
|
ModelName: "update-model",
|
|
MinReplicas: 2,
|
|
MaxReplicas: 5,
|
|
}
|
|
Expect(registry.SetModelScheduling(context.Background(), config2)).To(Succeed())
|
|
|
|
fetched, err := registry.GetModelScheduling(context.Background(), "update-model")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(fetched.MinReplicas).To(Equal(2))
|
|
Expect(fetched.MaxReplicas).To(Equal(5))
|
|
})
|
|
|
|
It("persists and updates route policy and thresholds", func() {
|
|
err := registry.SetModelScheduling(context.Background(), &ModelSchedulingConfig{
|
|
ModelName: "prefix-cache-model", RoutePolicy: "prefix_cache",
|
|
BalanceAbsThreshold: 3, BalanceRelThreshold: 2.0, MinPrefixMatch: 0.4,
|
|
})
|
|
Expect(err).ToNot(HaveOccurred())
|
|
|
|
got, err := registry.GetModelScheduling(context.Background(), "prefix-cache-model")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(got.RoutePolicy).To(Equal("prefix_cache"))
|
|
Expect(got.BalanceAbsThreshold).To(Equal(3))
|
|
Expect(got.BalanceRelThreshold).To(BeNumerically("==", 2.0))
|
|
Expect(got.MinPrefixMatch).To(BeNumerically("==", 0.4))
|
|
|
|
// Update must not be dropped on conflict.
|
|
Expect(registry.SetModelScheduling(context.Background(), &ModelSchedulingConfig{
|
|
ModelName: "prefix-cache-model", RoutePolicy: "round_robin",
|
|
})).ToNot(HaveOccurred())
|
|
|
|
got, err = registry.GetModelScheduling(context.Background(), "prefix-cache-model")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(got.RoutePolicy).To(Equal("round_robin"))
|
|
})
|
|
|
|
It("lists all configs", func() {
|
|
Expect(registry.SetModelScheduling(context.Background(), &ModelSchedulingConfig{ModelName: "list-a", MinReplicas: 1})).To(Succeed())
|
|
Expect(registry.SetModelScheduling(context.Background(), &ModelSchedulingConfig{ModelName: "list-b", MaxReplicas: 2})).To(Succeed())
|
|
|
|
configs, err := registry.ListModelSchedulings(context.Background())
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(len(configs)).To(BeNumerically(">=", 2))
|
|
})
|
|
|
|
It("lists only auto-scaling configs", func() {
|
|
Expect(registry.SetModelScheduling(context.Background(), &ModelSchedulingConfig{ModelName: "auto-a", MinReplicas: 2})).To(Succeed())
|
|
Expect(registry.SetModelScheduling(context.Background(), &ModelSchedulingConfig{ModelName: "auto-b", MaxReplicas: 3})).To(Succeed())
|
|
Expect(registry.SetModelScheduling(context.Background(), &ModelSchedulingConfig{ModelName: "no-auto", NodeSelector: `{"env":"prod"}`})).To(Succeed())
|
|
|
|
configs, err := registry.ListAutoScalingConfigs(context.Background())
|
|
Expect(err).ToNot(HaveOccurred())
|
|
|
|
names := make([]string, len(configs))
|
|
for i, c := range configs {
|
|
names[i] = c.ModelName
|
|
}
|
|
Expect(names).To(ContainElement("auto-a"))
|
|
Expect(names).To(ContainElement("auto-b"))
|
|
Expect(names).ToNot(ContainElement("no-auto"))
|
|
})
|
|
|
|
It("deletes a config", func() {
|
|
Expect(registry.SetModelScheduling(context.Background(), &ModelSchedulingConfig{ModelName: "delete-me", MinReplicas: 1})).To(Succeed())
|
|
|
|
Expect(registry.DeleteModelScheduling(context.Background(), "delete-me")).To(Succeed())
|
|
|
|
fetched, err := registry.GetModelScheduling(context.Background(), "delete-me")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(fetched).To(BeNil())
|
|
})
|
|
|
|
It("returns nil for non-existent model", func() {
|
|
fetched, err := registry.GetModelScheduling(context.Background(), "does-not-exist")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(fetched).To(BeNil())
|
|
})
|
|
})
|
|
|
|
Describe("CountLoadedReplicas", func() {
|
|
It("returns correct count of loaded replicas", func() {
|
|
n1 := makeNode("replica-node-1", "10.0.0.90:50051", 8_000_000_000)
|
|
n2 := makeNode("replica-node-2", "10.0.0.91:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), n1, true)).To(Succeed())
|
|
Expect(registry.Register(context.Background(), n2, true)).To(Succeed())
|
|
|
|
Expect(registry.SetNodeModel(context.Background(), n1.ID, "counted-model", 0, "loaded", "", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), n2.ID, "counted-model", 0, "loaded", "", 0)).To(Succeed())
|
|
|
|
count, err := registry.CountLoadedReplicas(context.Background(), "counted-model")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(count).To(Equal(int64(2)))
|
|
})
|
|
|
|
It("excludes non-loaded states", func() {
|
|
n1 := makeNode("replica-loaded", "10.0.0.92:50051", 8_000_000_000)
|
|
n2 := makeNode("replica-loading", "10.0.0.93:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), n1, true)).To(Succeed())
|
|
Expect(registry.Register(context.Background(), n2, true)).To(Succeed())
|
|
|
|
Expect(registry.SetNodeModel(context.Background(), n1.ID, "state-model", 0, "loaded", "", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), n2.ID, "state-model", 0, "loading", "", 0)).To(Succeed())
|
|
|
|
count, err := registry.CountLoadedReplicas(context.Background(), "state-model")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(count).To(Equal(int64(1)))
|
|
})
|
|
})
|
|
|
|
Describe("DecrementInFlight", func() {
|
|
It("does not go below zero", func() {
|
|
node := makeNode("dec-node", "10.0.0.50:50051", 4_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "dec-model", 0, "loaded", "", 0)).To(Succeed())
|
|
|
|
// in_flight starts at 0 — decrement should be a no-op
|
|
Expect(registry.DecrementInFlight(context.Background(), node.ID, "dec-model", 0)).To(Succeed())
|
|
|
|
nm, err := registry.GetNodeModel(context.Background(), node.ID, "dec-model", 0)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(nm.InFlight).To(Equal(0))
|
|
})
|
|
|
|
It("decrements correctly from a positive value", func() {
|
|
node := makeNode("dec-node-2", "10.0.0.51:50051", 4_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "dec-model-2", 0, "loaded", "", 0)).To(Succeed())
|
|
|
|
Expect(registry.IncrementInFlight(context.Background(), node.ID, "dec-model-2", 0)).To(Succeed())
|
|
Expect(registry.IncrementInFlight(context.Background(), node.ID, "dec-model-2", 0)).To(Succeed())
|
|
|
|
nm, err := registry.GetNodeModel(context.Background(), node.ID, "dec-model-2", 0)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(nm.InFlight).To(Equal(2))
|
|
|
|
Expect(registry.DecrementInFlight(context.Background(), node.ID, "dec-model-2", 0)).To(Succeed())
|
|
|
|
nm, err = registry.GetNodeModel(context.Background(), node.ID, "dec-model-2", 0)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(nm.InFlight).To(Equal(1))
|
|
})
|
|
})
|
|
|
|
Describe("Schema defaults", func() {
|
|
// These tests pin the GORM defaults that the multi-replica refactor
|
|
// relies on. If a future migration changes a default, the
|
|
// reconciler/router will silently misbehave (e.g. capacity 0 instead
|
|
// of 1) — these assertions catch that at the migration boundary.
|
|
It("BackendNode.MaxReplicasPerModel defaults to 1", func() {
|
|
node := makeNode("schema-default-mrpm", "10.0.0.200:50051", 4_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
fetched, err := registry.Get(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(fetched.MaxReplicasPerModel).To(Equal(1),
|
|
"old workers don't send the field; default must preserve single-replica behavior")
|
|
})
|
|
|
|
It("BackendNode.ReservedVRAM defaults to 0", func() {
|
|
node := makeNode("schema-default-reserved", "10.0.0.201:50051", 4_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
fetched, err := registry.Get(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(fetched.ReservedVRAM).To(Equal(uint64(0)))
|
|
})
|
|
|
|
It("NodeModel.ReplicaIndex defaults to 0", func() {
|
|
node := makeNode("schema-default-replica", "10.0.0.202:50051", 4_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "default-replica-model", 0, "loaded", "127.0.0.1:50100", 0)).To(Succeed())
|
|
|
|
nm, err := registry.GetNodeModel(context.Background(), node.ID, "default-replica-model", 0)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(nm).ToNot(BeNil())
|
|
Expect(nm.ReplicaIndex).To(Equal(0))
|
|
})
|
|
|
|
It("ModelSchedulingConfig.UnsatisfiableUntil is nullable and defaults to nil", func() {
|
|
cfg := &ModelSchedulingConfig{
|
|
ModelName: "schema-default-unsat",
|
|
MinReplicas: 1,
|
|
}
|
|
Expect(registry.SetModelScheduling(context.Background(), cfg)).To(Succeed())
|
|
|
|
fetched, err := registry.GetModelScheduling(context.Background(), "schema-default-unsat")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(fetched).ToNot(BeNil())
|
|
Expect(fetched.UnsatisfiableUntil).To(BeNil())
|
|
Expect(fetched.UnsatisfiableTicks).To(Equal(0))
|
|
})
|
|
})
|
|
|
|
Describe("Multi-replica registry", func() {
|
|
// PR2 tests: SetNodeModel with distinct replica indexes creates distinct
|
|
// rows; per-row mutations (Remove, Increment, Decrement, Touch) target
|
|
// only their indexed row so siblings are not orphaned.
|
|
It("SetNodeModel(replicaIndex=0) then SetNodeModel(replicaIndex=1) creates two distinct rows", func() {
|
|
node := makeNode("multi-1", "10.0.0.210:50051", 16_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "multi-model", 0, "loaded", "127.0.0.1:50100", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "multi-model", 1, "loaded", "127.0.0.1:50101", 0)).To(Succeed())
|
|
|
|
models, err := registry.GetNodeModels(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(models).To(HaveLen(2))
|
|
|
|
byIdx := map[int]NodeModel{}
|
|
for _, m := range models {
|
|
byIdx[m.ReplicaIndex] = m
|
|
}
|
|
Expect(byIdx[0].Address).To(Equal("127.0.0.1:50100"))
|
|
Expect(byIdx[1].Address).To(Equal("127.0.0.1:50101"))
|
|
Expect(byIdx[0].ID).ToNot(Equal(byIdx[1].ID))
|
|
})
|
|
|
|
It("RemoveNodeModel(replicaIndex=0) leaves replica 1 intact", func() {
|
|
node := makeNode("multi-2", "10.0.0.211:50051", 16_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "kept-model", 0, "loaded", "127.0.0.1:50110", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "kept-model", 1, "loaded", "127.0.0.1:50111", 0)).To(Succeed())
|
|
|
|
Expect(registry.RemoveNodeModel(context.Background(), node.ID, "kept-model", 0)).To(Succeed())
|
|
|
|
// Sibling replica must still exist — this was the latent bug pre-PR2:
|
|
// the WHERE clause matched both rows and orphaned the healthy sibling.
|
|
survivor, err := registry.GetNodeModel(context.Background(), node.ID, "kept-model", 1)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(survivor).ToNot(BeNil())
|
|
Expect(survivor.Address).To(Equal("127.0.0.1:50111"))
|
|
|
|
// Replica 0 is gone
|
|
_, err = registry.GetNodeModel(context.Background(), node.ID, "kept-model", 0)
|
|
Expect(err).To(HaveOccurred())
|
|
})
|
|
|
|
It("RemoveAllNodeModelReplicas deletes every replica of the model on the node", func() {
|
|
node := makeNode("multi-3", "10.0.0.212:50051", 16_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "purge-model", 0, "loaded", "a", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "purge-model", 1, "loaded", "b", 0)).To(Succeed())
|
|
|
|
Expect(registry.RemoveAllNodeModelReplicas(context.Background(), node.ID, "purge-model")).To(Succeed())
|
|
|
|
models, err := registry.GetNodeModels(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(models).To(BeEmpty())
|
|
})
|
|
|
|
It("IncrementInFlight only updates the targeted replica row", func() {
|
|
node := makeNode("multi-4", "10.0.0.213:50051", 16_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "infl-model", 0, "loaded", "a", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "infl-model", 1, "loaded", "b", 0)).To(Succeed())
|
|
|
|
Expect(registry.IncrementInFlight(context.Background(), node.ID, "infl-model", 1)).To(Succeed())
|
|
Expect(registry.IncrementInFlight(context.Background(), node.ID, "infl-model", 1)).To(Succeed())
|
|
|
|
r0, err := registry.GetNodeModel(context.Background(), node.ID, "infl-model", 0)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(r0.InFlight).To(Equal(0), "replica 0 must not have been incremented")
|
|
|
|
r1, err := registry.GetNodeModel(context.Background(), node.ID, "infl-model", 1)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(r1.InFlight).To(Equal(2))
|
|
})
|
|
|
|
It("CountReplicasOnNode returns the per-(node, model) row count", func() {
|
|
node := makeNode("multi-5", "10.0.0.214:50051", 16_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "count-model", 0, "loaded", "a", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "count-model", 1, "loaded", "b", 0)).To(Succeed())
|
|
|
|
n, err := registry.CountReplicasOnNode(context.Background(), node.ID, "count-model")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(n).To(Equal(2))
|
|
})
|
|
|
|
It("NextFreeReplicaIndex returns the lowest unused index < maxSlots", func() {
|
|
node := makeNode("multi-6", "10.0.0.215:50051", 16_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
// Slot 0 free initially
|
|
idx, err := registry.NextFreeReplicaIndex(context.Background(), node.ID, "slot-model", 4)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(idx).To(Equal(0))
|
|
|
|
// Occupy 0 and 2 — next free is 1 (lowest gap)
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "slot-model", 0, "loaded", "a", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "slot-model", 2, "loaded", "c", 0)).To(Succeed())
|
|
idx, err = registry.NextFreeReplicaIndex(context.Background(), node.ID, "slot-model", 4)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(idx).To(Equal(1), "must allocate the lowest free index for compactness")
|
|
|
|
// Fill all 4 — must return ErrNoFreeSlot
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "slot-model", 1, "loaded", "b", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "slot-model", 3, "loaded", "d", 0)).To(Succeed())
|
|
_, err = registry.NextFreeReplicaIndex(context.Background(), node.ID, "slot-model", 4)
|
|
Expect(err).To(MatchError(ErrNoFreeSlot))
|
|
|
|
// maxSlots=0 always returns ErrNoFreeSlot
|
|
_, err = registry.NextFreeReplicaIndex(context.Background(), node.ID, "no-slots-model", 0)
|
|
Expect(err).To(MatchError(ErrNoFreeSlot))
|
|
})
|
|
})
|
|
|
|
Describe("AddReplicaRemovedHook", func() {
|
|
type removed struct {
|
|
model, node string
|
|
replica int
|
|
}
|
|
|
|
It("fires once with the specific replica after RemoveNodeModel", func() {
|
|
node := makeNode("hook-remove-one", "10.0.0.230:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "hook-model", 1, "loaded", "a", 0)).To(Succeed())
|
|
|
|
var fired []removed
|
|
registry.AddReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
|
|
fired = append(fired, removed{model: modelName, node: nodeID, replica: replicaIndex})
|
|
})
|
|
|
|
// RemoveNodeModel(replica 1) must fire with the SPECIFIC replica index.
|
|
Expect(registry.RemoveNodeModel(context.Background(), node.ID, "hook-model", 1)).To(Succeed())
|
|
Expect(fired).To(HaveLen(1))
|
|
Expect(fired[0]).To(Equal(removed{model: "hook-model", node: node.ID, replica: 1}))
|
|
})
|
|
|
|
It("fires once with replica<0 after RemoveAllNodeModelReplicas", func() {
|
|
node := makeNode("hook-remove-all", "10.0.0.231:50051", 16_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "hook-all-model", 0, "loaded", "a", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "hook-all-model", 1, "loaded", "b", 0)).To(Succeed())
|
|
|
|
var fired []removed
|
|
registry.AddReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
|
|
fired = append(fired, removed{model: modelName, node: nodeID, replica: replicaIndex})
|
|
})
|
|
|
|
// One call covers all replicas of that model on the node: a negative
|
|
// replica index signals "all replicas", and the consumer's
|
|
// InvalidateNode drops every entry for the (model, node) pair.
|
|
Expect(registry.RemoveAllNodeModelReplicas(context.Background(), node.ID, "hook-all-model")).To(Succeed())
|
|
Expect(fired).To(HaveLen(1))
|
|
Expect(fired[0].model).To(Equal("hook-all-model"))
|
|
Expect(fired[0].node).To(Equal(node.ID))
|
|
Expect(fired[0].replica).To(BeNumerically("<", 0))
|
|
})
|
|
|
|
It("does not panic when no hook is set", func() {
|
|
node := makeNode("hook-unset", "10.0.0.232:50051", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "no-hook-model", 0, "loaded", "a", 0)).To(Succeed())
|
|
|
|
Expect(func() {
|
|
Expect(registry.RemoveNodeModel(context.Background(), node.ID, "no-hook-model", 0)).To(Succeed())
|
|
Expect(registry.RemoveAllNodeModelReplicas(context.Background(), node.ID, "no-hook-model")).To(Succeed())
|
|
}).ToNot(Panic())
|
|
})
|
|
|
|
// firedModelSet collects the distinct model names the hook saw for the
|
|
// given node. The bulk node-scoped deletes below remove every replica of
|
|
// every model on the node in one statement, so the chokepoint must fire
|
|
// the hook once per distinct model name (the consumer's Invalidate
|
|
// drops all entries for that (model, node) pair).
|
|
seedTwoModels := func(node *BackendNode) {
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "model-a", 0, "loaded", "a0", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "model-a", 1, "loaded", "a1", 0)).To(Succeed())
|
|
Expect(registry.SetNodeModel(context.Background(), node.ID, "model-b", 0, "loaded", "b0", 0)).To(Succeed())
|
|
}
|
|
|
|
It("fires once per distinct model after MarkOffline", func() {
|
|
node := makeNode("hook-offline", "10.0.0.240:50051", 8_000_000_000)
|
|
seedTwoModels(node)
|
|
|
|
fired := map[removed]int{}
|
|
registry.AddReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
|
|
// Bulk node-scoped deletes signal "all replicas" with replica<0.
|
|
Expect(replicaIndex).To(BeNumerically("<", 0))
|
|
fired[removed{model: modelName, node: nodeID}]++
|
|
})
|
|
|
|
Expect(registry.MarkOffline(context.Background(), node.ID)).To(Succeed())
|
|
Expect(fired).To(HaveLen(2))
|
|
Expect(fired[removed{model: "model-a", node: node.ID}]).To(Equal(1))
|
|
Expect(fired[removed{model: "model-b", node: node.ID}]).To(Equal(1))
|
|
})
|
|
|
|
It("fires once per distinct model after MarkDraining", func() {
|
|
node := makeNode("hook-draining", "10.0.0.241:50051", 8_000_000_000)
|
|
seedTwoModels(node)
|
|
|
|
fired := map[removed]int{}
|
|
registry.AddReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
|
|
// Bulk node-scoped deletes signal "all replicas" with replica<0.
|
|
Expect(replicaIndex).To(BeNumerically("<", 0))
|
|
fired[removed{model: modelName, node: nodeID}]++
|
|
})
|
|
|
|
Expect(registry.MarkDraining(context.Background(), node.ID)).To(Succeed())
|
|
Expect(fired).To(HaveLen(2))
|
|
Expect(fired[removed{model: "model-a", node: node.ID}]).To(Equal(1))
|
|
Expect(fired[removed{model: "model-b", node: node.ID}]).To(Equal(1))
|
|
})
|
|
|
|
It("fires once per distinct model after Deregister", func() {
|
|
node := makeNode("hook-deregister", "10.0.0.242:50051", 8_000_000_000)
|
|
seedTwoModels(node)
|
|
|
|
fired := map[removed]int{}
|
|
registry.AddReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
|
|
// Bulk node-scoped deletes signal "all replicas" with replica<0.
|
|
Expect(replicaIndex).To(BeNumerically("<", 0))
|
|
fired[removed{model: modelName, node: nodeID}]++
|
|
})
|
|
|
|
Expect(registry.Deregister(context.Background(), node.ID)).To(Succeed())
|
|
Expect(fired).To(HaveLen(2))
|
|
Expect(fired[removed{model: "model-a", node: node.ID}]).To(Equal(1))
|
|
Expect(fired[removed{model: "model-b", node: node.ID}]).To(Equal(1))
|
|
})
|
|
|
|
It("fires once per distinct model when re-registration clears stale rows", func() {
|
|
node := makeNode("hook-reregister", "10.0.0.243:50051", 8_000_000_000)
|
|
seedTwoModels(node)
|
|
|
|
fired := map[removed]int{}
|
|
registry.AddReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
|
|
// Bulk node-scoped deletes signal "all replicas" with replica<0.
|
|
Expect(replicaIndex).To(BeNumerically("<", 0))
|
|
fired[removed{model: modelName, node: nodeID}]++
|
|
})
|
|
|
|
// Re-register the same node (same name): the re-register path
|
|
// clears the stale model rows, which must fire the hook.
|
|
again := makeNode("hook-reregister", "10.0.0.243:50052", 8_000_000_000)
|
|
Expect(registry.Register(context.Background(), again, true)).To(Succeed())
|
|
Expect(fired).To(HaveLen(2))
|
|
Expect(fired[removed{model: "model-a", node: node.ID}]).To(Equal(1))
|
|
Expect(fired[removed{model: "model-b", node: node.ID}]).To(Equal(1))
|
|
})
|
|
|
|
// Atomicity: the bulk node-scoped delete in MarkOffline/MarkDraining/
|
|
// re-register now captures the model names and deletes the rows inside a
|
|
// single transaction. A true SetNodeModel-between-capture-and-delete race
|
|
// can't be forced deterministically here, but we can assert the
|
|
// post-condition the transaction guarantees: the set of fired hooks
|
|
// equals exactly the set of node_models rows the operation removed, with
|
|
// nothing left behind. If the capture and delete ever saw inconsistent
|
|
// snapshots, either a surviving row (delete missed it) or a missing hook
|
|
// (capture missed it) would break one of these assertions.
|
|
It("MarkOffline fires hooks for exactly the rows it deletes (consistent snapshot)", func() {
|
|
node := makeNode("hook-atomic-offline", "10.0.0.244:50051", 8_000_000_000)
|
|
seedTwoModels(node)
|
|
|
|
// Capture what the transaction should remove, straight from the DB,
|
|
// before running the operation.
|
|
before, err := registry.GetNodeModels(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
expectedModels := map[string]struct{}{}
|
|
for _, nm := range before {
|
|
expectedModels[nm.ModelName] = struct{}{}
|
|
}
|
|
Expect(expectedModels).To(HaveLen(2), "seed should create two distinct models")
|
|
|
|
fired := map[string]struct{}{}
|
|
registry.AddReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
|
|
Expect(nodeID).To(Equal(node.ID))
|
|
Expect(replicaIndex).To(BeNumerically("<", 0))
|
|
fired[modelName] = struct{}{}
|
|
})
|
|
|
|
Expect(registry.MarkOffline(context.Background(), node.ID)).To(Succeed())
|
|
|
|
// Hooks fired for exactly the distinct models that existed.
|
|
Expect(fired).To(Equal(expectedModels),
|
|
"hooks must fire for exactly the set of models the transaction deleted")
|
|
|
|
// And the delete actually emptied the node_models rows for the node:
|
|
// no row survives that did not get a hook.
|
|
after, err := registry.GetNodeModels(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(after).To(BeEmpty(), "no node_models row should survive the bulk delete")
|
|
})
|
|
})
|
|
|
|
Describe("ApplyAutoLabels", func() {
|
|
It("mirrors MaxReplicasPerModel as the node.replica-slots label", func() {
|
|
node := makeNode("auto-label-replicas", "10.0.0.220:50051", 16_000_000_000)
|
|
node.MaxReplicasPerModel = 4
|
|
node.GPUVendor = "nvidia"
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
registry.ApplyAutoLabels(context.Background(), node.ID, node)
|
|
|
|
labels, err := registry.GetNodeLabels(context.Background(), node.ID)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
byKey := map[string]string{}
|
|
for _, l := range labels {
|
|
byKey[l.Key] = l.Value
|
|
}
|
|
Expect(byKey).To(HaveKeyWithValue("node.replica-slots", "4"),
|
|
"selectors targeting fat nodes need this auto-label")
|
|
Expect(byKey).To(HaveKeyWithValue("gpu.vendor", "nvidia"))
|
|
})
|
|
|
|
It("defaults node.replica-slots to 1 when MaxReplicasPerModel is unset", func() {
|
|
node := makeNode("auto-label-default", "10.0.0.221:50051", 4_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
// Fetch back; default should be 1 (PR1 schema test)
|
|
fetched, _ := registry.Get(context.Background(), node.ID)
|
|
Expect(fetched.MaxReplicasPerModel).To(Equal(1))
|
|
|
|
registry.ApplyAutoLabels(context.Background(), node.ID, fetched)
|
|
|
|
labels, _ := registry.GetNodeLabels(context.Background(), node.ID)
|
|
byKey := map[string]string{}
|
|
for _, l := range labels {
|
|
byKey[l.Key] = l.Value
|
|
}
|
|
Expect(byKey).To(HaveKeyWithValue("node.replica-slots", "1"))
|
|
})
|
|
})
|
|
|
|
Describe("VRAM soft-reservation (PR5)", func() {
|
|
// These tests pin the soft-reservation contract: ReserveVRAM is the
|
|
// admission gate that prevents two concurrent scheduling decisions
|
|
// from over-committing the same node within one heartbeat window.
|
|
It("ReserveVRAM atomically deducts from effectively-free VRAM", func() {
|
|
node := makeNode("reserve-1", "10.0.0.230:50051", 10_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
Expect(registry.ReserveVRAM(context.Background(), node.ID, 3_000_000_000)).To(Succeed())
|
|
|
|
fetched, _ := registry.Get(context.Background(), node.ID)
|
|
Expect(fetched.ReservedVRAM).To(Equal(uint64(3_000_000_000)))
|
|
})
|
|
|
|
It("ReserveVRAM rejects when effectively-free VRAM is insufficient", func() {
|
|
node := makeNode("reserve-2", "10.0.0.231:50051", 5_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
// First reservation fits.
|
|
Expect(registry.ReserveVRAM(context.Background(), node.ID, 4_000_000_000)).To(Succeed())
|
|
// Second is too big — only 1 GB effectively free.
|
|
err := registry.ReserveVRAM(context.Background(), node.ID, 2_000_000_000)
|
|
Expect(err).To(MatchError(ErrInsufficientVRAM))
|
|
|
|
fetched, _ := registry.Get(context.Background(), node.ID)
|
|
Expect(fetched.ReservedVRAM).To(Equal(uint64(4_000_000_000)),
|
|
"failed reservation must not bump the column")
|
|
})
|
|
|
|
It("ReserveVRAM with bytes=0 is a no-op", func() {
|
|
node := makeNode("reserve-3", "10.0.0.232:50051", 1_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.ReserveVRAM(context.Background(), node.ID, 0)).To(Succeed())
|
|
|
|
fetched, _ := registry.Get(context.Background(), node.ID)
|
|
Expect(fetched.ReservedVRAM).To(Equal(uint64(0)))
|
|
})
|
|
|
|
It("ReleaseVRAM returns reserved bytes to the pool", func() {
|
|
node := makeNode("release-1", "10.0.0.233:50051", 10_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.ReserveVRAM(context.Background(), node.ID, 4_000_000_000)).To(Succeed())
|
|
|
|
Expect(registry.ReleaseVRAM(context.Background(), node.ID, 1_000_000_000)).To(Succeed())
|
|
|
|
fetched, _ := registry.Get(context.Background(), node.ID)
|
|
Expect(fetched.ReservedVRAM).To(Equal(uint64(3_000_000_000)))
|
|
})
|
|
|
|
It("ReleaseVRAM cannot underflow past zero", func() {
|
|
node := makeNode("release-underflow", "10.0.0.234:50051", 1_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
// No reservation; release is a guarded no-op rather than wrapping
|
|
// uint64 to a huge positive number.
|
|
Expect(registry.ReleaseVRAM(context.Background(), node.ID, 5_000_000_000)).To(Succeed())
|
|
|
|
fetched, _ := registry.Get(context.Background(), node.ID)
|
|
Expect(fetched.ReservedVRAM).To(Equal(uint64(0)))
|
|
})
|
|
|
|
It("Heartbeat with available_vram resets reserved_vram to 0", func() {
|
|
node := makeNode("heartbeat-reset", "10.0.0.235:50051", 10_000_000_000)
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.ReserveVRAM(context.Background(), node.ID, 5_000_000_000)).To(Succeed())
|
|
|
|
fresh := uint64(8_000_000_000)
|
|
Expect(registry.Heartbeat(context.Background(), node.ID, &HeartbeatUpdate{AvailableVRAM: &fresh})).To(Succeed())
|
|
|
|
fetched, _ := registry.Get(context.Background(), node.ID)
|
|
Expect(fetched.AvailableVRAM).To(Equal(fresh),
|
|
"heartbeat must overwrite available_vram with the worker's reading")
|
|
Expect(fetched.ReservedVRAM).To(Equal(uint64(0)),
|
|
"heartbeat must clear the soft reservation — worker is the source of truth")
|
|
})
|
|
|
|
It("UpdateMaxReplicasPerModel marks the value as a sticky override", func() {
|
|
// The original UX bug: workers default the flag to 1, so every
|
|
// re-registration silently reverted the admin's UI value. This
|
|
// test pins the fix.
|
|
node := &BackendNode{
|
|
Name: "override-survives",
|
|
NodeType: NodeTypeBackend,
|
|
Address: "10.0.0.240:50051",
|
|
MaxReplicasPerModel: 1,
|
|
}
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
|
|
// Admin sets capacity to 4 via the UI.
|
|
Expect(registry.UpdateMaxReplicasPerModel(context.Background(), node.ID, 4)).To(Succeed())
|
|
|
|
fetched, _ := registry.Get(context.Background(), node.ID)
|
|
Expect(fetched.MaxReplicasPerModel).To(Equal(4))
|
|
Expect(fetched.MaxReplicasPerModelManuallySet).To(BeTrue())
|
|
|
|
// Worker re-registers with its default of 1 (operator never set the flag).
|
|
restart := &BackendNode{
|
|
Name: "override-survives",
|
|
NodeType: NodeTypeBackend,
|
|
Address: "10.0.0.240:50051",
|
|
MaxReplicasPerModel: 1,
|
|
}
|
|
Expect(registry.Register(context.Background(), restart, true)).To(Succeed())
|
|
|
|
// Override must have survived.
|
|
fetched, _ = registry.Get(context.Background(), node.ID)
|
|
Expect(fetched.MaxReplicasPerModel).To(Equal(4),
|
|
"admin override must not be overwritten by worker re-registration")
|
|
Expect(fetched.MaxReplicasPerModelManuallySet).To(BeTrue())
|
|
})
|
|
|
|
It("ResetMaxReplicasPerModel hands control back to the worker", func() {
|
|
node := &BackendNode{
|
|
Name: "override-reset",
|
|
NodeType: NodeTypeBackend,
|
|
Address: "10.0.0.241:50051",
|
|
MaxReplicasPerModel: 1,
|
|
}
|
|
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
|
Expect(registry.UpdateMaxReplicasPerModel(context.Background(), node.ID, 4)).To(Succeed())
|
|
|
|
Expect(registry.ResetMaxReplicasPerModel(context.Background(), node.ID)).To(Succeed())
|
|
|
|
// Reset only flips the flag; the value stays until the worker
|
|
// re-registers (we don't presume to know what the worker wants).
|
|
fetched, _ := registry.Get(context.Background(), node.ID)
|
|
Expect(fetched.MaxReplicasPerModelManuallySet).To(BeFalse())
|
|
|
|
// Now worker re-registers with 8.
|
|
restart := &BackendNode{
|
|
Name: "override-reset",
|
|
NodeType: NodeTypeBackend,
|
|
Address: "10.0.0.241:50051",
|
|
MaxReplicasPerModel: 8,
|
|
}
|
|
Expect(registry.Register(context.Background(), restart, true)).To(Succeed())
|
|
|
|
fetched, _ = registry.Get(context.Background(), node.ID)
|
|
Expect(fetched.MaxReplicasPerModel).To(Equal(8),
|
|
"after reset, the worker's value should apply")
|
|
})
|
|
|
|
It("FindNodeWithVRAM honors the reservation", func() {
|
|
small := makeNode("find-vram-small", "10.0.0.236:50051", 5_000_000_000)
|
|
big := makeNode("find-vram-big", "10.0.0.237:50051", 20_000_000_000)
|
|
Expect(registry.Register(context.Background(), small, true)).To(Succeed())
|
|
Expect(registry.Register(context.Background(), big, true)).To(Succeed())
|
|
|
|
// Reserve almost all of the big node so its effective free
|
|
// drops below the request — small isn't big enough either —
|
|
// the call must return an error.
|
|
Expect(registry.ReserveVRAM(context.Background(), big.ID, 18_000_000_000)).To(Succeed())
|
|
|
|
_, err := registry.FindNodeWithVRAM(context.Background(), 8_000_000_000)
|
|
Expect(err).To(HaveOccurred(),
|
|
"reserved capacity must remove a node from VRAM-aware candidates")
|
|
})
|
|
})
|
|
|
|
Describe("ModelLoadInfo persistence (Bug-1)", func() {
|
|
It("survives every NodeModel row being removed", func() {
|
|
ctx := context.Background()
|
|
|
|
// One node with one loaded replica + per-replica blob (the legacy path).
|
|
node := makeNode("li-1", "10.0.1.1:50051", 8_000_000_000)
|
|
Expect(registry.Register(ctx, node, true)).To(Succeed())
|
|
Expect(registry.SetNodeModel(ctx, node.ID, "load-info-model", 0, "loaded", node.Address, 0)).To(Succeed())
|
|
Expect(registry.SetNodeModelLoadInfo(ctx, node.ID, "load-info-model", 0, "llama-cpp", []byte("opts-v1"))).To(Succeed())
|
|
|
|
// Persist per-model via the new path (the dispatch hook does this).
|
|
Expect(registry.UpsertModelLoadInfo(ctx, "load-info-model", "llama-cpp", []byte("opts-v1"))).To(Succeed())
|
|
|
|
// Simulate worker death + MarkOffline reaping: every NodeModel row gone.
|
|
Expect(registry.RemoveAllNodeModelReplicas(ctx, node.ID, "load-info-model")).To(Succeed())
|
|
|
|
bt, blob, err := registry.GetModelLoadInfo(ctx, "load-info-model")
|
|
Expect(err).ToNot(HaveOccurred(),
|
|
"per-model load info must survive every NodeModel row going away")
|
|
Expect(bt).To(Equal("llama-cpp"))
|
|
Expect(blob).To(Equal([]byte("opts-v1")))
|
|
})
|
|
|
|
It("ON CONFLICT updates backend type and opts (last-write-wins)", func() {
|
|
ctx := context.Background()
|
|
|
|
Expect(registry.UpsertModelLoadInfo(ctx, "lww", "llama-cpp", []byte("v1"))).To(Succeed())
|
|
Expect(registry.UpsertModelLoadInfo(ctx, "lww", "vllm", []byte("v2"))).To(Succeed())
|
|
|
|
bt, blob, err := registry.GetModelLoadInfo(ctx, "lww")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(bt).To(Equal("vllm"))
|
|
Expect(blob).To(Equal([]byte("v2")))
|
|
})
|
|
|
|
It("falls back to legacy NodeModel blob when no per-model row exists", func() {
|
|
// Pre-fix rolling-upgrade path: a frontend that ran before the new
|
|
// table existed only wrote the per-replica blob. The new
|
|
// GetModelLoadInfo must still find it so an upgrade doesn't
|
|
// regress the reconciler for already-loaded models.
|
|
ctx := context.Background()
|
|
|
|
node := makeNode("li-legacy", "10.0.1.2:50051", 8_000_000_000)
|
|
Expect(registry.Register(ctx, node, true)).To(Succeed())
|
|
Expect(registry.SetNodeModel(ctx, node.ID, "legacy-model", 0, "loaded", node.Address, 0)).To(Succeed())
|
|
Expect(registry.SetNodeModelLoadInfo(ctx, node.ID, "legacy-model", 0, "llama-cpp", []byte("legacy-opts"))).To(Succeed())
|
|
|
|
bt, blob, err := registry.GetModelLoadInfo(ctx, "legacy-model")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(bt).To(Equal("llama-cpp"))
|
|
Expect(blob).To(Equal([]byte("legacy-opts")))
|
|
})
|
|
|
|
It("returns ErrRecordNotFound when neither source has the model", func() {
|
|
ctx := context.Background()
|
|
_, _, err := registry.GetModelLoadInfo(ctx, "never-loaded")
|
|
Expect(err).To(MatchError(gorm.ErrRecordNotFound))
|
|
})
|
|
|
|
It("rejects empty model names", func() {
|
|
err := registry.UpsertModelLoadInfo(context.Background(), "", "llama-cpp", []byte("x"))
|
|
Expect(err).To(HaveOccurred())
|
|
})
|
|
})
|
|
})
|
|
|
|
var _ = Describe("ModelScheduling spread + seeding", func() {
|
|
var (
|
|
db *gorm.DB
|
|
registry *NodeRegistry
|
|
)
|
|
|
|
BeforeEach(func() {
|
|
if runtime.GOOS == "darwin" {
|
|
Skip("testcontainers requires Docker, not available on macOS CI")
|
|
}
|
|
db = testutil.SetupTestDB()
|
|
var err error
|
|
registry, err = NewNodeRegistry(db)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
})
|
|
|
|
It("persists and round-trips SpreadAll", func() {
|
|
Expect(registry.SetModelScheduling(context.Background(), &ModelSchedulingConfig{
|
|
ModelName: "m", SpreadAll: true,
|
|
})).To(Succeed())
|
|
got, err := registry.GetModelScheduling(context.Background(), "m")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(got.SpreadAll).To(BeTrue())
|
|
})
|
|
|
|
It("includes SpreadAll configs in ListAutoScalingConfigs", func() {
|
|
Expect(registry.SetModelScheduling(context.Background(), &ModelSchedulingConfig{
|
|
ModelName: "m", SpreadAll: true,
|
|
})).To(Succeed())
|
|
configs, err := registry.ListAutoScalingConfigs(context.Background())
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(configs).To(HaveLen(1))
|
|
Expect(configs[0].ModelName).To(Equal("m"))
|
|
})
|
|
|
|
It("seeds configs with authoritative upsert", func() {
|
|
Expect(registry.SetModelScheduling(context.Background(), &ModelSchedulingConfig{
|
|
ModelName: "m", MinReplicas: 9,
|
|
})).To(Succeed())
|
|
|
|
err := registry.SeedModelScheduling(context.Background(), []ModelSchedulingConfig{
|
|
{ModelName: "m", MinReplicas: 1, MaxReplicas: 2},
|
|
{ModelName: "n", SpreadAll: true},
|
|
})
|
|
Expect(err).ToNot(HaveOccurred())
|
|
|
|
m, _ := registry.GetModelScheduling(context.Background(), "m")
|
|
Expect(m.MinReplicas).To(Equal(1))
|
|
Expect(m.MaxReplicas).To(Equal(2))
|
|
Expect(m.SpreadAll).To(BeFalse())
|
|
|
|
n, _ := registry.GetModelScheduling(context.Background(), "n")
|
|
Expect(n.SpreadAll).To(BeTrue())
|
|
})
|
|
})
|