llmgw/internal/agent/provision/python.go

package provision

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"slices"
	"strconv"
	"strings"
	"syscall"
	"time"
)

// A Python engine is a virtual environment built by uv: an interpreter uv
// fetches, and the engine with every dependency from PyPI.
//
// Two things make it as trustworthy as an archive checked against its
// publisher's checksum. The dependencies are resolved as PyPI stood a day
// after the engine was published (--exclude-newer), so a node built next
// month gets what one built today got, not whatever a dependency released
// since — the usual way a working vLLM stops working. And the resolution is
// written out with every file's sha256 from PyPI and installed with
// --require-hashes, so a file that does not match what PyPI published is
// refused. The resolved list is kept in the environment, as the record of
// exactly what is installed.

// pythonInstallTimeout bounds one installation: several gigabytes, on a
// network that may be slow, but not forever.
const pythonInstallTimeout = time.Hour

// lockName is the resolved requirements, kept in the environment.
const lockName = "llmgw-requirements.txt"

// pylockName is the lock of a build with a second index: PEP 751's, with
// every file's address.
const pylockName = "pylock.llmgw.toml"

func (in *Installer) installPython(ctx context.Context, r Recipe, w Want, b Build, dir string, st *Status) error {
	ctx, cancel := context.WithTimeout(ctx, pythonInstallTimeout)
	defer cancel()
	// The version this build is, which for a GPU older than every build of
	// the one asked for is an older one: its publication date is the one
	// the dependencies are resolved as of.
	w.Version = b.versionOr(w.Version)

	if err := in.packages(ctx, b.Packages, st); err != nil {
		return err
	}
	if err := os.MkdirAll(in.root(), 0o755); err != nil {
		return err
	}
	if b.FreeGB > 0 {
		if free, err := in.free(in.root()); err == nil && free < float64(b.FreeGB) {
			return fmt.Errorf("%s needs about %d GB free under %s and there are %.1f GB",
				w.Engine, b.FreeGB, in.root(), free)
		}
	}

	in.progress(st, "installing uv %s", UVVersion)
	uv := in.one(ctx, Want{Engine: "uv", Version: UVVersion})
	if uv.State != Ready {
		return fmt.Errorf("installing uv, which installs %s: %s", w.Engine, uv.Detail)
	}

	var asOf string
	if !b.AsOf.IsZero() {
		asOf = b.AsOf.UTC().Format(time.RFC3339)
	} else {
		in.progress(st, "asking PyPI when %s %s was published", r.PyPI, w.Version)
		published, err := in.published(ctx, r.PyPI, w.Version)
		if err != nil {
			return err
		}
		asOf = published.Add(24 * time.Hour).UTC().Format(time.RFC3339)
	}

	cache := filepath.Join(in.root(), ".uv-cache")
	env := []string{
		"UV_CACHE_DIR=" + cache,
		"UV_PYTHON_INSTALL_DIR=" + filepath.Join(in.root(), "python"),
		// Only the interpreter uv fetched: the system's may be the wrong
		// version, or change under the environment with the next upgrade.
		"UV_PYTHON_PREFERENCE=only-managed",
		// Nobody's uv.toml or pip.conf decides what is installed here.
		"UV_NO_CONFIG=1",
	}
	tmp := dir + ".partial"
	_ = os.RemoveAll(tmp)
	defer os.RemoveAll(tmp) //nolint:errcheck // gone either way once renamed

	in.progress(st, "creating a Python %s environment for %s", b.Python, w.Engine)
	// Relocatable, because it is built under one name and renamed into place
	// once it works: its scripts find their interpreter relative to
	// themselves rather than by the path they were built at.
	if out, err := in.run(ctx, env, uv.Binary, "venv", "--relocatable", "--python", b.Python, tmp); err != nil {
		return fmt.Errorf("uv venv: %v: %s", err, tail(out))
	}
	python := filepath.Join(tmp, "bin", "python")

	reqs := slices.Clone(b.Requirements)
	if b.Source != nil {
		more, err := in.unpackSource(ctx, w, b, tmp, st)
		if err != nil {
			return err
		}
		reqs = append(reqs, more...)
	}
	if len(b.Assets) > 0 {
		if err := in.unpackAssets(ctx, r, w, b, tmp, st); err != nil {
			return err
		}
	}
	if b.Wheel != nil {
		// The engine from its GitHub release, pinned to the digest GitHub
		// publishes for the file: the same trust as an archive, and uv
		// refuses a download that does not match it.
		in.progress(st, "asking GitHub for %s %s's %s…%s", b.Wheel.Repo, b.Wheel.Tag, b.Wheel.Prefix, b.Wheel.Suffix)
		sums, err := in.checksums(ctx, b.Wheel.Repo, b.Wheel.Tag)
		if err != nil {
			return err
		}
		var asset, sum string
		for name, s := range sums {
			if strings.HasPrefix(name, b.Wheel.Prefix) && strings.HasSuffix(name, b.Wheel.Suffix) {
				asset, sum = name, s
			}
		}
		if asset == "" {
			return fmt.Errorf("%s %s publishes no %s…%s", b.Wheel.Repo, b.Wheel.Tag, b.Wheel.Prefix, b.Wheel.Suffix)
		}
		url := strings.TrimSuffix(in.downloadBase(), "/") + "/" + b.Wheel.Repo + "/releases/download/" +
			b.Wheel.Tag + "/" + asset
		reqs = append([]string{b.Wheel.Project + " @ " + url + "#sha256=" + sum}, reqs...)
	}
	in.progress(st, "resolving %s against PyPI as of %s", strings.Join(reqs, " "), asOf[:10])
	lock, err := in.resolve(ctx, env, uv.Binary, python, tmp, asOf, reqs, b)
	if err != nil {
		return err
	}

	in.progress(st, "downloading and installing %s and its dependencies: several gigabytes, which takes minutes",
		strings.Join(reqs, " "))
	if out, err := in.run(ctx, env, uv.Binary, "pip", "install", "--python", python,
		"--require-hashes", "--no-deps", "-r", lock); err != nil {
		return fmt.Errorf("installing %s: %v: %s", strings.Join(reqs, " "), err, tail(out))
	}

	in.progress(st, "checking that %s runs", w.Engine)
	if _, _, err := in.verify(ctx, tmp, b); err != nil {
		return err
	}
	if b.GPUCheck && in.Host.NVIDIADriver != "" {
		in.progress(st, "checking that %s (%s) computes on this node's GPU", w.Engine, b.Variant)
		if err := in.gpuCheck(ctx, python, w, b); err != nil {
			return err
		}
	}
	m, _ := json.MarshalIndent(markerFile{
		Engine: w.Engine, Version: w.Version, Variant: b.Variant,
		Binary: filepath.Join("bin", b.Binary), Requirements: reqs,
		ExcludeNewer: asOf, UV: UVVersion, Installed: time.Now().UTC(),
	}, "", "  ")
	if err := os.WriteFile(filepath.Join(tmp, marker), m, 0o644); err != nil { //nolint:gosec // not secret
		return err
	}
	_ = os.RemoveAll(dir)
	if err := os.Rename(tmp, dir); err != nil {
		return err
	}
	// The cache is as large as the environment. Kept after a failure, so the
	// next attempt does not download it all again; dropped after success.
	_ = os.RemoveAll(cache)
	in.log().Info("installed an engine", "engine", w.Engine, "version", w.Version,
		"variant", b.Variant, "dir", dir, "as_of", asOf)
	return nil
}

