📦 Shareable, Scalable Python Environments¶
Sourcery Review
## Reviewer's Guide This PR reorganizes the utils package, adds robust distributed tarball transfer scripts and CLI, and refactors the core launch/PBS routines for more flexible command construction, parameter validation, and enhanced logging/filtering. #### Sequence diagram for distributed tarball transfer and extractionsequenceDiagram
participant User as actor
participant YeetEnv
participant YeetTarball
participant ezpz
User->>YeetEnv: Run ezpz-yeet-env CLI
YeetEnv->>ezpz: setup_torch()
YeetEnv->>ezpz.pbs: get_pbs_launch_cmd(ngpu_per_host=1)
YeetEnv->>ezpz.launch: launch(cmd_to_launch=python3 -m ezpz.utils.yeet_tarball ...)
ezpz.launch->>YeetTarball: python3 -m ezpz.utils.yeet_tarball --src ... --dst ... --d
YeetTarball->>ezpz: get_rank(), get_world_size()
YeetTarball->>YeetTarball: Transfer(src, dst, decompress, chunk_size, flags)
YeetTarball->>ezpz.dist: broadcast(data, root=0)
YeetTarball->>OS: Write file, decompress tarball
classDiagram
class YeetTarball {
+parse_args() argparse.Namespace
+bcast_chunk(A, chunk_size)
+Transfer(src, dst, decompress, chunk_size, flags)
+main()
}
class YeetEnv {
+main()
}
YeetEnv --> YeetTarball : uses
YeetTarball ..> ezpz : depends on
YeetEnv ..> ezpz : depends on
YeetEnv ..> ezpz.pbs : uses
YeetEnv ..> ezpz.launch : uses
classDiagram
class Launch {
+launch(launch_cmd, cmd_to_launch, filters, include_python)
+build_launch_cmd(hostfile)
+get_aurora_filters(additional_filters)
+run_command(command, filters)
}
class PBS {
+get_pbs_launch_cmd(ngpus, ngpu_per_host, nhosts, hostfile)
+build_launch_cmd(hostfile)
+get_pbs_env()
}
Launch ..> PBS : uses
Launch ..> ezpz : uses
PBS ..> ezpz : uses
classDiagram
class Dist {
+setup_torch_DDP(rank, world_size, timeout, backend)
}
Dist ..> torch.distributed : uses
- Introduce yeet_tarball.py with chunked broadcast and extraction logic
- Add yeet_env.py as a CLI entrypoint for launching distributed transfers
- Register ezpz-yeet-env in pyproject.toml
- Restructure utils package by renaming utils.py to utils/__init__.py
`src/ezpz/utils/yeet_env.py`
`pyproject.toml`
`src/ezpz/utils/__init__.py` | | Refactor PBS launch command computation with robust parameters and logging |
- Compute ngpus, nhosts, and ngpu_per_host with validation assertions and fallbacks
- Unify hostfile fallback resolution
- Add info/warning logs about GPU usage and mismatches
- Update build_launch_cmd to accept a hostfile argument and delegate to get_pbs_launch_cmd
- Introduce EZPZ_LOG_LEVEL and adapt filtering logic in get_aurora_filters
- Log number of filters in run_command
- Extend filtering to include Sunspot and conditional debug behavior
- Update launch() signature to accept launch_cmd and include_python flag
- Dynamically prepend Python executable only when necessary
- Import torch.distributed and guard init_process_group with is_initialized() check
- Modify test.py to call launch(cmd_to_launch=...) instead of positional args
Tips and commands
#### Interacting with Sourcery - **Trigger a new review:** Comment `@sourcery-ai review` on the pull request. - **Continue discussions:** Reply directly to Sourcery's review comments. - **Generate a GitHub issue from a review comment:** Ask Sourcery to create an issue from a review comment by replying to it. You can also reply to a review comment with `@sourcery-ai issue` to create an issue from it. - **Generate a pull request title:** Write `@sourcery-ai` anywhere in the pull request title to generate a title at any time. You can also comment `@sourcery-ai title` on the pull request to (re-)generate the title at any time. - **Generate a pull request summary:** Write `@sourcery-ai summary` anywhere in the pull request body to generate a PR summary at any time exactly where you want it. You can also comment `@sourcery-ai summary` on the pull request to (re-)generate the summary at any time. - **Generate reviewer's guide:** Comment `@sourcery-ai guide` on the pull request to (re-)generate the reviewer's guide at any time. - **Resolve all Sourcery comments:** Comment `@sourcery-ai resolve` on the pull request to resolve all Sourcery comments. Useful if you've already addressed all the comments and don't want to see them anymore. - **Dismiss all Sourcery reviews:** Comment `@sourcery-ai dismiss` on the pull request to dismiss all existing Sourcery reviews. Especially useful if you want to start fresh with a new review - don't forget to comment `@sourcery-ai review` to trigger a new review! #### Customizing Your Experience Access your [dashboard](https://app.sourcery.ai) to: - Enable or disable review features such as the Sourcery-generated pull request summary, the reviewer's guide, and others. - Change the review language. - Add, remove or edit custom review instructions. - Adjust other review settings. #### Getting Help - [Contact our support team](mailto:support@sourcery.ai) for questions or feedback. - Visit our [documentation](https://docs.sourcery.ai) for detailed guides and information. - Keep in touch with the Sourcery team by following us on [X/Twitter](https://x.com/SourceryAI), [LinkedIn](https://www.linkedin.com/company/sourcery-ai/) or [GitHub](https://github.com/sourcery-ai).👀 Overview¶
We provide mechanisms for:
ezpz-tar-env: Creating tarball from {conda, virtual} environmentezpz-yeet-env: Copying an environment tarball to/tmp/on all worker nodes
🔍 Details¶
ezpz-tar-env: Creating tarball from {conda, virtual} environment
By default, this will try and create a tarball from the environment which it was invoked and place it in the working directory.
ezpz-yeet-env: Copying an environment tarball to/tmp/on all worker nodes
Following from (1.):
Example:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | |