Take only torch from PyTorch's index, and lock its builds with every file's address

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

internal/agent/provision/python.go

@@ -8,6 +8,8 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"slices"
+	"strconv"
 	"strings"
 	"syscall"
 	"time"
@@ -33,6 +35,10 @@ 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()
@@ -113,28 +119,16 @@ func (in *Installer) installPython(ctx context.Context, r Recipe, w Want, b Buil
 			b.Wheel.Tag + "/" + asset
 		reqs = append([]string{b.Wheel.Project + " @ " + url + "#sha256=" + sum}, reqs...)
 	}
-	// A second index, for the torch the engine was built on. Every package
-	// is taken from whichever of the two has the version resolved — both are
-	// the publishers' own — and each file is still pinned by its hash.
-	var index []string
-	if b.ExtraIndex != "" {
-		index = []string{"--extra-index-url", b.ExtraIndex, "--index-strategy", "unsafe-best-match"}
-	}
-
 	in.progress(st, "resolving %s against PyPI as of %s", strings.Join(reqs, " "), asOf[:10])
-	lock := filepath.Join(tmp, lockName)
-	args := append([]string{"pip", "compile", "--python", python, "--exclude-newer", asOf,
-		"--generate-hashes", "--quiet", "--output-file", lock}, index...)
-	args = append(args, "-")
-	if out, err := in.runStdin(ctx, env, strings.Join(reqs, "\n")+"\n", uv.Binary, args...); err != nil {
-		return fmt.Errorf("resolving %s: %v: %s", strings.Join(reqs, " "), err, tail(out))
+	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, " "))
-	install := append([]string{"pip", "install", "--python", python,
-		"--require-hashes", "--no-deps", "-r", lock}, index...)
-	if out, err := in.run(ctx, env, uv.Binary, install...); err != nil {
+	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))
 	}
 
@@ -168,6 +162,63 @@ func (in *Installer) installPython(ctx context.Context, r Recipe, w Want, b Buil
 	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"