llmgw/internal/agent/provision/comfyui_test.go

package provision

import (
	"crypto/sha256"
	"encoding/hex"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
)

// comfyuiSource is a ComfyUI release archive as GitHub makes one, and the
// recipe pointed at it.
func comfyuiSource(t *testing.T, in *Installer, rel *release, digest string) {
	t.Helper()
	archive := tarGz(t, map[string]string{
		"ComfyUI-0.39.2/main.py":            "print('comfy')\n",
		"ComfyUI-0.39.2/comfyui_version.py": "__version__ = \"0.39.2\"\n",
		"ComfyUI-0.39.2/requirements.txt": "comfyui-frontend-package==1.53.10\ntorch\ntorchvision\n" +
			"numpy>=1.25.0\n\n#non essential dependencies:\nkornia>=0.7.1\n",
	}, nil)
	rel.assets["v0.39.2.tar.gz"] = archive
	sum := sha256.Sum256(archive)
	if digest == "" {
		digest = hex.EncodeToString(sum[:])
	}
	old := comfyuiReleases
	comfyuiReleases = map[string]struct {
		sha256    string
		published string
	}{"0.39.2": {digest, "2026-10-07T21:58:07Z"}}
	t.Cleanup(func() { comfyuiReleases = old })
}

var turingNode = discovery.Host{OSID: "ubuntu", OSVersion: "24.04", Arch: "amd64",
	NVIDIADriver: "580.82.07", MinComputeCapability: 7.5}

func TestComfyUIIsInstalledFromItsSourceCheckedAgainstTheRecipe(t *testing.T) {
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}}
	in, rel := pythonInstaller(t, turingNode, sys)
	comfyuiSource(t, in, rel, "")

	got := in.Converge(t.Context(), []Want{{Engine: "comfyui", Version: "0.39.2"}})
	if got[0].State != Ready {
		t.Fatalf("not ready: %+v", got[0])
	}
	dir := filepath.Join(in.Root, "comfyui", "0.39.2-cu13")
	if got[0].Binary != filepath.Join(dir, "bin", "python") {
		t.Errorf("binary %s, want the environment's interpreter", got[0].Binary)
	}
	if _, err := os.Stat(filepath.Join(dir, "comfyui", "main.py")); err != nil {
		t.Errorf("ComfyUI's code is not beside the interpreter: %v", err)
	}
	all := strings.Join(sys.commands, "\n")
	for _, want := range []string{
		// Not PyPI's date, which ComfyUI has none of: the recipe's, a day on.
		"--exclude-newer 2026-10-08T21:58:07Z",
		"comfyui-frontend-package==1.53.10", "numpy>=1.25.0", "kornia>=0.7.1", "torch",
		"--require-hashes",
	} {
		if !strings.Contains(all, want) {
			t.Errorf("no %q in:\n%s", want, all)
		}
	}
	if strings.Contains(all, "/pypi/") {
		t.Error("PyPI was asked when ComfyUI was published")
	}
}

func TestComfyUIOnAnOlderDriverTakesPyTorchsCUDA12Torch(t *testing.T) {
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}}
	host := turingNode
	host.NVIDIADriver = "550.163.01"
	in, rel := pythonInstaller(t, host, sys)
	comfyuiSource(t, in, rel, "")
	got := in.Converge(t.Context(), []Want{{Engine: "comfyui", Version: "0.39.2"}})
	if got[0].State != Ready || got[0].Variant != "cu129" {
		t.Fatalf("got %+v", got[0])
	}
	var project string
	for _, c := range sys.commands {
		if strings.HasPrefix(c, "pyproject: ") {
			project = c
		}
	}
	// torch from PyTorch's index alone, named once, as the build names it.
	if !strings.Contains(project, `torch = { index = "extra" }`) || strings.Count(project, `"torch"`) != 1 {
		t.Errorf("torch is not from PyTorch's index alone:\n%s", project)
	}
}

