Skip to content

Migrate TensorBoard builds to Bzlmod on Bazel 7.7.0 - #7145

Open
psamanoelton wants to merge 4 commits into
tensorflow:masterfrom
psamanoelton:workspace_upgrade
Open

Migrate TensorBoard builds to Bzlmod on Bazel 7.7.0#7145
psamanoelton wants to merge 4 commits into
tensorflow:masterfrom
psamanoelton:workspace_upgrade

Conversation

@psamanoelton

Copy link
Copy Markdown
Contributor

Summary

This change migrates TensorBoard’s build to Bazel Modules (Bzlmod) while keeping Bazel pinned to 7.7.0.

The migration preserves compatibility with TensorFlow 2.21.0 and Protobuf 6.31.1, makes Bzlmod the only supported dependency-resolution mode, and adds a hermetic Bazel-side Python environment.

This is intentionally not a Bazel 8 upgrade.

Motivation

TensorBoard’s TensorFlow 2.21.0 release required several non-trivial dependency updates and compatibility workarounds. Continuing to manage the dependency graph primarily through WORKSPACE makes future upgrades difficult because versions and transitive dependencies are not resolved consistently.

Moving the supported build to Bzlmod provides:

  • Explicit module versions and dependency resolution
  • A checked-in module lockfile
  • Better isolation from host Python installations
  • Fewer manually initialized transitive repositories
  • A clearer path for future dependency upgrades

Main changes

Enable Bzlmod on Bazel 7.7.0

  • Add MODULE.bazel and MODULE.bazel.lock.
  • Pin Bazel to 7.7.0.
  • Enable Bzlmod by default in .bazelrc.
  • Mark legacy --noenable_bzlmod builds as unsupported.
  • Explicitly constrain the module graph to Bazel 7.x.
  • Update CI configuration to use Bazel 7.7.0.

Move core dependencies to modules

The module graph now directly manages:

  • bazel_skylib 1.7.1
  • rules_cc 0.1.1
  • rules_java 8.6.1
  • rules_python 1.0.0
  • Protobuf 31.1 / release 6.31.1
  • gRPC 1.74.0
  • grpc-java 1.69.0
  • aspect_rules_js 2.1.0
  • rules_webtesting 0.4.1
  • rules_web_testing_python 0.4.1

Migrating rules_webtesting also removes the old explicit setup for rules_go, Gazelle, legacy rules_python, Bazel Skylib, and browser repositories.

Add a hermetic Python environment

  • Configure a downloaded, checksummed Python 3.10 toolchain.
  • Add a locked pip dependency graph for Bazel-built Python targets.
  • Prevent host PYTHONPATH, user-site packages, and virtualenv state from
    leaking into Bazel tests.
  • Align the Python web-testing dependencies with Python 3.10.
  • Move TensorBoard’s Markdown, Bleach, and Webencodings source repositories
    behind a local module extension.
  • Replace the legacy Werkzeug repository with the locked pip dependency.
  • Remove the unused standalone urllib3 repository.

The pip-package smoke test continues to create its own isolated environment because it validates the produced TensorBoard wheel rather than Bazel targets.

Preserve TensorFlow 2.21 / Protobuf compatibility

  • Consume Protobuf 6.31.1 through a module source override.
  • Preserve the required Protobuf Java and Python compatibility patches.
  • Use rules_cc 0.1.1, which includes the required cc_proto_library
    correction.
  • Correct the grpc-java module dependency graph.
  • Regenerate the Rust data-server descriptor with the current Protobuf schema.
  • Update affected BUILD dependencies and compatibility targets.

All compatibility patches and their expected removal conditions are documented in patches/README.md.

Update Bzlmod runfiles handling

Bzlmod changes the runfiles layout for the root module and module-extension repositories.

The affected shell utilities now support the _main layout, including:

  • Pip package building and extraction
  • Example plugin smoke tests
  • Logo generation
  • HParams utilities

The wheel builder resolves module-extension repositories through Bazel’s _repo_mapping file. This allows vendored Bleach and Webencodings sources to work without hardcoding their canonical Bzlmod repository names.

Keep a transitional WORKSPACE.bzlmod

Some dependencies cannot be migrated safely without substantially expanding the scope of this change. WORKSPACE.bzlmod remains as a compatibility bridge for:

  • The patched Closure/Soy stack
  • Legacy rules_nodejs 5.8.1, yarn_install, and concatjs
  • rules_sass
  • The existing rules_rust and cargo-raze graph
  • Closure-dependent font and JavaScript repositories
  • Remaining Java artifacts and local compatibility repositories

The original WORKSPACE remains for repository compatibility and external tooling, but it is not a supported build mode.

Validation

The following workflows pass locally on Linux:

pip install "tensorflow==2.21.0"
pip install \
  -r ./tensorboard/pip_package/requirements.txt \
  -r ./tensorboard/pip_package/requirements_dev.txt

pip install --force-reinstall "urllib3==1.26.20" "six>=1.12,<2"

bazel build //tensorboard/... --keep_going

bazel test //tensorboard/... --test_output=errors

bazel run //tensorboard/pip_package:test_pip_package -- \
  --tf-version "tensorflow==2.21"

bazel build //tensorboard/pip_package:build_pip_package

mkdir -p /tmp/tb_wheel
./bazel-bin/tensorboard/pip_package/build_pip_package /tmp/tb_wheel
pip install /tmp/tb_wheel/tensorboard-*.whl

Additional validation included:

  • bazel mod tidy
  • Full //tensorboard/... target analysis
  • Functional web-test target analysis with Python 3.10
  • Frontend target analysis
  • Representative Python tests
  • Shell syntax and whitespace checks

