def test_selectTasksToStealOfferSplitting(self):
offers = [self.generateOffer(8)]
pendingTasks = [self.newTaskChunk(4)]
stolen = self.stealingScheduler.selectTasksToSteal(self.driver, offers, pendingTasks)
self.assertEqual(1, len(stolen))
self.assertTrue(offers[0].id.value in stolen)
taskChunks = stolen[offers[0].id.value]
self.assertEqual(2, len(taskChunks))
# First stolen chunk.
taskChunk = taskChunks[0]
self.assertEqual(2, chunk_utils.numSubTasks(taskChunk))
stolenTasks = [subTask for subTask in chunk_utils.subTaskIterator(taskChunk)]
self.assertEqual("task_id_2", stolenTasks[0].task_id.value)
self.assertEqual("task_id_3", stolenTasks[1].task_id.value)
# Second stolen chunk.
taskChunk = taskChunks[1]
self.assertEqual(1, chunk_utils.numSubTasks(taskChunk))
stolenTasks = [subTask for subTask in chunk_utils.subTaskIterator(taskChunk)]
self.assertEqual("task_id_1", stolenTasks[0].task_id.value)
def stealSubTasks(self, taskChunk, numToSteal, stealFromBack):
"""
Returns a list of stolenSubTasks and removes them from the taskChunk
"""
subTasks = [subTask for subTask in
chunk_utils.subTaskIterator(taskChunk)]
if stealFromBack:
startIndex = len(subTasks) - numToSteal
stolenTasks = subTasks[startIndex:]
else:
stolenTasks = subTasks[:numToSteal]
for stolenTask in stolenTasks:
chunk_utils.removeSubTask(taskChunk, stolenTask.task_id)
return stolenTasks
def resourceOffersStealing(self, driver, offers):
"""
Steals currently pending tasks and launches them on the new offer.
"""
print "\tTrying to find tasks to steal for {0} offers".format(len(offers))
tasksToSteal = self.selectTasksToSteal(driver, offers, driver.pendingTasks)
for offerIdValue, taskChunks in tasksToSteal.iteritems():
offerId = mesos_pb2.OfferID()
offerId.value = offerIdValue
for taskChunk in taskChunks:
subTaskIds = [task.task_id for task in subTaskIterator(taskChunk)]
self.stealSubTasks(driver, subTaskIds)
print "\tStolen tasks to be run in new task chunk: {0}".format(
taskChunk.task_id.value)
driver.launchTasks(offerId, taskChunks)
def killTask(self, driver, taskId):
"""
Kills task specified using taskId.
taskId can refer to task or TaskChunks.
"""
if taskId in self.pendingTaskChunks:
origTaskId = taskId
task = self.pendingTaskChunks[taskId]
while chunk_utils.isTaskChunk(task):
for subTask in chunk_utils.subTaskIterator(task):
if self.pendingTaskChunks.isActive(subTask.task_id):
taskId = subTask.task_id
task = self.pendingTaskChunks[taskId]
break
update = mesos_pb2.TaskStatus()
update.task_id.value = origTaskId.value
update.state = mesos_pb2.TASK_KILLED
driver.sendStatusUpdate(update)
chunk_utils.ExecutorWrapper.killTask(self, driver, taskId)
def test_selectTasksToSteal(self):
offers = [self.generateOffer()]
pendingTasks = [self.newTaskChunk(4)]
self.stealingScheduler.generateTaskId = MagicMock()
self.stealingScheduler.generateTaskId.return_value = "chunk_id_0"
stolen = self.stealingScheduler.selectTasksToSteal(self.driver, offers, pendingTasks)
self.stealingScheduler.generateTaskId.assert_called_once()
self.assertEqual(1, len(stolen))
self.assertTrue(offers[0].id.value in stolen)
taskChunks = stolen[offers[0].id.value]
self.assertEqual(1, len(taskChunks))
taskChunk = taskChunks[0]
self.assertEqual(2, chunk_utils.numSubTasks(taskChunk))
stolenTasks = [subTask for subTask in chunk_utils.subTaskIterator(taskChunk)]
self.assertEqual("task_id_2", stolenTasks[0].task_id.value)
self.assertEqual("task_id_3", stolenTasks[1].task_id.value)
