Added option to pass arbitrary OPT pass to Tetrisched (#96)

* Added option to pass arbitrary OPT pass

* Added sample config

configs/dagsched_dyn_arrival_0.025_0.013.conf

@@ -0,0 +1,49 @@
+# Output configuration.
+--log_dir=experiments/timing_check/final
+
+--log_file_name=alibaba_trace_replay.log
+--csv_file_name=alibaba_trace_replay.csv
+--log_level=debug
+
+# Worker configuration.
+--worker_profile_path=profiles/workers/alibaba_cluster_30k_slots.yaml
+
+# Workload configuration.
+--execution_mode=replay
+--replay_trace=alibaba
+--workload_profile_paths=traces/alibaba-cluster-trace-v2018/alibaba_filtered_new_alind_easy_dags_till_30k_cpu_usage.pkl,traces/alibaba-cluster-trace-v2018/medium_filtered.pkl,traces/alibaba-cluster-trace-v2018/hard_filtered.pkl
+--workload_profile_path_labels=easy,medium,hard
+--override_release_policies=poisson,poisson,poisson
+--override_num_invocations=0,350,650
+--override_poisson_arrival_rates=0.0075,0.013,0.025
+--randomize_start_time_max=50
+--min_deadline_variances=25,50,10
+--max_deadline_variances=50,100,25
+
+# Loader configuration.
+--alibaba_loader_task_cpu_usage_random
+--alibaba_loader_task_cpu_usage_min=120
+--alibaba_loader_task_cpu_usage_max=1500
+--alibaba_loader_task_cpu_multiplier=1
+--alibaba_loader_min_critical_path_runtimes=200,500,600
+--alibaba_loader_max_critical_path_runtimes=500,1000,1000
+
+# Scheduler configuration.
+--scheduler=TetriSched
+--release_taskgraphs
+--enforce_deadlines
+--scheduler_time_discretization=1
+# --scheduler_enable_optimization_pass
+--retract_schedules
+# --scheduler_dynamic_discretization 
+--opt_passes=CRITICAL_PATH_PASS
+--opt_passes=DYNAMIC_DISCRETIZATION_PASS
+--opt_passes=CAPACITY_CONSTRAINT_PURGE_PASS
+--scheduler_max_time_discretization=8
+--scheduler_max_occupancy_threshold=0.999
+--finer_discretization_at_prev_solution
+--finer_discretization_window=4
+--scheduler_selective_rescheduling
+--scheduler_reconsideration_period=0.6
+--scheduler_runtime=0
+--random_seed=42066545

main.py

@@ -64,6 +64,12 @@
     "Sets the mode in which the log file is opened. If 'append', the log file is "
     "opened in append mode, and if 'write', the log file is opened in write mode. ",
 )
+flags.DEFINE_multi_enum(
+    "opt_passes",
+    [], # default
+    ["CRITICAL_PATH_PASS", "DYNAMIC_DISCRETIZATION_PASS", "CAPACITY_CONSTRAINT_PURGE_PASS"], # choices
+    help="Specify the optimization passes that needs to be enabled once the STRL is generated, default: ['CRITICAL_PATH_PASS', 'CAPACITY_CONSTRAINT_PURGE_PASS']",
+)
 flags.DEFINE_string(
     "csv_file_name",
     None,
@@ -340,13 +346,13 @@
     "If `True`, the scheduler creates space-time matrix non-uniformly. "
     "The discretization is finer initially, and coarser at the end. (default: False)",
 )
-flags.DEFINE_bool(
-    "scheduler_dynamic_discretization",
-    False,
-    "If `True`, the scheduler creates space-time matrix non-uniformly. "
-    "The discretization is dynamically decided based on the occupancy request for "
-    "each time slice. (default: False)",
-)
+# flags.DEFINE_bool(
+#     "scheduler_dynamic_discretization",
+#     False,
+#     "If `True`, the scheduler creates space-time matrix non-uniformly. "
+#     "The discretization is dynamically decided based on the occupancy request for "
+#     "each time slice. (default: False)",
+# )
 flags.DEFINE_bool(
     "finer_discretization_at_prev_solution",
     False,
@@ -471,12 +477,12 @@
     "If `True`, the scheduler is allowed to batch tasks "
     "that share a WorkProfile together.",
 )
