Keep requirements.txt in sync for pipeline models (#2515)

* Always download and update requirements.txt

* lint

* Overwrite requirements.txt before validation

* skip test that expects validation errors since we're always overwriting the file

* lint fixes

* use CogBuildArtifactsFolder global

Author: michaeldwan

go.sum

@@ -175,8 +175,6 @@ github.com/dlclark/regexp2 v1.11.0 h1:G/nrcoOa7ZXlpoa/91N3X7mM3r8eIlMBBJZvsz/mxK
 github.com/dlclark/regexp2 v1.11.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
 github.com/dnephin/pflag v1.0.7 h1:oxONGlWxhmUct0YzKTgrpQv9AUA1wtPBn7zuSjJqptk=
 github.com/dnephin/pflag v1.0.7/go.mod h1:uxE91IoWURlOiTUIA8Mq5ZZkAv3dPUfZNaT80Zm7OQE=
-github.com/docker/cli v28.1.1+incompatible h1:eyUemzeI45DY7eDPuwUcmDyDj1pM98oD5MdSpiItp8k=
-github.com/docker/cli v28.1.1+incompatible/go.mod h1:JLrzqnKDaYBop7H2jaqPtU4hHvMKP+vjCwu2uszcLI8=
 github.com/docker/cli v28.3.0+incompatible h1:s+ttruVLhB5ayeuf2BciwDVxYdKi+RoUlxmwNHV3Vfo=
 github.com/docker/cli v28.3.0+incompatible/go.mod h1:JLrzqnKDaYBop7H2jaqPtU4hHvMKP+vjCwu2uszcLI8=
 github.com/docker/distribution v2.8.3+incompatible h1:AtKxIZ36LoNK51+Z6RpzLpddBirtxJnzDrHLEKxTAYk=

pkg/cli/init.go

@@ -5,12 +5,14 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"net/url"
 	"os"
 	"path"
 	"time"
 
 	"github.com/spf13/cobra"
 
+	"github.com/replicate/cog/pkg/env"
 	"github.com/replicate/cog/pkg/util/console"
 	"github.com/replicate/cog/pkg/util/files"
 )
