get_preemption_info

ray.train.get_preemption_info() → PreemptionInfo | None

Return the imminent preemption info for the current worker, or None.

Returns None until a node hosting one of the workers is being preempted (e.g. a spot instance reclaim). The recommended reaction is to save a just-in-time checkpoint and keep training: when the node is actually preempted, Ray Train restarts the run and resumes it from the latest checkpoint, retrying against FailureConfig.max_preemption_failures. A run that returns cleanly always finishes, whether or not a preemption is in progress.

Warning

All workers must call ray.train.get_preemption_info the same number of times so that Ray Train can agree on a single value across all workers. This method acts as a barrier across all workers, so be sure that every worker reaches this method.

Examples

>>> import ray.train
>>> def train_func(config):
...     saved_on_preemption = False
...     for step in range(config["total_steps"]):
...         # ... normal training step (with your usual periodic
...         # checkpointing) ...
...
...         preemption_info = ray.train.get_preemption_info()
...         if preemption_info is not None and not saved_on_preemption:
...             ray.train.report(metrics, checkpoint=checkpoint)
...             saved_on_preemption = True

Returns:

A PreemptionInfo with the affected node ids / world ranks and the reclaim deadline, or None if no preemption has been detected.

Return type:

PreemptionInfo | None

PublicAPI (alpha): This API is in alpha and may change before becoming stable.