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 stealTasks(self, offer, numToSteal=None, minNumTasks=2, stealFromBack=True):
"""
Give an offer to the queue. If accepted, it will return a
new task chunk containing the subTasks it stole.
If rejected, returns None
"""
popped_tasks = []
stolenTasksChunk = None
while self.queue.hasNext():
task = self.queue.pop()
if (chunk_utils.numSubTasks(task) > minNumTasks and fitsIn(task, offer)):
taskCopy = mesos_pb2.TaskInfo()
taskCopy.CopyFrom(task)
if numToSteal is None: #Steal Half
numToSteal = chunk_utils.numSubTasks(taskCopy) / 2
stolenTasks = self.stealSubTasks(taskCopy, numToSteal, stealFromBack)
stolenTasksChunk = chunk_utils.newTaskChunk(offer.slave_id,
executor=taskCopy.executor, subTasks=stolenTasks)
popped_tasks.append(taskCopy)
break
popped_tasks.append(task)
for task in popped_tasks:
self.queue.push(task)
return stolenTasksChunk
def runNextSubTask(self, driver, taskChunkId):
"""
Launches next subtask from current taskChunk.
"""
taskChunk = self.pendingTaskChunks[taskChunkId]
if chunk_utils.numSubTasks(taskChunk) > 0:
nextSubTaskId = chunk_utils.nextSubTask(taskChunk)
self.launchTask(driver, nextSubTaskId)
else:
update = mesos_pb2.TaskStatus()
update.task_id.value = taskChunkId.value
update.state = mesos_pb2.TASK_FINISHED
driver.sendStatusUpdate(update)
def resourceOffers(self, driver, offers):
print "Got %d resource offers" % len(offers)
# Initialize the tasks.
if self.all_tasks is None:
if ENABLE_LEVELS:
self.all_tasks = []
for i in xrange(NUM_TASK_CHUNKS):
tasks = self.initializeTasks(offers, i)
taskChunk = chunk_utils.newTaskChunk(offers[0].slave_id,
executor=self.executor,
subTasks = tasks)
self.all_tasks.append(taskChunk)
taskChunk.task_id.value = "Initial_taskChunk_"+str(i)
taskChunk.name = "Initial_taskChunk_"+str(i)
else:
self.all_tasks = self.initializeTasks(offers)
taskChunk = chunk_utils.newTaskChunk(offers[0].slave_id,
executor=self.executor, subTasks = self.all_tasks)
# dict <offerId, listofTaskChunks>
dictOffers = task_utils.selectTasksforOffers(offers, [taskChunk],
len(self.all_tasks), self.num_slaves,
distribution=self.distribution, isTaskChunk=ENABLE_TASK_CHUNKING)
for offer in offers:
subTasks = dictOffers[offer.id.value]
print "offer_id = %s, numSubTasks" % (offer.id.value), len(subTasks)
if subTasks:
offerTotalTasks = 0
for subTask in subTasks:
offerTotalTasks += chunk_utils.numSubTasks(subTask)
self.tasksLaunched += offerTotalTasks
index = min(offerTotalTasks,len(self.all_tasks))
del self.all_tasks[:index]
print "Accepting offer on %s to start task chunk" % offer.hostname
driver.launchTasks(offer.id, subTasks)
else:
print "Rejecting offer {0}".format(offer.id.value)
driver.launchTasks(offer.id, [])
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)
def __init__(self, pending_tasks):
self.queue = PriorityQueue(pending_tasks,
sort_key = lambda task: -chunk_utils.numSubTasks(task),
mapper = lambda offer: offer.task_id.value)