Add support for local Docker images (#1114)

Author: RobertLucian

pkg/lib/configreader/validators.go

@@ -22,6 +22,7 @@ import (
 	"time"
 
 	"github.com/cortexlabs/cortex/pkg/lib/aws"
+	"github.com/cortexlabs/cortex/pkg/lib/docker"
 	"github.com/cortexlabs/cortex/pkg/lib/files"
 	"github.com/cortexlabs/cortex/pkg/lib/urls"
 )
@@ -155,12 +156,7 @@ func ValidateImageVersion(image, cortexVersion string) (string, error) {
 		return image, nil
 	}
 
-	var tag string
-
-	if colonIndex := strings.LastIndex(image, ":"); colonIndex != -1 {
-		tag = image[colonIndex+1:]
-	}
-
+	tag := docker.ExtractImageTag(image)
 	// in docker, missing tag implies "latest"
 	if tag == "" {
 		tag = "latest"

pkg/lib/docker/docker.go

@@ -33,6 +33,7 @@ import (
 	"github.com/cortexlabs/cortex/pkg/lib/exit"
 	"github.com/cortexlabs/cortex/pkg/lib/parallel"
 	"github.com/cortexlabs/cortex/pkg/lib/print"
+	"github.com/cortexlabs/cortex/pkg/lib/slices"
 	dockertypes "github.com/docker/docker/api/types"
 	dockerclient "github.com/docker/docker/client"
 	"github.com/docker/docker/pkg/jsonmessage"
@@ -142,17 +143,8 @@ func PullImage(image string, encodedAuthConfig string, pullVerbosity PullVerbosi
 		return false, err
 	}
 
-	existingImages, err := dockerClient.ImageList(context.Background(), dockertypes.ImageListOptions{})
-	if err != nil {
-		return false, WrapDockerError(err)
-	}
-
-	for _, existingImage := range existingImages {
-		for _, tag := range existingImage.RepoTags {
-			if tag == image {
-				return false, nil
-			}
-		}
+	if err := CheckLocalImageAccessible(dockerClient, image); err == nil {
+		return false, nil
 	}
 
 	pullOutput, err := dockerClient.ImagePull(context.Background(), image, dockertypes.ImagePullOptions{
@@ -245,3 +237,29 @@ func CheckImageAccessible(dockerClient *Client, dockerImage, authConfig string)
 	}
 	return nil
 }
+
+func CheckLocalImageAccessible(dockerClient *Client, dockerImage string) error {
+	images, err := dockerClient.ImageList(context.Background(), dockertypes.ImageListOptions{})
+	if err != nil {
+		return WrapDockerError(err)
+	}
+
+	// in docker, missing tag implies "latest"
+	if ExtractImageTag(dockerImage) == "" {
+		dockerImage = fmt.Sprintf("%s:latest", dockerImage)
+	}
+
+	for _, image := range images {
+		if slices.HasString(image.RepoTags, dockerImage) {
+			return nil
+		}
+	}
+	return ErrorImageInaccessible(dockerImage, nil)
+}
+
+func ExtractImageTag(dockerImage string) string {
+	if colonIndex := strings.LastIndex(dockerImage, ":"); colonIndex != -1 {
+		return dockerImage[colonIndex+1:]
+	}
+	return ""
+}

pkg/lib/docker/errors.go

@@ -57,10 +57,14 @@ func ErrorDockerPermissions(err error) error {
 }
 
 func ErrorImageInaccessible(image string, cause error) error {
-	dockerErrMsg := errors.Message(cause)
+	message := fmt.Sprintf("%s is not accessible", image)
+	if cause != nil {
+		message += "\n" + errors.Message(cause) // add \n because docker client errors are
+	}
+
 	return errors.WithStack(&errors.Error{
 		Kind:    ErrImageInaccessible,
-		Message: fmt.Sprintf("%s is not accessible\n%s", image, dockerErrMsg), // add \n because docker client errors are verbose
+		Message: message,
 		Cause:   cause,
 	})
 }

pkg/types/spec/validations.go

@@ -471,7 +471,7 @@ func validatePredictor(predictor *userconfig.Predictor, projectFiles ProjectFile
 		if err := validateTensorFlowPredictor(predictor, providerType, projectFiles, awsClient); err != nil {
 			return err
 		}
-		if err := validateDockerImagePath(predictor.TensorFlowServingImage, awsClient); err != nil {
+		if err := validateDockerImagePath(predictor.TensorFlowServingImage, providerType, awsClient); err != nil {
 			return errors.Wrap(err, userconfig.TensorFlowServingImageKey)
 		}
 	case userconfig.ONNXPredictorType:
@@ -480,7 +480,7 @@ func validatePredictor(predictor *userconfig.Predictor, projectFiles ProjectFile
 		}
 	}
 
-	if err := validateDockerImagePath(predictor.Image, awsClient); err != nil {
+	if err := validateDockerImagePath(predictor.Image, providerType, awsClient); err != nil {
 		return errors.Wrap(err, userconfig.ImageKey)
 	}
 
@@ -828,14 +828,26 @@ func FindDuplicateNames(apis []userconfig.API) []userconfig.API {
 	return nil
 }
 
-func validateDockerImagePath(image string, awsClient *aws.Client) error {
+func validateDockerImagePath(image string, providerType types.ProviderType, awsClient *aws.Client) error {
 	if consts.DefaultImagePathsSet.Has(image) {
 		return nil
 	}
 	if _, err := cr.ValidateImageVersion(image, consts.CortexVersion); err != nil {
 		return err
 	}
 
+	dockerClient, err := docker.GetDockerClient()
+	if err != nil {
+		return err
+	}
+
+	if providerType == types.LocalProviderType {
+		// short circuit if the image is already available locally
+		if err := docker.CheckLocalImageAccessible(dockerClient, image); err == nil {
+			return nil
+		}
+	}
+
 	dockerAuth := docker.NoAuth
 	if regex.IsValidECRURL(image) {
 		if awsClient.IsAnonymous {
@@ -871,13 +883,9 @@ func validateDockerImagePath(image string, awsClient *aws.Client) error {
 		}
 	}
 
-	dockerClient, err := docker.GetDockerClient()
-	if err != nil {
-		return err
-	}
-
 	if err := docker.CheckImageAccessible(dockerClient, image, dockerAuth); err != nil {
 		return err
 	}
+
 	return nil
 }