def get_all_tasks_from_state(
mesos_state: MesosState,
include_orphans: bool = False) -> Sequence[MesosTask]:
"""Given a mesos state, find the tasks from all frameworks.
:param mesos_state: the mesos_state
:returns: a list of tasks
"""
tasks = [
task for framework in mesos_state.get("frameworks", [])
for task in framework.get("tasks", [])
]
if include_orphans:
tasks += mesos_state.get("orphan_tasks", [])
return tasks
def has_registered_slaves(mesos_state: MesosState, ) -> bool:
""" Return a boolean indicating if there are any slaves registered
to the master according to the mesos state.
:param mesos_state: the mesos state from the master
:returns: a boolean, indicating if there are > 0 slaves
"""
return len(mesos_state.get("slaves", [])) > 0
def get_resource_utilization_by_grouping(
grouping_func: Callable[..., _KeyFuncRetT],
mesos_state: MesosState,
filters: List[Callable]=[],
sort_func=None,
) -> Dict[_KeyFuncRetT, ResourceUtilizationDict]:
""" Given a function used to group slaves and mesos state, calculate
resource utilization for each value of a given attribute.
:grouping_func: a function that given a slave, will return the value of an
attribute to group by.
:param mesos_state: the mesos state
:param filters: filters to apply to the slaves in the calculation, with
filtering preformed by filter_slaves
:returns: a dict of {attribute_value: resource_usage}, where resource usage
is the dict returned by ``calculate_resource_utilization_for_slaves`` for
slaves grouped by attribute value.
"""
slaves = mesos_state.get('slaves', [])
slaves = filter_slaves(slaves, filters)
if not has_registered_slaves(mesos_state):
raise ValueError("There are no slaves registered in the mesos state.")
tasks = get_all_tasks_from_state(mesos_state, include_orphans=True)
non_terminal_tasks = [task for task in tasks if not is_task_terminal(task)]
slave_groupings = group_slaves_by_key_func(grouping_func, slaves, sort_func)
return {
attribute_value: calculate_resource_utilization_for_slaves(
slaves=slaves,
tasks=filter_tasks_for_slaves(slaves, non_terminal_tasks),
)
for attribute_value, slaves in slave_groupings.items()
}
async def get_mesos_task_count_by_slave(
mesos_state: MesosState,
slaves_list: Sequence[Dict]=None,
pool: Optional[str]=None,
) -> List[Dict]:
"""Get counts of running tasks per mesos slave. Also include separate count of chronos tasks
:param mesos_state: mesos state dict
:param slaves_list: a list of slave dicts to count running tasks for.
:param pool: pool of slaves to return (None means all)
:returns: list of slave dicts {'task_count': SlaveTaskCount}
"""
all_mesos_tasks = await get_all_running_tasks() # empty string = all app ids
slaves = {
slave['id']: {'count': 0, 'slave': slave, 'batch_count': 0} for slave in mesos_state.get('slaves', [])
}
for task in all_mesos_tasks:
try:
task_slave = await task.slave()
if task_slave['id'] not in slaves:
log.debug("Slave {} not found for task".format(task_slave['id']))
continue
else:
slaves[task_slave['id']]['count'] += 1
task_framework = await task.framework()
log.debug(f"Task framework: {task_framework.name}")
# Marathon is only framework that runs service. Others are batch.
if not task_framework.name.startswith(MARATHON_FRAMEWORK_NAME_PREFIX):
slaves[task_slave['id']]['batch_count'] += 1
except SlaveDoesNotExist:
log.debug("Tried to get mesos slaves for task {}, but none existed.".format(task['id']))
continue
if slaves_list:
for slave in slaves_list:
slave['task_counts'] = SlaveTaskCount(**slaves[slave['task_counts'].slave['id']])
slaves_with_counts = list(slaves_list)
elif pool:
slaves_with_counts = [
{'task_counts': SlaveTaskCount(**slave_counts)} for slave_counts in slaves.values()
if slave_counts['slave']['attributes'].get('pool', 'default') == pool
]
else:
slaves_with_counts = [{'task_counts': SlaveTaskCount(**slave_counts)} for slave_counts in slaves.values()]
for slave in slaves_with_counts:
log.debug("Slave: {}, running {} tasks, "
"including {} chronos tasks".format(
slave['task_counts'].slave['hostname'],
slave['task_counts'].count,
slave['task_counts'].batch_count,
))
return slaves_with_counts
