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.

Example

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.