Skip to content

Fixes KeyError for get_two_hop_neighbors when called with a small start vertices list#3778

Merged
rapids-bot[bot] merged 2 commits intorapidsai:branch-23.10from
rlratzel:branch-23.10-mg_keyerror_fix
Aug 12, 2023
Merged

Fixes KeyError for get_two_hop_neighbors when called with a small start vertices list#3778
rapids-bot[bot] merged 2 commits intorapidsai:branch-23.10from
rlratzel:branch-23.10-mg_keyerror_fix

Conversation

@rlratzel
Copy link
Contributor

@rlratzel rlratzel commented Aug 10, 2023

closes #3745

This PR adds updates to replace the get_distributed_data() call with persist_dask_df_equal_parts_per_worker() and get_persisted_df_worker_map() to avoid a problem where get_distributed_data() does not distribute data properly across all workers. This resulted in a KeyError when the data was accessed via worker, when that worker was not a key in the map.

More details are in the linked issue.

This PR also does minor refactoring in get_two_hop_neighbors() and reorganizes the imports according to PEP 8.

Tested manually on a 4-GPU system, where the problem described in #3745 was reproduced, the change in the PR applied and re-run, and the error no longer occurring.

@rlratzel rlratzel added bug Something isn't working non-breaking Non-breaking change labels Aug 10, 2023
@rlratzel rlratzel added this to the 23.10 milestone Aug 10, 2023
@rlratzel rlratzel self-assigned this Aug 10, 2023
…_per_worker() and get_persisted_df_worker_map(), minor refactor to get_two_hop_neighbors(), clean up imports
@rlratzel rlratzel marked this pull request as ready for review August 11, 2023 00:15
@rlratzel rlratzel requested a review from a team as a code owner August 11, 2023 00:15
@rlratzel rlratzel requested a review from VibhuJawa August 11, 2023 00:15
Copy link
Member

@VibhuJawa VibhuJawa left a comment

Choose a reason for hiding this comment

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

LGTM

@BradReesWork
Copy link
Member

/merge

@rapids-bot rapids-bot bot merged commit a1a85a7 into rapidsai:branch-23.10 Aug 12, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working non-breaking Non-breaking change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

MG get_two_hop_neighbors fails with KeyError when accessing start_vertices

3 participants