code review

.github/container/nsys_jax/nsys_jax/scripts/install_flamegraph.py

@@ -12,7 +12,7 @@ def main():
     # TODO: add a default to (with confirmation) install in the same prefix as this script is installed to
     parser = argparse.ArgumentParser("Fetch the flamegraph.pl script")
     parser.add_argument(
-        "prefix", help="Output prefix under which to install protoc", type=str
+        "prefix", help="Output prefix under which to install flamegraph.pl", type=str
     )
     args = parser.parse_args()
     install_dir = os.path.join(args.prefix, "bin")

.github/container/nsys_jax/nsys_jax/scripts/nsys_jax.py

@@ -91,6 +91,7 @@ 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.
     """
+    # setuptools_scm produces a shortened sha with a `g` prefix (for git)
     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))

docs/nsys-jax.md

@@ -41,7 +41,8 @@ $ pip install --src /checkout-dir -e 'git+https://github.com/NVIDIA/JAX-Toolbox.
 You may want to include this in a global `pip-compile`-based dependency resolution (as is done in the containers built
 from this repository), rather than running too many ad-hoc `pip install` commands.
 
-This will install all of the components mentioned so far, but does not currently include all (implicit) dependencies:
+This will install all of the components mentioned so far, but does not currently include the following implicit
+dependencies:
 - `protoc` must be installed at a version compatible with the `google.protobuf` runtime library; `nsys-jax` includes a
   helper script that can be run after `pip install`, e.g. to install `/usr/local/bin/protoc`, run
   `install-protoc /usr/local`.