Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use shell script to drive jaxlib build #909

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

yhtang
Copy link

@yhtang yhtang commented Mar 30, 2023

Together with google/jax#15205, this should eliminate the need to apply the build_jaxlib/update_build_scripts.patch each time we rebase. Note the --editable option will only be usable once alpa is rebased onto JAX 0.4.8.

The original update_build_scripts.patch file is meant to provide two functionalities:

  1. build jaxlib as an 'editable' package. This can now be done by
cd build
./build.sh  # build wheel
./build.sh --editable  # build editable installation
  1. Ask Bazel to prefer local installations NCCL. This can now be done by
./build.sh --nccl local  # prefer local NCCL
./build.sh --nccl download  # prefer to use XLA's pinned version

@yhtang yhtang changed the title Use Shell script to drive jaxlib build Use shell script to drive jaxlib build Mar 30, 2023
Copy link
Member

@gjoliver gjoliver left a comment

Choose a reason for hiding this comment

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

nice cleanup.
let's figure out when we should merge this.


if [[ "${EDITABLE}" == "1" ]]; then
BUILD_PARAM="${BUILD_PARAM} --editable"
fi
Copy link
Member

Choose a reason for hiding this comment

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

why don't we do this directly in the flag parsing while loop above?
you then don't need the additional EDITABLE variable.


if [[ "$DEBUG" == "1" ]]; then
BUILD_PARAM="${BUILD_PARAM} --bazel_options=-c --bazel_options=dbg --bazel_options=--strip=never --bazel_options=--cxxopt=-g --bazel_options=--cxxopt=-O0"
fi
Copy link
Member

Choose a reason for hiding this comment

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

same for --debug I guess ...

print_var TF_NCCL_VERSION
print_var CC_OPT_FLAGS

print_var BUILD_PARAM
Copy link
Member

Choose a reason for hiding this comment

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

BUILD_PARAM is printed twice. get rid of the one above I guess?

@@ -0,0 +1,30 @@
//usr/bin/env nvcc -odir `mktemp -d` -x cu --run "$0" ${1:+--run-args "${@:1}"} ; exit $?
Copy link
Member

Choose a reason for hiding this comment

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

is this a new util?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants