[WIP] Some improvements to the simulator

cmd/simulator/cmd/root.go

@@ -3,8 +3,10 @@ package cmd
 import (
 	"math"
 	"os"
+	"runtime/pprof"
 
 	"github.com/pkg/errors"
+	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"golang.org/x/exp/maps"
 
@@ -28,6 +30,7 @@ func RootCmd() *cobra.Command {
 	cmd.Flags().Bool("showSchedulerLogs", false, "Show scheduler logs.")
 	cmd.Flags().Int("logInterval", 0, "Log summary statistics every this many events. Disabled if 0.")
 	cmd.Flags().String("eventsOutputFilePath", "", "Path of file to write events to.")
+	cmd.Flags().String("cycleStatsOutputFilePath", "", "Path of file to write cycle stats to.")
 	cmd.Flags().Bool("enableFastForward", false, "Skips schedule events when we're in a steady state")
 	cmd.Flags().Int("hardTerminationMinutes", math.MaxInt, "Limit the time simulated")
 	cmd.Flags().Int("schedulerCyclePeriodSeconds", 10, "How often we should trigger schedule events")
@@ -56,7 +59,11 @@ func runSimulations(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return err
 	}
-	filePath, err := cmd.Flags().GetString("eventsOutputFilePath")
+	eventsOutputFilePath, err := cmd.Flags().GetString("eventsOutputFilePath")
+	if err != nil {
+		return err
+	}
+	cycleStatsOutputFilePath, err := cmd.Flags().GetString("cycleStatsOutputFilePath")
 	if err != nil {
 		return err
 	}
@@ -95,8 +102,8 @@ func runSimulations(cmd *cobra.Command, args []string) error {
 			return err
 		}
 	}
-	if len(clusterSpecs)*len(workloadSpecs)*len(schedulingConfigsByFilePath) > 1 && filePath != "" {
-		return errors.Errorf("cannot save multiple simulations to file")
+	if len(clusterSpecs)*len(workloadSpecs)*len(schedulingConfigsByFilePath) > 1 && eventsOutputFilePath != "" {
+		return errors.Errorf("cannot save multiple simulations to eventsOutputFile")
 	}
 
 	ctx := armadacontext.Background()
@@ -106,13 +113,25 @@ func runSimulations(cmd *cobra.Command, args []string) error {
 	ctx.Infof("SchedulingConfigs: %v", maps.Keys(schedulingConfigsByFilePath))
 
 	var fileWriter *simulator.Writer
-	file, err := os.Create(filePath)
+	eventsOutputFile, err := os.Create(eventsOutputFilePath)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		if err = eventsOutputFile.Close(); err != nil {
+			ctx.Errorf("failed to close eventsOutputFile: %s", err)
+			return
+		}
+	}()
+
+	var statsWriter *simulator.StatsWriter
+	cycleStatsOutputFile, err := os.Create(cycleStatsOutputFilePath)
 	if err != nil {
 		return err
 	}
 	defer func() {
-		if err = file.Close(); err != nil {
-			ctx.Errorf("failed to close file: %s", err)
+		if err = cycleStatsOutputFile.Close(); err != nil {
+			ctx.Errorf("failed to close cycleStatsOutputFile: %s", err)
 			return
 		}
 	}()
@@ -138,19 +157,35 @@ func runSimulations(cmd *cobra.Command, args []string) error {
 					mc.LogSummaryInterval = logInterval
 					metricsCollectors = append(metricsCollectors, mc)
 
-					if filePath != "" {
-						fw, err := simulator.NewWriter(file, s.StateTransitions())
+					if eventsOutputFilePath != "" {
+						fw, err := simulator.NewWriter(eventsOutputFile, s.StateTransitions())
 						if err != nil {
 							return errors.WithStack(err)
 						}
 						fileWriter = fw
 					}
+					if cycleStatsOutputFilePath != "" {
+						sw, err := simulator.NewStatsWriter(cycleStatsOutputFile, s.CycleMetrics())
+						if err != nil {
+							return errors.WithStack(err)
+						}
+						statsWriter = sw
+					}
 					stateTransitionChannels = append(stateTransitionChannels, s.StateTransitions())
 					schedulingConfigPaths = append(schedulingConfigPaths, schedulingConfigPath)
 				}
 			}
 		}
 	}
+	f, err := os.Create("profile")
+	if err != nil {
+		log.Fatal(err)
+	}
+	err = pprof.StartCPUProfile(f)
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer pprof.StopCPUProfile()
 
 	// Run simulators.
 	g, ctx := armadacontext.ErrGroup(ctx)
@@ -179,11 +214,16 @@ func runSimulations(cmd *cobra.Command, args []string) error {
 		})
 	}
 
-	// Run file writer
+	// Run eventsOutputFile writer
 	g.Go(func() error {
 		return fileWriter.Run(ctx)
 	})
 
+	// Run stats writer
+	g.Go(func() error {
+		return statsWriter.Run(ctx)
+	})
+
 	// Run metric collectors.
 	for _, mc := range metricsCollectors {
 		mc := mc

internal/scheduler/simulator/cycle_stats.go

@@ -0,0 +1,11 @@
+package simulator
+
+type CycleStats struct {
+	Time          int64   `parquet:"name=time, type=INT64"`
+	Queue         string  `parquet:"name=queue, type=BYTE_ARRAY, convertedtype=UTF8, encoding=PLAIN_DICTIONARY"`
+	Pool          string  `parquet:"name=pool, type=BYTE_ARRAY, convertedtype=UTF8, encoding=PLAIN_DICTIONARY"`
+	PriorityClass string  `parquet:"name=priority_class, type=BYTE_ARRAY, convertedtype=UTF8, encoding=PLAIN_DICTIONARY"`
+	Cpu           float64 `parquet:"name=cpu, type=DOUBLE"`
+	Memory        float64 `parquet:"name=memory, type=DOUBLE"`
+	Gpu           float64 `parquet:"name=gpu, type=DOUBLE"`
+}

internal/scheduler/simulator/metrics.go

@@ -19,6 +19,7 @@ type MetricsCollector struct {
 	MetricsByQueue map[string]MetricsVector
 	// If non-zero, log a summary every this many events.
 	LogSummaryInterval int
+	BaseTime           time.Time
 }
 
 type MetricsVector struct {
@@ -38,6 +39,7 @@ func NewMetricsCollector(c <-chan StateTransition) *MetricsCollector {
 	return &MetricsCollector{
 		c:              c,
 		MetricsByQueue: make(map[string]MetricsVector),
+		BaseTime:       time.Unix(0, 0),
 	}
 }
 
@@ -94,7 +96,7 @@ func (mc *MetricsCollector) addEventSequence(eventSequence *armadaevents.EventSe
 	perQueueMetrics := mc.MetricsByQueue[queue]
 	perQueueMetrics.NumEvents += 1
 	for _, event := range eventSequence.Events {
-		d := protoutil.ToStdTime(event.Created).Sub(time.Time{})
+		d := protoutil.ToStdTime(event.Created).Sub(mc.BaseTime)
 		mc.OverallMetrics.TimeOfMostRecentEvent = d
 		perQueueMetrics.TimeOfMostRecentEvent = d
 		switch event.GetEvent().(type) {