()
| 61 | |
| 62 | |
| 63 | def test_hlg_expr_sequence_finalize(): |
| 64 | hlgx = HighLevelGraph( |
| 65 | xlayer := {"xlayer": MaterializedLayer({"x": DataNode("x", 1)})}, |
| 66 | dependencies=(xdeps := {"xlayer": set()}), |
| 67 | ) |
| 68 | ylayer = {"ylayer": MaterializedLayer({"y": Task("y", func, TaskRef("x"))})} |
| 69 | ylayer.update(xlayer) |
| 70 | ydeps = {"ylayer": {"xlayer"}} |
| 71 | ydeps.update(xdeps) |
| 72 | hlgy = HighLevelGraph(ylayer, dependencies=ydeps) |
| 73 | zlayer = {"zlayer": MaterializedLayer({"z": Task("z", func, TaskRef("x"))})} |
| 74 | zlayer.update(xlayer) |
| 75 | zdeps = {"zlayer": {"xlayer"}} |
| 76 | |
| 77 | zdeps.update(xdeps) |
| 78 | hlgz = HighLevelGraph(zlayer, dependencies=zdeps) |
| 79 | hlgexprx = HLGExpr( |
| 80 | hlgx, |
| 81 | low_level_optimizer=optimizer, |
| 82 | output_keys=["x"], |
| 83 | ) |
| 84 | hlgexpry = HLGExpr( |
| 85 | hlgy, |
| 86 | low_level_optimizer=optimizer, |
| 87 | output_keys=["y"], |
| 88 | ) |
| 89 | hlgexprz = HLGExpr( |
| 90 | hlgz, |
| 91 | low_level_optimizer=optimizer, |
| 92 | output_keys=["z"], |
| 93 | ) |
| 94 | dskx = hlgexprx.finalize_compute().optimize().__dask_graph__() |
| 95 | assert isinstance(dskx, dict) |
| 96 | assert len(dskx) == 1 |
| 97 | assert "x" in dskx |
| 98 | assert dskx["x"] is hlgy.layers["xlayer"]["x"] |
| 99 | |
| 100 | dsky = hlgexpry.finalize_compute().optimize().__dask_graph__() |
| 101 | assert isinstance(dsky, dict) |
| 102 | # Linear low level fusion |
| 103 | assert len(dsky) == 1 |
| 104 | assert "y" in dsky |
| 105 | assert dsky["y"] != hlgy.layers["ylayer"]["y"] |
| 106 | |
| 107 | expryz_opt = _ExprSequence(hlgexprz, hlgexpry).finalize_compute().optimize() |
| 108 | keys_yz = expryz_opt.__dask_keys__() |
| 109 | assert len(keys_yz) == 2 |
| 110 | |
| 111 | dskyz = expryz_opt.__dask_graph__() |
| 112 | assert isinstance(dskyz, dict) |
| 113 | expected = {} |
| 114 | expected.update(next(iter(hlgx.layers.values())).mapping) |
| 115 | expected.update(next(iter(hlgy.layers.values())).mapping) |
| 116 | expected.update(next(iter(hlgz.layers.values())).mapping) |
| 117 | # This is building the graph properly without fusing anything |
| 118 | assert dskyz == expected |
| 119 | |
| 120 | hlgexprz_different_optimizer = HLGExpr( |
nothing calls this directly
no test coverage detected