mirror of
https://github.com/vale981/ray
synced 2025-03-06 10:31:39 -05:00
d699351748
This uses the generators introduced in #25247 to reduce memory usage during the merge stage in push-based shuffle. These tasks merge groups of map outputs, so it fits a generator pattern where we want to return merged outputs one at a time. Verified that this allows for merging more/larger objects at a time than the current list-based version. I also tried this for the map stage in random_shuffle, but it didn't seem to make a difference in memory usage for Arrow blocks. I think this is probably because Arrow is already doing some zero-copy optimizations when selecting rows? Also adds a new line to Dataset stats for memory usage. Unfortunately it's hard to get an accurate reading of physical memory usage in Python and this value will probably be an overestimate in a lot of cases. I didn't see a difference before and after this PR for the merge stage, for example. Arguably this field should be opt-in. For 100MB partitions, for example: ``` Substage 0 read->random_shuffle_map: 10/10 blocks executed * Remote wall time: 1.44s min, 3.32s max, 2.57s mean, 25.74s total * Remote cpu time: 1.42s min, 2.53s max, 2.03s mean, 20.25s total * Worker memory usage (MB): 462 min, 864 max, 552 mean * Output num rows: 12500000 min, 12500000 max, 12500000 mean, 125000000 total * Output size bytes: 101562500 min, 101562500 max, 101562500 mean, 1015625000 total * Tasks per node: 10 min, 10 max, 10 mean; 1 nodes used Substage 1 random_shuffle_reduce: 10/10 blocks executed * Remote wall time: 1.47s min, 2.94s max, 2.17s mean, 21.69s total * Remote cpu time: 1.45s min, 1.88s max, 1.71s mean, 17.09s total * Worker memory usage (MB): 462 min, 1047 max, 831 mean * Output num rows: 12500000 min, 12500000 max, 12500000 mean, 125000000 total * Output size bytes: 101562500 min, 101562500 max, 101562500 mean, 1015625000 total * Tasks per node: 10 min, 10 max, 10 mean; 1 nodes used ``` ## Checks - [ ] I've run `scripts/format.sh` to lint the changes in this PR. - [ ] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [ ] Unit tests - [ ] Release tests - [ ] This PR is not tested :( Co-authored-by: Eric Liang <ekhliang@gmail.com> |
||
|---|---|---|
| .. | ||
| ray | ||
| requirements | ||
| asv.conf.json | ||
| build-wheel-macos-arm64.sh | ||
| build-wheel-macos.sh | ||
| build-wheel-manylinux2014.sh | ||
| build-wheel-windows.sh | ||
| MANIFEST.in | ||
| README-building-wheels.md | ||
| requirements.txt | ||
| requirements_linters.txt | ||
| requirements_ml_docker.txt | ||
| setup.py | ||