Skip to content

[train][jax_trainer] add jax.distributed.shutdown() for JaxBackend#57802

Merged
matthewdeng merged 7 commits intoray-project:masterfrom
liulehui:shutdown_jaxtrainer
Oct 20, 2025
Merged

[train][jax_trainer] add jax.distributed.shutdown() for JaxBackend#57802
matthewdeng merged 7 commits intoray-project:masterfrom
liulehui:shutdown_jaxtrainer

Conversation

@liulehui
Copy link
Contributor

@liulehui liulehui commented Oct 16, 2025 •
edited
Loading

Description

  1. This PR added the jax.distributed.shutdown() for JaxBackend in order to free up any leaked resources on TPU RayTrainWorkers.
  2. if jax.distributed is not on, it is a noop: https://docs.jax.dev/en/latest/_autosummary/jax.distributed.shutdown.html
  3. Tested on Anyscale workspace.
image

Related issues

Types of change

  • Bug fix 🐛
  • New feature ✨
  • Enhancement 🚀
  • Code refactoring 🔧
  • Documentation update 📖
  • Chore 🧹
  • Style 🎨

Checklist

Does this PR introduce breaking changes?

  • Yes ⚠️
  • No

Testing:

  • Added/updated tests for my changes
  • Tested the changes manually
  • This PR is not tested ❌ (please explain why)

Code Quality:

  • Signed off every commit (git commit -s)
  • Ran pre-commit hooks (setup guide)

Documentation:

  • Updated documentation (if applicable) (contribution guide)
  • Added new APIs to doc/source/ (if applicable)

Additional context

@liulehui
Add shutdown for jaxtrainer
3d09927 
Signed-off-by: Lehui Liu <lehui@anyscale.com>
@liulehui
change unit test to capture debug log
ed041ec 
Signed-off-by: Lehui Liu <lehui@anyscale.com>
@liulehui liulehui added the go add ONLY when ready to merge, run all tests label Oct 16, 2025
@liulehui
add debug log check in controller
0b92a06 
Signed-off-by: Lehui Liu <lehui@anyscale.com>
@liulehui
not asserting logging since it is ray actors
b1f131f 
Signed-off-by: Lehui Liu <lehui@anyscale.com>
matthewdeng
matthewdeng reviewed Oct 16, 2025
View reviewed changes
python/ray/train/v2/jax/config.py
Copy link
Contributor

@matthewdeng matthewdeng Oct 16, 2025

Choose a reason for hiding this comment

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

Any way we could add a unit test for this?

Copy link
Contributor Author

@liulehui liulehui Oct 17, 2025

Choose a reason for hiding this comment

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

I tried a few way including using caplog and try to check jax.process_count()..
but seems Jax does not well support the cpu distributed env, I am adding a release test for it, would it be ok to check there: #57815

