[Simulator] Allow JobTemplates To Be Repeated (#4031)

* [Simulator] Allow JobTemplates To Be Repeated (#267) * wip * wip * wip * F/chrisma/simulator repeated submission (#269) * wip * wip * wip * lint * Update simulator.go --------- Co-authored-by: Christopher Martin <[email protected]>
armadaproject · Oct 31, 2024 · 77f8a0a · 77f8a0a
1 parent 133c490
commit 77f8a0a
Show file tree

Hide file tree

Showing 3 changed files with 420 additions and 98 deletions.
diff --git a/internal/scheduler/simulator/simulator.go b/internal/scheduler/simulator/simulator.go
@@ -10,6 +10,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/gogo/protobuf/proto"
 	"github.com/pkg/errors"
 	"golang.org/x/exp/maps"
 	"golang.org/x/exp/slices"
@@ -143,6 +144,7 @@ func NewSimulator(
 	if err := validateWorkloadSpec(workloadSpec); err != nil {
 		return nil, err
 	}
+	workloadSpec = expandRepeatingTemplates(workloadSpec)
 	jobDb := jobdb.NewJobDb(
 		schedulingConfig.PriorityClasses,
 		schedulingConfig.DefaultPriorityClassName,
@@ -292,6 +294,14 @@ func validateWorkloadSpec(workloadSpec *WorkloadSpec) error {
 			if template.GangCardinality != 0 && int(template.Number)%int(template.GangCardinality) != 0 {
 				return errors.Errorf("template.Number [%d] is not exactly divisible by template.GangCardinality [%d]", template.Number, template.GangCardinality)
 			}
+			if template.Repeat != nil {
+				if template.Repeat.Period == nil {
+					return errors.Errorf("template.Repeat.Period is unset")
+				}
+				if template.Repeat.NumTimes < 1 {
+					return errors.Errorf("template.Repeat.NumTimes must be greater than 0")
+				}
+			}
 		}
 	}
 	return nil
@@ -1035,3 +1045,26 @@ func (s *Simulator) removeJobFromDemand(job *jobdb.Job) {
 		r.SubV1ResourceList(job.PodRequirements().ResourceRequirements.Requests)
 	}
 }
+
+func expandRepeatingTemplates(w *WorkloadSpec) *WorkloadSpec {
+	workload := proto.Clone(w).(*WorkloadSpec)
+	for _, q := range workload.GetQueues() {
+		var templates []*JobTemplate
+		for _, template := range q.GetJobTemplates() {
+			if template.Repeat != nil {
+				period := *template.Repeat.Period
+				for i := 0; i < int(template.Repeat.NumTimes); i++ {
+					t := proto.Clone(template).(*JobTemplate)
+					t.Repeat = nil
+					t.Id = fmt.Sprintf("%s-repeat-%d", t.Id, i)
+					t.EarliestSubmitTime = t.EarliestSubmitTime + time.Duration(i)*period
+					templates = append(templates, t)
+				}
+			} else {
+				templates = append(templates, template)
+			}
+		}
+		q.JobTemplates = templates
+	}
+	return workload
+}