Backport patching functionality from 24.06#41
Merged
raydouglass merged 3 commits intorapidsai:branch-24.04from Apr 5, 2024
Merged
Backport patching functionality from 24.06#41raydouglass merged 3 commits intorapidsai:branch-24.04from
raydouglass merged 3 commits intorapidsai:branch-24.04from
Conversation
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
bdice
approved these changes
Apr 4, 2024
Contributor
bdice
left a comment
There was a problem hiding this comment.
Let's backport as-is, and tweak in follow-up PRs if needed.
Member
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 |
rjzamora
approved these changes
Apr 4, 2024
Contributor
|
@rjzamora I agree we should merge. I’ll request an admin merge from ops since the branch is frozen. |
Member
|
Thanks @vyasr @bdice and @raydouglass ! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR backports #27, #37, and #39 to 24.04