From 8452301105cd930ba468ae8ab4b794768aa4e332 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 7 Apr 2026 21:35:27 +0000 Subject: [PATCH] fix(autoscaling): extract load model from Route() and use as well when doing autoscale Signed-off-by: Ettore Di Giacinto --- core/services/nodes/interfaces.go | 2 + core/services/nodes/model_router_test.go | 7 ++ core/services/nodes/registry.go | 45 ++++++-- core/services/nodes/router.go | 137 +++++++++++++++-------- core/services/nodes/router_test.go | 8 ++ 5 files changed, 145 insertions(+), 54 deletions(-) diff --git a/core/services/nodes/interfaces.go b/core/services/nodes/interfaces.go index 02a6e7a920ed..5fadc50a7758 100644 --- a/core/services/nodes/interfaces.go +++ b/core/services/nodes/interfaces.go @@ -15,6 +15,8 @@ type ModelRouter interface { RemoveNodeModel(ctx context.Context, nodeID, modelName string) error TouchNodeModel(ctx context.Context, nodeID, modelName string) SetNodeModel(ctx context.Context, nodeID, modelName, state, address string, initialInFlight int) error + SetNodeModelLoadInfo(ctx context.Context, nodeID, modelName, backendType string, optsBlob []byte) error + GetModelLoadInfo(ctx context.Context, modelName string) (backendType string, optsBlob []byte, err error) FindNodeWithVRAM(ctx context.Context, minBytes uint64) (*BackendNode, error) FindIdleNode(ctx context.Context) (*BackendNode, error) FindLeastLoadedNode(ctx context.Context) (*BackendNode, error) diff --git a/core/services/nodes/model_router_test.go b/core/services/nodes/model_router_test.go index 946dccc235fa..5ae8e6a3cc7c 100644 --- a/core/services/nodes/model_router_test.go +++ b/core/services/nodes/model_router_test.go @@ -2,6 +2,7 @@ package nodes import ( "context" + "fmt" "sync" . "github.com/onsi/ginkgo/v2" @@ -49,6 +50,12 @@ func (f *fakeModelRouterForSmartRouter) TouchNodeModel(_ context.Context, _, _ s func (f *fakeModelRouterForSmartRouter) SetNodeModel(_ context.Context, _, _, _, _ string, _ int) error { return nil } +func (f *fakeModelRouterForSmartRouter) SetNodeModelLoadInfo(_ context.Context, _, _, _ string, _ []byte) error { + return nil +} +func (f *fakeModelRouterForSmartRouter) GetModelLoadInfo(_ context.Context, _ string) (string, []byte, error) { + return "", nil, fmt.Errorf("not found") +} func (f *fakeModelRouterForSmartRouter) FindNodeWithVRAM(_ context.Context, _ uint64) (*BackendNode, error) { return nil, nil } diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index b67e401f6a39..f56fcedccaeb 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -56,16 +56,18 @@ const ( // NodeModel tracks which models are loaded on which nodes. type NodeModel struct { - ID string `gorm:"primaryKey;size:36" json:"id"` - NodeID string `gorm:"index;size:36" json:"node_id"` - ModelName string `gorm:"index;size:255" json:"model_name"` - Address string `gorm:"size:255" json:"address"` // gRPC address for this model's backend process - State string `gorm:"size:32;default:idle" json:"state"` // loading, loaded, unloading, idle - InFlight int `json:"in_flight"` // number of active requests - LastUsed time.Time `json:"last_used"` - LoadingBy string `gorm:"size:36" json:"loading_by,omitempty"` // frontend ID that triggered loading - CreatedAt time.Time `json:"created_at"` - UpdatedAt time.Time `json:"updated_at"` + ID string `gorm:"primaryKey;size:36" json:"id"` + NodeID string `gorm:"index;size:36" json:"node_id"` + ModelName string `gorm:"index;size:255" json:"model_name"` + Address string `gorm:"size:255" json:"address"` // gRPC address for this model's backend process + State string `gorm:"size:32;default:idle" json:"state"` // loading, loaded, unloading, idle + InFlight int `json:"in_flight"` // number of active requests + LastUsed time.Time `json:"last_used"` + LoadingBy string `gorm:"size:36" json:"loading_by,omitempty"` // frontend ID that triggered loading + BackendType string `gorm:"size:128" json:"backend_type,omitempty"` // e.g. "llama-cpp"; used by reconciler to replicate loads + ModelOptsBlob []byte `gorm:"type:bytea" json:"-"` // serialized pb.ModelOptions for replica scale-ups + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` } // NodeLabel is a key-value label on a node (like K8s labels). @@ -414,6 +416,29 @@ func (r *NodeRegistry) SetNodeModel(ctx context.Context, nodeID, modelName, stat return result.Error } +// SetNodeModelLoadInfo stores the backend type and serialized model options on +// an existing NodeModel record. This metadata is used by the reconciler to +// replicate model loads during scale-up. +func (r *NodeRegistry) SetNodeModelLoadInfo(ctx context.Context, nodeID, modelName, backendType string, optsBlob []byte) error { + return r.db.WithContext(ctx).Model(&NodeModel{}). + Where("node_id = ? AND model_name = ?", nodeID, modelName). + Updates(map[string]any{"backend_type": backendType, "model_opts_blob": optsBlob}).Error +} + +// GetModelLoadInfo retrieves the stored backend type and serialized model +// options from any existing loaded replica. Returns gorm.ErrRecordNotFound +// if no replica has stored options. +func (r *NodeRegistry) GetModelLoadInfo(ctx context.Context, modelName string) (backendType string, optsBlob []byte, err error) { + var nm NodeModel + err = r.db.WithContext(ctx). + Where("model_name = ? AND state = ? AND model_opts_blob IS NOT NULL", modelName, "loaded"). + First(&nm).Error + if err != nil { + return "", nil, err + } + return nm.BackendType, nm.ModelOptsBlob, nil +} + // RemoveNodeModel removes a model association from a node. func (r *NodeRegistry) RemoveNodeModel(ctx context.Context, nodeID, modelName string) error { return r.db.WithContext(ctx).Where("node_id = ? AND model_name = ?", nodeID, modelName). diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index 8e7cc0359f75..e480a8bbb5d6 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -70,16 +70,97 @@ func NewSmartRouter(registry ModelRouter, opts SmartRouterOptions) *SmartRouter // Unloader returns the remote unloader adapter for external use. func (r *SmartRouter) Unloader() NodeCommandSender { return r.unloader } +// scheduleLoadResult holds the result of scheduling and loading a model on a node. +type scheduleLoadResult struct { + Node *BackendNode + Client grpc.Backend + BackendAddr string +} + +// scheduleAndLoad is the shared core for loading a model on a new node. +// Used by both Route() (for first-time loads) and ScheduleAndLoadModel() (for reconciler scale-ups). +// +// Steps: pick node → install backend → stage files → LoadModel → SetNodeModel. +func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, trackingKey, modelName string, + modelOpts *pb.ModelOptions, parallel bool, initialInFlight int) (*scheduleLoadResult, error) { + + node, backendAddr, err := r.scheduleNewModel(ctx, backendType, trackingKey, modelOpts) + if err != nil { + return nil, fmt.Errorf("no available nodes: %w", err) + } + + // Pre-stage model files via FileStager before loading + loadOpts := modelOpts + if r.fileStager != nil && modelOpts != nil { + staged, err := r.stageModelFiles(ctx, node, modelOpts, trackingKey) + if err != nil { + return nil, fmt.Errorf("staging model files for node %s: %w", node.Name, err) + } + loadOpts = staged + } + + client := r.buildClientForAddr(node, backendAddr, parallel) + + // Load the model on the remote node + if loadOpts != nil { + xlog.Info("Loading model on remote node", "node", node.Name, "model", modelName, "addr", backendAddr) + + loadCtx, cancel := context.WithTimeout(ctx, 5*time.Minute) + defer cancel() + + res, err := client.LoadModel(loadCtx, loadOpts) + if err != nil { + return nil, fmt.Errorf("loading model %s on node %s: %w", modelName, node.Name, err) + } + if !res.Success { + return nil, fmt.Errorf("loading model %s on node %s: %s", modelName, node.Name, res.Message) + } + } + + // Record the model as loaded on this node + if err := r.registry.SetNodeModel(ctx, node.ID, trackingKey, "loaded", backendAddr, initialInFlight); err != nil { + xlog.Warn("Failed to record model on node", "node", node.Name, "model", trackingKey, "error", err) + } + + // Store load metadata for future replica scale-ups by the reconciler + if modelOpts != nil { + if optsBlob, marshalErr := proto.Marshal(modelOpts); marshalErr == nil { + if storeErr := r.registry.SetNodeModelLoadInfo(ctx, node.ID, trackingKey, backendType, optsBlob); storeErr != nil { + xlog.Warn("Failed to store model load info", "node", node.Name, "model", trackingKey, "error", storeErr) + } + } + } + + return &scheduleLoadResult{Node: node, Client: client, BackendAddr: backendAddr}, nil +} + // ScheduleAndLoadModel implements ModelScheduler for the reconciler. -// It schedules a model on a suitable node (optionally from candidates) and loads it. +// It retrieves stored model options from an existing replica and performs the +// full load sequence (stage files, LoadModel, SetNodeModel) on a new node. func (r *SmartRouter) ScheduleAndLoadModel(ctx context.Context, modelName string, candidateNodeIDs []string) (*BackendNode, error) { - // Use scheduleNewModel with empty backend type and nil model options. - // The reconciler doesn't know the backend type — it will be determined by the model config. - node, _, err := r.scheduleNewModel(ctx, "", modelName, nil) + // Get load info from an existing replica (stored when Route() first loaded the model) + backendType, optsBlob, err := r.registry.GetModelLoadInfo(ctx, modelName) + if err != nil { + // No existing replica with stored opts — fall back to install-only. + // This happens on the very first load (before Route() has stored opts). + xlog.Warn("No stored model load info for reconciler scale-up, falling back to backend install only", + "model", modelName, "error", err) + node, _, schedErr := r.scheduleNewModel(ctx, "", modelName, nil) + return node, schedErr + } + + // Deserialize the stored model options + var modelOpts pb.ModelOptions + if err := proto.Unmarshal(optsBlob, &modelOpts); err != nil { + return nil, fmt.Errorf("unmarshalling stored model options for %s: %w", modelName, err) + } + + // initialInFlight=0: reconciler is pre-loading, not serving a request + result, err := r.scheduleAndLoad(ctx, backendType, modelName, modelName, &modelOpts, false, 0) if err != nil { return nil, err } - return node, nil + return result.Node, nil } // RouteResult contains the routing decision. @@ -195,53 +276,21 @@ func (r *SmartRouter) Route(ctx context.Context, modelID, modelName, backendType } } - // Still not loaded — proceed with scheduling - node, backendAddr, err := r.scheduleNewModel(ctx, backendType, trackingKey, modelOpts) + // Still not loaded — use shared schedule-and-load logic + result, err := r.scheduleAndLoad(ctx, backendType, trackingKey, modelName, modelOpts, parallel, 1) if err != nil { - return nil, fmt.Errorf("no available nodes: %w", err) - } - - // Pre-stage model files via FileStager before loading - if r.fileStager != nil && modelOpts != nil { - stagedOpts, err := r.stageModelFiles(ctx, node, modelOpts, trackingKey) - if err != nil { - return nil, fmt.Errorf("staging model files for node %s: %w", node.Name, err) - } - modelOpts = stagedOpts - } - - client := r.buildClientForAddr(node, backendAddr, parallel) - - // Load the model on this node - if modelOpts != nil { - xlog.Info("Loading model on remote node", "node", node.Name, "model", modelName, "addr", backendAddr) - - loadCtx, cancel := context.WithTimeout(ctx, 5*time.Minute) - defer cancel() - - res, err := client.LoadModel(loadCtx, modelOpts) - if err != nil { - return nil, fmt.Errorf("loading model %s on node %s: %w", modelName, node.Name, err) - } - if !res.Success { - return nil, fmt.Errorf("loading model %s on node %s: %s", modelName, node.Name, res.Message) - } - } - - // Record the model as loaded on this node with its per-process address - if err := r.registry.SetNodeModel(ctx, node.ID, trackingKey, "loaded", backendAddr, 1); err != nil { - xlog.Warn("Failed to record model on node", "node", node.Name, "model", trackingKey, "error", err) + return nil, err } - tracked := NewInFlightTrackingClient(client, r.registry, node.ID, trackingKey) + tracked := NewInFlightTrackingClient(result.Client, r.registry, result.Node.ID, trackingKey) tracked.OnFirstComplete(func() { - r.registry.DecrementInFlight(context.Background(), node.ID, trackingKey) + r.registry.DecrementInFlight(context.Background(), result.Node.ID, trackingKey) }) return &RouteResult{ - Node: node, + Node: result.Node, Client: tracked, Release: func() { - closeClient(client) + closeClient(result.Client) }, }, nil } diff --git a/core/services/nodes/router_test.go b/core/services/nodes/router_test.go index 7d366833791c..7f2840ba648b 100644 --- a/core/services/nodes/router_test.go +++ b/core/services/nodes/router_test.go @@ -142,6 +142,14 @@ func (f *fakeModelRouter) SetNodeModel(_ context.Context, nodeID, modelName, sta return nil } +func (f *fakeModelRouter) SetNodeModelLoadInfo(_ context.Context, _, _, _ string, _ []byte) error { + return nil +} + +func (f *fakeModelRouter) GetModelLoadInfo(_ context.Context, _ string) (string, []byte, error) { + return "", nil, fmt.Errorf("not found") +} + func (f *fakeModelRouter) FindNodeWithVRAM(_ context.Context, _ uint64) (*BackendNode, error) { return f.findVRAMNode, f.findVRAMErr }