Scope and follow-up work

This PR intentionally does not include:

  • Migration from Yarn/concatjs to modern rules_js package management
  • Replacement of rules_sass
  • Migration from cargo-raze to crate-universe
  • Complete removal of WORKSPACE.bzlmod

Known limitation

TensorBoard uses Bzlmod as its supported dependency-resolution entry point, but some dependencies are still provided by the transitional WORKSPACE.bzlmod.

Consequently, builds using --noenable_workspace are not supported yet and are expected to fail, currently beginning with the legacy Closure/Soy repository.

Removing WORKSPACE.bzlmod should be handled as follow-up work alongside the Closure, Node/Yarn, Sass, and Rust migrations.

@arcra arcra left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Ah, I think I was hoping (dreaming?) that it would be a cleaner transition to new packages, but I realize now that this is an intermediate, incremental step towards updating some dependencies.

I left several comments, but generally LGTM. We can continue iterating on whatever else needs to be updated later.

If it's a matter of using Bazel 8, I think we don't have to necesarily use the same one as TF... our dependency on TF is via a specific binary version, we're not building TF code as part of our bazel setup, so I don't think this should be a restriction.

Anyway, we can discuss more, but this is fine for an incremental step, I think.

Comment thread patches/README.md
install-time invocation was less reliable than applying the generated patch
files directly.
At build time, `WORKSPACE` and the transitional `WORKSPACE.bzlmod` apply the
generated patch artifacts via `yarn_install(post_install_patches = ...)`

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

I was hoping that by transitioning to BAZELMOD, we'd be able to get rid of the patches, instead of introducing more. Do you currently think this would be solved with Bazel 8? Or do you think we have dependencies on things that are too old and there are no new versions that are compatible with other things we're using?

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Ah, I think I'm realizing this is sort of a first step, where we're still mixing some old libraries with newer ones, just to get it to work with Bazel Modules, and so little by little we would be updating dependencies (and code, when necessary, for example if we need to use a different, newer dependency), to remove this patches. Is this right?

trap cleanup EXIT

cp -LR "${TEST_SRCDIR}/org_tensorflow_tensorboard/tensorboard/examples/plugins/example_basic/" \
workspace_runfiles="${TEST_SRCDIR}/${TEST_WORKSPACE:-org_tensorflow_tensorboard}"

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

What is this syntax with the :- characters?

tensorflow==2.21.0
# Match the source dependency vendored by third_party/python.bzl. Keeping a
# single version avoids non-deterministic runfiles import ordering.
urllib3==1.26.20

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Do you know why this is necessary? If we're vendoring this, then it shouldn't be required in the runtime... that's the point of vendoring it, AFAIU.

Comment thread patches/README.md
- Restores protobuf's system-Python repository when protobuf 6.31.1 is consumed
through its source `MODULE.bazel`. The upstream module currently aliases that
name to a rules_python toolchain repository, which does not provide the
`version.bzl` and Python-header targets that protobuf's public Python build

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

This is hard to follow... is this saying that the protobuf python package expects to find some header file that is not provided when protobuf is installed/provided via bazel?

So protobuf cannot be used with bazel "out of the box"?

What does "upstream module" (from above) refers to, in this context?

Comment thread patches/README.md
`rules_web_testing_python` as Bazel modules while retaining the existing
Bazel-7-compatible Closure/Soy setup.

Removal is planned when TensorBoard moves to a module-native Closure release

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Is this meant to be a follow-up step, only not done here to reduce scope of this change? Or is there a blocker to move to another version?

EOF
}

# Resolve an apparent repository path through Bazel's Bzlmod runfiles mapping.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

This was fairly confusing to me until I realized these are concepts in Bazel.

Please include a link to https://bazel.build/external/overview#concepts, so that whoever reads this can more easily understand what we're doing.

# Keep the TensorFlow version aligned with the release compatibility baseline.
# The pip-package smoke test separately validates the CI-selected TensorFlow
# package (currently tf-nightly).
-r requirements.txt

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Why not just include whatever is needed in requirements.txt (or the _dev file)?

Comment thread MODULE.bazel
# TensorBoard's Linux CI builds execute inside a root-owned container.
ignore_root_user_error = True,
is_default = True,
python_version = "3.10",

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Will this help remove some of the "hacky" things we had to set up python in the host runner? Or will that still be needed to make it available to bazel, and this is just using that?

Comment thread tensorboard/backend/BUILD
deps = [
":json_util",
"@org_pocoo_werkzeug",
"@tensorboard_pip_deps//werkzeug",

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

In MODULE.bazel, I see where this "repo name" is defined, but I don't see any references to werkzeug there... does that mean that any pip package can be referenced this way?

Can we / should we then update the "expect_library" targets to use this syntax? Or how are they different?

Comment thread DEVELOPMENT.md
TensorBoard builds are done with [Bazel](https://bazel.build), so you may need to [install Bazel](https://docs.bazel.build/versions/master/install.html). The Bazel build will automatically "vulcanize" all the HTML files and generate a "binary" launcher script. When HTML is vulcanized, it means all the script tags and HTML imports are inlined into one big HTML file. Then the Bazel build puts that index.html file inside a static assets zip. The python HTTP server then reads static assets from that zip while serving.
TensorBoard builds are done with [Bazel](https://bazel.build). The supported
version is pinned in `.bazelversion` (currently Bazel 7.7.0), and Bazel 8 is
intentionally unsupported. TensorBoard uses Bzlmod for dependency resolution;

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Can we add some details for why we pin to a specific version?

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.

2 participants