Add fast pusher for fast builds (#2114)

* Add fast pusher for fast builds

* Separates push into standard and fast
* Generates the necessary tar files to push
* Calls the backend to create a file to push

* Add Login Token to HTTP POST request

* Remove redundant tar files

* When building/pushing make sure we remove
redundant tar files we have created previously.

* Add tests to fast push

* Fix injecting docker command into push

* Add LoadLoginToken to docker command object

* Add a create src tarfile error

* Mock CreateTarFile

* Format generate apt tarball errors

* Inject the docker command object

* Allow mocking of docker exec commands

* Use command to create apt tar file

* Fix bad variable assignment

* Remove aptTarPath assignment

* Add a comment on the use of image on CreateTarFile

* Change the config arg to a packages array arg

* Use snakeCase for apt tarball prefix/suffix consts

* Use snakeCase in fast push consts

* Use error groups to wait for multiple uploads

* Remove ineffectual assignment

* Do not use url as a variable name

* Remove size from the backend call

* Set test environment with t.Setenv

* Put MockCommand in dockertest package

* Fix pip_freeze integration test

Author: 8W9aG

pkg/cli/baseimage.go

@@ -148,5 +148,6 @@ func baseImageGeneratorFromFlags() (*dockerfile.BaseImageGenerator, error) {
 		baseImageCUDAVersion,
 		baseImagePythonVersion,
 		baseImageTorchVersion,
+		docker.NewDockerCommand(),
 	)
 }

pkg/cli/debug.go

@@ -6,6 +6,7 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/replicate/cog/pkg/config"
+	"github.com/replicate/cog/pkg/docker"
 	"github.com/replicate/cog/pkg/dockerfile"
 	"github.com/replicate/cog/pkg/global"
 	"github.com/replicate/cog/pkg/util/console"
@@ -37,7 +38,8 @@ func cmdDockerfile(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	generator, err := dockerfile.NewGenerator(cfg, projectDir, false)
+	command := docker.NewDockerCommand()
+	generator, err := dockerfile.NewGenerator(cfg, projectDir, false, command)
 	if err != nil {
 		return fmt.Errorf("Error creating Dockerfile generator: %w", err)
 	}

pkg/cli/push.go

@@ -60,7 +60,6 @@ func push(cmd *cobra.Command, args []string) error {
 	}
 
 	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile, DetermineUseCogBaseImage(cmd), buildStrip, buildPrecompile, buildFast); err != nil {
-
 		return err
 	}
 
@@ -69,7 +68,8 @@ func push(cmd *cobra.Command, args []string) error {
 		console.Info("Fast push enabled.")
 	}
 
