<s> from collections import defaultdict <EOL> try : <EOL> from collections import OrderedDict <EOL> except ImportError : <EOL> from ordereddict import OrderedDict <EOL> import networkx as nx <EOL> import six <EOL> from nailgun import consts <EOL> from nailgun import errors <EOL> from nailgun . logger import logger <EOL> from nailgun import objects <EOL> from nailgun . orchestrator import priority_serializers as ps <EOL> from nailgun . orchestrator . tasks_serializer import TaskSerializers <EOL> from nailgun . policy . name_match import NameMatchingPolicy <EOL> from nailgun . utils . role_resolver import RoleResolver <EOL> class GraphSolver ( nx . DiGraph ) : <EOL> """<STR_LIT>""" <EOL> node_dict_factory = OrderedDict <EOL> adjlist_dict_factory = OrderedDict <EOL> def __init__ ( self , tasks = None , * args , ** kwargs ) : <EOL> super ( GraphSolver , self ) . __init__ ( * args , ** kwargs ) <EOL> self . node = self . node_dict_factory ( ) <EOL> self . adj = self . adjlist_dict_factory ( ) <EOL> self . pred = self . adjlist_dict_factory ( ) <EOL> self . succ = self . adj <EOL> if tasks is not None : <EOL> self . add_tasks ( tasks ) <EOL> def add_node ( self , n , ** attr ) : <EOL> if n not in self . succ : <EOL> self . succ [ n ] = self . adjlist_dict_factory ( ) <EOL> self . pred [ n ] = self . adjlist_dict_factory ( ) <EOL> self . node [ n ] = attr <EOL> super ( GraphSolver , self ) . add_node ( n , ** attr ) <EOL> def add_edge ( self , u , v , ** attr ) : <EOL> if u not in self . succ : <EOL> self . succ [ u ] = self . adjlist_dict_factory ( ) <EOL> self . pred [ u ] = self . adjlist_dict_factory ( ) <EOL> self . node [ u ] = self . node_dict_factory ( ) <EOL> if v not in self . succ : <EOL> self . succ [ v ] = self . adjlist_dict_factory ( ) <EOL> self . pred [ v ] = self . adjlist_dict_factory ( ) <EOL> self . node [ v ] = self . node_dict_factory ( ) <EOL> super ( GraphSolver , self ) . add_edge ( u , v , ** attr ) <EOL> def add_tasks ( self , tasks ) : <EOL> for task in tasks : <EOL> self . add_task ( task ) <EOL> self . _update_dependencies ( ) <EOL> def add_task ( self , task ) : <EOL> self . add_node ( task [ '<STR_LIT:id>' ] , ** task ) <EOL> for req in task . get ( '<STR_LIT>' , ( ) ) : <EOL> self . add_edge ( task [ '<STR_LIT:id>' ] , req ) <EOL> for req in task . get ( '<STR_LIT>' , ( ) ) : <EOL> self . add_edge ( req , task [ '<STR_LIT:id>' ] ) <EOL> if '<STR_LIT>' in task : <EOL> self . add_edge ( task [ '<STR_LIT:id>' ] , task [ '<STR_LIT>' ] ) <EOL> def _update_dependencies ( self ) : <EOL> """<STR_LIT>""" <EOL> for task in six . itervalues ( self . node ) : <EOL> available_groups = self . get_groups_subgraph ( ) . nodes ( ) <EOL> for group in task . get ( '<STR_LIT>' , ( ) ) : <EOL> pattern = NameMatchingPolicy . create ( group ) <EOL> not_matched = [ ] <EOL> for available_group in available_groups : <EOL> if pattern . match ( available_group ) : <EOL> self . add_edge ( task [ '<STR_LIT:id>' ] , available_group ) <EOL> else : <EOL> not_matched . append ( available_group ) <EOL> if len ( available_groups ) == len ( not_matched ) : <EOL> self . add_edge ( task [ '<STR_LIT:id>' ] , group ) <EOL> logger . warning ( <EOL> '<STR_LIT>' , group ) <EOL> available_groups = not_matched <EOL> for req in task . get ( '<STR_LIT>' , ( ) ) : <EOL> self . add_edge ( req , task [ '<STR_LIT:id>' ] ) <EOL> def is_acyclic ( self ) : <EOL> """<STR_LIT>""" <EOL> return nx . is_directed_acyclic_graph ( self ) <EOL> def get_next_groups ( self , processed_nodes ) : <EOL> """<STR_LIT>""" <EOL> result = [ ] <EOL> for node in self . nodes ( ) : <EOL> if node in processed_nodes : <EOL> continue <EOL> predecessors = nx . dfs_predecessors ( self . reverse ( ) , node ) <EOL> if ( set ( predecessors . keys ( ) ) <= processed_nodes ) : <EOL> result . append ( node ) <EOL> return result <EOL> def get_groups_subgraph ( self ) : <EOL> """<STR_LIT>""" <EOL> roles = [ t [ '<STR_LIT:id>' ] for t in six . itervalues ( self . node ) <EOL> if t . get ( '<STR_LIT:type>' ) == consts . ORCHESTRATOR_TASK_TYPES . group ] <EOL> return self . subgraph ( roles ) <EOL> def get_group_tasks ( self , group_name ) : <EOL> rst = [ ] <EOL> predecessors = self . predecessors ( group_name ) <EOL> for task in nx . topological_sort ( self ) : <EOL> if task not in predecessors : <EOL> continue <EOL> elif self . should_exclude_task ( task ) : <EOL> continue <EOL> rst . append ( self . node [ task ] ) <EOL> return rst <EOL> def should_exclude_task ( self , task ) : <EOL> """<STR_LIT>""" <EOL> if self . node [ task ] [ '<STR_LIT:type>' ] in consts . INTERNAL_TASKS : <EOL> logger . debug ( <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' , task ) <EOL> return True <EOL> elif self . node [ task ] . get ( '<STR_LIT>' ) : <EOL> logger . debug ( '<STR_LIT>' , task ) <EOL> return True <EOL> return False <EOL> @ property <EOL> def topology ( self ) : <EOL> return map ( lambda t : self . node [ t ] , nx . topological_sort ( self ) ) <EOL> def make_skipped_task ( self , task ) : <EOL> """<STR_LIT>""" <EOL> if task [ '<STR_LIT:type>' ] in consts . INTERNAL_TASKS : <EOL> logger . debug ( <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' , task ) <EOL> return <EOL> task [ '<STR_LIT>' ] = True <EOL> def only_tasks ( self , task_ids ) : <EOL> """<STR_LIT>""" <EOL> if not task_ids : <EOL> return <EOL> for task in self . node . values ( ) : <EOL> if task [ '<STR_LIT:id>' ] not in task_ids : <EOL> self . make_skipped_task ( task ) <EOL> else : <EOL> task [ '<STR_LIT>' ] = False <EOL> def reexecutable_tasks ( self , task_filter ) : <EOL> """<STR_LIT>""" <EOL> if not task_filter : <EOL> return <EOL> task_filter = set ( task_filter ) <EOL> for task in six . itervalues ( self . node ) : <EOL> reexecute_on = task . get ( '<STR_LIT>' ) <EOL> if reexecute_on is not None and task_filter . issubset ( reexecute_on ) : <EOL> task [ '<STR_LIT>' ] = False <EOL> else : <EOL> self . make_skipped_task ( task ) <EOL> def find_subgraph ( self , start = None , end = None ) : <EOL> """<STR_LIT>""" <EOL> working_graph = self <EOL> if start : <EOL> working_graph = self . subgraph ( <EOL> nx . dfs_postorder_nodes ( working_graph , start ) ) <EOL> if end : <EOL> working_graph = self . subgraph ( nx . dfs_postorder_nodes ( <EOL> working_graph . reverse ( ) , end ) ) <EOL> return working_graph <EOL> def filter_subgraph ( self , start = None , end = None , include = ( ) ) : <EOL> """<STR_LIT>""" <EOL> wgraph = self . find_subgraph ( start = start , end = end ) <EOL> for task in wgraph : <EOL> if self . should_exclude_task ( task ) and task not in include : <EOL> wgraph . remove_node ( task ) <EOL> return wgraph <EOL> class AstuteGraph ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , cluster , tasks = None ) : <EOL> self . cluster = cluster <EOL> self . tasks = tasks or objects . Cluster . get_deployment_tasks ( cluster ) <EOL> self . graph = GraphSolver ( ) <EOL> self . graph . add_tasks ( self . tasks ) <EOL> self . serializers = TaskSerializers ( ) <EOL> def only_tasks ( self , task_ids ) : <EOL> self . graph . only_tasks ( task_ids ) <EOL> def reexecutable_tasks ( self , task_filter ) : <EOL> self . graph . reexecutable_tasks ( task_filter ) <EOL> def group_nodes_by_roles ( self , nodes ) : <EOL> """<STR_LIT>""" <EOL> res = defaultdict ( list ) <EOL> for node in nodes : <EOL> res [ node [ '<STR_LIT>' ] ] . append ( node ) <EOL> return res <EOL> def get_nodes_with_roles ( self , grouped_nodes , roles ) : <EOL> """<STR_LIT>""" <EOL> result = [ ] <EOL> for role in roles : <EOL> if role in grouped_nodes : <EOL> result . extend ( grouped_nodes [ role ] ) <EOL> return result <EOL> def assign_parallel_nodes ( self , priority , nodes ) : <EOL> """<STR_LIT>""" <EOL> current_nodes = nodes <EOL> while current_nodes : <EOL> next_nodes = [ ] <EOL> group = [ ] <EOL> added_uids = [ ] <EOL> for node in current_nodes : <EOL> if '<STR_LIT>' not in node or '<STR_LIT>' not in node : <EOL> raise errors . InvalidSerializedNode ( <EOL> '<STR_LIT>' . format ( <EOL> node ) ) <EOL> if node [ '<STR_LIT>' ] not in added_uids : <EOL> group . append ( node ) <EOL> added_uids . append ( node [ '<STR_LIT>' ] ) <EOL> else : <EOL> next_nodes . append ( node ) <EOL> priority . in_parallel ( group ) <EOL> current_nodes = next_nodes <EOL> def process_parallel_nodes ( self , priority , parallel_groups , grouped_nodes ) : <EOL> """<STR_LIT>""" <EOL> parallel_nodes = [ ] <EOL> for group in parallel_groups : <EOL> nodes = self . get_nodes_with_roles ( grouped_nodes , group [ '<STR_LIT>' ] ) <EOL> if '<STR_LIT>' in group [ '<STR_LIT>' ] [ '<STR_LIT>' ] : <EOL> priority . in_parallel_by ( <EOL> nodes , <EOL> group [ '<STR_LIT>' ] [ '<STR_LIT>' ] [ '<STR_LIT>' ] ) <EOL> else : <EOL> parallel_nodes . extend ( nodes ) <EOL> if parallel_nodes : <EOL> self . assign_parallel_nodes ( priority , parallel_nodes ) <EOL> def add_priorities ( self , nodes ) : <EOL> """<STR_LIT>""" <EOL> priority = ps . PriorityStrategy ( ) <EOL> groups_subgraph = self . graph . get_groups_subgraph ( ) <EOL> all_groups = groups_subgraph . nodes ( ) <EOL> grouped_nodes = self . group_nodes_by_roles ( nodes ) <EOL> processed_groups = set ( all_groups ) - set ( grouped_nodes . keys ( ) ) <EOL> current_groups = groups_subgraph . get_next_groups ( processed_groups ) <EOL> while current_groups : <EOL> one_by_one = [ ] <EOL> parallel = [ ] <EOL> for r in current_groups : <EOL> group = self . graph . node [ r ] <EOL> if ( group [ '<STR_LIT>' ] [ '<STR_LIT>' ] [ '<STR_LIT:type>' ] <EOL> == consts . DEPLOY_STRATEGY . one_by_one ) : <EOL> one_by_one . append ( group ) <EOL> elif ( group [ '<STR_LIT>' ] [ '<STR_LIT>' ] [ '<STR_LIT:type>' ] <EOL> == consts . DEPLOY_STRATEGY . parallel ) : <EOL> parallel . append ( group ) <EOL> for group in one_by_one : <EOL> nodes = self . get_nodes_with_roles ( grouped_nodes , group [ '<STR_LIT>' ] ) <EOL> priority . one_by_one ( nodes ) <EOL> self . process_parallel_nodes ( priority , parallel , grouped_nodes ) <EOL> processed_groups . update ( current_groups ) <EOL> current_groups = groups_subgraph . get_next_groups ( processed_groups ) <EOL> def stage_tasks_serialize ( self , tasks , nodes ) : <EOL> """<STR_LIT>""" <EOL> serialized = [ ] <EOL> role_resolver = RoleResolver ( nodes ) <EOL> for task in tasks : <EOL> if self . graph . should_exclude_task ( task [ '<STR_LIT:id>' ] ) : <EOL> continue <EOL> serializer = self . serializers . get_stage_serializer ( task ) ( <EOL> task , self . cluster , nodes , role_resolver = role_resolver ) <EOL> if not serializer . should_execute ( ) : <EOL> continue <EOL> serialized . extend ( serializer . serialize ( ) ) <EOL> return serialized <EOL> def post_tasks_serialize ( self , nodes ) : <EOL> """<STR_LIT>""" <EOL> if '<STR_LIT>' in self . graph : <EOL> subgraph = self . graph . find_subgraph ( start = '<STR_LIT>' ) <EOL> else : <EOL> errors . NotEnoughInformation ( <EOL> '<STR_LIT>' ) <EOL> return self . stage_tasks_serialize ( subgraph . topology , nodes ) <EOL> def pre_tasks_serialize ( self , nodes ) : <EOL> """<STR_LIT>""" <EOL> if '<STR_LIT>' in self . graph : <EOL> subgraph = self . graph . find_subgraph ( end = '<STR_LIT>' ) <EOL> else : <EOL> raise errors . NotEnoughInformation ( <EOL> '<STR_LIT>' ) <EOL> return self . stage_tasks_serialize ( subgraph . topology , nodes ) <EOL> def deploy_task_serialize ( self , node ) : <EOL> """<STR_LIT>""" <EOL> serialized = [ ] <EOL> priority = ps . PriorityStrategy ( ) <EOL> for task in self . graph . get_group_tasks ( node [ '<STR_LIT>' ] ) : <EOL> serializer = self . serializers . get_deploy_serializer ( task ) ( <EOL> task , self . cluster , node ) <EOL> if not serializer . should_execute ( ) : <EOL> continue <EOL> serialized . extend ( serializer . serialize ( ) ) <EOL> priority . one_by_one ( serialized ) <EOL> return serialized <EOL> def check ( self ) : <EOL> if not self . graph . is_acyclic ( ) : <EOL> err = "<STR_LIT>" <EOL> err += '<STR_LIT:U+002CU+0020>' . join ( six . moves . map ( str , <EOL> nx . simple_cycles ( <EOL> nx . DiGraph ( self . graph ) ) ) ) <EOL> err += '<STR_LIT:\n>' <EOL> raise errors . InvalidData ( err ) <EOL> non_existing_tasks = [ ] <EOL> invalid_tasks = [ ] <EOL> for node_key , node_value in six . iteritems ( self . graph . node ) : <EOL> if not node_value . get ( '<STR_LIT:id>' ) : <EOL> successors = self . graph . successors ( node_key ) <EOL> predecessors = self . graph . predecessors ( node_key ) <EOL> neighbors = successors + predecessors <EOL> non_existing_tasks . append ( node_key ) <EOL> invalid_tasks . extend ( neighbors ) <EOL> if non_existing_tasks : <EOL> raise errors . InvalidData ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" . format ( <EOL> non_existing_tasks = '<STR_LIT:U+002CU+0020>' . join ( <EOL> str ( x ) for x in sorted ( non_existing_tasks ) ) , <EOL> invalid_tasks = '<STR_LIT:U+002CU+0020>' . join ( <EOL> str ( x ) for x in sorted ( set ( invalid_tasks ) ) ) ) ) </s>
