(elastic_manager: manager.Manager)
| 319 | |
| 320 | |
| 321 | def wait_for_all_slices(elastic_manager: manager.Manager) -> None: |
| 322 | elastic_manager.good_slice_indices = elastic_manager.get_slice_availability() |
| 323 | while len(elastic_manager.good_slice_indices) < elastic_manager.total_slice_count: |
| 324 | max_logging.log( |
| 325 | f"Only {elastic_manager.good_slice_count} slices out of {elastic_manager.total_slice_count} available. " |
| 326 | "Sleeping for 5 seconds." |
| 327 | ) |
| 328 | time.sleep(5) |
| 329 | elastic_manager.good_slice_indices = elastic_manager.get_slice_availability() |
| 330 | max_logging.log("All slices are available") |
| 331 | |
| 332 | |
| 333 | def elastic_initialize(devices: Sequence[jax.Device]) -> manager.Manager: |
no outgoing calls
no test coverage detected