| 859 | |
| 860 | |
| 861 | def test_linear_fusion(): |
| 862 | tasks = [ |
| 863 | io := Task("foo", func, 1), |
| 864 | second := Task("second", func, io.ref()), |
| 865 | Task("third", func, second.ref()), |
| 866 | ] |
| 867 | dsk = {t.key: t for t in tasks} |
| 868 | result = fuse_linear_task_spec(dsk, {"third"}) |
| 869 | assert len(result) == 2 |
| 870 | assert isinstance(result["third"], Alias) |
| 871 | assert ( |
| 872 | isinstance(result["foo-second-third"], Task) |
| 873 | and funcname(result["foo-second-third"].func) == "_execute_subgraph" |
| 874 | ) |
| 875 | |
| 876 | # Data Nodes don't get fused |
| 877 | |
| 878 | tasks = [ |
| 879 | io := DataNode("foo", 1), |
| 880 | second := Task("second", func, io.ref()), |
| 881 | third := Task("third", func, second.ref()), |
| 882 | Task("fourth", func, third.ref()), |
| 883 | ] |
| 884 | dsk = {t.key: t for t in tasks} |
| 885 | result = fuse_linear_task_spec(dsk, {"fourth"}) |
| 886 | assert len(result) == 2 |
| 887 | assert isinstance(result["fourth"], Alias) |
| 888 | assert ( |
| 889 | isinstance(result["foo-second-third-fourth"], Task) |
| 890 | and funcname(result["foo-second-third-fourth"].func) == "_execute_subgraph" |
| 891 | ) |
| 892 | assert "foo" not in result |
| 893 | |
| 894 | # Branch, so no fusion |
| 895 | tasks.append(Task("branch", func, third.ref())) |
| 896 | dsk = {t.key: t for t in tasks} |
| 897 | result = fuse_linear_task_spec(dsk, {"fourth"}) |
| 898 | assert len(result) == 4 |
| 899 | assert "foo-second-third" in result |
| 900 | assert isinstance(result["third"], Alias) |
| 901 | |
| 902 | # Branch, so no fusion at all |
| 903 | tasks.append(Task("branch2", func, second.ref())) |
| 904 | dsk = {t.key: t for t in tasks} |
| 905 | result = fuse_linear_task_spec(dsk, {"fourth"}) |
| 906 | assert len(result) == 6 |
| 907 | assert not any("-" in k for k in dsk) |
| 908 | |
| 909 | |
| 910 | def test_linear_fusion_intermediate_branch(): |