-	err = docker.Push(imageName)
+	command := docker.NewDockerCommand()
+	err = docker.Push(imageName, buildFast, projectDir, command)
 	if err != nil {
 		if strings.Contains(err.Error(), "NAME_UNKNOWN") {
 			return fmt.Errorf("Unable to find existing Replicate model for %s. "+

pkg/config/config_test.go

@@ -670,41 +670,6 @@ func TestBlankBuild(t *testing.T) {
 	require.Equal(t, false, config.Build.GPU)
 }
 
-func TestSplitPinnedPythonRequirement(t *testing.T) {
-	testCases := []struct {
-		input                  string
-		expectedName           string
-		expectedVersion        string
-		expectedFindLinks      []string
-		expectedExtraIndexURLs []string
-		expectedError          bool
-	}{
-		{"package1==1.0.0", "package1", "1.0.0", nil, nil, false},
-		{"package1==1.0.0+alpha", "package1", "1.0.0+alpha", nil, nil, false},
-		{"--find-links=link1 --find-links=link2 package3==3.0.0", "package3", "3.0.0", []string{"link1", "link2"}, nil, false},
-		{"package4==4.0.0 --extra-index-url=url1 --extra-index-url=url2", "package4", "4.0.0", nil, []string{"url1", "url2"}, false},
-		{"-f link1 --find-links=link2 package5==5.0.0 --extra-index-url=url1 --extra-index-url=url2", "package5", "5.0.0", []string{"link1", "link2"}, []string{"url1", "url2"}, false},
-		{"package6 --find-links=link1 --find-links=link2 --extra-index-url=url1 --extra-index-url=url2", "", "", nil, nil, true},
-		{"invalid package", "", "", nil, nil, true},
-		{"package8==", "", "", nil, nil, true},
-		{"==8.0.0", "", "", nil, nil, true},
-	}
-
-	for _, tc := range testCases {
-		name, version, findLinks, extraIndexURLs, err := SplitPinnedPythonRequirement(tc.input)
-
-		if tc.expectedError {
-			require.Error(t, err)
-		} else {
-			require.NoError(t, err)
-			require.Equal(t, tc.expectedName, name, "input: "+tc.input)
-			require.Equal(t, tc.expectedVersion, version, "input: "+tc.input)
-			require.Equal(t, tc.expectedFindLinks, findLinks, "input: "+tc.input)
-			require.Equal(t, tc.expectedExtraIndexURLs, extraIndexURLs, "input: "+tc.input)
-		}
-	}
-}
-
 func TestPythonRequirementsForArchWithAddedPackage(t *testing.T) {
 	config := &Config{
 		Build: &Build{

pkg/config/requirements.go

@@ -1,85 +1,10 @@
 package config
 
 import (
-	"bufio"
-	"errors"
 	"fmt"
-	"os"
-	"path/filepath"
 	"regexp"
-	"sort"
 )
 
-func GenerateRequirements(tmpDir string, config *Config) (string, error) {
-	// Deduplicate packages between the requirements.txt and the python packages directive.
-	packageNames := make(map[string]string)
-
-	// Read the python packages configuration.
-	for _, requirement := range config.Build.PythonPackages {
-		packageName, err := PackageName(requirement)
-		if err != nil {
-			return "", err
-		}
-		packageNames[packageName] = requirement
-	}
-
-	// Read the python requirements.
-	if config.Build.PythonRequirements != "" {
-		fh, err := os.Open(config.Build.PythonRequirements)
-		if err != nil {
-			return "", err
-		}
-		scanner := bufio.NewScanner(fh)
-		for scanner.Scan() {
-			requirement := scanner.Text()
-			packageName, err := PackageName(requirement)
-			if err != nil {
-				return "", err
-			}
-			packageNames[packageName] = requirement
-		}
-	}
-
-	// If we don't have any packages skip further processing
-	if len(packageNames) == 0 {
-		return "", nil
-	}
-
-	// Sort the package names by alphabetical order.
-	keys := make([]string, 0, len(packageNames))
-	for k := range packageNames {
-		keys = append(keys, k)
-	}
-	sort.Strings(keys)
-
-	// Render the expected contents
-	requirementsContent := ""
-	for _, k := range keys {
-		requirementsContent += packageNames[k] + "\n"
-	}
-
-	// Check against the old requirements contents
-	requirementsFile := filepath.Join(tmpDir, "requirements.txt")
-	_, err := os.Stat(requirementsFile)
-	if !errors.Is(err, os.ErrNotExist) {
-		bytes, err := os.ReadFile(requirementsFile)
-		if err != nil {
-			return "", err
-		}
-		oldRequirementsContents := string(bytes)
-		if oldRequirementsContents == requirementsFile {
-			return requirementsFile, nil
-		}
-	}
-
-	// Write out a new requirements file
-	err = os.WriteFile(requirementsFile, []byte(requirementsContent), 0o644)
-	if err != nil {
-		return "", err
-	}
-	return requirementsFile, nil
-}
-
 // SplitPinnedPythonRequirement returns the name, version, findLinks, and extraIndexURLs from a requirements.txt line
 // in the form name==version [--find-links=<findLink>] [-f <findLink>] [--extra-index-url=<extraIndexURL>]
 func SplitPinnedPythonRequirement(requirement string) (name string, version string, findLinks []string, extraIndexURLs []string, err error) {

pkg/config/requirements_test.go

@@ -1,21 +1,42 @@
 package config
 
 import (
-	"path/filepath"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 )
 
-func TestGenerateRequirements(t *testing.T) {
-	tmpDir := t.TempDir()
-	build := Build{
-		PythonPackages: []string{"torch==2.5.1"},
+func TestSplitPinnedPythonRequirement(t *testing.T) {
+	testCases := []struct {
+		input                  string
+		expectedName           string
+		expectedVersion        string
+		expectedFindLinks      []string
+		expectedExtraIndexURLs []string
+		expectedError          bool
+	}{
+		{"package1==1.0.0", "package1", "1.0.0", nil, nil, false},
+		{"package1==1.0.0+alpha", "package1", "1.0.0+alpha", nil, nil, false},
+		{"--find-links=link1 --find-links=link2 package3==3.0.0", "package3", "3.0.0", []string{"link1", "link2"}, nil, false},
+		{"package4==4.0.0 --extra-index-url=url1 --extra-index-url=url2", "package4", "4.0.0", nil, []string{"url1", "url2"}, false},
+		{"-f link1 --find-links=link2 package5==5.0.0 --extra-index-url=url1 --extra-index-url=url2", "package5", "5.0.0", []string{"link1", "link2"}, []string{"url1", "url2"}, false},
+		{"package6 --find-links=link1 --find-links=link2 --extra-index-url=url1 --extra-index-url=url2", "", "", nil, nil, true},
+		{"invalid package", "", "", nil, nil, true},
+		{"package8==", "", "", nil, nil, true},
+		{"==8.0.0", "", "", nil, nil, true},
 	}
-	config := Config{
-		Build: &build,
+
+	for _, tc := range testCases {
+		name, version, findLinks, extraIndexURLs, err := SplitPinnedPythonRequirement(tc.input)
+
+		if tc.expectedError {
+			require.Error(t, err)
+		} else {
+			require.NoError(t, err)
+			require.Equal(t, tc.expectedName, name, "input: "+tc.input)
+			require.Equal(t, tc.expectedVersion, version, "input: "+tc.input)
+			require.Equal(t, tc.expectedFindLinks, findLinks, "input: "+tc.input)
+			require.Equal(t, tc.expectedExtraIndexURLs, extraIndexURLs, "input: "+tc.input)
+		}
 	}
-	requirementsFile, err := GenerateRequirements(tmpDir, &config)
-	require.NoError(t, err)
-	require.Equal(t, filepath.Join(tmpDir, "requirements.txt"), requirementsFile)
 }

pkg/docker/apt.go

@@ -0,0 +1,78 @@
+package docker
+
+import (
+	"crypto/sha256"
+	"encoding/hex"
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+	"sort"
+	"strings"
+)
+
+const aptTarballPrefix = "apt."
+const aptTarballSuffix = ".tar.zst"
+
+func CreateAptTarball(tmpDir string, command Command, packages ...string) (string, error) {
+	if len(packages) > 0 {
+		sort.Strings(packages)
+		hash := sha256.New()
+		hash.Write([]byte(strings.Join(packages, " ")))
+		hexHash := hex.EncodeToString(hash.Sum(nil))
+		aptTarFile := aptTarballPrefix + hexHash + aptTarballSuffix
+		aptTarPath := filepath.Join(tmpDir, aptTarFile)
+
+		if _, err := os.Stat(aptTarPath); errors.Is(err, os.ErrNotExist) {
+			// Remove previous apt tar files.
+			err = removeAptTarballs(tmpDir)
+			if err != nil {
+				return "", err
+			}
+
+			// Create the apt tar file
+			_, err = command.CreateAptTarFile(tmpDir, aptTarFile, packages...)
+			if err != nil {
+				return "", err
+			}
+		}
+
+		return aptTarFile, nil
+	}
+	return "", nil
+}
+
+func CurrentAptTarball(tmpDir string) (string, error) {
+	files, err := os.ReadDir(tmpDir)
+	if err != nil {
+		return "", fmt.Errorf("os read dir error: %w", err)
+	}
+
+	for _, file := range files {
+		fileName := file.Name()
+		if strings.HasPrefix(fileName, aptTarballPrefix) && strings.HasSuffix(fileName, aptTarballSuffix) {
+			return filepath.Join(tmpDir, fileName), nil
+		}
+	}
+
+	return "", nil
+}
+
+func removeAptTarballs(tmpDir string) error {
+	files, err := os.ReadDir(tmpDir)
+	if err != nil {
+		return err
+	}
+
+	for _, file := range files {
+		fileName := file.Name()
+		if strings.HasPrefix(fileName, aptTarballPrefix) && strings.HasSuffix(fileName, aptTarballSuffix) {
+			err = os.Remove(filepath.Join(tmpDir, fileName))
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}

pkg/docker/apt_test.go

@@ -0,0 +1,17 @@
+package docker
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/replicate/cog/pkg/docker/dockertest"
+)
+
+func TestCreateAptTarball(t *testing.T) {
+	dir := t.TempDir()
+	command := dockertest.NewMockCommand()
+	tarball, err := CreateAptTarball(dir, command, []string{}...)
+	require.NoError(t, err)
+	require.Equal(t, "", tarball)
+}

pkg/docker/build.go

@@ -8,7 +8,6 @@ import (
 	"strings"
 
 	"github.com/replicate/cog/pkg/config"
-	"github.com/replicate/cog/pkg/dockerfile"
 
 	"github.com/replicate/cog/pkg/util"
 	"github.com/replicate/cog/pkg/util/console"
@@ -17,7 +16,7 @@ import (
 func Build(dir, dockerfileContents, imageName string, secrets []string, noCache bool, progressOutput string, epoch int64) error {
 	var args []string
 
-	userCache, err := dockerfile.UserCache()
+	userCache, err := UserCache()
 	if err != nil {
 		return err
 	}

pkg/docker/command.go

@@ -0,0 +1,14 @@
+package docker
+
+type Command interface {
+	Push(string) error
+	LoadLoginToken(string) (string, error)
+	CreateTarFile(string, string, string, string) (string, error)
+	CreateAptTarFile(string, string, ...string) (string, error)
+}
+
+type CredentialHelperInput struct {
+	Username  string
+	Secret    string
+	ServerURL string
+}