// A V100 on driver 580: CUDA 13 runs on the driver, but neither its torch
// nor CUDA 12.9's has kernels for Volta. Tried in turn, each was installed,
// 10 GB at a time, and failed on the card.
func TestComfyUIOnAVoltaTakesPyTorchsCUDA126Torch(t *testing.T) {
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}}
	host := turingNode
	host.MinComputeCapability = 7.0
	in, rel := pythonInstaller(t, host, sys)
	comfyuiSource(t, in, rel, "")
	got := in.Converge(t.Context(), []Want{{Engine: "comfyui", Version: "0.39.2"}})
	if got[0].State != Ready || got[0].Variant != "cu126" {
		t.Fatalf("got %+v", got[0])
	}
	all := strings.Join(sys.commands, "\n")
	if !strings.Contains(all, pytorchCU126) || strings.Contains(all, pytorchCU129) {
		t.Errorf("torch is not CUDA 12.6's:\n%s", all)
	}
}

func TestComfyUIWhoseSourceIsNotTheOneTheRecipeNamesIsNotInstalled(t *testing.T) {
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}}
	in, rel := pythonInstaller(t, turingNode, sys)
	comfyuiSource(t, in, rel, strings.Repeat("ab", 32))
	got := in.Converge(t.Context(), []Want{{Engine: "comfyui", Version: "0.39.2"}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "does not match") {
		t.Fatalf("got %s %q", got[0].State, got[0].Detail)
	}
	for _, c := range sys.commands {
		if strings.Contains(c, "pip install") {
			t.Errorf("installed from a source that did not match: %s", c)
		}
	}
}

func TestComfyUIIsInstalledOnlyWhereItCanRun(t *testing.T) {
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}}
	in, _ := pythonInstaller(t, ubuntu, sys)
	if got := in.Converge(t.Context(), []Want{{Engine: "comfyui", Version: "0.39.2"}}); got[0].State != Skipped {
		t.Errorf("a node with no GPU: %+v", got[0])
	}
	in, _ = pythonInstaller(t, turingNode, sys)
	got := in.Converge(t.Context(), []Want{{Engine: "comfyui", Version: "0.1.0"}})
	if got[0].State != Unsupported || !strings.Contains(got[0].Detail, "knows the source digest of ComfyUI 0.39.2") {
		t.Errorf("a version without a digest: %s %q", got[0].State, got[0].Detail)
	}
}

func TestARequirementsFileIsReadForPackagesAndNothingElse(t *testing.T) {
	got, err := requirementsIn("torch\ntorchvision>=0.20\nPillow\n# a comment\naiohttp>=3.11 # inline\n", pytorchPackages)
	if err != nil || strings.Join(got, ",") != "Pillow,aiohttp>=3.11" {
		t.Errorf("%v %v", got, err)
	}
	for _, bad := range []string{"-r other.txt", "--index-url https://evil", "x @ https://evil/x.whl", "-e ."} {
		if _, err := requirementsIn(bad+"\n", nil); err == nil {
			t.Errorf("%q was accepted", bad)
		}
	}
}

// A card every build runs on keeps the order it had: a Turing node's
// installed cu13 is still the first found, and is not installed again
// because a build for older cards was added after it.
func TestComfyUIBuildsAreTriedBestFirst(t *testing.T) {
	variants := func(h discovery.Host) string {
		builds, err := comfyuiCandidates("0.39.2", "", h)
		if err != nil {
			return err.Error()
		}
		var v []string
		for _, b := range builds {
			v = append(v, b.Variant)
		}
		return strings.Join(v, ",")
	}
	volta := turingNode
	volta.MinComputeCapability = 7.0
	older := turingNode
	older.NVIDIADriver = "550.163.01"
	for h, want := range map[*discovery.Host]string{&turingNode: "cu13,cu129,cu126", &volta: "cu126", &older: "cu129,cu126"} {
		if got := variants(*h); got != want {
			t.Errorf("cc %.1f driver %s: %s, want %s", h.MinComputeCapability, h.NVIDIADriver, got, want)
		}
	}
}