Read and validate from buckets in any regions (#1059)

vishalbollu · deliahu · commit b5b8766f6735 · 2020-05-14T19:01:24.000-07:00
(cherry picked from commit 468dce0)
diff --git a/cli/local/api.go b/cli/local/api.go
@@ -50,7 +50,7 @@ func UpdateAPI(apiConfig *userconfig.API, cortexYAMLPath string, projectID strin
 	if apiConfig.Predictor.Model != nil {
 		localModelCache, err := CacheModel(*apiConfig.Predictor.Model, awsClient)
 		if err != nil {
-			return nil, "", errors.Wrap(err, userconfig.ModelKey, userconfig.PredictorKey, apiConfig.Identify())
+			return nil, "", errors.Wrap(err, apiConfig.Identify(), userconfig.PredictorKey, userconfig.ModelKey)
 		}
 		apiSpec.LocalModelCache = localModelCache
 	}
diff --git a/cli/local/model_cache.go b/cli/local/model_cache.go
@@ -36,12 +36,17 @@ import (
 func CacheModel(modelPath string, awsClient *aws.Client) (*spec.LocalModelCache, error) {
 	localModelCache := spec.LocalModelCache{}
 
+	awsClientForBucket, err := aws.NewFromClientS3Path(modelPath, awsClient)
+	if err != nil {
+		return nil, err
+	}
+
 	if strings.HasPrefix(modelPath, "s3://") {
 		bucket, prefix, err := aws.SplitS3Path(modelPath)
 		if err != nil {
 			return nil, err
 		}
-		hash, err := awsClient.HashS3Dir(bucket, prefix, nil)
+		hash, err := awsClientForBucket.HashS3Dir(bucket, prefix, nil)
 		if err != nil {
 			return nil, err
 		}
@@ -61,13 +66,13 @@ func CacheModel(modelPath string, awsClient *aws.Client) (*spec.LocalModelCache,
 		return &localModelCache, nil
 	}
 
-	err := ResetModelCacheDir(modelDir)
+	err = ResetModelCacheDir(modelDir)
 	if err != nil {
 		return nil, err
 	}
 
 	if strings.HasPrefix(modelPath, "s3://") {
-		err := downloadModel(modelPath, modelDir, awsClient)
+		err := downloadModel(modelPath, modelDir, awsClientForBucket)
 		if err != nil {
 			return nil, err
 		}
@@ -105,7 +110,7 @@ func CacheModel(modelPath string, awsClient *aws.Client) (*spec.LocalModelCache,
 	return &localModelCache, nil
 }
 
-func downloadModel(modelPath string, modelDir string, awsClient *aws.Client) error {
+func downloadModel(modelPath string, modelDir string, awsClientForBucket *aws.Client) error {
 	fmt.Printf("￮ downloading model %s ", modelPath)
 	defer fmt.Print(" ✓\n")
 	dotCron := cron.Run(print.Dot, nil, 2*time.Second)
@@ -118,7 +123,7 @@ func downloadModel(modelPath string, modelDir string, awsClient *aws.Client) err
 
 	if strings.HasSuffix(modelPath, ".zip") || strings.HasSuffix(modelPath, ".onnx") {
 		localPath := filepath.Join(modelDir, filepath.Base(modelPath))
-		err := awsClient.DownloadFileFromS3(bucket, prefix, localPath)
+		err := awsClientForBucket.DownloadFileFromS3(bucket, prefix, localPath)
 		if err != nil {
 			return err
 		}
@@ -134,7 +139,7 @@ func downloadModel(modelPath string, modelDir string, awsClient *aws.Client) err
 		}
 	} else {
 		tfModelVersion := filepath.Base(prefix)
-		err := awsClient.DownloadDirFromS3(bucket, prefix, filepath.Join(modelDir, tfModelVersion), true, nil)
+		err := awsClientForBucket.DownloadDirFromS3(bucket, prefix, filepath.Join(modelDir, tfModelVersion), true, nil)
 		if err != nil {
 			return err
 		}
diff --git a/pkg/lib/aws/aws.go b/pkg/lib/aws/aws.go
@@ -23,6 +23,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/credentials"
 	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/cortexlabs/cortex/pkg/lib/errors"
 )
 
 type Client struct {
@@ -43,6 +44,22 @@ func NewFromCreds(region string, accessKeyID string, secretAccessKey string) (*C
 	return New(region, creds)
 }
 
+func NewFromClientS3Path(s3Path string, awsClient *Client) (*Client, error) {
+	if !awsClient.IsAnonymous {
+		if awsClient.AccessKeyID() == nil || awsClient.SecretAccessKey() == nil {
+			return nil, ErrorUnexpectedMissingCredentials(awsClient.AccessKeyID(), awsClient.SecretAccessKey())
+		}
+		return NewFromCredsS3Path(s3Path, *awsClient.AccessKeyID(), *awsClient.SecretAccessKey())
+	}
+
+	region, err := GetBucketRegionFromS3Path(s3Path)
+	if err != nil {
+		return nil, err
+	}
+
+	return NewAnonymousClientWithRegion(region)
+}
+
 func NewFromEnvS3Path(s3Path string) (*Client, error) {
 	bucket, _, err := SplitS3Path(s3Path)
 	if err != nil {
@@ -85,7 +102,7 @@ func New(region string, creds *credentials.Credentials) (*Client, error) {
 	})
 
 	if err != nil {
-		return nil, err
+		return nil, errors.WithStack(err)
 	}
 
 	return &Client{
diff --git a/pkg/lib/aws/errors.go b/pkg/lib/aws/errors.go
@@ -25,16 +25,17 @@ import (
 )
 
 const (
-	ErrInvalidAWSCredentials    = "aws.invalid_aws_credentials"
-	ErrInvalidS3aPath           = "aws.invalid_s3a_path"
-	ErrInvalidS3Path            = "aws.invalid_s3_path"
-	ErrAuth                     = "aws.auth"
-	ErrBucketInaccessible       = "aws.bucket_inaccessible"
-	ErrBucketNotFound           = "aws.bucket_not_found"
-	ErrInstanceTypeLimitIsZero  = "aws.instance_type_limit_is_zero"
-	ErrNoValidSpotPrices        = "aws.no_valid_spot_prices"
-	ErrReadCredentials          = "aws.read_credentials"
-	ErrECRExtractingCredentials = "aws.ecr_failed_credentials"
+	ErrInvalidAWSCredentials        = "aws.invalid_aws_credentials"
+	ErrInvalidS3aPath               = "aws.invalid_s3a_path"
+	ErrInvalidS3Path                = "aws.invalid_s3_path"
+	ErrUnexpectedMissingCredentials = "aws.unexpected_missing_credentials"
+	ErrAuth                         = "aws.auth"
+	ErrBucketInaccessible           = "aws.bucket_inaccessible"
+	ErrBucketNotFound               = "aws.bucket_not_found"
+	ErrInstanceTypeLimitIsZero      = "aws.instance_type_limit_is_zero"
+	ErrNoValidSpotPrices            = "aws.no_valid_spot_prices"
+	ErrReadCredentials              = "aws.read_credentials"
+	ErrECRExtractingCredentials     = "aws.ecr_failed_credentials"
 )
 
 func IsNotFoundErr(err error) bool {
@@ -91,6 +92,22 @@ func ErrorInvalidS3Path(provided string) error {
 	})
 }
 
+func ErrorUnexpectedMissingCredentials(awsAccessKeyID *string, awsSecretAccessKey *string) error {
+	var msg string
+	if awsAccessKeyID == nil && awsSecretAccessKey == nil {
+		msg = "aws access key id and aws secret access key are missing"
+	} else if awsAccessKeyID == nil {
+		msg = "aws access key id is missing"
+	} else if awsSecretAccessKey == nil {
+		msg = "aws secret access key is missing"
+	}
+
+	return errors.WithStack(&errors.Error{
+		Kind:    ErrUnexpectedMissingCredentials,
+		Message: msg,
+	})
+}
+
 func ErrorAuth() error {
 	return errors.WithStack(&errors.Error{
 		Kind:    ErrAuth,
diff --git a/pkg/lib/aws/s3.go b/pkg/lib/aws/s3.go
@@ -122,6 +122,15 @@ func IsValidS3aPath(s3aPath string) bool {
 	return true
 }
 
+func GetBucketRegionFromS3Path(s3Path string) (string, error) {
+	bucket, _, err := SplitS3Path(s3Path)
+	if err != nil {
+		return "", err
+	}
+
+	return GetBucketRegion(bucket)
+}
+
 func GetBucketRegion(bucket string) (string, error) {
 	sess := session.Must(session.NewSession()) // credentials are not necessary for this request, and will not be used
 	region, err := s3manager.GetBucketRegion(aws.BackgroundContext(), sess, bucket, endpoints.UsWest2RegionID)
diff --git a/pkg/types/spec/validations.go b/pkg/types/spec/validations.go
@@ -517,18 +517,24 @@ func validateTensorFlowPredictor(predictor *userconfig.Predictor, providerType t
 	}
 
 	model := *predictor.Model
+
+	awsClientForBucket, err := aws.NewFromClientS3Path(model, awsClient)
+	if err != nil {
+		return errors.Wrap(err, userconfig.ModelKey)
+	}
+
 	if strings.HasPrefix(model, "s3://") {
 		model, err := cr.S3PathValidator(model)
 		if err != nil {
 			return errors.Wrap(err, userconfig.ModelKey)
 		}
 
 		if strings.HasSuffix(model, ".zip") {
-			if ok, err := awsClient.IsS3PathFile(model); err != nil || !ok {
+			if ok, err := awsClientForBucket.IsS3PathFile(model); err != nil || !ok {
 				return errors.Wrap(ErrorS3FileNotFound(model), userconfig.ModelKey)
 			}
 		} else {
-			path, err := getTFServingExportFromS3Path(model, awsClient)
+			path, err := getTFServingExportFromS3Path(model, awsClientForBucket)
 			if err != nil {
 				return errors.Wrap(err, userconfig.ModelKey)
 			} else if path == "" {
@@ -584,13 +590,18 @@ func validateONNXPredictor(predictor *userconfig.Predictor, providerType types.P
 		return errors.Wrap(ErrorInvalidONNXModelPath(), userconfig.ModelKey, model)
 	}
 
+	awsClientForBucket, err := aws.NewFromClientS3Path(model, awsClient)
+	if err != nil {
+		return errors.Wrap(err, userconfig.ModelKey)
+	}
+
 	if strings.HasPrefix(model, "s3://") {
 		model, err := cr.S3PathValidator(model)
 		if err != nil {
 			return errors.Wrap(err, userconfig.ModelKey)
 		}
 
-		if ok, err := awsClient.IsS3PathFile(model); err != nil || !ok {
+		if ok, err := awsClientForBucket.IsS3PathFile(model); err != nil || !ok {
 			return errors.Wrap(ErrorS3FileNotFound(model), userconfig.ModelKey)
 		}
 	} else {
@@ -620,8 +631,8 @@ func validateONNXPredictor(predictor *userconfig.Predictor, providerType types.P
 	return nil
 }
 
-func getTFServingExportFromS3Path(path string, awsClient *aws.Client) (string, error) {
-	if isValidTensorFlowS3Directory(path, awsClient) {
+func getTFServingExportFromS3Path(path string, awsClientForBucket *aws.Client) (string, error) {
+	if isValidTensorFlowS3Directory(path, awsClientForBucket) {
 		return path, nil
 	}
 
@@ -630,7 +641,7 @@ func getTFServingExportFromS3Path(path string, awsClient *aws.Client) (string, e
 		return "", err
 	}
 
-	objects, err := awsClient.ListS3PathDir(path, false, pointer.Int64(1000))
+	objects, err := awsClientForBucket.ListS3PathDir(path, false, pointer.Int64(1000))
 	if err != nil {
 		return "", err
 	} else if len(objects) == 0 {
@@ -652,7 +663,7 @@ func getTFServingExportFromS3Path(path string, awsClient *aws.Client) (string, e
 		}
 
 		possiblePath := "s3://" + filepath.Join(bucket, filepath.Join(keyParts[:len(keyParts)-1]...))
-		if version >= highestVersion && isValidTensorFlowS3Directory(possiblePath, awsClient) {
+		if version >= highestVersion && isValidTensorFlowS3Directory(possiblePath, awsClientForBucket) {
 			highestVersion = version
 			highestPath = possiblePath
 		}
@@ -668,15 +679,15 @@ func getTFServingExportFromS3Path(path string, awsClient *aws.Client) (string, e
 //		- variables/
 //			- variables.index
 //			- variables.data-00000-of-00001 (there are a variable number of these files)
-func isValidTensorFlowS3Directory(path string, awsClient *aws.Client) bool {
-	if valid, err := awsClient.IsS3PathFile(
+func isValidTensorFlowS3Directory(path string, awsClientForBucket *aws.Client) bool {
+	if valid, err := awsClientForBucket.IsS3PathFile(
 		aws.JoinS3Path(path, "saved_model.pb"),
 		aws.JoinS3Path(path, "variables/variables.index"),
 	); err != nil || !valid {
 		return false
 	}
 
-	if valid, err := awsClient.IsS3PathPrefix(
+	if valid, err := awsClientForBucket.IsS3PathPrefix(
 		aws.JoinS3Path(path, "variables/variables.data-00000-of"),
 	); err != nil || !valid {
 		return false

Original file line number	Diff line number	Diff line change
`@@ -50,7 +50,7 @@ func UpdateAPI(apiConfig *userconfig.API, cortexYAMLPath string, projectID strin`
`50`	`50`	`if apiConfig.Predictor.Model != nil {`
`51`	`51`	`localModelCache, err := CacheModel(*apiConfig.Predictor.Model, awsClient)`
`52`	`52`	`if err != nil {`
`53`		`- return nil, "", errors.Wrap(err, userconfig.ModelKey, userconfig.PredictorKey, apiConfig.Identify())`
	`53`	`+ return nil, "", errors.Wrap(err, apiConfig.Identify(), userconfig.PredictorKey, userconfig.ModelKey)`
`54`	`54`	`}`
`55`	`55`	`apiSpec.LocalModelCache = localModelCache`
`56`	`56`	`}`