llmgw/internal/comfy/nodes.go

package comfy

import (
	"context"
	"embed"
	"encoding/json"
	"fmt"
	"io/fs"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	"runtime"
	"slices"
	"strconv"
	"strings"
	"time"
)

// Custom nodes, put in a replica's own base directory at the commit the
// catalog names, their packages installed beside the replica rather than
// into ComfyUI's environment: one model's node cannot break another
// model's replicas, or the pictures every replica makes.

// Node is a custom node as the catalog gives it.
type Node struct {
	Repo         string   `json:"repo"`
	Commit       string   `json:"commit"`
	Requirements []string `json:"requirements,omitempty"`
	// Build are directories of the node that are packages to compile.
	Build []string `json:"build,omitempty"`
}

//go:embed all:nodes
var ownNodes embed.FS

// installOwnNodes writes llmgw's own nodes into a replica's custom_nodes:
// what every workflow may use, from the binary rather than a repository,
// written again at every start so a replica runs the gate's own version.
func installOwnNodes(state string) error {
	return fs.WalkDir(ownNodes, "nodes", func(p string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		dst := filepath.Join(state, "custom_nodes", strings.TrimPrefix(p, "nodes"))
		if d.IsDir() {
			return os.MkdirAll(dst, 0o755)
		}
		b, err := ownNodes.ReadFile(p)
		if err != nil {
			return err
		}
		return os.WriteFile(dst, b, 0o644) //nolint:gosec // code ComfyUI reads, not secret
	})
}

// NodesEnv is where the agent hands the gate a model's custom nodes.
const NodesEnv = "LLMGW_COMFYUI_NODES"

var nodeDir = regexp.MustCompile(`[^A-Za-z0-9._-]`)

// installNodes puts each node in state/custom_nodes at its commit and its
// packages in state/site, and returns the PYTHONPATH ComfyUI is started
// with: empty when no node wants a package. Run as the engines' user, in
// the replica's own directory, before ComfyUI starts.
func installNodes(ctx context.Context, state, python string, nodes []Node, uv string, out func(string)) (string, error) {
	if len(nodes) == 0 {
		return "", nil
	}
	ctx, cancel := context.WithTimeout(ctx, 20*time.Minute)
	defer cancel()
	var reqs, built []string
	for _, n := range nodes {
		name := nodeDir.ReplaceAllString(strings.TrimSuffix(path.Base(n.Repo), ".git"), "_")
		dir := filepath.Join(state, "custom_nodes", name)
		if head, err := gitOutput(ctx, dir, "rev-parse", "HEAD"); err == nil && head == n.Commit {
			reqs = append(reqs, n.Requirements...)
			wheels, err := buildNode(ctx, n, dir, python, uv, out)
			if err != nil {
				return "", err
			}
			built = append(built, wheels...)
			continue
		}
		out(fmt.Sprintf("custom node %s at %s", n.Repo, n.Commit[:12]))
		_ = os.RemoveAll(dir)
		for _, args := range [][]string{
			{"init", "-q", dir},
			{"-C", dir, "remote", "add", "origin", n.Repo},
			{"-C", dir, "fetch", "-q", "--depth", "1", "origin", n.Commit},
			{"-C", dir, "checkout", "-q", "FETCH_HEAD"},
		} {
			if b, err := exec.CommandContext(ctx, "git", args...).CombinedOutput(); err != nil { //nolint:gosec // the catalog's repository, checked as https
				return "", fmt.Errorf("custom node %s at %s: git %s: %v: %s", n.Repo, n.Commit, args[len(args)-2], err, lastLines(b))
			}
		}
		reqs = append(reqs, n.Requirements...)
		wheels, err := buildNode(ctx, n, dir, python, uv, out)
		if err != nil {
			return "", err
		}
		built = append(built, wheels...)
	}
	if len(reqs) == 0 && len(built) == 0 {
		return "", nil
	}
	if uv == "" {
		return "", fmt.Errorf("custom nodes need %s, and no uv is installed to install them with", strings.Join(reqs, ", "))
	}
	// A built package is named by its wheel; ComfyUI's own copy, if one was
	// ever put there, must not pin it to another.
	named := slices.Clone(reqs)
	for _, w := range built {
		named = append(named, wheelPackage(w))
	}
	reqs = append(reqs, built...)
	site := filepath.Join(state, "site")
	marker := filepath.Join(site, ".llmgw-requirements")
	want := strings.Join(reqs, "\n")
	if b, err := os.ReadFile(marker); err == nil && string(b) == want {
		return site, nil
	}
	out("installing " + strings.Join(reqs, " ") + " for the custom nodes")
	// What ComfyUI's environment has, as constraints: a package both want is
	// installed at the version ComfyUI runs, so the replica's copy and
	// ComfyUI's never disagree.
	freeze, err := exec.CommandContext(ctx, uv, "pip", "freeze", "--python", python).Output() //nolint:gosec // the agent's own uv
	if err != nil {
		return "", fmt.Errorf("reading ComfyUI's packages: %w", err)
	}
	constraints := filepath.Join(state, "constraints.txt")
	if err := os.WriteFile(constraints, withoutNamed(freeze, named), 0o644); err != nil { //nolint:gosec // not secret
		return "", err
	}
	_ = os.RemoveAll(site)
	args := append([]string{"pip", "install", "--python", python, "--target", site, "-c", constraints, "--quiet"}, reqs...)
	cmd := exec.CommandContext(ctx, uv, args...) //nolint:gosec // requirements checked by the catalog as package specs
	cmd.Env = append(os.Environ(), "UV_NO_CONFIG=1", "UV_CACHE_DIR="+filepath.Join(state, ".uv-cache"))
	if b, err := cmd.CombinedOutput(); err != nil {
		return "", fmt.Errorf("installing the custom nodes' packages (%s): %v: %s", strings.Join(reqs, " "), err, lastLines(b))
	}
	// The replica's copies come first on ComfyUI's path, but of anything
	// both have they are the same version, so which is found is the same.
	if err := os.WriteFile(marker, []byte(want), 0o644); err != nil { //nolint:gosec // not secret
		return "", err
	}
	return site, nil
}

