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.

Sign up for GitHub

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

Jump to bottom

Enable rocm nightly builds on tensorflow-upstream fork #2498

Merged
merged 2 commits into from
May 2, 2024
Merged

Enable rocm nightly builds on tensorflow-upstream fork #2498

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
acbf96b
Enable rocm nightly builds on tensorflow-upstream fork
hsharsha Apr 22, 2024
c718cba
Add OUTPUT_PATH repo env to set destination for artifacts
hsharsha Apr 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
54 changes: 36 additions & 18 deletions build_rocm_python3
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ while getopts "hrn" opt; do
restriction=true
;;
n)
nightly=true
nightly=true
esac
done
shift "$((OPTIND-1))"
Expand All @@ -47,22 +47,40 @@ else
fi

if [ -f /usertools/rocm.bazelrc ]; then
# Use the bazelrc files in /usertools if available
# Also, this is likely a tensorflow-build container so put the whl in /tf/pkg
TF_PKG_LOC=/tf/pkg
# Remove any previous builds and build nightly
rm -f $TF_PKG_LOC/tensorflow*.whl
python3 tensorflow/tools/ci_build/update_version.py --nightly --rocm_version &&
bazel --bazelrc=/usertools/rocm.bazelrc build $RESOURCE_OPTION --config=rocm --action_env=TF_PYTHON_VERSION=$PYTHON_VERSION tensorflow/tools/pip_package:wheel --verbose_failures &&
cp ./bazel-bin/tensorflow/tools/pip_package/wheel_house/tensorflow*.whl $TF_PKG_LOC/ &&
pip3 install --upgrade $TF_PKG_LOC/tensorflow*.whl
# Use the bazelrc files in /usertools if available
# Also, this is likely a tensorflow-build container so put the whl in /tf/pkg
TF_PKG_LOC=/tf/pkg
if [[ -n $nightly ]]; then
# Remove any previous builds and build nightly
rm -f $TF_PKG_LOC/tf_nightly_rocm*.whl
export project_name=tf_nightly_rocm
python3 tensorflow/tools/ci_build/update_version.py --nightly --rocm_version &&
bazel --bazelrc=/usertools/rocm.bazelrc build $RESOURCE_OPTION --config=rocm --repo_env=WHEEL_NAME=tf_nightly_rocm --action_env=project_name=tf_nightly_rocm --action_env=TF_PYTHON_VERSION=$PYTHON_VERSION tensorflow/tools/pip_package:wheel --verbose_failures &&
cp ./bazel-bin/tensorflow/tools/pip_package/wheel_house/tf_nightly_rocm*.whl $TF_PKG_LOC/ &&
pip3 install --upgrade $TF_PKG_LOC/tf_nightly_rocm*.whl
else
# Remove any previous builds and build release
rm -f $TF_PKG_LOC/tensorflow*.whl
python3 tensorflow/tools/ci_build/update_version.py --rocm_version &&
bazel --bazelrc=/usertools/rocm.bazelrc build $RESOURCE_OPTION --config=rocm --repo_env=WHEEL_NAME=tensorflow_rocm --action_env=project_name=tensorflow_rocm --action_env=TF_PYTHON_VERSION=$PYTHON_VERSION tensorflow/tools/pip_package:wheel --verbose_failures &&
cp ./bazel-bin/tensorflow/tools/pip_package/wheel_house/tensorflow*.whl $TF_PKG_LOC/ &&
pip3 install --upgrade $TF_PKG_LOC/tensorflow*.whl
fi
else
# Legacy style: run configure then build
TF_PKG_LOC=/tmp/tensorflow_pkg
yes "" | TF_NEED_CLANG=0 ROCM_PATH=$ROCM_INSTALL_DIR TF_NEED_ROCM=1 PYTHON_BIN_PATH=/usr/bin/python3 ./configure &&
# Remove any previous builds and build nightly
rm -f $TF_PKG_LOC/tensorflow*.whl
bazel build $RESOURCE_OPTION --config=opt --config=rocm //tensorflow/tools/pip_package:wheel --verbose_failures &&
cp ./bazel-bin/tensorflow/tools/pip_package/wheel_house/tensorflow*.whl $TF_PKG_LOC/ &&
pip3 install --upgrade $TF_PKG_LOC/tf_nightly_rocm*.whl
# Legacy style: run configure then build
TF_PKG_LOC=/tmp/tensorflow_pkg
yes "" | TF_NEED_CLANG=0 ROCM_PATH=$ROCM_INSTALL_DIR TF_NEED_ROCM=1 PYTHON_BIN_PATH=/usr/bin/python3 ./configure &&
if [[ -n $nightly ]]; then
# Remove any previous builds and build nightly
rm -f $TF_PKG_LOC/tf_nightly_rocm*.whl
bazel build $RESOURCE_OPTION --config=opt --config=rocm --repo_env=WHEEL_NAME=tf_nightly_rocm --action_env=project_name=tf_nightly_rocm //tensorflow/tools/pip_package:wheel --verbose_failures &&
cp ./bazel-bin/tensorflow/tools/pip_package/wheel_house/tf_nightly_rocm*.whl $TF_PKG_LOC/ &&
pip3 install --upgrade $TF_PKG_LOC/tf_nightly_rocm*.whl
else
# Remove any previous builds and build release
rm -f $TF_PKG_LOC/tensorflow*.whl
bazel build $RESOURCE_OPTION --config=opt --config=rocm --repo_env=WHEEL_NAME=tensorflow_rocm --action_env=project_name=tensorflow_rocm //tensorflow/tools/pip_package:wheel --verbose_failures &&
cp ./bazel-bin/tensorflow/tools/pip_package/wheel_house/tensorflow*.whl $TF_PKG_LOC/ &&
pip3 install --upgrade $TF_PKG_LOC/tensorflow*.whl
fi
fi
5 changes: 3 additions & 2 deletions tensorflow/tools/pip_package/utils/tf_wheel.bzl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,17 +23,18 @@ Should be set via --repo_env=WHEEL_NAME=tensorflow_cpu.
6) `--xla_aot` - paths to files that should be in xla_aot directory.
"""

load("@python_version_repo//:py_version.bzl", "WHEEL_COLLAB", "WHEEL_NAME")
load("@python_version_repo//:py_version.bzl", "WHEEL_COLLAB", "WHEEL_NAME", "OUTPUT_PATH")
load("//tensorflow:tensorflow.bzl", "VERSION")

def _tf_wheel_impl(ctx):
executable = ctx.executable.wheel_binary

output = ctx.actions.declare_directory("wheel_house")
output_path = OUTPUT_PATH if OUTPUT_PATH else output.path
args = ctx.actions.args()
args.add("--project-name", WHEEL_NAME)
args.add("--collab", str(WHEEL_COLLAB))
args.add("--output-name", output.path)
args.add("--output-name", output_path)
args.add("--version", VERSION)

headers = ctx.files.headers[:]
Expand Down
7 changes: 5 additions & 2 deletions tensorflow/tools/toolchains/python/python_repo.bzl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ Version can be set via build parameter "--repo_env=TF_PYTHON_VERSION=3.10"
Defaults to 3.10.

To set wheel name, add "--repo_env=WHEEL_NAME=tensorflow_cpu"
To set output path, add "--repo_env=OUTPUT_PATH=/tf/pkg"
"""

