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