// parseNodes reads the nodes the agent handed over.
func parseNodes(s string) ([]Node, error) {
	if strings.TrimSpace(s) == "" {
		return nil, nil
	}
	var nodes []Node
	if err := json.Unmarshal([]byte(s), &nodes); err != nil {
		return nil, fmt.Errorf("%s: %w", NodesEnv, err)
	}
	return nodes, nil
}

// findUV is the uv the agent installed engines with.
func findUV() string {
	if m, _ := filepath.Glob("/opt/llmgw/engines/uv/*/*/uv"); len(m) > 0 {
		return m[len(m)-1]
	}
	if p, err := exec.LookPath("uv"); err == nil {
		return p
	}
	return ""
}

func gitOutput(ctx context.Context, dir string, args ...string) (string, error) {
	b, err := exec.CommandContext(ctx, "git", append([]string{"-C", dir}, args...)...).Output() //nolint:gosec // fixed arguments
	return strings.TrimSpace(string(b)), err
}

func lastLines(b []byte) string {
	lines := strings.Split(strings.TrimSpace(string(b)), "\n")
	if len(lines) > 4 {
		lines = lines[len(lines)-4:]
	}
	return strings.Join(lines, " | ")
}

// withoutNamed is ComfyUI's packages less those a node names itself: of a
// package a node asks for by name, the node's range decides — TripoSR's
// checkpoint loads with transformers 4 and not 5 — and of everything else,
// what ComfyUI runs. A node that names a package must name a range ComfyUI
// can run with too: the copy it gets is first on ComfyUI's path.
func withoutNamed(freeze []byte, reqs []string) []byte {
	named := map[string]bool{}
	for _, r := range reqs {
		named[packageName(r)] = true
	}
	var out []string
	for _, line := range strings.Split(string(freeze), "\n") {
		if strings.TrimSpace(line) == "" || named[packageName(line)] {
			continue
		}
		out = append(out, line)
	}
	return []byte(strings.Join(out, "\n") + "\n")
}

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

// Packages a node's repository has that no index does — C++ and CUDA
// extensions — are built here, for the card the replica runs on: a CUDA
// kernel compiled for one architecture does not load on another, and a
// site's cards are not the PoC's. A build is minutes; its wheel is kept on
// the node, by the node's commit, the directory and the architecture, for
// the next replica on the same kind of card.

// cudaArch is the compute capability of the card the replica sees, as
// TORCH_CUDA_ARCH_LIST wants it ("7.5"). A var for the tests.
var cudaArch = func(ctx context.Context, python string) (string, error) {
	cmd := exec.CommandContext(ctx, python, "-c", //nolint:gosec // the agent's own Python
		"import torch; print('%d.%d' % torch.cuda.get_device_capability())")
	b, err := cmd.Output()
	if err != nil {
		return "", fmt.Errorf("asking torch for the card's compute capability: %v", err)
	}
	return strings.TrimSpace(string(b)), nil
}

