def _add_logging(dsk, ignore=None): """ Add logging to a Dask graph. @param dsk: The Dask graph. @return: New Dask graph. """ ctx = current_action() result = {} # Use topological sort to ensure Eliot actions are in logical order of # execution in Dask: keys = toposort(dsk) # Give each key a string name. Some keys are just aliases to other # keys, so make sure we have underlying key available. Later on might # want to shorten them as well. def simplify(k): if isinstance(k, str): return k return "-".join(str(o) for o in k) key_names = {} for key in keys: value = dsk[key] if not callable(value) and value in keys: # It's an alias for another key: key_names[key] = key_names[value] else: key_names[key] = simplify(key) # 2. Create Eliot child Actions for each key, in topological order: key_to_action_id = { key: str(ctx.serialize_task_id(), "utf-8") for key in keys } # 3. Replace function with wrapper that logs appropriate Action: for key in keys: func = dsk[key][0] args = dsk[key][1:] if not callable(func): # This key is just an alias for another key, no need to add # logging: result[key] = dsk[key] continue wrapped_func = _RunWithEliotContext( task_id=key_to_action_id[key], func=func, key=key_names[key], dependencies=[key_names[k] for k in get_dependencies(dsk, key)], ) result[key] = (wrapped_func, ) + tuple(args) assert result.keys() == dsk.keys() return result
def _add_logging(dsk, ignore=None): """ Add logging to a Dask graph. @param dsk: The Dask graph. @return: New Dask graph. """ ctx = current_action() result = {} # Use topological sort to ensure Eliot actions are in logical order of # execution in Dask: keys = toposort(dsk) # Give each key a string name. Some keys are just aliases to other # keys, so make sure we have underlying key available. Later on might # want to shorten them as well. def simplify(k): if isinstance(k, str): return k return "-".join(str(o) for o in k) key_names = {} for key in keys: value = dsk[key] if not callable(value) and value in keys: # It's an alias for another key: key_names[key] = key_names[value] else: key_names[key] = simplify(key) # 2. Create Eliot child Actions for each key, in topological order: key_to_action_id = {key: str(ctx.serialize_task_id(), "utf-8") for key in keys} # 3. Replace function with wrapper that logs appropriate Action: for key in keys: func = dsk[key][0] args = dsk[key][1:] if not callable(func): # This key is just an alias for another key, no need to add # logging: result[key] = dsk[key] continue wrapped_func = _RunWithEliotContext( task_id=key_to_action_id[key], func=func, key=key_names[key], dependencies=[key_names[k] for k in get_dependencies(dsk, key)], ) result[key] = (wrapped_func,) + tuple(args) assert result.keys() == dsk.keys() return result
def compile_subgraph(self, subgraph: Dict, inputs: List[Hashable], output: Hashable) -> Callable: """Fuse a subgraph of tasks into a single compiled function. It is assumed that the function will be called with values corresponding to `inputs` in the order they are given. """ tbl = SymbolTable() # must populate the symbol table in toposort order toposort_keys = list(toposort(subgraph)) # register the inputs as variables for key in inputs: tbl.register_var(key) # register each function in the subgraph for key in toposort_keys: task = subgraph[key] # all metagraph tasks are in (func, args, kwargs) format delayed_algo, args, kwargs = task if isinstance(kwargs, tuple): # FIXME: why are dictionaries represented this way in the DAG? kwargs = kwargs[0](kwargs[1]) if len(kwargs) != 0: raise CompileError( "NumbaCompiler only supports functions with bound kwargs.\n" f"When compiling:\n{delayed_algo.func_label}\nfound unbound kwargs:\n{kwargs}" ) # for maximum optimization, inline every task function during # compilation of the wrapper jit_func = numba.jit(inline="always")(delayed_algo.algo.func) tbl.register_func(key, jit_func, args) # generate the wrapper subgraph_wrapper_name = "subgraph" + str(self._subgraph_count) self._subgraph_count += 1 wrapper_text, wrapper_globals = construct_call_wrapper_text( wrapper_name=subgraph_wrapper_name, symbol_table=tbl, input_keys=inputs, execute_keys=toposort_keys, output_key=output, ) wrapper_func = compile_wrapper(subgraph_wrapper_name, wrapper_text, wrapper_globals) return wrapper_func
def inline(dsk, keys=None, inline_constants=True, dependencies=None): """Return new dask with the given keys inlined with their values. Inlines all constants if ``inline_constants`` keyword is True. Note that the constant keys will remain in the graph, to remove them follow ``inline`` with ``cull``. Examples -------- >>> def inc(x): ... return x + 1 >>> def add(x, y): ... return x + y >>> d = {'x': 1, 'y': (inc, 'x'), 'z': (add, 'x', 'y')} >>> inline(d) # doctest: +ELLIPSIS {'x': 1, 'y': (<function inc at ...>, 1), 'z': (<function add at ...>, 1, 'y')} >>> inline(d, keys='y') # doctest: +ELLIPSIS {'x': 1, 'y': (<function inc at ...>, 1), 'z': (<function add at ...>, 1, (<function inc at ...>, 1))} >>> inline(d, keys='y', inline_constants=False) # doctest: +ELLIPSIS {'x': 1, 'y': (<function inc at ...>, 'x'), 'z': (<function add at ...>, 'x', (<function inc at ...>, 'x'))} """ if dependencies and isinstance(next(iter(dependencies.values())), list): dependencies = {k: set(v) for k, v in dependencies.items()} keys = _flat_set(keys) if dependencies is None: dependencies = {k: get_dependencies(dsk, k) for k in dsk} if inline_constants: keys.update( k for k, v in dsk.items() if (ishashable(v) and v in dsk) or (not dependencies[k] and not istask(v)) ) # Keys may depend on other keys, so determine replace order with toposort. # The values stored in `keysubs` do not include other keys. replaceorder = toposort( {k: dsk[k] for k in keys if k in dsk}, dependencies=dependencies ) keysubs = {} for key in replaceorder: val = dsk[key] for dep in keys & dependencies[key]: if dep in keysubs: replace = keysubs[dep] else: replace = dsk[dep] val = subs(val, dep, replace) keysubs[key] = val # Make new dask with substitutions dsk2 = keysubs.copy() for key, val in dsk.items(): if key not in dsk2: for item in keys & dependencies[key]: val = subs(val, item, keysubs[item]) dsk2[key] = val return dsk2
def _add_logging(dsk, ignore=None): """ Add logging to a Dask graph. @param dsk: The Dask graph. @return: New Dask graph. """ ctx = current_action() result = {} # Use topological sort to ensure Eliot actions are in logical order of # execution in Dask: keys = toposort(dsk) # Give each key a string name. Some keys are just aliases to other # keys, so make sure we have underlying key available. Later on might # want to shorten them as well. def simplify(k): if isinstance(k, str): return k return "-".join(str(o) for o in k) key_names = {} for key in keys: value = dsk[key] if not callable(value) and ishashable(value) and value in keys: # It's an alias for another key: key_names[key] = key_names[value] else: key_names[key] = simplify(key) # Values in the graph can be either: # # 1. A list of other values. # 2. A tuple, where first value might be a callable, aka a task. # 3. A literal of some sort. def maybe_wrap(key, value): if isinstance(value, list): return [maybe_wrap(key, v) for v in value] elif isinstance(value, tuple): func = value[0] args = value[1:] if not callable(func): # Not a callable, so nothing to wrap. return value wrapped_func = _RunWithEliotContext( task_id=str(ctx.serialize_task_id(), "utf-8"), func=func, key=key_names[key], dependencies=[ key_names[k] for k in get_dependencies(dsk, key) ], ) return (wrapped_func, ) + args else: return value # Replace function with wrapper that logs appropriate Action; iterate in # topological order so action task levels are in reasonable order. for key in keys: result[key] = maybe_wrap(key, dsk[key]) assert set(result.keys()) == set(dsk.keys()) return result