diff --git a/internal/cli/ctl.go b/internal/cli/ctl.go index 158056cf..fca09b6f 100644 --- a/internal/cli/ctl.go +++ b/internal/cli/ctl.go @@ -135,6 +135,7 @@ func newCtlCmd() *cobra.Command { newStatusCmd(client, &asJSON), newModelCmd(client), newScaleCmd(client), + newRestartCmd(client), newEngineCheckCmd(), newToolCheckCmd(), newWeightsCmd(client, &asJSON), @@ -370,6 +371,9 @@ func newReplicasCmd(client func() *ctlClient, asJSON *bool) *cobra.Command { if r.Stale { state += " (stale profile)" } + if r.Restart != "" { + state += " (restart " + r.Restart + ")" + } fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n", r.ID, r.Model, state, r.Pool, join(nodes), (time.Duration(r.AgeSec) * time.Second).String(), diff --git a/internal/cli/ctl_restart.go b/internal/cli/ctl_restart.go new file mode 100644 index 00000000..3b5febdb --- /dev/null +++ b/internal/cli/ctl_restart.go @@ -0,0 +1,48 @@ +package cli + +import ( + "encoding/json" + "errors" + "fmt" + + "github.com/spf13/cobra" + + "gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi" +) + +func newRestartCmd(client func() *ctlClient) *cobra.Command { + var ( + models []string + all bool + ) + c := &cobra.Command{ + Use: "restart [REPLICA...]", + Short: "Replace replicas with new ones, starting each replacement before stopping the old", + Long: "How a running engine picks up what only applies at start: an engine key,\n" + + "an edited profile, a new engine version.\n\n" + + "Each replacement is placed beside the replica it replaces, and the old one\n" + + "keeps serving until the new one is ready. Replicas of one model are\n" + + "replaced one at a time. Where nothing else has room for a second copy —\n" + + "the usual case on a small CPU fleet — the old one is stopped after the\n" + + "drain grace period (5 minutes by default) and the new one starts in its\n" + + "place; a model with only that replica is unserved while it loads.\n\n" + + " llmgw ctl restart qwen2.5-7b-3f2a91c0\n" + + " llmgw ctl restart --model qwen2.5-7b\n" + + " llmgw ctl restart --all", + RunE: func(cmd *cobra.Command, args []string) error { + if len(args) == 0 && len(models) == 0 && !all { + return errors.New("name replicas, --model, or --all") + } + body, _ := json.Marshal(adminapi.RestartRequest{Replicas: args, Models: models, All: all}) + var res adminapi.RestartResult + if err := client().post(cmd.Context(), "/api/v1/restart", body, &res); err != nil { + return err + } + fmt.Printf("%d replica(s) queued for restart; watch with: llmgw ctl replicas\n", len(res.Queued)) + return nil + }, + } + c.Flags().StringArrayVar(&models, "model", nil, "every replica of this model; repeatable") + c.Flags().BoolVar(&all, "all", false, "every ready replica in the cluster") + return c +} diff --git a/internal/control/adminapi/api.go b/internal/control/adminapi/api.go index ed041df7..447e09a1 100644 --- a/internal/control/adminapi/api.go +++ b/internal/control/adminapi/api.go @@ -77,6 +77,7 @@ func (a *API) Register(mux *http.ServeMux) { mux.HandleFunc("DELETE /api/v1/nodes/{id}", a.removeNode) mux.HandleFunc("POST /api/v1/fabric-test", a.fabricTest) mux.HandleFunc("GET /api/v1/nodes/{id}/logs", a.nodeLogs) + mux.HandleFunc("POST /api/v1/restart", a.restart) } // --- responses --- @@ -243,8 +244,11 @@ type ReplicaView struct { // and anyone else holding the admin listener who needs to reach an engine // directly. The client-facing pages never render it. EngineKey string `json:"engine_key,omitempty"` - Model string `json:"model"` - Profile string `json:"profile"` + // Restart is "queued" or "replacing" while an operator's restart of this + // replica is under way. + Restart string `json:"restart,omitempty"` + Model string `json:"model"` + Profile string `json:"profile"` // Stale is true when the profile has been edited since this replica was // placed, so it is running a definition the catalog no longer holds. // @@ -476,6 +480,7 @@ func (a *API) Replicas(model string) []ReplicaView { State: rep.State.String(), Pool: rep.Pool, Fabric: rep.Fabric, Endpoint: rep.Endpoint, Err: rep.Err, EngineKey: a.ctl.EngineKey(rep.ID), + Restart: a.ctl.Reconciler().RestartState(rep.ID), Stale: a.profileChanged(rep.Model, rep.Profile, rep.ProfileFP), AgeSec: int64(time.Since(rep.CreatedAt).Seconds()), IdleSec: -1, @@ -1044,3 +1049,39 @@ func (a *API) Topology(id string) (TopologyView, bool) { func (a *API) NVLink() nvlink.Report { return nvlink.Analyze(a.ctl.Inventory().Nodes()) } + +// RestartRequest names what to restart: replicas by id, every replica of some +// models, or — All — every ready replica in the cluster. +type RestartRequest struct { + Replicas []string `json:"replicas,omitempty"` + Models []string `json:"models,omitempty"` + All bool `json:"all,omitempty"` +} + +// RestartResult lists the replicas queued. +type RestartResult struct { + Queued []string `json:"queued"` +} + +func (a *API) restart(w http.ResponseWriter, r *http.Request) { + var req RestartRequest + if err := json.NewDecoder(io.LimitReader(r.Body, 1<<16)).Decode(&req); err != nil { + writeError(w, http.StatusBadRequest, "the restart request is not valid JSON: "+err.Error()) + return + } + ids := slices.Clone(req.Replicas) + for _, rep := range a.ctl.Reconciler().ReadyReplicas() { + if req.All || slices.Contains(req.Models, rep.Model) { + ids = append(ids, rep.ID) + } + } + if len(ids) == 0 { + writeError(w, http.StatusBadRequest, "nothing ready matches; name replicas, models, or all") + return + } + if err := a.ctl.Reconciler().Restart(ids...); err != nil { + writeError(w, http.StatusNotFound, err.Error()) + return + } + writeJSON(w, http.StatusAccepted, RestartResult{Queued: ids}) +} diff --git a/internal/control/reconciler/reconciler.go b/internal/control/reconciler/reconciler.go index c266799e..173ba2bb 100644 --- a/internal/control/reconciler/reconciler.go +++ b/internal/control/reconciler/reconciler.go @@ -3,6 +3,7 @@ package reconciler import ( "context" crypterand "crypto/rand" + "errors" "fmt" "log/slog" "math/rand/v2" @@ -204,6 +205,10 @@ type Reconciler struct { // because it is state about this moment and not about what the model is. // A model nobody has wanted has no entry and is simply not running. wanted map[string]time.Time + // restartQueue holds, per model, the replicas an operator asked to be + // restarted and not yet begun. One is replaced at a time: replacing them + // all at once would need the memory of every replica twice over. + restartQueue map[string][]string wake chan struct{} stop chan struct{} @@ -232,6 +237,7 @@ func New(cfg Config, cat *catalog.Live, inv *inventory.Store, disp Dispatcher, l attempts: map[string]int{}, reservedPorts: map[int]struct{}{}, wanted: map[string]time.Time{}, + restartQueue: map[string][]string{}, lastExplain: map[string]placement.Result{}, wake: make(chan struct{}, 1), stop: make(chan struct{}), @@ -708,19 +714,37 @@ func (r *Reconciler) reconcileModel(ctx context.Context, m catalog.Model) { // pointer stops being safe the moment it is in the map: an agent's status // report can mutate it from another goroutine at any time. r.mu.Lock() - var liveIDs, drainingIDs, avoid []string + var liveIDs, drainingIDs, retiringIDs, avoid []string + readyLive := 0 for _, rep := range r.replicas { if rep.Model != m.Name { continue } switch { + case rep.State.Counts() && !rep.RetiringAt.IsZero(): + retiringIDs = append(retiringIDs, rep.ID) + avoid = append(avoid, rep.NodeIDs()...) case rep.State.Counts(): liveIDs = append(liveIDs, rep.ID) avoid = append(avoid, rep.NodeIDs()...) + if rep.State == Ready { + readyLive++ + } case rep.State == Draining: drainingIDs = append(drainingIDs, rep.ID) } } + // The next restart an operator queued, once the last one is done. + if len(retiringIDs) == 0 { + if id, ok := r.nextRestartLocked(m.Name); ok { + r.replicas[id].RetiringAt = r.now() + retiringIDs = append(retiringIDs, id) + liveIDs = slices.DeleteFunc(liveIDs, func(x string) bool { return x == id }) + readyLive-- + r.log.Info("restarting a replica: placing its replacement first", + "replica", id, "model", m.Name) + } + } blockedUntil := r.nextAttempt[m.Name] r.mu.Unlock() @@ -763,6 +787,112 @@ func (r *Reconciler) reconcileModel(ctx context.Context, m catalog.Model) { r.stopReplica(ctx, id) } } + + r.retire(ctx, retiringIDs, readyLive, len(liveIDs), desired) +} + +// retire finishes restarting replicas. +// +// A retiring replica keeps serving until its replacement is ready, and is then +// drained like any other. When no replacement could be placed at all — +// nowhere else has the memory, which on a small CPU fleet is the usual case — +// it is stopped after DrainGrace so its replacement can take its place: a +// restart in place, with the model unserved while it loads if it had no other +// replica. Said in the log, since that is the one case with an outage. +func (r *Reconciler) retire(ctx context.Context, ids []string, readyLive, live, desired int) { + for _, id := range ids { + r.mu.Lock() + rep, ok := r.replicas[id] + var since time.Time + if ok { + since = rep.RetiringAt + } + r.mu.Unlock() + if !ok { + continue + } + switch { + case readyLive >= desired: + r.drainForRestart(ctx, id) + case live < desired && r.now().Sub(since) > r.cfg.DrainGrace: + r.log.Warn("restarting a replica in place: no room for its replacement beside it", + "replica", id, "after", r.cfg.DrainGrace) + r.drainForRestart(ctx, id) + r.stopReplica(ctx, id) + } + } +} + +func (r *Reconciler) drainForRestart(ctx context.Context, id string) { + r.mu.Lock() + rep, ok := r.replicas[id] + var shards []Shard + if ok { + rep.State, rep.DrainingAt, rep.RetiringAt = Draining, r.now(), time.Time{} + shards = slices.Clone(rep.Shards) + } + r.mu.Unlock() + for _, s := range shards { + if err := r.disp.DrainShard(ctx, s.NodeID, s.ID); err != nil { + r.log.Warn("draining a shard failed", "shard", s.ID, "err", err) + } + } +} + +// nextRestartLocked pops the next queued replica of a model that is still +// ready to be restarted. Called with mu held. +func (r *Reconciler) nextRestartLocked(model string) (string, bool) { + q := r.restartQueue[model] + for len(q) > 0 { + id := q[0] + q = q[1:] + if rep, ok := r.replicas[id]; ok && rep.State == Ready && rep.RetiringAt.IsZero() { + r.restartQueue[model] = q + return id, true + } + } + delete(r.restartQueue, model) + return "", false +} + +// ErrNoSuchReplica means no ready replica has that id. +var ErrNoSuchReplica = errors.New("no ready replica with that id") + +// Restart queues replicas to be replaced, each by a new one started beside it +// before it is stopped — to pick up an edited profile, an engine key, a new +// engine version. Replicas of one model are replaced one at a time. +func (r *Reconciler) Restart(ids ...string) error { + r.mu.Lock() + for _, id := range ids { + rep, ok := r.replicas[id] + if !ok || rep.State != Ready { + r.mu.Unlock() + return fmt.Errorf("%w: %s", ErrNoSuchReplica, id) + } + if !slices.Contains(r.restartQueue[rep.Model], id) && rep.RetiringAt.IsZero() { + r.restartQueue[rep.Model] = append(r.restartQueue[rep.Model], id) + } + } + r.mu.Unlock() + r.Notify() + return nil +} + +// RestartState says where a replica's restart stands: "queued", "replacing", +// or "" when none was asked for. +func (r *Reconciler) RestartState(id string) string { + r.mu.Lock() + defer r.mu.Unlock() + rep, ok := r.replicas[id] + switch { + case !ok: + return "" + case !rep.RetiringAt.IsZero(): + return "replacing" + case slices.Contains(r.restartQueue[rep.Model], id): + return "queued" + } + return "" } // scaleUp places and starts one replica per pass. diff --git a/internal/control/reconciler/reconciler_test.go b/internal/control/reconciler/reconciler_test.go index a2df6161..237695fe 100644 --- a/internal/control/reconciler/reconciler_test.go +++ b/internal/control/reconciler/reconciler_test.go @@ -1741,3 +1741,98 @@ func TestAnEngineIsStartedWithItsReplicasKey(t *testing.T) { } } } + +func replicaByID(h *harness, id string) (Replica, bool) { + for _, rep := range h.r.Replicas() { + if rep.ID == id { + return rep, true + } + } + return Replica{}, false +} + +func TestARestartedReplicaServesUntilItsReplacementIsReady(t *testing.T) { + // Restarting is how a running engine picks up a key, an edited profile or + // a new engine version. Stopping first would leave a one-replica model + // unserved for as long as its replacement takes to load. + h := newHarness(t, testModel(1), gpuNode("1", 4), gpuNode("2", 4)) + h.r.Reconcile(t.Context()) + h.markReady() + old := h.r.ReadyReplicas()[0].ID + + if err := h.r.Restart(old); err != nil { + t.Fatal(err) + } + h.r.Reconcile(t.Context()) + if rep, _ := replicaByID(h, old); rep.State != Ready || h.r.RestartState(old) != "replacing" { + t.Fatalf("the old replica stopped serving before its replacement was up: %+v", rep) + } + if n := len(h.r.Replicas()); n != 2 { + t.Fatalf("no replacement was placed beside it: %d replicas", n) + } + + h.markReady() + h.r.Reconcile(t.Context()) + h.r.Reconcile(t.Context()) + if _, still := replicaByID(h, old); still { + if rep, _ := replicaByID(h, old); rep.State != Draining && rep.State != Stopped { + t.Errorf("the old replica is %s after its replacement became ready", rep.State) + } + } + if got := len(h.r.ReadyReplicas()); got != 1 { + t.Errorf("%d ready replicas at the end; the model asks for 1", got) + } +} + +func TestARestartWithNoRoomBesideItHappensInPlace(t *testing.T) { + // One node, room for one replica: the replacement can only start once the + // old one has gone. Waiting forever would mean never restarting at all. + h := newHarness(t, testModel(1), gpuNode("1", 4)) + h.r.Reconcile(t.Context()) + h.markReady() + old := h.r.ReadyReplicas()[0].ID + if err := h.r.Restart(old); err != nil { + t.Fatal(err) + } + h.r.Reconcile(t.Context()) + if rep, _ := replicaByID(h, old); rep.State != Ready { + t.Fatalf("stopped before the grace period: %s", rep.State) + } + + h.advance(h.r.cfg.DrainGrace + time.Second) + h.r.Reconcile(t.Context()) + if rep, ok := replicaByID(h, old); ok && rep.State == Ready { + t.Fatal("still serving after the grace period with no replacement possible") + } + for i := 0; i < 3; i++ { + h.r.Reconcile(t.Context()) + h.markReady() + } + ready := h.r.ReadyReplicas() + if len(ready) != 1 || ready[0].ID == old { + t.Errorf("expected one new replica in its place, got %+v", ready) + } +} + +func TestReplicasOfAModelAreRestartedOneAtATime(t *testing.T) { + h := newHarness(t, testModel(2), gpuNode("1", 8), gpuNode("2", 8), gpuNode("3", 8)) + for i := 0; i < 3; i++ { + h.r.Reconcile(t.Context()) + h.markReady() + } + var ids []string + for _, rep := range h.r.ReadyReplicas() { + ids = append(ids, rep.ID) + } + if err := h.r.Restart(ids...); err != nil { + t.Fatal(err) + } + h.r.Reconcile(t.Context()) + states := []string{h.r.RestartState(ids[0]), h.r.RestartState(ids[1])} + if !slices.Contains(states, "replacing") || !slices.Contains(states, "queued") { + t.Errorf("expected one replacing and one queued, got %v", states) + } + if err := h.r.Restart("no-such"); err == nil { + t.Error("an unknown replica was accepted") + } +} diff --git a/internal/control/reconciler/replica.go b/internal/control/reconciler/replica.go index c9ff2ffc..da9fcdcc 100644 --- a/internal/control/reconciler/replica.go +++ b/internal/control/reconciler/replica.go @@ -156,6 +156,10 @@ type Replica struct { // DrainingAt is when the replica stopped taking new work, so the wait for a // replacement can be bounded. DrainingAt time.Time + // RetiringAt is set while the replica is being restarted: it keeps serving + // but no longer counts toward the model's replicas, so a replacement is + // placed beside it, and it is drained only once that replacement is ready. + RetiringAt time.Time // Attempt counts how many times this model has failed to come up in a row, // carried across rebuilds so the backoff does not reset on every attempt. Attempt int