def launchTasks(self, offerId, tasks, filters=None):
"""
Keeps track of the task and its launched resources before launching.
"""
consumedResources = self.consumedResources[offerId.value]
for task in tasks:
# Update the resource offer with the task resources.
chunk_utils.incrementResources(consumedResources, task.resources)
# Add the task to the tracked tasked.
self.pendingTasks.addTask(task)
TaskChunkSchedulerDriver.launchTasks(self, offerId, tasks, filters)
def __init__(self, scheduler, framework, master, outerScheduler=None):
if not isinstance(scheduler, TaskStealingScheduler):
scheduler = TaskStealingScheduler(scheduler)
# Wrap the outer scheduler with the driver overriding wrapper.
if not outerScheduler:
outerScheduler = chunk_utils.DriverOverridingScheduler(scheduler, self)
TaskChunkSchedulerDriver.__init__(self, scheduler, framework, master, outerScheduler)
self.consumedResources = defaultdict(lambda: mesos_pb2.Offer().resources)
self.pendingTasks = chunk_utils.TaskTable()
def killSubTasks(self, subTaskIds):
"""
Kills the sub tasks with the given IDs, and clears them from the table.
"""
subTasks = []
parentIds = []
for subTaskId in subTaskIds:
# Backwards compatibility for passing sub tasks instead of IDs.
if isinstance(subTaskId, mesos_pb2.TaskInfo):
subTaskId = subTaskId.task_id
parent = self.pendingTasks.getParent(subTaskId)
parentIds.append(parent.task_id)
subTask = self.pendingTasks[subTaskId]
subTasks.append(subTask)
del self.pendingTasks[subTaskId]
TaskChunkSchedulerDriver.killSubTasks(self, subTasks)
return zip(parentIds, subTasks)
