nsys-jax: re-work to be more pip-install-able #1165
Conversation
This file contains 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
olupton
force-pushed
the
olupton/pip-install-nsys-jax
branch
2 times, most recently
from
1d0b327
to
80c4ca0
Compare
olupton
commented
Nov 25, 2024
There was a problem hiding this comment.
This one is new (but trivial)
There was a problem hiding this comment.
This isn't really new, it's the old install-protoc.
.github/container/nsys-jax
Outdated
There was a problem hiding this comment.
This isn't really gone, it became .github/container/nsys_jax/nsys_jax/scripts/nsys_jax.py
There was a problem hiding this comment.
This isn't really new, it's the old nsys-jax
Comment on lines
49
to
86
| install_script_template = r"""#!/bin/bash | ||
| # | ||
| # Usage: ./install.sh [optional arguments to virtualenv] | ||
| # | ||
| # If it doesn't already exist, this creates a virtual environment named | ||
| # `nsys_jax_env` in the current directory and installs Jupyter Lab and the | ||
| # dependencies of the Analysis.ipynb notebook that is shipped alongside this | ||
| # script inside the output archives of the `nsys-jax` wrapper. | ||
| # | ||
| # The expectation is that those archives will be copied and extracted on a | ||
| # laptop or workstation, and this installation script will be run there, while | ||
| # the `nsys-jax` wrapper is executed on a remote GPU cluster. | ||
| set -ex | ||
| SCRIPT_DIR=$(cd -- "$( dirname -- "${{BASH_SOURCE[0]}}" )" &> /dev/null && pwd) | ||
| VIRTUALENV="${{SCRIPT_DIR}}/nsys_jax_venv" | ||
| BIN="${{VIRTUALENV}}/bin" | ||
| if [[ ! -d "${{VIRTUALENV}}" ]]; then | ||
| # Let `virtualenv` find/choose a Python. Currently >=3.10 is supported. | ||
| virtualenv -p 3.13 -p 3.12 -p 3.11 -p 3.10 "$@" "${{VIRTUALENV}}" | ||
| "${{BIN}}/pip" install -U pip | ||
| "${{BIN}}/pip" install 'nsys-jax[jupyter] @ git+https://github.com/NVIDIA/JAX-Toolbox.git@{jax_toolbox_commit}#subdirectory=.github/container/nsys_jax' | ||
| "${{BIN}}/install-flamegraph" "${{VIRTUALENV}}" | ||
| "${{BIN}}/install-protoc" "${{VIRTUALENV}}" | ||
| else | ||
| echo "Virtual environment already exists, not installing anything..." | ||
| fi | ||
| if [ -z ${{NSYS_JAX_INSTALL_SKIP_LAUNCH+x}} ]; then | ||
| # Pick up the current profile data by default | ||
| export NSYS_JAX_DEFAULT_PREFIX="${{PWD}}" | ||
| # https://setuptools.pypa.io/en/latest/userguide/datafiles.html#accessing-data-files-at-runtime | ||
| NOTEBOOK=$("${{BIN}}/python" -c 'from importlib.resources import files; print(files("nsys_jax") / "analyses" / "Analysis.ipynb")') | ||
| echo "Launching: cd ${{SCRIPT_DIR}} && ${{BIN}}/jupyter-lab ${{NOTEBOOK}}" | ||
| cd "${{SCRIPT_DIR}}" && "${{BIN}}/jupyter-lab" "${{NOTEBOOK}}" | ||
| else | ||
| echo "Skipping launch of Jupyter Lab due to NSYS_JAX_INSTALL_SKIP_LAUNCH" | ||
| fi | ||
| """ |
There was a problem hiding this comment.
This replaces the old install.sh, otherwise the content of this file is mostly the old nsys-jax.
Comment on lines
+88
to
+96
| def create_install_script(output_queue): | ||
| """ | ||
| Write an install.sh to the output archive that installs nsys-jax at the same | ||
| version/commit that the current execution is using. | ||
| """ | ||
| jax_toolbox_sha = jax_toolbox_sha_with_prefix[1:] | ||
| install_script = install_script_template.format(jax_toolbox_commit=jax_toolbox_sha) | ||
| output_queue.put(("install.sh", install_script.encode(), COMPRESS_DEFLATE)) |
There was a problem hiding this comment.
This drives the install script generation. The 1: slice is because the hash is prefixed with g (for git).
| assert mirror_dir is not None | ||
| # Execute post-processing recipes and add any outputs to `ofile` | ||
| for analysis in args.analysis: | ||
| result, output_prefix = execute_analysis_script( |
There was a problem hiding this comment.
FWIW this has been factored out and shared between nsys-jax and nsys-jax-combine
| version_file_template = """\ | ||
| __version__ = version = {version!r} | ||
| __version_tuple__ = version_tuple = {version_tuple!r} | ||
| __sha__ = {scm_version.node!r} |
There was a problem hiding this comment.
This is propagated so that the install.sh in the nsys-jax[-combine] output .zip archives can install the same nsys-jax commit as was used to generate the archive.
olupton
force-pushed
the
olupton/pip-install-nsys-jax
branch
2 times, most recently
from
bd221c5
to
b66834f
Compare
olupton
force-pushed
the
olupton/pip-install-nsys-jax
branch
from
b66834f
to
faed0fc
Compare
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.
The overarching goal of this PR is to get closer to a world where the
nsys-jaxtooling is straightforwardlypip install-able. While the diff looks scary, it's mostly re-organisation.Substantive changes:
nsys-jaxno longer bundles Python code in the output archives, theinstall.shscript provided for users to run on local machines becomes, loosely,install 'pip nsys-jax[jupyter] @ git+https://github.com/NVIDIA/JAX-Toolbox.git@COMMIT#subdirectory=.github/container/nsys_jax', whereCOMMITcorresponds to thensys-jaxcommand that produced the archive. For theghcr.io/nvidia/jaxcontainers, this is the commit of JAX-Toolbox that triggered the container build.Changes included:
/opt/pip-tools-post-install.d, whichpip-finalize.shwill execute the contents of after installing thepip-managed worldinstall-protocto use this, sopip-finalize.shcan forget about that detail.nvtx_gpu_proj_tracePython code in Nsight Systems 2024.5 and 2024.6 via this.nsys-jaxinstallation (specifically for the containers) intoinstall-nsys-jax.shand thereby clean upinstall-nsight.sh. The new script has to be told the git commit hash of JAX-Toolbox that is being built, becausensys-jaxbakes this into an installation script in its output.ziparchives to ensure the local environment matches the profile-collection environment.nsys-jax,nsys-jax-combineandinstall-protocare now handled via[project.scripts]inpyproject.tomlinstead of being standalone Python scripts. This is "more standard", and also makes it easier to share code betweennsys-jaxandnsys-jax-combine.jax_nsystonsys_jaxfor consistency.NSYS_JAX_DEFAULT_PREFIXenvironment variable; previously the default was the current working directory, but that can be inconvenient to steer in Jupyter environments.