// buildWheel builds the package in dir as a wheel into outDir, with
// ComfyUI's Python and what it already has (torch, setuptools): a CUDA
// extension builds against the torch it will load with. A var for the tests.
var buildWheel = func(ctx context.Context, uv, python, dir, outDir string, env []string) ([]byte, error) {
	cmd := exec.CommandContext(ctx, uv, "build", "--wheel", "--no-build-isolation", //nolint:gosec // the catalog's node, at its commit
		"--python", python, "--out-dir", outDir, dir)
	cmd.Env = env
	return cmd.CombinedOutput()
}

// buildCache is where built wheels are kept: under the engines' home, which
// outlives a replica.
func buildCache() string {
	if d, err := os.UserCacheDir(); err == nil {
		return filepath.Join(d, "llmgw", "node-builds")
	}
	return filepath.Join(os.TempDir(), "llmgw-node-builds")
}

// buildNode builds each of a node's packages that are compiled, or finds the
// wheel built before for this commit and card, and returns the wheels.
func buildNode(ctx context.Context, n Node, dir, python, uv string, out func(string)) ([]string, error) {
	if len(n.Build) == 0 {
		return nil, nil
	}
	if uv == "" {
		return nil, fmt.Errorf("custom node %s has packages to build, and no uv is installed to build them with", n.Repo)
	}
	arch, err := cudaArch(ctx, python)
	if err != nil {
		return nil, fmt.Errorf("custom node %s: %w", n.Repo, err)
	}
	cudaHome := os.Getenv("CUDA_HOME")
	if cudaHome == "" {
		for _, c := range []string{"/usr/local/cuda"} {
			if _, err := os.Stat(filepath.Join(c, "bin", "nvcc")); err == nil {
				cudaHome = c
			}
		}
	}
	var wheels []string
	for _, sub := range n.Build {
		cached := filepath.Join(buildCache(), n.Commit[:12], "sm"+strings.ReplaceAll(arch, ".", ""), nodeDir.ReplaceAllString(sub, "_"))
		if w, ok := firstWheel(cached); ok {
			wheels = append(wheels, w)
			continue
		}
		if cudaHome == "" {
			return nil, fmt.Errorf("custom node %s: %s is compiled, and this node has no CUDA toolkit (no nvcc in "+
				"CUDA_HOME or /usr/local/cuda) to compile it with", n.Repo, sub)
		}
		out(fmt.Sprintf("building %s of %s for compute %s", sub, n.Repo, arch))
		// Built beside where it is kept and moved there whole: two replicas
		// starting at once on one node must not see half a wheel.
		if err := os.MkdirAll(filepath.Dir(cached), 0o755); err != nil {
			return nil, err
		}
		tmp, err := os.MkdirTemp(filepath.Dir(cached), ".build-")
		if err != nil {
			return nil, err
		}
		env := append(os.Environ(), "TORCH_CUDA_ARCH_LIST="+arch, "CUDA_HOME="+cudaHome,
			"PATH="+filepath.Join(cudaHome, "bin")+string(os.PathListSeparator)+os.Getenv("PATH"),
			"UV_NO_CONFIG=1", "MAX_JOBS="+strconv.Itoa(max(runtime.NumCPU(), 1)))
		b, err := buildWheel(ctx, uv, python, filepath.Join(dir, sub), tmp, env)
		if err != nil {
			_ = os.RemoveAll(tmp)
			return nil, fmt.Errorf("custom node %s: building %s: %v: %s", n.Repo, sub, err, lastLines(b))
		}
		if _, ok := firstWheel(tmp); !ok {
			_ = os.RemoveAll(tmp)
			return nil, fmt.Errorf("custom node %s: building %s made no wheel: %s", n.Repo, sub, lastLines(b))
		}
		if err := os.Rename(tmp, cached); err != nil {
			_ = os.RemoveAll(tmp)
			// Another replica got there first: its wheel is as good.
			if _, ok := firstWheel(cached); !ok {
				return nil, err
			}
		}
		w, _ := firstWheel(cached)
		wheels = append(wheels, w)
	}
	return wheels, nil
}

func firstWheel(dir string) (string, bool) {
	m, _ := filepath.Glob(filepath.Join(dir, "*.whl"))
	if len(m) == 0 {
		return "", false
	}
	return m[0], true
}

// wheelPackage is the package a wheel file is of: its name up to the first
// '-', as PyPI compares names.
func wheelPackage(wheel string) string {
	name, _, _ := strings.Cut(filepath.Base(wheel), "-")
	return packageName(name)
}