wehub-resource-sync
122 lines
5.2 KiB
Go
122 lines
5.2 KiB
Go
package worker
|
|
|
|
import (
|
|
. "github.com/onsi/ginkgo/v2"
|
|
. "github.com/onsi/gomega"
|
|
)
|
|
|
|
var _ = Describe("Worker per-replica process keying", func() {
|
|
Describe("buildProcessKey", func() {
|
|
// Pin the supervisor's keying contract: distinct replica indexes for
|
|
// the same modelID produce distinct process keys, so the supervisor
|
|
// map can hold multiple processes for one model. Dropping the suffix
|
|
// would re-introduce the original flap (one model, one slot, churn).
|
|
DescribeTable("produces stable, distinct keys",
|
|
func(modelID, backend string, replica int, want string) {
|
|
Expect(buildProcessKey(modelID, backend, replica)).To(Equal(want))
|
|
},
|
|
Entry("modelID present, replica 0", "Qwen3-35B", "llama-cpp", 0, "Qwen3-35B#0"),
|
|
Entry("modelID present, replica 1", "Qwen3-35B", "llama-cpp", 1, "Qwen3-35B#1"),
|
|
Entry("falls back to backend when modelID empty", "", "llama-cpp", 0, "llama-cpp#0"),
|
|
Entry("backend fallback with replica 2", "", "llama-cpp", 2, "llama-cpp#2"),
|
|
)
|
|
|
|
It("makes replicas distinguishable", func() {
|
|
r0 := buildProcessKey("model-a", "llama-cpp", 0)
|
|
r1 := buildProcessKey("model-a", "llama-cpp", 1)
|
|
Expect(r0).ToNot(Equal(r1), "replicas of the same model must produce distinct keys")
|
|
})
|
|
})
|
|
|
|
Describe("registrationBody", func() {
|
|
It("includes max_replicas_per_model and the auto-label", func() {
|
|
cfg := &Config{
|
|
Addr: "worker.example.com:50051",
|
|
MaxReplicasPerModel: 4,
|
|
}
|
|
body := cfg.registrationBody()
|
|
|
|
Expect(body).To(HaveKey("max_replicas_per_model"))
|
|
Expect(body["max_replicas_per_model"]).To(Equal(4))
|
|
|
|
labels, ok := body["labels"].(map[string]string)
|
|
Expect(ok).To(BeTrue(), "labels must be present so selectors can target the slot count")
|
|
Expect(labels).To(HaveKeyWithValue("node.replica-slots", "4"))
|
|
})
|
|
|
|
It("coerces zero/unset MaxReplicasPerModel to 1", func() {
|
|
cfg := &Config{Addr: "worker.example.com:50051"}
|
|
body := cfg.registrationBody()
|
|
Expect(body["max_replicas_per_model"]).To(Equal(1),
|
|
"unset must default to single-replica behavior, not capacity 0")
|
|
|
|
labels := body["labels"].(map[string]string)
|
|
Expect(labels).To(HaveKeyWithValue("node.replica-slots", "1"))
|
|
})
|
|
|
|
It("preserves user-provided labels alongside the auto-label", func() {
|
|
cfg := &Config{
|
|
Addr: "worker.example.com:50051",
|
|
MaxReplicasPerModel: 2,
|
|
NodeLabels: "tier=fast,gpu=a100",
|
|
}
|
|
body := cfg.registrationBody()
|
|
labels := body["labels"].(map[string]string)
|
|
Expect(labels).To(HaveKeyWithValue("tier", "fast"))
|
|
Expect(labels).To(HaveKeyWithValue("gpu", "a100"))
|
|
Expect(labels).To(HaveKeyWithValue("node.replica-slots", "2"))
|
|
})
|
|
})
|
|
|
|
Describe("Process map lookup by bare model name", func() {
|
|
// Regression: PR #9583 changed the supervisor's map key from
|
|
// `modelID` to `modelID#replicaIndex`. The NATS backend.stop
|
|
// handler kept passing the bare modelID, so the lookup silently
|
|
// no-op'd — the worker process stayed alive after an admin
|
|
// "Unload model" click, and subsequent chats kept being served
|
|
// by the leftover process. The registry rows were gone, so the
|
|
// UI reported "no models loaded" while the model kept
|
|
// responding. resolveProcessKeys must turn a bare modelID into
|
|
// the actual replica process keys so stop/isRunning find the
|
|
// running processes.
|
|
It("resolves a bare modelID to its replica process keys", func() {
|
|
s := &backendSupervisor{
|
|
processes: map[string]*backendProcess{
|
|
"qwen3.6-35B#0": {addr: "127.0.0.1:50051"},
|
|
"qwen3.6-35B#1": {addr: "127.0.0.1:50052"},
|
|
"other-model#0": {addr: "127.0.0.1:50053"},
|
|
},
|
|
}
|
|
keys := s.resolveProcessKeys("qwen3.6-35B")
|
|
Expect(keys).To(ConsistOf("qwen3.6-35B#0", "qwen3.6-35B#1"),
|
|
"bare modelID must match all replica process keys")
|
|
|
|
// Bare modelID for a model with no live processes returns nothing.
|
|
Expect(s.resolveProcessKeys("not-loaded")).To(BeEmpty())
|
|
|
|
// Full processKey resolves to itself (per-replica callers stay precise).
|
|
Expect(s.resolveProcessKeys("qwen3.6-35B#0")).To(ConsistOf("qwen3.6-35B#0"))
|
|
|
|
// A processKey that doesn't exist returns nothing — no spurious
|
|
// prefix fallback when the caller was explicit.
|
|
Expect(s.resolveProcessKeys("qwen3.6-35B#9")).To(BeEmpty())
|
|
})
|
|
|
|
It("isRunning returns false when no replica matches", func() {
|
|
// We can only test the not-found path without a real *process.Process
|
|
// (IsAlive() requires PID introspection). That's enough to pin the
|
|
// regression — pre-fix, isRunning("qwen3.6-35B") would always
|
|
// return false because the map was keyed by "qwen3.6-35B#0".
|
|
// Post-fix, isRunning calls resolveProcessKeys first, so the
|
|
// per-replica lookup is exercised before the IsAlive probe.
|
|
s := &backendSupervisor{processes: map[string]*backendProcess{}}
|
|
Expect(s.isRunning("qwen3.6-35B")).To(BeFalse())
|
|
// resolveProcessKeys finds the replica entries (the lookup contract
|
|
// is what the backend.delete handler relies on); the IsAlive probe
|
|
// itself is exercised by the integration path in distributed mode.
|
|
s.processes["qwen3.6-35B#0"] = &backendProcess{addr: "127.0.0.1:50051"}
|
|
Expect(s.resolveProcessKeys("qwen3.6-35B")).To(ConsistOf("qwen3.6-35B#0"))
|
|
})
|
|
})
|
|
})
|