llmgw/internal/control/catalog/environments.go

package catalog

import (
	"fmt"
	"regexp"
	"slices"
	"strings"
)

// Environments: the ComfyUI models that may share one replica.
//
// Every ComfyUI model used to run in a replica of its own, so two models on
// one card took turns with a cold start each time: flux, then a mesh model,
// then flux again, each loading ComfyUI and its weights from nothing. ComfyUI
// itself moves between models' weights in one process; what keeps two models
// out of one process is their custom nodes, which share one Python. An
// environment is a set of models whose nodes can: the catalog checks it, and
// placement may then run one replica for all of them.

// Environment is one environment's models and what a replica of it carries.
type Environment struct {
	Name string `json:"name"`
	// Models are the catalog models in it, in catalog order.
	Models []string `json:"models"`
	// Nodes are every custom node any of them names, once each.
	Nodes []CustomNode `json:"nodes,omitempty"`
	// Sources are every weights source its models load, once each: what a
	// replica of it must have on its node.
	Sources []string `json:"sources,omitempty"`
}

// ComfyEnvironment is the environment a ComfyUI model is in, or "" for one in
// none, which is an environment of its own.
func (m Model) ComfyEnvironment() string {
	if m.ComfyUI == nil {
		return ""
	}
	return m.ComfyUI.Environment
}

var envName = regexp.MustCompile(`^[a-z0-9][a-z0-9._-]{0,62}$`)

// Environments are the catalog's named environments, by name.
func (c *Catalog) Environments() []Environment {
	envs, _ := c.environments()
	return envs
}

// EnvironmentOf is the environment a model is in, when it is in a named one.
func (c *Catalog) EnvironmentOf(model string) (Environment, bool) {
	m, ok := c.Resolve(model)
	if !ok || m.ComfyEnvironment() == "" {
		return Environment{}, false
	}
	for _, e := range c.Environments() {
		if e.Name == m.ComfyEnvironment() {
			return e, true
		}
	}
	return Environment{}, false
}

// environments builds the named environments and refuses one whose models
// cannot share a replica: two commits of one node, or two ranges of one
// package, of which one Python can hold only one.
func (c *Catalog) environments() ([]Environment, error) {
	byName := map[string]*Environment{}
	var order []string
	// Who brought each node and each package into each environment, to name
	// both models when two disagree.
	nodeFrom := map[string]map[string]string{} // env -> repo -> model
	nodeAt := map[string]map[string]string{}   // env -> repo -> commit
	pkgFrom := map[string]map[string]string{}  // env -> package -> model
	pkgSpec := map[string]map[string]string{}  // env -> package -> spec
	for _, name := range c.order {
		m := c.models[name]
		env := m.ComfyEnvironment()
		if env == "" || m.IsEnvironment() {
			continue
		}
		if !envName.MatchString(env) {
			return nil, fmt.Errorf("model %q: comfyui.environment %q is not a name of lowercase letters, digits, '.', '_' and '-'",
				name, env)
		}
		if !m.ServesComfyUI() || m.External != nil {
			return nil, fmt.Errorf("model %q: comfyui.environment is for a model the cluster's ComfyUI serves", name)
		}
		e := byName[env]
		if e == nil {
			e = &Environment{Name: env}
			byName[env] = e
			order = append(order, env)
			nodeFrom[env], nodeAt[env] = map[string]string{}, map[string]string{}
			pkgFrom[env], pkgSpec[env] = map[string]string{}, map[string]string{}
		}
		e.Models = append(e.Models, name)
		for _, src := range m.WeightSources() {
			if !slices.Contains(e.Sources, src) {
				e.Sources = append(e.Sources, src)
			}
		}
		for _, n := range m.ComfyUI.Nodes {
			repo := strings.TrimSuffix(strings.TrimRight(n.Repo, "/"), ".git")
			if at, seen := nodeAt[env][repo]; seen {
				if at != n.Commit {
					return nil, fmt.Errorf("environment %q: models %q and %q name custom node %s at different commits (%s, %s): "+
						"one replica holds one of them", env, nodeFrom[env][repo], name, n.Repo, short(at), short(n.Commit))
				}
			} else {
				nodeAt[env][repo], nodeFrom[env][repo] = n.Commit, name
				e.Nodes = append(e.Nodes, CustomNode{Repo: n.Repo, Commit: n.Commit})
			}
			i := slices.IndexFunc(e.Nodes, func(x CustomNode) bool {
				return strings.TrimSuffix(strings.TrimRight(x.Repo, "/"), ".git") == repo
			})
			for _, b := range n.Build {
				if !slices.Contains(e.Nodes[i].Build, b) {
					e.Nodes[i].Build = append(e.Nodes[i].Build, b)
				}
			}
			for _, r := range n.Requirements {
				pkg, spec := requirementName(r), normalizeSpec(r)
				if prev, seen := pkgSpec[env][pkg]; seen {
					if prev != spec {
						return nil, fmt.Errorf("environment %q: models %q and %q ask for %s as %q and %q: "+
							"one Python holds one version; name the same range in both",
							env, pkgFrom[env][pkg], name, pkg, prev, spec)
					}
					continue
				}
				pkgSpec[env][pkg], pkgFrom[env][pkg] = spec, name
				// The requirement goes with the node that first named it:
				// what the replica installs is the union, and each once.
				e.Nodes[i].Requirements = append(e.Nodes[i].Requirements, strings.TrimSpace(r))
			}
		}
	}
	out := make([]Environment, 0, len(order))
	for _, n := range order {
		out = append(out, *byName[n])
	}
	return out, nil
}

// requirementName is a requirement's package, as PyPI compares names.
func requirementName(spec string) string {
	s := strings.TrimSpace(spec)
	if i := strings.IndexAny(s, "<>=!~;[ "); i >= 0 {
		s = s[:i]
	}
	return strings.NewReplacer("_", "-", ".", "-").Replace(strings.ToLower(s))
}

// normalizeSpec is a requirement without its spaces and with its package
// named as PyPI compares it, so "transformers >=4.50, <5" and
// "Transformers>=4.50,<5" are one range.
func normalizeSpec(spec string) string {
	s := strings.Join(strings.Fields(spec), "")
	name := requirementName(s)
	i := strings.IndexAny(s, "<>=!~;[")
	if i < 0 {
		return name
	}
	return name + s[i:]
}

func short(commit string) string {
	if len(commit) > 12 {
		return commit[:12]
	}
	return commit
}