VERSIONS = ["3.9", "3.10", "3.11", "3.12"]
Expand All @@ -21,22 +22,24 @@ TF_PYTHON_VERSION = "{}"
HERMETIC_PYTHON_VERSION = "{}"
WHEEL_NAME = "{}"
WHEEL_COLLAB = "{}"
OUTPUT_PATH = "{}"
"""

def _python_repository_impl(repository_ctx):
repository_ctx.file("BUILD", "")
version = repository_ctx.os.environ.get("TF_PYTHON_VERSION", "")
wheel_name = repository_ctx.os.environ.get("WHEEL_NAME", "tensorflow")
wheel_collab = repository_ctx.os.environ.get("WHEEL_COLLAB", False)
output_path = repository_ctx.os.environ.get("OUTPUT_PATH", None)
if version not in VERSIONS:
print(WARNING) # buildifier: disable=print
version = DEFAULT_VERSION
repository_ctx.file(
"py_version.bzl",
content.format(version, version, wheel_name, wheel_collab),
content.format(version, version, wheel_name, wheel_collab, output_path),
)

python_repository = repository_rule(
implementation = _python_repository_impl,
environ = ["TF_PYTHON_VERSION", "WHEEL_NAME", "WHEEL_COLLAB"],
environ = ["TF_PYTHON_VERSION", "WHEEL_NAME", "WHEEL_COLLAB", "OUTPUT_PATH"],
)