def generate_status_json(self, status_dict):
"""Generates the portion of the status JSON that describes the resources
:param status_dict: The status JSON dict
:type status_dict: dict
"""
num_offers = 0
total_running = NodeResources()
total_offered = NodeResources()
total_watermark = NodeResources()
total_resources = NodeResources()
with self._agent_resources_lock:
for node_dict in status_dict['nodes']:
agent_id = node_dict['agent_id']
is_active = node_dict['is_active']
if agent_id in self._agent_resources:
agent_resources = self._agent_resources[agent_id]
if is_active:
num_offers += agent_resources.generate_status_json(
node_dict, total_running, total_offered,
total_watermark, total_resources)
else:
agent_resources.generate_status_json(node_dict)
free_resources = total_watermark.copy()
free_resources.subtract(total_running)
free_resources.subtract(total_offered)
unavailable_resources = total_resources.copy()
unavailable_resources.subtract(total_watermark)
resources_dict = {}
total_running.round_values()
total_offered.round_values()
free_resources.round_values()
unavailable_resources.round_values()
total_resources.round_values()
total_running.generate_status_json(resources_dict, 'running')
total_offered.generate_status_json(resources_dict, 'offered')
free_resources.generate_status_json(resources_dict, 'free')
unavailable_resources.generate_status_json(resources_dict,
'unavailable')
total_resources.generate_status_json(resources_dict, 'total')
# Fill in any missing values
for resource in total_resources.resources:
resource_dict = resources_dict[resource.name]
if 'running' not in resource_dict:
resource_dict['running'] = 0.0
if 'offered' not in resource_dict:
resource_dict['offered'] = 0.0
if 'free' not in resource_dict:
resource_dict['free'] = 0.0
if 'unavailable' not in resource_dict:
resource_dict['unavailable'] = 0.0
status_dict['num_offers'] = num_offers
status_dict['resources'] = resources_dict
class AgentResources(object):
"""This class represents an agent's set of resource offers."""
def __init__(self, agent_id):
"""Constructor
:param agent_id: The agent ID
:type agent_id: string
"""
self.agent_id = agent_id
self._offers = {} # {Offer ID: ResourceOffer}
self._recent_watermark_resources = NodeResources() # Recent watermark, used to provide a rolling watermark
self._task_resources = NodeResources() # Total resources for current tasks
self._watermark_resources = NodeResources() # Highest level of offer + task resources
self._offer_resources = None # Resources from offers
self._shortage_resources = None # Resources that agent needs to fulfill current obligations
self._total_resources = None
self._update_resources()
def allocate_offers(self, resources, when):
"""Directs the agent to allocate offers sufficient to match the given resources. Any offers that have been held
too long will automatically be included. It's possible that the offer resources returned are less than
requested.
:param resources: The requested resources
:type resources: :class:`node.resources.node_resources.NodeResources`
:param when: The current time
:type when: :class:`datetime.datetime`
:returns: The list of allocated offers
:rtype: [:class:`scheduler.resources.offer.ResourceOffer`]
"""
allocated_offers = {}
allocated_resources = NodeResources()
available_offer_ids = set(self._offers.keys())
# Automatically include all offers that have been held too long
for offer in self._offers.values():
if when - offer.received >= MAX_OFFER_HOLD_DURATION:
allocated_offers[offer.id] = offer
allocated_resources.add(offer.resources)
available_offer_ids.discard(offer.id)
if self._offer_resources.is_sufficient_to_meet(resources):
# We have enough resources to meet the request, so keep allocating offers until we get enough
while not allocated_resources.is_sufficient_to_meet(resources):
if len(available_offer_ids) == 0:
# We unexpectedly ran out of offers (can occur due to resource rounding error)
break
offer_id = available_offer_ids.pop()
offer = self._offers[offer_id]
allocated_offers[offer_id] = offer
allocated_resources.add(offer.resources)
# Remove allocated offers and return them
for offer in allocated_offers.values():
del self._offers[offer.id]
self._update_resources()
return allocated_offers.values()
def generate_status_json(self, node_dict, total_running=None, total_offered=None, total_watermark=None, total=None):
"""Generates the portion of the status JSON that describes the resources for this agent
:param node_dict: The dict for this agent's node within the status JSON
:type node_dict: dict
:param total_running: The total running resources to add up, possibly None
:type total_running: :class:`node.resources.node_resources.NodeResources`
:param total_offered: The total offered resources to add up, possibly None
:type total_offered: :class:`node.resources.node_resources.NodeResources`
:param total_watermark: The total watermark resources to add up, possibly None
:type total_watermark: :class:`node.resources.node_resources.NodeResources`
:param total: The total resources to add up, possibly None
:type total: :class:`node.resources.node_resources.NodeResources`
:returns: The total number of offers this agent has
:rtype: int
"""
if self._total_resources:
total_resources = self._total_resources
else:
total_resources = self._watermark_resources
free_resources = self._watermark_resources.copy()
free_resources.subtract(self._task_resources)
free_resources.subtract(self._offer_resources)
free_resources.round_values()
unavailable_resources = total_resources.copy()
unavailable_resources.subtract(self._watermark_resources)
unavailable_resources.round_values()
resources_dict = {}
if total_running:
total_running.add(self._task_resources)
if total_offered:
total_offered.add(self._offer_resources)
if total_watermark:
total_watermark.add(self._watermark_resources)
if total:
total.add(total_resources)
self._task_resources.generate_status_json(resources_dict, 'running')
self._offer_resources.generate_status_json(resources_dict, 'offered')
free_resources.generate_status_json(resources_dict, 'free')
unavailable_resources.generate_status_json(resources_dict, 'unavailable')
total_resources.generate_status_json(resources_dict, 'total')
# Fill in any missing values
for resource in total_resources.resources:
resource_dict = resources_dict[resource.name]
if 'running' not in resource_dict:
resource_dict['running'] = 0.0
if 'offered' not in resource_dict:
resource_dict['offered'] = 0.0
if 'free' not in resource_dict:
resource_dict['free'] = 0.0
if 'unavailable' not in resource_dict:
resource_dict['unavailable'] = 0.0
num_offers = len(self._offers)
node_dict['num_offers'] = num_offers
node_dict['resources'] = resources_dict
return num_offers
def decline_offers(self):
"""Removes offers that haven't been allocated
and returns them for the scheduler to decline them.
:returns: list of declined offers
:rtype: list
"""
# Decline old offers
declined_offers = []
declined_ids = []
for offer in self._offers.values():
declined_offers.append(offer)
declined_ids.append(offer.id)
for id in declined_ids:
if id in self._offers:
del self._offers[id]
self._update_resources()
return declined_offers
def has_total_resources(self):
"""Indicates whether this agent knows its total resources or not
:returns: True if agent knows its total resources, False otherwise
:rtype: bool
"""
return self._total_resources is not None
def refresh_resources(self, offers, tasks):
"""Refreshes the agent's resources by setting the current running tasks and adding new resource offers. Returns
a copy of the set of resources for the agent.
:param offers: The new resource offers to add
:type offers: [:class:`scheduler.resources.offer.ResourceOffer`]
:param tasks: The current tasks running on the agent
:type tasks: [:class:`job.tasks.base_task.Task`]
:returns: A copy of the set of agent resources
:rtype: :class:`scheduler.resources.agent.ResourceSet`
"""
# Add new offers
for offer in offers:
if offer.id not in self._offers:
self._offers[offer.id] = offer
self._update_resources(tasks)
offered_resources = self._offer_resources.copy()
task_resources = self._task_resources.copy()
watermark_resources = self._watermark_resources.copy()
return ResourceSet(offered_resources, task_resources, watermark_resources)
def rescind_offers(self, offer_ids):
"""Rescinds the offers with the given IDs
:param offer_ids: The list of IDs of the offers to remove
:type offer_ids: [str]
"""
for offer_id in offer_ids:
if offer_id in self._offers:
offer = self._offers[offer_id]
del self._offers[offer_id]
self._update_resources()
def reset_watermark(self):
"""Resets the agent's watermark to the highest recent value
"""
self._watermark_resources = self._recent_watermark_resources
self._recent_watermark_resources = NodeResources()
self._update_resources()
def set_shortage(self, shortage_resources=None):
"""Sets the resource shortage for the agent, if any
:param shortage_resources: The resource shortage
:type shortage_resources: :class:`node.resources.node_resources.NodeResources`
"""
if shortage_resources:
shortage_resources.round_values()
self._shortage_resources = shortage_resources
def get_max_resources(self):
"""Gets the maximum resources for the agent
:returns: The maximum resources for an agent
:rtype: :class:`node.resources.node_resources.NodeResources`
"""
return self._total_resources if self._total_resources else self._watermark_resources
def set_total(self, total_resources):
"""Sets the total resources for the agent
:param total_resources: The total resources
:type total_resources: :class:`node.resources.node_resources.NodeResources`
"""
self._total_resources = total_resources.copy()
def _update_resources(self, tasks=None):
"""Updates the agent's resources from its current offers and tasks
:param tasks: The new list of current tasks running on the agent, possibly None
:type tasks: list
"""
# Add up offered resources
self._offer_resources = NodeResources()
for offer in self._offers.values():
self._offer_resources.add(offer.resources)
# Recalculate task resources if needed
if tasks is not None:
self._task_resources = NodeResources()
for task in tasks:
self._task_resources.add(task.get_resources())
# Increase watermark if needed
available_resources = self._offer_resources.copy()
available_resources.add(self._task_resources)
self._watermark_resources.increase_up_to(available_resources)
self._recent_watermark_resources.increase_up_to(available_resources)
# Make sure watermark does not exceed total (can happen when we get task resources back before task update)
if self._total_resources and not self._total_resources.is_sufficient_to_meet(self._watermark_resources):
self._watermark_resources.limit_to(self._total_resources)
self._recent_watermark_resources.limit_to(self._total_resources)
# Since watermark was limited to not be higher than total, we're going to limit offered resources so that
# offered + task = watermark
max_offered = self._watermark_resources.copy()
max_offered.subtract(self._task_resources)
self._offer_resources.limit_to(max_offered)
# Round values to deal with float precision issues
self._offer_resources.round_values()
self._task_resources.round_values()
self._watermark_resources.round_values()