From 5628f8ce18c02ff159edae9316216784408b6dfc Mon Sep 17 00:00:00 2001 From: asahoo Date: Tue, 7 Jul 2026 16:10:01 -0500 Subject: [PATCH 1/5] Support local package artifacts in requirements --- architecture/05-build-system.md | 2 +- docs/llms.txt | 16 +- docs/yaml.md | 16 +- .../local_python_requirement_artifact.txtar | 34 ++++ pkg/config/config.go | 95 ++++++++++++ pkg/config/config_test.go | 146 ++++++++++++++++++ pkg/config/data/config_schema_v1.0.json | 2 +- pkg/dockerfile/standard_generator.go | 53 ++++++- pkg/dockerfile/standard_generator_test.go | 42 +++++ pkg/requirements/local_artifact.go | 109 +++++++++++++ pkg/requirements/requirements_test.go | 40 +++++ 11 files changed, 549 insertions(+), 6 deletions(-) create mode 100644 integration-tests/tests/local_python_requirement_artifact.txtar create mode 100644 pkg/requirements/local_artifact.go diff --git a/architecture/05-build-system.md b/architecture/05-build-system.md index 8a79198d75..0b8f0b9627 100644 --- a/architecture/05-build-system.md +++ b/architecture/05-build-system.md @@ -217,7 +217,7 @@ Resolution follows a 3-tier priority for each wheel: | 2. Auto-detect `dist/coglet-*.whl` | Dev builds only | | 3. Default | Install from PyPI | -Local wheel files are copied into `.cog/build/` and referenced via the `cog_build` named build context, then `COPY --from=cog_build`'d and `pip install`'d in the Dockerfile. +Local wheel files are copied into `.cog/build/` and referenced via the `cog_build` named build context, then `COPY --from=cog_build`'d and `pip install`'d in the Dockerfile. User-provided local wheel and source archive requirements use the same staging boundary: Cog copies only the referenced artifacts into `.cog/build/` before the requirements install, preserving the later `COPY . /src` source layer. --- diff --git a/docs/llms.txt b/docs/llms.txt index e62311b0d6..20aca719ce 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -3470,6 +3470,20 @@ Your `cog.yaml` file can set either `python_packages` or `python_requirements`, This follows the standard [requirements.txt](https://pip.pypa.io/en/stable/reference/requirements-file-format/) format. +Requirements files can also reference local Python package artifacts, such as wheels and source archives: + +`requirements.txt`: + +``` +./dist/mylib-0.1.0-py3-none-any.whl +./vendor/helperlib.zip +./packages/localpkg.tar.gz +``` + +Local artifact paths are resolved relative to the requirements file, and the referenced files must be inside your project directory. Cog stages these artifacts before installing requirements, so this is the supported way to install a local package artifact during `cog build`. + +Local package directories, `name @ file:...` requirements, local `--find-links` directories, recursive local artifact includes, and inline hashes or options on local artifact lines are not supported. + To install Git-hosted Python packages, add `git` to the `system_packages` list, then use the `git+https://` syntax to specify the package name. For example: `cog.yaml`: @@ -3547,7 +3561,7 @@ build: - cd cowsay-3.7.0 && make install ``` -Your code is _not_ available to commands in `run`. This is so we can build your image efficiently when running locally. +Your code is _not_ available to commands in `run`. This is so we can build your image efficiently when running locally. To install a local wheel or source archive, list it in your `python_requirements` file instead of running `pip install ./artifact.zip` from `run`. Each command in `run` can be either a string or a dictionary in the following format: diff --git a/docs/yaml.md b/docs/yaml.md index e7d98b80e1..f8eb942165 100644 --- a/docs/yaml.md +++ b/docs/yaml.md @@ -59,6 +59,20 @@ Your `cog.yaml` file can set either `python_packages` or `python_requirements`, This follows the standard [requirements.txt](https://pip.pypa.io/en/stable/reference/requirements-file-format/) format. +Requirements files can also reference local Python package artifacts, such as wheels and source archives: + +`requirements.txt`: + +``` +./dist/mylib-0.1.0-py3-none-any.whl +./vendor/helperlib.zip +./packages/localpkg.tar.gz +``` + +Local artifact paths are resolved relative to the requirements file, and the referenced files must be inside your project directory. Cog stages these artifacts before installing requirements, so this is the supported way to install a local package artifact during `cog build`. + +Local package directories, `name @ file:...` requirements, local `--find-links` directories, recursive local artifact includes, and inline hashes or options on local artifact lines are not supported. + To install Git-hosted Python packages, add `git` to the `system_packages` list, then use the `git+https://` syntax to specify the package name. For example: `cog.yaml`: @@ -136,7 +150,7 @@ build: - cd cowsay-3.7.0 && make install ``` -Your code is _not_ available to commands in `run`. This is so we can build your image efficiently when running locally. +Your code is _not_ available to commands in `run`. This is so we can build your image efficiently when running locally. To install a local wheel or source archive, list it in your `python_requirements` file instead of running `pip install ./artifact.zip` from `run`. Each command in `run` can be either a string or a dictionary in the following format: diff --git a/integration-tests/tests/local_python_requirement_artifact.txtar b/integration-tests/tests/local_python_requirement_artifact.txtar new file mode 100644 index 0000000000..089d5083f9 --- /dev/null +++ b/integration-tests/tests/local_python_requirement_artifact.txtar @@ -0,0 +1,34 @@ +# Local package artifacts listed in requirements.txt are available during build. + +mkdir dist +exec python3 -c 'import zipfile; z = zipfile.ZipFile("dist/local_pkg-0.1.0.zip", "w"); z.write("setup.py"); z.write("local_pkg/__init__.py"); z.close()' + +cog build -t $TEST_IMAGE +cog predict $TEST_IMAGE +stdout 'hello from local package' + +-- cog.yaml -- +build: + python_version: "3.12" + python_requirements: requirements.txt +predict: predict.py:Predictor + +-- requirements.txt -- +./dist/local_pkg-0.1.0.zip + +-- setup.py -- +from setuptools import setup + +setup(name="local-pkg", version="0.1.0", packages=["local_pkg"]) + +-- local_pkg/__init__.py -- +MESSAGE = "hello from local package" + +-- predict.py -- +from cog import BasePredictor +from local_pkg import MESSAGE + + +class Predictor(BasePredictor): + def predict(self) -> str: + return MESSAGE diff --git a/pkg/config/config.go b/pkg/config/config.go index 621b4bf49f..db62a9482e 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -1,8 +1,11 @@ package config import ( + "crypto/sha256" + "encoding/hex" "encoding/json" "fmt" + "os" "path/filepath" "regexp" "slices" @@ -60,6 +63,14 @@ type Build struct { SDKVersion string `json:"sdk_version,omitempty" yaml:"sdk_version,omitempty"` pythonRequirementsContent []string + localPackageArtifacts []LocalPackageArtifact +} + +type LocalPackageArtifact struct { + Requirement string + SourcePath string + StagedDir string + Filename string } type Concurrency struct { @@ -277,6 +288,9 @@ func (c *Config) Complete(projectDir string) error { return fmt.Errorf("failed to open python_requirements file: %w", err) } c.Build.pythonRequirementsContent = reqs + if err := c.loadLocalPackageArtifacts(projectDir, requirementsFilePath); err != nil { + return err + } } else if len(c.Build.PythonPackages) > 0 { // Backwards compatibility: if using deprecated python_packages, populate requirements content c.Build.pythonRequirementsContent = c.Build.PythonPackages @@ -297,6 +311,80 @@ func (c *Config) Complete(projectDir string) error { return nil } +func (c *Config) loadLocalPackageArtifacts(projectDir string, requirementsFilePath string) error { + projectRoot, err := filepath.Abs(projectDir) + if err != nil { + return fmt.Errorf("failed to resolve project directory: %w", err) + } + projectRoot, err = filepath.EvalSymlinks(projectRoot) + if err != nil { + return fmt.Errorf("failed to resolve project directory symlinks: %w", err) + } + + requirementsDir := filepath.Dir(requirementsFilePath) + seen := map[string]LocalPackageArtifact{} + artifacts := []LocalPackageArtifact{} + for _, line := range c.Build.pythonRequirementsContent { + artifactPath, ok, err := requirements.ParseLocalArtifactRequirement(line) + if err != nil { + return err + } + if !ok { + continue + } + + resolvedPath := artifactPath + if !filepath.IsAbs(resolvedPath) { + resolvedPath = filepath.Join(requirementsDir, resolvedPath) + } + absPath, err := filepath.Abs(resolvedPath) + if err != nil { + return fmt.Errorf("failed to resolve local Python package artifact %q: %w", artifactPath, err) + } + canonicalPath, err := filepath.EvalSymlinks(absPath) + if err != nil { + return fmt.Errorf("local Python package artifact %q not found: %w", artifactPath, err) + } + info, err := os.Stat(canonicalPath) + if err != nil { + return fmt.Errorf("failed to inspect local Python package artifact %q: %w", artifactPath, err) + } + if !info.Mode().IsRegular() { + return fmt.Errorf("local Python package artifact %q must be a regular file", artifactPath) + } + if !pathWithin(projectRoot, canonicalPath) { + return fmt.Errorf("local Python package artifact %q must be inside the project directory", artifactPath) + } + + if artifact, ok := seen[canonicalPath]; ok { + artifact.Requirement = line + artifacts = append(artifacts, artifact) + continue + } + + relPath, err := filepath.Rel(projectRoot, canonicalPath) + if err != nil { + return fmt.Errorf("failed to resolve local Python package artifact %q relative to project directory: %w", artifactPath, err) + } + hash := sha256.Sum256([]byte(filepath.ToSlash(relPath))) + artifact := LocalPackageArtifact{ + Requirement: line, + SourcePath: canonicalPath, + StagedDir: hex.EncodeToString(hash[:])[:16], + Filename: filepath.Base(absPath), + } + seen[canonicalPath] = artifact + artifacts = append(artifacts, artifact) + } + c.Build.localPackageArtifacts = artifacts + return nil +} + +func pathWithin(root string, target string) bool { + rel, err := filepath.Rel(root, target) + return err == nil && rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator)) +} + // PythonRequirementsForArch returns a requirements.txt file with all the GPU packages resolved for given OS and architecture. func (c *Config) PythonRequirementsForArch(goos string, goarch string, includePackages []string) (string, error) { packages := []string{} @@ -546,6 +634,13 @@ func (c *Config) RequirementsFile(projectDir string) string { return filepath.Join(projectDir, c.Build.PythonRequirements) } +func (c *Config) LocalPackageArtifacts() []LocalPackageArtifact { + if c.Build == nil { + return nil + } + return slices.Clone(c.Build.localPackageArtifacts) +} + func (c *Config) ParsedEnvironment() map[string]string { return c.parsedEnvironment } diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index 49164e087e..0494cf3595 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -196,6 +196,152 @@ flask>0.4 } +func TestPythonRequirementsLocalPackageArtifacts(t *testing.T) { + tmpDir := t.TempDir() + require.NoError(t, os.MkdirAll(path.Join(tmpDir, "requirements", "dist"), 0o755)) + wheelPath := path.Join(tmpDir, "requirements", "dist", "local_pkg-0.1.0-py3-none-any.whl") + archivePath := path.Join(tmpDir, "requirements", "mylibpackage.zip") + require.NoError(t, os.WriteFile(wheelPath, []byte("wheel"), 0o644)) + require.NoError(t, os.WriteFile(archivePath, []byte("zip"), 0o644)) + require.NoError(t, os.WriteFile(path.Join(tmpDir, "requirements", "requirements.txt"), []byte(`torch==1.13.1 +./dist/local_pkg-0.1.0-py3-none-any.whl +mylibpackage.zip`), 0o644)) + + config := &Config{ + Build: &Build{ + GPU: true, + PythonVersion: "3.10", + PythonRequirements: "requirements/requirements.txt", + }, + } + require.NoError(t, config.Complete(tmpDir)) + + artifacts := config.LocalPackageArtifacts() + canonicalWheelPath, err := filepath.EvalSymlinks(wheelPath) + require.NoError(t, err) + canonicalArchivePath, err := filepath.EvalSymlinks(archivePath) + require.NoError(t, err) + require.Len(t, artifacts, 2) + require.Equal(t, "./dist/local_pkg-0.1.0-py3-none-any.whl", artifacts[0].Requirement) + require.Equal(t, canonicalWheelPath, artifacts[0].SourcePath) + require.Equal(t, "local_pkg-0.1.0-py3-none-any.whl", artifacts[0].Filename) + require.Len(t, artifacts[0].StagedDir, 16) + require.Equal(t, "mylibpackage.zip", artifacts[1].Requirement) + require.Equal(t, canonicalArchivePath, artifacts[1].SourcePath) + + requirements, err := config.PythonRequirementsForArch("linux", "amd64", []string{}) + require.NoError(t, err) + require.Equal(t, `--extra-index-url https://download.pytorch.org/whl/cu117/ +torch==1.13.1 +./dist/local_pkg-0.1.0-py3-none-any.whl +mylibpackage.zip`, requirements) +} + +func TestPythonRequirementsLocalPackageArtifactDuplicateAliases(t *testing.T) { + tmpDir := t.TempDir() + require.NoError(t, os.MkdirAll(path.Join(tmpDir, "dist"), 0o755)) + wheelPath := path.Join(tmpDir, "dist", "local_pkg-0.1.0-py3-none-any.whl") + require.NoError(t, os.WriteFile(wheelPath, []byte("wheel"), 0o644)) + require.NoError(t, os.WriteFile(path.Join(tmpDir, "requirements.txt"), []byte("./dist/local_pkg-0.1.0-py3-none-any.whl\ndist/local_pkg-0.1.0-py3-none-any.whl"), 0o644)) + + config := &Config{Build: &Build{PythonVersion: "3.10", PythonRequirements: "requirements.txt"}} + require.NoError(t, config.Complete(tmpDir)) + + artifacts := config.LocalPackageArtifacts() + require.Len(t, artifacts, 2) + require.Equal(t, "./dist/local_pkg-0.1.0-py3-none-any.whl", artifacts[0].Requirement) + require.Equal(t, "dist/local_pkg-0.1.0-py3-none-any.whl", artifacts[1].Requirement) + require.Equal(t, artifacts[0].SourcePath, artifacts[1].SourcePath) + require.Equal(t, artifacts[0].StagedDir, artifacts[1].StagedDir) + require.Equal(t, artifacts[0].Filename, artifacts[1].Filename) +} + +func TestPythonRequirementsLocalPackageArtifactUnsupportedLocalOptions(t *testing.T) { + for _, line := range []string{ + "--find-links ./wheels", + "--find-links=./wheels", + "-r requirements-local.txt", + "--requirement requirements-local.txt", + } { + t.Run(line, func(t *testing.T) { + projectDir := t.TempDir() + require.NoError(t, os.WriteFile(path.Join(projectDir, "requirements.txt"), []byte(line), 0o644)) + config := &Config{Build: &Build{PythonVersion: "3.10", PythonRequirements: "requirements.txt"}} + err := config.Complete(projectDir) + require.Error(t, err) + require.Contains(t, err.Error(), "local requirements option") + }) + } +} + +func TestPythonRequirementsLocalPackageArtifactValidation(t *testing.T) { + testCases := []struct { + name string + line string + setup func(t *testing.T, projectDir string) string + expectedErr string + }{ + { + name: "MissingArtifact", + line: "./missing.zip", + expectedErr: "not found", + }, + { + name: "DirectoryArtifact", + line: "./pkg.zip", + setup: func(t *testing.T, projectDir string) string { + require.NoError(t, os.Mkdir(path.Join(projectDir, "pkg.zip"), 0o755)) + return "" + }, + expectedErr: "regular file", + }, + { + name: "OutsideProject", + setup: func(t *testing.T, projectDir string) string { + outsideDir := t.TempDir() + outsidePath := path.Join(outsideDir, "pkg.zip") + require.NoError(t, os.WriteFile(outsidePath, []byte("zip"), 0o644)) + return outsidePath + }, + expectedErr: "inside the project directory", + }, + { + name: "SymlinkOutsideProject", + line: "./pkg.zip", + setup: func(t *testing.T, projectDir string) string { + outsideDir := t.TempDir() + outsidePath := path.Join(outsideDir, "pkg.zip") + require.NoError(t, os.WriteFile(outsidePath, []byte("zip"), 0o644)) + require.NoError(t, os.Symlink(outsidePath, path.Join(projectDir, "pkg.zip"))) + return "" + }, + expectedErr: "inside the project directory", + }, + { + name: "InlineHashUnsupported", + line: "./pkg.whl --hash=sha256:abc", + expectedErr: "inline options or hashes", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + projectDir := t.TempDir() + line := tc.line + if tc.setup != nil { + if setupLine := tc.setup(t, projectDir); setupLine != "" { + line = setupLine + } + } + require.NoError(t, os.WriteFile(path.Join(projectDir, "requirements.txt"), []byte(line), 0o644)) + config := &Config{Build: &Build{PythonVersion: "3.10", PythonRequirements: "requirements.txt"}} + err := config.Complete(projectDir) + require.Error(t, err) + require.Contains(t, err.Error(), tc.expectedErr) + }) + } +} + func TestValidateAndCompleteCUDAForAllTF(t *testing.T) { for _, compat := range TFCompatibilityMatrix { config := &Config{ diff --git a/pkg/config/data/config_schema_v1.0.json b/pkg/config/data/config_schema_v1.0.json index 27de32d113..9259c49b03 100644 --- a/pkg/config/data/config_schema_v1.0.json +++ b/pkg/config/data/config_schema_v1.0.json @@ -74,7 +74,7 @@ "python_requirements": { "$id": "#/properties/build/properties/python_requirements", "type": "string", - "description": "A pip requirements file specifying the Python packages to install." + "description": "A pip requirements file specifying the Python packages to install. Local wheel and source archive paths in the requirements file are supported when they point to files inside the project." }, "system_packages": { "$id": "#/properties/build/properties/system_packages", diff --git a/pkg/dockerfile/standard_generator.go b/pkg/dockerfile/standard_generator.go index 6282d42fb1..4613f8a6ab 100644 --- a/pkg/dockerfile/standard_generator.go +++ b/pkg/dockerfile/standard_generator.go @@ -14,6 +14,7 @@ import ( "github.com/replicate/cog/pkg/registry" "github.com/replicate/cog/pkg/requirements" "github.com/replicate/cog/pkg/util/console" + "github.com/replicate/cog/pkg/util/files" "github.com/replicate/cog/pkg/util/version" "github.com/replicate/cog/pkg/weightslegacy" "github.com/replicate/cog/pkg/wheels" @@ -902,6 +903,11 @@ func (g *StandardGenerator) pipInstalls() (string, error) { // Strip cog/coglet from user requirements — we always install them ourselves // via installCog(). Leaving them in would cause pip to overwrite our version. g.pythonRequirementsContents = g.filterManagedPackages(g.pythonRequirementsContents) + artifactCopyLines, err := g.stageLocalPackageArtifacts() + if err != nil { + return "", err + } + g.pythonRequirementsContents = g.rewriteLocalPackageArtifacts(g.pythonRequirementsContents) if strings.Trim(g.pythonRequirementsContents, "") == "" { return "", nil @@ -917,12 +923,55 @@ func (g *StandardGenerator) pipInstalls() (string, error) { if g.strip { pipInstallLine += " && " + StripDebugSymbolsCommand } - return strings.Join([]string{ + lines := []string{} + lines = append(lines, artifactCopyLines...) + lines = append(lines, copyLine[0], CFlags, pipInstallLine, "ENV CFLAGS=", - }, "\n"), nil + ) + return strings.Join(lines, "\n"), nil +} + +func (g *StandardGenerator) stageLocalPackageArtifacts() ([]string, error) { + artifacts := g.Config.LocalPackageArtifacts() + if len(artifacts) == 0 { + return nil, nil + } + staged := map[string]bool{} + for _, artifact := range artifacts { + dst := filepath.Join(g.tmpDir, "local_package_artifacts", artifact.StagedDir, artifact.Filename) + if staged[dst] { + continue + } + if err := files.Copy(artifact.SourcePath, dst); err != nil { + return nil, fmt.Errorf("failed to stage local Python package artifact %s: %w", artifact.SourcePath, err) + } + staged[dst] = true + } + return []string{"COPY --from=cog_build local_package_artifacts/ /tmp/local_package_artifacts/"}, nil +} + +func (g *StandardGenerator) rewriteLocalPackageArtifacts(reqContents string) string { + artifacts := map[string]string{} + for _, artifact := range g.Config.LocalPackageArtifacts() { + containerPath := path.Join("/tmp/local_package_artifacts", artifact.StagedDir, artifact.Filename) + artifacts[artifact.Requirement] = containerPath + } + if len(artifacts) == 0 { + return reqContents + } + + lines := []string{} + for line := range strings.SplitSeq(reqContents, "\n") { + if replacement, ok := artifacts[strings.TrimSpace(line)]; ok { + lines = append(lines, replacement) + } else { + lines = append(lines, line) + } + } + return strings.Join(lines, "\n") } func (g *StandardGenerator) runCommands() (string, error) { diff --git a/pkg/dockerfile/standard_generator_test.go b/pkg/dockerfile/standard_generator_test.go index dc7e8c8b82..d4f7082189 100644 --- a/pkg/dockerfile/standard_generator_test.go +++ b/pkg/dockerfile/standard_generator_test.go @@ -5,6 +5,7 @@ import ( "os" "path" "path/filepath" + "strings" "testing" "time" @@ -338,6 +339,47 @@ build: require.Contains(t, actual, `uv run pip install --cache-dir /root/.cache/pip -r /tmp/requirements.txt`) } +func TestPythonRequirementsLocalPackageArtifact(t *testing.T) { + tmpDir := t.TempDir() + require.NoError(t, os.MkdirAll(path.Join(tmpDir, "requirements", "dist"), 0o755)) + artifactPath := path.Join(tmpDir, "requirements", "dist", "local_pkg-0.1.0-py3-none-any.whl") + require.NoError(t, os.WriteFile(artifactPath, []byte("wheel"), 0o644)) + require.NoError(t, os.WriteFile(path.Join(tmpDir, "requirements", "requirements.txt"), []byte("./dist/local_pkg-0.1.0-py3-none-any.whl\ndist/local_pkg-0.1.0-py3-none-any.whl"), 0o644)) + + conf, err := config.FromYAML([]byte(` +build: + python_version: "3.12" + python_requirements: "requirements/requirements.txt" +`)) + require.NoError(t, err) + require.NoError(t, conf.Complete(tmpDir)) + command := dockertest.NewMockCommand() + client := registrytest.NewMockRegistryClient() + buildDir := t.TempDir() + gen, err := NewStandardGenerator(conf, tmpDir, buildDir, "", command, client, true) + require.NoError(t, err) + gen.SetUseCogBaseImage(false) + pypiWheels(gen) + _, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), "r8.im/replicate/cog-test") + require.NoError(t, err) + + artifact := conf.LocalPackageArtifacts()[0] + copyArtifacts := "COPY --from=cog_build local_package_artifacts/ /tmp/local_package_artifacts/" + copyRequirements := "COPY --from=cog_build requirements.txt /tmp/requirements.txt" + pipInstall := "uv run pip install --cache-dir /root/.cache/pip -r /tmp/requirements.txt" + require.Contains(t, actual, copyArtifacts) + require.Less(t, strings.Index(actual, copyArtifacts), strings.Index(actual, copyRequirements)) + require.Less(t, strings.Index(actual, copyRequirements), strings.Index(actual, pipInstall)) + + stagedArtifact, err := os.ReadFile(path.Join(buildDir, "local_package_artifacts", artifact.StagedDir, artifact.Filename)) + require.NoError(t, err) + require.Equal(t, []byte("wheel"), stagedArtifact) + + requirements, err := os.ReadFile(path.Join(buildDir, "requirements.txt")) + require.NoError(t, err) + require.Equal(t, "/tmp/local_package_artifacts/"+artifact.StagedDir+"/local_pkg-0.1.0-py3-none-any.whl\n/tmp/local_package_artifacts/"+artifact.StagedDir+"/local_pkg-0.1.0-py3-none-any.whl", string(requirements)) +} + // GPU builds on nvidia/cuda base images install Python via `uv python install` // (in installPythonCUDA), which marks it as externally managed (PEP 668). All // pip install lines must include --break-system-packages. diff --git a/pkg/requirements/local_artifact.go b/pkg/requirements/local_artifact.go new file mode 100644 index 0000000000..4da4388e03 --- /dev/null +++ b/pkg/requirements/local_artifact.go @@ -0,0 +1,109 @@ +package requirements + +import ( + "fmt" + "path/filepath" + "strings" +) + +var localArtifactSuffixes = []string{ + ".whl", + ".zip", + ".tar.gz", + ".tgz", + ".tar.bz2", + ".tar.xz", +} + +// ParseLocalArtifactRequirement identifies simple local wheel/source-archive +// requirement lines. It intentionally does not parse full pip requirement +// syntax; callers should reject unsupported local forms with clear errors. +func ParseLocalArtifactRequirement(line string) (string, bool, error) { + line = strings.TrimSpace(line) + if line == "" { + return "", false, nil + } + + if strings.HasPrefix(line, "file:") || strings.Contains(line, " @ file:") { + return "", false, fmt.Errorf("local file URL requirements are not supported: %s", line) + } + if option, ok := parseUnsupportedLocalOption(line); ok { + return "", false, fmt.Errorf("local requirements option %q is not supported: %s", option, line) + } + if strings.HasPrefix(line, "-") || isRemoteRequirement(line) { + return "", false, nil + } + + fields := strings.Fields(line) + if len(fields) > 1 { + if hasLocalArtifactSuffix(fields[0]) || isLocalPath(fields[0]) { + return "", false, fmt.Errorf("local package artifact requirements do not support inline options or hashes: %s", line) + } + return "", false, nil + } + + if !isLocalPath(line) && !hasLocalArtifactSuffix(line) { + return "", false, nil + } + if !hasLocalArtifactSuffix(line) { + return "", false, fmt.Errorf("local package requirement %q is not a supported wheel or source archive", line) + } + + return line, true, nil +} + +func parseUnsupportedLocalOption(line string) (string, bool) { + for _, option := range []string{"--find-links", "--requirement"} { + if value, ok := optionValue(line, option); ok && isLocalOptionValue(value) { + return option, true + } + } + for _, option := range []string{"-f", "-r"} { + if value, ok := shortOptionValue(line, option); ok && isLocalOptionValue(value) { + return option, true + } + } + return "", false +} + +func optionValue(line string, option string) (string, bool) { + if value, ok := strings.CutPrefix(line, option+"="); ok { + return strings.TrimSpace(value), true + } + if value, ok := strings.CutPrefix(line, option+" "); ok { + return strings.TrimSpace(value), true + } + return "", false +} + +func shortOptionValue(line string, option string) (string, bool) { + if value, ok := strings.CutPrefix(line, option+" "); ok { + return strings.TrimSpace(value), true + } + return "", false +} + +func isLocalOptionValue(value string) bool { + if value == "" || isRemoteRequirement(value) || strings.HasPrefix(value, "file:") { + return false + } + return true +} + +func isRemoteRequirement(line string) bool { + return strings.Contains(line, "://") || strings.HasPrefix(line, "git+") +} + +func isLocalPath(line string) bool { + return filepath.IsAbs(line) || strings.HasPrefix(line, "./") || strings.HasPrefix(line, "../") +} + +func hasLocalArtifactSuffix(path string) bool { + path = strings.ToLower(path) + for _, suffix := range localArtifactSuffixes { + if strings.HasSuffix(path, suffix) { + return true + } + } + return false +} diff --git a/pkg/requirements/requirements_test.go b/pkg/requirements/requirements_test.go index dabbfe9de3..17c0f200d9 100644 --- a/pkg/requirements/requirements_test.go +++ b/pkg/requirements/requirements_test.go @@ -20,6 +20,46 @@ func TestReadRequirements(t *testing.T) { require.Equal(t, []string{"torch==2.5.1"}, requirements) } +func TestParseLocalArtifactRequirement(t *testing.T) { + testCases := []struct { + name string + line string + expected string + expectedOK bool + expectsErr bool + }{ + {name: "Wheel", line: "./dist/pkg-0.1.0-py3-none-any.whl", expected: "./dist/pkg-0.1.0-py3-none-any.whl", expectedOK: true}, + {name: "Zip", line: "mylibpackage.zip", expected: "mylibpackage.zip", expectedOK: true}, + {name: "TarGz", line: "../pkg-0.1.0.tar.gz", expected: "../pkg-0.1.0.tar.gz", expectedOK: true}, + {name: "TarBz2", line: "/tmp/pkg-0.1.0.tar.bz2", expected: "/tmp/pkg-0.1.0.tar.bz2", expectedOK: true}, + {name: "Package", line: "torch==2.5.1"}, + {name: "URL", line: "https://example.com/pkg.zip"}, + {name: "VCS", line: "git+https://example.com/repo.git"}, + {name: "FileURL", line: "name @ file:./pkg.whl", expectsErr: true}, + {name: "DirectFileURL", line: "file:///tmp/pkg.whl", expectsErr: true}, + {name: "NamedFileURL", line: "name @ file:///tmp/pkg.whl", expectsErr: true}, + {name: "LocalFindLinks", line: "--find-links ./wheels", expectsErr: true}, + {name: "LocalFindLinksEquals", line: "--find-links=./wheels", expectsErr: true}, + {name: "RemoteFindLinks", line: "--find-links https://example.com/wheels"}, + {name: "LocalRecursiveRequirement", line: "-r requirements-local.txt", expectsErr: true}, + {name: "InlineHash", line: "./pkg.whl --hash=sha256:abc", expectsErr: true}, + {name: "LocalDirectory", line: "./pkg", expectsErr: true}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + actual, ok, err := ParseLocalArtifactRequirement(tc.line) + if tc.expectsErr { + require.Error(t, err) + return + } + require.NoError(t, err) + require.Equal(t, tc.expectedOK, ok) + require.Equal(t, tc.expected, actual) + }) + } +} + func TestReadRequirementsLineContinuations(t *testing.T) { srcDir := t.TempDir() reqFile := path.Join(srcDir, "requirements.txt") From c8109cc9ba44604626ddbe0beddc2c79e089ff8f Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Thu, 16 Jul 2026 12:42:44 -0500 Subject: [PATCH 2/5] fix: reject local direct requirement references --- pkg/requirements/local_artifact.go | 6 ++++++ pkg/requirements/requirements_test.go | 4 ++++ 2 files changed, 10 insertions(+) diff --git a/pkg/requirements/local_artifact.go b/pkg/requirements/local_artifact.go index 4da4388e03..068168d448 100644 --- a/pkg/requirements/local_artifact.go +++ b/pkg/requirements/local_artifact.go @@ -27,6 +27,12 @@ func ParseLocalArtifactRequirement(line string) (string, bool, error) { if strings.HasPrefix(line, "file:") || strings.Contains(line, " @ file:") { return "", false, fmt.Errorf("local file URL requirements are not supported: %s", line) } + if _, path, ok := strings.Cut(line, " @ "); ok { + path = strings.TrimSpace(path) + if !isRemoteRequirement(path) && (isLocalPath(path) || hasLocalArtifactSuffix(path)) { + return "", false, fmt.Errorf("local direct reference requirements (\"name @ path\") are not supported; list the path directly instead: %s", line) + } + } if option, ok := parseUnsupportedLocalOption(line); ok { return "", false, fmt.Errorf("local requirements option %q is not supported: %s", option, line) } diff --git a/pkg/requirements/requirements_test.go b/pkg/requirements/requirements_test.go index 17c0f200d9..08fe57503c 100644 --- a/pkg/requirements/requirements_test.go +++ b/pkg/requirements/requirements_test.go @@ -44,6 +44,10 @@ func TestParseLocalArtifactRequirement(t *testing.T) { {name: "LocalRecursiveRequirement", line: "-r requirements-local.txt", expectsErr: true}, {name: "InlineHash", line: "./pkg.whl --hash=sha256:abc", expectsErr: true}, {name: "LocalDirectory", line: "./pkg", expectsErr: true}, + {name: "DirectRefRelative", line: "name @ ./pkg.whl", expectsErr: true}, + {name: "DirectRefAbsolute", line: "name @ /tmp/pkg.whl", expectsErr: true}, + {name: "DirectRefBareRelative", line: "name @ pkg.whl", expectsErr: true}, + {name: "DirectRefRemote", line: "name @ https://example.com/pkg.whl"}, } for _, tc := range testCases { From 1f5080c1f17f22b097f6555306f556a9a88a06b8 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Mon, 3 Aug 2026 16:52:43 -0500 Subject: [PATCH 3/5] refactor: simplify local package artifacts --- architecture/05-build-system.md | 2 +- docs/llms.txt | 8 +- docs/yaml.md | 8 +- .../local_python_requirement_artifact.txtar | 15 ++- pkg/config/config.go | 30 ++--- pkg/config/config_test.go | 108 +++--------------- pkg/dockerfile/standard_generator.go | 88 +++++++++----- pkg/dockerfile/standard_generator_test.go | 57 +++++++-- pkg/requirements/local_artifact.go | 103 +++++++---------- pkg/requirements/requirements.go | 3 +- pkg/requirements/requirements_test.go | 56 +++++---- 11 files changed, 223 insertions(+), 255 deletions(-) diff --git a/architecture/05-build-system.md b/architecture/05-build-system.md index 0b8f0b9627..83fdf0394d 100644 --- a/architecture/05-build-system.md +++ b/architecture/05-build-system.md @@ -217,7 +217,7 @@ Resolution follows a 3-tier priority for each wheel: | 2. Auto-detect `dist/coglet-*.whl` | Dev builds only | | 3. Default | Install from PyPI | -Local wheel files are copied into `.cog/build/` and referenced via the `cog_build` named build context, then `COPY --from=cog_build`'d and `pip install`'d in the Dockerfile. User-provided local wheel and source archive requirements use the same staging boundary: Cog copies only the referenced artifacts into `.cog/build/` before the requirements install, preserving the later `COPY . /src` source layer. +Cog stages local SDK, Coglet, and user-provided wheel or source archive requirements before installing dependencies. The later `COPY . /src` remains a separate source layer. --- diff --git a/docs/llms.txt b/docs/llms.txt index 7227b6e420..b2b38c7e50 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -3483,7 +3483,7 @@ Your `cog.yaml` file can set either `python_packages` or `python_requirements`, This follows the standard [requirements.txt](https://pip.pypa.io/en/stable/reference/requirements-file-format/) format. -Requirements files can also reference local Python package artifacts, such as wheels and source archives: +Requirements files can list a local wheel or source archive: `requirements.txt`: @@ -3493,9 +3493,9 @@ Requirements files can also reference local Python package artifacts, such as wh ./packages/localpkg.tar.gz ``` -Local artifact paths are resolved relative to the requirements file, and the referenced files must be inside your project directory. Cog stages these artifacts before installing requirements, so this is the supported way to install a local package artifact during `cog build`. +Cog supports `.whl`, `.zip`, `.tar.gz`, `.tgz`, `.tar.bz2`, and `.tar.xz` files. Paths are resolved relative to the requirements file and must stay inside the project directory. -Local package directories, `name @ file:...` requirements, local `--find-links` directories, recursive local artifact includes, and inline hashes or options on local artifact lines are not supported. +Only bare paths are supported. Local directories, direct references, and options, hashes, extras, or markers on a local artifact line are rejected. Use `build.sdk_version` or `COG_SDK_WHEEL` for Cog, and `COGLET_WHEEL` for Coglet. To install Git-hosted Python packages, add `git` to the `system_packages` list, then use the `git+https://` syntax to specify the package name. For example: @@ -3574,7 +3574,7 @@ build: - cd cowsay-3.7.0 && make install ``` -Your code is _not_ available to commands in `run`. This is so we can build your image efficiently when running locally. To install a local wheel or source archive, list it in your `python_requirements` file instead of running `pip install ./artifact.zip` from `run`. +Your source code is not available to `run` commands. List local wheels and source archives in `python_requirements` instead. Each command in `run` can be either a string or a dictionary in the following format: diff --git a/docs/yaml.md b/docs/yaml.md index f8eb942165..1a45b55b4a 100644 --- a/docs/yaml.md +++ b/docs/yaml.md @@ -59,7 +59,7 @@ Your `cog.yaml` file can set either `python_packages` or `python_requirements`, This follows the standard [requirements.txt](https://pip.pypa.io/en/stable/reference/requirements-file-format/) format. -Requirements files can also reference local Python package artifacts, such as wheels and source archives: +Requirements files can list a local wheel or source archive: `requirements.txt`: @@ -69,9 +69,9 @@ Requirements files can also reference local Python package artifacts, such as wh ./packages/localpkg.tar.gz ``` -Local artifact paths are resolved relative to the requirements file, and the referenced files must be inside your project directory. Cog stages these artifacts before installing requirements, so this is the supported way to install a local package artifact during `cog build`. +Cog supports `.whl`, `.zip`, `.tar.gz`, `.tgz`, `.tar.bz2`, and `.tar.xz` files. Paths are resolved relative to the requirements file and must stay inside the project directory. -Local package directories, `name @ file:...` requirements, local `--find-links` directories, recursive local artifact includes, and inline hashes or options on local artifact lines are not supported. +Only bare paths are supported. Local directories, direct references, and options, hashes, extras, or markers on a local artifact line are rejected. Use `build.sdk_version` or `COG_SDK_WHEEL` for Cog, and `COGLET_WHEEL` for Coglet. To install Git-hosted Python packages, add `git` to the `system_packages` list, then use the `git+https://` syntax to specify the package name. For example: @@ -150,7 +150,7 @@ build: - cd cowsay-3.7.0 && make install ``` -Your code is _not_ available to commands in `run`. This is so we can build your image efficiently when running locally. To install a local wheel or source archive, list it in your `python_requirements` file instead of running `pip install ./artifact.zip` from `run`. +Your source code is not available to `run` commands. List local wheels and source archives in `python_requirements` instead. Each command in `run` can be either a string or a dictionary in the following format: diff --git a/integration-tests/tests/local_python_requirement_artifact.txtar b/integration-tests/tests/local_python_requirement_artifact.txtar index 089d5083f9..842bd9cd8a 100644 --- a/integration-tests/tests/local_python_requirement_artifact.txtar +++ b/integration-tests/tests/local_python_requirement_artifact.txtar @@ -1,11 +1,13 @@ -# Local package artifacts listed in requirements.txt are available during build. - mkdir dist -exec python3 -c 'import zipfile; z = zipfile.ZipFile("dist/local_pkg-0.1.0.zip", "w"); z.write("setup.py"); z.write("local_pkg/__init__.py"); z.close()' +exec python3 -m zipfile -c dist/local_pkg-0.1.0.zip setup.py local_pkg/__init__.py + +# Ensure the import depends on the artifact being installed. +rm local_pkg +rm setup.py cog build -t $TEST_IMAGE cog predict $TEST_IMAGE -stdout 'hello from local package' +stdout 'hello from local package 0.1.0' -- cog.yaml -- build: @@ -25,10 +27,13 @@ setup(name="local-pkg", version="0.1.0", packages=["local_pkg"]) MESSAGE = "hello from local package" -- predict.py -- +from importlib.metadata import version + from cog import BasePredictor + from local_pkg import MESSAGE class Predictor(BasePredictor): def predict(self) -> str: - return MESSAGE + return f"{MESSAGE} {version('local-pkg')}" diff --git a/pkg/config/config.go b/pkg/config/config.go index db62a9482e..0e56ad514d 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -1,8 +1,6 @@ package config import ( - "crypto/sha256" - "encoding/hex" "encoding/json" "fmt" "os" @@ -66,11 +64,11 @@ type Build struct { localPackageArtifacts []LocalPackageArtifact } +// LocalPackageArtifact is a local requirement staged into the Docker build context. type LocalPackageArtifact struct { - Requirement string - SourcePath string - StagedDir string - Filename string + Requirement string + SourcePath string + RelativePath string } type Concurrency struct { @@ -322,10 +320,10 @@ func (c *Config) loadLocalPackageArtifacts(projectDir string, requirementsFilePa } requirementsDir := filepath.Dir(requirementsFilePath) - seen := map[string]LocalPackageArtifact{} artifacts := []LocalPackageArtifact{} for _, line := range c.Build.pythonRequirementsContent { - artifactPath, ok, err := requirements.ParseLocalArtifactRequirement(line) + requirement := strings.TrimSpace(line) + artifactPath, ok, err := requirements.ParseLocalArtifactRequirement(requirement) if err != nil { return err } @@ -356,25 +354,11 @@ func (c *Config) loadLocalPackageArtifacts(projectDir string, requirementsFilePa return fmt.Errorf("local Python package artifact %q must be inside the project directory", artifactPath) } - if artifact, ok := seen[canonicalPath]; ok { - artifact.Requirement = line - artifacts = append(artifacts, artifact) - continue - } - relPath, err := filepath.Rel(projectRoot, canonicalPath) if err != nil { return fmt.Errorf("failed to resolve local Python package artifact %q relative to project directory: %w", artifactPath, err) } - hash := sha256.Sum256([]byte(filepath.ToSlash(relPath))) - artifact := LocalPackageArtifact{ - Requirement: line, - SourcePath: canonicalPath, - StagedDir: hex.EncodeToString(hash[:])[:16], - Filename: filepath.Base(absPath), - } - seen[canonicalPath] = artifact - artifacts = append(artifacts, artifact) + artifacts = append(artifacts, LocalPackageArtifact{requirement, canonicalPath, relPath}) } c.Build.localPackageArtifacts = artifacts return nil diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index 0494cf3595..4c80f869f4 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -200,78 +200,20 @@ func TestPythonRequirementsLocalPackageArtifacts(t *testing.T) { tmpDir := t.TempDir() require.NoError(t, os.MkdirAll(path.Join(tmpDir, "requirements", "dist"), 0o755)) wheelPath := path.Join(tmpDir, "requirements", "dist", "local_pkg-0.1.0-py3-none-any.whl") - archivePath := path.Join(tmpDir, "requirements", "mylibpackage.zip") require.NoError(t, os.WriteFile(wheelPath, []byte("wheel"), 0o644)) - require.NoError(t, os.WriteFile(archivePath, []byte("zip"), 0o644)) - require.NoError(t, os.WriteFile(path.Join(tmpDir, "requirements", "requirements.txt"), []byte(`torch==1.13.1 -./dist/local_pkg-0.1.0-py3-none-any.whl -mylibpackage.zip`), 0o644)) + require.NoError(t, os.Symlink(wheelPath, path.Join(tmpDir, "requirements", "dist", "latest.whl"))) + require.NoError(t, os.WriteFile(path.Join(tmpDir, "requirements", "requirements.txt"), []byte("./dist/latest.whl # vendored helper"), 0o644)) - config := &Config{ - Build: &Build{ - GPU: true, - PythonVersion: "3.10", - PythonRequirements: "requirements/requirements.txt", - }, - } + config := &Config{Build: &Build{PythonVersion: "3.10", PythonRequirements: "requirements/requirements.txt"}} require.NoError(t, config.Complete(tmpDir)) artifacts := config.LocalPackageArtifacts() canonicalWheelPath, err := filepath.EvalSymlinks(wheelPath) require.NoError(t, err) - canonicalArchivePath, err := filepath.EvalSymlinks(archivePath) - require.NoError(t, err) - require.Len(t, artifacts, 2) - require.Equal(t, "./dist/local_pkg-0.1.0-py3-none-any.whl", artifacts[0].Requirement) - require.Equal(t, canonicalWheelPath, artifacts[0].SourcePath) - require.Equal(t, "local_pkg-0.1.0-py3-none-any.whl", artifacts[0].Filename) - require.Len(t, artifacts[0].StagedDir, 16) - require.Equal(t, "mylibpackage.zip", artifacts[1].Requirement) - require.Equal(t, canonicalArchivePath, artifacts[1].SourcePath) - - requirements, err := config.PythonRequirementsForArch("linux", "amd64", []string{}) - require.NoError(t, err) - require.Equal(t, `--extra-index-url https://download.pytorch.org/whl/cu117/ -torch==1.13.1 -./dist/local_pkg-0.1.0-py3-none-any.whl -mylibpackage.zip`, requirements) -} - -func TestPythonRequirementsLocalPackageArtifactDuplicateAliases(t *testing.T) { - tmpDir := t.TempDir() - require.NoError(t, os.MkdirAll(path.Join(tmpDir, "dist"), 0o755)) - wheelPath := path.Join(tmpDir, "dist", "local_pkg-0.1.0-py3-none-any.whl") - require.NoError(t, os.WriteFile(wheelPath, []byte("wheel"), 0o644)) - require.NoError(t, os.WriteFile(path.Join(tmpDir, "requirements.txt"), []byte("./dist/local_pkg-0.1.0-py3-none-any.whl\ndist/local_pkg-0.1.0-py3-none-any.whl"), 0o644)) - - config := &Config{Build: &Build{PythonVersion: "3.10", PythonRequirements: "requirements.txt"}} - require.NoError(t, config.Complete(tmpDir)) - - artifacts := config.LocalPackageArtifacts() - require.Len(t, artifacts, 2) - require.Equal(t, "./dist/local_pkg-0.1.0-py3-none-any.whl", artifacts[0].Requirement) - require.Equal(t, "dist/local_pkg-0.1.0-py3-none-any.whl", artifacts[1].Requirement) - require.Equal(t, artifacts[0].SourcePath, artifacts[1].SourcePath) - require.Equal(t, artifacts[0].StagedDir, artifacts[1].StagedDir) - require.Equal(t, artifacts[0].Filename, artifacts[1].Filename) -} - -func TestPythonRequirementsLocalPackageArtifactUnsupportedLocalOptions(t *testing.T) { - for _, line := range []string{ - "--find-links ./wheels", - "--find-links=./wheels", - "-r requirements-local.txt", - "--requirement requirements-local.txt", - } { - t.Run(line, func(t *testing.T) { - projectDir := t.TempDir() - require.NoError(t, os.WriteFile(path.Join(projectDir, "requirements.txt"), []byte(line), 0o644)) - config := &Config{Build: &Build{PythonVersion: "3.10", PythonRequirements: "requirements.txt"}} - err := config.Complete(projectDir) - require.Error(t, err) - require.Contains(t, err.Error(), "local requirements option") - }) - } + require.Len(t, artifacts, 1) + assert.Equal(t, "./dist/latest.whl", artifacts[0].Requirement) + assert.Equal(t, canonicalWheelPath, artifacts[0].SourcePath) + assert.Equal(t, path.Join("requirements", "dist", "local_pkg-0.1.0-py3-none-any.whl"), artifacts[0].RelativePath) } func TestPythonRequirementsLocalPackageArtifactValidation(t *testing.T) { @@ -281,46 +223,30 @@ func TestPythonRequirementsLocalPackageArtifactValidation(t *testing.T) { setup func(t *testing.T, projectDir string) string expectedErr string }{ + {name: "Missing", line: "./missing.zip", expectedErr: "not found"}, { - name: "MissingArtifact", - line: "./missing.zip", - expectedErr: "not found", - }, - { - name: "DirectoryArtifact", - line: "./pkg.zip", + name: "Directory", line: "./pkg.zip", expectedErr: "regular file", setup: func(t *testing.T, projectDir string) string { require.NoError(t, os.Mkdir(path.Join(projectDir, "pkg.zip"), 0o755)) return "" }, - expectedErr: "regular file", }, { - name: "OutsideProject", + name: "OutsideProject", expectedErr: "inside the project directory", setup: func(t *testing.T, projectDir string) string { - outsideDir := t.TempDir() - outsidePath := path.Join(outsideDir, "pkg.zip") + outsidePath := path.Join(t.TempDir(), "pkg.zip") require.NoError(t, os.WriteFile(outsidePath, []byte("zip"), 0o644)) return outsidePath }, - expectedErr: "inside the project directory", }, { - name: "SymlinkOutsideProject", - line: "./pkg.zip", + name: "SymlinkOutsideProject", line: "./pkg.zip", expectedErr: "inside the project directory", setup: func(t *testing.T, projectDir string) string { - outsideDir := t.TempDir() - outsidePath := path.Join(outsideDir, "pkg.zip") + outsidePath := path.Join(t.TempDir(), "pkg.zip") require.NoError(t, os.WriteFile(outsidePath, []byte("zip"), 0o644)) require.NoError(t, os.Symlink(outsidePath, path.Join(projectDir, "pkg.zip"))) return "" }, - expectedErr: "inside the project directory", - }, - { - name: "InlineHashUnsupported", - line: "./pkg.whl --hash=sha256:abc", - expectedErr: "inline options or hashes", }, } @@ -329,15 +255,13 @@ func TestPythonRequirementsLocalPackageArtifactValidation(t *testing.T) { projectDir := t.TempDir() line := tc.line if tc.setup != nil { - if setupLine := tc.setup(t, projectDir); setupLine != "" { - line = setupLine + if value := tc.setup(t, projectDir); value != "" { + line = value } } require.NoError(t, os.WriteFile(path.Join(projectDir, "requirements.txt"), []byte(line), 0o644)) config := &Config{Build: &Build{PythonVersion: "3.10", PythonRequirements: "requirements.txt"}} - err := config.Complete(projectDir) - require.Error(t, err) - require.Contains(t, err.Error(), tc.expectedErr) + require.ErrorContains(t, config.Complete(projectDir), tc.expectedErr) }) } } diff --git a/pkg/dockerfile/standard_generator.go b/pkg/dockerfile/standard_generator.go index 4613f8a6ab..1cd1283009 100644 --- a/pkg/dockerfile/standard_generator.go +++ b/pkg/dockerfile/standard_generator.go @@ -30,6 +30,11 @@ const uvBreakSystemPackages = "--break-system-packages" const PrecompilePythonCommand = "RUN find / -type f -name \"*.py[co]\" -delete && find / -type f -name \"*.py\" -exec touch -t 197001010000 {} \\; && find / -type f -name \"*.py\" -printf \"%h\\n\" | sort -u | /usr/bin/python3 -m compileall --invalidation-mode timestamp -o 2 -j 0" const STANDARD_GENERATOR_NAME = "STANDARD_GENERATOR" +// localArtifactsDir is the directory local package artifacts are staged into, +// both under the build context on the host and under /tmp in the container. +const localArtifactsDir = "local_package_artifacts" +const localArtifactsContainerDir = "/tmp/" + localArtifactsDir + type StandardGenerator struct { Config *config.Config Dir string @@ -864,7 +869,7 @@ func (g *StandardGenerator) filterManagedPackages(reqContents string) string { "Remove it from requirements and use build.sdk_version in cog.yaml or %s to control the version.", trimmed, override(baseName), - map[string]string{"cog": "COG_SDK_WHEEL", "coglet": "COGLET_WHEEL"}[baseName], + map[string]string{"cog": wheels.CogSDKWheelEnvVar, "coglet": wheels.CogletWheelEnvVar}[baseName], ) continue } @@ -903,11 +908,11 @@ func (g *StandardGenerator) pipInstalls() (string, error) { // Strip cog/coglet from user requirements — we always install them ourselves // via installCog(). Leaving them in would cause pip to overwrite our version. g.pythonRequirementsContents = g.filterManagedPackages(g.pythonRequirementsContents) - artifactCopyLines, err := g.stageLocalPackageArtifacts() + var artifactCopyLine string + g.pythonRequirementsContents, artifactCopyLine, err = g.stageLocalPackageArtifacts(g.pythonRequirementsContents) if err != nil { return "", err } - g.pythonRequirementsContents = g.rewriteLocalPackageArtifacts(g.pythonRequirementsContents) if strings.Trim(g.pythonRequirementsContents, "") == "" { return "", nil @@ -923,55 +928,76 @@ func (g *StandardGenerator) pipInstalls() (string, error) { if g.strip { pipInstallLine += " && " + StripDebugSymbolsCommand } - lines := []string{} - lines = append(lines, artifactCopyLines...) - lines = append(lines, + return strings.Join(filterEmpty([]string{ + artifactCopyLine, copyLine[0], CFlags, pipInstallLine, "ENV CFLAGS=", - ) - return strings.Join(lines, "\n"), nil + }), "\n"), nil } -func (g *StandardGenerator) stageLocalPackageArtifacts() ([]string, error) { +func (g *StandardGenerator) stageLocalPackageArtifacts(reqContents string) (string, string, error) { artifacts := g.Config.LocalPackageArtifacts() if len(artifacts) == 0 { - return nil, nil + return reqContents, "", nil } + + containerPaths := map[string]string{} staged := map[string]bool{} for _, artifact := range artifacts { - dst := filepath.Join(g.tmpDir, "local_package_artifacts", artifact.StagedDir, artifact.Filename) - if staged[dst] { - continue - } - if err := files.Copy(artifact.SourcePath, dst); err != nil { - return nil, fmt.Errorf("failed to stage local Python package artifact %s: %w", artifact.SourcePath, err) + filename := filepath.Base(artifact.RelativePath) + switch { + case isVersionedArtifact(filename, "cog"): + return "", "", fmt.Errorf("local cog artifact %q is not supported; use build.sdk_version or %s", artifact.Requirement, wheels.CogSDKWheelEnvVar) + case isVersionedArtifact(filename, "coglet"): + return "", "", fmt.Errorf("local coglet artifact %q is not supported; use %s", artifact.Requirement, wheels.CogletWheelEnvVar) } - staged[dst] = true - } - return []string{"COPY --from=cog_build local_package_artifacts/ /tmp/local_package_artifacts/"}, nil -} -func (g *StandardGenerator) rewriteLocalPackageArtifacts(reqContents string) string { - artifacts := map[string]string{} - for _, artifact := range g.Config.LocalPackageArtifacts() { - containerPath := path.Join("/tmp/local_package_artifacts", artifact.StagedDir, artifact.Filename) - artifacts[artifact.Requirement] = containerPath - } - if len(artifacts) == 0 { - return reqContents + if !staged[artifact.SourcePath] { + dst := filepath.Join(g.tmpDir, localArtifactsDir, artifact.RelativePath) + if err := files.Copy(artifact.SourcePath, dst); err != nil { + return "", "", fmt.Errorf("failed to stage local Python package artifact %s: %w", artifact.SourcePath, err) + } + staged[artifact.SourcePath] = true + } + containerPaths[artifact.Requirement] = path.Join(localArtifactsContainerDir, filepath.ToSlash(artifact.RelativePath)) } - lines := []string{} for line := range strings.SplitSeq(reqContents, "\n") { - if replacement, ok := artifacts[strings.TrimSpace(line)]; ok { + requirement := strings.TrimSpace(line) + if replacement, ok := containerPaths[requirement]; ok { lines = append(lines, replacement) } else { lines = append(lines, line) } } - return strings.Join(lines, "\n") + return strings.Join(lines, "\n"), fmt.Sprintf("COPY --from=cog_build %s/ %s/", localArtifactsDir, localArtifactsContainerDir), nil +} + +func isVersionedArtifact(filename, name string) bool { + version, ok := strings.CutPrefix(strings.ToLower(filename), name+"-") + if !ok { + return false + } + major, rest, ok := strings.Cut(version, ".") + if !ok || !isDigits(major) { + return false + } + minor, _, _ := strings.Cut(rest, ".") + return isDigits(minor) +} + +func isDigits(value string) bool { + if value == "" { + return false + } + for _, char := range value { + if char < '0' || char > '9' { + return false + } + } + return true } func (g *StandardGenerator) runCommands() (string, error) { diff --git a/pkg/dockerfile/standard_generator_test.go b/pkg/dockerfile/standard_generator_test.go index d4f7082189..2555dd8e0b 100644 --- a/pkg/dockerfile/standard_generator_test.go +++ b/pkg/dockerfile/standard_generator_test.go @@ -9,6 +9,7 @@ import ( "testing" "time" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/replicate/cog/pkg/config" @@ -344,7 +345,8 @@ func TestPythonRequirementsLocalPackageArtifact(t *testing.T) { require.NoError(t, os.MkdirAll(path.Join(tmpDir, "requirements", "dist"), 0o755)) artifactPath := path.Join(tmpDir, "requirements", "dist", "local_pkg-0.1.0-py3-none-any.whl") require.NoError(t, os.WriteFile(artifactPath, []byte("wheel"), 0o644)) - require.NoError(t, os.WriteFile(path.Join(tmpDir, "requirements", "requirements.txt"), []byte("./dist/local_pkg-0.1.0-py3-none-any.whl\ndist/local_pkg-0.1.0-py3-none-any.whl"), 0o644)) + require.NoError(t, os.Symlink(artifactPath, path.Join(tmpDir, "requirements", "dist", "latest.whl"))) + require.NoError(t, os.WriteFile(path.Join(tmpDir, "requirements", "requirements.txt"), []byte("./dist/latest.whl # vendored helper\ndist/latest.whl"), 0o644)) conf, err := config.FromYAML([]byte(` build: @@ -365,19 +367,60 @@ build: artifact := conf.LocalPackageArtifacts()[0] copyArtifacts := "COPY --from=cog_build local_package_artifacts/ /tmp/local_package_artifacts/" - copyRequirements := "COPY --from=cog_build requirements.txt /tmp/requirements.txt" pipInstall := "uv run pip install --cache-dir /root/.cache/pip -r /tmp/requirements.txt" require.Contains(t, actual, copyArtifacts) - require.Less(t, strings.Index(actual, copyArtifacts), strings.Index(actual, copyRequirements)) - require.Less(t, strings.Index(actual, copyRequirements), strings.Index(actual, pipInstall)) + require.Contains(t, actual, pipInstall) + assert.Less(t, strings.Index(actual, copyArtifacts), strings.Index(actual, pipInstall)) - stagedArtifact, err := os.ReadFile(path.Join(buildDir, "local_package_artifacts", artifact.StagedDir, artifact.Filename)) + stagedArtifact, err := os.ReadFile(path.Join(buildDir, localArtifactsDir, artifact.RelativePath)) require.NoError(t, err) - require.Equal(t, []byte("wheel"), stagedArtifact) + assert.Equal(t, []byte("wheel"), stagedArtifact) requirements, err := os.ReadFile(path.Join(buildDir, "requirements.txt")) require.NoError(t, err) - require.Equal(t, "/tmp/local_package_artifacts/"+artifact.StagedDir+"/local_pkg-0.1.0-py3-none-any.whl\n/tmp/local_package_artifacts/"+artifact.StagedDir+"/local_pkg-0.1.0-py3-none-any.whl", string(requirements)) + containerPath := path.Join(localArtifactsContainerDir, filepath.ToSlash(artifact.RelativePath)) + assert.Equal(t, containerPath+"\n"+containerPath, string(requirements)) +} + +func TestLocalPackageArtifactRejectsManagedPackages(t *testing.T) { + testCases := []struct { + name string + filename string + expectedErr string + }{ + {name: "Cog", filename: "cog-0.1.0-py3-none-any.whl", expectedErr: wheels.CogSDKWheelEnvVar}, + {name: "Coglet", filename: "COGLET-1.0.0.tar.gz", expectedErr: wheels.CogletWheelEnvVar}, + {name: "Cog2FA", filename: "cog-2fa-1.0.tar.gz"}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + tmpDir := t.TempDir() + require.NoError(t, os.MkdirAll(path.Join(tmpDir, "dist"), 0o755)) + require.NoError(t, os.WriteFile(path.Join(tmpDir, "dist", tc.filename), []byte("artifact"), 0o644)) + require.NoError(t, os.WriteFile(path.Join(tmpDir, "requirements.txt"), []byte("./dist/"+tc.filename), 0o644)) + + conf, err := config.FromYAML([]byte(` +build: + python_version: "3.12" + python_requirements: "requirements.txt" +`)) + require.NoError(t, err) + require.NoError(t, conf.Complete(tmpDir)) + gen, err := NewStandardGenerator(conf, tmpDir, t.TempDir(), "", dockertest.NewMockCommand(), registrytest.NewMockRegistryClient(), true) + require.NoError(t, err) + gen.SetUseCogBaseImage(false) + pypiWheels(gen) + + _, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), "r8.im/replicate/cog-test") + if tc.expectedErr != "" { + require.ErrorContains(t, err, tc.expectedErr) + return + } + require.NoError(t, err) + assert.Contains(t, actual, localArtifactsContainerDir+"/") + }) + } } // GPU builds on nvidia/cuda base images install Python via `uv python install` diff --git a/pkg/requirements/local_artifact.go b/pkg/requirements/local_artifact.go index 068168d448..62ca83022a 100644 --- a/pkg/requirements/local_artifact.go +++ b/pkg/requirements/local_artifact.go @@ -6,94 +6,75 @@ import ( "strings" ) -var localArtifactSuffixes = []string{ - ".whl", - ".zip", - ".tar.gz", - ".tgz", - ".tar.bz2", - ".tar.xz", -} - -// ParseLocalArtifactRequirement identifies simple local wheel/source-archive -// requirement lines. It intentionally does not parse full pip requirement -// syntax; callers should reject unsupported local forms with clear errors. +// ParseLocalArtifactRequirement returns a bare local wheel or source archive +// path. Other pip requirements pass through unchanged; unsupported local forms +// return an error before the Docker build starts. func ParseLocalArtifactRequirement(line string) (string, bool, error) { line = strings.TrimSpace(line) if line == "" { return "", false, nil } - - if strings.HasPrefix(line, "file:") || strings.Contains(line, " @ file:") { + if strings.HasPrefix(line, "file:") { return "", false, fmt.Errorf("local file URL requirements are not supported: %s", line) } - if _, path, ok := strings.Cut(line, " @ "); ok { - path = strings.TrimSpace(path) - if !isRemoteRequirement(path) && (isLocalPath(path) || hasLocalArtifactSuffix(path)) { - return "", false, fmt.Errorf("local direct reference requirements (\"name @ path\") are not supported; list the path directly instead: %s", line) + if name, target, ok := strings.Cut(line, "@"); ok { + spaced := name != strings.TrimSpace(name) || target != strings.TrimSpace(target) + name, target = strings.TrimSpace(name), strings.TrimSpace(target) + if name != "" && PackageName(name) == name && strings.HasPrefix(target, "file:") { + return "", false, fmt.Errorf("local file URL requirements are not supported: %s", line) + } + if name != "" && PackageName(name) == name && (isLocalPath(target) || spaced && isLocalArtifact(target)) { + return "", false, fmt.Errorf("local direct reference requirements (`name @ path`) are not supported; list the path directly instead: %s", line) } } - if option, ok := parseUnsupportedLocalOption(line); ok { - return "", false, fmt.Errorf("local requirements option %q is not supported: %s", option, line) + if isRemoteRequirement(line) { + return "", false, nil } - if strings.HasPrefix(line, "-") || isRemoteRequirement(line) { + if strings.HasPrefix(line, "-") { + if option := unsupportedLocalOption(line); option != "" { + return "", false, fmt.Errorf("local requirements option %q is not supported: %s", option, line) + } return "", false, nil } + if base, _, ok := strings.Cut(line, ";"); ok && isLocalArtifact(strings.TrimSpace(base)) { + return "", false, fmt.Errorf("environment markers are not supported on local package artifact requirements: %s", line) + } + if base, _, ok := strings.Cut(line, "["); ok && strings.HasSuffix(line, "]") && isLocalArtifact(strings.TrimSpace(base)) { + return "", false, fmt.Errorf("extras are not supported on local package artifact requirements: %s", line) + } fields := strings.Fields(line) if len(fields) > 1 { - if hasLocalArtifactSuffix(fields[0]) || isLocalPath(fields[0]) { + if isLocalArtifact(fields[0]) { return "", false, fmt.Errorf("local package artifact requirements do not support inline options or hashes: %s", line) } return "", false, nil } - if !isLocalPath(line) && !hasLocalArtifactSuffix(line) { return "", false, nil } if !hasLocalArtifactSuffix(line) { return "", false, fmt.Errorf("local package requirement %q is not a supported wheel or source archive", line) } - return line, true, nil } -func parseUnsupportedLocalOption(line string) (string, bool) { - for _, option := range []string{"--find-links", "--requirement"} { - if value, ok := optionValue(line, option); ok && isLocalOptionValue(value) { - return option, true +func unsupportedLocalOption(line string) string { + for _, option := range []string{"--find-links", "--requirement", "-f", "-r"} { + value, ok := strings.CutPrefix(line, option+" ") + if !ok && strings.HasPrefix(option, "--") { + value, ok = strings.CutPrefix(line, option+"=") } - } - for _, option := range []string{"-f", "-r"} { - if value, ok := shortOptionValue(line, option); ok && isLocalOptionValue(value) { - return option, true + value = strings.TrimSpace(value) + if ok && value != "" && !isRemoteRequirement(value) && !strings.HasPrefix(value, "file:") { + return option } } - return "", false + return "" } -func optionValue(line string, option string) (string, bool) { - if value, ok := strings.CutPrefix(line, option+"="); ok { - return strings.TrimSpace(value), true - } - if value, ok := strings.CutPrefix(line, option+" "); ok { - return strings.TrimSpace(value), true - } - return "", false -} - -func shortOptionValue(line string, option string) (string, bool) { - if value, ok := strings.CutPrefix(line, option+" "); ok { - return strings.TrimSpace(value), true - } - return "", false -} - -func isLocalOptionValue(value string) bool { - if value == "" || isRemoteRequirement(value) || strings.HasPrefix(value, "file:") { - return false - } - return true +func isLocalArtifact(path string) bool { + return !isRemoteRequirement(path) && (isLocalPath(path) || hasLocalArtifactSuffix(path)) } func isRemoteRequirement(line string) bool { @@ -106,10 +87,10 @@ func isLocalPath(line string) bool { func hasLocalArtifactSuffix(path string) bool { path = strings.ToLower(path) - for _, suffix := range localArtifactSuffixes { - if strings.HasSuffix(path, suffix) { - return true - } - } - return false + return strings.HasSuffix(path, ".whl") || + strings.HasSuffix(path, ".zip") || + strings.HasSuffix(path, ".tar.gz") || + strings.HasSuffix(path, ".tgz") || + strings.HasSuffix(path, ".tar.bz2") || + strings.HasSuffix(path, ".tar.xz") } diff --git a/pkg/requirements/requirements.go b/pkg/requirements/requirements.go index b3b8ac59ab..d9810a513a 100644 --- a/pkg/requirements/requirements.go +++ b/pkg/requirements/requirements.go @@ -35,9 +35,8 @@ func ReadRequirements(path string) ([]string, error) { continue } - // Remove any trailing comments if idx := strings.Index(line, "#"); idx >= 0 { - line = line[:idx] + line = strings.TrimSpace(line[:idx]) } if line != "" { diff --git a/pkg/requirements/requirements_test.go b/pkg/requirements/requirements_test.go index 08fe57503c..364b253d28 100644 --- a/pkg/requirements/requirements_test.go +++ b/pkg/requirements/requirements_test.go @@ -6,6 +6,7 @@ import ( "strings" "testing" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -22,44 +23,49 @@ func TestReadRequirements(t *testing.T) { func TestParseLocalArtifactRequirement(t *testing.T) { testCases := []struct { - name string - line string - expected string - expectedOK bool - expectsErr bool + name string + line string + expected string + expectedOK bool + expectedErr string }{ - {name: "Wheel", line: "./dist/pkg-0.1.0-py3-none-any.whl", expected: "./dist/pkg-0.1.0-py3-none-any.whl", expectedOK: true}, + {name: "Wheel", line: " ./dist/PKG-0.1.0-PY3-NONE-ANY.WHL ", expected: "./dist/PKG-0.1.0-PY3-NONE-ANY.WHL", expectedOK: true}, {name: "Zip", line: "mylibpackage.zip", expected: "mylibpackage.zip", expectedOK: true}, {name: "TarGz", line: "../pkg-0.1.0.tar.gz", expected: "../pkg-0.1.0.tar.gz", expectedOK: true}, + {name: "Tgz", line: "./dist/pkg-0.1.0.tgz", expected: "./dist/pkg-0.1.0.tgz", expectedOK: true}, {name: "TarBz2", line: "/tmp/pkg-0.1.0.tar.bz2", expected: "/tmp/pkg-0.1.0.tar.bz2", expectedOK: true}, - {name: "Package", line: "torch==2.5.1"}, - {name: "URL", line: "https://example.com/pkg.zip"}, + {name: "TarXz", line: "./dist/pkg-0.1.0.tar.xz", expected: "./dist/pkg-0.1.0.tar.xz", expectedOK: true}, + {name: "AtInPath", line: "./dist/pkg@1.0.tar.gz", expected: "./dist/pkg@1.0.tar.gz", expectedOK: true}, + {name: "AtInFilename", line: "pkg@1.0.tar.gz", expected: "pkg@1.0.tar.gz", expectedOK: true}, + {name: "PackageWithExtras", line: "torch[all]==2.5.1"}, + {name: "PackageWithMarker", line: `torch==2.5.1; python_version < "3.11"`}, + {name: "URLWithHash", line: "https://user:pass@example.com/pkg.whl --hash=sha256:abc"}, + {name: "URLWithAtSign", line: "https://token@file:443/pkg.whl"}, {name: "VCS", line: "git+https://example.com/repo.git"}, - {name: "FileURL", line: "name @ file:./pkg.whl", expectsErr: true}, - {name: "DirectFileURL", line: "file:///tmp/pkg.whl", expectsErr: true}, - {name: "NamedFileURL", line: "name @ file:///tmp/pkg.whl", expectsErr: true}, - {name: "LocalFindLinks", line: "--find-links ./wheels", expectsErr: true}, - {name: "LocalFindLinksEquals", line: "--find-links=./wheels", expectsErr: true}, + {name: "RemoteDirectReference", line: "name @ https://example.com/pkg.whl"}, {name: "RemoteFindLinks", line: "--find-links https://example.com/wheels"}, - {name: "LocalRecursiveRequirement", line: "-r requirements-local.txt", expectsErr: true}, - {name: "InlineHash", line: "./pkg.whl --hash=sha256:abc", expectsErr: true}, - {name: "LocalDirectory", line: "./pkg", expectsErr: true}, - {name: "DirectRefRelative", line: "name @ ./pkg.whl", expectsErr: true}, - {name: "DirectRefAbsolute", line: "name @ /tmp/pkg.whl", expectsErr: true}, - {name: "DirectRefBareRelative", line: "name @ pkg.whl", expectsErr: true}, - {name: "DirectRefRemote", line: "name @ https://example.com/pkg.whl"}, + {name: "DirectFileURL", line: "file:///tmp/pkg.whl", expectedErr: "local file URL requirements are not supported"}, + {name: "NamedFileURLNoSpace", line: "name@file:./pkg.whl", expectedErr: "local file URL requirements are not supported"}, + {name: "LocalFindLinks", line: "--find-links ./wheels", expectedErr: `local requirements option "--find-links" is not supported`}, + {name: "LocalRequirement", line: "-r requirements-local.txt", expectedErr: `local requirements option "-r" is not supported`}, + {name: "InlineHash", line: "./pkg.whl --hash=sha256:abc", expectedErr: "do not support inline options or hashes"}, + {name: "LocalDirectory", line: "./pkg", expectedErr: "is not a supported wheel or source archive"}, + {name: "ArtifactWithMarker", line: `./dist/pkg.whl; python_version < "3.11"`, expectedErr: "environment markers are not supported"}, + {name: "ArtifactWithExtras", line: "./dist/pkg.whl[extra]", expectedErr: "extras are not supported"}, + {name: "DirectRefRelative", line: "name @ ./pkg.whl", expectedErr: "local direct reference requirements"}, + {name: "DirectRefNoSpace", line: "name@./pkg.whl", expectedErr: "local direct reference requirements"}, } for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { actual, ok, err := ParseLocalArtifactRequirement(tc.line) - if tc.expectsErr { - require.Error(t, err) + if tc.expectedErr != "" { + require.ErrorContains(t, err, tc.expectedErr) return } require.NoError(t, err) - require.Equal(t, tc.expectedOK, ok) - require.Equal(t, tc.expected, actual) + assert.Equal(t, tc.expectedOK, ok) + assert.Equal(t, tc.expected, actual) }) } } @@ -78,7 +84,7 @@ func TestReadRequirementsLineContinuations(t *testing.T) { func TestReadRequirementsStripComments(t *testing.T) { srcDir := t.TempDir() reqFile := path.Join(srcDir, "requirements.txt") - err := os.WriteFile(reqFile, []byte("torch==\\\n2.5.1# Heres my comment\ntorchvision==2.5.1\n# Heres a beginning of line comment"), 0o644) + err := os.WriteFile(reqFile, []byte("torch==\\\n2.5.1 # Heres my comment\ntorchvision==2.5.1\n# Heres a beginning of line comment"), 0o644) require.NoError(t, err) requirements, err := ReadRequirements(reqFile) From fa4890bccf6ea45ee30134c74dc5dcfd45082ead Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Tue, 4 Aug 2026 10:16:10 -0500 Subject: [PATCH 4/5] fix: reject uppercase file URLs and use tar.gz sdist in integration test Make the file: URL rejection case-insensitive so FILE:/// forms are not mistaken for remote requirements, and reorder the option check ahead of the remote short-circuit so file:-scheme options are rejected too. Switch the integration test sdist from python3 -m zipfile (which flattens paths and drops the package directory) to a tar.gz that preserves the package layout, so pip install succeeds. --- .../local_python_requirement_artifact.txtar | 4 ++-- pkg/requirements/local_artifact.go | 16 ++++++++++------ pkg/requirements/requirements_test.go | 4 ++++ 3 files changed, 16 insertions(+), 8 deletions(-) diff --git a/integration-tests/tests/local_python_requirement_artifact.txtar b/integration-tests/tests/local_python_requirement_artifact.txtar index 842bd9cd8a..1f9a2835b1 100644 --- a/integration-tests/tests/local_python_requirement_artifact.txtar +++ b/integration-tests/tests/local_python_requirement_artifact.txtar @@ -1,5 +1,5 @@ mkdir dist -exec python3 -m zipfile -c dist/local_pkg-0.1.0.zip setup.py local_pkg/__init__.py +exec tar -czf dist/local_pkg-0.1.0.tar.gz setup.py local_pkg # Ensure the import depends on the artifact being installed. rm local_pkg @@ -16,7 +16,7 @@ build: predict: predict.py:Predictor -- requirements.txt -- -./dist/local_pkg-0.1.0.zip +./dist/local_pkg-0.1.0.tar.gz -- setup.py -- from setuptools import setup diff --git a/pkg/requirements/local_artifact.go b/pkg/requirements/local_artifact.go index 62ca83022a..d0cc797bf7 100644 --- a/pkg/requirements/local_artifact.go +++ b/pkg/requirements/local_artifact.go @@ -14,28 +14,28 @@ func ParseLocalArtifactRequirement(line string) (string, bool, error) { if line == "" { return "", false, nil } - if strings.HasPrefix(line, "file:") { + if isFileURL(line) { return "", false, fmt.Errorf("local file URL requirements are not supported: %s", line) } if name, target, ok := strings.Cut(line, "@"); ok { spaced := name != strings.TrimSpace(name) || target != strings.TrimSpace(target) name, target = strings.TrimSpace(name), strings.TrimSpace(target) - if name != "" && PackageName(name) == name && strings.HasPrefix(target, "file:") { + if name != "" && PackageName(name) == name && isFileURL(target) { return "", false, fmt.Errorf("local file URL requirements are not supported: %s", line) } if name != "" && PackageName(name) == name && (isLocalPath(target) || spaced && isLocalArtifact(target)) { return "", false, fmt.Errorf("local direct reference requirements (`name @ path`) are not supported; list the path directly instead: %s", line) } } - if isRemoteRequirement(line) { - return "", false, nil - } if strings.HasPrefix(line, "-") { if option := unsupportedLocalOption(line); option != "" { return "", false, fmt.Errorf("local requirements option %q is not supported: %s", option, line) } return "", false, nil } + if isRemoteRequirement(line) { + return "", false, nil + } if base, _, ok := strings.Cut(line, ";"); ok && isLocalArtifact(strings.TrimSpace(base)) { return "", false, fmt.Errorf("environment markers are not supported on local package artifact requirements: %s", line) } @@ -66,7 +66,7 @@ func unsupportedLocalOption(line string) string { value, ok = strings.CutPrefix(line, option+"=") } value = strings.TrimSpace(value) - if ok && value != "" && !isRemoteRequirement(value) && !strings.HasPrefix(value, "file:") { + if ok && value != "" && (isFileURL(value) || !isRemoteRequirement(value)) { return option } } @@ -81,6 +81,10 @@ func isRemoteRequirement(line string) bool { return strings.Contains(line, "://") || strings.HasPrefix(line, "git+") } +func isFileURL(line string) bool { + return strings.HasPrefix(strings.ToLower(line), "file:") +} + func isLocalPath(line string) bool { return filepath.IsAbs(line) || strings.HasPrefix(line, "./") || strings.HasPrefix(line, "../") } diff --git a/pkg/requirements/requirements_test.go b/pkg/requirements/requirements_test.go index 364b253d28..4a39172965 100644 --- a/pkg/requirements/requirements_test.go +++ b/pkg/requirements/requirements_test.go @@ -45,9 +45,13 @@ func TestParseLocalArtifactRequirement(t *testing.T) { {name: "RemoteDirectReference", line: "name @ https://example.com/pkg.whl"}, {name: "RemoteFindLinks", line: "--find-links https://example.com/wheels"}, {name: "DirectFileURL", line: "file:///tmp/pkg.whl", expectedErr: "local file URL requirements are not supported"}, + {name: "UppercaseFileURL", line: "FILE:///tmp/pkg.whl", expectedErr: "local file URL requirements are not supported"}, {name: "NamedFileURLNoSpace", line: "name@file:./pkg.whl", expectedErr: "local file URL requirements are not supported"}, + {name: "NamedUppercaseFileURL", line: "name@FILE:./pkg.whl", expectedErr: "local file URL requirements are not supported"}, {name: "LocalFindLinks", line: "--find-links ./wheels", expectedErr: `local requirements option "--find-links" is not supported`}, + {name: "FileURLFindLinks", line: "--find-links file:///wheels", expectedErr: `local requirements option "--find-links" is not supported`}, {name: "LocalRequirement", line: "-r requirements-local.txt", expectedErr: `local requirements option "-r" is not supported`}, + {name: "UppercaseFileURLRequirement", line: "-r FILE:///tmp/requirements.txt", expectedErr: `local requirements option "-r" is not supported`}, {name: "InlineHash", line: "./pkg.whl --hash=sha256:abc", expectedErr: "do not support inline options or hashes"}, {name: "LocalDirectory", line: "./pkg", expectedErr: "is not a supported wheel or source archive"}, {name: "ArtifactWithMarker", line: `./dist/pkg.whl; python_version < "3.11"`, expectedErr: "environment markers are not supported"}, From e2340438354b15fed43b413fdf9c3a69accf7f6f Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 5 Aug 2026 16:23:00 -0500 Subject: [PATCH 5/5] fix: harden local package artifact handling --- docs/llms.txt | 4 +- docs/yaml.md | 4 +- .../local_python_requirement_artifact.txtar | 63 +++++- pkg/config/config.go | 59 ++++- pkg/config/config_test.go | 57 ++++- pkg/config/data/config_schema_v1.0.json | 2 +- pkg/config/validate.go | 3 + pkg/config/validate_test.go | 14 ++ pkg/dockerfile/standard_generator.go | 195 +++++++++++++---- pkg/dockerfile/standard_generator_test.go | 206 ++++++++++++++---- pkg/requirements/local_artifact.go | 54 ++++- pkg/requirements/requirements_test.go | 8 + 12 files changed, 559 insertions(+), 110 deletions(-) diff --git a/docs/llms.txt b/docs/llms.txt index b2b38c7e50..1d8ad9e182 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -3493,9 +3493,9 @@ Requirements files can list a local wheel or source archive: ./packages/localpkg.tar.gz ``` -Cog supports `.whl`, `.zip`, `.tar.gz`, `.tgz`, `.tar.bz2`, and `.tar.xz` files. Paths are resolved relative to the requirements file and must stay inside the project directory. +Cog supports `.whl`, `.zip`, `.tar.gz`, `.tgz`, `.tar.bz2`, and `.tar.xz` files. Paths may contain spaces, are resolved relative to the requirements file, and must stay inside the project directory. -Only bare paths are supported. Local directories, direct references, and options, hashes, extras, or markers on a local artifact line are rejected. Use `build.sdk_version` or `COG_SDK_WHEEL` for Cog, and `COGLET_WHEEL` for Coglet. +Only bare paths are supported. Local directories, local direct references such as `name @ path`, and options, hashes, extras, or markers on a local artifact line are rejected. Remote direct references remain supported. Cog overrides any `cog` or `coglet` distribution installed by a local artifact; use `build.sdk_version` or `COG_SDK_WHEEL` for Cog, and `COGLET_WHEEL` for Coglet. To install Git-hosted Python packages, add `git` to the `system_packages` list, then use the `git+https://` syntax to specify the package name. For example: diff --git a/docs/yaml.md b/docs/yaml.md index 1a45b55b4a..ceebf35d95 100644 --- a/docs/yaml.md +++ b/docs/yaml.md @@ -69,9 +69,9 @@ Requirements files can list a local wheel or source archive: ./packages/localpkg.tar.gz ``` -Cog supports `.whl`, `.zip`, `.tar.gz`, `.tgz`, `.tar.bz2`, and `.tar.xz` files. Paths are resolved relative to the requirements file and must stay inside the project directory. +Cog supports `.whl`, `.zip`, `.tar.gz`, `.tgz`, `.tar.bz2`, and `.tar.xz` files. Paths may contain spaces, are resolved relative to the requirements file, and must stay inside the project directory. -Only bare paths are supported. Local directories, direct references, and options, hashes, extras, or markers on a local artifact line are rejected. Use `build.sdk_version` or `COG_SDK_WHEEL` for Cog, and `COGLET_WHEEL` for Coglet. +Only bare paths are supported. Local directories, local direct references such as `name @ path`, and options, hashes, extras, or markers on a local artifact line are rejected. Remote direct references remain supported. Cog overrides any `cog` or `coglet` distribution installed by a local artifact; use `build.sdk_version` or `COG_SDK_WHEEL` for Cog, and `COGLET_WHEEL` for Coglet. To install Git-hosted Python packages, add `git` to the `system_packages` list, then use the `git+https://` syntax to specify the package name. For example: diff --git a/integration-tests/tests/local_python_requirement_artifact.txtar b/integration-tests/tests/local_python_requirement_artifact.txtar index 1f9a2835b1..c64086bfbf 100644 --- a/integration-tests/tests/local_python_requirement_artifact.txtar +++ b/integration-tests/tests/local_python_requirement_artifact.txtar @@ -1,13 +1,18 @@ mkdir dist exec tar -czf dist/local_pkg-0.1.0.tar.gz setup.py local_pkg +exec python3 make_artifacts.py -# Ensure the import depends on the artifact being installed. +# Ensure the imports depend on the artifacts being installed. rm local_pkg +rm local_zip_pkg +rm local_wheel_pkg rm setup.py +rm zip_setup.py +rm make_artifacts.py cog build -t $TEST_IMAGE cog predict $TEST_IMAGE -stdout 'hello from local package 0.1.0' +stdout 'tar package 0.1.0; zip package 0.2.0; wheel package 0.3.0' -- cog.yaml -- build: @@ -17,6 +22,8 @@ predict: predict.py:Predictor -- requirements.txt -- ./dist/local_pkg-0.1.0.tar.gz +./dist/local zip pkg-0.2.0.zip +./dist/local_wheel_pkg-0.3.0-py3-none-any.whl -- setup.py -- from setuptools import setup @@ -24,16 +31,62 @@ from setuptools import setup setup(name="local-pkg", version="0.1.0", packages=["local_pkg"]) -- local_pkg/__init__.py -- -MESSAGE = "hello from local package" +MESSAGE = "tar package" + +-- zip_setup.py -- +from setuptools import setup + +setup(name="local-zip-pkg", version="0.2.0", packages=["local_zip_pkg"]) + +-- local_zip_pkg/__init__.py -- +MESSAGE = "zip package" + +-- local_wheel_pkg/__init__.py -- +MESSAGE = "wheel package" + +-- make_artifacts.py -- +from zipfile import ZIP_DEFLATED, ZipFile + + +with ZipFile("dist/local zip pkg-0.2.0.zip", "w", ZIP_DEFLATED) as archive: + archive.write("zip_setup.py", "setup.py") + archive.write("local_zip_pkg/__init__.py", "local_zip_pkg/__init__.py") + +dist_info = "local_wheel_pkg-0.3.0.dist-info" +with ZipFile( + "dist/local_wheel_pkg-0.3.0-py3-none-any.whl", "w", ZIP_DEFLATED +) as wheel: + wheel.write("local_wheel_pkg/__init__.py", "local_wheel_pkg/__init__.py") + wheel.writestr( + f"{dist_info}/METADATA", + "Metadata-Version: 2.1\nName: local-wheel-pkg\nVersion: 0.3.0\n", + ) + wheel.writestr( + f"{dist_info}/WHEEL", + "Wheel-Version: 1.0\nGenerator: cog-test\nRoot-Is-Purelib: true\nTag: py3-none-any\n", + ) + wheel.writestr( + f"{dist_info}/RECORD", + "local_wheel_pkg/__init__.py,,\n" + f"{dist_info}/METADATA,,\n" + f"{dist_info}/WHEEL,,\n" + f"{dist_info}/RECORD,,\n", + ) -- predict.py -- from importlib.metadata import version from cog import BasePredictor -from local_pkg import MESSAGE +from local_pkg import MESSAGE as TAR_MESSAGE +from local_wheel_pkg import MESSAGE as WHEEL_MESSAGE +from local_zip_pkg import MESSAGE as ZIP_MESSAGE class Predictor(BasePredictor): def predict(self) -> str: - return f"{MESSAGE} {version('local-pkg')}" + return ( + f"{TAR_MESSAGE} {version('local-pkg')}; " + f"{ZIP_MESSAGE} {version('local-zip-pkg')}; " + f"{WHEEL_MESSAGE} {version('local-wheel-pkg')}" + ) diff --git a/pkg/config/config.go b/pkg/config/config.go index 0e56ad514d..6871887968 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -66,9 +66,9 @@ type Build struct { // LocalPackageArtifact is a local requirement staged into the Docker build context. type LocalPackageArtifact struct { - Requirement string - SourcePath string - RelativePath string + Requirement string // Normalized requirements-file line. + SourcePath string // Canonical source path used for containment checks and reads. + RelativePath string // Project-relative staging path preserving the requirement filename. } type Concurrency struct { @@ -236,6 +236,9 @@ func (c *Config) cudaFromTF() (tfVersion string, tfCUDA string, tfCuDNN string, func (c *Config) pythonPackageVersion(name string) (version string, ok bool) { for _, pkg := range c.Build.pythonRequirementsContent { + if isLocalPackageArtifactRequirement(pkg) { + continue + } pkgName := requirements.PackageName(pkg) if pkgName == name { versions := requirements.Versions(pkg) @@ -270,6 +273,9 @@ func splitPythonVersion(version string) (major int, minor int, err error) { // Use this when building a Config struct directly (not from YAML). // For configs loaded from YAML, use Load() instead which handles validation and completion. func (c *Config) Complete(projectDir string) error { + c.Build.pythonRequirementsContent = nil + c.Build.localPackageArtifacts = nil + // Validate mutual exclusion of python_packages and python_requirements if len(c.Build.PythonPackages) > 0 && c.Build.PythonRequirements != "" { return fmt.Errorf("only one of python_packages or python_requirements can be set in your cog.yaml, not both") @@ -286,9 +292,6 @@ func (c *Config) Complete(projectDir string) error { return fmt.Errorf("failed to open python_requirements file: %w", err) } c.Build.pythonRequirementsContent = reqs - if err := c.loadLocalPackageArtifacts(projectDir, requirementsFilePath); err != nil { - return err - } } else if len(c.Build.PythonPackages) > 0 { // Backwards compatibility: if using deprecated python_packages, populate requirements content c.Build.pythonRequirementsContent = c.Build.PythonPackages @@ -309,6 +312,21 @@ func (c *Config) Complete(projectDir string) error { return nil } +// ResolveLocalPackageArtifacts validates local requirements for a generated Dockerfile +// and replaces the artifacts returned by LocalPackageArtifacts. +func (c *Config) ResolveLocalPackageArtifacts(projectDir string) error { + c.Build.localPackageArtifacts = nil + if c.Build.PythonRequirements == "" { + return nil + } + + requirementsFilePath := c.Build.PythonRequirements + if !filepath.IsAbs(requirementsFilePath) { + requirementsFilePath = filepath.Join(projectDir, requirementsFilePath) + } + return c.loadLocalPackageArtifacts(projectDir, requirementsFilePath) +} + func (c *Config) loadLocalPackageArtifacts(projectDir string, requirementsFilePath string) error { projectRoot, err := filepath.Abs(projectDir) if err != nil { @@ -339,6 +357,14 @@ func (c *Config) loadLocalPackageArtifacts(projectDir string, requirementsFilePa if err != nil { return fmt.Errorf("failed to resolve local Python package artifact %q: %w", artifactPath, err) } + resolvedParent, err := filepath.EvalSymlinks(filepath.Dir(absPath)) + if err != nil { + return fmt.Errorf("local Python package artifact %q not found: %w", artifactPath, err) + } + stagedPath := filepath.Join(resolvedParent, filepath.Base(absPath)) + if !pathWithin(projectRoot, stagedPath) { + return fmt.Errorf("local Python package artifact %q must be inside the project directory", artifactPath) + } canonicalPath, err := filepath.EvalSymlinks(absPath) if err != nil { return fmt.Errorf("local Python package artifact %q not found: %w", artifactPath, err) @@ -354,11 +380,15 @@ func (c *Config) loadLocalPackageArtifacts(projectDir string, requirementsFilePa return fmt.Errorf("local Python package artifact %q must be inside the project directory", artifactPath) } - relPath, err := filepath.Rel(projectRoot, canonicalPath) + relPath, err := filepath.Rel(projectRoot, stagedPath) if err != nil { return fmt.Errorf("failed to resolve local Python package artifact %q relative to project directory: %w", artifactPath, err) } - artifacts = append(artifacts, LocalPackageArtifact{requirement, canonicalPath, relPath}) + artifacts = append(artifacts, LocalPackageArtifact{ + Requirement: requirement, + SourcePath: canonicalPath, + RelativePath: relPath, + }) } c.Build.localPackageArtifacts = artifacts return nil @@ -399,7 +429,10 @@ func (c *Config) PythonRequirementsForArch(goos string, goarch string, includePa } } - packageName := requirements.PackageName(archPkg) + packageName := "" + if !isLocalPackageArtifactRequirement(archPkg) { + packageName = requirements.PackageName(archPkg) + } if packageName != "" { foundIdx := -1 for i, includePkg := range includePackageNames { @@ -434,9 +467,17 @@ func (c *Config) PythonRequirementsForArch(goos string, goarch string, includePa return strings.Join(lines, "\n"), nil } +func isLocalPackageArtifactRequirement(requirement string) bool { + _, ok, err := requirements.ParseLocalArtifactRequirement(requirement) + return err != nil || ok +} + // pythonPackageForArch takes a package==version line and // returns a package==version and index URL resolved to the correct GPU package for the given OS and architecture func (c *Config) pythonPackageForArch(pkg, goos, goarch string) (actualPackage string, findLinksList []string, extraIndexURLs []string, err error) { + if isLocalPackageArtifactRequirement(pkg) { + return pkg, []string{}, []string{}, nil + } name, version, findLinksList, extraIndexURLs, err := requirements.SplitPinnedPythonRequirement(pkg) if err != nil { // It's not pinned, so just return the line verbatim diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index 4c80f869f4..049c3cb2c6 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -199,23 +199,50 @@ flask>0.4 func TestPythonRequirementsLocalPackageArtifacts(t *testing.T) { tmpDir := t.TempDir() require.NoError(t, os.MkdirAll(path.Join(tmpDir, "requirements", "dist"), 0o755)) - wheelPath := path.Join(tmpDir, "requirements", "dist", "local_pkg-0.1.0-py3-none-any.whl") + require.NoError(t, os.Mkdir(path.Join(tmpDir, "objects"), 0o755)) + wheelPath := path.Join(tmpDir, "objects", "blob") + wheelLink := path.Join(tmpDir, "requirements", "dist", "local_pkg-0.1.0-py3-none-any.whl") require.NoError(t, os.WriteFile(wheelPath, []byte("wheel"), 0o644)) - require.NoError(t, os.Symlink(wheelPath, path.Join(tmpDir, "requirements", "dist", "latest.whl"))) - require.NoError(t, os.WriteFile(path.Join(tmpDir, "requirements", "requirements.txt"), []byte("./dist/latest.whl # vendored helper"), 0o644)) + require.NoError(t, os.Symlink(wheelPath, wheelLink)) + require.NoError(t, os.WriteFile(path.Join(tmpDir, "requirements", "requirements.txt"), []byte("./dist/local_pkg-0.1.0-py3-none-any.whl # vendored helper"), 0o644)) config := &Config{Build: &Build{PythonVersion: "3.10", PythonRequirements: "requirements/requirements.txt"}} require.NoError(t, config.Complete(tmpDir)) + require.NoError(t, config.ResolveLocalPackageArtifacts(tmpDir)) artifacts := config.LocalPackageArtifacts() canonicalWheelPath, err := filepath.EvalSymlinks(wheelPath) require.NoError(t, err) require.Len(t, artifacts, 1) - assert.Equal(t, "./dist/latest.whl", artifacts[0].Requirement) + assert.Equal(t, "./dist/local_pkg-0.1.0-py3-none-any.whl", artifacts[0].Requirement) assert.Equal(t, canonicalWheelPath, artifacts[0].SourcePath) assert.Equal(t, path.Join("requirements", "dist", "local_pkg-0.1.0-py3-none-any.whl"), artifacts[0].RelativePath) } +func TestPythonRequirementsDoesNotParseLocalArtifactAsPackage(t *testing.T) { + projectDir := t.TempDir() + requirement := "torch==custom.whl" + require.NoError(t, os.WriteFile(path.Join(projectDir, "requirements.txt"), []byte(requirement), 0o644)) + config := &Config{Build: &Build{PythonVersion: "3.10", PythonRequirements: "requirements.txt"}} + require.NoError(t, config.Complete(projectDir)) + + _, ok := config.TorchVersion() + assert.False(t, ok) + resolved, err := config.PythonRequirementsForArch("linux", "amd64", []string{"torch==2.0.0"}) + require.NoError(t, err) + assert.Equal(t, requirement+"\ntorch==2.0.0", resolved) +} + +func TestCompleteDefersInvalidLocalFrameworkRequirement(t *testing.T) { + projectDir := t.TempDir() + requirement := "torch @ file:./torch.whl" + require.NoError(t, os.WriteFile(path.Join(projectDir, "requirements.txt"), []byte(requirement), 0o644)) + config := &Config{Build: &Build{GPU: true, PythonVersion: "3.10", PythonRequirements: "requirements.txt"}} + + require.NoError(t, config.Complete(projectDir)) + require.ErrorContains(t, config.ResolveLocalPackageArtifacts(projectDir), "local file URL requirements are not supported") +} + func TestPythonRequirementsLocalPackageArtifactValidation(t *testing.T) { testCases := []struct { name string @@ -261,7 +288,27 @@ func TestPythonRequirementsLocalPackageArtifactValidation(t *testing.T) { } require.NoError(t, os.WriteFile(path.Join(projectDir, "requirements.txt"), []byte(line), 0o644)) config := &Config{Build: &Build{PythonVersion: "3.10", PythonRequirements: "requirements.txt"}} - require.ErrorContains(t, config.Complete(projectDir), tc.expectedErr) + require.NoError(t, config.Complete(projectDir)) + require.ErrorContains(t, config.ResolveLocalPackageArtifacts(projectDir), tc.expectedErr) + }) + } +} + +func TestCompleteDefersLocalPackageArtifactValidation(t *testing.T) { + testCases := []string{ + "./generated-later.whl", + "-r requirements-local.txt", + "name @ file:./package.whl", + } + + for _, requirement := range testCases { + t.Run(requirement, func(t *testing.T) { + projectDir := t.TempDir() + require.NoError(t, os.WriteFile(path.Join(projectDir, "requirements.txt"), []byte(requirement), 0o644)) + config := &Config{Build: &Build{PythonVersion: "3.10", PythonRequirements: "requirements.txt"}} + + require.NoError(t, config.Complete(projectDir)) + assert.Empty(t, config.LocalPackageArtifacts()) }) } } diff --git a/pkg/config/data/config_schema_v1.0.json b/pkg/config/data/config_schema_v1.0.json index 9259c49b03..6395c6e322 100644 --- a/pkg/config/data/config_schema_v1.0.json +++ b/pkg/config/data/config_schema_v1.0.json @@ -74,7 +74,7 @@ "python_requirements": { "$id": "#/properties/build/properties/python_requirements", "type": "string", - "description": "A pip requirements file specifying the Python packages to install. Local wheel and source archive paths in the requirements file are supported when they point to files inside the project." + "description": "A pip requirements file specifying the Python packages to install. Bare local wheel and source archive paths in the requirements file are supported when they point to files inside the project." }, "system_packages": { "$id": "#/properties/build/properties/system_packages", diff --git a/pkg/config/validate.go b/pkg/config/validate.go index cd16752af0..29e5d32cc9 100644 --- a/pkg/config/validate.go +++ b/pkg/config/validate.go @@ -395,6 +395,9 @@ func validateFrameworkCompatibility(cfg *configFile, reqs []string, result *Vali // findPackageVersion finds a package version in requirements. func findPackageVersion(reqs []string, name string) string { for _, req := range reqs { + if isLocalPackageArtifactRequirement(req) { + continue + } pkgName := requirements.PackageName(req) if pkgName == name { versions := requirements.Versions(req) diff --git a/pkg/config/validate_test.go b/pkg/config/validate_test.go index 316698ecf0..bc7f0f20d9 100644 --- a/pkg/config/validate_test.go +++ b/pkg/config/validate_test.go @@ -22,6 +22,20 @@ func TestValidateConfigFile(t *testing.T) { require.False(t, result.HasErrors(), "expected no errors, got: %v", result.Errors) } +func TestValidateConfigFileIgnoresLocalArtifactPackageNames(t *testing.T) { + cfg := &configFile{ + Build: &buildFile{ + GPU: new(true), + PythonVersion: new("3.10"), + PythonPackages: []string{"tensorflow==2.15.0 source.tar.gz"}, + CUDA: new("11.8"), + }, + } + + result := ValidateConfigFile(cfg) + require.False(t, result.HasErrors(), "expected no errors, got: %v", result.Errors) +} + func TestValidateConfigFileSuccess(t *testing.T) { cfg := &configFile{ Build: &buildFile{ diff --git a/pkg/dockerfile/standard_generator.go b/pkg/dockerfile/standard_generator.go index 1cd1283009..e1af7fe7c4 100644 --- a/pkg/dockerfile/standard_generator.go +++ b/pkg/dockerfile/standard_generator.go @@ -1,8 +1,11 @@ package dockerfile import ( + "bytes" "context" + "errors" "fmt" + "io" "os" "path" "path/filepath" @@ -14,7 +17,6 @@ import ( "github.com/replicate/cog/pkg/registry" "github.com/replicate/cog/pkg/requirements" "github.com/replicate/cog/pkg/util/console" - "github.com/replicate/cog/pkg/util/files" "github.com/replicate/cog/pkg/util/version" "github.com/replicate/cog/pkg/weightslegacy" "github.com/replicate/cog/pkg/wheels" @@ -50,7 +52,9 @@ type StandardGenerator struct { precompile bool // absolute path to the build cache dir (.cog/build/) - tmpDir string + tmpDir string + buildRootInfo os.FileInfo + buildRoot *os.Root fileWalker weightslegacy.FileWalker @@ -79,6 +83,16 @@ func NewStandardGenerator(config *config.Config, dir string, buildCacheDir strin if configFilename == "" { configFilename = "cog.yaml" } + if err := config.ResolveLocalPackageArtifacts(dir); err != nil { + return nil, err + } + buildRootInfo, err := os.Lstat(buildCacheDir) + if err != nil { + return nil, fmt.Errorf("failed to inspect build directory: %w", err) + } + if !buildRootInfo.IsDir() { + return nil, fmt.Errorf("build path %s must be a directory", buildCacheDir) + } return &StandardGenerator{ Config: config, @@ -89,6 +103,7 @@ func NewStandardGenerator(config *config.Config, dir string, buildCacheDir strin GOOS: "linux", GOARCH: "amd64", tmpDir: buildCacheDir, + buildRootInfo: buildRootInfo, fileWalker: filepath.Walk, useCudaBaseImage: true, useCogBaseImage: nil, @@ -154,6 +169,11 @@ func (g *StandardGenerator) uvPipInstallFlags(flags string) string { } func (g *StandardGenerator) GenerateInitialSteps(ctx context.Context) (string, error) { + if err := g.openBuildRoot(); err != nil { + return "", err + } + defer g.closeBuildRoot() + baseImage, err := g.BaseImage(ctx) if err != nil { return "", err @@ -239,6 +259,31 @@ func (g *StandardGenerator) GenerateInitialSteps(ctx context.Context) (string, e return joinStringsWithoutLineSpace(steps), nil } +func (g *StandardGenerator) openBuildRoot() error { + root, err := os.OpenRoot(g.tmpDir) + if err != nil { + return fmt.Errorf("failed to open build directory: %w", err) + } + info, err := root.Stat(".") + if err != nil { + _ = root.Close() + return fmt.Errorf("failed to inspect build directory: %w", err) + } + if !os.SameFile(g.buildRootInfo, info) { + _ = root.Close() + return fmt.Errorf("build directory changed during Dockerfile generation") + } + g.buildRoot = root + return nil +} + +func (g *StandardGenerator) closeBuildRoot() { + if g.buildRoot != nil { + _ = g.buildRoot.Close() + g.buildRoot = nil + } +} + func (g *StandardGenerator) GenerateModelBase(ctx context.Context) (string, error) { initialSteps, err := g.GenerateInitialSteps(ctx) if err != nil { @@ -905,14 +950,15 @@ func (g *StandardGenerator) pipInstalls() (string, error) { return "", err } - // Strip cog/coglet from user requirements — we always install them ourselves - // via installCog(). Leaving them in would cause pip to overwrite our version. - g.pythonRequirementsContents = g.filterManagedPackages(g.pythonRequirementsContents) var artifactCopyLine string g.pythonRequirementsContents, artifactCopyLine, err = g.stageLocalPackageArtifacts(g.pythonRequirementsContents) if err != nil { return "", err } + // Strip cog/coglet from user requirements — we always install them ourselves + // via installCog(). Local paths have already been rewritten and cannot be + // mistaken for managed package names. + g.pythonRequirementsContents = g.filterManagedPackages(g.pythonRequirementsContents) if strings.Trim(g.pythonRequirementsContents, "") == "" { return "", nil @@ -934,70 +980,146 @@ func (g *StandardGenerator) pipInstalls() (string, error) { CFlags, pipInstallLine, "ENV CFLAGS=", + g.resetManagedPackages(), }), "\n"), nil } +func (g *StandardGenerator) resetManagedPackages() string { + if !g.requiresCog || len(g.Config.LocalPackageArtifacts()) == 0 { + return "" + } + flags := "" + if g.needsBreakSystemPackages() { + flags = " " + uvBreakSystemPackages + } + return "RUN " + uvPip + " uninstall" + flags + " cog coglet" +} + func (g *StandardGenerator) stageLocalPackageArtifacts(reqContents string) (string, string, error) { artifacts := g.Config.LocalPackageArtifacts() if len(artifacts) == 0 { return reqContents, "", nil } + root, err := os.OpenRoot(g.Dir) + if err != nil { + return "", "", fmt.Errorf("failed to open project directory: %w", err) + } + defer root.Close() + + projectRoot, err := filepath.Abs(g.Dir) + if err != nil { + return "", "", fmt.Errorf("failed to resolve project directory: %w", err) + } + projectRoot, err = filepath.EvalSymlinks(projectRoot) + if err != nil { + return "", "", fmt.Errorf("failed to resolve project directory symlinks: %w", err) + } + openedRootInfo, err := root.Stat(".") + if err != nil { + return "", "", fmt.Errorf("failed to inspect project directory: %w", err) + } + resolvedRootInfo, err := os.Stat(projectRoot) + if err != nil { + return "", "", fmt.Errorf("failed to inspect resolved project directory: %w", err) + } + if !os.SameFile(openedRootInfo, resolvedRootInfo) { + return "", "", fmt.Errorf("project directory changed while staging local Python package artifacts") + } + if g.buildRoot == nil { + return "", "", fmt.Errorf("build directory is not open") + } + containerPaths := map[string]string{} staged := map[string]bool{} for _, artifact := range artifacts { - filename := filepath.Base(artifact.RelativePath) - switch { - case isVersionedArtifact(filename, "cog"): - return "", "", fmt.Errorf("local cog artifact %q is not supported; use build.sdk_version or %s", artifact.Requirement, wheels.CogSDKWheelEnvVar) - case isVersionedArtifact(filename, "coglet"): - return "", "", fmt.Errorf("local coglet artifact %q is not supported; use %s", artifact.Requirement, wheels.CogletWheelEnvVar) + sourcePath, err := filepath.Rel(projectRoot, artifact.SourcePath) + if err != nil || !pathStaysWithinRoot(sourcePath) { + return "", "", fmt.Errorf("local Python package artifact %q must be inside the project directory", artifact.Requirement) } - - if !staged[artifact.SourcePath] { - dst := filepath.Join(g.tmpDir, localArtifactsDir, artifact.RelativePath) - if err := files.Copy(artifact.SourcePath, dst); err != nil { + if !pathStaysWithinRoot(artifact.RelativePath) { + return "", "", fmt.Errorf("invalid staged path for local Python package artifact %q", artifact.Requirement) + } + if !staged[artifact.RelativePath] { + destinationPath := filepath.Join(localArtifactsDir, artifact.RelativePath) + if err := copyLocalPackageArtifact(root, sourcePath, g.buildRoot, destinationPath); err != nil { return "", "", fmt.Errorf("failed to stage local Python package artifact %s: %w", artifact.SourcePath, err) } - staged[artifact.SourcePath] = true + staged[artifact.RelativePath] = true } containerPaths[artifact.Requirement] = path.Join(localArtifactsContainerDir, filepath.ToSlash(artifact.RelativePath)) } lines := []string{} + matched := map[string]bool{} for line := range strings.SplitSeq(reqContents, "\n") { requirement := strings.TrimSpace(line) if replacement, ok := containerPaths[requirement]; ok { lines = append(lines, replacement) + matched[requirement] = true } else { lines = append(lines, line) } } + for requirement := range containerPaths { + if !matched[requirement] { + return "", "", fmt.Errorf("failed to rewrite local Python package artifact %q", requirement) + } + } return strings.Join(lines, "\n"), fmt.Sprintf("COPY --from=cog_build %s/ %s/", localArtifactsDir, localArtifactsContainerDir), nil } -func isVersionedArtifact(filename, name string) bool { - version, ok := strings.CutPrefix(strings.ToLower(filename), name+"-") - if !ok { - return false +func copyLocalPackageArtifact(sourceRoot *os.Root, sourcePath string, destinationRoot *os.Root, destinationPath string) error { + source, err := sourceRoot.Open(sourcePath) + if err != nil { + return err } - major, rest, ok := strings.Cut(version, ".") - if !ok || !isDigits(major) { - return false + defer source.Close() + info, err := source.Stat() + if err != nil { + return err } - minor, _, _ := strings.Cut(rest, ".") - return isDigits(minor) + if !info.Mode().IsRegular() { + return fmt.Errorf("source must be a regular file") + } + return writeBuildFile(destinationRoot, destinationPath, source) } -func isDigits(value string) bool { - if value == "" { - return false +func writeBuildFile(root *os.Root, destinationPath string, source io.Reader) error { + if !pathStaysWithinRoot(destinationPath) { + return fmt.Errorf("invalid build path %q", destinationPath) } - for _, char := range value { - if char < '0' || char > '9' { - return false - } + if err := root.MkdirAll(filepath.Dir(destinationPath), 0o755); err != nil { + return err } - return true + tmpPath := destinationPath + ".tmp" + if err := root.Remove(tmpPath); err != nil && !errors.Is(err, os.ErrNotExist) { + return err + } + destination, err := root.OpenFile(tmpPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o644) + if err != nil { + return err + } + if _, err := io.Copy(destination, source); err != nil { + _ = destination.Close() + _ = root.Remove(tmpPath) + return err + } + if err := destination.Close(); err != nil { + _ = root.Remove(tmpPath) + return err + } + if err := root.Rename(tmpPath, destinationPath); err != nil { + _ = root.Remove(tmpPath) + return err + } + return nil +} + +func pathStaysWithinRoot(path string) bool { + if path == "" || filepath.IsAbs(path) || path == ".." { + return false + } + return !strings.HasPrefix(path, ".."+string(filepath.Separator)) } func (g *StandardGenerator) runCommands() (string, error) { @@ -1055,11 +1177,10 @@ func (g *StandardGenerator) installCACert() (string, error) { // referenced via the "cog_build" named build context so the path is // relative to .cog/build/, not the project root. func (g *StandardGenerator) writeTemp(filename string, contents []byte) ([]string, string, error) { - path := filepath.Join(g.tmpDir, filename) - if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { - return []string{}, "", fmt.Errorf("Failed to write %s: %w", filename, err) + if g.buildRoot == nil { + return []string{}, "", fmt.Errorf("build directory is not open") } - if err := os.WriteFile(path, contents, 0o644); err != nil { + if err := writeBuildFile(g.buildRoot, filename, bytes.NewReader(contents)); err != nil { return []string{}, "", fmt.Errorf("Failed to write %s: %w", filename, err) } return []string{fmt.Sprintf("COPY --from=cog_build %s /tmp/%s", filename, filename)}, "/tmp/" + filename, nil diff --git a/pkg/dockerfile/standard_generator_test.go b/pkg/dockerfile/standard_generator_test.go index 2555dd8e0b..ca94fdb883 100644 --- a/pkg/dockerfile/standard_generator_test.go +++ b/pkg/dockerfile/standard_generator_test.go @@ -342,11 +342,18 @@ build: func TestPythonRequirementsLocalPackageArtifact(t *testing.T) { tmpDir := t.TempDir() + require.NoError(t, os.MkdirAll(path.Join(tmpDir, "requirements", "cog"), 0o755)) require.NoError(t, os.MkdirAll(path.Join(tmpDir, "requirements", "dist"), 0o755)) - artifactPath := path.Join(tmpDir, "requirements", "dist", "local_pkg-0.1.0-py3-none-any.whl") + require.NoError(t, os.Mkdir(path.Join(tmpDir, "objects"), 0o755)) + artifactPath := path.Join(tmpDir, "objects", "blob") + wheelName := "local_pkg-0.1.0-py3-none-any.whl" + wheelLink := path.Join(tmpDir, "requirements", "cog", wheelName) + archiveName := "local helper-0.1.0.tar.gz" require.NoError(t, os.WriteFile(artifactPath, []byte("wheel"), 0o644)) - require.NoError(t, os.Symlink(artifactPath, path.Join(tmpDir, "requirements", "dist", "latest.whl"))) - require.NoError(t, os.WriteFile(path.Join(tmpDir, "requirements", "requirements.txt"), []byte("./dist/latest.whl # vendored helper\ndist/latest.whl"), 0o644)) + require.NoError(t, os.Symlink(artifactPath, wheelLink)) + require.NoError(t, os.WriteFile(path.Join(tmpDir, "requirements", "dist", archiveName), []byte("archive"), 0o644)) + requirementsContents := "cog/" + wheelName + " # vendored helper\n./cog/" + wheelName + "\n./dist/" + archiveName + require.NoError(t, os.WriteFile(path.Join(tmpDir, "requirements", "requirements.txt"), []byte(requirementsContents), 0o644)) conf, err := config.FromYAML([]byte(` build: @@ -366,6 +373,7 @@ build: require.NoError(t, err) artifact := conf.LocalPackageArtifacts()[0] + assert.Equal(t, path.Join("requirements", "cog", wheelName), artifact.RelativePath) copyArtifacts := "COPY --from=cog_build local_package_artifacts/ /tmp/local_package_artifacts/" pipInstall := "uv run pip install --cache-dir /root/.cache/pip -r /tmp/requirements.txt" require.Contains(t, actual, copyArtifacts) @@ -379,48 +387,172 @@ build: requirements, err := os.ReadFile(path.Join(buildDir, "requirements.txt")) require.NoError(t, err) containerPath := path.Join(localArtifactsContainerDir, filepath.ToSlash(artifact.RelativePath)) - assert.Equal(t, containerPath+"\n"+containerPath, string(requirements)) + archivePath := path.Join(localArtifactsContainerDir, "requirements", "dist", archiveName) + assert.Equal(t, containerPath+"\n"+containerPath+"\n"+archivePath, string(requirements)) } -func TestLocalPackageArtifactRejectsManagedPackages(t *testing.T) { - testCases := []struct { - name string - filename string - expectedErr string - }{ - {name: "Cog", filename: "cog-0.1.0-py3-none-any.whl", expectedErr: wheels.CogSDKWheelEnvVar}, - {name: "Coglet", filename: "COGLET-1.0.0.tar.gz", expectedErr: wheels.CogletWheelEnvVar}, - {name: "Cog2FA", filename: "cog-2fa-1.0.tar.gz"}, - } +func TestLocalPackageArtifactsResetManagedPackages(t *testing.T) { + tmpDir := t.TempDir() + require.NoError(t, os.MkdirAll(path.Join(tmpDir, "dist"), 0o755)) + require.NoError(t, os.WriteFile(path.Join(tmpDir, "dist", "renamed-runtime.tar.gz"), []byte("artifact"), 0o644)) + require.NoError(t, os.WriteFile(path.Join(tmpDir, "requirements.txt"), []byte("./dist/renamed-runtime.tar.gz"), 0o644)) + + conf, err := config.FromYAML([]byte(` +build: + python_version: "3.12" + python_requirements: "requirements.txt" +`)) + require.NoError(t, err) + require.NoError(t, conf.Complete(tmpDir)) + gen, err := NewStandardGenerator(conf, tmpDir, t.TempDir(), "", dockertest.NewMockCommand(), registrytest.NewMockRegistryClient(), true) + require.NoError(t, err) + gen.SetUseCogBaseImage(false) + pypiWheels(gen) - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - tmpDir := t.TempDir() - require.NoError(t, os.MkdirAll(path.Join(tmpDir, "dist"), 0o755)) - require.NoError(t, os.WriteFile(path.Join(tmpDir, "dist", tc.filename), []byte("artifact"), 0o644)) - require.NoError(t, os.WriteFile(path.Join(tmpDir, "requirements.txt"), []byte("./dist/"+tc.filename), 0o644)) + _, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), "r8.im/replicate/cog-test") + require.NoError(t, err) + userInstall := "uv run pip install --cache-dir /root/.cache/pip -r /tmp/requirements.txt" + resetManaged := "RUN uv pip uninstall cog coglet" + managedInstall := "uv pip install --no-cache cog" + require.Contains(t, actual, userInstall) + require.Contains(t, actual, resetManaged) + require.Contains(t, actual, managedInstall) + assert.Less(t, strings.Index(actual, userInstall), strings.Index(actual, resetManaged)) + assert.Less(t, strings.Index(actual, resetManaged), strings.Index(actual, managedInstall)) +} - conf, err := config.FromYAML([]byte(` +func TestLocalPackageArtifactsResetManagedPackagesWithGPU(t *testing.T) { + projectDir := t.TempDir() + require.NoError(t, os.WriteFile(path.Join(projectDir, "package-0.1.0.tar.gz"), []byte("artifact"), 0o644)) + require.NoError(t, os.WriteFile(path.Join(projectDir, "requirements.txt"), []byte("./package-0.1.0.tar.gz"), 0o644)) + conf, err := config.FromYAML([]byte(` build: + gpu: true python_version: "3.12" python_requirements: "requirements.txt" `)) - require.NoError(t, err) - require.NoError(t, conf.Complete(tmpDir)) - gen, err := NewStandardGenerator(conf, tmpDir, t.TempDir(), "", dockertest.NewMockCommand(), registrytest.NewMockRegistryClient(), true) - require.NoError(t, err) - gen.SetUseCogBaseImage(false) - pypiWheels(gen) - - _, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), "r8.im/replicate/cog-test") - if tc.expectedErr != "" { - require.ErrorContains(t, err, tc.expectedErr) - return - } - require.NoError(t, err) - assert.Contains(t, actual, localArtifactsContainerDir+"/") - }) - } + require.NoError(t, err) + require.NoError(t, conf.Complete(projectDir)) + gen, err := NewStandardGenerator(conf, projectDir, t.TempDir(), "", dockertest.NewMockCommand(), registrytest.NewMockRegistryClient(), true) + require.NoError(t, err) + gen.SetUseCogBaseImage(false) + pypiWheels(gen) + + _, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), "r8.im/replicate/cog-test") + require.NoError(t, err) + require.Contains(t, actual, "RUN uv pip uninstall --break-system-packages cog coglet") +} + +func TestLocalPackageArtifactRejectsSymlinkSwapOutsideProject(t *testing.T) { + projectDir := t.TempDir() + artifactPath := path.Join(projectDir, "package-0.1.0.tar.gz") + require.NoError(t, os.WriteFile(artifactPath, []byte("inside"), 0o644)) + require.NoError(t, os.WriteFile(path.Join(projectDir, "requirements.txt"), []byte("./package-0.1.0.tar.gz"), 0o644)) + conf, err := config.FromYAML([]byte(` +build: + python_version: "3.12" + python_requirements: "requirements.txt" +`)) + require.NoError(t, err) + require.NoError(t, conf.Complete(projectDir)) + gen, err := NewStandardGenerator(conf, projectDir, t.TempDir(), "", dockertest.NewMockCommand(), registrytest.NewMockRegistryClient(), true) + require.NoError(t, err) + gen.SetUseCogBaseImage(false) + pypiWheels(gen) + + outsidePath := path.Join(t.TempDir(), "outside.tar.gz") + require.NoError(t, os.WriteFile(outsidePath, []byte("outside"), 0o644)) + require.NoError(t, os.Remove(artifactPath)) + require.NoError(t, os.Symlink(outsidePath, artifactPath)) + + _, _, _, err = gen.GenerateModelBaseWithSeparateWeights(t.Context(), "r8.im/replicate/cog-test") + require.ErrorContains(t, err, "failed to stage local Python package artifact") +} + +func TestLocalPackageArtifactDoesNotFollowStagingSymlink(t *testing.T) { + projectDir := t.TempDir() + artifactName := "package-0.1.0.tar.gz" + require.NoError(t, os.WriteFile(path.Join(projectDir, artifactName), []byte("artifact"), 0o644)) + require.NoError(t, os.WriteFile(path.Join(projectDir, "requirements.txt"), []byte("./"+artifactName), 0o644)) + conf, err := config.FromYAML([]byte(` +build: + python_version: "3.12" + python_requirements: "requirements.txt" +`)) + require.NoError(t, err) + require.NoError(t, conf.Complete(projectDir)) + buildDir := t.TempDir() + gen, err := NewStandardGenerator(conf, projectDir, buildDir, "", dockertest.NewMockCommand(), registrytest.NewMockRegistryClient(), true) + require.NoError(t, err) + gen.SetUseCogBaseImage(false) + pypiWheels(gen) + + victimPath := path.Join(t.TempDir(), "victim") + require.NoError(t, os.WriteFile(victimPath, []byte("untouched"), 0o644)) + stagedPath := path.Join(buildDir, localArtifactsDir, artifactName) + require.NoError(t, os.MkdirAll(filepath.Dir(stagedPath), 0o755)) + require.NoError(t, os.Symlink(victimPath, stagedPath)) + + _, _, _, err = gen.GenerateModelBaseWithSeparateWeights(t.Context(), "r8.im/replicate/cog-test") + require.NoError(t, err) + victim, err := os.ReadFile(victimPath) + require.NoError(t, err) + assert.Equal(t, []byte("untouched"), victim) + staged, err := os.ReadFile(stagedPath) + require.NoError(t, err) + assert.Equal(t, []byte("artifact"), staged) +} + +func TestPythonRequirementsDoesNotFollowBuildSymlink(t *testing.T) { + projectDir := t.TempDir() + require.NoError(t, os.WriteFile(path.Join(projectDir, "requirements.txt"), []byte("packaging==26.0"), 0o644)) + conf, err := config.FromYAML([]byte(` +build: + python_version: "3.12" + python_requirements: "requirements.txt" +`)) + require.NoError(t, err) + require.NoError(t, conf.Complete(projectDir)) + buildDir := t.TempDir() + gen, err := NewStandardGenerator(conf, projectDir, buildDir, "", dockertest.NewMockCommand(), registrytest.NewMockRegistryClient(), true) + require.NoError(t, err) + gen.SetUseCogBaseImage(false) + pypiWheels(gen) + + victimPath := path.Join(t.TempDir(), "victim") + require.NoError(t, os.WriteFile(victimPath, []byte("untouched"), 0o644)) + require.NoError(t, os.Symlink(victimPath, path.Join(buildDir, "requirements.txt"))) + + _, _, _, err = gen.GenerateModelBaseWithSeparateWeights(t.Context(), "r8.im/replicate/cog-test") + require.NoError(t, err) + victim, err := os.ReadFile(victimPath) + require.NoError(t, err) + assert.Equal(t, []byte("untouched"), victim) + requirements, err := os.ReadFile(path.Join(buildDir, "requirements.txt")) + require.NoError(t, err) + assert.Equal(t, []byte("packaging==26.0"), requirements) +} + +func TestGenerateRejectsReplacedBuildDirectory(t *testing.T) { + projectDir := t.TempDir() + conf, err := config.FromYAML([]byte(` +build: + python_version: "3.12" +`)) + require.NoError(t, err) + require.NoError(t, conf.Complete(projectDir)) + parentDir := t.TempDir() + buildDir := path.Join(parentDir, "build") + require.NoError(t, os.Mkdir(buildDir, 0o755)) + gen, err := NewStandardGenerator(conf, projectDir, buildDir, "", dockertest.NewMockCommand(), registrytest.NewMockRegistryClient(), true) + require.NoError(t, err) + gen.SetUseCogBaseImage(false) + pypiWheels(gen) + + require.NoError(t, os.Rename(buildDir, path.Join(parentDir, "old-build"))) + require.NoError(t, os.Mkdir(buildDir, 0o755)) + _, _, _, err = gen.GenerateModelBaseWithSeparateWeights(t.Context(), "r8.im/replicate/cog-test") + require.ErrorContains(t, err, "build directory changed") } // GPU builds on nvidia/cuda base images install Python via `uv python install` diff --git a/pkg/requirements/local_artifact.go b/pkg/requirements/local_artifact.go index d0cc797bf7..dd57900637 100644 --- a/pkg/requirements/local_artifact.go +++ b/pkg/requirements/local_artifact.go @@ -33,6 +33,9 @@ func ParseLocalArtifactRequirement(line string) (string, bool, error) { } return "", false, nil } + if hasInlineOption(line) { + return "", false, fmt.Errorf("local package artifact requirements do not support inline options or hashes: %s", line) + } if isRemoteRequirement(line) { return "", false, nil } @@ -43,13 +46,6 @@ func ParseLocalArtifactRequirement(line string) (string, bool, error) { return "", false, fmt.Errorf("extras are not supported on local package artifact requirements: %s", line) } - fields := strings.Fields(line) - if len(fields) > 1 { - if isLocalArtifact(fields[0]) { - return "", false, fmt.Errorf("local package artifact requirements do not support inline options or hashes: %s", line) - } - return "", false, nil - } if !isLocalPath(line) && !hasLocalArtifactSuffix(line) { return "", false, nil } @@ -61,11 +57,7 @@ func ParseLocalArtifactRequirement(line string) (string, bool, error) { func unsupportedLocalOption(line string) string { for _, option := range []string{"--find-links", "--requirement", "-f", "-r"} { - value, ok := strings.CutPrefix(line, option+" ") - if !ok && strings.HasPrefix(option, "--") { - value, ok = strings.CutPrefix(line, option+"=") - } - value = strings.TrimSpace(value) + value, ok := requirementOptionValue(line, option) if ok && value != "" && (isFileURL(value) || !isRemoteRequirement(value)) { return option } @@ -73,6 +65,44 @@ func unsupportedLocalOption(line string) string { return "" } +func requirementOptionValue(line string, option string) (string, bool) { + rest, ok := strings.CutPrefix(line, option) + if !ok || rest == "" { + return "", false + } + if strings.HasPrefix(option, "--") { + switch rest[0] { + case '=': + rest = rest[1:] + case ' ', '\t': + default: + return "", false + } + } + return strings.TrimSpace(rest), true +} + +func hasInlineOption(line string) bool { + for i := 0; i < len(line); i++ { + if line[i] != ' ' && line[i] != '\t' { + continue + } + rest := strings.TrimLeft(line[i:], " \t") + if !isLocalArtifact(strings.TrimSpace(line[:i])) { + continue + } + if strings.HasPrefix(rest, "--") { + return true + } + for _, option := range []string{"-f", "-r"} { + if _, ok := requirementOptionValue(rest, option); ok { + return true + } + } + } + return false +} + func isLocalArtifact(path string) bool { return !isRemoteRequirement(path) && (isLocalPath(path) || hasLocalArtifactSuffix(path)) } diff --git a/pkg/requirements/requirements_test.go b/pkg/requirements/requirements_test.go index 4a39172965..c504c312c7 100644 --- a/pkg/requirements/requirements_test.go +++ b/pkg/requirements/requirements_test.go @@ -37,6 +37,8 @@ func TestParseLocalArtifactRequirement(t *testing.T) { {name: "TarXz", line: "./dist/pkg-0.1.0.tar.xz", expected: "./dist/pkg-0.1.0.tar.xz", expectedOK: true}, {name: "AtInPath", line: "./dist/pkg@1.0.tar.gz", expected: "./dist/pkg@1.0.tar.gz", expectedOK: true}, {name: "AtInFilename", line: "pkg@1.0.tar.gz", expected: "pkg@1.0.tar.gz", expectedOK: true}, + {name: "PathWithSpaces", line: "./vendor/my helper-1.0.tar.gz", expected: "./vendor/my helper-1.0.tar.gz", expectedOK: true}, + {name: "BarePathWithSpaces", line: "vendor/my helper-1.0.tar.gz", expected: "vendor/my helper-1.0.tar.gz", expectedOK: true}, {name: "PackageWithExtras", line: "torch[all]==2.5.1"}, {name: "PackageWithMarker", line: `torch==2.5.1; python_version < "3.11"`}, {name: "URLWithHash", line: "https://user:pass@example.com/pkg.whl --hash=sha256:abc"}, @@ -51,8 +53,14 @@ func TestParseLocalArtifactRequirement(t *testing.T) { {name: "LocalFindLinks", line: "--find-links ./wheels", expectedErr: `local requirements option "--find-links" is not supported`}, {name: "FileURLFindLinks", line: "--find-links file:///wheels", expectedErr: `local requirements option "--find-links" is not supported`}, {name: "LocalRequirement", line: "-r requirements-local.txt", expectedErr: `local requirements option "-r" is not supported`}, + {name: "LocalRequirementNoSpace", line: "-rrequirements-local.txt", expectedErr: `local requirements option "-r" is not supported`}, + {name: "LocalFindLinksNoSpace", line: "-f./wheels", expectedErr: `local requirements option "-f" is not supported`}, + {name: "LocalRequirementTab", line: "-r\trequirements-local.txt", expectedErr: `local requirements option "-r" is not supported`}, + {name: "RemoteRequirementNoSpace", line: "-rhttps://example.com/requirements.txt"}, {name: "UppercaseFileURLRequirement", line: "-r FILE:///tmp/requirements.txt", expectedErr: `local requirements option "-r" is not supported`}, {name: "InlineHash", line: "./pkg.whl --hash=sha256:abc", expectedErr: "do not support inline options or hashes"}, + {name: "InlineRemoteFindLinks", line: "./pkg.whl --find-links=https://example.com", expectedErr: "do not support inline options or hashes"}, + {name: "InlineShortFindLinks", line: "./pkg.whl -f https://example.com", expectedErr: "do not support inline options or hashes"}, {name: "LocalDirectory", line: "./pkg", expectedErr: "is not a supported wheel or source archive"}, {name: "ArtifactWithMarker", line: `./dist/pkg.whl; python_version < "3.11"`, expectedErr: "environment markers are not supported"}, {name: "ArtifactWithExtras", line: "./dist/pkg.whl[extra]", expectedErr: "extras are not supported"},