TEST(SorterTest, Update)
{
DRFSorter sorter;
sorter.add("a");
sorter.add("b");
sorter.add(Resources::parse("cpus:10;mem:10;disk:10").get());
sorter.allocated("a", Resources::parse("cpus:10;mem:10;disk:10").get());
// Construct an offer operation.
Resource volume = Resources::parse("disk", "5", "*").get();
volume.mutable_disk()->mutable_persistence()->set_id("ID");
volume.mutable_disk()->mutable_volume()->set_container_path("data");
Offer::Operation create;
create.set_type(Offer::Operation::CREATE);
create.mutable_create()->add_volumes()->CopyFrom(volume);
// Compute the updated allocation.
Resources allocation = sorter.allocation("a");
Try<Resources> newAllocation = allocation.apply(create);
ASSERT_SOME(newAllocation);
// Update the resources for the client.
sorter.update("a", allocation, newAllocation.get());
EXPECT_EQ(newAllocation.get(), sorter.allocation("a"));
}
TEST(SorterTest, UpdateAllocation)
{
DRFSorter sorter;
SlaveID slaveId;
slaveId.set_value("agentId");
sorter.add("a");
sorter.add("b");
sorter.add(slaveId, Resources::parse("cpus:10;mem:10;disk:10").get());
sorter.allocated(
"a", slaveId, Resources::parse("cpus:10;mem:10;disk:10").get());
// Construct an offer operation.
Resource volume = Resources::parse("disk", "5", "*").get();
volume.mutable_disk()->mutable_persistence()->set_id("ID");
volume.mutable_disk()->mutable_volume()->set_container_path("data");
// Compute the updated allocation.
Resources oldAllocation = sorter.allocation("a", slaveId);
Try<Resources> newAllocation = oldAllocation.apply(CREATE(volume));
ASSERT_SOME(newAllocation);
// Update the resources for the client.
sorter.update("a", slaveId, oldAllocation, newAllocation.get());
hashmap<SlaveID, Resources> allocation = sorter.allocation("a");
EXPECT_EQ(1u, allocation.size());
EXPECT_EQ(newAllocation.get(), allocation[slaveId]);
EXPECT_EQ(newAllocation.get(), sorter.allocation("a", slaveId));
}
// We aggregate resources from multiple slaves into the sorter. Since
// non-scalar resources don't aggregate well across slaves, we need to
// keep track of the SlaveIDs of the resources. This tests that no
// resources vanish in the process of aggregation by performing update
// allocations from unreserved to reserved resources.
TEST(SorterTest, MultipleSlavesUpdateAllocation)
{
DRFSorter sorter;
SlaveID slaveA;
slaveA.set_value("agentA");
SlaveID slaveB;
slaveB.set_value("agentB");
sorter.add("framework");
Resources slaveResources =
Resources::parse("cpus:2;mem:512;disk:10;ports:[31000-32000]").get();
sorter.add(slaveA, slaveResources);
sorter.add(slaveB, slaveResources);
sorter.allocated("framework", slaveA, slaveResources);
sorter.allocated("framework", slaveB, slaveResources);
// Construct an offer operation.
Resource volume = Resources::parse("disk", "5", "*").get();
volume.mutable_disk()->mutable_persistence()->set_id("ID");
volume.mutable_disk()->mutable_volume()->set_container_path("data");
// Compute the updated allocation.
Try<Resources> newAllocation = slaveResources.apply(CREATE(volume));
ASSERT_SOME(newAllocation);
// Update the resources for the client.
sorter.update("framework", slaveA, slaveResources, newAllocation.get());
sorter.update("framework", slaveB, slaveResources, newAllocation.get());
EXPECT_EQ(2u, sorter.allocation("framework").size());
EXPECT_EQ(newAllocation.get(), sorter.allocation("framework", slaveA));
EXPECT_EQ(newAllocation.get(), sorter.allocation("framework", slaveB));
}
virtual void resourceOffers(SchedulerDriver* driver,
const vector<Offer>& offers)
{
foreach (const Offer& offer, offers) {
LOG(INFO) << "Received offer " << offer.id() << " with "
<< offer.resources();
// If the framework got this offer for the first time, the state is
// `State::INIT`; framework will reserve it (sending RESERVE operation
// to master) in this loop.
if (!states.contains(offer.slave_id())) {
// If all tasks were launched, do not reserve more resources; wait
// for them to finish and unreserve resources.
if (tasksLaunched == totalTasks) {
continue;
}
states[offer.slave_id()] = State::INIT;
}
const State state = states[offer.slave_id()];
Filters filters;
filters.set_refuse_seconds(0);
switch (state) {
case State::INIT: {
// Framework reserves resources from this offer for only one task;
// the task'll be dispatched when reserved resources are re-offered
// to this framework.
Resources resources = offer.resources();
Offer::Operation reserve = RESERVE(taskResources);
Try<Resources> apply = resources.apply(reserve);
if (apply.isError()) {
LOG(INFO) << "Failed to reserve resources for task in offer "
<< stringify(offer.id()) << ": " << apply.error();
break;
}
driver->acceptOffers({offer.id()}, {reserve}, filters);
states[offer.slave_id()] = State::RESERVING;
break;
}
case State::RESERVING: {
Resources resources = offer.resources();
Resources reserved = resources.reserved(role);
if (!reserved.contains(taskResources)) {
break;
}
states[offer.slave_id()] = State::RESERVED;
// We fallthrough here to save an offer cycle.
}
case State::RESERVED: {
Resources resources = offer.resources();
Resources reserved = resources.reserved(role);
CHECK(reserved.contains(taskResources));
// If all tasks were launched, unreserve those resources.
if (tasksLaunched == totalTasks) {
driver->acceptOffers(
{offer.id()}, {UNRESERVE(taskResources)}, filters);
states[offer.slave_id()] = State::UNRESERVING;
break;
}
// Framework dispatches task on the reserved resources.
CHECK(tasksLaunched < totalTasks);
// Launch tasks on reserved resources.
const string& taskId = stringify(tasksLaunched++);
LOG(INFO) << "Launching task " << taskId << " using offer "
<< offer.id();
TaskInfo task;
task.set_name("Task " + taskId + ": " + command);
task.mutable_task_id()->set_value(taskId);
task.mutable_slave_id()->MergeFrom(offer.slave_id());
task.mutable_command()->set_shell(true);
task.mutable_command()->set_value(command);
task.mutable_resources()->MergeFrom(taskResources);
driver->launchTasks(offer.id(), {task}, filters);
states[offer.slave_id()] = State::TASK_RUNNING;
break;
}
case State::TASK_RUNNING:
LOG(INFO) << "The task on " << offer.slave_id()
<< " is running, waiting for task done";
break;
case State::UNRESERVING: {
Resources resources = offer.resources();
Resources reserved = resources.reserved(role);
if (!reserved.contains(taskResources)) {
states[offer.slave_id()] = State::UNRESERVED;
}
break;
}
case State::UNRESERVED:
// If state of slave is UNRESERVED, ignore it. The driver is stopped
// when all tasks are done and all resources are unreserved.
break;
}
}
