llmgw/internal/control/catalog/shared.go

package catalog

import (
	"time"
)

// A shared replica: one ComfyUI serving every model of an environment.
//
// Each ComfyUI model ran in a replica of its own, so on one card a picture
// model and a 3D model took turns, each stopping the other and loading
// ComfyUI and its weights from nothing. An environment of two or more models
// gets a hidden catalog entry of its own, whose replica carries every
// member's weights and custom nodes and answers for each of them; the
// members get no replicas of their own. Placement, idleness, eviction and
// adoption then work as for any model, on that entry.

// RoleEnvironment is the role of an environment's hidden entry.
const RoleEnvironment = "environment"

// EnvModelPrefix names an environment's hidden entry: comfyui-env-<name>.
const EnvModelPrefix = "comfyui-env-"

// SharedComfy is what an environment's entry carries of its members.
type SharedComfy struct {
	Environment string         `json:"environment"`
	Members     []SharedMember `json:"members"`
}

// SharedMember is one model a shared replica answers for.
type SharedMember struct {
	Name    string   `json:"name"`
	Aliases []string `json:"aliases,omitempty"`
	// Source is the member's weights, which the replica's node fetches.
	Source string `json:"source"`
	// HFRepo is the repository its node loads the weights by name from.
	HFRepo string `json:"hf_repo,omitempty"`
	// AllowPickle is the member model's allow_pickle.
	AllowPickle bool `json:"allow_pickle,omitempty"`
	// Workflow and Checkpoint are what a request for it is sent.
	Workflow   string `json:"workflow"`
	Checkpoint string `json:"checkpoint"`
}

// IsEnvironment reports whether a model is an environment's hidden entry.
func (m Model) IsEnvironment() bool { return m.Role == RoleEnvironment }

// SharedBy is the hidden entry whose replica serves a model, when it is in
// an environment of two or more.
func (c *Catalog) SharedBy(model string) (string, bool) {
	m, ok := c.Resolve(model)
	if !ok || m.IsEnvironment() || m.ComfyEnvironment() == "" {
		return "", false
	}
	name := EnvModelPrefix + m.ComfyEnvironment()
	if _, ok := c.models[name]; !ok {
		return "", false
	}
	return name, true
}

// envModel is the hidden entry for an environment of two or more models: on
// the first member's profiles, needing the most memory any of them does,
// on demand for the longest idle timeout of theirs.
func envModel(env Environment, members []Model) Model {
	first := members[0]
	profs := make([]Profile, len(first.Profiles))
	copy(profs, first.Profiles)
	mem := 0
	for _, m := range members {
		for _, p := range m.Profiles {
			mem = max(mem, p.MinMemoryGB)
		}
	}
	shared := &SharedComfy{Environment: env.Name}
	var idle time.Duration
	minN, maxN := 0, 1
	class := ""
	for _, m := range members {
		workflow, ckpt := m.ComfyWorkflow()
		hf := ""
		if m.ComfyUI != nil {
			hf = m.ComfyUI.HuggingFaceRepo
		}
		shared.Members = append(shared.Members, SharedMember{Name: m.Name, Aliases: m.Aliases,
			Source: m.Source, HFRepo: hf, AllowPickle: m.AllowPickle, Workflow: workflow, Checkpoint: ckpt})
		idle = max(idle, m.Scaling.IdleTimeout)
		minN, maxN = max(minN, m.Scaling.Min), max(maxN, m.Scaling.Max)
		if ClassRank(m.EffectiveClass()) > ClassRank(class) || class == "" {
			class = m.EffectiveClass()
		}
	}
	for i := range profs {
		profs[i].MinMemoryGB = max(profs[i].MinMemoryGB, mem)
		profs[i].ComfyUI = nil // the environment's nodes, below
		profs[i].Source = ""
	}
	return Model{
		Name:        EnvModelPrefix + env.Name,
		Role:        RoleEnvironment,
		Description: "the shared ComfyUI replica of environment " + env.Name,
		Source:      first.Source,
		AllowPickle: first.AllowPickle,
		ComfyUI:     &ComfyUI{Environment: env.Name, Nodes: env.Nodes, Output: OutputEnvironment},
		Profiles:    profs,
		Scaling:     Scaling{Min: minN, Max: maxN, IdleTimeout: idle},
		Class:       class,
		Shared:      shared,
	}
}

// OutputEnvironment is an environment entry's output: it makes nothing of
// its own, its members do.
const OutputEnvironment = "environment"