// resolve pins every package the requirements need, each file by its hash,
// into a lock in the environment, and returns its path.
//
// With a second index — PyTorch's, for the torch a build was made against —
// the packages it is for come from it and nowhere else, and every other
// package from PyPI alone. Letting both indexes offer every package is how
// torchtriton was taken in 2022: a name one index had, registered on the
// other with a higher version. So the second index is explicit, named only
// for its packages, and the lock is PEP 751's, which records each file's
// address as well as its hash: installing it asks no index anything.
func (in *Installer) resolve(ctx context.Context, env []string, uv, python, tmp, asOf string,
	reqs []string, b Build,
) (string, error) {
	if b.ExtraIndex == "" {
		lock := filepath.Join(tmp, lockName)
		if out, err := in.runStdin(ctx, env, strings.Join(reqs, "\n")+"\n", uv, "pip", "compile",
			"--python", python, "--exclude-newer", asOf, "--generate-hashes", "--quiet",
			"--output-file", lock, "-"); err != nil {
			return "", fmt.Errorf("resolving %s: %v: %s", strings.Join(reqs, " "), err, tail(out))
		}
		return lock, nil
	}

	project := tmp + ".resolve"
	if err := os.MkdirAll(project, 0o755); err != nil {
		return "", err
	}
	defer os.RemoveAll(project)                                                                                            //nolint:errcheck // scratch
	if err := os.WriteFile(filepath.Join(project, "pyproject.toml"), []byte(resolveProject(reqs, b)), 0o644); err != nil { //nolint:gosec // not secret
		return "", err
	}
	lock := filepath.Join(tmp, pylockName)
	if out, err := in.run(ctx, env, uv, "pip", "compile", "--python", python, "--exclude-newer", asOf,
		"--format", "pylock.toml", "--quiet", "--output-file", lock,
		filepath.Join(project, "pyproject.toml")); err != nil {
		return "", fmt.Errorf("resolving %s: %v: %s", strings.Join(reqs, " "), err, tail(out))
	}
	return lock, nil
}

