Support pipeline pushes

pkg/cli/push.go

@@ -18,6 +18,8 @@ import (
 	"github.com/replicate/cog/pkg/util/console"
 )
 
+var pushPipeline bool
+
 func newPushCommand() *cobra.Command {
 	cmd := &cobra.Command{
 		Use: "push [IMAGE]",
@@ -40,6 +42,7 @@ func newPushCommand() *cobra.Command {
 	addFastFlag(cmd)
 	addLocalImage(cmd)
 	addConfigFlag(cmd)
+	addPipelineImage(cmd)
 
 	return cmd
 }
@@ -111,6 +114,7 @@ func push(cmd *cobra.Command, args []string) error {
 	err = docker.Push(ctx, imageName, buildFast, projectDir, dockerClient, docker.BuildInfo{
 		BuildTime: buildDuration,
 		BuildID:   buildID.String(),
+		Pipeline:  pushPipeline,
 	}, client)
 	if err != nil {
 		if strings.Contains(err.Error(), "404") {
@@ -140,3 +144,9 @@ func push(cmd *cobra.Command, args []string) error {
 
 	return nil
 }
+
+func addPipelineImage(cmd *cobra.Command) {
+	const pipeline = "x-pipeline"
+	cmd.Flags().BoolVar(&pushPipeline, pipeline, false, "Whether to use the experimental pipeline push feature")
+	_ = cmd.Flags().MarkHidden(pipeline)
+}

pkg/docker/pipeline_push.go

@@ -0,0 +1,78 @@
+package docker
+
+import (
+	"archive/tar"
+	"bytes"
+	"context"
+	"io"
+	"os"
+	"path/filepath"
+
+	"github.com/replicate/cog/pkg/dockerignore"
+	"github.com/replicate/cog/pkg/web"
+)
+
+func PipelinePush(ctx context.Context, image string, projectDir string, webClient *web.Client) error {
+	tarball, err := createTarball(projectDir)
+	if err != nil {
+		return err
+	}
+	return webClient.PostNewPipeline(ctx, image, tarball)
+}
+
+func createTarball(folder string) (*bytes.Buffer, error) {
+	var buf bytes.Buffer
+	tw := tar.NewWriter(&buf)
+
+	matcher, err := dockerignore.CreateMatcher(folder)
+	if err != nil {
+		return nil, err
+	}
+
+	err = dockerignore.Walk(folder, matcher, func(path string, info os.FileInfo, err error) error {
+		if err != nil {
+			return err
+		}
+
+		if info.IsDir() {
+			return nil
+		}
+
+		relPath, err := filepath.Rel(folder, path)
+		if err != nil {
+			return err
+		}
+
+		file, err := os.Open(path)
+		if err != nil {
+			return err
+		}
+		defer file.Close()
+
+		header, err := tar.FileInfoHeader(info, info.Name())
+		if err != nil {
+			return err
+		}
+		header.Name = relPath
+
+		err = tw.WriteHeader(header)
+		if err != nil {
+			return err
+		}
+
+		_, err = io.Copy(tw, file)
+		if err != nil {
+			return err
+		}
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	if err := tw.Close(); err != nil {
+		return nil, err
+	}
+
+	return &buf, nil
+}

pkg/docker/pipeline_push_test.go

@@ -0,0 +1,49 @@
+package docker
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/replicate/cog/pkg/docker/dockertest"
+	"github.com/replicate/cog/pkg/env"
+	cogHttp "github.com/replicate/cog/pkg/http"
+	"github.com/replicate/cog/pkg/web"
+)
+
+func TestPipelinePush(t *testing.T) {
+	// Setup mock http server
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		output := "{\"version\":\"user/test:53c740f17ce88a61c3da5b0c20e48fd48e2da537c3a1276dec63ab11fbad6bcb\"}"
+		w.WriteHeader(http.StatusCreated)
+		w.Write([]byte(output))
+	}))
+	defer server.Close()
+	url, err := url.Parse(server.URL)
+	require.NoError(t, err)
+	t.Setenv(env.SchemeEnvVarName, url.Scheme)
+	t.Setenv(web.WebHostEnvVarName, url.Host)
+
+	dir := t.TempDir()
+
+	// Create mock predict
+	predictPyPath := filepath.Join(dir, "predict.py")
+	handle, err := os.Create(predictPyPath)
+	require.NoError(t, err)
+	handle.WriteString("import cog")
+	dockertest.MockCogConfig = "{\"build\":{\"python_version\":\"3.12\",\"python_packages\":[\"torch==2.5.0\",\"beautifulsoup4==4.12.3\"],\"system_packages\":[\"git\"]},\"image\":\"test\",\"predict\":\"" + predictPyPath + ":Predictor\"}"
+
+	// Setup mock command
+	command := dockertest.NewMockCommand()
+	client, err := cogHttp.ProvideHTTPClient(t.Context(), command)
+	require.NoError(t, err)
+	webClient := web.NewClient(command, client)
+
+	err = PipelinePush(t.Context(), "r8.im/username/modelname", dir, webClient)
+	require.NoError(t, err)
+}

pkg/docker/push.go

@@ -14,11 +14,16 @@ import (
 type BuildInfo struct {
 	BuildTime time.Duration
 	BuildID   string
+	Pipeline  bool
 }
 
 func Push(ctx context.Context, image string, fast bool, projectDir string, command command.Command, buildInfo BuildInfo, client *http.Client) error {
 	webClient := web.NewClient(command, client)
 
+	if buildInfo.Pipeline {
+		return PipelinePush(ctx, image, projectDir, webClient)
+	}
+
 	if err := webClient.PostPushStart(ctx, buildInfo.BuildID, buildInfo.BuildTime); err != nil {
 		console.Warnf("Failed to send build timings to server: %v", err)
 	}

pkg/docker/push_test.go

@@ -152,3 +152,36 @@ func TestPushWithWeight(t *testing.T) {
 	err = Push(t.Context(), "r8.im/username/modelname", true, dir, command, BuildInfo{}, client)
 	require.NoError(t, err)
 }
+
+func TestPushPipeline(t *testing.T) {
+	// Setup mock http server
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		output := "{\"version\":\"user/test:53c740f17ce88a61c3da5b0c20e48fd48e2da537c3a1276dec63ab11fbad6bcb\"}"
+		w.WriteHeader(http.StatusCreated)
+		w.Write([]byte(output))
+	}))
+	defer server.Close()
+	url, err := url.Parse(server.URL)
+	require.NoError(t, err)
+	t.Setenv(env.SchemeEnvVarName, url.Scheme)
+	t.Setenv(web.WebHostEnvVarName, url.Host)
+
+	dir := t.TempDir()
+
+	// Create mock predict
+	predictPyPath := filepath.Join(dir, "predict.py")
+	handle, err := os.Create(predictPyPath)
+	require.NoError(t, err)
+	handle.WriteString("import cog")
+	dockertest.MockCogConfig = "{\"build\":{\"python_version\":\"3.12\",\"python_packages\":[\"torch==2.5.0\",\"beautifulsoup4==4.12.3\"],\"system_packages\":[\"git\"]},\"image\":\"test\",\"predict\":\"" + predictPyPath + ":Predictor\"}"
+
+	// Setup mock command
+	command := dockertest.NewMockCommand()
+	client, err := cogHttp.ProvideHTTPClient(t.Context(), command)
+	require.NoError(t, err)
+
+	err = Push(t.Context(), "r8.im/username/modelname", false, dir, command, BuildInfo{
+		Pipeline: true,
+	}, client)
+	require.NoError(t, err)
+}

pkg/dockerignore/dockerignore.go

@@ -10,8 +10,10 @@ import (
 	"github.com/replicate/cog/pkg/util/files"
 )
 
+const DockerIgnoreFilename = ".dockerignore"
+
 func CreateMatcher(dir string) (*ignore.GitIgnore, error) {
-	dockerIgnorePath := filepath.Join(dir, ".dockerignore")
+	dockerIgnorePath := filepath.Join(dir, DockerIgnoreFilename)
 	dockerIgnoreExists, err := files.Exists(dockerIgnorePath)
 	if err != nil {
 		return nil, err
@@ -27,6 +29,32 @@ func CreateMatcher(dir string) (*ignore.GitIgnore, error) {
 	return ignore.CompileIgnoreLines(patterns...), nil
 }
 
+func Walk(root string, ignoreMatcher *ignore.GitIgnore, fn filepath.WalkFunc) error {
+	return filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
+		if err != nil {
+			return err
+		}
+
+		// We ignore files ignored by .dockerignore
+		if ignoreMatcher != nil && ignoreMatcher.MatchesPath(path) {
+			if info.IsDir() {
+				return filepath.SkipDir
+			}
+			return nil
+		}
+
+		if info.IsDir() && info.Name() == ".cog" {
+			return filepath.SkipDir
+		}
+
+		if info.Name() == DockerIgnoreFilename {
+			return nil
+		}
+
+		return fn(path, info, err)
+	})
+}
+
 func readDockerIgnore(dockerIgnorePath string) ([]string, error) {
 	var patterns []string
 	file, err := os.Open(dockerIgnorePath)

pkg/dockerignore/dockerignore_test.go

@@ -0,0 +1,56 @@
+package dockerignore
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestWalk(t *testing.T) {
+	dir := t.TempDir()
+
+	predictOtherPyFilename := "predict_other.py"
+	predictOtherPyFilepath := filepath.Join(dir, predictOtherPyFilename)
+	predictOtherPyHandle, err := os.Create(predictOtherPyFilepath)
+	require.NoError(t, err)
+	predictOtherPyHandle.WriteString("import cog")
+
+	dockerIgnorePath := filepath.Join(dir, ".dockerignore")
+	dockerIgnoreHandle, err := os.Create(dockerIgnorePath)
+	require.NoError(t, err)
+	dockerIgnoreHandle.WriteString(predictOtherPyFilename)
+
+	predictPyFilename := "predict.py"
+	predictPyFilepath := filepath.Join(dir, predictPyFilename)
+	predictPyHandle, err := os.Create(predictPyFilepath)
+	require.NoError(t, err)
+	predictPyHandle.WriteString("import cog")
+
+	matcher, err := CreateMatcher(dir)
+	require.NoError(t, err)
+
+	foundFiles := []string{}
+	err = Walk(dir, matcher, func(path string, info os.FileInfo, err error) error {
+		if err != nil {
+			return err
+		}
+
+		if info.IsDir() {
+			return nil
+		}
+
+		relPath, err := filepath.Rel(dir, path)
+		if err != nil {
+			return err
+		}
+
+		foundFiles = append(foundFiles, relPath)
+
+		return nil
+	})
+	require.NoError(t, err)
+
+	require.Equal(t, []string{predictPyFilename}, foundFiles)
+}