@liulehui
address comment
f18748c 
Signed-off-by: Lehui Liu <lehui@anyscale.com>
@liulehui
add jax backend timeout
477407d 
Signed-off-by: Lehui Liu <lehui@anyscale.com>
@liulehui
address comments
3c1e0a8 
Signed-off-by: Lehui Liu <lehui@anyscale.com>
@liulehui liulehui marked this pull request as ready for review October 17, 2025 22:56
@liulehui liulehui requested a review from a team as a code owner October 17, 2025 22:56
@ray-gardener ray-gardener bot added the train Ray Train Related Issue label Oct 18, 2025
@matthewdeng matthewdeng merged commit bf206b2 into ray-project:master Oct 20, 2025
6 checks passed
xinyuangui2 pushed a commit to xinyuangui2/ray that referenced this pull request Oct 22, 2025
@liulehui @xinyuangui2
[train][jax_trainer] add jax.distributed.shutdown() for JaxBackend (r…
f15a6bf 
…ay-project#57802)

## Description

1. This PR added the `jax.distributed.shutdown()` for JaxBackend in
order to free up any leaked resources on TPU RayTrainWorkers.
2. if `jax.distributed` is not on, it is a noop:
https://docs.jax.dev/en/latest/_autosummary/jax.distributed.shutdown.html
3. Tested on Anyscale workspace.
<img width="1264" height="62" alt="image"
src="https://github.com/user-attachments/assets/f28102ff-f6d1-4da0-b41a-6cc785603e72"
/>

Signed-off-by: xgui <xgui@anyscale.com>
elliot-barn pushed a commit that referenced this pull request Oct 23, 2025
@liulehui @elliot-barn
[train][jax_trainer] add jax.distributed.shutdown() for JaxBackend (#…
6e1f067 
…57802)

## Description

1. This PR added the `jax.distributed.shutdown()` for JaxBackend in
order to free up any leaked resources on TPU RayTrainWorkers.
2. if `jax.distributed` is not on, it is a noop:
https://docs.jax.dev/en/latest/_autosummary/jax.distributed.shutdown.html
3. Tested on Anyscale workspace.
<img width="1264" height="62" alt="image"
src="https://github.com/user-attachments/assets/f28102ff-f6d1-4da0-b41a-6cc785603e72"
/>

Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>
landscapepainter pushed a commit to landscapepainter/ray that referenced this pull request Nov 17, 2025
@liulehui @landscapepainter
[train][jax_trainer] add jax.distributed.shutdown() for JaxBackend (r…
0930f30 
…ay-project#57802)

## Description

1. This PR added the `jax.distributed.shutdown()` for JaxBackend in
order to free up any leaked resources on TPU RayTrainWorkers.
2. if `jax.distributed` is not on, it is a noop:
https://docs.jax.dev/en/latest/_autosummary/jax.distributed.shutdown.html
3. Tested on Anyscale workspace.
<img width="1264" height="62" alt="image"
src="https://github.com/user-attachments/assets/f28102ff-f6d1-4da0-b41a-6cc785603e72"
/>
Aydin-ab pushed a commit to Aydin-ab/ray-aydin that referenced this pull request Nov 19, 2025
@liulehui @Aydin-ab
[train][jax_trainer] add jax.distributed.shutdown() for JaxBackend (r…
e50ae20 
…ay-project#57802)

## Description

1. This PR added the `jax.distributed.shutdown()` for JaxBackend in
order to free up any leaked resources on TPU RayTrainWorkers.
2. if `jax.distributed` is not on, it is a noop:
https://docs.jax.dev/en/latest/_autosummary/jax.distributed.shutdown.html
3. Tested on Anyscale workspace.
<img width="1264" height="62" alt="image"
src="https://github.com/user-attachments/assets/f28102ff-f6d1-4da0-b41a-6cc785603e72"
/>

Signed-off-by: Aydin Abiar <aydin@anyscale.com>
Future-Outlier pushed a commit to Future-Outlier/ray that referenced this pull request Dec 7, 2025
@liulehui @Future-Outlier
[train][jax_trainer] add jax.distributed.shutdown() for JaxBackend (r…
a59f776 
…ay-project#57802)

## Description

1. This PR added the `jax.distributed.shutdown()` for JaxBackend in
order to free up any leaked resources on TPU RayTrainWorkers.
2. if `jax.distributed` is not on, it is a noop:
https://docs.jax.dev/en/latest/_autosummary/jax.distributed.shutdown.html
3. Tested on Anyscale workspace.
<img width="1264" height="62" alt="image"
src="https://github.com/user-attachments/assets/f28102ff-f6d1-4da0-b41a-6cc785603e72"
/>

Signed-off-by: Future-Outlier <eric901201@gmail.com>
Blaze-DSP pushed a commit to Blaze-DSP/ray that referenced this pull request Dec 18, 2025
@liulehui
[train][jax_trainer] add jax.distributed.shutdown() for JaxBackend (r…
16ca8bc 
…ay-project#57802)

## Description

1. This PR added the `jax.distributed.shutdown()` for JaxBackend in
order to free up any leaked resources on TPU RayTrainWorkers.
2. if `jax.distributed` is not on, it is a noop:
https://docs.jax.dev/en/latest/_autosummary/jax.distributed.shutdown.html
3. Tested on Anyscale workspace.
<img width="1264" height="62" alt="image"
src="https://github.com/user-attachments/assets/f28102ff-f6d1-4da0-b41a-6cc785603e72"
/>
peterxcli pushed a commit to peterxcli/ray that referenced this pull request Feb 25, 2026
@liulehui @peterxcli
[train][jax_trainer] add jax.distributed.shutdown() for JaxBackend (r…
9520365 
…ay-project#57802)

## Description

1. This PR added the `jax.distributed.shutdown()` for JaxBackend in
order to free up any leaked resources on TPU RayTrainWorkers.
2. if `jax.distributed` is not on, it is a noop:
https://docs.jax.dev/en/latest/_autosummary/jax.distributed.shutdown.html
3. Tested on Anyscale workspace.
<img width="1264" height="62" alt="image"
src="https://github.com/user-attachments/assets/f28102ff-f6d1-4da0-b41a-6cc785603e72"
/>

Signed-off-by: peterxcli <peterxcli@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@matthewdeng matthewdeng matthewdeng approved these changes

Assignees

No one assigned

Labels

go add ONLY when ready to merge, run all tests train Ray Train Related Issue

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@liulehui @matthewdeng