Skip to content

speculators.train.recovery

Classes:

BatchRecoveryCoordinator 

BatchRecoveryCoordinator(
    phase: str, *, device: device | int | None = None
)

Consume batch metadata and periodically coordinate recovery across ranks.

Methods:

  • consume –

    Remove recovery metadata and optionally synchronize pending state.

  • synchronize –

    Aggregate pending failures and raise a circuit breaker on every rank.

Source code in speculators/train/recovery.py 
148
149
150
151
152
153
154
155
156
def __init__(
    self,
    phase: str,
    *,
    device: torch.device | int | None = None,
) -> None:
    self.phase = phase
    self.device = device if device is not None else self._default_device()
    self._pending = RecoveryMetadata()

consume 

consume(
    batch: dict[str, Any], *, synchronize: bool = False
) -> None

Remove recovery metadata and optionally synchronize pending state.

Source code in speculators/train/recovery.py 
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
def consume(self, batch: dict[str, Any], *, synchronize: bool = False) -> None:
    """Remove recovery metadata and optionally synchronize pending state."""
    metadata = batch.pop(RECOVERY_METADATA_KEY, RecoveryMetadata())
    if not isinstance(metadata, RecoveryMetadata):
        raise TypeError(
            f"{RECOVERY_METADATA_KEY} must contain RecoveryMetadata, "
            f"got {type(metadata).__name__}"
        )
    if leftover := sorted(key for key in batch if key.startswith("__")):
        raise RuntimeError(f"Unconsumed data recovery metadata: {leftover}")

    if metadata.failure_count or metadata.locally_empty:
        logger.warning(
            "%s data recovery on rank %d: failed_samples=%d, "
            "locally_empty_batch=%s",
            self.phase,
            get_rank(),
            metadata.failure_count,
            metadata.locally_empty,
        )

    self._pending = RecoveryMetadata(
        failure_count=(self._pending.failure_count + metadata.failure_count),
        locally_empty=self._pending.locally_empty or metadata.locally_empty,
        fatal=self._pending.fatal or metadata.fatal,
        error=self._pending.error or metadata.error,
    )

    if synchronize or not is_distributed():
        self.synchronize()

synchronize 

synchronize() -> None

Aggregate pending failures and raise a circuit breaker on every rank.

Source code in speculators/train/recovery.py 
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
def synchronize(self) -> None:
    """Aggregate pending failures and raise a circuit breaker on every rank."""
    pending = self._pending
    self._pending = RecoveryMetadata()

    if is_distributed():
        status = torch.tensor(
            [int(pending.fatal), pending.failure_count],
            dtype=torch.int64,
            device=self.device,
        )
        dist.all_reduce(status, op=dist.ReduceOp.SUM)
        fatal_ranks, total_failures = (int(value) for value in status.tolist())
    else:
        fatal_ranks = int(pending.fatal)
        total_failures = pending.failure_count

    rank = get_rank()
    if total_failures and rank == 0:
        logger.warning(
            "%s hidden-state generation dropped %d sample(s) across ranks; "
            "affected ranks continue with remaining or locally empty data.",
            self.phase,
            total_failures,
        )
    if fatal_ranks:
        detail = f" Local error: {pending.error}" if pending.error else ""
        raise RuntimeError(
            "Hidden-state generation circuit breaker tripped on "
            f"{fatal_ranks} rank(s) during {self.phase}.{detail}"
        )

GenerationRecoveryGuard 

GenerationRecoveryGuard(
    retries: int, max_consecutive_failures: int
)

Retry generation operations and bound consecutive exhausted samples.

DataLoader workers each receive their own dataset copy, so a guard stored on the dataset naturally tracks failures independently per persistent worker.

Methods:

  • run –

    Run operation until it succeeds or all attempts are exhausted.

Attributes:

Source code in speculators/train/recovery.py 
66
67
68
69
70
71
72
73
def __init__(self, retries: int, max_consecutive_failures: int):
    if retries < 0:
        raise ValueError("retries must be >= 0")
    if max_consecutive_failures <= 0:
        raise ValueError("max_consecutive_failures must be > 0")
    self.retries = retries
    self.max_consecutive_failures = max_consecutive_failures
    self._consecutive_failures = 0

consecutive_failures property 

consecutive_failures: int

Number of consecutive samples which exhausted all attempts.

run 

run(
    operation: Callable[[], T], *, description: str
) -> T | SampleUnavailable

Run operation until it succeeds or all attempts are exhausted.

Source code in speculators/train/recovery.py 
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
def run(
    self,
    operation: Callable[[], T],
    *,
    description: str,
) -> T | SampleUnavailable:
    """Run ``operation`` until it succeeds or all attempts are exhausted."""
    total_attempts = self.retries + 1
    for attempt in range(1, total_attempts + 1):
        try:
            result = operation()
        except Exception as error:  # noqa: BLE001 - worker recovery boundary
            if attempt == total_attempts:
                return self._record_exhausted(
                    description,
                    attempt=attempt,
                    total_attempts=total_attempts,
                    error=error,
                )
            logger.warning(
                "%s (attempt %d/%d): %s: %s",
                description,
                attempt,
                total_attempts,
                type(error).__name__,
                error,
            )
            continue

        self._record_success()
        return result

    raise RuntimeError("Generation retry loop completed without a result")

RecoveryMetadata dataclass 

RecoveryMetadata(
    failure_count: int = 0,
    locally_empty: bool = False,
    fatal: bool = False,
    error: str = "",
)

Recovery state carried from the collator to the trainer.

Methods:

from_unavailable classmethod 

from_unavailable(
    unavailable: Sequence[SampleUnavailable],
    *,
    locally_empty: bool,
) -> RecoveryMetadata

Summarize unavailable samples for one collated batch.

Source code in speculators/train/recovery.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
@classmethod
def from_unavailable(
    cls,
    unavailable: Sequence[SampleUnavailable],
    *,
    locally_empty: bool,
) -> RecoveryMetadata:
    """Summarize unavailable samples for one collated batch."""
    failures = [sample for sample in unavailable if sample.counts_as_failure]
    fatal = next((sample for sample in failures if sample.fatal), None)
    return cls(
        failure_count=len(failures),
        locally_empty=locally_empty,
        fatal=fatal is not None,
        error=fatal.reason if fatal is not None else "",
    )

SampleUnavailable dataclass 

SampleUnavailable(
    reason: str = "",
    counts_as_failure: bool = False,
    consecutive_failures: int = 0,
    fatal: bool = False,
)

Picklable result for a sample the dataset could not provide.