Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
260 commits
Select commit Hold shift + click to select a range
222542d
change docker image for bazel local gpu test workflow
nitins17 Aug 28, 2024
8fe81fb
change image to tag instead of sha
nitins17 Aug 28, 2024
6666472
set default Python version for all builds
nitins17 Aug 28, 2024
f7b47cf
set default Python version for jax-cuda-pjrt and read CLI argument fo…
nitins17 Aug 28, 2024
057b8e3
Create a matrix for building JAX artifacts
nitins17 Aug 28, 2024
843afe3
add default python to matrix for jax-cuda-pjrt
nitins17 Aug 28, 2024
82f0454
exlude pjrt from the python matrix
nitins17 Aug 28, 2024
23fb8c9
add workflow and scripts for running pytests
nitins17 Aug 29, 2024
875a271
make pytest script executable
nitins17 Aug 29, 2024
6384f73
add workflow for pytest
nitins17 Aug 29, 2024
320bf83
change workflow name
nitins17 Aug 29, 2024
204c029
try fixing workflow
nitins17 Aug 29, 2024
cc4e210
fix runner name in pytest workflow
nitins17 Aug 29, 2024
4474934
add connection wait for pytest workflow and fix typo in container fie…
nitins17 Aug 29, 2024
6ae0257
dummy change; push to make github actions detect labels
nitins17 Aug 29, 2024
d7ab9e8
dummy change; push to make github actions detect labels
nitins17 Aug 29, 2024
4858b38
set clang path for all linux ci builds
nitins17 Aug 29, 2024
30a3951
install wheels with system python only for pytests
nitins17 Aug 29, 2024
ff2e5fe
Merge branch 'main' of https://github.com/google-ml-infra/jax-fork in…
nitins17 Aug 30, 2024
60f1231
install missing docker dependencies
nitins17 Aug 30, 2024
e17ca80
install missing docker deps specifically with python3
nitins17 Aug 30, 2024
d9bfc98
fix workflows failing
nitins17 Aug 30, 2024
b1b0953
update docker images in the artifacts workflow
nitins17 Sep 3, 2024
045f4b1
no-op; make github detect label change
nitins17 Sep 3, 2024
d68dcc2
add remaining scripts and workflows for running tests
nitins17 Sep 4, 2024
6105043
update workflow names
nitins17 Sep 4, 2024
22ce7ac
fix typos and add missing env
nitins17 Sep 4, 2024
b1a6ace
disable artifact build workflow
nitins17 Sep 4, 2024
ea3cba3
pass hermetic python env correctly in bazel cpu workflow
nitins17 Sep 4, 2024
077d76b
update gpu workfow runner name
nitins17 Sep 4, 2024
5981ec9
rework some envs
nitins17 Sep 4, 2024
b0b3171
rename JAXCI_USE_RBE to JAXCI_BUILD_ARTIFACT_WITH_RBE
nitins17 Sep 4, 2024
6cba3bd
Fix typo
nitins17 Sep 4, 2024
4c8afb9
add windows runner to build artifacts workflow and disable gpu workflows
nitins17 Sep 4, 2024
6630dc0
clean up bazel gpu tests
nitins17 Sep 4, 2024
4c6090e
rework how envs are passed to docker container inside setup_docker.sh
nitins17 Sep 4, 2024
850fa06
disable running in container to test windows builds
nitins17 Sep 4, 2024
958010d
enable bazel gpu jobs
nitins17 Sep 4, 2024
1f91951
tweak some setup settings
nitins17 Sep 4, 2024
1de1480
change how envs are passed to the container in setup_docker.sh
nitins17 Sep 5, 2024
3be81b6
add wait for connection to gpu workflow for testing and fix small iss…
nitins17 Sep 5, 2024
ff0062d
install wheels locally for bazel test local
nitins17 Sep 5, 2024
f57f3fb
disable cpu test workflows and move wait for connection in bazel gpu …
nitins17 Sep 5, 2024
6544ca9
make github detect label change
nitins17 Sep 7, 2024
cbdcbbf
Explictly make the shell as bash for script to run on Windows
nitins17 Sep 7, 2024
50d3380
fix issue with scripts when running on windows runners
nitins17 Sep 8, 2024
6674827
build without rbe to test windows runners
nitins17 Sep 8, 2024
dbeca71
Merge branch 'main' of https://github.com/google-ml-infra/jax-fork in…
nitins17 Sep 8, 2024
f4cb21e
remove root .bazelrc to avoid conflicts
nitins17 Sep 8, 2024
007356d
fixes for building a windows wheel
nitins17 Sep 8, 2024
662ddf6
fix wheel build command for windows
nitins17 Sep 8, 2024
f7d503a
enable artifact build on Linux runners
nitins17 Sep 8, 2024
59a0efa
apply container only for linux jobs
nitins17 Sep 8, 2024
1d01b3b
comment out container settings to fix artifact workflow and make gith…
nitins17 Sep 9, 2024
7171258
comment out linux runners
nitins17 Sep 9, 2024
0ed41d7
remove path to bash
nitins17 Sep 9, 2024
4e551a0
update condition to select rbe config
nitins17 Sep 9, 2024
042f6c5
make github detect label change
nitins17 Sep 9, 2024
46a9b30
enable rbe for artifact build
nitins17 Sep 9, 2024
f1fb463
make github detect label change
nitins17 Sep 12, 2024
44c7ad2
Merge branch 'main' into srnitin/task-jax-ci-rework
nitins17 Sep 12, 2024
5d766c9
Add more functionality to the build CLI
nitins17 Sep 13, 2024
cd8f7a9
Merge branch 'srnitin/task-jax-ci-rework' of https://github.com/googl…
nitins17 Sep 13, 2024
64a17eb
clean up scripts
nitins17 Sep 13, 2024
d2c899d
indent image field in build artifacts workflow
nitins17 Sep 13, 2024
57c1bf1
fix syntax error
nitins17 Sep 13, 2024
3dcf44f
another try at fixing syntax error
nitins17 Sep 13, 2024
4fb642d
fix envs
nitins17 Sep 13, 2024
8c69f79
test cpu test workflows
nitins17 Sep 14, 2024
83140e3
fix env in workflow
nitins17 Sep 14, 2024
7cc5bfc
some fixes
nitins17 Sep 14, 2024
92cecea
fix syntax error
nitins17 Sep 14, 2024
0f41ab4
explictly set shell to bash
nitins17 Sep 14, 2024
31ddd2d
convert msys paths to windows paths
nitins17 Sep 14, 2024
e2f4e63
make github detect label change
nitins17 Sep 14, 2024
6ca9a8a
Change how mysys paths are converted to windows-like paths
nitins17 Sep 14, 2024
f4c6761
minor fix
nitins17 Sep 14, 2024
f83da93
fix error when args are None
nitins17 Sep 14, 2024
bfc334d
make github detect label change
nitins17 Sep 14, 2024
9db6fda
remove quotes around local xla path to fix windows issue
nitins17 Sep 14, 2024
13fe6f9
fix path issue on windows when installing wheels locally
nitins17 Sep 15, 2024
ae65260
clean up
nitins17 Sep 16, 2024
e7e7dca
update scripts
nitins17 Sep 16, 2024
515c397
update env var name
nitins17 Sep 16, 2024
87d277a
update default env file
nitins17 Sep 16, 2024
fa21235
update env var
nitins17 Sep 16, 2024
877b12a
update scripts
nitins17 Sep 17, 2024
88ebac3
update .bazelrc
nitins17 Sep 17, 2024
1ba0bf6
update envs
nitins17 Sep 17, 2024
faa98f3
update default env
nitins17 Sep 17, 2024
0bfc868
Merge branch 'main' of https://github.com/google-ml-infra/jax-fork in…
nitins17 Sep 17, 2024
36d04de
update scripts
nitins17 Sep 19, 2024
3b55033
update scripts
nitins17 Sep 19, 2024
5627330
update pytest workflow
nitins17 Sep 19, 2024
31a6767
make github detect label change
nitins17 Sep 19, 2024
3a910bd
update bazel gpu local workflow
nitins17 Sep 19, 2024
766738f
enable bazel gpu local workflow
nitins17 Sep 19, 2024
5f3ee2c
fix syntax error
nitins17 Sep 19, 2024
12590a8
update env var
nitins17 Sep 19, 2024
534bc1a
update env
nitins17 Sep 19, 2024
7ba4a37
update linux arm64 container image
nitins17 Sep 19, 2024
7a2cc1c
disable import jax command in bazel gpu local script
nitins17 Sep 19, 2024
8b65c70
add env var setting to pytest cpu env file
nitins17 Sep 19, 2024
dc11f12
enable pytest gpu workflow
nitins17 Sep 20, 2024
ab6f3ff
update bazel cpu workflow
nitins17 Sep 20, 2024
e58035e
fix error in pytest gpu script
nitins17 Sep 20, 2024
637b0a5
fix installing wheels correctly
nitins17 Sep 20, 2024
460cb11
fix issue when installing wheels
nitins17 Sep 20, 2024
c65c2b2
modify scripts to run tpu tests
nitins17 Sep 20, 2024
c448ea6
set jaxci python env var
nitins17 Sep 20, 2024
5b450af
set the tpu cores in the env file
nitins17 Sep 20, 2024
1418ba8
enable bazel gpu rbe tests
nitins17 Sep 20, 2024
9f739ab
make github detect label change
nitins17 Sep 20, 2024
391a913
change to python 3.10
nitins17 Sep 23, 2024
3b98f9b
install the right version of python deps for pytest cpu
nitins17 Sep 23, 2024
dd41e8d
fix syntax error
nitins17 Sep 23, 2024
ed31404
Debug GPU failures
nitins17 Sep 23, 2024
04537e7
debug gpu vms getting killed
nitins17 Sep 23, 2024
b880ed5
debug gpu vms getting killed (2)
nitins17 Sep 23, 2024
a8716ce
set cuda compute arg correctly
nitins17 Sep 24, 2024
de602b5
set cuda compute arg correctly (2)
nitins17 Sep 24, 2024
3783c84
Merge branch 'main' of https://github.com/google-ml-infra/jax-fork in…
nitins17 Sep 24, 2024
bc6d9cb
do not override xla repository when running gpu tests locally
nitins17 Sep 24, 2024
b8978b5
bump device memory limit when running tests locally on single gpu
nitins17 Sep 25, 2024
26b1f59
change bazel build to bazel run
nitins17 Sep 26, 2024
7b63e0c
increase test timeout
nitins17 Sep 26, 2024
a7663fa
update runner for bazel local gpu workflow
nitins17 Sep 27, 2024
84c3daa
enable build artifact workflow
nitins17 Sep 27, 2024
b27028d
Add continue-on-error to build artifacts workflow
nitins17 Sep 27, 2024
d83cf08
update container images in build artifacts workflow
nitins17 Sep 27, 2024
a600e8c
change local bazel gpu test workflow to default gpu runner
nitins17 Sep 30, 2024
3bf0bff
Fix case when _DIR variables are empty
nitins17 Sep 30, 2024
c346896
make github detect label change
nitins17 Oct 1, 2024
730f746
trigger test runs
nitins17 Oct 2, 2024
0666e80
install k8s in pytest gpu
nitins17 Oct 2, 2024
d8efebb
install dependencies from test-requirements for pytest gpu
nitins17 Oct 2, 2024
c3e23a0
Merge branch 'main' of https://github.com/google-ml-infra/jax-fork in…
nitins17 Oct 2, 2024
e61ff65
remove || true from bazel test since we have now fixed the gpu test t…
nitins17 Oct 4, 2024
4454776
Merge branch 'main' of https://github.com/google-ml-infra/jax-fork in…
nitins17 Oct 4, 2024
49e5de8
switch the linux x86 build artifact workflow to use the ml-build cont…
nitins17 Oct 8, 2024
b527a68
update ml-build container
nitins17 Oct 9, 2024
09d6d8f
Merge branch 'main' of https://github.com/google-ml-infra/jax-fork in…
nitins17 Oct 9, 2024
960b28c
updaet ml-build container to latest tagged image
nitins17 Oct 9, 2024
286c91e
Experiment with resuable workflows
nitins17 Oct 16, 2024
504c9ad
make the pytest job depend on the build artifact job in the resuable …
nitins17 Oct 16, 2024
902b721
Remove matrix definitions from called workflow job
nitins17 Oct 16, 2024
e71eea9
remove artifact from caller workflow matrix strategy
nitins17 Oct 16, 2024
3aedc18
Enable running bazel cpu rbe workflow for windows
nitins17 Oct 23, 2024
cf4a68e
Merge branch 'main' of https://github.com/google-ml-infra/jax-fork in…
nitins17 Oct 23, 2024
6a374bf
remove ci/.bazelrc from the bazel build script
nitins17 Oct 23, 2024
3b61881
Make the default shell in bazel cpu workflow to be bash
nitins17 Oct 23, 2024
99d4098
Add conditional to correctly pass os value for windows
nitins17 Oct 23, 2024
0c10c84
fix syntax issue
nitins17 Oct 23, 2024
c9b63dc
Correctly pass architecture value for Windows runners to be able to r…
nitins17 Oct 23, 2024
02b79ef
make github detect label change
nitins17 Oct 23, 2024
d722633
Update build/build.py, scripts, envs to be simpler
nitins17 Oct 29, 2024
c18d386
make github detect label change
nitins17 Oct 29, 2024
6107d2b
fix unbound variable errors
nitins17 Oct 29, 2024
5ca7866
Fix how xla is cloned at head for workflows
nitins17 Oct 29, 2024
1bf2ca0
delete the ci/cli folder
nitins17 Oct 29, 2024
478696c
add the modules used by build.py
nitins17 Oct 29, 2024
434a753
Fix cuda version parsing and rename remaining envs with the .env ext
nitins17 Oct 30, 2024
18e64dd
Pass in the envs correctly (with .env ext) in the github workflows
nitins17 Oct 30, 2024
2c7fc89
add missing .env ext in some envs
nitins17 Oct 30, 2024
b64c59c
update cuda-pjrt env to fix sourcing default envs and remove redundan…
nitins17 Oct 30, 2024
fd97546
fix bazel cpu rbe env var name in bazel test script
nitins17 Oct 30, 2024
966fdd1
Produce verbose output when building artifacts
nitins17 Oct 30, 2024
daf31aa
replace custom shell execute function with subprocess.run()
nitins17 Oct 30, 2024
89833fe
Make debug logs simpler and fix unbbound variable error
nitins17 Oct 30, 2024
b15621e
update linux x86 docker image
nitins17 Oct 30, 2024
43aef52
Re-factor how docker envs are set and docker container gets set up
nitins17 Oct 30, 2024
3646516
Fix formating
nitins17 Oct 30, 2024
c8743b6
rename use_rbe flag to request_rbe
nitins17 Oct 30, 2024
55f53dd
fix formating
nitins17 Oct 30, 2024
0d7271d
Update --request_rbe help string
nitins17 Oct 30, 2024
46415cc
Change how flags are parsed if ci_mode is not set
nitins17 Oct 30, 2024
c3c6251
update --clang_path comment
nitins17 Oct 30, 2024
d72f55e
update rbe flag name
nitins17 Oct 30, 2024
e26333f
Make cuda args to be None by default and add function to retrive cuda…
nitins17 Oct 30, 2024
a5bf923
Ignore --build_cuda_with_clang when --ci_mode is set
nitins17 Oct 30, 2024
91d5ab9
fix typo
nitins17 Oct 30, 2024
a793371
Add missing auditwheel env var to jaxlib env and fix certain comments
nitins17 Nov 1, 2024
47fdbae
Move halt for connection step to debug auditwheel failure
nitins17 Nov 1, 2024
e484332
Switch python3 with python
nitins17 Nov 1, 2024
b04da6e
fix env var name
nitins17 Nov 1, 2024
68a8353
Update GPU/TPU test scripts
nitins17 Nov 2, 2024
a2ecde3
update python site packages path
nitins17 Nov 2, 2024
22f0c0c
update env vars
nitins17 Nov 2, 2024
be77781
change how ci mode flags are parsed
nitins17 Nov 2, 2024
fa91bef
Update Pytest envs and make minor change to get_ci_bazelrc_config
nitins17 Nov 4, 2024
2217b86
Parse ci bazelrc flags correctly
nitins17 Nov 4, 2024
240d284
Change how arch is detected on Windows machines
nitins17 Nov 4, 2024
330b76a
Fix path issues on Windows and fix comments
nitins17 Nov 6, 2024
ceea75e
Fix rocm case if use_ci_bazelrcs is passed and enable color output in CI
nitins17 Nov 6, 2024
92cbc6c
Fix comment
nitins17 Nov 6, 2024
ae819fa
update help string
nitins17 Nov 6, 2024
5a9c5f2
Enable color output by default
nitins17 Nov 6, 2024
b9a3140
update how bazel options are based for writing to .jax_configure
nitins17 Nov 6, 2024
783fdc8
update how subprocess command are constructed
nitins17 Nov 7, 2024
c605dab
do not apply --config=clang on Mac
nitins17 Nov 8, 2024
f43ea9c
GPU runner stress testing
nitins17 Nov 8, 2024
c51441e
GPU runner stress testing (2)
nitins17 Nov 8, 2024
0873ac8
Merge branch 'main' of https://github.com/google-ml-infra/jax-fork in…
nitins17 Nov 8, 2024
33ac8a2
nvidia-smi issue testing (1)
nitins17 Nov 8, 2024
55812be
nvidia-smi issue testing (2)
nitins17 Nov 8, 2024
970ffd5
split test commands into separate scripts
nitins17 Nov 8, 2024
97b9c95
Use the updated ci connect
nitins17 Nov 8, 2024
b39cea9
add set commands and remove setup_jaxci_envs.sh
nitins17 Nov 9, 2024
536219a
remove duplicated set commands
nitins17 Nov 9, 2024
d021e8a
refactor build.py to allow for building multiple artifacts with a sin…
nitins17 Nov 11, 2024
93025a5
update commen
nitins17 Nov 11, 2024
620eb99
update ci_connection step
nitins17 Nov 11, 2024
25ef454
update workflows to clone main xla via yaml after test env removal
nitins17 Nov 11, 2024
a0ee3ff
Merge branch 'main' of https://github.com/google-ml-infra/jax-fork in…
nitins17 Nov 11, 2024
f796ee6
update artifact build script with new build.py invocation
nitins17 Nov 11, 2024
98bcdc8
Set local XLA path if it already exists at /usr/local/google/home/srn…
nitins17 Nov 11, 2024
4ecdc9f
remove unusued envs and update comments
nitins17 Nov 11, 2024
8b2688a
Install test dependencies from build/requirements.in
nitins17 Nov 11, 2024
042a3dd
remove empty line
nitins17 Nov 11, 2024
08da390
Clone XLA at head if path to local XLA is not provided
nitins17 Nov 11, 2024
fb1c05f
set tpu cores
nitins17 Nov 11, 2024
cea3a12
accept pre-set value for tpu env
nitins17 Nov 11, 2024
481ed8d
Add local_xla_path and target_cpu to compile options and change some …
nitins17 Nov 11, 2024
f70e9e2
change to -n 4 to debug gpu worker crashes
nitins17 Nov 11, 2024
d977dca
add clone_main_xla env to each step that needs it
nitins17 Nov 11, 2024
250adf0
Add values for testing on CUDA 12.1 image
nitins17 Nov 11, 2024
2e83e5f
split build and test jobs
nitins17 Nov 11, 2024
acaade6
revert change to -n
nitins17 Nov 11, 2024
56eae2e
use the smaller runner when running with rbe
nitins17 Nov 11, 2024
2d4b21c
Update docker env and scripts
nitins17 Nov 12, 2024
8b1bad2
source default env
nitins17 Nov 12, 2024
4b7dc97
add set commands
nitins17 Nov 12, 2024
412bd83
make script executable
nitins17 Nov 12, 2024
10023fb
fix path
nitins17 Nov 12, 2024
bd71d94
disable workflows from running on prs while debugging docker scripts
nitins17 Nov 12, 2024
3692ac8
fix path
nitins17 Nov 12, 2024
ed3ce17
fix path
nitins17 Nov 12, 2024
5306637
fix path (final..)
nitins17 Nov 12, 2024
c2a482e
Remove jobs not needed for testing
MichaelHudgins Nov 12, 2024
d0fc20f
Allow on push
MichaelHudgins Nov 12, 2024
250e7af
Change clone
MichaelHudgins Nov 12, 2024
63b570f
Revert this test to only be a single runner
MichaelHudgins Nov 12, 2024
6a4877f
Flip the container
MichaelHudgins Nov 12, 2024
b7d8296
Pip install cuda
MichaelHudgins Nov 12, 2024
017f843
Move halt and halt for testing
MichaelHudgins Nov 12, 2024
e54045e
Move dispatch event
MichaelHudgins Nov 12, 2024
28b90f0
Recomment this
MichaelHudgins Nov 13, 2024
8fbe206
Merge branch 'main' into test-nvidia-smi-fail
MichaelHudgins Nov 14, 2024
c0c86b4
Mod the test slightly
MichaelHudgins Nov 14, 2024
3979f69
Remove spmm test due to sefault
MichaelHudgins Nov 15, 2024
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
46 changes: 0 additions & 46 deletions .github/workflows/bazel_cpu_rbe.yml

This file was deleted.

51 changes: 0 additions & 51 deletions .github/workflows/bazel_gpu_non_rbe.yml

This file was deleted.

40 changes: 0 additions & 40 deletions .github/workflows/bazel_gpu_rbe.yml

This file was deleted.

74 changes: 0 additions & 74 deletions .github/workflows/build_artifacts.yml

This file was deleted.

Loading
Loading