-flags.DEFINE_bool(
-    "scheduler_enable_optimization_pass",
-    False,
-    "If `True`, the scheduler runs pre/post-translation optimization passes"
-    "when registering STRL expression.",
-)
+# flags.DEFINE_bool(
+#     "scheduler_enable_optimization_pass",
+#     False,
+#     "If `True`, the scheduler runs pre/post-translation optimization passes"
+#     "when registering STRL expression.",
+# )
 flags.DEFINE_bool(
     "scheduler_selective_rescheduling",
     False,
@@ -819,7 +825,6 @@ def main(args):
             max_time_discretization=EventTime(
                 FLAGS.scheduler_max_time_discretization, EventTime.Unit.US
             ),
-            dynamic_discretization=FLAGS.scheduler_dynamic_discretization,
             max_occupancy_threshold=FLAGS.scheduler_max_occupancy_threshold,
             finer_discretization_at_prev_solution=finer_discretization,
             finer_discretization_window=EventTime(

schedulers/tetrisched/include/tetrisched/OptimizationPasses.hpp

@@ -17,6 +17,31 @@ enum OptimizationPassType {
   POST_TRANSLATION_PASS = 1,
 };
 
+enum OptimizationPassCategory {
+    CRITICAL_PATH_PASS = 0,
+    DYNAMIC_DISCRETIZATION_PASS = 1,
+    CAPACITY_CONSTRAINT_PURGE_PASS = 2,
+};
+
+/// The `OptimizationPassConfig` structure represents the configuration of the
+/// opt passes. This config is used to inform the choice of how to perform fidelity
+/// and non fidelity preserving OPT passes.
+struct OptimizationPassConfig {
+    /// configs for dynamic discretization
+    Time minDiscretization = 1;
+    Time maxDiscretization = 5;
+    float maxOccupancyThreshold = 0.8;
+    bool finerDiscretizationAtPrevSolution = false;
+    Time finerDiscretizationWindow = 5;
+    std::string toString() const {
+        std::stringstream ss;
+        ss << "{minDisc: " << minDiscretization << ", maxDisc: " << maxDiscretization << ", maxOccupancyThreshold: " << maxOccupancyThreshold << ", finerDiscretizationAtPrevSolution: " << finerDiscretizationAtPrevSolution << ", finerDiscretizationWindow: " << finerDiscretizationWindow << "\n";
+        return ss.str();
+    }
+};
+using OptimizationPassConfig = struct OptimizationPassConfig;
+using OptimizationPassConfigPtr = std::shared_ptr<OptimizationPassConfig>;
+
 /// An `OptimizationPass` is a base class for all Optimization passes that
 /// run on the STRL tree.
 class OptimizationPass {
@@ -159,17 +184,12 @@ class OptimizationPassRunner {
   bool debug;
   /// A list of optimization passes to run.
   std::vector<OptimizationPassPtr> registeredPasses;
-  /// flag for dynamic discretization optimization pass
-  bool enableDynamicDiscretization;
+
+  OptimizationPassConfigPtr optConfig;
 
- public:
+  public :
   /// Initialize the OptimizationPassRunner.
-  OptimizationPassRunner(bool debug = false,
-                         bool enableDynamicDiscretization = false,
-                         Time minDiscretization = 1, Time maxDiscretization = 5,
-                         float maxOccupancyThreshold = 0.8,
-                         bool finerDiscretizationAtPrevSolution = false,
-                         Time finerDiscretizationWindow = 5);
+  OptimizationPassRunner(bool debug = false, OptimizationPassConfigPtr optConfig = nullptr);
 
   /// Run the pre-translation optimization passes on the given STRL expression.
   void runPreTranslationPasses(Time currentTime, ExpressionPtr strlExpression,
@@ -178,6 +198,9 @@ class OptimizationPassRunner {
   /// Run the post-translation optimization passes on the given STRL expression.
   void runPostTranslationPasses(Time currentTime, ExpressionPtr strlExpression,
                                 CapacityConstraintMapPtr capacityConstraints);
+
+  /// Add the optimization passes
+  void addOptimizationPass(OptimizationPassCategory optPass);  
 };
 }  // namespace tetrisched
 #endif  // _TETRISCHED_OPTIMIZATION_PASSES_HPP_

schedulers/tetrisched/include/tetrisched/Scheduler.hpp

@@ -44,32 +44,32 @@ class Scheduler {
 
  public:
   /// Initialize the scheduler with a solver backend.
-  Scheduler(Time discretization, SolverBackendType solverBackend,
-            std::string logDir = "./", bool enableDynamicDiscretization = false,
-            Time maxDiscretization = 5, float maxOccupancyThreshold = 0.8,
-            bool finerDiscretizationAtPrevSolution = false,
-            Time finerDiscretizationWindow = 5);
+   Scheduler(Time discretization, SolverBackendType solverBackend,
+             std::string logDir = "./", OptimizationPassConfigPtr optConfig = nullptr);
 
-  /// Registers the STRL expression for the scheduler to schedule from
-  /// and parses it to populate the SolverModel.
-  void registerSTRL(
-      ExpressionPtr expression, Partitions availablePartitions,
-      Time currentTime, SchedulerConfigPtr schedulerConfig,
-      std::vector<std::pair<TimeRange, Time>> timeRangeToGranularities = {});
+   /// Add OptimizationPasses before STRL registration
+   void addOptimizationPass(OptimizationPassCategory optPass);
 
-  /// Invokes the solver to schedule the registered STRL expression
-  /// on the given partitions at the given time.
-  /// Use expression->getSolution() to retrieve the solution.
-  void schedule(Time currentTime);
+   /// Registers the STRL expression for the scheduler to schedule from
+   /// and parses it to populate the SolverModel.
+   void registerSTRL(
+       ExpressionPtr expression, Partitions availablePartitions,
+       Time currentTime, SchedulerConfigPtr schedulerConfig,
+       std::vector<std::pair<TimeRange, Time>> timeRangeToGranularities = {});
 
-  /// Retrieve the solution from the last invocation of the solver.
-  SolverSolutionPtr getLastSolverSolution() const;
+   /// Invokes the solver to schedule the registered STRL expression
+   /// on the given partitions at the given time.
+   /// Use expression->getSolution() to retrieve the solution.
+   void schedule(Time currentTime);
 
-  /// Exports the model from the last invocation of the solver.
-  void exportLastSolverModel(const std::string &fileName) const;
+   /// Retrieve the solution from the last invocation of the solver.
+   SolverSolutionPtr getLastSolverSolution() const;
 
-  /// Exports the solution from the last invocation of the solver.
-  void exportLastSolverSolution(const std::string &fileName) const;
+   /// Exports the model from the last invocation of the solver.
+   void exportLastSolverModel(const std::string &fileName) const;
+
+   /// Exports the solution from the last invocation of the solver.
+   void exportLastSolverSolution(const std::string &fileName) const;
 };
 }  // namespace tetrisched
 

schedulers/tetrisched/python/TetrischedPy.cpp

@@ -8,6 +8,12 @@
 namespace py = pybind11;
 
 void defineBasicTypes(py::module_& tetrisched_m) {
+  // define OPT pass enum
+  py::enum_<tetrisched::OptimizationPassCategory>(tetrisched_m, "OptimizationPassCategory")
+      .value("CRITICAL_PATH_PASS", tetrisched::OptimizationPassCategory::CRITICAL_PATH_PASS)
+      .value("DYNAMIC_DISCRETIZATION_PASS", tetrisched::OptimizationPassCategory::DYNAMIC_DISCRETIZATION_PASS)
+      .value("CAPACITY_CONSTRAINT_PURGE_PASS", tetrisched::OptimizationPassCategory::CAPACITY_CONSTRAINT_PURGE_PASS)
+      .export_values();
   // Define the Partition type.
   py::class_<tetrisched::Partition, tetrisched::PartitionPtr>(
       tetrisched_m, "Partition", py::dynamic_attr())
@@ -54,6 +60,20 @@ void defineBasicTypes(py::module_& tetrisched_m) {
 
 /// Define the Scheduler interface.
 void defineScheduler(py::module_& tetrisched_m) {
+  // Define the Config for Optimization Passes
+  py::class_<tetrisched::OptimizationPassConfig, tetrisched::OptimizationPassConfigPtr>(tetrisched_m, "OptimizationPassConfig")
+      .def(py::init<>(), "Initializes an empty OptimizationPassConfig.")
+      .def_readwrite("minDiscretization", &tetrisched::OptimizationPassConfig::minDiscretization,
+                     "the min discretization for Dynamic Discretiazation OPT pass.")
+      .def_readwrite("maxDiscretization", &tetrisched::OptimizationPassConfig::maxDiscretization,
+                     "the max discretization for Dynamic Discretiazation OPT pass.")
+      .def_readwrite("maxOccupancyThreshold", &tetrisched::OptimizationPassConfig::maxOccupancyThreshold,
+                     "the max occupancy threshold beyond which dynamic discretization is always min.")
+      .def_readwrite("finerDiscretizationAtPrevSolution", &tetrisched::OptimizationPassConfig::finerDiscretizationAtPrevSolution,
+                     "Whether to enabled finer discretization at solved solutions.")
+      .def_readwrite("finerDiscretizationWindow", &tetrisched::OptimizationPassConfig::finerDiscretizationWindow,
+                     "The window upto which finer discretization should be enabled around previously solved solutions.")
+      .def("toString", &tetrisched::OptimizationPassConfig::toString, "Print String Representation");
   // Define the Config for the Scheduler.
   py::class_<tetrisched::SchedulerConfig, tetrisched::SchedulerConfigPtr>(
       tetrisched_m, "SchedulerConfig")
@@ -73,33 +93,18 @@ void defineScheduler(py::module_& tetrisched_m) {
   // Define the interface to the Scheduler.
   py::class_<tetrisched::Scheduler>(tetrisched_m, "Scheduler")
       .def(
-          py::init<tetrisched::Time, tetrisched::SolverBackendType, std::string,
-                   bool, tetrisched::Time, float, bool, tetrisched::Time>(),
+          py::init<tetrisched::Time, tetrisched::SolverBackendType, std::string, tetrisched::OptimizationPassConfigPtr>(),
           "Initializes the Scheduler with the given backend.\n"
           "\nArgs:\n"
           "  discretization (int): The time discretization to use for the "
           "scheduler.\n"
           "  solverBackend (SolverBackendType): The solver backend to use for "
           "the scheduler.\n"
           "  logDir (str): The directory where the logs are to be output.\n"
-          "  enableDynamicDiscretization (bool): Whether to enable dynamic "
-          "discretization.\n"
-          "  maxDiscretization (int): The maximum discretization to use for "
-          "dynamic discretization.\n"
-          "  maxOccupancyThreshold (float): The maximum occupancy threshold "
-          "to use for dynamic discretization.\n",
-          "  finerDiscretizationAtPrevSolution (bool): Enables finer "
-          "discretization "
-          "At previous solution.\n",
-          "  finerDiscretizationWindow (int): The discretization around prev "
-          "solution until which the discretization would be 1 ",
+         "OptimizationPass Config for OPT passes",
           py::arg("discretization"), py::arg("solverBackend"),
           py::arg("logDir") = "./",
-          py::arg("enableDynamicDiscretization") = false,
-          py::arg("maxDiscretization") = 5,
-          py::arg("maxOccupancyThreshold") = 0.8,
-          py::arg("finerDiscretizationAtPrevSolution") = false,
-          py::arg("finerDiscretizationWindow") = 5)
+          py::arg("optConfig") = nullptr)
       .def(
           "registerSTRL", &tetrisched::Scheduler::registerSTRL,
           "Registers the STRL expression for the scheduler to schedule from.\n"
@@ -116,6 +121,12 @@ void defineScheduler(py::module_& tetrisched_m) {
           py::arg("currentTime"), py::arg("schedulerConfig"),
           py::arg("timeRangeToGranularities") =
               std::vector<std::pair<tetrisched::TimeRange, tetrisched::Time>>())
+      .def(
+          "addOptimizationPass", &tetrisched::Scheduler::addOptimizationPass,
+          "Add Optimization passes to Tetrisched Based Declarative Scheduler. \n"
+          "\nArgs:\n"
+          "  optPass (OptimizationPassCategory): The Optimization Pass that will get enabled.\n",
+          py::arg("optPass"))
       .def("schedule", &tetrisched::Scheduler::schedule,
            "Invokes the solver to schedule the registered STRL expression.\n"
            "\nArgs:\n"

schedulers/tetrisched/src/OptimizationPasses.cpp

@@ -1129,24 +1129,48 @@ void CapacityConstraintMapPurgingOptimizationPass::clean() { cliques.clear(); }
 
 /* Methods for OptimizationPassRunner */
 OptimizationPassRunner::OptimizationPassRunner(
-    bool debug, bool enableDynamicDiscretization, Time minDiscretization,
-    Time maxDiscretization, float maxOccupancyThreshold,
-    bool finerDiscretizationAtPrevSolution, Time finerDiscretizationWindow)
-    : debug(debug), enableDynamicDiscretization(enableDynamicDiscretization) {
+    bool debug, OptimizationPassConfigPtr optConfig)
+    : debug(debug), optConfig(optConfig)
+{
   // Register the Critical Path optimization pass.
-  registeredPasses.push_back(std::make_shared<CriticalPathOptimizationPass>());
-
-  if (enableDynamicDiscretization) {
-    // Register the DiscretizationGenerator pass.
-    registeredPasses.push_back(
-        std::make_shared<DiscretizationSelectorOptimizationPass>(
-            minDiscretization, maxDiscretization, maxOccupancyThreshold,
-            finerDiscretizationAtPrevSolution, finerDiscretizationWindow));
-  }
+  // registeredPasses.push_back(std::make_shared<CriticalPathOptimizationPass>());
+
+  // if (enableDynamicDiscretization) {
+  //   // Register the DiscretizationGenerator pass.
+  //   registeredPasses.push_back(
+  //       std::make_shared<DiscretizationSelectorOptimizationPass>(
+  //           minDiscretization, maxDiscretization, maxOccupancyThreshold,
+  //           finerDiscretizationAtPrevSolution, finerDiscretizationWindow));
+  // }
+
+  // // Register the CapacityConstraintMapPurging optimization pass.
+  // registeredPasses.push_back(
+  //     std::make_shared<CapacityConstraintMapPurgingOptimizationPass>());
+}
 
-  // Register the CapacityConstraintMapPurging optimization pass.
-  registeredPasses.push_back(
-      std::make_shared<CapacityConstraintMapPurgingOptimizationPass>());
+void OptimizationPassRunner::addOptimizationPass(OptimizationPassCategory optPass) {
+  // Register the Critical Path optimization pass.
+
+  switch (optPass)
+  {
+    case OptimizationPassCategory::CRITICAL_PATH_PASS:
+      registeredPasses.push_back(std::make_shared<CriticalPathOptimizationPass>());
+      break;
+    case OptimizationPassCategory::DYNAMIC_DISCRETIZATION_PASS:
+      if (optConfig == nullptr) {
+        optConfig = std::make_shared<OptimizationPassConfig>();
+      }
+      registeredPasses.push_back(std::make_shared<DiscretizationSelectorOptimizationPass>(
+          optConfig->minDiscretization, optConfig->maxDiscretization, optConfig->maxOccupancyThreshold,
+          optConfig->finerDiscretizationAtPrevSolution, optConfig->finerDiscretizationWindow));
+      break;
+    case OptimizationPassCategory::CAPACITY_CONSTRAINT_PURGE_PASS:
+      registeredPasses.push_back(std::make_shared<CapacityConstraintMapPurgingOptimizationPass>());
+      break;
+
+    default:
+      break;
+  }
 }
 
 void OptimizationPassRunner::runPreTranslationPasses(