def import_global(module, check_depth=True, block=True): try: # We should only be doing something for this if we're the top-level task if c.legion_task_get_depth(top_level.task[0]) > 0 and check_depth: return None except AttributeError: raise RuntimeError( '"import_global" must be called in a legion_python task') if isinstance(module, str): name = module elif isinstance(module, unicode): name = module elif isinstance(module, types.ModuleType): name = module.__name__ else: raise TypeError( '"module" arg to "import_global" must be a ModuleType or str type') mapper = c.legion_runtime_generate_library_mapper_ids( top_level.runtime[0], _unique_name.encode('utf-8'), 1) future = c.legion_runtime_select_tunable_value(top_level.runtime[0], top_level.context[0], 0, mapper, 0) num_python_procs = struct.unpack_from( 'i', ffi.buffer(c.legion_future_get_untyped_pointer(future), 4))[0] c.legion_future_destroy(future) assert num_python_procs > 0 # Launch an index space task across all the python # processors to import the module in every interpreter task_id = c.legion_runtime_generate_library_task_ids( top_level.runtime[0], _unique_name.encode('utf-8'), 3) + 1 rect = ffi.new('legion_rect_1d_t *') rect[0].lo.x[0] = 0 rect[0].hi.x[0] = num_python_procs - 1 domain = c.legion_domain_from_rect_1d(rect[0]) packed = name.encode('utf-8') arglen = len(packed) array = ffi.new('char[]', arglen) ffi.buffer(array, arglen)[:] = packed args = ffi.new('legion_task_argument_t *') args[0].args = array args[0].arglen = arglen argmap = c.legion_argument_map_create() launcher = c.legion_index_launcher_create(task_id, domain, args[0], argmap, c.legion_predicate_true(), False, mapper, 0) future = c.legion_index_launcher_execute_reduction( top_level.runtime[0], top_level.context[0], launcher, c.LEGION_REDOP_SUM_INT32) c.legion_index_launcher_destroy(launcher) c.legion_argument_map_destroy(argmap) if block: result = struct.unpack_from( 'i', ffi.buffer(c.legion_future_get_untyped_pointer(future), 4))[0] c.legion_future_destroy(future) if result > 0: raise ImportError('failed to globally import ' + name + ' on ' + str(result) + ' nodes') return None else: return future
def legion_python_main(raw_args, user_data, proc): raw_arg_ptr = ffi.new('char[]', bytes(raw_args)) raw_arg_size = len(raw_args) # Execute preamble to obtain Legion API context. task = ffi.new('legion_task_t *') raw_regions = ffi.new('legion_physical_region_t **') num_regions = ffi.new('unsigned *') context = ffi.new('legion_context_t *') runtime = ffi.new('legion_runtime_t *') c.legion_task_preamble( raw_arg_ptr, raw_arg_size, proc, task, raw_regions, num_regions, context, runtime) top_level.runtime, top_level.context, top_level.task = runtime, context, task # Run user's script. args = input_args(True) start = 1 if len(args) > 1 and args[1] == '--nocr': start += 1 local_cleanup = False else: local_cleanup = True if len(args) < (start+1): sys.argv = [''] run_repl() elif args[start] == '-': sys.argv = args[start:] run_repl() elif args[start] == '-h': print('usage: legion_python [--nocr] [-c cmd | -m mod | file | -] [arg] ...') print('--nocr : disable control replication for multi-node execution') print('-c cmd : program passed in as string (terminates option list)') print('-h : print this help message and exit') print('-m mod : run library module as a script (terminates option list)') print('file : program read from script file') print('- : program read from stdin (default; interactive mode if a tty)') print('arg ...: arguments passed to program in sys.argv[1:]') elif args[start] == '-c': if len(args) > (start+1): sys.argv = ['-c'] + list(args[start+2:]) run_cmd(args[start+1], run_name='__main__') else: print('Argument expected for the -c option') c.legion_runtime_set_return_code(1) elif args[start] == '-m': if len(args) > (start+1): filename = args[start+1] + '.py' found = False for path in sys.path: for root,dirs,files in os.walk(path): if filename not in files: continue module = os.path.join(root, filename) sys.argv = [module] + list(args[start+2:]) run_path(module, run_name='__main__') found = True break if found: break if not found: print('No module named '+args[start+1]) c.legion_runtime_set_return_code(1) else: print('Argument expected for the -m option') c.legion_runtime_set_return_code(1) else: assert start < len(args) sys.argv = list(args[start:]) run_path(args[start], run_name='__main__') if local_cleanup: # If we were control replicated then we just need to do our cleanup # Do it in reverse order so modules get FILO properties for cleanup in reversed(cleanup_items): cleanup() else: # Otherwise, run a task on every node to perform the cleanup mapper = c.legion_runtime_generate_library_mapper_ids( top_level.runtime[0], _unique_name.encode('utf-8'), 1) future = c.legion_runtime_select_tunable_value( top_level.runtime[0], top_level.context[0], 0, mapper, 0) num_python_procs = struct.unpack_from('i', ffi.buffer(c.legion_future_get_untyped_pointer(future),4))[0] c.legion_future_destroy(future) assert num_python_procs > 0 # Launch an index space task across all the python # processors to import the module in every interpreter task_id = c.legion_runtime_generate_library_task_ids( top_level.runtime[0], _unique_name.encode('utf-8'), 3) + 2 rect = ffi.new('legion_rect_1d_t *') rect[0].lo.x[0] = 0 rect[0].hi.x[0] = num_python_procs - 1 domain = c.legion_domain_from_rect_1d(rect[0]) args = ffi.new('legion_task_argument_t *') args[0].args = ffi.NULL args[0].arglen = 0 argmap = c.legion_argument_map_create() launcher = c.legion_index_launcher_create(task_id, domain, args[0], argmap, c.legion_predicate_true(), False, mapper, 0) future_map = c.legion_index_launcher_execute(top_level.runtime[0], top_level.context[0], launcher) c.legion_index_launcher_destroy(launcher) c.legion_argument_map_destroy(argmap) # Wait for all the cleanup tasks to be done c.legion_future_map_wait_all_results(future_map) c.legion_future_map_destroy(future_map) del top_level.runtime del top_level.context del top_level.task # Force a garbage collection so that we know that all objects which can # be collected are actually collected before we exit the top-level task gc.collect() # Execute postamble. c.legion_task_postamble(runtime[0], context[0], ffi.NULL, 0)
def legion_python_main(raw_args, user_data, proc): raw_arg_ptr = ffi.new('char[]', bytes(raw_args)) raw_arg_size = len(raw_args) # Execute preamble to obtain Legion API context. task = ffi.new('legion_task_t *') raw_regions = ffi.new('legion_physical_region_t **') num_regions = ffi.new('unsigned *') context = ffi.new('legion_context_t *') runtime = ffi.new('legion_runtime_t *') c.legion_task_preamble( raw_arg_ptr, raw_arg_size, proc, task, raw_regions, num_regions, context, runtime) top_level.runtime, top_level.context, top_level.task = runtime, context, task # Run user's script. args = input_args(True) start = 1 if len(args) > 1 and args[1] == '--nocr': start += 1 local_cleanup = False else: local_cleanup = True if len(args) < (start+1): sys.argv = [''] run_repl() elif args[start] == '-': sys.argv = args[start:] run_repl() elif args[start] == '-c': assert len(args) >= 3 sys.argv = ['-c'] + list(args[start+2:]) run_cmd(args[start+1], run_name='__main__') else: assert start < len(args) sys.argv = list(args[start:]) run_path(args[start], run_name='__main__') if local_cleanup: # If we were control replicated then we just need to do our cleanup for cleanup in cleanup_items: cleanup() else: # Otherwise, run a task on every node to perform the cleanup mapper = c.legion_runtime_generate_library_mapper_ids( top_level.runtime[0], _unique_name.encode('utf-8'), 1) future = c.legion_runtime_select_tunable_value( top_level.runtime[0], top_level.context[0], 0, mapper, 0) num_python_procs = struct.unpack_from('i', ffi.buffer(c.legion_future_get_untyped_pointer(future),4))[0] c.legion_future_destroy(future) assert num_python_procs > 0 # Launch an index space task across all the python # processors to import the module in every interpreter task_id = c.legion_runtime_generate_library_task_ids( top_level.runtime[0], _unique_name.encode('utf-8'), 3) + 2 rect = ffi.new('legion_rect_1d_t *') rect[0].lo.x[0] = 0 rect[0].hi.x[0] = num_python_procs - 1 domain = c.legion_domain_from_rect_1d(rect[0]) args = ffi.new('legion_task_argument_t *') args[0].args = ffi.NULL args[0].arglen = 0 argmap = c.legion_argument_map_create() launcher = c.legion_index_launcher_create(task_id, domain, args[0], argmap, c.legion_predicate_true(), False, mapper, 0) future_map = c.legion_index_launcher_execute(top_level.runtime[0], top_level.context[0], launcher) c.legion_index_launcher_destroy(launcher) c.legion_argument_map_destroy(argmap) # Wait for all the cleanup tasks to be done c.legion_future_map_wait_all_results(future_map) c.legion_future_map_destroy(future_map) del top_level.runtime del top_level.context del top_level.task # Force a garbage collection so that we know that all objects which can # be collected are actually collected before we exit the top-level task gc.collect() # Execute postamble. c.legion_task_postamble(runtime[0], context[0], ffi.NULL, 0)