From 113b51eda246d2882ad00cda8cb2530ceb4a484e Mon Sep 17 00:00:00 2001 From: LordMathis Date: Sat, 18 Oct 2025 00:33:16 +0200 Subject: [PATCH] Refactor instance node handling to use a map --- pkg/instance/instance.go | 10 ++++++---- pkg/instance/instance_test.go | 8 ++++---- pkg/instance/options.go | 23 ++++++++++++++++++++++- pkg/manager/manager.go | 17 +++++++++++++---- pkg/manager/operations.go | 26 ++++++++++++++------------ pkg/server/handlers_instances.go | 15 ++++++++++++--- pkg/server/handlers_openai.go | 15 ++++++++++++--- 7 files changed, 83 insertions(+), 31 deletions(-) diff --git a/pkg/instance/instance.go b/pkg/instance/instance.go index da7cda7..1407eba 100644 --- a/pkg/instance/instance.go +++ b/pkg/instance/instance.go @@ -193,8 +193,10 @@ func (i *Instance) GetProxy() (*httputil.ReverseProxy, error) { // Remote instances should not use local proxy - they are handled by RemoteInstanceProxy opts := i.GetOptions() - if opts != nil && len(opts.Nodes) > 0 && opts.Nodes[0] != i.localNodeName { - return nil, fmt.Errorf("instance %s is a remote instance and should not use local proxy", i.Name) + if opts != nil && len(opts.Nodes) > 0 { + if _, isLocal := opts.Nodes[i.localNodeName]; !isLocal { + return nil, fmt.Errorf("instance %s is a remote instance and should not use local proxy", i.Name) + } } return i.proxy.GetProxy() @@ -303,8 +305,8 @@ func (i *Instance) IsRemote() bool { return false } - // If the first node is the local node, treat it as a local instance - if opts.Nodes[0] == i.localNodeName { + // If the local node is in the nodes map, treat it as a local instance + if _, isLocal := opts.Nodes[i.localNodeName]; isLocal { return false } diff --git a/pkg/instance/instance_test.go b/pkg/instance/instance_test.go index 7413b58..5260769 100644 --- a/pkg/instance/instance_test.go +++ b/pkg/instance/instance_test.go @@ -209,7 +209,7 @@ func TestSetOptions_PreservesNodes(t *testing.T) { // Create instance with initial nodes initialOptions := &instance.Options{ BackendType: backends.BackendTypeLlamaCpp, - Nodes: []string{"worker1"}, + Nodes: map[string]struct{}{"worker1": {}}, LlamaServerOptions: &llamacpp.LlamaServerOptions{ Model: "/path/to/model.gguf", Port: 8080, @@ -222,7 +222,7 @@ func TestSetOptions_PreservesNodes(t *testing.T) { // Try to update with different nodes updatedOptions := &instance.Options{ BackendType: backends.BackendTypeLlamaCpp, - Nodes: []string{"worker2"}, // Attempt to change node + Nodes: map[string]struct{}{"worker2": {}}, // Attempt to change node LlamaServerOptions: &llamacpp.LlamaServerOptions{ Model: "/path/to/new-model.gguf", Port: 8081, @@ -233,8 +233,8 @@ func TestSetOptions_PreservesNodes(t *testing.T) { opts := inst.GetOptions() // Nodes should remain unchanged - if len(opts.Nodes) != 1 || opts.Nodes[0] != "worker1" { - t.Errorf("Expected nodes to remain ['worker1'], got %v", opts.Nodes) + if _, exists := opts.Nodes["worker1"]; len(opts.Nodes) != 1 || !exists { + t.Errorf("Expected nodes to contain 'worker1', got %v", opts.Nodes) } // Other options should be updated diff --git a/pkg/instance/options.go b/pkg/instance/options.go index c5db58b..f8af9ee 100644 --- a/pkg/instance/options.go +++ b/pkg/instance/options.go @@ -10,6 +10,7 @@ import ( "llamactl/pkg/config" "log" "maps" + "slices" "sync" ) @@ -29,7 +30,7 @@ type Options struct { BackendType backends.BackendType `json:"backend_type"` BackendOptions map[string]any `json:"backend_options,omitempty"` - Nodes []string `json:"nodes,omitempty"` + Nodes map[string]struct{} `json:"-"` // Backend-specific options LlamaServerOptions *llamacpp.LlamaServerOptions `json:"-"` @@ -87,6 +88,7 @@ func (c *Options) UnmarshalJSON(data []byte) error { // Use anonymous struct to avoid recursion type Alias Options aux := &struct { + Nodes []string `json:"nodes,omitempty"` // Accept JSON array *Alias }{ Alias: (*Alias)(c), @@ -96,6 +98,14 @@ func (c *Options) UnmarshalJSON(data []byte) error { return err } + // Convert nodes array to map + if len(aux.Nodes) > 0 { + c.Nodes = make(map[string]struct{}, len(aux.Nodes)) + for _, node := range aux.Nodes { + c.Nodes[node] = struct{}{} + } + } + // Parse backend-specific options switch c.BackendType { case backends.BackendTypeLlamaCpp: @@ -147,11 +157,22 @@ func (c *Options) MarshalJSON() ([]byte, error) { // Use anonymous struct to avoid recursion type Alias Options aux := struct { + Nodes []string `json:"nodes,omitempty"` // Output as JSON array *Alias }{ Alias: (*Alias)(c), } + // Convert nodes map to array (sorted for consistency) + if len(c.Nodes) > 0 { + aux.Nodes = make([]string, 0, len(c.Nodes)) + for node := range c.Nodes { + aux.Nodes = append(aux.Nodes, node) + } + // Sort for consistent output + slices.Sort(aux.Nodes) + } + // Convert backend-specific options back to BackendOptions map for JSON switch c.BackendType { case backends.BackendTypeLlamaCpp: diff --git a/pkg/manager/manager.go b/pkg/manager/manager.go index 34ec2e5..b4378a7 100644 --- a/pkg/manager/manager.go +++ b/pkg/manager/manager.go @@ -275,9 +275,19 @@ func (im *instanceManager) loadInstance(name, path string) error { options := persistedInstance.GetOptions() - // Check if this is a remote instance - // An instance is remote if Nodes is specified AND the first node is not the local node - isRemote := options != nil && len(options.Nodes) > 0 && options.Nodes[0] != im.localNodeName + // Check if this is a remote instance (local node not in the Nodes set) + var isRemote bool + var nodeName string + if options != nil { + if _, isLocal := options.Nodes[im.localNodeName]; !isLocal { + // Get the first node from the set + for node := range options.Nodes { + nodeName = node + isRemote = true + break + } + } + } var statusCallback func(oldStatus, newStatus instance.Status) if !isRemote { @@ -296,7 +306,6 @@ func (im *instanceManager) loadInstance(name, path string) error { // Handle remote instance mapping if isRemote { - nodeName := options.Nodes[0] nodeConfig, exists := im.nodeConfigMap[nodeName] if !exists { return fmt.Errorf("node %s not found for remote instance %s", nodeName, name) diff --git a/pkg/manager/operations.go b/pkg/manager/operations.go index 81e374f..5931cfa 100644 --- a/pkg/manager/operations.go +++ b/pkg/manager/operations.go @@ -3,7 +3,6 @@ package manager import ( "fmt" "llamactl/pkg/backends" - "llamactl/pkg/config" "llamactl/pkg/instance" "llamactl/pkg/validation" "os" @@ -27,10 +26,12 @@ func (im *instanceManager) updateLocalInstanceFromRemote(localInst *instance.Ins // Preserve the Nodes field from the local instance localOptions := localInst.GetOptions() - var preservedNodes []string + var preservedNodes map[string]struct{} if localOptions != nil && len(localOptions.Nodes) > 0 { - preservedNodes = make([]string, len(localOptions.Nodes)) - copy(preservedNodes, localOptions.Nodes) + preservedNodes = make(map[string]struct{}, len(localOptions.Nodes)) + for node := range localOptions.Nodes { + preservedNodes[node] = struct{}{} + } } // Create a copy of remote options and restore the Nodes field @@ -98,16 +99,17 @@ func (im *instanceManager) CreateInstance(name string, options *instance.Options return nil, fmt.Errorf("instance with name %s already exists", name) } - // Check if this is a remote instance - // An instance is remote if Nodes is specified AND the first node is not the local node - isRemote := len(options.Nodes) > 0 && options.Nodes[0] != im.localNodeName - var nodeConfig *config.NodeConfig + // Check if this is a remote instance (local node not in the Nodes set) + if _, isLocal := options.Nodes[im.localNodeName]; !isLocal && len(options.Nodes) > 0 { + // Get the first node from the set + var nodeName string + for node := range options.Nodes { + nodeName = node + break + } - if isRemote { // Validate that the node exists - nodeName := options.Nodes[0] // Use first node for now - var exists bool - nodeConfig, exists = im.nodeConfigMap[nodeName] + nodeConfig, exists := im.nodeConfigMap[nodeName] if !exists { return nil, fmt.Errorf("node %s not found", nodeName) } diff --git a/pkg/server/handlers_instances.go b/pkg/server/handlers_instances.go index 9cb887e..9e41190 100644 --- a/pkg/server/handlers_instances.go +++ b/pkg/server/handlers_instances.go @@ -394,12 +394,21 @@ func (h *Handler) ProxyToInstance() http.HandlerFunc { func (h *Handler) RemoteInstanceProxy(w http.ResponseWriter, r *http.Request, name string, inst *instance.Instance) { // Get the node name from instance options options := inst.GetOptions() - if options == nil || len(options.Nodes) == 0 { - http.Error(w, "Instance has no node configured", http.StatusInternalServerError) + if options == nil { + http.Error(w, "Instance has no options configured", http.StatusInternalServerError) return } - nodeName := options.Nodes[0] + // Get the first node from the set + var nodeName string + for node := range options.Nodes { + nodeName = node + break + } + if nodeName == "" { + http.Error(w, "Instance has no node configured", http.StatusInternalServerError) + return + } // Check if we have a cached proxy for this node h.remoteProxiesMu.RLock() diff --git a/pkg/server/handlers_openai.go b/pkg/server/handlers_openai.go index 2849841..075f651 100644 --- a/pkg/server/handlers_openai.go +++ b/pkg/server/handlers_openai.go @@ -155,12 +155,21 @@ func (h *Handler) OpenAIProxy() http.HandlerFunc { func (h *Handler) RemoteOpenAIProxy(w http.ResponseWriter, r *http.Request, modelName string, inst *instance.Instance) { // Get the node name from instance options options := inst.GetOptions() - if options == nil || len(options.Nodes) == 0 { - http.Error(w, "Instance has no node configured", http.StatusInternalServerError) + if options == nil { + http.Error(w, "Instance has no options configured", http.StatusInternalServerError) return } - nodeName := options.Nodes[0] + // Get the first node from the set + var nodeName string + for node := range options.Nodes { + nodeName = node + break + } + if nodeName == "" { + http.Error(w, "Instance has no node configured", http.StatusInternalServerError) + return + } // Check if we have a cached proxy for this node h.remoteProxiesMu.RLock()