def merged_dag_content(
ep_proto: 'EndpointProtocol',
components: Dict[str, 'ModelComponent']) -> 'MergedJSON':
init = _process_initial(ep_proto, components)
dsk_connections = connections_from_components_map(components)
epjson = endpoint_protocol_content(ep_proto)
merged = {**init.merged_dsk, **init.payload_tasks_dsk}
dependencies, _ = get_deps(merged)
merged_proto = defaultdict(list)
for task_name, task in merged.items():
for parent in dependencies[task_name]:
merged_proto[task_name].append(parent)
for request_name, task_key in init.payload_dsk_map.items():
cluster, *_ = task_key.split(".")
merged_proto[task_key[:-len(".serial")]].append(task_key)
merged_proto[task_key].append(request_name)
merged_proto = dict(merged_proto)
dependencies, dependents = get_deps(merged_proto)
dependents = dict(dependents)
functions_merged = valmap(functions_of, merged)
function_names_merged = {
k: tuple(map(funcname, v))
for k, v in functions_merged.items()
}
return MergedJSON(
dependencies=dependencies,
dependents=dependents,
funcnames=function_names_merged,
connections=dsk_connections,
endpoint=epjson,
)
def test_get_deps():
"""
>>> dsk = {'a': 1, 'b': (inc, 'a'), 'c': (inc, 'b')}
>>> dependencies, dependents = get_deps(dsk)
>>> dependencies
{'a': set(), 'b': {'a'}, 'c': {'b'}}
>>> dict(dependents)
{'a': {'b'}, 'b': {'c'}, 'c': set()}
"""
dsk = {
"a": [1, 2, 3],
"b": "a",
"c": [1, (inc, 1)],
"d": [(sum, "c")],
"e": ["b", "zzz", "b"],
"f": [["a", "b"], 2, 3],
}
dependencies, dependents = get_deps(dsk)
assert dependencies == {
"a": set(),
"b": {"a"},
"c": set(),
"d": {"c"},
"e": {"b"},
"f": {"a", "b"},
}
assert dependents == {
"a": {"b", "f"},
"b": {"e", "f"},
"c": {"d"},
"d": set(),
"e": set(),
"f": set(),
}
def component_dag_content(components: Dict[str, "ModelComponent"]) -> "ComponentJSON":
dsk_connections = connections_from_components_map(components)
comp_dependencies, comp_dependents, comp_funcnames = {}, {}, {}
for comp_name, comp in components.items():
functions_comp = valmap(functions_of, comp._flashserve_meta_.dsk)
function_names_comp = {k: sorted(set(map(funcname, v))) for k, v in functions_comp.items()}
comp_funcnames[comp_name] = function_names_comp
_dependencies, _dependents = get_deps(comp._flashserve_meta_.dsk)
_dependents = dict(_dependents)
comp_dependencies[comp_name] = _dependencies
comp_dependents[comp_name] = _dependents
return ComponentJSON(
component_dependencies=comp_dependencies,
component_dependents=comp_dependents,
component_funcnames=comp_funcnames,
connections=dsk_connections,
)
def visualize(
tc: 'TaskComposition',
fhandle: BytesIO = None,
format: str = "png",
*,
no_optimization: bool = False,
):
"""Visualize a graph"""
dsk = tc.pre_optimization_dsk if no_optimization else tc.dsk
dependencies, dependents = get_deps(dsk)
g = _dag_to_graphviz(
dag=dsk,
dependencies=dependencies,
request_data=tc.ep_dsk_input_keys,
response_data=tc.ep_dsk_output_keys,
no_optimization=no_optimization,
)
if fhandle is not None:
data = g.pipe(format=format)
fhandle.seek(0)
fhandle.write(data)
return
return g
def test_stacklimit(abcde):
dsk = dict(("x%s" % (i + 1), (inc, "x%s" % i)) for i in range(10000))
dependencies, dependents = get_deps(dsk)
ndependencies(dependencies, dependents)
def test_stacklimit(abcde):
dsk = {"x%s" % (i + 1): (inc, "x%s" % i) for i in range(10000)}
dependencies, dependents = get_deps(dsk)
ndependencies(dependencies, dependents)