Add wait_for_devices_ready utility to pathwaysutils. - #314
Add wait_for_devices_ready utility to pathwaysutils.#314copybara-service[bot] wants to merge 1 commit into
Conversation
0daf83d to
cf95576
Compare
|
|
||
| _logger.info("Waiting for %d devices to be ready.", len(devices)) | ||
| results = [jax.jit(lambda x: x + 1, device=d)(0) for d in devices] | ||
| jax.block_until_ready(results) |
There was a problem hiding this comment.
Would it make sense to add a timeout here (controlled through an arg)
E.g. users might want to fail fast if devices are not ready after 30 mins.
Workloads connecting to Pathways or initializing JAX may need to ensure all target devices are placed, warm, and ready to accept computations before starting training or inference. PiperOrigin-RevId: 968248469
cf95576 to
7e0f0b0
Compare
| ) | ||
|
|
||
|
|
||
| def wait_for_devices_ready( |
There was a problem hiding this comment.
This is pretty much duplicating the existing wait_for_slices API at https://github.com/AI-Hypercomputer/pathways-utils/blob/main/pathwaysutils%2Felastic%2Felastic.py#L173
Please clarify why you cannot use that API or propose a modification to that API to fit your desired use case
There was a problem hiding this comment.
The API is different and implementation is as well. The elastic wait for slices does more things.
I wanted an API that waits_for_devices_ready and requires no arguments to pass or figure out how many slices there are. In addition I don't want to have to check the return status. I just want it to wait on for all devices to be ready and continue when they are ready. Or raise an exception when timeout is specified and timeout is reached.
There was a problem hiding this comment.
I would suggest writing a wrapper over that API so you do not do duplicate work like checking for timeout, test program to check health etc.
3 participants
Add wait_for_devices_ready utility to pathwaysutils.
Workloads connecting to Pathways or initializing JAX may need to ensure all target devices are placed, warm, and ready to accept computations before starting training or inference.