Improve batch job submit validation efficiency (#2179)

deliahu · web-flow · commit 8226f42fe197 · 2021-05-22T00:48:07.000-07:00
diff --git a/pkg/enqueuer/enqueuer.go b/pkg/enqueuer/enqueuer.go
@@ -56,9 +56,10 @@ type ItemList struct {
 }
 
 type S3Lister struct {
-	S3Paths  []string `json:"s3_paths"` // s3://<bucket_name>/key
-	Includes []string `json:"includes"`
-	Excludes []string `json:"excludes"`
+	S3Paths    []string `json:"s3_paths"` // s3://<bucket_name>/key
+	Includes   []string `json:"includes"`
+	Excludes   []string `json:"excludes"`
+	MaxResults *int64   `json:"-"` // this is not currently exposed to the user (it's used for validations)
 }
 
 type FilePathLister struct {
@@ -246,7 +247,7 @@ func (e *Enqueuer) enqueueS3Paths(s3PathsLister *FilePathLister) (int, error) {
 	var s3PathList []string
 	uploader := newSQSBatchUploader(e.envConfig.APIName, e.envConfig.JobID, e.queueURL, e.aws.SQS())
 
-	err := s3IteratorFromLister(e.aws, s3PathsLister.S3Lister, func(bucket string, s3Obj *s3.Object) (bool, error) {
+	_, err := s3IteratorFromLister(e.aws, s3PathsLister.S3Lister, func(bucket string, s3Obj *s3.Object) (bool, error) {
 		s3Path := awslib.S3Path(bucket, *s3Obj.Key)
 
 		s3PathList = append(s3PathList, s3Path)
@@ -290,7 +291,7 @@ func (e *Enqueuer) enqueueS3FileContents(delimitedFiles *DelimitedFiles) (int, e
 	uploader := newSQSBatchUploader(e.envConfig.APIName, e.envConfig.JobID, e.queueURL, e.aws.SQS())
 
 	bytesBuffer := bytes.NewBuffer([]byte{})
-	err := s3IteratorFromLister(e.aws, delimitedFiles.S3Lister, func(bucket string, s3Obj *s3.Object) (bool, error) {
+	_, err := s3IteratorFromLister(e.aws, delimitedFiles.S3Lister, func(bucket string, s3Obj *s3.Object) (bool, error) {
 		s3Path := awslib.S3Path(bucket, *s3Obj.Key)
 		log.Info("enqueuing contents from file", zap.String("path", s3Path))
 
diff --git a/pkg/enqueuer/helpers.go b/pkg/enqueuer/helpers.go
@@ -64,13 +64,13 @@ func addJSONObjectsToQueue(uploader *sqsBatchUploader, jsonMessageList *jsonBuff
 	return nil
 }
 
-func s3IteratorFromLister(awsClient *awslib.Client, s3Lister S3Lister, fn func(string, *s3.Object) (bool, error)) error {
+func s3IteratorFromLister(awsClient *awslib.Client, s3Lister S3Lister, fn func(string, *s3.Object) (bool, error)) (int64, error) {
 	includeGlobPatterns := make([]glob.Glob, 0, len(s3Lister.Includes))
 
 	for _, includePattern := range s3Lister.Includes {
 		globExpression, err := glob.Compile(includePattern, '/')
 		if err != nil {
-			return errors.Wrap(err, "failed to interpret glob pattern", includePattern)
+			return 0, errors.Wrap(err, "failed to interpret glob pattern", includePattern)
 		}
 		includeGlobPatterns = append(includeGlobPatterns, globExpression)
 	}
@@ -79,20 +79,22 @@ func s3IteratorFromLister(awsClient *awslib.Client, s3Lister S3Lister, fn func(s
 	for _, excludePattern := range s3Lister.Excludes {
 		globExpression, err := glob.Compile(excludePattern, '/')
 		if err != nil {
-			return errors.Wrap(err, "failed to interpret glob pattern", excludePattern)
+			return 0, errors.Wrap(err, "failed to interpret glob pattern", excludePattern)
 		}
 		excludeGlobPatterns = append(excludeGlobPatterns, globExpression)
 	}
 
+	var numResults int64
+
 	for _, s3Path := range s3Lister.S3Paths {
 		bucket, key, err := awslib.SplitS3Path(s3Path)
 		if err != nil {
-			return err
+			return 0, err
 		}
 
 		awsClientForBucket, err := awslib.NewFromClientS3Path(s3Path, awsClient)
 		if err != nil {
-			return err
+			return 0, err
 		}
 
 		err = awsClientForBucket.S3Iterator(bucket, key, false, nil, nil, func(s3Obj *s3.Object) (bool, error) {
@@ -117,15 +119,24 @@ func s3IteratorFromLister(awsClient *awslib.Client, s3Lister S3Lister, fn func(s
 			}
 
 			if !shouldSkip {
-				return fn(bucket, s3Obj)
+				shouldContinue, err := fn(bucket, s3Obj)
+				numResults++
+				if s3Lister.MaxResults != nil && numResults >= *s3Lister.MaxResults {
+					shouldContinue = false
+				}
+				return shouldContinue, err
 			}
 
 			return true, nil
 		})
 		if err != nil {
-			return err
+			return 0, err
+		}
+
+		if s3Lister.MaxResults != nil && numResults >= *s3Lister.MaxResults {
+			return numResults, nil
 		}
 	}
 
-	return nil
+	return numResults, nil
 }
diff --git a/pkg/operator/resources/job/batchapi/s3_iterator.go b/pkg/operator/resources/job/batchapi/s3_iterator.go
@@ -26,13 +26,13 @@ import (
 )
 
 // Takes in a function(shouldSkip, bucketName, s3.Object)
-func s3IteratorFromLister(s3Lister schema.S3Lister, fn func(string, *s3.Object) (bool, error)) error {
+func s3IteratorFromLister(s3Lister schema.S3Lister, fn func(string, *s3.Object) (bool, error)) (int64, error) {
 	includeGlobPatterns := make([]glob.Glob, 0, len(s3Lister.Includes))
 
 	for _, includePattern := range s3Lister.Includes {
 		globExpression, err := glob.Compile(includePattern, '/')
 		if err != nil {
-			return errors.Wrap(err, "failed to interpret glob pattern", includePattern)
+			return 0, errors.Wrap(err, "failed to interpret glob pattern", includePattern)
 		}
 		includeGlobPatterns = append(includeGlobPatterns, globExpression)
 	}
@@ -41,20 +41,22 @@ func s3IteratorFromLister(s3Lister schema.S3Lister, fn func(string, *s3.Object)
 	for _, excludePattern := range s3Lister.Excludes {
 		globExpression, err := glob.Compile(excludePattern, '/')
 		if err != nil {
-			return errors.Wrap(err, "failed to interpret glob pattern", excludePattern)
+			return 0, errors.Wrap(err, "failed to interpret glob pattern", excludePattern)
 		}
 		excludeGlobPatterns = append(excludeGlobPatterns, globExpression)
 	}
 
+	var numResults int64
+
 	for _, s3Path := range s3Lister.S3Paths {
 		bucket, key, err := aws.SplitS3Path(s3Path)
 		if err != nil {
-			return err
+			return 0, err
 		}
 
 		awsClientForBucket, err := aws.NewFromClientS3Path(s3Path, config.AWS)
 		if err != nil {
-			return err
+			return 0, err
 		}
 
 		err = awsClientForBucket.S3Iterator(bucket, key, false, nil, nil, func(s3Obj *s3.Object) (bool, error) {
@@ -79,15 +81,24 @@ func s3IteratorFromLister(s3Lister schema.S3Lister, fn func(string, *s3.Object)
 			}
 
 			if !shouldSkip {
-				return fn(bucket, s3Obj)
+				shouldContinue, err := fn(bucket, s3Obj)
+				numResults++
+				if s3Lister.MaxResults != nil && numResults >= *s3Lister.MaxResults {
+					shouldContinue = false
+				}
+				return shouldContinue, err
 			}
 
 			return true, nil
 		})
 		if err != nil {
-			return err
+			return 0, err
+		}
+
+		if s3Lister.MaxResults != nil && numResults >= *s3Lister.MaxResults {
+			return numResults, nil
 		}
 	}
 
-	return nil
+	return numResults, nil
 }
diff --git a/pkg/operator/resources/job/batchapi/validations.go b/pkg/operator/resources/job/batchapi/validations.go
@@ -24,6 +24,7 @@ import (
 	awslib "github.com/cortexlabs/cortex/pkg/lib/aws"
 	cr "github.com/cortexlabs/cortex/pkg/lib/configreader"
 	"github.com/cortexlabs/cortex/pkg/lib/errors"
+	"github.com/cortexlabs/cortex/pkg/lib/pointer"
 	"github.com/cortexlabs/cortex/pkg/operator/resources/job"
 	"github.com/cortexlabs/cortex/pkg/operator/schema"
 	"github.com/gobwas/glob"
@@ -143,26 +144,30 @@ func validateS3Lister(s3Lister *schema.S3Lister) error {
 		}
 	}
 
-	filesFound := 0
 	for _, s3Path := range s3Lister.S3Paths {
 		if !awslib.IsValidS3Path(s3Path) {
 			return awslib.ErrorInvalidS3Path(s3Path)
 		}
+	}
 
-		err := s3IteratorFromLister(*s3Lister, func(objPath string, s3Obj *s3.Object) (bool, error) {
-			filesFound++
-			return false, nil
-		})
-		if err != nil {
-			return errors.Wrap(err, s3Path)
-		}
+	shortCircuitLister := schema.S3Lister{
+		S3Paths:    s3Lister.S3Paths,
+		Includes:   s3Lister.Includes,
+		Excludes:   s3Lister.Excludes,
+		MaxResults: pointer.Int64(1),
+	}
+	numResults, err := s3IteratorFromLister(shortCircuitLister, func(objPath string, s3Obj *s3.Object) (bool, error) {
+		return false, nil
+	})
+	if err != nil {
+		return err
+	}
 
-		if filesFound > 0 {
-			return nil
-		}
+	if numResults == 0 {
+		return ErrorNoS3FilesFound()
 	}
 
-	return ErrorNoS3FilesFound()
+	return nil
 }
 
 func listFilesDryRun(s3Lister *schema.S3Lister) ([]string, error) {
@@ -171,15 +176,14 @@ func listFilesDryRun(s3Lister *schema.S3Lister) ([]string, error) {
 		if !awslib.IsValidS3Path(s3Path) {
 			return nil, awslib.ErrorInvalidS3Path(s3Path)
 		}
+	}
 
-		err := s3IteratorFromLister(*s3Lister, func(bucket string, s3Obj *s3.Object) (bool, error) {
-			s3Files = append(s3Files, awslib.S3Path(bucket, *s3Obj.Key))
-			return true, nil
-		})
-
-		if err != nil {
-			return nil, errors.Wrap(err, s3Path)
-		}
+	_, err := s3IteratorFromLister(*s3Lister, func(bucket string, s3Obj *s3.Object) (bool, error) {
+		s3Files = append(s3Files, awslib.S3Path(bucket, *s3Obj.Key))
+		return true, nil
+	})
+	if err != nil {
+		return nil, err
 	}
 
 	if len(s3Files) == 0 {
diff --git a/pkg/operator/schema/job_submission.go b/pkg/operator/schema/job_submission.go
@@ -28,9 +28,10 @@ type ItemList struct {
 }
 
 type S3Lister struct {
-	S3Paths  []string `json:"s3_paths"` // s3://<bucket_name>/key
-	Includes []string `json:"includes"`
-	Excludes []string `json:"excludes"`
+	S3Paths    []string `json:"s3_paths"` // s3://<bucket_name>/key
+	Includes   []string `json:"includes"`
+	Excludes   []string `json:"excludes"`
+	MaxResults *int64   `json:"-"` // this is not currently exposed to the user (it's used for validations)
 }
 
 type FilePathLister struct {

Original file line number	Diff line number	Diff line change
`@@ -64,13 +64,13 @@ func addJSONObjectsToQueue(uploader sqsBatchUploader, jsonMessageList jsonBuff`
`64`	`64`	`return nil`
`65`	`65`	`}`
`66`	`66`
`67`		`-func s3IteratorFromLister(awsClient awslib.Client, s3Lister S3Lister, fn func(string, s3.Object) (bool, error)) error {`
	`67`	`+func s3IteratorFromLister(awsClient awslib.Client, s3Lister S3Lister, fn func(string, s3.Object) (bool, error)) (int64, error) {`
`68`	`68`	`includeGlobPatterns := make([]glob.Glob, 0, len(s3Lister.Includes))`
`69`	`69`
`70`	`70`	`for _, includePattern := range s3Lister.Includes {`
`71`	`71`	`globExpression, err := glob.Compile(includePattern, '/')`
`72`	`72`	`if err != nil {`
`73`		`- return errors.Wrap(err, "failed to interpret glob pattern", includePattern)`
	`73`	`+ return 0, errors.Wrap(err, "failed to interpret glob pattern", includePattern)`
`74`	`74`	`}`
`75`	`75`	`includeGlobPatterns = append(includeGlobPatterns, globExpression)`
`76`	`76`	`}`
`@@ -79,20 +79,22 @@ func s3IteratorFromLister(awsClient *awslib.Client, s3Lister S3Lister, fn func(s`
`79`	`79`	`for _, excludePattern := range s3Lister.Excludes {`
`80`	`80`	`globExpression, err := glob.Compile(excludePattern, '/')`
`81`	`81`	`if err != nil {`
`82`		`- return errors.Wrap(err, "failed to interpret glob pattern", excludePattern)`
	`82`	`+ return 0, errors.Wrap(err, "failed to interpret glob pattern", excludePattern)`
`83`	`83`	`}`
`84`	`84`	`excludeGlobPatterns = append(excludeGlobPatterns, globExpression)`
`85`	`85`	`}`
`86`	`86`
	`87`	`+ var numResults int64`
	`88`	`+`
`87`	`89`	`for _, s3Path := range s3Lister.S3Paths {`
`88`	`90`	`bucket, key, err := awslib.SplitS3Path(s3Path)`
`89`	`91`	`if err != nil {`
`90`		`- return err`
	`92`	`+ return 0, err`
`91`	`93`	`}`
`92`	`94`
`93`	`95`	`awsClientForBucket, err := awslib.NewFromClientS3Path(s3Path, awsClient)`
`94`	`96`	`if err != nil {`
`95`		`- return err`
	`97`	`+ return 0, err`
`96`	`98`	`}`
`97`	`99`
`98`	`100`	`err = awsClientForBucket.S3Iterator(bucket, key, false, nil, nil, func(s3Obj *s3.Object) (bool, error) {`
`@@ -117,15 +119,24 @@ func s3IteratorFromLister(awsClient *awslib.Client, s3Lister S3Lister, fn func(s`
`117`	`119`	`}`
`118`	`120`
`119`	`121`	`if !shouldSkip {`
`120`		`- return fn(bucket, s3Obj)`
	`122`	`+ shouldContinue, err := fn(bucket, s3Obj)`
	`123`	`+ numResults++`
	`124`	`+ if s3Lister.MaxResults != nil && numResults >= *s3Lister.MaxResults {`
	`125`	`+ shouldContinue = false`
	`126`	`+ }`
	`127`	`+ return shouldContinue, err`
`121`	`128`	`}`
`122`	`129`
`123`	`130`	`return true, nil`
`124`	`131`	`})`
`125`	`132`	`if err != nil {`
`126`		`- return err`
	`133`	`+ return 0, err`
	`134`	`+ }`
	`135`	`+`
	`136`	`+ if s3Lister.MaxResults != nil && numResults >= *s3Lister.MaxResults {`
	`137`	`+ return numResults, nil`
`127`	`138`	`}`
`128`	`139`	`}`
`129`	`140`
`130`		`- return nil`
	`141`	`+ return numResults, nil`
`131`	`142`	`}`
Original file line number	Diff line number	Diff line change
`@@ -26,13 +26,13 @@ import (`
`26`	`26`	`)`
`27`	`27`
`28`	`28`	`// Takes in a function(shouldSkip, bucketName, s3.Object)`
`29`		`-func s3IteratorFromLister(s3Lister schema.S3Lister, fn func(string, *s3.Object) (bool, error)) error {`
	`29`	`+func s3IteratorFromLister(s3Lister schema.S3Lister, fn func(string, *s3.Object) (bool, error)) (int64, error) {`
`30`	`30`	`includeGlobPatterns := make([]glob.Glob, 0, len(s3Lister.Includes))`
`31`	`31`
`32`	`32`	`for _, includePattern := range s3Lister.Includes {`
`33`	`33`	`globExpression, err := glob.Compile(includePattern, '/')`
`34`	`34`	`if err != nil {`
`35`		`- return errors.Wrap(err, "failed to interpret glob pattern", includePattern)`
	`35`	`+ return 0, errors.Wrap(err, "failed to interpret glob pattern", includePattern)`
`36`	`36`	`}`
`37`	`37`	`includeGlobPatterns = append(includeGlobPatterns, globExpression)`
`38`	`38`	`}`
`@@ -41,20 +41,22 @@ func s3IteratorFromLister(s3Lister schema.S3Lister, fn func(string, *s3.Object)`
`41`	`41`	`for _, excludePattern := range s3Lister.Excludes {`
`42`	`42`	`globExpression, err := glob.Compile(excludePattern, '/')`
`43`	`43`	`if err != nil {`
`44`		`- return errors.Wrap(err, "failed to interpret glob pattern", excludePattern)`
	`44`	`+ return 0, errors.Wrap(err, "failed to interpret glob pattern", excludePattern)`
`45`	`45`	`}`
`46`	`46`	`excludeGlobPatterns = append(excludeGlobPatterns, globExpression)`
`47`	`47`	`}`
`48`	`48`
	`49`	`+ var numResults int64`
	`50`	`+`
`49`	`51`	`for _, s3Path := range s3Lister.S3Paths {`
`50`	`52`	`bucket, key, err := aws.SplitS3Path(s3Path)`
`51`	`53`	`if err != nil {`
`52`		`- return err`
	`54`	`+ return 0, err`
`53`	`55`	`}`
`54`	`56`
`55`	`57`	`awsClientForBucket, err := aws.NewFromClientS3Path(s3Path, config.AWS)`
`56`	`58`	`if err != nil {`
`57`		`- return err`
	`59`	`+ return 0, err`
`58`	`60`	`}`
`59`	`61`
`60`	`62`	`err = awsClientForBucket.S3Iterator(bucket, key, false, nil, nil, func(s3Obj *s3.Object) (bool, error) {`
`@@ -79,15 +81,24 @@ func s3IteratorFromLister(s3Lister schema.S3Lister, fn func(string, *s3.Object)`
`79`	`81`	`}`
`80`	`82`
`81`	`83`	`if !shouldSkip {`
`82`		`- return fn(bucket, s3Obj)`
	`84`	`+ shouldContinue, err := fn(bucket, s3Obj)`
	`85`	`+ numResults++`
	`86`	`+ if s3Lister.MaxResults != nil && numResults >= *s3Lister.MaxResults {`
	`87`	`+ shouldContinue = false`
	`88`	`+ }`
	`89`	`+ return shouldContinue, err`
`83`	`90`	`}`
`84`	`91`
`85`	`92`	`return true, nil`
`86`	`93`	`})`
`87`	`94`	`if err != nil {`
`88`		`- return err`
	`95`	`+ return 0, err`
	`96`	`+ }`
	`97`	`+`
	`98`	`+ if s3Lister.MaxResults != nil && numResults >= *s3Lister.MaxResults {`
	`99`	`+ return numResults, nil`
`89`	`100`	`}`
`90`	`101`	`}`
`91`	`102`
`92`		`- return nil`
	`103`	`+ return numResults, nil`
`93`	`104`	`}`
Original file line number	Diff line number	Diff line change
`@@ -28,9 +28,10 @@ type ItemList struct {`
`28`	`28`	`}`
`29`	`29`
`30`	`30`	`type S3Lister struct {`
`31`		- S3Paths []string `json:"s3_paths"` // s3://<bucket_name>/key
`32`		- Includes []string `json:"includes"`
`33`		- Excludes []string `json:"excludes"`
	`31`	+ S3Paths []string `json:"s3_paths"` // s3://<bucket_name>/key
	`32`	+ Includes []string `json:"includes"`
	`33`	+ Excludes []string `json:"excludes"`
	`34`	+ MaxResults *int64 `json:"-"` // this is not currently exposed to the user (it's used for validations)
`34`	`35`	`}`
`35`	`36`
`36`	`37`	`type FilePathLister struct {`