Skip to content

[DNM] Reduce memory footprint of P2P shuffle #8128

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
hendrikmakait wants to merge 2 commits into from
Closed

[DNM] Reduce memory footprint of P2P shuffle #8128

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
320de59
Reduce final footprint
hendrikmakait Aug 23, 2023
aa2711f
Experiment
hendrikmakait Aug 24, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions distributed/shuffle/_arrow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ def convert_partition(data: bytes, meta: pd.DataFrame) -> pd.DataFrame:
while file.tell() < end:
sr = pa.RecordBatchStreamReader(file)
shards.append(sr.read_all())
table = pa.concat_tables(shards, promote=True)
table = pa.concat_tables(shards, promote=True).combine_chunks()

df = from_pyarrow_table_dispatch(meta, table, self_destruct=True)
return df.astype(meta.dtypes, copy=False)
Expand All @@ -68,7 +68,7 @@ def list_of_buffers_to_table(data: list[bytes]) -> pa.Table:

return pa.concat_tables(
(deserialize_table(buffer) for buffer in data), promote=True
)
).combine_chunks()

@phofl phofl Aug 23, 2023

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This one isn’t necessary, only the one above

@hendrikmakait hendrikmakait Aug 23, 2023

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When I'm done, it hopefully will :)



def serialize_table(table: pa.Table) -> bytes:
Expand Down
9 changes: 5 additions & 4 deletions distributed/shuffle/_worker_plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,10 +104,11 @@ async def shuffle_inputs_done(self, shuffle_id: ShuffleId, run_id: int) -> None:
await shuffle.inputs_done()

async def _close_shuffle_run(self, shuffle: ShuffleRun) -> None:
await shuffle.close()
async with self._runs_cleanup_condition:
self._runs.remove(shuffle)
self._runs_cleanup_condition.notify_all()
pass
# await shuffle.close()
# async with self._runs_cleanup_condition:
# self._runs.remove(shuffle)
# self._runs_cleanup_condition.notify_all()

def shuffle_fail(self, shuffle_id: ShuffleId, run_id: int, message: str) -> None:
"""Fails the shuffle run with the message as exception and triggers cleanup.
Expand Down