Get job status from batch crd for in progress jobs (#2146)

Author: vishalbollu

go.mod

@@ -33,6 +33,7 @@ require (
 	github.com/onsi/gomega v1.10.2
 	github.com/opencontainers/go-digest v1.0.0 // indirect
 	github.com/opencontainers/image-spec v1.0.1 // indirect
+	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.7.1
 	github.com/prometheus/common v0.10.0

go.sum

@@ -507,6 +507,9 @@ github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3I
 github.com/opencontainers/image-spec v1.0.1 h1:JMemWkRwHx4Zj+fVxWoMCFm/8sYGGrUVojFA6h/TRcI=
 github.com/opencontainers/image-spec v1.0.1/go.mod h1:BtxoFyWECRxE4U/7sNtV5W15zMzWCbyJoFRP3s7yZA0=
 github.com/pascaldekloe/goe v0.0.0-20180627143212-57f6aae5913c/go.mod h1:lzWF7FIEvWOWxwDKqyGYQf6ZUaNfKdP144TG7ZOy1lc=
+github.com/patrickmn/go-cache v1.0.0 h1:3gD5McaYs9CxjyK5AXGcq8gdeCARtd/9gJDUvVeaZ0Y=
+github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
+github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/pborman/uuid v1.2.0/go.mod h1:X/NO0urCmaxf9VXbdlT7C2Yzkj2IKimNn4k+gtPdI/k=
 github.com/pelletier/go-toml v1.2.0/go.mod h1:5z9KED0ma1S8pY6P1sdut58dfprrGBbd/94hg7ilaic=
 github.com/peterbourgon/diskv v2.0.1+incompatible/go.mod h1:uqqh8zWWbv1HBMNONnaR/tNboyR3/BZd58JJSHlUSCU=

manager/manifests/prometheus-monitoring.yaml.j2

@@ -224,7 +224,7 @@ spec:
     - port: metrics
       scheme: http
       path: /metrics
-      interval: 30s
+      interval: 20s
   namespaceSelector:
     any: true
   selector:

pkg/crds/apis/batch/v1alpha1/batchjob_types.go

@@ -83,15 +83,15 @@ type BatchJobStatus struct {
 	// Job ID
 	ID string `json:"id,omitempty"`
 
-	// Processing start timestamp
-	StartTime *kmeta.Time `json:"start_time,omitempty"`
-
 	// Processing ending timestamp
 	EndTime *kmeta.Time `json:"end_time,omitempty"`
 
 	// URL for the used SQS queue
 	QueueURL string `json:"queue_url,omitempty"`
 
+	// Total batch count
+	TotalBatchCount int `json:"total_batch_count,omitempty"`
+
 	// +kubebuilder:validation:Type=string
 	// Status of the batch job
 	Status status.JobCode `json:"status,omitempty"`

pkg/crds/apis/batch/v1alpha1/zz_generated.deepcopy.go

@@ -125,13 +125,12 @@ spec:
               queue_url:
                 description: URL for the used SQS queue
                 type: string
-              start_time:
-                description: Processing start timestamp
-                format: date-time
-                type: string
               status:
                 description: Status of the batch job
                 type: string
+              total_batch_count:
+                description: Total batch count
+                type: integer
               worker_counts:
                 description: Detailed worker counts with respective status
                 properties:

pkg/crds/config/crd/bases/batch.cortex.dev_batchjobs.yaml

@@ -6,6 +6,14 @@ metadata:
   creationTimestamp: null
   name: manager-role
 rules:
+- apiGroups:
+  - ""
+  resources:
+  - pods
+  verbs:
+  - get
+  - list
+  - watch
 - apiGroups:
   - batch
   resources:

pkg/crds/config/rbac/role.yaml

@@ -56,6 +56,7 @@ type BatchJobReconciler struct {
 // +kubebuilder:rbac:groups=batch.cortex.dev,resources=batchjobs/status,verbs=get;update;patch
 // +kubebuilder:rbac:groups=batch.cortex.dev,resources=batchjobs/finalizers,verbs=update
 // +kubebuilder:rbac:groups=batch,resources=jobs,verbs=get;list;watch;create;update;patch
+// +kubebuilder:rbac:groups="",resources=pods,verbs=get;list;watch
 
 // Reconcile is part of the main kubernetes reconciliation loop which aims to
 // move the current state of the cluster closer to the desired state.
@@ -150,12 +151,18 @@ func (r *BatchJobReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 		return ctrl.Result{}, err
 	}
 
+	var totalBatchCount int
+	if enqueuingStatus == batch.EnqueuingDone {
+		totalBatchCount, err = r.Config.GetTotalBatchCount(r, batchJob)
+	}
+
 	workerJobExists := workerJob != nil
 	statusInfo := batchJobStatusInfo{
 		QueueExists:     queueExists,
 		EnqueuingStatus: enqueuingStatus,
 		EnqueuerJob:     enqueuerJob,
 		WorkerJob:       workerJob,
+		TotalBatchCount: totalBatchCount,
 	}
 
 	log.V(1).Info("status data successfully acquired",

pkg/crds/controllers/batch/batchjob_controller.go

@@ -23,16 +23,16 @@ import (
 
 // BatchJobReconcilerConfig reconciler config for the BatchJob kind. Allows for mocking specific methods
 type BatchJobReconcilerConfig struct {
-	GetMaxBatchCount func(r *BatchJobReconciler, batchJob batch.BatchJob) (int, error)
-	GetMetrics       func(r *BatchJobReconciler, batchJob batch.BatchJob) (metrics.BatchMetrics, error)
-	SaveJobMetrics   func(r *BatchJobReconciler, batchJob batch.BatchJob) error
-	SaveJobStatus    func(r *BatchJobReconciler, batchJob batch.BatchJob) error
+	GetTotalBatchCount func(r *BatchJobReconciler, batchJob batch.BatchJob) (int, error)
+	GetMetrics         func(r *BatchJobReconciler, batchJob batch.BatchJob) (metrics.BatchMetrics, error)
+	SaveJobMetrics     func(r *BatchJobReconciler, batchJob batch.BatchJob) error
+	SaveJobStatus      func(r *BatchJobReconciler, batchJob batch.BatchJob) error
 }
 
 // ApplyDefaults sets the defaults for BatchJobReconcilerConfig
 func (c BatchJobReconcilerConfig) ApplyDefaults() BatchJobReconcilerConfig {
-	if c.GetMaxBatchCount == nil {
-		c.GetMaxBatchCount = getMaxBatchCount
+	if c.GetTotalBatchCount == nil {
+		c.GetTotalBatchCount = getTotalBatchCount
 	}
 
 	if c.GetMetrics == nil {

pkg/crds/controllers/batch/batchjob_controller_config.go

@@ -42,6 +42,7 @@ import (
 	"github.com/cortexlabs/cortex/pkg/types/userconfig"
 	"github.com/cortexlabs/cortex/pkg/workloads"
 	"github.com/cortexlabs/yaml"
+	cache "github.com/patrickmn/go-cache"
 	kbatch "k8s.io/api/batch/v1"
 	kcore "k8s.io/api/core/v1"
 	kerrors "k8s.io/apimachinery/pkg/api/errors"
@@ -52,13 +53,21 @@ import (
 const (
 	_enqueuerContainerName  = "enqueuer"
 	_deadlineExceededReason = "DeadlineExceeded"
+	_cacheDuration          = 60 * time.Second
 )
 
+var totalBatchCountCache *cache.Cache
+
+func init() {
+	totalBatchCountCache = cache.New(_cacheDuration, _cacheDuration)
+}
+
 type batchJobStatusInfo struct {
 	QueueExists     bool
 	EnqueuingStatus batch.EnqueuingStatus
 	EnqueuerJob     *kbatch.Job
 	WorkerJob       *kbatch.Job
+	TotalBatchCount int
 }
 
 func (r *BatchJobReconciler) checkIfQueueExists(batchJob batch.BatchJob) (bool, error) {
@@ -381,11 +390,12 @@ func (r *BatchJobReconciler) updateStatus(ctx context.Context, batchJob *batch.B
 	case batch.EnqueuingFailed:
 		batchJob.Status.Status = status.JobEnqueueFailed
 		batchJob.Status.EndTime = statusInfo.EnqueuerJob.Status.CompletionTime
+	case batch.EnqueuingDone:
+		batchJob.Status.TotalBatchCount = statusInfo.TotalBatchCount
 	}
 
 	worker := statusInfo.WorkerJob
 	if worker != nil {
-		batchJob.Status.StartTime = worker.Status.StartTime    // assign right away, because it's a pointer
 		batchJob.Status.EndTime = worker.Status.CompletionTime // assign right away, because it's a pointer
 
 		if worker.Status.Failed == batchJob.Spec.Workers {
@@ -423,9 +433,7 @@ func (r *BatchJobReconciler) updateStatus(ctx context.Context, batchJob *batch.B
 			batchJob.Status.Status = status.JobRunning
 		}
 
-		pendingWorkers := batchJob.Spec.Workers - (worker.Status.Active + worker.Status.Succeeded + worker.Status.Failed)
 		batchJob.Status.WorkerCounts = &status.WorkerCounts{
-			Pending:   pendingWorkers,
 			Running:   worker.Status.Active,
 			Succeeded: worker.Status.Succeeded,
 			Failed:    worker.Status.Failed,
@@ -495,7 +503,7 @@ func (r *BatchJobReconciler) uploadJobSpec(batchJob batch.BatchJob, api spec.API
 		timeout = pointer.Int(int(batchJob.Spec.Timeout.Seconds()))
 	}
 
-	maxBatchCount, err := r.Config.GetMaxBatchCount(r, batchJob)
+	totalBatchCount, err := r.Config.GetTotalBatchCount(r, batchJob)
 	if err != nil {
 		return nil, err
 	}
@@ -513,11 +521,9 @@ func (r *BatchJobReconciler) uploadJobSpec(batchJob batch.BatchJob, api spec.API
 			Config:             config,
 		},
 		APIID:           api.ID,
-		SpecID:          api.SpecID,
-		HandlerID:       api.HandlerID,
 		SQSUrl:          queueURL,
 		StartTime:       batchJob.CreationTimestamp.Time,
-		TotalBatchCount: maxBatchCount,
+		TotalBatchCount: totalBatchCount,
 	}
 
 	if err = r.AWS.UploadJSONToS3(&jobSpec, r.ClusterConfig.Bucket, r.jobSpecKey(batchJob)); err != nil {
@@ -569,19 +575,27 @@ func (r *BatchJobReconciler) persistJobToS3(batchJob batch.BatchJob) error {
 	)
 }
 
-func getMaxBatchCount(r *BatchJobReconciler, batchJob batch.BatchJob) (int, error) {
+func getTotalBatchCount(r *BatchJobReconciler, batchJob batch.BatchJob) (int, error) {
 	key := spec.JobBatchCountKey(r.ClusterConfig.ClusterUID, userconfig.BatchAPIKind, batchJob.Spec.APIName, batchJob.Name)
-	maxBatchCountBytes, err := r.AWS.ReadBytesFromS3(r.ClusterConfig.Bucket, key)
-	if err != nil {
-		return 0, err
-	}
+	cachedTotalBatchCount, found := totalBatchCountCache.Get(key)
+	var totalBatchCount int
+	if !found {
+		totalBatchCountBytes, err := r.AWS.ReadBytesFromS3(r.ClusterConfig.Bucket, key)
+		if err != nil {
+			return 0, err
+		}
 
-	maxBatchCount, err := strconv.Atoi(string(maxBatchCountBytes))
-	if err != nil {
-		return 0, err
+		totalBatchCount, err = strconv.Atoi(string(totalBatchCountBytes))
+		if err != nil {
+			return 0, err
+		}
+	} else {
+		totalBatchCount = cachedTotalBatchCount.(int)
 	}
 
-	return maxBatchCount, nil
+	totalBatchCountCache.Set(key, totalBatchCount, _cacheDuration)
+
+	return totalBatchCount, nil
 }
 
 func getMetrics(r *BatchJobReconciler, batchJob batch.BatchJob) (metrics.BatchMetrics, error) {
@@ -611,14 +625,24 @@ func saveJobMetrics(r *BatchJobReconciler, batchJob batch.BatchJob) error {
 }
 
 func saveJobStatus(r *BatchJobReconciler, batchJob batch.BatchJob) error {
-	jobStatus := batchJob.Status.Status.String()
-	key := filepath.Join(
-		spec.JobAPIPrefix(r.ClusterConfig.ClusterUID, userconfig.BatchAPIKind, batchJob.Spec.APIName),
-		batchJob.Name,
-		jobStatus,
+	return parallel.RunFirstErr(
+		func() error {
+			stoppedStatusKey := filepath.Join(
+				spec.JobAPIPrefix(r.ClusterConfig.ClusterUID, userconfig.BatchAPIKind, batchJob.Spec.APIName),
+				batchJob.Name,
+				status.JobStopped.String(),
+			)
+			return r.AWS.UploadStringToS3("", r.ClusterConfig.Bucket, stoppedStatusKey)
+
+		},
+		func() error {
+			jobStatus := batchJob.Status.Status.String()
+			key := filepath.Join(
+				spec.JobAPIPrefix(r.ClusterConfig.ClusterUID, userconfig.BatchAPIKind, batchJob.Spec.APIName),
+				batchJob.Name,
+				jobStatus,
+			)
+			return r.AWS.UploadStringToS3("", r.ClusterConfig.Bucket, key)
+		},
 	)
-	if err := r.AWS.UploadStringToS3("", r.ClusterConfig.Bucket, key); err != nil {
-		return err
-	}
-	return nil
 }

pkg/crds/controllers/batch/batchjob_controller_helpers.go

@@ -124,7 +124,7 @@ var _ = BeforeSuite(func(done Done) {
 
 	// mock certain methods of the reconciler
 	reconcilerConfig := batchcontrollers.BatchJobReconcilerConfig{
-		GetMaxBatchCount: func(r *batchcontrollers.BatchJobReconciler, batchJob batch.BatchJob) (int, error) {
+		GetTotalBatchCount: func(r *batchcontrollers.BatchJobReconciler, batchJob batch.BatchJob) (int, error) {
 			return 1, nil
 		},
 		GetMetrics: func(r *batchcontrollers.BatchJobReconciler, batchJob batch.BatchJob) (metrics.BatchMetrics, error) {

pkg/crds/controllers/batch/suite_test.go

@@ -123,8 +123,10 @@ func GetAllAPIs(virtualServices []istioclientnetworking.VirtualService, batchJob
 	batchAPIsMap := map[string]*schema.APIResponse{}
 
 	jobIDToBatchJobMap := map[string]*batch.BatchJob{}
+	apiNameToBatchJobsMap := map[string][]*batch.BatchJob{}
 	for i, batchJob := range batchJobList {
 		jobIDToBatchJobMap[batchJob.Name] = &batchJobList[i]
+		apiNameToBatchJobsMap[batchJob.Spec.APIName] = append(apiNameToBatchJobsMap[batchJob.Spec.APIName], &batchJobList[i])
 	}
 
 	for _, virtualService := range virtualServices {
@@ -141,19 +143,33 @@ func GetAllAPIs(virtualServices []istioclientnetworking.VirtualService, batchJob
 			return nil, err
 		}
 
-		jobStates, err := job.GetMostRecentlySubmittedJobStates(apiName, 1, userconfig.BatchAPIKind)
-		if err != nil {
-			return nil, err
-		}
-
 		var jobStatuses []status.BatchJobStatus
-		if len(jobStates) > 0 {
-			jobStatus, err := getJobStatusFromJobState(jobStates[0], jobIDToBatchJobMap[jobStates[0].ID])
+		batchJobs := apiNameToBatchJobsMap[apiName]
+
+		if len(batchJobs) == 0 {
+			jobStates, err := job.GetMostRecentlySubmittedJobStates(apiName, 1, userconfig.BatchAPIKind)
 			if err != nil {
 				return nil, err
 			}
 
-			jobStatuses = append(jobStatuses, *jobStatus)
+			if len(jobStates) > 0 {
+				jobStatus, err := getJobStatusFromJobState(jobStates[0])
+				if err != nil {
+					return nil, err
+				}
+
+				jobStatuses = append(jobStatuses, *jobStatus)
+			}
+		} else {
+			for i := range batchJobs {
+				batchJob := batchJobs[i]
+				jobStatus, err := getJobStatusFromBatchJob(*batchJob)
+				if err != nil {
+					return nil, err
+				}
+
+				jobStatuses = append(jobStatuses, *jobStatus)
+			}
 		}
 
 		batchAPIsMap[apiName] = &schema.APIResponse{
@@ -199,7 +215,7 @@ func GetAPIByName(deployedResource *operator.DeployedResource) ([]schema.APIResp
 	var jobStatuses []status.BatchJobStatus
 	jobIDSet := strset.New()
 	for _, batchJob := range batchJobList.Items {
-		jobStatus, err := getJobStatusFromK8sBatchJob(batchJob)
+		jobStatus, err := getJobStatusFromBatchJob(batchJob)
 		if err != nil {
 			return nil, err
 		}
@@ -218,7 +234,7 @@ func GetAPIByName(deployedResource *operator.DeployedResource) ([]schema.APIResp
 			}
 			jobIDSet.Add(jobState.ID)
 
-			jobStatus, err := getJobStatusFromJobState(jobState, nil)
+			jobStatus, err := getJobStatusFromJobState(jobState)
 			if err != nil {
 				return nil, err
 			}