Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Add nightly docker tag and pull instructions #54

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

Merged
mrodden merged 1 commit into from
Jul 9, 2025
Merged

Add nightly docker tag and pull instructions #54

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions .github/workflows/build-docker.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,9 @@ on:
required: false
default: 'plugin_wheels'
type: string
extra-cr-tag:
required: false
type: string

jobs:
build-docker:
Expand Down Expand Up @@ -59,6 +62,7 @@ jobs:
| docker login ghcr.io -u ${{ github.actor }} --password-stdin
- name: Push docker images
env:
EXTRA_CR_TAG: ${{ inputs.extra-cr-tag }}
ROCM_VERSION: ${{ inputs.rocm-version }}
run: |
ubu22_img="ghcr.io/rocm/jax-ubu22.rocm${ROCM_VERSION//.}:${GITHUB_SHA}"
Expand All @@ -69,4 +73,18 @@ jobs:
docker tag "jax-ubu24.rocm${ROCM_VERSION//.}" "${ubu24_img}"
docker push "${ubu22_img}"
docker push "${ubu24_img}"
- name: Push extra tags
if: ${{ inputs.extra-cr-tag }}
env:
EXTRA_CR_TAG: ${{ inputs.extra-cr-tag }}
ROCM_VERSION: ${{ inputs.rocm-version }}
run: |
ubu22_img="ghcr.io/rocm/jax-ubu22.rocm${ROCM_VERSION//.}:${EXTRA_CR_TAG}"
ubu24_img="ghcr.io/rocm/jax-ubu24.rocm${ROCM_VERSION//.}:${EXTRA_CR_TAG}"
echo "Ubuntu 22 image name: ${ubu22_img}"
echo "Ubuntu 24 image name: ${ubu24_img}"
docker tag "jax-ubu22.rocm${ROCM_VERSION//.}" "${ubu22_img}"
docker tag "jax-ubu24.rocm${ROCM_VERSION//.}" "${ubu24_img}"
docker push "${ubu22_img}"
docker push "${ubu24_img}"

2 changes: 2 additions & 0 deletions .github/workflows/nightly.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,10 +38,12 @@ jobs:
include:
- rocm-version: "6.4.1"
runner-label: "mi-250"
extra-cr-tag: "nightly"
- rocm-version: "7.0"
rocm-build-job: "compute-rocm-dkms-no-npi-hipclang"
rocm-build-num: "16322"
runner-label: "internal"
extra-cr-tag: "nightly"
uses: ./.github/workflows/build-docker.yml
with:
rocm-version: ${{ matrix.rocm-version }}
Expand Down
33 changes: 32 additions & 1 deletion README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@

`rocm-jax` contains sources for the ROCm plugin for JAX, as well as Dockerfiles used to build AMDs `rocm/jax` images.


# development setup

Run stack.py develop to clone jax/xla
Expand Down Expand Up @@ -141,3 +140,35 @@ Use make to build the plugin
```
(cd jax_rocm_plugin && make clean dist)
```

# Nightly Builds

We build rocm-jax nightly with [a Github Actions workflow](https://github.com/ROCm/rocm-jax/actions/workflows/nightly.yml).

## Docker

Nightly docker images are kept in the Github Container Registry

```
echo <MY_GITHUB_ACCESS_TOKEN> | docker login ghcr.io -u <USERNAME> --password-stdin
docker pull ghcr.io/rocm/jax-ubu24.rocm70:nightly
```

You can also find nightly images for other Ubuntu versions and ROCm version as well as older nightly images on the [packages page](https://github.com/orgs/ROCm/packages?repo_name=rocm-jax). Images get tagged with the git commit hash of the commit that the image was built from.

### Authenticating to the Container Registry

Pull access to the Github CR is done by a personal access token (classic) with the `read:packages` permission. To create one, click your profile picture in the top-right of Github, select Settings > Developer settings > Personal access tokens > Tokens (classic) and then select the option to generate a new token. Make sure you select the classic token option and git it the `read:packages` permission.

Once your token has been created, go back to the Tokens (classic) page and set your token's SSO settings to allow access to the ROCm Github organization.
Copy link
Contributor

@mrodden mrodden Jul 8, 2025

Choose a reason for hiding this comment

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

+1


Once your token has been set up to use SSO, you can log in with the `docker` command line by running,

```
echo <MY_GITHUB_ACCESS_TOKEN> | docker login ghcr.io -u <USERNAME> --password-stdin
```

## Wheels

Wheels get saved as artifacts to each run of the nightly workflow. Go to the [nightly workflow](https://github.com/ROCm/rocm-jax/actions/workflows/nightly.yml), select the run you want to get wheels from, and scroll down to the bottom of the page to find the build artifacts. Each artifact is a zip file that contains all of the wheels built for a specific ROCm version.

Loading