@@ -119,8 +121,10 @@ func processTemplateFile(fs embed.FS, templateDir, filename, cwd string) error {
 
 	var content []byte
 
-	// Special handling for AGENTS.md - try to download from Replicate docs
-	if filename == "AGENTS.md" {
+	// Special handling for specific template files
+	switch {
+	case filename == "AGENTS.md":
+		// Try to download from Replicate docs
 		downloadedContent, err := downloadAgentsFile()
 		if err != nil {
 			console.Infof("Failed to download AGENTS.md: %v", err)
@@ -133,7 +137,21 @@ func processTemplateFile(fs embed.FS, templateDir, filename, cwd string) error {
 		} else {
 			content = downloadedContent
 		}
-	} else {
+	case filename == "requirements.txt" && pipelineTemplate:
+		// Special handling for requirements.txt in pipeline templates - download from runtime
+		downloadedContent, err := downloadPipelineRequirementsFile()
+		if err != nil {
+			console.Infof("Failed to download pipeline requirements.txt: %v", err)
+			console.Infof("Using template version instead...")
+			// Fall back to template version
+			content, err = fs.ReadFile(path.Join(templateDir, filename))
+			if err != nil {
+				return fmt.Errorf("Error reading template %s: %w", filename, err)
+			}
+		} else {
+			content = downloadedContent
+		}
+	default:
 		// Regular template file processing
 		content, err = fs.ReadFile(path.Join(templateDir, filename))
 		if err != nil {
@@ -174,6 +192,40 @@ func downloadAgentsFile() ([]byte, error) {
 	return content, nil
 }
 
+func downloadPipelineRequirementsFile() ([]byte, error) {
+	requirementsURL := pipelinesRuntimeRequirementsURL()
+
+	client := &http.Client{
+		Timeout: 10 * time.Second,
+	}
+
+	resp, err := client.Get(requirementsURL.String())
+	if err != nil {
+		return nil, fmt.Errorf("%w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("HTTP %d", resp.StatusCode)
+	}
+
+	content, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read response body: %w", err)
+	}
+
+	return content, nil
+}
+
+func pipelinesRuntimeRequirementsURL() url.URL {
+	baseURL := url.URL{
+		Scheme: env.SchemeFromEnvironment(),
+		Host:   env.PipelinesRuntimeHostFromEnvironment(),
+	}
+	baseURL.Path = "requirements.txt"
+	return baseURL
+}
+
 func addPipelineInit(cmd *cobra.Command) {
 	const pipeline = "x-pipeline"
 	cmd.Flags().BoolVar(&pipelineTemplate, pipeline, false, "Initialize a pipeline template")

pkg/docker/pipeline_push.go

@@ -8,27 +8,29 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/replicate/cog/pkg/api"
 	"github.com/replicate/cog/pkg/config"
 	"github.com/replicate/cog/pkg/dockerignore"
+	"github.com/replicate/cog/pkg/global"
 	"github.com/replicate/cog/pkg/procedure"
 )
 
 func PipelinePush(ctx context.Context, image string, projectDir string, apiClient *api.Client, client *http.Client, cfg *config.Config) error {
-	err := procedure.Validate(projectDir, client, cfg, false)
+	err := procedure.Validate(projectDir, client, cfg, true)
 	if err != nil {
 		return err
 	}
 
-	tarball, err := createTarball(projectDir)
+	tarball, err := createTarball(projectDir, cfg)
 	if err != nil {
 		return err
 	}
 	return apiClient.PostNewPipeline(ctx, image, tarball)
 }
 
-func createTarball(folder string) (*bytes.Buffer, error) {
+func createTarball(folder string, cfg *config.Config) (*bytes.Buffer, error) {
 	var buf bytes.Buffer
 	tw := tar.NewWriter(&buf)
 
@@ -37,6 +39,25 @@ func createTarball(folder string) (*bytes.Buffer, error) {
 		return nil, err
 	}
 
+	// Track if we need to add downloaded requirements to the tarball
+	var downloadedRequirementsPath string
+	var downloadedRequirementsContent []byte
+
+	// If config points to downloaded requirements (outside project directory),
+	// we need to include them in the tarball as requirements.txt
+	if cfg.Build.PythonRequirements != "" {
+		reqPath := cfg.RequirementsFile(folder)
+		if !strings.HasPrefix(reqPath, folder) || strings.Contains(reqPath, global.CogBuildArtifactsFolder) {
+			// This is a downloaded requirements file, read its content
+			content, err := os.ReadFile(reqPath)
+			if err != nil {
+				return nil, err
+			}
+			downloadedRequirementsPath = "requirements.txt"
+			downloadedRequirementsContent = content
+		}
+	}
+
 	err = dockerignore.Walk(folder, matcher, func(path string, info os.FileInfo, err error) error {
 		if err != nil {
 			return err
@@ -51,6 +72,12 @@ func createTarball(folder string) (*bytes.Buffer, error) {
 			return err
 		}
 
+		// If this is the local requirements.txt and we have downloaded requirements,
+		// skip the local one (we'll add the downloaded version instead)
+		if downloadedRequirementsPath != "" && relPath == "requirements.txt" {
+			return nil
+		}
+
 		file, err := os.Open(path)
 		if err != nil {
 			return err
@@ -78,6 +105,25 @@ func createTarball(folder string) (*bytes.Buffer, error) {
 		return nil, err
 	}
 
+	// Add downloaded requirements as requirements.txt if we have them
+	if downloadedRequirementsPath != "" {
+		header := &tar.Header{
+			Name: downloadedRequirementsPath,
+			Mode: 0o644,
+			Size: int64(len(downloadedRequirementsContent)),
+		}
+
+		err = tw.WriteHeader(header)
+		if err != nil {
+			return nil, err
+		}
+
+		_, err = tw.Write(downloadedRequirementsContent)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	if err := tw.Close(); err != nil {
 		return nil, err
 	}

pkg/docker/pipeline_push_test.go

@@ -90,6 +90,7 @@ func TestPipelinePush(t *testing.T) {
 }
 
 func TestPipelinePushFailWithExtraRequirements(t *testing.T) {
+	t.Skip("Skipping for now, requirements.txt is always overwritten, and hopefully we replace that with support for custom requirements, if not this test comes back")
 	// Setup mock web server for cog.replicate.com (token exchange)
 	webServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		switch r.URL.Path {

pkg/procedure/validate.go

@@ -2,6 +2,7 @@ package procedure
 
 import (
 	"errors"
+	"fmt"
 	"io"
 	"net/http"
 	"net/url"
@@ -118,6 +119,15 @@ func validateRequirements(projectDir string, client *http.Client, cfg *config.Co
 		return err
 	}
 
+	// Update local requirements.txt to match production before validation if filling
+	if fill {
+		err := updateLocalRequirementsFile(projectDir, requirementsFilePath)
+		if err != nil {
+			// Log warning but don't fail the build - the downloaded requirements will still be used
+			console.Warn(fmt.Sprintf("Failed to update local requirements.txt: %v", err))
+		}
+	}
+
 	if cfg.Build.PythonRequirements != "" {
 		pipelineRequirements, err := requirements.ReadRequirements(filepath.Join(projectDir, requirementsFilePath))
 		if err != nil {
@@ -243,6 +253,27 @@ func downloadRequirements(projectDir string, client *http.Client) (string, error
 	return requirementsFilePath, nil
 }
 
+// updateLocalRequirementsFile copies the downloaded requirements to the local requirements.txt file
+// This keeps the local file in sync with what's actually available in the runtime
+func updateLocalRequirementsFile(projectDir, downloadedRequirementsPath string) error {
+	// Read the downloaded requirements
+	downloadedPath := filepath.Join(projectDir, downloadedRequirementsPath)
+	downloadedContent, err := os.ReadFile(downloadedPath)
+	if err != nil {
+		return fmt.Errorf("failed to read downloaded requirements: %w", err)
+	}
+
+	// Write to local requirements.txt
+	localRequirementsPath := filepath.Join(projectDir, "requirements.txt")
+	err = os.WriteFile(localRequirementsPath, downloadedContent, 0o644)
+	if err != nil {
+		return fmt.Errorf("failed to write local requirements.txt: %w", err)
+	}
+
+	console.Infof("Updated local requirements.txt with runtime requirements")
+	return nil
+}
+
 func requirementsURL() url.URL {
 	requirementsURL := pipelinesRuntimeBaseURL()
 	requirementsURL.Path = "requirements.txt"

test-integration/test_integration/fixtures/pipeline-requirements-project/.gitignore

@@ -0,0 +1,4 @@
+# Ignore dynamically generated requirements.txt file during tests
+# This file is created by test_predict_pipeline_downloaded_requirements 
+# to simulate out-of-sync scenario and gets cleaned up automatically
+requirements.txt

test-integration/test_integration/fixtures/pipeline-requirements-project/cog.yaml

@@ -0,0 +1,4 @@
+build:
+  python_version: "3.13"
+  python_requirements: "requirements.txt"
+predict: "predict.py:Predictor"

test-integration/test_integration/fixtures/pipeline-requirements-project/predict.py

@@ -0,0 +1,21 @@
+from cog import BasePredictor
+import importlib.metadata
+
+class Predictor(BasePredictor):
+    def predict(self) -> str:
+        """Test function that verifies downloaded requirements are available"""
+        
+        try:
+            # Get all installed packages and their versions
+            packages = []
+            for dist in importlib.metadata.distributions():
+                packages.append(f"{dist.metadata['name']}=={dist.version}")
+            
+            # Sort for consistent output
+            packages.sort()
+            
+            # Create output with prompt and all packages listed
+            return '\n'.join(packages)
+            
+        except Exception as e:
+            return f"ERROR: Unexpected error - {str(e)}"