Skip to content

Backport patching functionality from 24.06#41

Merged
raydouglass merged 3 commits intorapidsai:branch-24.04from
vyasr:backport/patching_plus
Apr 5, 2024
Merged

Backport patching functionality from 24.06#41
raydouglass merged 3 commits intorapidsai:branch-24.04from
vyasr:backport/patching_plus

Conversation

@vyasr
Copy link
Copy Markdown
Contributor

@vyasr vyasr commented Apr 4, 2024

This PR backports #27, #37, and #39 to 24.04

vyasr and others added 3 commits April 4, 2024 20:01
This PR adds the ability to monkey-patch imports of dask and distributed
whenever those imports occur by simply installing
rapids-dask-dependency.

There's a tiny bit of scope creep here because this PR added real Python
code to the repo for the first time, so I also added pre-commit hooks
that in turn modified some unrelated files (only minimally, though).

TODO:
- [x] Update conda CI and packaging
- [ ] Stress test extensively

---------

Signed-off-by: Vyas Ramasubramani <vyasr@nvidia.com>
Co-authored-by: Richard (Rick) Zamora <rzamora217@gmail.com>
Currently, the tests added in rapidsai#27 (backported to 24.04 in rapidsai#36) do not
check the exit codes of their subprocesses. This means that failing
tests are not caught. This PR fixes the test utilities to check the exit
codes and print any stdout/stderr outputs.
This PR makes a number of significant changes to the patching infrastructure:
1. This PR reorganizes the patching logic to be based on a more principled approach. Rather than maintaining lists of patch functions that are each responsible for filtering modules to apply themselves to, patches are organized in the patches directory in a tree structure matching dask itself. Patches are found and run by importing the same relative paths within the `patches` directory corresponding to a particular dask or distributed module.
2. It adds proper support for patching submodules. Previously the loader was being disabled whenever a real dask module was being imported, but this is problematic because if some dask modules import others they will pre-populate `sys.modules` with the real modules and therefore the loader will never be used for loading a patched version of the submodule.
3. Patches are no longer just functions applied to modules, they are arbitrary functions executed when a module is imported. As a result, a wider range of modifications is possible than was previously allowed. In particular:
4. The more general functions allow the entire module import to be hijacked and redirected to other modules.
5. The new framework is used to vendor a patched version of the accessor.py module in dask, which resolves the issues observed in dask/dask#11035.

Authors:
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - Richard (Rick) Zamora (https://github.com/rjzamora)

URL: rapidsai#39
@vyasr vyasr added feature request New feature or request non-breaking Introduces a non-breaking change labels Apr 4, 2024
@vyasr vyasr self-assigned this Apr 4, 2024
@vyasr vyasr requested a review from a team as a code owner April 4, 2024 20:04
Copy link
Copy Markdown
Contributor

@bdice bdice left a comment

Choose a reason for hiding this comment

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

Let's backport as-is, and tweak in follow-up PRs if needed.

@rjzamora
Copy link
Copy Markdown
Member

rjzamora commented Apr 4, 2024

Let's backport as-is, and tweak in follow-up PRs if needed.

Seems fine to me. The easiest way to test is honestly to get this merged. The only tweak I have in mind is to avoid installing DaskLoader for python<3.11.9, but I'd like to test with the patching in place either way.

Copy link
Copy Markdown
Member

@rjzamora rjzamora left a comment

Choose a reason for hiding this comment

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

Local testing looks good to me. Thanks @vyasr !

I think it makes sense to merge this now. We could submit a test PR to cudf first to "check CI", but I don't think that saves us any time given that 24.04 is now frozen anyway. Thoughts @bdice ?

@bdice
Copy link
Copy Markdown
Contributor

bdice commented Apr 5, 2024

@rjzamora I agree we should merge. I’ll request an admin merge from ops since the branch is frozen.

@raydouglass raydouglass merged commit 99d37eb into rapidsai:branch-24.04 Apr 5, 2024
@rjzamora
Copy link
Copy Markdown
Member

rjzamora commented Apr 5, 2024

Thanks @vyasr @bdice and @raydouglass !

@vyasr vyasr deleted the backport/patching_plus branch July 5, 2024 01:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

feature request New feature or request non-breaking Introduces a non-breaking change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants