Replace docker.Xxx helpers with functions on command.Command

Makefile

@@ -7,7 +7,8 @@ BINDIR = $(PREFIX)/bin
 INSTALL := install -m 0755
 
 GO ?= go
-GORELEASER := $(GO) tool goreleaser
+# GORELEASER := $(GO) tool goreleaser
+GORELEASER := $(GO) run github.com/goreleaser/goreleaser/v2@latest
 GOIMPORTS := $(GO) tool goimports
 GOLINT := $(GO) tool golangci-lint
 
@@ -67,10 +68,11 @@ clean:
 
 .PHONY: test-go
 test-go: pkg/dockerfile/embed/.wheel
-	$(GO) tool gotestsum -- -timeout 1200s -parallel 5 ./... $(ARGS)
+	$(GO) tool gotestsum -- -short -timeout 1200s -parallel 5 ./... $(ARGS)
 
 .PHONY: test-integration
 test-integration: $(COG_BINARIES)
+	$(GO) test ./pkg/docker/...
 	PATH="$(PWD):$(PATH)" $(TOX) -e integration
 
 .PHONY: test-python

pkg/cli/predict.go

@@ -69,6 +69,8 @@ the prediction on that.`,
 func cmdPredict(cmd *cobra.Command, args []string) error {
 	ctx := cmd.Context()
 
+	dockerCommand := docker.NewDockerCommand()
+
 	imageName := ""
 	volumes := []docker.Volume{}
 	gpus := gpusFlag
@@ -112,17 +114,12 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
 			return fmt.Errorf("Invalid image name '%s'. Did you forget `-i`?", imageName)
 		}
 
-		exists, err := docker.ImageExists(ctx, imageName)
+		inspectResp, err := dockerCommand.Pull(ctx, imageName, false)
 		if err != nil {
-			return fmt.Errorf("Failed to determine if %s exists: %w", imageName, err)
-		}
-		if !exists {
-			console.Infof("Pulling image: %s", imageName)
-			if err := docker.Pull(ctx, imageName); err != nil {
-				return fmt.Errorf("Failed to pull %s: %w", imageName, err)
-			}
+			return fmt.Errorf("Failed to pull image %q: %w", imageName, err)
 		}
-		conf, err := image.GetConfig(ctx, imageName)
+
+		conf, err := image.CogConfigFromManifest(ctx, inspectResp)
 		if err != nil {
 			return err
 		}
@@ -136,7 +133,6 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
 
 	console.Info("")
 	console.Infof("Starting Docker image %s and running setup()...", imageName)
-	dockerCommand := docker.NewDockerCommand()
 
 	predictor, err := predict.NewPredictor(ctx, docker.RunOptions{
 		GPUs:    gpus,

pkg/cli/train.go

@@ -55,6 +55,8 @@ Otherwise, it will build the model in the current directory and train it.`,
 func cmdTrain(cmd *cobra.Command, args []string) error {
 	ctx := cmd.Context()
 
+	dockerCommand := docker.NewDockerCommand()
+
 	imageName := ""
 	volumes := []docker.Volume{}
 	gpus := gpusFlag
@@ -88,17 +90,12 @@ func cmdTrain(cmd *cobra.Command, args []string) error {
 		// Use existing image
 		imageName = args[0]
 
-		exists, err := docker.ImageExists(ctx, imageName)
+		inspectResp, err := dockerCommand.Pull(ctx, imageName, false)
 		if err != nil {
-			return fmt.Errorf("Failed to determine if %s exists: %w", imageName, err)
-		}
-		if !exists {
-			console.Infof("Pulling image: %s", imageName)
-			if err := docker.Pull(ctx, imageName); err != nil {
-				return fmt.Errorf("Failed to pull %s: %w", imageName, err)
-			}
+			return fmt.Errorf("Failed to pull image %q: %w", imageName, err)
 		}
-		conf, err := image.GetConfig(ctx, imageName)
+
+		conf, err := image.CogConfigFromManifest(ctx, inspectResp)
 		if err != nil {
 			return err
 		}
@@ -112,7 +109,6 @@ func cmdTrain(cmd *cobra.Command, args []string) error {
 
 	console.Info("")
 	console.Infof("Starting Docker image %s...", imageName)
-	dockerCommand := docker.NewDockerCommand()
 
 	predictor, err := predict.NewPredictor(ctx, docker.RunOptions{
 		GPUs:    gpus,

pkg/docker/command/command.go

@@ -1,12 +1,25 @@
 package command
 
-import "context"
+import (
+	"context"
+	"io"
+
+	"github.com/docker/docker/api/types/container"
+	"github.com/docker/docker/api/types/image"
+)
 
 type Command interface {
-	Pull(ctx context.Context, ref string) error
+	// Pull pulls an image from a remote registry and returns the inspect response for the local image.
+	// If the image already exists, it will return the inspect response for the local image without pulling.
+	// When force is true, it will always attempt to pull the image.
+	Pull(ctx context.Context, ref string, force bool) (*image.InspectResponse, error)
 	Push(ctx context.Context, ref string) error
 	LoadUserInformation(ctx context.Context, registryHost string) (*UserInfo, error)
 	CreateTarFile(ctx context.Context, ref string, tmpDir string, tarFile string, folder string) (string, error)
 	CreateAptTarFile(ctx context.Context, tmpDir string, aptTarFile string, packages ...string) (string, error)
-	Inspect(ctx context.Context, ref string) (*Manifest, error)
+	Inspect(ctx context.Context, ref string) (*image.InspectResponse, error)
+	ImageExists(ctx context.Context, ref string) (bool, error)
+	ContainerLogs(ctx context.Context, containerID string, w io.Writer) error
+	ContainerInspect(ctx context.Context, id string) (*container.InspectResponse, error)
+	ContainerStop(ctx context.Context, containerID string) error
 }

pkg/docker/command/errors.go

@@ -0,0 +1,31 @@
+package command
+
+import (
+	"errors"
+	"fmt"
+)
+
+// NotFoundError represents “object <ref> wasn’t found” inside the Docker engine.
+type NotFoundError struct {
+	// Ref is a unique identifier, such as an image reference, container ID, etc.
+	Ref string
+	// Object is the ref type, such as "container", "image", "volume", etc.
+	Object string
+}
+
+func (e *NotFoundError) Error() string {
+	objType := e.Object
+	if objType == "" {
+		objType = "object"
+	}
+	return fmt.Sprintf("%s not found: %q", objType, e.Ref)
+}
+
+func (e *NotFoundError) Is(target error) bool {
+	_, ok := target.(*NotFoundError)
+	return ok
+}
+
+func IsNotFoundError(err error) bool {
+	return errors.Is(err, &NotFoundError{})
+}