diff --git a/core/http/endpoints/localai/nodes.go b/core/http/endpoints/localai/nodes.go index d6c44e383..e91eda6f4 100644 --- a/core/http/endpoints/localai/nodes.go +++ b/core/http/endpoints/localai/nodes.go @@ -60,7 +60,10 @@ func GetNodeEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc { return func(c echo.Context) error { ctx := c.Request().Context() id := c.Param("id") - node, err := registry.Get(ctx, id) + // GetWithExtras (not Get) so the response carries the node's labels, + // loaded-model count, and in-flight total — the bare BackendNode keeps + // labels in a separate table, leaving the detail view's label list empty. + node, err := registry.GetWithExtras(ctx, id) if err != nil { return c.JSON(http.StatusNotFound, nodeError(http.StatusNotFound, "node not found")) } diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index aafee13cb..8470c67ed 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -673,6 +673,49 @@ func (r *NodeRegistry) Get(ctx context.Context, nodeID string) (*BackendNode, er return &node, nil } +// GetWithExtras returns a single node enriched with the same computed fields as +// ListWithExtras (labels, loaded-model count, in-flight total). The plain Get +// returns a bare BackendNode whose Labels live in a separate table, so the node +// detail view needs this to show a node's existing labels and live counts. +func (r *NodeRegistry) GetWithExtras(ctx context.Context, nodeID string) (*NodeWithExtras, error) { + node, err := r.Get(ctx, nodeID) + if err != nil { + return nil, err + } + + labels := make(map[string]string) + nodeLabels, err := r.GetNodeLabels(ctx, nodeID) + if err != nil { + xlog.Warn("GetWithExtras: failed to get labels", "node", nodeID, "error", err) + } else { + for _, l := range nodeLabels { + labels[l.Key] = l.Value + } + } + + var modelCount int64 + if err := r.db.WithContext(ctx).Model(&NodeModel{}). + Where("node_id = ? AND state = ?", nodeID, "loaded"). + Count(&modelCount).Error; err != nil { + xlog.Warn("GetWithExtras: failed to get model count", "node", nodeID, "error", err) + } + + var inFlight struct{ Total int } + if err := r.db.WithContext(ctx).Model(&NodeModel{}). + Select("COALESCE(SUM(in_flight), 0) as total"). + Where("node_id = ? AND state IN ?", nodeID, []string{"loaded", "unloading"}). + Scan(&inFlight).Error; err != nil { + xlog.Warn("GetWithExtras: failed to get in-flight count", "node", nodeID, "error", err) + } + + return &NodeWithExtras{ + BackendNode: *node, + ModelCount: int(modelCount), + InFlightCount: inFlight.Total, + Labels: labels, + }, nil +} + // GetByName returns a single node by name. func (r *NodeRegistry) GetByName(ctx context.Context, name string) (*BackendNode, error) { var node BackendNode diff --git a/core/services/nodes/registry_test.go b/core/services/nodes/registry_test.go index 6f43706db..589a62fa3 100644 --- a/core/services/nodes/registry_test.go +++ b/core/services/nodes/registry_test.go @@ -646,6 +646,38 @@ var _ = Describe("NodeRegistry", func() { }) }) + 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)