Open
Conversation
* Build wheels on the fly * Remove uneccessary script * Unify sources for wheel creation * Fix * Fix lint * Fix lint * Follow different logic if srcs is set * Sync with master * Simplify * Clean-up * Reformat * Suppress warning * Drop copy_file from build_utils * Fix copyright * Added TODO, remove due to the upstreamed changes * Fix lint warning
Remove Dockerfile.base-ubu22 as part of consolidating to Ubuntu 24.04 only. This is the first step in the Ubuntu 24.04 and multi-Python migration. Related to Phase 1 of the migration plan.
Add support for all Python versions (3.11-3.14 including nogil variants) in a single Docker image. Changes: - Install all Python versions from deadsnakes PPA: - python3.11, python3.12, python3.13, python3.14 - python3.13-nogil, python3.14-nogil - All with -dev and -venv packages - Set Python 3.12 as default via update-alternatives - Update image label to list all available Python versions (removed single python_version label, added python_versions) Benefits: - Single image contains all Python versions - Users can choose Python version at runtime - Simpler CI (build once, test all versions) - No Python version in image name This completes Phase 1 of the migration plan.
Remove Dockerfile.jax-ubu22 as part of consolidating to Ubuntu 24.04 only. This continues the Ubuntu 24.04 and multi-Python migration. Related to Phase 2 of the migration plan.
Install JAX and all dependencies for all Python versions (3.11-3.14) in a single Docker image. Changes: - Install common dependencies for all Python versions using loop: - numpy, build, wheel, six, auditwheel, scipy - pytest and related testing tools - cloudpickle, portpicker, matplotlib, absl-py, flatbuffers, hypothesis - Install JAX and jaxlib from requirements.txt for all Python versions - Install ROCm JAX wheels (plugin, pjrt, jaxlib) for all Python versions - Update image label to list all Python versions: Added: com.amdgpu.python_versions="3.11,3.12,3.13,3.14" Benefits: - Single image contains JAX for all Python versions - Users can choose Python version at runtime - No Python version in image name - Simpler CI workflow Note: Using loop with python3.11, python3.12, python3.13, python3.14 to install packages for each version. Nogil variants (3.13-nogil, 3.14-nogil) are available in the base image but JAX installation for them will be added in future iterations if needed. This completes Phase 2 of the migration plan.
Update build_wheels.py to build wheels for Python 3.11, 3.12, 3.13, and 3.14 by default. Changes: - Update default --python-versions from "3.11.13,3.12" to "3.11,3.12,3.13,3.14" Note: The manylinux_2_28 base image already contains all required Python versions (3.8-3.14, including 3.13t and 3.14t free-threaded variants) in /opt/python/. No Dockerfile changes are needed. Requirements lock files already exist for all versions: - requirements_lock_3_11.txt - requirements_lock_3_12.txt - requirements_lock_3_13.txt - requirements_lock_3_13_ft.txt (free-threaded) - requirements_lock_3_14.txt - requirements_lock_3_14_ft.txt (free-threaded) This completes Phase 3 of the migration plan.
Update all CI workflows to use Ubuntu 24.04 images exclusively and support Python 3.11, 3.12, 3.13, and 3.14. Changes: 1. build-base-docker.yml: - Remove ubuntu-version "22" from matrix - Hardcode filter to "ubu24" - Remove UBUNTU_VERSION env var from push step - Update image tags to always use ubu24 2. build-docker.yml: - Remove all ubu22 image references - Only build and push ubu24 images - Update both main push and extra-tag push steps 3. ci.yml: - Update python-versions to "3.11,3.12,3.13,3.14" - Change all image references from ubu22 to ubu24 - Remove TODO comment about Python 3.13 4. nightly.yml: - Update python-versions to "3.11,3.12,3.13,3.14" - Remove ubuntu-version "22" from test matrix - Hardcode ubuntu-version to "24" in test-and-upload call 5. Other workflows verified (no changes needed): - build-wheels.yml: Already parameterized - nightly-rbe.yml: Already tests 3.11-3.14 - rocm-perf.yml: Already uses ubu24 - llama-perf.yml: Already uses ubu24 - test-and-upload.yml: Accepts ubuntu-version as input - upstream-ci-watcher.yml: No Docker image references Impact: - All CI workflows now use single multi-Python Ubuntu 24.04 images - Wheels built for Python 3.11, 3.12, 3.13, 3.14 - No more Ubuntu 22.04 references in CI This completes Phase 4 of the migration plan.
Update all build scripts to use Ubuntu 24.04 exclusively and support Python 3.11, 3.12, 3.13, and 3.14 by default. Changes: 1. build/ci_build: - Update --python-versions default from "3.12" to "3.11,3.12,3.13,3.14" - Update docstring examples from ubu22 to ubu24 - Update comments to reflect ubu24 as the standard 2. build/ci.sh: - Update --python-versions from "3.12" to "3.11,3.12,3.13,3.14" - Update test image reference from "jax-ubu22.rocm7100" to "jax-ubu24.rocm720" - Aligns with current ROCm 7.2.0 default 3. docker/Makefile: - Remove jax-ubu22 from "all" target - Remove clean-jax-ubu22 from "clean" target - Only build and clean ubu24 images Impact: - All build scripts now default to Python 3.11-3.14 - Ubuntu 22.04 completely removed from build infrastructure - Simplified Makefile with single Ubuntu version - Consistent defaults across all build tools This completes Phase 5 of the migration plan.
Install gh (GitHub CLI) and gcloud (Google Cloud CLI) tools in the Ubuntu 24.04 base container for improved CI/CD capabilities. Changes: - Add GitHub CLI (gh) installation: - Install from official GitHub CLI repository - Add GPG key and apt source - Install gh package - Add Google Cloud CLI (gcloud) installation: - Install from official Google Cloud SDK repository - Add GPG key and apt source - Install google-cloud-cli package Both tools are commonly used in CI/CD pipelines for: - gh: GitHub API interactions, PR management, release automation - gcloud: GCS bucket access, artifact storage, cloud resource management Dependencies added: curl, gpg, apt-transport-https, ca-certificates, gnupg
Add gcloud configuration to disable credential requirements, enabling access to public GCS buckets from upstream projects without login. Changes: - Add "gcloud config set auth/disable_credentials True" command - This allows reading from public Google Cloud Storage buckets - Useful for accessing upstream artifacts, datasets, and resources - No authentication required for public bucket access This configuration is particularly helpful for CI/CD pipelines that need to fetch public resources from GCS without managing service account keys.
Install python-dev packages for Python 3.11, 3.12, 3.13, and 3.14 to enable building Python C extensions and native modules. Changes: - Add python3.11-dev - Add python3.12-dev - Add python3.13-dev - Add python3.14-dev The -dev packages provide: - Python header files (Python.h) - Static libraries for linking - Development tools for building extensions - Required for compiling packages with native code (numpy, scipy, etc.) Note: Free-threaded variants (python3.13-nogil, python3.14-nogil) do not have separate -dev packages in the deadsnakes PPA.
* Introduce ci pr check tests pipeline, remove test execution from the regular pipeline * Restore buidl container CI job
* Update runner labels in Llama perf workflow * Add flaky, timeout handling and ignore failing tests --------- Co-authored-by: psanal35 <pakize.sanal@amd.com>
* Introduce asan build * Remove duplicate entry * Move tsan/asan targets into rocm-jax repo * Apply testing options * Add more failing tests * Address review comments * Address review comments * Add docu * Fix lint warning
Switch from http_archive to git_repository for pulling XLA and JAX dependencies. This eliminates the need to compute SHA256 hashes when updating to new commits. Changes: - third_party/xla/workspace.bzl: Use git_repository with commit= - third_party/jax/workspace.bzl: Use git_repository with commit=, removed patch 0006 - WORKSPACE: Remove external_deps_repository (was added by patch 0006) - build/rocm/ci_build: Remove GIT_DIR/GIT_WORK_TREE env vars that conflicted with git_repository No SHA256 computation required.
* Restore patches * Trigger CI/CD pipelinee * Fix commit * Fix patch
Ruturaj4
force-pushed
the
v0.9.0-merger
branch
from
92ae5a4 to
8d95c2e
Compare
This file contains hidden or 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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
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.
Motivation
merge
Submission Checklist