Fixing correct flow of Placement solutions.

schedulers/tetrisched/include/tetrisched/Expression.hpp

@@ -3,6 +3,7 @@
 
 #include <functional>
 #include <optional>
+#include <set>
 #include <unordered_map>
 
 #include "tetrisched/Partition.hpp"
@@ -65,13 +66,19 @@ class Placement {
   bool placed;
   /// The start time of the Placement.
   std::optional<Time> startTime;
-  /// The ID of the Partition that the Task is placed on.
-  std::unordered_map<uint32_t, TETRISCHED_ILP_TYPE> partitionToResources;
+  /// The end time of the Placement.
+  std::optional<Time> endTime;
+  /// A <PartitionID, <Time, Allocation>> vector that represents the
+  /// allocation of resources from each Partition to this Placement at a
+  /// particular time. Note that the last Partition assignments are valid
+  /// until the end of the Placement.
+  std::unordered_map<uint32_t, std::set<std::pair<Time, uint32_t>>>
+      partitionToResourceAllocations;
 
  public:
   /// Initialize a Placement with the given Task name, start time and Partition
   /// ID.
-  Placement(std::string taskName, Time startTime);
+  Placement(std::string taskName, Time startTime, Time endTime);
 
   /// Initialize a Placement with the given Task name signifying that the Task
   /// was not actually placed.
@@ -80,24 +87,22 @@ class Placement {
   /// Check if the Task was actually placed.
   bool isPlaced() const;
 
-  /// Add a Partition along with the resources it is contributing to this
-  /// Placement.
-  void addPartition(uint32_t partitionId, TETRISCHED_ILP_TYPE resources);
+  /// Add an allocation for the Partition at the given time.
+  void addPartitionAllocation(uint32_t partitionId, Time time,
+                              uint32_t allocation);
 
   /// Retrieve the name of the Task.
   std::string getName() const;
 
   /// Retrieve the start time of the Placement, if available.
   std::optional<Time> getStartTime() const;
 
-  /// Retrieve an assignment from the Partition ID to the resources it is
-  /// contributing to this Placement.
-  std::vector<std::pair<uint32_t, TETRISCHED_ILP_TYPE>>
-  getPartitionAssignments() const;
+  /// Retrieve the end time of the Placement, if available.
+  std::optional<Time> getEndTime() const;
 
-  /// Retrieve the total resources contributed by all Partitions to this
-  /// Placement.
-  TETRISCHED_ILP_TYPE getTotalResources() const;
+  /// Retrieve the allocations.
+  const std::unordered_map<uint32_t, std::set<std::pair<Time, uint32_t>>>&
+  getPartitionAllocations() const;
 };
 using PlacementPtr = std::shared_ptr<Placement>;
 
@@ -359,6 +364,7 @@ class MalleableChooseExpression : public Expression {
                        Partitions availablePartitions,
                        CapacityConstraintMap& capacityConstraints,
                        Time currentTime) override;
+  SolutionResultPtr populateResults(SolverModelPtr solverModel) override;
 };
 
 /// An `AllocationExpression` represents the allocation of the given number of

schedulers/tetrisched/python/Expressions.cpp

