Skip to content

Commit d6ec8f4

Browse files
jbdalidojbdalido
committed
revamp the workflow in preparation for nightly runs
1 parent f8d704c commit d6ec8f4

19 files changed

+96
-203
lines changed

‎.github/workflows/build.yaml‎

Lines changed: 96 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -2,182 +2,189 @@ name: PJRT GPU library
22
on:
33
push:
44
tags:
5-
- '*'
5+
- "*"
66
pull_request:
7-
7+
workflow_dispatch:
8+
inputs:
9+
xla_commit:
10+
required: true
11+
type: string
12+
default: "main"
13+
814
concurrency:
915
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
1016
cancel-in-progress: true
1117

1218
permissions:
13-
contents: write
19+
contents: write
20+
21+
env:
22+
XLA_COMMIT: ${{ inputs.xla_commit || 'cc075beb6148c2777da2b6749c63830856ee6c2a' }}
23+
1424
jobs:
15-
pjrt-osx-artifacts:
16-
strategy:
17-
matrix:
18-
pjrt:
19-
- target: cpu
20-
config: "--config=release_macos_arm64"
21-
artifact: bazel-bin/xla/pjrt/c/libpjrt_c_api_cpu_plugin.dylib
22-
runs_on: ['self-hosted', 'macOS']
23-
platform: darwin-arm64
24-
bazel_target: //xla/pjrt/c:pjrt_c_api_cpu_plugin
25-
runs-on: ${{ matrix.pjrt.runs_on }}
25+
setup_openxla:
26+
runs-on: ubuntu-latest
2627
steps:
2728
- name: "Checking out repository"
2829
uses: actions/checkout@v4
2930
with:
30-
path: 'pjrt-artifacts'
31-
32-
- name: Extract OpenXLA commit id
33-
id: extract
34-
run: |
35-
commitid=`grep "OPENXLA_COMMIT = " ./pjrt-artifacts/third_party/openxla/xla.bzl | awk -F ' = ' '{print $2}' | sed 's/"//g'`
36-
echo "commitid=$commitid" >> $GITHUB_OUTPUT
37-
31+
path: "pjrt-artifacts"
3832
- name: "Checking out openxla repository"
3933
uses: actions/checkout@v4
4034
with:
35+
ref: ${{ env.XLA_COMMIT }}
4136
repository: openxla/xla
42-
path: 'xla'
43-
ref: ${{ steps.extract.outputs.commitid }}
44-
45-
- name: Apply patches
37+
path: "xla"
38+
- name: Apply patches to openxla
4639
working-directory: ./xla
4740
run: |
48-
for patch in $(ls ../pjrt-artifacts/third_party/openxla/patches/*.patch | sort); do
41+
echo "::notice Applying patches to openxla ${{ github.sha }}"
42+
43+
for patch in $(ls ../pjrt-artifacts/openxla/patches/*.patch | sort); do
4944
echo "Applying patch $patch"
5045
git apply "$patch"
5146
done
52-
53-
- name: Setup Bazel
54-
run: |
55-
wget https://github.com/bazelbuild/bazelisk/releases/download/v1.25.0/bazelisk-linux-amd64 -O /usr/local/bin/bazel
56-
chmod +x /usr/local/bin/bazel
57-
58-
- name: "Build ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }}"
59-
working-directory: ./xla
60-
run: |
61-
bazelisk \
62-
build \
63-
--disk_cache=/tmp/cache \
64-
${{ matrix.pjrt.config }} \
65-
${{ matrix.pjrt.bazel_target }}
66-
67-
- name: Strip binary and rpath
68-
working-directory: ./xla
69-
run: |
70-
strip -x -o libpjrt_cpu.dylib ${{ matrix.pjrt.artifact }}
71-
install_name_tool -id '@rpath/libpjrt_cpu.dylib' libpjrt_cpu.dylib
72-
73-
- name: Create compressed ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} release file
74-
uses: a7ul/tar-action@v1.2.0
75-
with:
76-
command: c
77-
cwd: ./xla
78-
files: |
79-
./libpjrt_cpu.dylib
80-
outPath: pjrt-${{ matrix.pjrt.target }}_${{ matrix.pjrt.platform }}.tar.gz
81-
82-
- name: Upload ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} artifacts
47+
- name: Upload openxla repository artifact
8348
uses: actions/upload-artifact@v4
8449
with:
85-
name: pjrt-${{ matrix.pjrt.target }}_${{ matrix.pjrt.platform }}.tar.gz
86-
path: pjrt-${{ matrix.pjrt.target }}_${{ matrix.pjrt.platform }}.tar.gz
87-
50+
include-hidden-files: true
51+
name: xla-${{ env.XLA_COMMIT }}
52+
path: ./xla
53+
8854
pjrt-artifacts:
55+
runs-on: ${{ matrix.pjrt.runs_on }}
8956
strategy:
9057
matrix:
9158
pjrt:
9259
- target: cuda
9360
config: "--config=cuda"
9461
artifact: libpjrt_c_api_gpu_plugin.so
9562
renamed_artifact: libpjrt_cuda.so
96-
runs_on: ['runs-on','runner=32cpu-linux-x64', 'image=ubuntu24-amd64']
63+
runs_on: ["self-hosted", "Linux", "X64"]
9764
platform: linux-amd64
9865
bazel_target: //xla/pjrt/c:pjrt_c_api_gpu_plugin
9966
- target: rocm
10067
config: "--config=rocm"
10168
artifact: libpjrt_c_api_gpu_plugin.so
10269
renamed_artifact: libpjrt_rocm.so
103-
runs_on: ['runs-on','runner=32cpu-linux-x64', 'image=ubuntu24-amd64']
70+
runs_on:
71+
["runs-on", "runner=32cpu-linux-x64", "image=ubuntu24-amd64"]
10472
platform: linux-amd64
10573
bazel_target: //xla/pjrt/c:pjrt_c_api_gpu_plugin
74+
- target: cpu
75+
config: "--config=release_macos_arm64"
76+
artifact: bazel-bin/xla/pjrt/c/libpjrt_c_api_cpu_plugin.dylib
77+
runs_on: ["self-hosted", "macOS"]
78+
renamed_artifact: libpjrt_cpu.dylib
79+
platform: darwin-arm64
80+
bazel_target: //xla/pjrt/c:pjrt_c_api_cpu_plugin
10681
- target: cpu
10782
config: ""
10883
artifact: libpjrt_c_api_cpu_plugin.so
10984
renamed_artifact: libpjrt_cpu.so
110-
runs_on: ['runs-on','runner=32cpu-linux-x64', 'image=ubuntu24-amd64']
85+
runs_on:
86+
["runs-on", "runner=32cpu-linux-x64", "image=ubuntu24-amd64"]
11187
platform: linux-amd64
11288
bazel_target: //xla/pjrt/c:pjrt_c_api_cpu_plugin
113-
114-
runs-on: ${{ matrix.pjrt.runs_on }}
89+
needs: ["setup_openxla"]
11590
steps:
11691
- name: "Checking out repository"
11792
uses: actions/checkout@v4
93+
with:
94+
path: "pjrt-artifacts"
95+
96+
- name: Download xla artifact
97+
uses: actions/download-artifact@v4
98+
with:
99+
name: xla-${{ env.XLA_COMMIT }}
100+
path: xla
118101

119102
- name: Setup Bazel
120103
run: |
121-
wget https://github.com/bazelbuild/bazelisk/releases/download/v1.25.0/bazelisk-linux-amd64 -O ~/bin/bazel
104+
mkdir -p ~/bin/
105+
wget https://github.com/bazelbuild/bazelisk/releases/download/v1.25.0/bazelisk-${{ matrix.pjrt.platform }} -O ~/bin/bazel
122106
chmod +x ~/bin/bazel
123107
124-
- name: Force remove container
108+
- name: "Setup bazelrc for openxla"
125109
run: |
126-
docker rm -f pjrt-${{ matrix.pjrt.target }}
110+
cp pjrt-artifacts/openxla/bazelrc/common.bazelrc xla/
111+
cp pjrt-artifacts/openxla/bazelrc/${{ matrix.pjrt.target }}.bazelrc xla/xla_configure.bazelrc
127112
128113
- name: "Build ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} image"
114+
if: matrix.pjrt.target != 'cpu'
115+
working-directory: ./pjrt-artifacts
129116
run: |
117+
docker rm -f pjrt-${{ matrix.pjrt.target }}
130118
~/bin/bazel run //${{ matrix.pjrt.target }}:${{ matrix.pjrt.target }}_stripped_tarball
131119
132120
- name: "Run ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} image"
121+
if: matrix.pjrt.target != 'cpu'
133122
run: |
134123
docker run \
135124
--name=pjrt-${{ matrix.pjrt.target }} \
136125
-w=/xla \
137126
--net=host \
127+
-v=$(pwd)/xla:/xla \
138128
-v=/tmp/cache:/tmp/cache \
139129
distroless/${{ matrix.pjrt.target }}_builder:latest \
140130
bazelisk \
141131
build \
142-
--remote_cache=grpc://127.0.0.1:15501 \
143132
--disk_cache=/tmp/cache \
144133
${{ matrix.pjrt.config }} \
145134
${{ matrix.pjrt.bazel_target }}
146135
147136
- name: "Retrieve ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} artifacts"
137+
if: matrix.pjrt.target != 'cpu'
138+
run: |
139+
docker cp pjrt-${{ matrix.pjrt.target }}:/xla/bazel-bin/xla/pjrt/c/${{ matrix.pjrt.artifact }} ${{ matrix.pjrt.renamed_artifact }}
140+
docker rm -f pjrt-${{ matrix.pjrt.target }}
141+
142+
- name: "Build ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }}"
143+
if: matrix.pjrt.target == 'cpu'
144+
working-directory: ./xla
148145
run: |
149-
docker cp pjrt-${{ matrix.pjrt.target }}:/xla/bazel-bin/xla/pjrt/c/${{ matrix.pjrt.artifact }} ${{ matrix.pjrt.renamed_artifact }}
150-
docker rm -f pjrt-${{ matrix.pjrt.target }}
146+
~/bin/bazel \
147+
build \
148+
--disk_cache=/tmp/cache \
149+
${{ matrix.pjrt.config }} \
150+
${{ matrix.pjrt.bazel_target }}
151+
152+
- name: Strip binary and rpath
153+
if: matrix.pjrt.target == 'cpu'
154+
working-directory: ./xla
155+
run: |
156+
strip -x -o libpjrt_cpu.dylib ${{ matrix.pjrt.artifact }}
157+
install_name_tool -id '@rpath/libpjrt_cpu.dylib' libpjrt_cpu.dylib
151158
152159
- name: Create compressed ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} release file
153160
uses: a7ul/tar-action@v1.2.0
154161
with:
155162
command: c
156-
cwd: .
163+
cwd: ./
157164
files: |
158165
./${{ matrix.pjrt.renamed_artifact }}
159166
outPath: pjrt-${{ matrix.pjrt.target }}_${{ matrix.pjrt.platform }}.tar.gz
160-
167+
161168
- name: Upload ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} artifacts
162169
uses: actions/upload-artifact@v4
163170
with:
164171
name: pjrt-${{ matrix.pjrt.target }}_${{ matrix.pjrt.platform }}.tar.gz
165172
path: pjrt-${{ matrix.pjrt.target }}_${{ matrix.pjrt.platform }}.tar.gz
166173

167174
release:
168-
needs: ["pjrt-artifacts","pjrt-osx-artifacts"]
169-
runs-on: ['runs-on','runner=2cpu-linux-x64']
175+
needs: ["pjrt-artifacts"]
176+
runs-on: ubuntu-latest
170177
steps:
171-
- run: rm -rf pjrt*.tar.gz
172-
- name: Download all artifacts
173-
uses: actions/download-artifact@v4
174-
if: startsWith(github.ref, 'refs/tags/')
175-
- name: Release
176-
uses: softprops/action-gh-release@v2
177-
if: startsWith(github.ref, 'refs/tags/')
178-
with:
179-
files: |
180-
pjrt-rocm_linux-amd64.tar.gz/pjrt-rocm_linux-amd64.tar.gz
181-
pjrt-cuda_linux-amd64.tar.gz/pjrt-cuda_linux-amd64.tar.gz
182-
pjrt-cpu_linux-amd64.tar.gz/pjrt-cpu_linux-amd64.tar.gz
183-
pjrt-cpu_darwin-arm64.tar.gz/pjrt-cpu_darwin-arm64.tar.gz
178+
- run: rm -rf pjrt*.tar.gz
179+
- name: Download all artifacts
180+
uses: actions/download-artifact@v4
181+
if: startsWith(github.ref, 'refs/tags/')
182+
- name: Release
183+
uses: softprops/action-gh-release@v2
184+
if: startsWith(github.ref, 'refs/tags/')
185+
with:
186+
files: |
187+
pjrt-rocm_linux-amd64.tar.gz/pjrt-rocm_linux-amd64.tar.gz
188+
pjrt-cuda_linux-amd64.tar.gz/pjrt-cuda_linux-amd64.tar.gz
189+
pjrt-cpu_linux-amd64.tar.gz/pjrt-cpu_linux-amd64.tar.gz
190+
pjrt-cpu_darwin-arm64.tar.gz/pjrt-cpu_darwin-arm64.tar.gz

‎cpu/BUILD.bazel‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,8 +65,6 @@ oci_image(
6565
"//base/jammy:packages",
6666
":strip_packages",
6767
":symlinks",
68-
"//third_party/openxla:openxla_tar",
69-
"//third_party/openxla:cpu_configuration_tar",
7068
],
7169
env = {
7270
"MAX_NUM_WORKERS": -1,

‎cuda/BUILD.bazel‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,8 +67,6 @@ oci_image(
6767
"//base/noble:bazelisk",
6868
":strip_packages",
6969
":symlinks",
70-
"//third_party/openxla:openxla_tar",
71-
"//third_party/openxla:cuda_configuration_tar",
7270
],
7371
env = {
7472
"MAX_NUM_WORKERS": -1,

third_party/openxla/common.bazelrc renamed to openxla/bazelrc/common.bazelrc

File renamed without changes.

‎openxla/bazelrc/cpu.bazelrc‎

Whitespace-only changes.

third_party/openxla/cuda.bazelrc renamed to openxla/bazelrc/cuda.bazelrc

File renamed without changes.

third_party/openxla/rocm.bazelrc renamed to openxla/bazelrc/rocm.bazelrc

File renamed without changes.

third_party/openxla/patches/20240131-001-Expose-PJRT-mlir_to_hlo-to-public.patch renamed to openxla/patches/20240131-001-Expose-PJRT-mlir_to_hlo-to-public.patch

File renamed without changes.

third_party/openxla/patches/20240318-001-PJRT-C-API-Ensure-C-compliance-for-Profiler-Extension.patch renamed to openxla/patches/20240318-001-PJRT-C-API-Ensure-C-compliance-for-Profiler-Extension.patch

File renamed without changes.

third_party/openxla/patches/20240901-001-Do-not-force-DEVELOPER_DIR-on-macOS.patch renamed to openxla/patches/20240901-001-Do-not-force-DEVELOPER_DIR-on-macOS.patch

File renamed without changes.

0 commit comments

Comments
 (0)