AI-Hypercomputer / AI-Hypercomputer/ray-tpu

Maybe fork a separate process to monitor execution

Open
#13 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
15
Forks
8
PR merge metrics
No merged PRs in 30d

Description

A main use case for Ray-on-TPU is for orchestrating JAX training jobs. XLA likes to abort the process when things go wrong, whichh ends the process outside of Python control flow and means that Ray can't tell the difference between the Ray process crashing due to some internal Ray error or XLA killing it. To help with retry logic, we have found it useful to fork a process to run the main payload function.

We're happy to keep this in our library if you don't think it's necessary. Just raising it for now.

Our implementation and justification: https://github.com/stanford-crfm/levanter/blob/94afdc17f6091249e70e90cb27b4378d0553ff56/src/levanter/infra/ray_tpu.py#L459

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.