diff --git a/CHANGELOG.md b/CHANGELOG.md index 6892a3a..5a16e5c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -164,6 +164,13 @@ APIs change between minors without a deprecation period. archive to them with age, as `.tar.gz.age`: an archive holds the CA's key and every secret in the store, and is meant to be copied off the machine. Open one with `age -d -i `. +- **PyTorch's index gives torch and nothing else.** vLLM's CUDA 12.9 build + was resolved with PyTorch's index beside PyPI and `unsafe-best-match`, so + either index could supply any package — how torchtriton was hijacked in + 2022. Now that index is explicit, named only for torch, torchvision and + torchaudio; every other package comes from PyPI alone, and the lock is + PEP 751's (`pylock.llmgw.toml`), pinning each file's address as well as + its hash. Checked against vLLM 0.30.0: the same 195 packages resolve. ### Everything that can be changed on a running cluster, from the console diff --git a/internal/agent/provision/python.go b/internal/agent/provision/python.go index 3cb74c4..7eb848d 100644 --- a/internal/agent/provision/python.go +++ b/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" diff --git a/internal/agent/provision/python_test.go b/internal/agent/provision/python_test.go index 6df8e23..9de3edf 100644 --- a/internal/agent/provision/python_test.go +++ b/internal/agent/provision/python_test.go @@ -40,6 +40,14 @@ func (f *fakePython) run(ctx context.Context, env []string, name string, args .. return []byte("uv 0.12.20\n"), nil case strings.HasSuffix(name, "/uv") && args[0] == "venv": return nil, os.MkdirAll(filepath.Join(args[len(args)-1], "bin"), 0o755) + case strings.HasSuffix(name, "/uv") && args[0] == "pip" && args[1] == "compile": + // From a pyproject: kept in the record, as the request was. + project, _ := os.ReadFile(args[len(args)-1]) + f.pmu.Lock() + f.commands = append(f.commands, "pyproject: "+string(project)) + f.pmu.Unlock() + out := args[slices.Index(args, "--output-file")+1] + return nil, os.WriteFile(out, []byte("lock-version = \"1.0\"\n"), 0o644) case strings.HasSuffix(name, "/uv") && args[0] == "pip" && args[1] == "install": if f.failPip { return []byte("error: Hash mismatch for torch\n"), errExit @@ -220,8 +228,12 @@ func TestAnOlderDriverGetsTheSameVLLMBuiltForCUDA12(t *testing.T) { // Pinned to the digest GitHub publishes for the wheel. "vllm @ " + in.DownloadBase + "/vllm-project/vllm/releases/download/v0.30.0/" + wheel + "#sha256=" + strings.TrimPrefix(gh.digests[wheel], "sha256:"), - "--extra-index-url https://download.pytorch.org/whl/cu129 --index-strategy unsafe-best-match -", - "--require-hashes --no-deps -r", + // The torch packages from PyTorch's index and only those; the lock + // records every file's address. + "--format pylock.toml", + `url = "https://download.pytorch.org/whl/cu129"` + "\nexplicit = true", + `torch = { index = "extra" }`, `torchvision = { index = "extra" }`, + "--require-hashes --no-deps -r " + filepath.Join(in.Root, "vllm", "0.30.0-cu129.partial", pylockName), // Computing on the card, not only importing. "/bin/python -c import torch", } { @@ -230,8 +242,8 @@ func TestAnOlderDriverGetsTheSameVLLMBuiltForCUDA12(t *testing.T) { } } for _, c := range sys.commands { - if strings.Contains(c, "pip install") && !strings.Contains(c, "download.pytorch.org/whl/cu129") { - t.Errorf("installed without the index its torch comes from: %s", c) + if strings.Contains(c, "unsafe-best-match") || strings.Contains(c, "--extra-index-url") { + t.Errorf("every index offered every package: %s", c) } } } diff --git a/internal/agent/provision/recipe.go b/internal/agent/provision/recipe.go index 76c6c39..c76ca16 100644 --- a/internal/agent/provision/recipe.go +++ b/internal/agent/provision/recipe.go @@ -50,10 +50,12 @@ type Build struct { GPUBackend string // Wheel, when set, is the engine itself as a wheel from a GitHub // release rather than from PyPI, checked against the digest GitHub - // publishes for it. ExtraIndex is a package index besides PyPI that the - // rest resolves from too: the CUDA 12.9 torch the wheel was built on. - Wheel *Wheel - ExtraIndex string + // publishes for it. ExtraIndex is a package index besides PyPI, for + // IndexPackages only: the CUDA 12.9 torch the wheel was built on. Every + // other package comes from PyPI alone. + Wheel *Wheel + ExtraIndex string + IndexPackages []string // GPUCheck, on a node with an NVIDIA GPU, has the installed environment // compute something on one before the build counts as installed. A build // with no kernels for the card — a GPU generation its torch dropped — @@ -234,6 +236,10 @@ var VLLMVariants = []string{"cu13", "cu129"} // torch is built for CUDA 13; the cu129 vLLM wheel needs this one. const pytorchCU129 = "https://download.pytorch.org/whl/cu129" +// pytorchPackages are what is taken from PyTorch's index: the builds for one +// CUDA version, which PyPI does not have. +var pytorchPackages = []string{"torch", "torchvision", "torchaudio"} + // vllmCUDA13Since is the first vLLM whose PyPI build is CUDA 13, read from // its releases: 0.19.1 publishes "+cu130" beside a CUDA 12 default, 0.20.0 // "+cu129" beside a CUDA 13 one. @@ -332,6 +338,7 @@ func vllmCandidates(version, variant string, h discovery.Host) ([]Build, error) Prefix: "vllm-" + v + "+cu129-", Suffix: "_" + map[string]string{"amd64": "x86_64", "arm64": "aarch64"}[h.Arch] + ".whl", } cu129.ExtraIndex = pytorchCU129 + cu129.IndexPackages = pytorchPackages return pick("vLLM "+v, variant, []candidate{ {pypi("cu13"), gpuNeeds{minDriverForCUDA13, minCC}, "vLLM " + v + " built for CUDA 13"}, {cu129, gpuNeeds{minDriverForCUDA12, minCC}, "vLLM " + v + " built for CUDA 12.9"},