@@ -26,8 +26,8 @@ void defineSTRLExpressions(py::module_& tetrisched_m) {
       .def_property_readonly("startTime", &tetrisched::Placement::getStartTime)
       .def("isPlaced", &tetrisched::Placement::isPlaced,
            "Returns true if the Placement was placed, false otherwise.")
-      .def("getPartitionAssignments",
-           &tetrisched::Placement::getPartitionAssignments,
+      .def("getPartitionAllocations",
+           &tetrisched::Placement::getPartitionAllocations,
            "Returns the Partition assignments for this Placement.");
 
   // Define the SolutionResult.

schedulers/tetrisched/src/Expression.cpp

@@ -5,37 +5,37 @@
 namespace tetrisched {
 
 /* Method definitions for Placement */
-Placement::Placement(std::string taskName, Time startTime)
-    : taskName(taskName), startTime(startTime), placed(true) {}
+Placement::Placement(std::string taskName, Time startTime, Time endTime)
+    : taskName(taskName),
+      startTime(startTime),
+      endTime(endTime),
+      placed(true) {}
 
 Placement::Placement(std::string taskName)
     : taskName(taskName), startTime(std::nullopt), placed(false) {}
 
 bool Placement::isPlaced() const { return placed; }
 
-void Placement::addPartition(uint32_t partitionId, TETRISCHED_ILP_TYPE usage) {
-  partitionToResources[partitionId] = usage;
+void Placement::addPartitionAllocation(uint32_t partitionId, Time time,
+                                       uint32_t allocation) {
+  if (partitionToResourceAllocations.find(partitionId) ==
+      partitionToResourceAllocations.end()) {
+    partitionToResourceAllocations[partitionId] =
+        std::set<std::pair<Time, uint32_t>>();
+  }
+  partitionToResourceAllocations[partitionId].insert(
+      std::make_pair(time, allocation));
 }
 
 std::string Placement::getName() const { return taskName; }
 
 std::optional<Time> Placement::getStartTime() const { return startTime; }
 
-std::vector<std::pair<uint32_t, TETRISCHED_ILP_TYPE>>
-Placement::getPartitionAssignments() const {
-  std::vector<std::pair<uint32_t, TETRISCHED_ILP_TYPE>> partitionAssignments;
-  for (const auto& [partitionId, usage] : partitionToResources) {
-    partitionAssignments.push_back(std::make_pair(partitionId, usage));
-  }
-  return partitionAssignments;
-}
+std::optional<Time> Placement::getEndTime() const { return endTime; }
 
-TETRISCHED_ILP_TYPE Placement::getTotalResources() const {
-  TETRISCHED_ILP_TYPE totalResources = 0;
-  for (const auto& [_, usage] : partitionToResources) {
-    totalResources += usage;
-  }
-  return totalResources;
+const std::unordered_map<uint32_t, std::set<std::pair<Time, uint32_t>>>&
+Placement::getPartitionAllocations() const {
+  return partitionToResourceAllocations;
 }
 
 /* Method definitions for CapacityConstraintMap */
@@ -318,8 +318,8 @@ ParseResultPtr ChooseExpression::parse(
                    << name << " will be limited to "
                    << schedulablePartitions.size() << " partitions.");
   if (schedulablePartitions.size() == 0) {
-    // There are no schedulable partitions, this expression cannot be satisfied.
-    // and should provide 0 utility.
+    // There are no schedulable partitions, this expression cannot be
+    // satisfied. and should provide 0 utility.
     parsedResult->type = ParseResultType::EXPRESSION_NO_UTILITY;
     return parsedResult;
   }
@@ -396,25 +396,19 @@ SolutionResultPtr ChooseExpression::populateResults(
   }
 
   // Find the ID of the Partition that was chosen.
-  PlacementPtr placement =
-      std::make_shared<Placement>(name, solution->startTime.value());
+  PlacementPtr placement = std::make_shared<Placement>(
+      name, solution->startTime.value(), solution->endTime.value());
   for (const auto& [partitionId, variable] : partitionVariables) {
     auto variableValue = variable->getValue();
     if (variableValue == 0) {
       // This partition was not used.
       continue;
     }
     // This partition was used. Add it to the Placement.
-    placement->addPartition(partitionId, variableValue.value());
-  }
-
-  if (placement->getTotalResources() != numRequiredMachines) {
-    throw tetrisched::exceptions::ExpressionSolutionException(
-        "ChooseExpression for " + name +
-        " was satisfied but did not use the correct number of machines.");
-  } else {
-    solution->placements[name] = std::move(placement);
+    placement->addPartitionAllocation(partitionId, solution->startTime.value(),
+                                      variableValue.value());
   }
+  solution->placements[name] = std::move(placement);
   return solution;
 }
 
@@ -732,6 +726,36 @@ ParseResultPtr MalleableChooseExpression::parse(
   return parsedResult;
 }
 
+SolutionResultPtr MalleableChooseExpression::populateResults(
+    SolverModelPtr solverModel) {
+  // Populate the results for the SolverModel's variables (i.e, this
+  // Expression's utility, start time and end time) from the Base Expression
+  // class.
+  Expression::populateResults(solverModel);
+
+  // Populate the Placements from the SolverModel.
+  if (!solution->utility || solution->utility.value() == 0) {
+    // This Choose expression was not satisfied.
+    // No placements to populate.
+    return solution;
+  }
+
+  // Find the IDs of the Partitions that were chosen.
+  PlacementPtr placement = std::make_shared<Placement>(
+      name, solution->startTime.value(), solution->endTime.value());
+  for (const auto& [partitionPair, variable] : partitionVariables) {
+    auto variableValue = variable->getValue();
+    if (variableValue == 0) {
+      // This partition was not used.
+      continue;
+    }
+    auto& [partitionId, time] = partitionPair;
+    placement->addPartitionAllocation(partitionId, time, variableValue.value());
+  }
+  solution->placements[name] = std::move(placement);
+  return solution;
+}
+
 /* Method definitions for AllocationExpression */
 
 AllocationExpression::AllocationExpression(
@@ -913,7 +937,8 @@ ParseResultPtr LessThanExpression::parse(
   if (firstChildResult->type != ParseResultType::EXPRESSION_UTILITY ||
       secondChildResult->type != ParseResultType::EXPRESSION_UTILITY) {
     throw tetrisched::exceptions::ExpressionConstructionException(
-        "LessThanExpression must have two children that are being evaluated.");
+        "LessThanExpression must have two children that are being "
+        "evaluated.");
   }
 
   // Generate the result of parsing the expression.
@@ -1264,7 +1289,8 @@ ParseResultPtr ScaleExpression::parse(
 
     if (!childParseResult->utility) {
       throw tetrisched::exceptions::ExpressionConstructionException(
-          "ScaleExpression applied to a child that does not have any utility.");
+          "ScaleExpression applied to a child that does not have any "
+          "utility.");
     }
     TETRISCHED_DEBUG("The child utility is "
                      << childParseResult->utility.value()->toString());

schedulers/tetrisched/test/test_expression.cpp

@@ -120,17 +120,22 @@ TEST(Expression, TestLessThanEnforcesOrdering) {
   EXPECT_EQ(result->placements.size(), 2) << "There should be 2 placements.";
   EXPECT_TRUE(result->placements["task1"]->isPlaced())
       << "task1 should be placed";
-  EXPECT_EQ(result->placements["task1"]->getPartitionAssignments()[0].first,
-            partition->getPartitionId())
-      << "task1 should be placed on partition.";
+
+  auto allocationsForTask1 =
+      result->placements["task1"]->getPartitionAllocations();
+  EXPECT_TRUE(allocationsForTask1.find(partition->getPartitionId()) !=
+              allocationsForTask1.end())
+      << "task1 should be placed on partition 1.";
   EXPECT_EQ(result->placements["task1"]->getStartTime().value(), 0)
       << "task1 should start at 0.";
 
   EXPECT_TRUE(result->placements["task2"]->isPlaced())
       << "task2 should be placed";
-  EXPECT_EQ(result->placements["task2"]->getPartitionAssignments()[0].first,
-            partition->getPartitionId())
-      << "task2 should be placed on partition.";
+  auto allocationsForTask2 =
+      result->placements["task2"]->getPartitionAllocations();
+  EXPECT_TRUE(allocationsForTask2.find(partition->getPartitionId()) !=
+              allocationsForTask2.end())
+      << "task2 should be placed on partition 1.";
   EXPECT_EQ(result->placements["task2"]->getStartTime().value(), 200)
       << "task2 should start at 200.";
 }
@@ -426,7 +431,7 @@ TEST(Expression, TestMalleableChooseExpressionConstructsVariableRectangles) {
   // usage.
   tetrisched::ExpressionPtr malleableChooseExpression =
       std::make_shared<tetrisched::MalleableChooseExpression>(
-          "task1", partitions, 10, 4, 10, 1);
+          "task2", partitions, 10, 4, 10, 1);
 
   // Construct an ObjectiveExpression.
   tetrisched::ExpressionPtr objectiveExpression =
@@ -444,7 +449,23 @@ TEST(Expression, TestMalleableChooseExpressionConstructsVariableRectangles) {
                                       capacityConstraintMap, 0);
   solverModelPtr->exportModel("testMalleableChooseVariableRectangle.lp");
 
+  // Translate and Solve the model.
   gurobiSolver.translateModel();
-  gurobiSolver.exportModel("testblah.lp");
+  gurobiSolver.solveModel();
+
+  auto result = objectiveExpression->populateResults(solverModelPtr);
+  EXPECT_TRUE(result->utility) << "Result should have some utility.";
+  EXPECT_EQ(2, result->utility.value()) << "The utility should be 2.";
+  auto task2Placement = result->placements["task2"]->getPartitionAllocations();
+  EXPECT_TRUE(task2Placement.find(partition1->getPartitionId()) !=
+              task2Placement.end())
+      << "task2 should be placed on partition 1.";
+  auto partition1Allocation = task2Placement[partition1->getPartitionId()];
+  uint32_t totalAllocationSum = 0;
+  for (auto& [time, allocation] : partition1Allocation) {
+    totalAllocationSum += allocation;
+  }
+  EXPECT_EQ(totalAllocationSum, 10)
+      << "The total allocation on partition 1 should be 10.";
 }
 #endif