Skip to content

Commit a057dbe

Browse files
juliagmt-googlejuliagmt-google
authored
Update benchmarks.yml
1 parent 5afa062 commit a057dbe

File tree

1 file changed

+52
-52
lines changed

1 file changed

+52
-52
lines changed

‎.github/workflows/benchmarks.yml‎

Lines changed: 52 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -16,50 +16,50 @@ on:
1616
- 'no'
1717

1818
jobs:
19-
# jax-build:
20-
# strategy:
21-
# matrix:
22-
# runner: ["linux-x86-g2-48-l4-4gpu"]
19+
jax-build-and-test:
20+
strategy:
21+
matrix:
22+
runner: ["linux-x86-g2-48-l4-4gpu"]
2323

24-
# runs-on: ${{ matrix.runner }}
25-
# container:
26-
# image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
24+
runs-on: ${{ matrix.runner }}
25+
container:
26+
image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
2727

28-
# env:
29-
# JAXCI_HERMETIC_PYTHON_VERSION: 3.11
28+
env:
29+
JAXCI_HERMETIC_PYTHON_VERSION: 3.11
3030

31-
# steps:
32-
# - name: Checkout JAX Fork
33-
# uses: actions/checkout@v3
34-
# with:
35-
# repository: 'google-ml-infra/jax-fork'
36-
# path: jax-fork
31+
steps:
32+
- name: Checkout JAX Fork
33+
uses: actions/checkout@v3
34+
with:
35+
repository: 'google-ml-infra/jax-fork'
36+
path: jax-fork
3737

38-
# - name: Install JAX Dependencies
39-
# working-directory: jax-fork
40-
# run: |
41-
# python -m pip install --upgrade pip
42-
# pip install pytest
43-
# pip install absl-py
44-
# pip install "jax[cuda12_pip]" # Adjust CUDA version if needed
45-
# pip install google-benchmark
38+
- name: Install JAX Dependencies
39+
working-directory: jax-fork
40+
run: |
41+
python -m pip install --upgrade pip
42+
pip install pytest
43+
pip install absl-py
44+
pip install "jax[cuda12_pip]" # Adjust CUDA version if needed
45+
pip install google-benchmark
4646
47-
# - name: Run JAX Multiprocess GPU Test
48-
# working-directory: jax-fork
49-
# continue-on-error: true
50-
# run: python -m pytest tests/multiprocess_gpu_test.py
47+
- name: Run JAX Multiprocess GPU Test
48+
working-directory: jax-fork
49+
continue-on-error: true
50+
run: python -m pytest tests/multiprocess_gpu_test.py
5151

52-
# - name: Build XLA GPU Atomic Test
53-
# working-directory: xla
54-
# continue-on-error: true
55-
# run: bazel build -c opt --config=cuda //xla/service/gpu/tests:gpu_atomic_test
52+
- name: Build XLA GPU Atomic Test
53+
working-directory: xla
54+
continue-on-error: true
55+
run: bazel build -c opt --config=cuda //xla/service/gpu/tests:gpu_atomic_test
5656

57-
# - name: Run XLA GPU Atomic Test
58-
# working-directory: xla
59-
# continue-on-error: true
60-
# run: bazel test -c opt --config=cuda //xla/service/gpu/tests:gpu_atomic_test
57+
- name: Run XLA GPU Atomic Test
58+
working-directory: xla
59+
continue-on-error: true
60+
run: bazel test -c opt --config=cuda //xla/service/gpu/tests:gpu_atomic_test
6161

62-
build-xla-gpu:
62+
build-xla-gpu-and-test:
6363
runs-on: linux-x86-g2-48-l4-4gpu # Use a GPU-enabled runner
6464
container:
6565
image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
@@ -72,20 +72,20 @@ jobs:
7272
repository: openxla/xla # Replace with your fork if needed
7373
path: xla
7474

75-
# - name: Build XLA with GPU support with RBE
76-
# working-directory: xla
77-
# continue-on-error: true
78-
# run: bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,gpu,-rocm-only --test_tag_filters=-no_oss,requires-gpu-nvidia,gpu,-rocm-only,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-amd --config=warnings --config=rbe_linux_cuda_nvcc --run_under=//tools/ci_build/gpu_build:parallel_gpu_execute --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --@cuda_driver//:enable_forward_compatibility=true --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/... //build_tools/... @tsl//tsl/...
75+
- name: Build XLA with GPU support with RBE
76+
working-directory: xla
77+
continue-on-error: true
78+
run: bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,gpu,-rocm-only --test_tag_filters=-no_oss,requires-gpu-nvidia,gpu,-rocm-only,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-amd --config=warnings --config=rbe_linux_cuda_nvcc --run_under=//tools/ci_build/gpu_build:parallel_gpu_execute --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --@cuda_driver//:enable_forward_compatibility=true --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/... //build_tools/... @tsl//tsl/...
7979

80-
# - name: Run XLA tests with GPU support with RBE
81-
# working-directory: xla
82-
# continue-on-error: true
83-
# run: bazel test --build_tag_filters=-no_oss,requires-gpu-nvidia,gpu,-rocm-only --test_tag_filters=-no_oss,requires-gpu-nvidia,gpu,-rocm-only,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-amd --config=warnings --config=rbe_linux_cuda_nvcc --run_under=//tools/ci_build/gpu_build:parallel_gpu_execute --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --@cuda_driver//:enable_forward_compatibility=true --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/... //build_tools/... @tsl//tsl/...
80+
- name: Run XLA tests with GPU support with RBE
81+
working-directory: xla
82+
continue-on-error: true
83+
run: bazel test --build_tag_filters=-no_oss,requires-gpu-nvidia,gpu,-rocm-only --test_tag_filters=-no_oss,requires-gpu-nvidia,gpu,-rocm-only,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-amd --config=warnings --config=rbe_linux_cuda_nvcc --run_under=//tools/ci_build/gpu_build:parallel_gpu_execute --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --@cuda_driver//:enable_forward_compatibility=true --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/... //build_tools/... @tsl//tsl/...
8484

85-
# - name: Run Profile Analysis
86-
# working-directory: xla
87-
# continue-on-error: true
88-
# run: bazel analyze-profile profile.json.gz
85+
- name: Run Profile Analysis
86+
working-directory: xla
87+
continue-on-error: true
88+
run: bazel analyze-profile profile.json.gz
8989

9090
- name: Get GPU spec
9191
working-directory: xla
@@ -104,10 +104,10 @@ jobs:
104104
working-directory: xla
105105
run: bazel build -c opt --dynamic_mode=off xla/tools:run_hlo_module --config=cuda
106106

107-
- name: Wait For Connection
108-
uses: google-ml-infra/actions/ci_connection@main
109-
with:
110-
halt-dispatch-input: ${{ inputs.halt-for-connection }}
107+
# - name: Wait For Connection
108+
# uses: google-ml-infra/actions/ci_connection@main
109+
# with:
110+
# halt-dispatch-input: ${{ inputs.halt-for-connection }}
111111

112112
- name: Run HLO Module Benchmarks withg GPU in xla/tests/fuzz
113113
working-directory: xla

0 commit comments

Comments
 (0)