Spaces:
Running on Zero
Running on Zero
| """Small fail-closed checks around model-family distributed process groups.""" | |
| from __future__ import annotations | |
| import hashlib | |
| import time | |
| from collections.abc import Callable, Mapping, Sequence | |
| from datetime import timedelta | |
| from typing import Any, TypeVar | |
| from solarwm.config.loader import canonical_json | |
| from solarwm.errors import BackendContractError | |
| T = TypeVar("T") | |
| class DataReadiness: | |
| """Finish remote batch reads on every rank before entering GPU collectives. | |
| All raw ranks construct and call this at the same reader boundary. The | |
| separate CPU group keeps storage latency out of the NCCL watchdog window; | |
| it does not select samples, advance RNGs, or move batch tensors. | |
| """ | |
| def __init__(self, *, enabled: bool) -> None: | |
| self.group: Any = None | |
| if not enabled: | |
| return | |
| import torch | |
| import torch.distributed as dist | |
| if not dist.is_initialized() or dist.get_world_size() == 1: | |
| return | |
| self.torch = torch | |
| self.dist = dist | |
| self.group = dist.new_group(backend="gloo", timeout=timedelta(hours=2)) | |
| def read(self, call: Callable[[], T]) -> T: | |
| if self.group is None: | |
| return call() | |
| started = time.monotonic() | |
| result: T | None = None | |
| local_error: Exception | None = None | |
| try: | |
| result = call() | |
| except Exception as exc: | |
| local_error = exc | |
| rank = self.dist.get_rank() | |
| world = self.dist.get_world_size() | |
| status = self.torch.tensor( | |
| [rank if local_error is not None else world], dtype=self.torch.int64, device="cpu" | |
| ) | |
| self.dist.all_reduce(status, op=self.dist.ReduceOp.MIN, group=self.group) | |
| failed_rank = int(status.item()) | |
| if failed_rank != world: | |
| messages = [ | |
| f"{type(local_error).__name__}: {local_error}" if rank == failed_rank else "" | |
| ] | |
| self.dist.broadcast_object_list(messages, src=failed_rank, group=self.group) | |
| raise BackendContractError( | |
| f"training data read failed on rank {failed_rank}: {messages[0]}" | |
| ) from local_error | |
| elapsed = time.monotonic() - started | |
| if elapsed >= 30: | |
| print( | |
| f"[data-readiness][rank{rank}] all batches ready after {elapsed:.1f}s", flush=True | |
| ) | |
| return result # type: ignore[return-value] | |
| def _collective_coordinates( | |
| dist: Any | None, | |
| *, | |
| rank: int | None, | |
| world_size: int | None, | |
| ) -> tuple[int, int, bool]: | |
| """Resolve one raw-world collective without importing Torch. | |
| Callers pass ``torch.distributed`` in heavy runtimes and a small fake in | |
| CPU tests. An explicitly multi-rank call may never silently degrade to a | |
| local check when the process group is unavailable. | |
| """ | |
| initialized = bool(dist is not None and dist.is_initialized()) | |
| resolved_rank = ( | |
| int(dist.get_rank()) if rank is None and initialized else int(0 if rank is None else rank) | |
| ) | |
| resolved_world = ( | |
| int(dist.get_world_size()) | |
| if world_size is None and initialized | |
| else int(1 if world_size is None else world_size) | |
| ) | |
| if resolved_world < 1 or not 0 <= resolved_rank < resolved_world: | |
| raise BackendContractError( | |
| "collective failure propagation received an invalid rank/world_size" | |
| ) | |
| if resolved_world > 1 and not initialized: | |
| raise BackendContractError( | |
| "multi-rank failure propagation requires initialized torch.distributed" | |
| ) | |
| if initialized: | |
| actual_rank = int(dist.get_rank()) | |
| actual_world = int(dist.get_world_size()) | |
| if (resolved_rank, resolved_world) != (actual_rank, actual_world): | |
| raise BackendContractError( | |
| "collective failure propagation rank/world_size differs from the process group" | |
| ) | |
| return resolved_rank, resolved_world, initialized | |
| def propagate_collective_error( | |
| local_error: str | None, | |
| *, | |
| dist: Any | None, | |
| label: str, | |
| rank: int | None = None, | |
| world_size: int | None = None, | |
| error_type: type[Exception] = BackendContractError, | |
| ) -> None: | |
| """Raise the same rank-ordered failure on every raw-world peer. | |
| Only strings cross the process-group boundary. Exception instances are | |
| deliberately never pickled or broadcast. Every peer must call this at | |
| the same phase boundary. | |
| """ | |
| resolved_rank, resolved_world, initialized = _collective_coordinates( | |
| dist, | |
| rank=rank, | |
| world_size=world_size, | |
| ) | |
| message = str(local_error or "") | |
| local = (resolved_rank, message) | |
| gathered: list[Any] = [local] | |
| if initialized and resolved_world > 1: | |
| gathered = [None] * resolved_world | |
| dist.all_gather_object(gathered, local) | |
| normalized: list[tuple[int, str]] = [] | |
| for item in gathered: | |
| if not isinstance(item, (tuple, list)) or len(item) != 2 or isinstance(item[0], bool): | |
| raise error_type(f"{label} failure exchange returned a malformed rank record") | |
| try: | |
| peer_rank = int(item[0]) | |
| except (TypeError, ValueError) as exc: | |
| raise error_type(f"{label} failure exchange returned a malformed rank value") from exc | |
| peer_error = str(item[1] or "") | |
| if not 0 <= peer_rank < resolved_world: | |
| raise error_type(f"{label} failure exchange returned an out-of-range rank") | |
| normalized.append((peer_rank, peer_error)) | |
| if len(normalized) != resolved_world or {item[0] for item in normalized} != set( | |
| range(resolved_world) | |
| ): | |
| raise error_type(f"{label} failure exchange did not cover every raw rank") | |
| failures = [(peer_rank, error) for peer_rank, error in sorted(normalized) if error] | |
| if failures: | |
| detail = " | ".join(f"rank {peer_rank}: {error}" for peer_rank, error in failures) | |
| raise error_type(f"{label} failed collectively: {detail}") | |
| def collective_call( | |
| call: Callable[[], T], | |
| *, | |
| dist: Any | None, | |
| label: str, | |
| rank: int | None = None, | |
| world_size: int | None = None, | |
| error_type: type[Exception] = BackendContractError, | |
| ) -> T: | |
| """Run rank-local work, then exchange errors before any later collective.""" | |
| result: T | None = None | |
| local_exception: Exception | None = None | |
| local_error = "" | |
| try: | |
| result = call() | |
| except Exception as exc: | |
| local_exception = exc | |
| local_error = f"{type(exc).__name__}: {exc}" | |
| try: | |
| propagate_collective_error( | |
| local_error, | |
| dist=dist, | |
| label=label, | |
| rank=rank, | |
| world_size=world_size, | |
| error_type=error_type, | |
| ) | |
| except Exception as exc: | |
| if local_exception is not None: | |
| raise exc from local_exception | |
| raise | |
| return result # type: ignore[return-value] | |
| def collective_rank_zero_call( | |
| call: Callable[[], T], | |
| *, | |
| dist: Any | None, | |
| label: str, | |
| rank: int | None = None, | |
| world_size: int | None = None, | |
| error_type: type[Exception] = BackendContractError, | |
| ) -> T: | |
| """Run work on raw rank zero and broadcast only its successful value.""" | |
| resolved_rank, resolved_world, initialized = _collective_coordinates( | |
| dist, | |
| rank=rank, | |
| world_size=world_size, | |
| ) | |
| result = collective_call( | |
| call if resolved_rank == 0 else lambda: None, # type: ignore[arg-type] | |
| dist=dist, | |
| label=label, | |
| rank=resolved_rank, | |
| world_size=resolved_world, | |
| error_type=error_type, | |
| ) | |
| values: list[Any] = [result if resolved_rank == 0 else None] | |
| if initialized and resolved_world > 1: | |
| dist.broadcast_object_list(values, src=0) | |
| return values[0] # type: ignore[return-value] | |
| def identity_fingerprint(value: Mapping[str, Any]) -> str: | |
| return hashlib.blake2s(canonical_json(value)).hexdigest() | |
| def assert_peer_fingerprints(values: Sequence[str]) -> str: | |
| if not values or any(not value for value in values): | |
| raise BackendContractError("SP peer fingerprints are missing") | |
| if len(set(values)) != 1: | |
| raise BackendContractError(f"SP peers loaded different sample/RNG identities: {values}") | |
| return values[0] | |
| def gather_and_assert_sp_identity( | |
| value: Mapping[str, Any], *, sp_size: int, group: Any = None | |
| ) -> str: | |
| """Collect identity hashes in the already-created SP group and compare.""" | |
| local = identity_fingerprint(value) | |
| if sp_size == 1: | |
| return local | |
| try: | |
| import torch.distributed as dist | |
| except ImportError as exc: | |
| raise BackendContractError("torch is required for sequence parallelism") from exc | |
| if not dist.is_available() or not dist.is_initialized(): | |
| raise BackendContractError("SP identity check requires initialized torch.distributed") | |
| peers: list[str | None] = [None] * sp_size | |
| dist.all_gather_object(peers, local, group=group) | |
| return assert_peer_fingerprints([str(value or "") for value in peers]) | |