// resolveProject is the pyproject.toml a build with a second index is
// resolved from. Its packages are direct requirements too: an index is
// given to a package only where the project names it.
func resolveProject(reqs []string, b Build) string {
	var p strings.Builder
	p.WriteString("[project]\nname = \"llmgw-engine\"\nversion = \"0\"\ndependencies = [\n")
	for _, r := range append(slices.Clone(reqs), b.IndexPackages...) {
		fmt.Fprintf(&p, "  %s,\n", strconv.Quote(r))
	}
	p.WriteString("]\n\n[[tool.uv.index]]\nname = \"extra\"\n")
	fmt.Fprintf(&p, "url = %s\nexplicit = true\n\n[tool.uv.sources]\n", strconv.Quote(b.ExtraIndex))
	for _, name := range b.IndexPackages {
		fmt.Fprintf(&p, "%s = { index = \"extra\" }\n", name)
	}
	return p.String()
}

// published is when PyPI received a release's last file.
func (in *Installer) published(ctx context.Context, project, version string) (time.Time, error) {
	url := strings.TrimSuffix(in.pypiBase(), "/") + "/pypi/" + project + "/" + version + "/json"
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return time.Time{}, err
	}
	resp, err := in.client().Do(req)
	if err != nil {
		return time.Time{}, fmt.Errorf("asking PyPI for %s %s: %w", project, version, err)
	}
	defer resp.Body.Close() //nolint:errcheck // read-only
	if resp.StatusCode == http.StatusNotFound {
		return time.Time{}, fmt.Errorf("PyPI has no %s %s", project, version)
	}
	if resp.StatusCode != http.StatusOK {
		return time.Time{}, fmt.Errorf("asking PyPI for %s %s: %s", project, version, resp.Status)
	}
	var rel struct {
		URLs []struct {
			Uploaded time.Time `json:"upload_time_iso_8601"`
		} `json:"urls"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&rel); err != nil {
		return time.Time{}, fmt.Errorf("reading PyPI's answer for %s %s: %w", project, version, err)
	}
	var last time.Time
	for _, u := range rel.URLs {
		if u.Uploaded.After(last) {
			last = u.Uploaded
		}
	}
	if last.IsZero() {
		return time.Time{}, errors.New("PyPI lists no files for " + project + " " + version)
	}
	return last, nil
}

func (in *Installer) free(dir string) (float64, error) {
	if in.FreeSpace != nil {
		return in.FreeSpace(dir)
	}
	return FreeGB(dir)
}

func (in *Installer) pypiBase() string {
	if in.PyPIBase != "" {
		return in.PyPIBase
	}
	return "https://pypi.org"
}

// runStdin runs a command with input on its standard input.
func (in *Installer) runStdin(ctx context.Context, env []string, stdin, name string, args ...string) ([]byte, error) {
	if in.RunStdin != nil {
		return in.RunStdin(ctx, env, stdin, name, args...)
	}
	return runWithStdin(ctx, env, stdin, name, args...)
}

// FreeGB is the space available to an unprivileged writer under dir.
func FreeGB(dir string) (float64, error) {
	var s syscall.Statfs_t
	if err := syscall.Statfs(dir, &s); err != nil {
		return 0, err
	}
	return float64(s.Bavail) * float64(s.Bsize) / (1 << 30), nil
}

// tail is the last few lines of a command's output: uv explains a failed
// resolution over several, and the last alone is rarely the reason.
func tail(b []byte) string {
	lines := strings.Split(strings.TrimSpace(string(b)), "\n")
	if len(lines) > 4 {
		lines = lines[len(lines)-4:]
	}
	return strings.Join(lines, " / ")
}

// gpuProbe is what the environment is asked to run on the GPU: one small
// operation, which needs the torch to have kernels for the card.
const gpuProbe = `import torch
assert torch.cuda.is_available(), "torch sees no CUDA device"
x = torch.ones(64, device="cuda")
print(float((x * 2).sum()), torch.cuda.get_device_name(0), torch.cuda.get_device_capability(0))`

// gpuCheck runs gpuProbe in a freshly installed environment.
func (in *Installer) gpuCheck(ctx context.Context, python string, w Want, b Build) error {
	ctx, cancel := context.WithTimeout(ctx, 3*time.Minute)
	defer cancel()
	out, err := in.run(ctx, nil, python, "-c", gpuProbe)
	if err != nil {
		hint := ""
		if strings.Contains(string(out), "no kernel image") {
			hint = "; this build has no kernels for this GPU, so try the other build for this node"
		}
		return fmt.Errorf("%s %s (%s) installed but cannot compute on this node's GPU: %s%s",
			w.Engine, w.Version, b.Variant, tail(out), hint)
	}
	return nil
}

// unpackSource fetches an engine's source archive, checks it against the
// digest the recipe holds, puts its code in the environment, and returns the
// requirements its requirements file names — less the packages a second
// index gives, which the build names itself.
func (in *Installer) unpackSource(ctx context.Context, w Want, b Build, env string, st *Status) ([]string, error) {
	src := b.Source
	url := src.URL
	if strings.HasPrefix(url, "/") {
		url = strings.TrimSuffix(in.downloadBase(), "/") + url
	}
	scratch := env + ".source"
	_ = os.RemoveAll(scratch)
	if err := os.MkdirAll(scratch, 0o755); err != nil {
		return nil, err
	}
	defer os.RemoveAll(scratch) //nolint:errcheck // scratch
	archive := filepath.Join(scratch, "source.tar.gz")
	name := w.Engine + " " + w.Version + "'s source"
	if err := in.download(ctx, url, name, archive, src.SHA256, st); err != nil {
		return nil, err
	}
	in.progress(st, "unpacking %s", name)
	if err := untarGz(archive, scratch); err != nil {
		return nil, fmt.Errorf("unpacking %s: %w", name, err)
	}
	dst := filepath.Join(env, src.Dir)
	if err := os.Rename(filepath.Join(scratch, src.Top), dst); err != nil {
		return nil, fmt.Errorf("%s has no %s at its top: %w", name, src.Top, err)
	}
	if src.Requirements == "" {
		return nil, nil
	}
	f, err := os.ReadFile(filepath.Join(dst, src.Requirements))
	if err != nil {
		return nil, fmt.Errorf("%s: %w", name, err)
	}
	return requirementsIn(string(f), b.IndexPackages)
}

// unpackAssets fetches a Python engine's release binaries, each checked
// against the digest GitHub publishes for it, and unpacks them into the
// environment's AssetDir.
func (in *Installer) unpackAssets(ctx context.Context, r Recipe, w Want, b Build, env string, st *Status) error {
	tag := b.AssetTag
	if tag == "" {
		tag = w.Version
	}
	in.progress(st, "asking GitHub for the checksums of %s %s", r.Repo, tag)
	sums, err := in.checksums(ctx, r.Repo, tag)
	if err != nil {
		return err
	}
	dst := filepath.Join(env, b.AssetDir)
	if err := os.MkdirAll(dst, 0o755); err != nil {
		return err
	}
	for _, asset := range b.Assets {
		want, ok := sums[asset]
		if !ok {
			// Refused rather than installed unchecked: this is a binary that
			// will run on every node.
			return fmt.Errorf("%s %s publishes no checksum for %s, so it cannot be verified", r.Repo, tag, asset)
		}
		url := strings.TrimSuffix(in.downloadBase(), "/") + "/" + r.Repo + "/releases/download/" + tag + "/" + asset
		archive := filepath.Join(env+".source", asset)
		if err := os.MkdirAll(filepath.Dir(archive), 0o755); err != nil {
			return err
		}
		if err := in.download(ctx, url, asset, archive, want, st); err != nil {
			return err
		}
		in.progress(st, "unpacking %s", asset)
		err := untarGz(archive, dst)
		_ = os.RemoveAll(filepath.Dir(archive))
		if err != nil {
			return fmt.Errorf("unpacking %s: %w", asset, err)
		}
	}
	return nil
}

// requirementsIn reads a requirements file's lines. Options (-r, -e, --index-url
// and the like) are refused: they would fetch from somewhere other than PyPI,
// or install what the lock never saw.
func requirementsIn(text string, skip []string) ([]string, error) {
	var out []string
	for _, line := range strings.Split(text, "\n") {
		line, _, _ = strings.Cut(line, "#")
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		if strings.HasPrefix(line, "-") || strings.Contains(line, "://") || strings.Contains(line, " @ ") {
			return nil, fmt.Errorf("the requirements name %q, which is not a package from PyPI", line)
		}
		name := strings.ToLower(line)
		if i := strings.IndexAny(name, "<>=!~;[ "); i >= 0 {
			name = name[:i]
		}
		if slices.Contains(skip, name) {
			continue
		}
		out = append(out, line)
	}
	return out, nil
}