<s> import collections <EOL> import copy <EOL> from nailgun import consts <EOL> __all__ = [ '<STR_LIT>' ] <EOL> class NumaNode ( object ) : <EOL> def __init__ ( self , id_ , memory ) : <EOL> """<STR_LIT>""" <EOL> self . id = id_ <EOL> self . free_memory = memory <EOL> self . pages = collections . defaultdict ( int ) <EOL> def allocate ( self , component ) : <EOL> """<STR_LIT>""" <EOL> if component . is_done ( ) : <EOL> return True <EOL> allocated = False <EOL> for page_size , _ in component . pages ( ) : <EOL> available = self . free_memory // page_size <EOL> allocated_pages = component . allocate ( page_size , available ) <EOL> if allocated_pages : <EOL> self . pages [ page_size ] += allocated_pages <EOL> self . free_memory -= page_size * allocated_pages <EOL> allocated = True <EOL> return allocated <EOL> def report ( self ) : <EOL> result = [ ] <EOL> for size , count in sorted ( self . pages . items ( ) ) : <EOL> result . append ( { '<STR_LIT:count>' : count , <EOL> '<STR_LIT>' : self . id , <EOL> '<STR_LIT:size>' : size } ) <EOL> return result <EOL> class Component ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , pages_info ) : <EOL> """<STR_LIT>""" <EOL> pages = [ ( int ( size ) , count ) <EOL> for size , count in pages_info . items ( ) if count != <NUM_LIT:0> ] <EOL> pages . sort ( reverse = True ) <EOL> self . _pages = collections . OrderedDict ( pages ) <EOL> def allocate ( self , page_size , pages_available ) : <EOL> """<STR_LIT>""" <EOL> pages = self . _pages <EOL> if page_size not in pages : <EOL> return <NUM_LIT:0> <EOL> required = pages [ page_size ] <EOL> limit = min ( required , pages_available ) <EOL> pages [ page_size ] -= limit <EOL> if pages [ page_size ] == <NUM_LIT:0> : <EOL> pages . pop ( page_size ) <EOL> return limit <EOL> def pages ( self ) : <EOL> return self . _pages . items ( ) <EOL> def is_done ( self ) : <EOL> return not self . _pages <EOL> def distribute_hugepages ( numa_topology , components ) : <EOL> all_comps = [ Component ( comp ) for comp in components [ '<STR_LIT:all>' ] ] <EOL> any_comps = [ Component ( comp ) for comp in components [ '<STR_LIT>' ] ] <EOL> numa_nodes = [ ] <EOL> for numa_node in numa_topology [ '<STR_LIT>' ] : <EOL> memory = numa_node [ '<STR_LIT>' ] // <NUM_LIT> <EOL> if numa_node [ '<STR_LIT:id>' ] == <NUM_LIT:0> : <EOL> memory -= consts . MEMORY_RESERVED_FOR_OPERATING_SYSTEM // <NUM_LIT> <EOL> numa_nodes . append ( NumaNode ( numa_node [ '<STR_LIT:id>' ] , memory ) ) <EOL> _allocate_all ( numa_nodes , all_comps ) <EOL> _allocate_any ( numa_nodes , any_comps ) <EOL> return sum ( [ n . report ( ) for n in numa_nodes ] , [ ] ) <EOL> def _allocate_all ( numa_nodes , components ) : <EOL> """<STR_LIT>""" <EOL> for component in components : <EOL> for numa_node in numa_nodes : <EOL> comp = copy . deepcopy ( component ) <EOL> allocated = numa_node . allocate ( comp ) <EOL> if not ( allocated and comp . is_done ( ) ) : <EOL> raise ValueError ( '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> def _allocate_any ( numa_nodes , components ) : <EOL> """<STR_LIT>""" <EOL> comp = _merge_components ( components ) <EOL> for numa_node in numa_nodes : <EOL> numa_node . allocate ( comp ) <EOL> if comp . is_done ( ) : <EOL> break <EOL> else : <EOL> raise ValueError ( '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> def _merge_components ( components ) : <EOL> res = collections . defaultdict ( int ) <EOL> for comp in components : <EOL> for size , count in comp . pages ( ) : <EOL> res [ size ] += count <EOL> return Component ( res ) </s>
