llmgw/internal/comfy/nodes_test.go

package comfy

import (
	"context"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync/atomic"
	"testing"
)

// fakeUV is a uv that answers pip freeze with ComfyUI's packages and records
// what pip install was given.
func fakeUV(t *testing.T, freeze string) (uv, log string) {
	t.Helper()
	dir := t.TempDir()
	log = filepath.Join(dir, "install.log")
	uv = filepath.Join(dir, "uv")
	script := "#!/bin/sh\nif [ \"$2\" = freeze ]; then printf '%s' '" + freeze + "'; exit 0; fi\n" +
		"echo \"$@\" >> " + log + "\n" +
		"prev=; for a in \"$@\"; do if [ \"$prev\" = --target ]; then mkdir -p \"$a\"; fi; prev=$a; done\nexit 0\n"
	if err := os.WriteFile(uv, []byte(script), 0o755); err != nil {
		t.Fatal(err)
	}
	return uv, log
}

func gitRepo(t *testing.T, files map[string]string) (repo, commit string) {
	t.Helper()
	if _, err := exec.LookPath("git"); err != nil {
		t.Skip("no git")
	}
	repo = t.TempDir()
	run := func(args ...string) string {
		cmd := exec.Command("git", args...)
		cmd.Dir = repo
		cmd.Env = append(os.Environ(), "GIT_AUTHOR_NAME=t", "GIT_AUTHOR_EMAIL=t@t", "GIT_COMMITTER_NAME=t", "GIT_COMMITTER_EMAIL=t@t")
		b, err := cmd.CombinedOutput()
		if err != nil {
			t.Fatalf("git %v: %v %s", args, err, b)
		}
		return strings.TrimSpace(string(b))
	}
	run("init", "-q")
	run("config", "uploadpack.allowReachableSHA1InWant", "true")
	for name, body := range files {
		_ = os.MkdirAll(filepath.Dir(filepath.Join(repo, name)), 0o755)
		_ = os.WriteFile(filepath.Join(repo, name), []byte(body), 0o644)
	}
	run("add", ".")
	run("commit", "-qm", "one")
	return repo, run("rev-parse", "HEAD")
}

func TestANodesCompiledPackagesAreBuiltForTheCardOnceAndInstalledBesideTheReplica(t *testing.T) {
	repo, commit := gitRepo(t, map[string]string{"__init__.py": "", "texture_baker/setup.py": "", "uv_unwrapper/setup.py": ""})
	t.Setenv("XDG_CACHE_HOME", t.TempDir())
	t.Setenv("HOME", t.TempDir())
	cuda := t.TempDir()
	_ = os.MkdirAll(filepath.Join(cuda, "bin"), 0o755)
	_ = os.WriteFile(filepath.Join(cuda, "bin", "nvcc"), nil, 0o755)
	t.Setenv("CUDA_HOME", cuda)

	var builds atomic.Int32
	var envSeen []string
	oldArch, oldBuild := cudaArch, buildWheel
	t.Cleanup(func() { cudaArch, buildWheel = oldArch, oldBuild })
	cudaArch = func(context.Context, string) (string, error) { return "7.5", nil }
	buildWheel = func(_ context.Context, _, _, dir, out string, env []string) ([]byte, error) {
		builds.Add(1)
		envSeen = env
		name := strings.ReplaceAll(filepath.Base(dir), "_", "-")
		return nil, os.WriteFile(filepath.Join(out, strings.ReplaceAll(name, "-", "_")+"-0.0.1-cp312-cp312-linux_x86_64.whl"), []byte("w"), 0o644)
	}
	// ComfyUI's environment has an old hand-installed copy, which must not
	// pin the built one.
	uv, log := fakeUV(t, "numpy==2.5.3\ntexture-baker @ file:///home/someone/texture_baker-0.0.1.whl\n")
	node := Node{Repo: "file://" + repo, Commit: commit, Requirements: []string{"trimesh"},
		Build: []string{"texture_baker", "uv_unwrapper"}}

	state := t.TempDir()
	site, err := installNodes(t.Context(), state, "python3", []Node{node}, uv, func(string) {})
	if err != nil {
		t.Fatal(err)
	}
	if site == "" || builds.Load() != 2 {
		t.Fatalf("site %q, %d builds", site, builds.Load())
	}
	if !strings.Contains(strings.Join(envSeen, " "), "TORCH_CUDA_ARCH_LIST=7.5") {
		t.Error("the build was not told the card's architecture")
	}
	installed, _ := os.ReadFile(log)
	if !strings.Contains(string(installed), "trimesh") || !strings.Contains(string(installed), "sm75/texture_baker/texture_baker-0.0.1") ||
		!strings.Contains(string(installed), "uv_unwrapper-0.0.1") {
		t.Errorf("installed %s", installed)
	}
	constraints, _ := os.ReadFile(filepath.Join(state, "constraints.txt"))
	if strings.Contains(string(constraints), "texture-baker") || !strings.Contains(string(constraints), "numpy==2.5.3") {
		t.Errorf("constraints %s", constraints)
	}

	// A second replica on the same kind of card finds the wheels built.
	if _, err := installNodes(t.Context(), t.TempDir(), "python3", []Node{node}, uv, func(string) {}); err != nil {
		t.Fatal(err)
	}
	if builds.Load() != 2 {
		t.Errorf("built again: %d builds", builds.Load())
	}
	// Another kind of card builds its own.
	cudaArch = func(context.Context, string) (string, error) { return "8.9", nil }
	if _, err := installNodes(t.Context(), t.TempDir(), "python3", []Node{node}, uv, func(string) {}); err != nil {
		t.Fatal(err)
	}
	if builds.Load() != 4 {
		t.Errorf("an Ada used Turing's wheels: %d builds", builds.Load())
	}
}

