fix(autoscaling): extract load model from Route() and use as well when doing autoscale

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)

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
 }

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).

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
 	}

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
 }