Call the web when layers have been pushed

pkg/docker/apt.go

@@ -9,12 +9,14 @@ import (
 	"path/filepath"
 	"sort"
 	"strings"
+
+	"github.com/replicate/cog/pkg/docker/command"
 )
 
 const aptTarballPrefix = "apt."
 const aptTarballSuffix = ".tar.zst"
 
-func CreateAptTarball(tmpDir string, command Command, packages ...string) (string, error) {
+func CreateAptTarball(tmpDir string, dockerCommand command.Command, packages ...string) (string, error) {
 	if len(packages) > 0 {
 		sort.Strings(packages)
 		hash := sha256.New()
@@ -31,7 +33,7 @@ func CreateAptTarball(tmpDir string, command Command, packages ...string) (strin
 			}
 
 			// Create the apt tar file
-			_, err = command.CreateAptTarFile(tmpDir, aptTarFile, packages...)
+			_, err = dockerCommand.CreateAptTarFile(tmpDir, aptTarFile, packages...)
 			if err != nil {
 				return "", err
 			}

pkg/docker/apt_test.go

@@ -1,6 +1,7 @@
 package docker
 
 import (
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -15,3 +16,11 @@ func TestCreateAptTarball(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, "", tarball)
 }
+
+func TestCreateAptTarballWithPackages(t *testing.T) {
+	dir := t.TempDir()
+	command := dockertest.NewMockCommand()
+	tarball, err := CreateAptTarball(dir, command, []string{"git"}...)
+	require.NoError(t, err)
+	require.True(t, strings.HasPrefix(tarball, "apt."))
+}

pkg/docker/command.go → pkg/docker/command/command.go

@@ -1,14 +1,9 @@
-package docker
+package command
 
 type Command interface {
 	Push(string) error
-	LoadLoginToken(string) (string, error)
+	LoadUserInformation(string) (*UserInfo, error)
 	CreateTarFile(string, string, string, string) (string, error)
 	CreateAptTarFile(string, string, ...string) (string, error)
-}
-
-type CredentialHelperInput struct {
-	Username  string
-	Secret    string
-	ServerURL string
+	Inspect(string) (*Manifest, error)
 }

pkg/docker/command/manifest.go

@@ -0,0 +1,22 @@
+package command
+
+import "github.com/replicate/cog/pkg/global"
+
+type Config struct {
+	Labels map[string]string `json:"Labels"`
+	Env    []string          `json:"Env"`
+}
+
+type Manifest struct {
+	Config Config `json:"Config"`
+}
+
+const UvPythonInstallDirEnvVarName = "UV_PYTHON_INSTALL_DIR"
+const R8TorchVersionEnvVarName = "R8_TORCH_VERSION"
+const R8CudaVersionEnvVarName = "R8_CUDA_VERSION"
+const R8CudnnVersionEnvVarName = "R8_CUDNN_VERSION"
+const R8PythonVersionEnvVarName = "R8_PYTHON_VERSION"
+
+var CogConfigLabelKey = global.LabelNamespace + "config"
+var CogVersionLabelKey = global.LabelNamespace + "version"
+var CogOpenAPISchemaLabelKey = global.LabelNamespace + "openapi_schema"

pkg/docker/command/user_info.go

@@ -0,0 +1,6 @@
+package command
+
+type UserInfo struct {
+	Token    string
+	Username string
+}

pkg/docker/credential_helper_input.go

@@ -0,0 +1,7 @@
+package docker
+
+type CredentialHelperInput struct {
+	Username  string
+	Secret    string
+	ServerURL string
+}

pkg/docker/docker_command.go

@@ -1,6 +1,7 @@
 package docker
 
 import (
+	"bytes"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -11,7 +12,9 @@ import (
 
 	"github.com/docker/cli/cli/config"
 	"github.com/docker/cli/cli/config/configfile"
+	"github.com/docker/cli/cli/config/types"
 
+	"github.com/replicate/cog/pkg/docker/command"
 	"github.com/replicate/cog/pkg/util/console"
 )
 
@@ -22,16 +25,31 @@ func NewDockerCommand() *DockerCommand {
 }
 
 func (c *DockerCommand) Push(image string) error {
-	return c.exec("push", image)
+	_, err := c.exec("push", false, image)
+	return err
 }
 
-func (c *DockerCommand) LoadLoginToken(registryHost string) (string, error) {
+func (c *DockerCommand) LoadUserInformation(registryHost string) (*command.UserInfo, error) {
 	conf := config.LoadDefaultConfigFile(os.Stderr)
 	credsStore := conf.CredentialsStore
 	if credsStore == "" {
-		return loadAuthFromConfig(conf, registryHost)
+		authConf, err := loadAuthFromConfig(conf, registryHost)
+		if err != nil {
+			return nil, err
+		}
+		return &command.UserInfo{
+			Token:    authConf.Password,
+			Username: authConf.Username,
+		}, nil
 	}
-	return loadAuthFromCredentialsStore(credsStore, registryHost)
+	credsHelper, err := loadAuthFromCredentialsStore(credsStore, registryHost)
+	if err != nil {
+		return nil, err
+	}
+	return &command.UserInfo{
+		Token:    credsHelper.Secret,
+		Username: credsHelper.Username,
+	}, nil
 }
 
 func (c *DockerCommand) CreateTarFile(image string, tmpDir string, tarFile string, folder string) (string, error) {
@@ -45,7 +63,7 @@ func (c *DockerCommand) CreateTarFile(image string, tmpDir string, tarFile strin
 		"/",
 		folder,
 	}
-	err := c.exec("run", args...)
+	_, err := c.exec("run", false, args...)
 	if err != nil {
 		return "", err
 	}
@@ -66,25 +84,61 @@ func (c *DockerCommand) CreateAptTarFile(tmpDir string, aptTarFile string, packa
 		"/buildtmp/" + aptTarFile,
 	}
 	args = append(args, packages...)
-	return aptTarFile, c.exec("run", args...)
+	_, err := c.exec("run", false, args...)
+	if err != nil {
+		return "", err
+	}
+
+	return aptTarFile, nil
 }
 
-func (c *DockerCommand) exec(name string, args ...string) error {
+func (c *DockerCommand) Inspect(image string) (*command.Manifest, error) {
+	args := []string{
+		"inspect",
+		image,
+	}
+	manifestData, err := c.exec("image", false, args...)
+	if err != nil {
+		return nil, err
+	}
+
+	decoder := json.NewDecoder(bytes.NewReader(([]byte(manifestData))))
+	var manifest command.Manifest
+	err = decoder.Decode(&manifest)
+	if err != nil {
+		return nil, err
+	}
+
+	return &manifest, nil
+}
+
+func (c *DockerCommand) exec(name string, capture bool, args ...string) (string, error) {
 	cmdArgs := []string{name}
 	cmdArgs = append(cmdArgs, args...)
-	cmd := exec.Command("docker", cmdArgs...)
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
+	dockerCmd := DockerCommandFromEnvironment()
+	cmd := exec.Command(dockerCmd, cmdArgs...)
+	var out strings.Builder
+	if !capture {
+		cmd.Stdout = os.Stdout
+		cmd.Stderr = os.Stderr
+	} else {
+		cmd.Stdout = &out
+		cmd.Stderr = &out
+	}
 
 	console.Debug("$ " + strings.Join(cmd.Args, " "))
-	return cmd.Run()
+	err := cmd.Run()
+	if err != nil {
+		return "", err
+	}
+	return out.String(), nil
 }
 
-func loadAuthFromConfig(conf *configfile.ConfigFile, registryHost string) (string, error) {
-	return conf.AuthConfigs[registryHost].Password, nil
+func loadAuthFromConfig(conf *configfile.ConfigFile, registryHost string) (types.AuthConfig, error) {
+	return conf.AuthConfigs[registryHost], nil
 }
 
-func loadAuthFromCredentialsStore(credsStore string, registryHost string) (string, error) {
+func loadAuthFromCredentialsStore(credsStore string, registryHost string) (*CredentialHelperInput, error) {
 	var out strings.Builder
 	binary := DockerCredentialBinary(credsStore)
 	cmd := exec.Command(binary, "get")
@@ -93,34 +147,34 @@ func loadAuthFromCredentialsStore(credsStore string, registryHost string) (strin
 	cmd.Stderr = &out
 	stdin, err := cmd.StdinPipe()
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 	defer stdin.Close()
 	console.Debug("$ " + strings.Join(cmd.Args, " "))
 	err = cmd.Start()
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 	_, err = io.WriteString(stdin, registryHost)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 	err = stdin.Close()
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 	err = cmd.Wait()
 	if err != nil {
-		return "", fmt.Errorf("exec wait error: %w", err)
+		return nil, fmt.Errorf("exec wait error: %w", err)
 	}
 
 	var config CredentialHelperInput
 	err = json.Unmarshal([]byte(out.String()), &config)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 
-	return config.Secret, nil
+	return &config, nil
 }
 
 func DockerCredentialBinary(credsStore string) string {

pkg/docker/docker_command_test.go

@@ -0,0 +1,15 @@
+package docker
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestDockerPush(t *testing.T) {
+	t.Setenv(DockerCommandEnvVarName, "echo")
+
+	command := NewDockerCommand()
+	err := command.Push("test")
+	require.NoError(t, err)
+}

pkg/docker/dockertest/mock_command.go

@@ -3,9 +3,13 @@ package dockertest
 import (
 	"os"
 	"path/filepath"
+
+	"github.com/replicate/cog/pkg/docker/command"
 )
 
 var PushError error = nil
+var MockCogConfig string = "{\"build\":{\"python_version\":\"3.12\",\"python_packages\":[\"torch==2.5.0\",\"beautifulsoup4==4.12.3\"],\"system_packages\":[\"git\"]},\"image\":\"test\",\"predict\":\"predict.py:Predictor\"}"
+var MockOpenAPISchema string = "{}"
 
 type MockCommand struct{}
 
@@ -17,8 +21,12 @@ func (c *MockCommand) Push(image string) error {
 	return PushError
 }
 
-func (c *MockCommand) LoadLoginToken(registryHost string) (string, error) {
-	return "", nil
+func (c *MockCommand) LoadUserInformation(registryHost string) (*command.UserInfo, error) {
+	userInfo := command.UserInfo{
+		Token:    "",
+		Username: "",
+	}
+	return &userInfo, nil
 }
 
 func (c *MockCommand) CreateTarFile(image string, tmpDir string, tarFile string, folder string) (string, error) {
@@ -40,3 +48,23 @@ func (c *MockCommand) CreateAptTarFile(tmpDir string, aptTarFile string, package
 	}
 	return path, nil
 }
+
+func (c *MockCommand) Inspect(image string) (*command.Manifest, error) {
+	manifest := command.Manifest{
+		Config: command.Config{
+			Labels: map[string]string{
+				command.CogConfigLabelKey:        MockCogConfig,
+				command.CogOpenAPISchemaLabelKey: MockOpenAPISchema,
+				command.CogVersionLabelKey:       "0.11.3",
+			},
+			Env: []string{
+				command.UvPythonInstallDirEnvVarName + "=/tmp",
+				command.R8TorchVersionEnvVarName + "=2.5.0",
+				command.R8CudaVersionEnvVarName + "=2.4",
+				command.R8CudnnVersionEnvVarName + "=1.0",
+				command.R8PythonVersionEnvVarName + "=3.12",
+			},
+		},
+	}
+	return &manifest, nil
+}

pkg/docker/env.go

@@ -0,0 +1,13 @@
+package docker
+
+import "os"
+
+const DockerCommandEnvVarName = "R8_DOCKER_COMMAND"
+
+func DockerCommandFromEnvironment() string {
+	command := os.Getenv(DockerCommandEnvVarName)
+	if command == "" {
+		command = "docker"
+	}
+	return command
+}