func TestANodeToCompileOnAMachineWithoutCUDASaysSo(t *testing.T) {
	repo, commit := gitRepo(t, map[string]string{"__init__.py": "", "ext/setup.py": ""})
	t.Setenv("XDG_CACHE_HOME", t.TempDir())
	t.Setenv("HOME", t.TempDir())
	t.Setenv("CUDA_HOME", "")
	if _, err := os.Stat("/usr/local/cuda/bin/nvcc"); err == nil {
		t.Skip("this machine has a CUDA toolkit")
	}
	oldArch := cudaArch
	t.Cleanup(func() { cudaArch = oldArch })
	cudaArch = func(context.Context, string) (string, error) { return "8.0", nil }
	uv, _ := fakeUV(t, "")
	_, err := installNodes(t.Context(), t.TempDir(), "python3", []Node{{Repo: "file://" + repo, Commit: commit, Build: []string{"ext"}}},
		uv, func(string) {})
	if err == nil || !strings.Contains(err.Error(), "no CUDA toolkit") {
		t.Errorf("%v", err)
	}
}

func TestWeightsANodeLoadsByRepositoryNameAreInTheHuggingFaceCache(t *testing.T) {
	weights := t.TempDir()
	_ = os.WriteFile(filepath.Join(weights, "config.yaml"), []byte("c"), 0o644)
	_ = os.WriteFile(filepath.Join(weights, "model.safetensors"), []byte("m"), 0o644)
	hub := t.TempDir()
	if err := LinkHFCache(weights, "stabilityai/stable-fast-3d", hub); err != nil {
		t.Fatal(err)
	}
	root := filepath.Join(hub, "models--stabilityai--stable-fast-3d")
	ref, _ := os.ReadFile(filepath.Join(root, "refs", "main"))
	if len(ref) != 40 {
		t.Fatalf("refs/main %q is not a commit's form", ref)
	}
	b, err := os.ReadFile(filepath.Join(root, "snapshots", string(ref), "model.safetensors"))
	if err != nil || string(b) != "m" {
		t.Errorf("the weights through the cache: %q %v", b, err)
	}
	// Again, as a replica started a second time does.
	if err := LinkHFCache(weights, "stabilityai/stable-fast-3d", hub); err != nil {
		t.Errorf("a second time: %v", err)
	}
	for _, bad := range []string{"noslash", "a/b/c", "../x", "a/.."} {
		if err := LinkHFCache(weights, bad, hub); err == nil {
			t.Errorf("%q taken for a repository", bad)
		}
	}
}