"""Admission and local probes for fixed-membership distributed preflight.""" from __future__ import annotations import ctypes import hashlib import hmac import json import os import platform import re import signal import shutil import stat import threading import time from contextlib import contextmanager from dataclasses import asdict, dataclass from pathlib import Path from typing import Protocol import accelerate # type: ignore[import-untyped] import safetensors import torch import torch.distributed as dist import transformers from obliteratus.distributed.config import DistributedPreflightConfig from obliteratus.distributed.consensus import ( assert_rank_order, gloo_all_gather_records, require_record_consensus, unanimous_vote, ) from obliteratus.distributed.contracts import ( ContractError, RankInventory, RuntimeContractError, RuntimeStage, RunIdentity, StageMessage, TopologyPlan, Vote, advance_stage, canonical_record, contract_digest, ) from obliteratus.distributed.evidence import ( PreflightEvidence, read_evidence, read_stage_message, write_evidence, write_stage_message, ) from obliteratus.distributed.launcher import ( TorchrunEnvironment, control_group, validate_network_interface, ) from obliteratus.checkpoint_errors import CheckpointContractError from obliteratus.checkpoint_inspection import InspectionLimits, inspect_checkpoint _DIGEST_RE = re.compile(r"^[0-9a-f]{64}$") _COMMIT_RE = re.compile(r"^[0-9a-f]{40}$") _ALLOWED_SOURCE_SUFFIXES = frozenset({".json", ".model", ".safetensors", ".tiktoken", ".txt"}) _FORBIDDEN_SOURCE_SUFFIXES = frozenset({".bin", ".ckpt", ".pkl", ".pickle", ".pt", ".pth", ".py"}) HASH_CHUNK_BYTES = 8 * 1024 * 1024 _ACTIVE_DEADLINE_DEPTH = 0 def _digest(value: object, name: str, pattern: re.Pattern[str] = _DIGEST_RE) -> None: if not isinstance(value, str) or pattern.fullmatch(value) is None: raise ContractError(f"{name} has an invalid format") def _memory(value: object, name: str, *, allow_zero: bool = False) -> None: minimum = 0 if allow_zero else 1 if isinstance(value, bool) or not isinstance(value, int) or not minimum <= value <= 2**63 - 1: raise ContractError(f"{name} is outside its allowed range") @dataclass(frozen=True) class RankAttestation: """Redaction-safe, exact local facts submitted by one fixed worker.""" rank: int local_rank: int local_world_size: int group_rank: int world_size: int host_digest: str device_digest: str device_profile_digest: str device_config_digest: str device_kind: str total_device_memory_bytes: int free_device_memory_bytes: int total_host_memory_bytes: int free_host_memory_bytes: int free_staging_bytes: int software_digest: str storage_digest: str source_digest: str model_digest: str tokenizer_digest: str config_digest: str commit_sha: str code_digest: str placement_plan_digest: str network_interface_digest: str def __post_init__(self) -> None: RankInventory( rank=self.rank, local_rank=self.local_rank, world_size=self.world_size, host_digest=self.host_digest, device_digest=self.device_digest, device_kind=self.device_kind, total_memory_bytes=self.total_device_memory_bytes, free_memory_bytes=self.free_device_memory_bytes, software_digest=self.software_digest, storage_digest=self.storage_digest, ) _memory(self.total_host_memory_bytes, "total_host_memory_bytes") _memory(self.free_host_memory_bytes, "free_host_memory_bytes", allow_zero=True) if self.free_host_memory_bytes > self.total_host_memory_bytes: raise ContractError("free_host_memory_bytes cannot exceed total_host_memory_bytes") _memory(self.free_staging_bytes, "free_staging_bytes", allow_zero=True) for name in ( "device_profile_digest", "device_config_digest", "source_digest", "model_digest", "tokenizer_digest", "config_digest", "code_digest", "placement_plan_digest", "network_interface_digest", ): _digest(getattr(self, name), name) _digest(self.commit_sha, "commit_sha", _COMMIT_RE) def to_dict(self) -> dict[str, object]: return asdict(self) @classmethod def from_bytes(cls, payload: bytes) -> "RankAttestation": def unique_object(pairs: list[tuple[str, object]]) -> dict[str, object]: result: dict[str, object] = {} for key, value in pairs: if key in result: raise ContractError("rank attestation contains a duplicate field") result[key] = value return result try: value = json.loads(payload.decode("utf-8"), object_pairs_hook=unique_object) except (UnicodeError, json.JSONDecodeError) as exc: raise ContractError("rank attestation must be strict UTF-8 JSON") from exc if not isinstance(value, dict) or set(value) != set(cls.__dataclass_fields__): raise ContractError("rank attestation fields do not match the closed schema") result = cls(**value) if payload != canonical_record(result): raise ContractError("rank attestation is not canonical") return result def _stage_message_from_bytes(payload: bytes) -> StageMessage: try: return StageMessage.from_bytes(payload) except RuntimeContractError: raise except ContractError as exc: raise RuntimeContractError( "LMS_CONSENSUS_RECORD_INVALID", "lifecycle message is invalid" ) from exc def _validate_stage_messages( messages: tuple[StageMessage, ...], *, run_id: str, world_size: int, stage: RuntimeStage, sequence: int, identity_digest: str, vote: Vote | None, ) -> None: if len(messages) != world_size or {item.rank for item in messages} != set(range(world_size)): raise RuntimeContractError( "LMS_MEMBERSHIP_INVALID", "lifecycle records do not cover the fixed world" ) if any( item.run_id != run_id or item.stage is not stage or item.sequence != sequence or item.identity_digest != identity_digest or item.vote is not vote or item.error_code is not None for item in messages ): raise RuntimeContractError("LMS_LIFECYCLE_INVALID", "rank lifecycle records disagree") @dataclass(frozen=True) class SourceIdentity: source_digest: str model_digest: str tokenizer_digest: str file_count: int total_bytes: int @dataclass(frozen=True) class LocalSnapshot: host_identity: str device_identity: str device_name: str compute_capability: str device_kind: str total_device_memory_bytes: int free_device_memory_bytes: int total_host_memory_bytes: int free_host_memory_bytes: int free_staging_bytes: int storage_identity: str source: SourceIdentity software_versions: tuple[tuple[str, str], ...] commit_sha: str code_digest: str @dataclass(frozen=True) class PreflightResult: """Capability that later distributed code must require before allocation.""" identity: RunIdentity topology: TopologyPlan attestations: tuple[RankAttestation, ...] identity_digest: str attempt_path: Path source_identity: SourceIdentity lifecycle_identity_digest: str created_messages: tuple[StageMessage, ...] class PreflightProbes(Protocol): def collect( self, config: DistributedPreflightConfig, launch: TorchrunEnvironment ) -> LocalSnapshot: ... def _hash_file( path: Path, *, deadline: float, max_bytes: int, require_read_only: bool = True, ) -> tuple[int, str]: descriptor: int | None = None try: descriptor = os.open(path, os.O_RDONLY | os.O_NOFOLLOW) before = os.fstat(descriptor) if not stat.S_ISREG(before.st_mode) or (require_read_only and before.st_mode & 0o222): raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "source files must be immutable regular files", ) if before.st_nlink != 1: raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "source files cannot have hard-link aliases", ) if before.st_size > max_bytes: raise RuntimeContractError( "LMS_RESOURCE_ADMISSION_DENIED", "source file exceeds the configured byte bound", ) digest = hashlib.sha256() while chunk := os.read(descriptor, HASH_CHUNK_BYTES): if time.monotonic() > deadline: raise RuntimeContractError( "LMS_STAGE_TIMEOUT", "source inspection exceeded its explicit timeout", ) digest.update(chunk) after = os.fstat(descriptor) final_path = path.stat(follow_symlinks=False) except RuntimeContractError: raise except OSError as exc: raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "source files must be immutable regular files", ) from exc finally: if descriptor is not None: os.close(descriptor) identity = lambda item: ( item.st_dev, item.st_ino, item.st_size, item.st_mtime_ns, item.st_mode, ) if identity(before) != identity(after) or identity(after) != identity(final_path): raise RuntimeContractError( "LMS_SOURCE_CHANGED", "source changed during preflight inspection" ) return before.st_size, digest.hexdigest() @contextmanager def _stage_deadline(timeout_seconds: int, message: str): """Enforce a nestable stage bound with a process-level Linux timer.""" global _ACTIVE_DEADLINE_DEPTH if ( not hasattr(signal, "SIGALRM") or not hasattr(signal, "setitimer") or threading.current_thread() is not threading.main_thread() ): raise RuntimeContractError( "LMS_RUNTIME_PROFILE_MISMATCH", "the runtime cannot enforce the source inspection deadline", ) previous_timer = signal.getitimer(signal.ITIMER_REAL) if previous_timer[0] > 0 and _ACTIVE_DEADLINE_DEPTH == 0: raise RuntimeContractError( "LMS_RUNTIME_PROFILE_MISMATCH", "the stage deadline conflicts with an active runtime timer", ) previous_handler = signal.getsignal(signal.SIGALRM) started = time.monotonic() def expire(_signum: int, _frame: object) -> None: raise RuntimeContractError("LMS_STAGE_TIMEOUT", message) signal.signal(signal.SIGALRM, expire) duration = ( min(float(timeout_seconds), previous_timer[0]) if previous_timer[0] > 0 else float(timeout_seconds) ) signal.setitimer(signal.ITIMER_REAL, duration) _ACTIVE_DEADLINE_DEPTH += 1 try: yield finally: signal.setitimer(signal.ITIMER_REAL, 0) _ACTIVE_DEADLINE_DEPTH -= 1 signal.signal(signal.SIGALRM, previous_handler) if previous_timer[0] > 0: remaining = previous_timer[0] - (time.monotonic() - started) if remaining > 0: signal.setitimer(signal.ITIMER_REAL, remaining, previous_timer[1]) def inspect_source( path: Path, *, max_files: int = 100_000, max_total_bytes: int = 2**44, max_file_bytes: int = 2**41, timeout_seconds: int = 3600, ) -> SourceIdentity: """Hash an immutable local HF safetensors tree without loading tensor payloads.""" if ( isinstance(timeout_seconds, bool) or not isinstance(timeout_seconds, int) or not 1 <= timeout_seconds <= 3600 ): raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "timeout_seconds is outside the bounded source-inspection contract", ) with _stage_deadline(timeout_seconds, "source inspection exceeded its explicit timeout"): try: return _inspect_source( path, max_files=max_files, max_total_bytes=max_total_bytes, max_file_bytes=max_file_bytes, timeout_seconds=timeout_seconds, ) except RuntimeContractError: raise except ContractError as exc: raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "source violates the bounded immutable input contract", ) from exc def _inspect_source( path: Path, *, max_files: int, max_total_bytes: int, max_file_bytes: int, timeout_seconds: int, ) -> SourceIdentity: if path.is_symlink() or not path.is_dir(): raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "source must be a local non-symlink directory", ) if path.stat(follow_symlinks=False).st_mode & 0o222: raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "source directory must be immutable" ) for value, name, maximum in ( (max_files, "max_files", 1_000_000), (max_total_bytes, "max_total_bytes", 2**63 - 1), (max_file_bytes, "max_file_bytes", 2**63 - 1), (timeout_seconds, "timeout_seconds", 3600), ): if isinstance(value, bool) or not isinstance(value, int) or not 1 <= value <= maximum: raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", f"{name} is outside the bounded source-inspection contract", ) if max_file_bytes > max_total_bytes: raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "max_file_bytes cannot exceed max_total_bytes", ) deadline = time.monotonic() + timeout_seconds files: list[Path] = [] observed_bytes = 0 observed_inodes: set[tuple[int, int]] = set() for directory, directory_names, file_names in os.walk(path, topdown=True, followlinks=False): if time.monotonic() > deadline: raise RuntimeContractError( "LMS_STAGE_TIMEOUT", "source inspection exceeded its explicit timeout" ) directory_names.sort() file_names.sort() directory_path = Path(directory) directory_metadata = directory_path.stat(follow_symlinks=False) if not stat.S_ISDIR(directory_metadata.st_mode) or directory_metadata.st_mode & 0o222: raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "every source directory must be immutable", ) for name in (*directory_names, *file_names): item = directory_path / name metadata = item.stat(follow_symlinks=False) if stat.S_ISLNK(metadata.st_mode): raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "source cannot contain symbolic links", ) for name in file_names: item = directory_path / name metadata = item.stat(follow_symlinks=False) if not stat.S_ISREG(metadata.st_mode): raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "source can contain only regular files and directories", ) if metadata.st_nlink != 1: raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "source files cannot have hard-link aliases", ) files.append(item) if len(files) > max_files: raise RuntimeContractError( "LMS_RESOURCE_ADMISSION_DENIED", "source file count exceeds the configured bound", ) if metadata.st_size > max_file_bytes: raise RuntimeContractError( "LMS_RESOURCE_ADMISSION_DENIED", "source file exceeds the configured byte bound", ) observed_bytes += metadata.st_size if observed_bytes > max_total_bytes: raise RuntimeContractError( "LMS_RESOURCE_ADMISSION_DENIED", "source exceeds the configured total-byte bound", ) inode = (metadata.st_dev, metadata.st_ino) if inode in observed_inodes: raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "source files cannot alias the same inode", ) observed_inodes.add(inode) if not files: raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "source file count is outside the bounded contract", ) for item in files: suffix = item.suffix.lower() if suffix in _FORBIDDEN_SOURCE_SUFFIXES or suffix not in _ALLOWED_SOURCE_SUFFIXES: raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "source contains a serialization type outside the safetensors envelope", ) try: report = inspect_checkpoint( path, limits=InspectionLimits( max_files=max_files, max_directories=min(1_000_000, max_files + 1), max_total_bytes=max_total_bytes, max_json_bytes=min(max_file_bytes, 8 << 20), max_safetensors_header_bytes=min(max_file_bytes, 64 << 20), max_tensors=2_000_000, hash_chunk_bytes=HASH_CHUNK_BYTES, ), ) except CheckpointContractError as exc: if exc.code == "DCI_SOURCE_CHANGED": raise RuntimeContractError( "LMS_SOURCE_CHANGED", "source changed during safe-structure inspection" ) from exc if exc.code == "DCI_RESOURCE_LIMIT" and exc.detail == "deadline": raise RuntimeContractError( "LMS_STAGE_TIMEOUT", "source inspection exceeded its explicit timeout" ) from exc raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "source failed bounded safe-structure inspection", ) from exc if time.monotonic() > deadline: raise RuntimeContractError( "LMS_STAGE_TIMEOUT", "source inspection exceeded its explicit timeout" ) descriptor = report.to_dict() component_formats = {component["format"] for component in descriptor["components"]} if ( report.primary_format != "hf_safetensors" or report.support_decision != "canonical_hf_ready" or component_formats != {"hf_safetensors"} ): raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "source must be a canonical Hugging Face safetensors checkpoint without adapters", ) inventory = descriptor["source_inventory"] inventory_files = inventory["files"] if len(inventory_files) != len(files) or inventory["total_bytes"] != observed_bytes: raise RuntimeContractError( "LMS_SOURCE_CHANGED", "source changed during safe-structure inspection" ) records: list[dict[str, object]] = [] model_records: list[dict[str, object]] = [] tokenizer_records: list[dict[str, object]] = [] for observed in inventory_files: relative_path = str(observed["relative_path"]) record = { "path": relative_path, "size_bytes": observed["size_bytes"], "sha256": observed["sha256"], } records.append(record) if relative_path.lower().endswith(".safetensors"): model_records.append(record) name = Path(relative_path).name.lower() if any(marker in name for marker in ("token", "vocab", "merges")): tokenizer_records.append(record) if not model_records: raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "source must contain at least one safetensors file", ) if not tokenizer_records: raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "source must contain immutable tokenizer metadata", ) return SourceIdentity( source_digest=contract_digest(records), model_digest=contract_digest(model_records), tokenizer_digest=contract_digest(tokenizer_records), file_count=len(records), total_bytes=observed_bytes, ) def revalidate_source( config: DistributedPreflightConfig, expected: SourceIdentity ) -> SourceIdentity: """Re-hash immediately before #60 allocation; replacement fails capability use.""" observed = inspect_source( config.source_path, max_files=config.max_source_files, max_total_bytes=config.max_source_bytes, max_file_bytes=config.max_source_file_bytes, timeout_seconds=config.source_timeout_seconds, ) if observed != expected: raise RuntimeContractError( "LMS_SOURCE_CHANGED", "source identity changed after distributed preflight", ) return observed def _host_memory() -> tuple[int, int]: page_size = os.sysconf("SC_PAGE_SIZE") total = page_size * os.sysconf("SC_PHYS_PAGES") available = page_size * os.sysconf("SC_AVPHYS_PAGES") return total, available def _cuda_driver_version() -> str: try: library = ctypes.CDLL("libcuda.so.1") if library.cuInit(0) != 0: return "unavailable" value = ctypes.c_int() if library.cuDriverGetVersion(ctypes.byref(value)) != 0: return "unavailable" return str(value.value) except (AttributeError, OSError): return "unavailable" def _nccl_version() -> str: try: value = torch.cuda.nccl.version() except Exception: return "unavailable" if isinstance(value, tuple): return ".".join(str(item) for item in value) return str(value) def _bounded_text_file(path: Path, *, limit: int = 4096) -> str: descriptor: int | None = None try: descriptor = os.open(path, os.O_RDONLY | os.O_NOFOLLOW) metadata = os.fstat(descriptor) if not stat.S_ISREG(metadata.st_mode) or not 1 <= metadata.st_size <= limit: raise ContractError("checkout identity file is invalid") raw = os.read(descriptor, limit + 1) if len(raw) != metadata.st_size: raise ContractError("checkout identity changed while it was read") return raw.decode("utf-8").strip() except ContractError: raise except (OSError, UnicodeError) as exc: raise ContractError("checkout identity is unavailable") from exc finally: if descriptor is not None: os.close(descriptor) def checkout_commit(root: Path | None = None) -> str: """Resolve the checkout HEAD without invoking Git or another process.""" checkout = root or Path(__file__).resolve().parents[2] dot_git = checkout / ".git" if dot_git.is_dir() and not dot_git.is_symlink(): git_dir = dot_git else: pointer = _bounded_text_file(dot_git) prefix = "gitdir: " if not pointer.startswith(prefix): raise ContractError("checkout identity is unavailable") git_dir = Path(pointer[len(prefix) :]) if not git_dir.is_absolute(): git_dir = (checkout / git_dir).resolve() head = _bounded_text_file(git_dir / "HEAD") if _COMMIT_RE.fullmatch(head): return head prefix = "ref: " if not head.startswith(prefix): raise ContractError("checkout HEAD has an invalid format") reference = head[len(prefix) :] if not reference.startswith("refs/") or ".." in reference or "\\" in reference: raise ContractError("checkout HEAD reference has an invalid format") common_dir = git_dir common_pointer = git_dir / "commondir" if common_pointer.exists(): common_dir = (git_dir / _bounded_text_file(common_pointer)).resolve() reference_path = common_dir / reference if reference_path.exists(): commit = _bounded_text_file(reference_path) else: packed = _bounded_text_file(common_dir / "packed-refs", limit=1 << 20) matches = [ line.split(" ", 1)[0] for line in packed.splitlines() if not line.startswith(("#", "^")) and line.endswith(f" {reference}") ] if len(matches) != 1: raise ContractError("checkout HEAD reference is unavailable") commit = matches[0] if _COMMIT_RE.fullmatch(commit) is None: raise ContractError("checkout commit has an invalid format") return commit def checkout_code_digest(root: Path | None = None) -> str: """Hash the executable Python package independently of Git metadata.""" checkout = root or Path(__file__).resolve().parents[2] package = checkout / "obliteratus" records: list[dict[str, object]] = [] total_bytes = 0 for path in sorted(package.rglob("*.py")): if len(records) >= 10_000: raise RuntimeContractError( "LMS_RESOURCE_ADMISSION_DENIED", "executable source exceeds its file bound" ) if path.is_symlink() or not path.is_file(): raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "executable source must contain only regular Python files", ) size, digest = _hash_file( path, deadline=time.monotonic() + 60, max_bytes=8 << 20, require_read_only=False, ) total_bytes += size if total_bytes > 64 << 20: raise RuntimeContractError( "LMS_RESOURCE_ADMISSION_DENIED", "executable source exceeds its byte bound" ) records.append( { "path": path.relative_to(checkout).as_posix(), "size_bytes": size, "sha256": digest, } ) if not records: raise RuntimeContractError( "LMS_SOURCE_BOUNDARY_VIOLATION", "executable source inventory is empty" ) return contract_digest(records) def storage_mount_digest(path: Path) -> str: """Measure the exact Linux mount backing a local staging directory.""" try: resolved = path.resolve(strict=True) if resolved != path.absolute() or path.is_symlink() or not path.is_dir(): raise RuntimeContractError( "LMS_STORAGE_PROFILE_MISMATCH", "staging must be a local non-symlink directory", ) descriptor = os.open("/proc/self/mountinfo", os.O_RDONLY | os.O_NOFOLLOW) try: raw = os.read(descriptor, (4 << 20) + 1) finally: os.close(descriptor) if len(raw) > 4 << 20: raise RuntimeContractError( "LMS_RESOURCE_ADMISSION_DENIED", "mount inventory exceeds its byte bound" ) lines = raw.decode("utf-8").splitlines() except RuntimeContractError: raise except (OSError, UnicodeError) as exc: raise RuntimeContractError( "LMS_STORAGE_PROFILE_MISMATCH", "staging mount identity is unavailable" ) from exc def unescape(value: str) -> str: return ( value.replace("\\040", " ") .replace("\\011", "\t") .replace("\\012", "\n") .replace("\\134", "\\") ) candidates: list[tuple[int, dict[str, object]]] = [] for line in lines: if " - " not in line: continue left, right = line.split(" - ", 1) left_fields = left.split() right_fields = right.split() if len(left_fields) < 6 or len(right_fields) < 3: continue mountpoint = Path(unescape(left_fields[4])) try: resolved.relative_to(mountpoint) except ValueError: continue candidates.append( ( len(mountpoint.parts), { "root": unescape(left_fields[3]), "mountpoint": mountpoint.as_posix(), "mount_options": sorted(left_fields[5].split(",")), "filesystem": right_fields[0], "source": unescape(right_fields[1]), "super_options": sorted(right_fields[2].split(",")), }, ) ) if not candidates: raise RuntimeContractError( "LMS_STORAGE_PROFILE_MISMATCH", "staging mount identity is unavailable" ) return contract_digest(max(candidates, key=lambda item: item[0])[1]) class SystemProbes: """Production local probes; tests inject deterministic snapshots instead.""" def collect( self, config: DistributedPreflightConfig, launch: TorchrunEnvironment ) -> LocalSnapshot: source = inspect_source( config.source_path, max_files=config.max_source_files, max_total_bytes=config.max_source_bytes, max_file_bytes=config.max_source_file_bytes, timeout_seconds=config.source_timeout_seconds, ) if config.staging_path.is_symlink() or not config.staging_path.is_dir(): raise RuntimeContractError( "LMS_STORAGE_PROFILE_MISMATCH", "staging must be a local non-symlink directory", ) disk = shutil.disk_usage(config.staging_path) validate_network_interface( config.network_interface, config.allowed_master_cidrs, config.master_addr, coordinator=launch.rank == config.coordinator_rank, ) host_total, host_free = _host_memory() software = tuple( sorted( { "python": platform.python_version(), "platform": platform.platform(), "machine": platform.machine(), "torch": torch.__version__.split("+")[0], "transformers": transformers.__version__, "accelerate": accelerate.__version__, "safetensors": safetensors.__version__, "cuda": torch.version.cuda or "unavailable", "nccl": _nccl_version(), "driver": _cuda_driver_version(), }.items() ) ) if config.device_kind == "cuda": if not torch.cuda.is_available(): raise RuntimeContractError( "LMS_RUNTIME_PROFILE_MISMATCH", "the configured CUDA device is unavailable", ) torch.cuda.set_device(launch.local_rank) properties = torch.cuda.get_device_properties(launch.local_rank) device_uuid = getattr(properties, "uuid", None) if not device_uuid: raise RuntimeContractError( "LMS_RUNTIME_PROFILE_MISMATCH", "the CUDA device UUID is unavailable", ) free_device, total_device = torch.cuda.mem_get_info(launch.local_rank) device_identity = str(device_uuid) device_name = properties.name compute_capability = f"{properties.major}.{properties.minor}" else: total_device, free_device = host_total, host_free device_identity = f"cpu:{platform.node()}:{platform.machine()}:{launch.local_rank}" device_name = "cpu" compute_capability = "none" try: commit_sha = checkout_commit() code_digest = checkout_code_digest() except RuntimeContractError: raise except ContractError as exc: raise RuntimeContractError( "LMS_RUNTIME_PROFILE_MISMATCH", "checkout identity cannot be verified", ) from exc return LocalSnapshot( host_identity=platform.node(), device_identity=device_identity, device_name=device_name, compute_capability=compute_capability, device_kind=config.device_kind, total_device_memory_bytes=total_device, free_device_memory_bytes=free_device, total_host_memory_bytes=host_total, free_host_memory_bytes=host_free, free_staging_bytes=disk.free, storage_identity=storage_mount_digest(config.staging_path), source=source, software_versions=software, commit_sha=commit_sha, code_digest=code_digest, ) def _opaque_identity(value: str, rendezvous_id: str) -> str: return hmac.new(bytes.fromhex(rendezvous_id), value.encode(), hashlib.sha256).hexdigest() def _attestation( config: DistributedPreflightConfig, launch: TorchrunEnvironment, snapshot: LocalSnapshot, ) -> RankAttestation: source = snapshot.source return RankAttestation( rank=launch.rank, local_rank=launch.local_rank, local_world_size=launch.local_world_size, group_rank=launch.group_rank, world_size=launch.world_size, host_digest=_opaque_identity(snapshot.host_identity, launch.rendezvous_id), device_digest=_opaque_identity(snapshot.device_identity, launch.rendezvous_id), device_profile_digest=contract_digest( { "name": snapshot.device_name, "compute_capability": snapshot.compute_capability, "total_memory_bytes": snapshot.total_device_memory_bytes, } ), device_config_digest=contract_digest( { "kind": snapshot.device_kind, "name": snapshot.device_name, "compute_capability": snapshot.compute_capability, } ), device_kind=snapshot.device_kind, total_device_memory_bytes=snapshot.total_device_memory_bytes, free_device_memory_bytes=snapshot.free_device_memory_bytes, total_host_memory_bytes=snapshot.total_host_memory_bytes, free_host_memory_bytes=snapshot.free_host_memory_bytes, free_staging_bytes=snapshot.free_staging_bytes, software_digest=contract_digest(dict(snapshot.software_versions)), storage_digest=snapshot.storage_identity, source_digest=source.source_digest, model_digest=source.model_digest, tokenizer_digest=source.tokenizer_digest, config_digest=config.digest, commit_sha=snapshot.commit_sha, code_digest=snapshot.code_digest, placement_plan_digest=config.placement_plan_digest, network_interface_digest=_opaque_identity(config.network_interface, launch.rendezvous_id), ) def _prepare_attempt_directory( config: DistributedPreflightConfig, launch: TorchrunEnvironment ) -> Path: attempt = config.staging_path / config.run_id staging_descriptor: int | None = None attempt_descriptor: int | None = None try: staging_descriptor = os.open( config.staging_path, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW ) if launch.rank == config.coordinator_rank: try: os.mkdir(config.run_id, mode=0o700, dir_fd=staging_descriptor) except FileExistsError as exc: raise RuntimeContractError( "LMS_ATTEMPT_IDENTITY_CONFLICT", "distributed staging attempt cannot be reused", ) from exc dist.barrier() attempt_descriptor = os.open( config.run_id, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=staging_descriptor, ) metadata = os.fstat(attempt_descriptor) effective_uid = getattr(os, "geteuid", lambda: metadata.st_uid)() if metadata.st_uid != effective_uid or metadata.st_mode & 0o077: raise RuntimeContractError( "LMS_STORAGE_PROFILE_MISMATCH", "distributed staging attempt is not private and worker-owned", ) dist.barrier() return attempt except RuntimeContractError: raise except OSError as exc: raise RuntimeContractError( "LMS_STORAGE_PROFILE_MISMATCH", "distributed staging attempt is unavailable", ) from exc finally: if attempt_descriptor is not None: os.close(attempt_descriptor) if staging_descriptor is not None: os.close(staging_descriptor) def _probe_shared_staging(config: DistributedPreflightConfig, launch: TorchrunEnvironment) -> Path: attempt = config.staging_path / config.run_id pending = ".shared-probe.pending" ready = ".shared-probe.ready" nonce = contract_digest( {"run_id": config.run_id, "rendezvous_id": config.rendezvous_id} ).encode() staging_descriptor: int | None = None attempt_descriptor: int | None = None try: staging_descriptor = os.open( config.staging_path, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW ) try: attempt_descriptor = os.open( config.run_id, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=staging_descriptor, ) except OSError as exc: raise RuntimeContractError( "LMS_STORAGE_PROFILE_MISMATCH", "shared staging attempt is not visible to every rank", ) from exc if launch.rank == config.coordinator_rank: descriptor = os.open( pending, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600, dir_fd=attempt_descriptor, ) try: if os.write(descriptor, nonce) != len(nonce): raise RuntimeContractError( "LMS_STORAGE_PROFILE_MISMATCH", "shared staging probe write was incomplete", ) os.fsync(descriptor) finally: os.close(descriptor) os.rename( pending, ready, src_dir_fd=attempt_descriptor, dst_dir_fd=attempt_descriptor, ) os.fsync(attempt_descriptor) dist.barrier() descriptor = os.open(ready, os.O_RDONLY | os.O_NOFOLLOW, dir_fd=attempt_descriptor) try: observed = os.read(descriptor, len(nonce) + 1) finally: os.close(descriptor) if observed != nonce: raise RuntimeContractError( "LMS_STORAGE_PROFILE_MISMATCH", "shared staging probe identity disagrees", ) dist.barrier() if launch.rank == config.coordinator_rank: os.unlink(ready, dir_fd=attempt_descriptor) os.fsync(attempt_descriptor) dist.barrier() return attempt except ContractError: raise except OSError as exc: raise RuntimeContractError( "LMS_STORAGE_PROFILE_MISMATCH", "shared staging probe failed" ) from exc finally: cleanup_failed = False if attempt_descriptor is not None and launch.rank == config.coordinator_rank: for marker in (pending, ready): try: os.unlink(marker, dir_fd=attempt_descriptor) except FileNotFoundError: pass except OSError: cleanup_failed = True try: os.fsync(attempt_descriptor) except OSError: cleanup_failed = True if attempt_descriptor is not None: os.close(attempt_descriptor) if staging_descriptor is not None: os.close(staging_descriptor) if cleanup_failed: raise RuntimeContractError("LMS_CLEANUP_INCOMPLETE", "shared staging cleanup failed") def _run_identity(config: DistributedPreflightConfig) -> RunIdentity: return RunIdentity( run_id=config.run_id, config_digest=config.digest, source_digest=config.source_digest, model_digest=config.model_digest, tokenizer_digest=config.tokenizer_digest, commit_sha=config.commit_sha, world_size=config.world_size, ) def _begin_lifecycle( config: DistributedPreflightConfig, launch: TorchrunEnvironment, ) -> tuple[RunIdentity, str, tuple[StageMessage, ...]]: identity = _run_identity(config) try: require_record_consensus(identity) except RuntimeContractError: raise except ContractError as exc: raise RuntimeContractError("LMS_IDENTITY_MISMATCH", "rank run identities disagree") from exc lifecycle_identity_digest = contract_digest(identity) created = StageMessage( run_id=config.run_id, identity_digest=lifecycle_identity_digest, rank=launch.rank, sequence=0, stage=RuntimeStage.CREATED, ) created_messages = tuple( _stage_message_from_bytes(item) for item in gloo_all_gather_records(created) ) _validate_stage_messages( created_messages, run_id=config.run_id, world_size=config.world_size, stage=RuntimeStage.CREATED, sequence=0, identity_digest=lifecycle_identity_digest, vote=None, ) return identity, lifecycle_identity_digest, created_messages def run_preflight( config: DistributedPreflightConfig, launch: TorchrunEnvironment, *, probes: PreflightProbes | None = None, lifecycle: tuple[RunIdentity, str, tuple[StageMessage, ...]] | None = None, ) -> PreflightResult: """Admit the complete worker group; this function has no model allocation path.""" try: config.validate() except RuntimeContractError: raise except ContractError as exc: raise RuntimeContractError( "LMS_RUNTIME_PROFILE_MISMATCH", "distributed runtime profile validation failed", ) from exc identity, lifecycle_identity_digest, created_messages = lifecycle or _begin_lifecycle( config, launch ) attempt = _prepare_attempt_directory(config, launch) with _stage_deadline( min(config.source_timeout_seconds, config.collective_timeout_seconds), "local preflight probes exceeded their explicit timeout", ): snapshot = (probes or SystemProbes()).collect(config, launch) local = _attestation(config, launch, snapshot) try: gathered = gloo_all_gather_records(local) attestations = tuple(RankAttestation.from_bytes(item) for item in gathered) except RuntimeContractError: raise except ContractError as exc: raise RuntimeContractError( "LMS_CONSENSUS_RECORD_INVALID", "rank attestation record is invalid" ) from exc validate_attestations( attestations, world_size=config.world_size, tensor_parallel_size=config.tensor_parallel_size, dimension_divisors=config.dimension_divisors, expected_device_kind=config.device_kind, expected_device_config_digest=contract_digest( { "kind": config.device_kind, "name": config.device_name, "compute_capability": config.compute_capability, } ), expected_software_digest=config.software_digest, expected_source_digest=config.source_digest, expected_model_digest=config.model_digest, expected_tokenizer_digest=config.tokenizer_digest, expected_config_digest=config.digest, expected_commit_sha=config.commit_sha, expected_code_digest=config.code_digest, expected_placement_plan_digest=config.placement_plan_digest, expected_storage_digest=config.storage_digest, expected_network_interface_digest=_opaque_identity( config.network_interface, launch.rendezvous_id ), local_world_size=config.local_world_size, min_free_device_memory_bytes=config.min_free_device_memory_bytes, min_free_host_memory_bytes=config.min_free_host_memory_bytes, min_free_staging_bytes=config.min_free_staging_bytes, ) if _probe_shared_staging(config, launch) != attempt: raise RuntimeContractError( "LMS_STORAGE_PROFILE_MISMATCH", "staging attempt identity changed" ) if not unanimous_vote(0, True): raise ContractError("preflight admission was not unanimous") topology = TopologyPlan( world_size=config.world_size, coordinator_rank=config.coordinator_rank, backend="gloo", placement_plan_digest=config.placement_plan_digest, ) identity_digest = contract_digest( { "identity": identity, "topology": topology, "attestations": attestations, "created_messages": created_messages, } ) return PreflightResult( identity, topology, attestations, identity_digest, attempt, snapshot.source, lifecycle_identity_digest, created_messages, ) def _failure_code(error: BaseException) -> str: if isinstance(error, RuntimeContractError): return error.code if isinstance(error, KeyboardInterrupt): return "LMS_ATTEMPT_CANCELLED" if isinstance(error, TimeoutError): return "LMS_STAGE_TIMEOUT" if isinstance(error, ContractError): return "LMS_LIFECYCLE_INVALID" return "LMS_COLLECTIVE_FAILED" _FAILURE_PRIORITY = ( "LMS_CLEANUP_INCOMPLETE", "LMS_DIAGNOSTIC_REDACTION_FAILED", "LMS_SECURITY_BASELINE_REVOKED", "LMS_SECRET_INPUT_REJECTED", "LMS_SOURCE_CHANGED", "LMS_SOURCE_BOUNDARY_VIOLATION", "LMS_IDENTITY_MISMATCH", "LMS_STORAGE_PROFILE_MISMATCH", "LMS_NETWORK_PROFILE_DENIED", "LMS_RUNTIME_PROFILE_MISMATCH", "LMS_RESOURCE_ADMISSION_DENIED", "LMS_ATTEMPT_CANCELLED", "LMS_ATTEMPT_IDENTITY_CONFLICT", "LMS_MEMBERSHIP_INVALID", "LMS_MEMBERSHIP_TIMEOUT", "LMS_ELASTICITY_FORBIDDEN", "LMS_LAUNCH_IDENTITY_INVALID", "LMS_RANK_DEVICE_CONFLICT", "LMS_TOPOLOGY_UNSUPPORTED", "LMS_TRANSPORT_BOUNDARY_UNSATISFIED", "LMS_FORBIDDEN_RUNTIME_CAPABILITY", "LMS_ATOMIC_PROMOTION_UNAVAILABLE", "LMS_STAGE_TIMEOUT", "LMS_CONSENSUS_RECORD_INVALID", "LMS_LIFECYCLE_INVALID", "LMS_EVIDENCE_UNAVAILABLE", "LMS_EVIDENCE_SCOPE_MISMATCH", "LMS_CLOCK_PROFILE_MISMATCH", "LMS_DISTRIBUTED_INTENT_REQUIRED", "LMS_COLLECTIVE_FAILED", ) def _select_failure_code(codes: set[str]) -> str: for code in _FAILURE_PRIORITY: if code in codes: return code return "LMS_COLLECTIVE_FAILED" def _collective_aborting( config: DistributedPreflightConfig, launch: TorchrunEnvironment, identity_digest: str, local_code: str, ) -> tuple[str, tuple[StageMessage, ...]]: advance_stage(RuntimeStage.CREATED, RuntimeStage.ABORTING) local = StageMessage( run_id=config.run_id, identity_digest=identity_digest, rank=launch.rank, sequence=1, stage=RuntimeStage.ABORTING, vote=Vote.ABORT, error_code=local_code, ) try: messages = tuple(_stage_message_from_bytes(item) for item in gloo_all_gather_records(local)) except RuntimeContractError: raise except (ContractError, RuntimeError) as exc: raise RuntimeContractError( "LMS_LIFECYCLE_INVALID", "collective abort lifecycle could not be established", ) from exc if len(messages) != config.world_size or {item.rank for item in messages} != set( range(config.world_size) ): raise RuntimeContractError( "LMS_MEMBERSHIP_INVALID", "abort lifecycle records do not cover the fixed world", ) if any( item.run_id != config.run_id or item.identity_digest != identity_digest or item.sequence != 1 or item.stage is not RuntimeStage.ABORTING or item.vote is not Vote.ABORT or item.error_code is None for item in messages ): raise RuntimeContractError("LMS_LIFECYCLE_INVALID", "rank abort lifecycle records disagree") return _select_failure_code({item.error_code for item in messages if item.error_code}), messages def _write_local_terminal_receipts( attempt: Path, *, config: DistributedPreflightConfig, launch: TorchrunEnvironment, identity_digest: str, code: str, quarantined: bool, ) -> None: aborting = StageMessage( run_id=config.run_id, identity_digest=identity_digest, rank=launch.rank, sequence=1, stage=RuntimeStage.ABORTING, vote=Vote.ABORT, error_code=code, ) terminal_stage = RuntimeStage.QUARANTINED if quarantined else RuntimeStage.ABORTED advance_stage(RuntimeStage.ABORTING, terminal_stage) terminal = StageMessage( run_id=config.run_id, identity_digest=identity_digest, rank=launch.rank, sequence=2, stage=terminal_stage, vote=Vote.ABORT, error_code=code, ) write_stage_message(attempt / f".rank-{launch.rank}.aborting.stage.json", aborting) write_stage_message(attempt / f".rank-{launch.rank}.terminal.stage.json", terminal) def _read_terminal_receipts( attempt: Path, *, config: DistributedPreflightConfig, identity_digest: str, ) -> tuple[tuple[StageMessage, ...], tuple[StageMessage, ...]]: aborting = tuple( read_stage_message(attempt / f".rank-{rank}.aborting.stage.json") for rank in range(config.world_size) ) terminal = tuple( read_stage_message(attempt / f".rank-{rank}.terminal.stage.json") for rank in range(config.world_size) ) expected_ranks = list(range(config.world_size)) if [item.rank for item in aborting] != expected_ranks or [ item.rank for item in terminal ] != expected_ranks: raise RuntimeContractError( "LMS_MEMBERSHIP_INVALID", "terminal lifecycle receipts omit a fixed rank" ) for before, after in zip(aborting, terminal, strict=True): if ( before.run_id != config.run_id or after.run_id != config.run_id or before.identity_digest != identity_digest or after.identity_digest != identity_digest or before.sequence != 1 or after.sequence != 2 or before.stage is not RuntimeStage.ABORTING or after.stage not in {RuntimeStage.ABORTED, RuntimeStage.QUARANTINED} or before.vote is not Vote.ABORT or after.vote is not Vote.ABORT or before.error_code is None or after.error_code != before.error_code ): raise RuntimeContractError( "LMS_LIFECYCLE_INVALID", "terminal lifecycle receipts disagree" ) advance_stage(before.stage, after.stage) return aborting, terminal def _ensure_failure_attempt(config: DistributedPreflightConfig) -> Path: attempt = config.staging_path / config.run_id if attempt.exists(): return attempt descriptor: int | None = None try: descriptor = os.open(config.staging_path, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW) os.mkdir(config.run_id, mode=0o700, dir_fd=descriptor) os.fsync(descriptor) return attempt except FileExistsError: return attempt except OSError as exc: raise RuntimeContractError( "LMS_EVIDENCE_UNAVAILABLE", "failure evidence directory could not be established", ) from exc finally: if descriptor is not None: os.close(descriptor) def _wait_for_evidence(path: Path, *, timeout_seconds: int) -> PreflightEvidence: deadline = time.monotonic() + timeout_seconds while True: try: return read_evidence(path) except FileNotFoundError: if time.monotonic() >= deadline: raise RuntimeContractError( "LMS_EVIDENCE_UNAVAILABLE", "terminal evidence was not published within its bound", ) from None time.sleep(0.01) def _write_terminal_evidence(path: Path, evidence: PreflightEvidence) -> None: try: write_evidence(path, evidence) except Exception as exc: raise RuntimeContractError( "LMS_EVIDENCE_UNAVAILABLE", "terminal evidence could not be persisted" ) from exc def execute_preflight( config: DistributedPreflightConfig, launch: TorchrunEnvironment, *, probes: PreflightProbes | None = None, ) -> PreflightEvidence: """Own the control group and emit a redacted terminal preflight record.""" result: PreflightResult | None = None prepared_evidence: PreflightEvidence | None = None lifecycle_identity_digest = contract_digest(_run_identity(config)) terminal_identity: str | None = None local_code: str | None = None try: with control_group(config, launch): try: lifecycle = _begin_lifecycle(config, launch) result = run_preflight( config, launch, probes=probes, lifecycle=lifecycle, ) advance_stage(RuntimeStage.CREATED, RuntimeStage.PREFLIGHTED) preflighted = StageMessage( run_id=config.run_id, identity_digest=result.lifecycle_identity_digest, rank=launch.rank, sequence=1, stage=RuntimeStage.PREFLIGHTED, vote=Vote.PREPARED, ) preflighted_messages = tuple( _stage_message_from_bytes(item) for item in gloo_all_gather_records(preflighted) ) _validate_stage_messages( preflighted_messages, run_id=config.run_id, world_size=config.world_size, stage=RuntimeStage.PREFLIGHTED, sequence=1, identity_digest=result.lifecycle_identity_digest, vote=Vote.PREPARED, ) terminal_identity = contract_digest( { "admission_identity": result.identity_digest, "preflighted_messages": preflighted_messages, } ) prepared_evidence = PreflightEvidence.success( run_id=config.run_id, config_digest=config.digest, world_size=config.world_size, evidence_tier=config.evidence_tier, identity_digest=terminal_identity, ) prepared_marker = StageMessage( run_id=config.run_id, identity_digest=terminal_identity, rank=config.coordinator_rank, sequence=1, stage=RuntimeStage.PREFLIGHTED, vote=Vote.PREPARED, ) prepared_path = result.attempt_path / ".preflight.prepared.stage.json" local_sink_ready = True if launch.rank == config.coordinator_rank: try: write_stage_message(prepared_path, prepared_marker) except Exception: local_sink_ready = False dist.barrier() try: observed = read_stage_message(prepared_path) local_sink_ready = local_sink_ready and observed == prepared_marker except Exception: local_sink_ready = False if not unanimous_vote(2, local_sink_ready): raise RuntimeContractError( "LMS_EVIDENCE_UNAVAILABLE", "prepared lifecycle record was not verified by every rank", ) except BaseException as error: local_code = _failure_code(error) try: local_code, _messages = _collective_aborting( config, launch, lifecycle_identity_digest, local_code, ) except BaseException as lifecycle_error: local_code = _select_failure_code({local_code, _failure_code(lifecycle_error)}) raise RuntimeContractError(local_code, "distributed preflight aborted") from None except BaseException as error: local_code = _select_failure_code( {code for code in (local_code, _failure_code(error)) if code is not None} ) attempt = result.attempt_path if result is not None else config.staging_path / config.run_id if launch.rank == config.coordinator_rank and not attempt.exists(): try: attempt = _ensure_failure_attempt(config) except RuntimeContractError as error: local_code = _select_failure_code( {code for code in (local_code, error.code) if code is not None} ) if local_code is None and prepared_evidence is not None and terminal_identity is not None: acknowledgement = StageMessage( run_id=config.run_id, identity_digest=terminal_identity, rank=launch.rank, sequence=2, stage=RuntimeStage.PREFLIGHTED, vote=Vote.COMMITTED, ) try: write_stage_message( attempt / f".rank-{launch.rank}.teardown.stage.json", acknowledgement, ) except Exception: local_code = "LMS_CLEANUP_INCOMPLETE" else: deadline = time.monotonic() + config.teardown_timeout_seconds while not attempt.is_dir() and time.monotonic() < deadline: time.sleep(0.01) try: _write_local_terminal_receipts( attempt, config=config, launch=launch, identity_digest=lifecycle_identity_digest, code=local_code or "LMS_COLLECTIVE_FAILED", quarantined=local_code == "LMS_CLEANUP_INCOMPLETE", ) except Exception: local_code = "LMS_CLEANUP_INCOMPLETE" if launch.rank == config.coordinator_rank: terminal: PreflightEvidence if ( local_code is not None or result is None or prepared_evidence is None or terminal_identity is None ): deadline = time.monotonic() + config.teardown_timeout_seconds receipts: tuple[tuple[StageMessage, ...], tuple[StageMessage, ...]] | None = None while time.monotonic() < deadline: try: receipts = _read_terminal_receipts( attempt, config=config, identity_digest=lifecycle_identity_digest, ) break except FileNotFoundError: time.sleep(0.01) except Exception: break terminal_code = local_code or "LMS_COLLECTIVE_FAILED" if receipts is None: terminal_code = "LMS_CLEANUP_INCOMPLETE" else: _aborting, completed = receipts terminal_code = _select_failure_code( {item.error_code for item in completed if item.error_code} ) if any(item.stage is RuntimeStage.QUARANTINED for item in completed): terminal_code = "LMS_CLEANUP_INCOMPLETE" terminal = PreflightEvidence.failure( run_id=config.run_id, config_digest=config.digest, code=terminal_code, world_size=config.world_size, evidence_tier=config.evidence_tier, ) else: deadline = time.monotonic() + config.teardown_timeout_seconds all_acknowledged = False while time.monotonic() < deadline: try: acknowledgements = tuple( read_stage_message( result.attempt_path / f".rank-{rank}.teardown.stage.json" ) for rank in range(config.world_size) ) except FileNotFoundError: time.sleep(0.01) continue except Exception: break try: _validate_stage_messages( acknowledgements, run_id=config.run_id, world_size=config.world_size, stage=RuntimeStage.PREFLIGHTED, sequence=2, identity_digest=terminal_identity, vote=Vote.COMMITTED, ) all_acknowledged = True except RuntimeContractError: all_acknowledged = False break terminal = ( prepared_evidence if all_acknowledged else PreflightEvidence.failure( run_id=config.run_id, config_digest=config.digest, code="LMS_CLEANUP_INCOMPLETE", world_size=config.world_size, evidence_tier=config.evidence_tier, ) ) _write_terminal_evidence(config.evidence_path, terminal) terminal = _wait_for_evidence( config.evidence_path, timeout_seconds=min(3600, config.teardown_timeout_seconds * 2 + 1), ) if terminal.result != "preflighted" or local_code is not None: code = terminal.error_code or local_code or "LMS_LIFECYCLE_INVALID" raise RuntimeContractError(code, "distributed preflight failed") from None return terminal def _require_expected(records: tuple[RankAttestation, ...], field: str, expected: object) -> None: if any(getattr(record, field) != expected for record in records): codes = { "storage_digest": "LMS_STORAGE_PROFILE_MISMATCH", "network_interface_digest": "LMS_NETWORK_PROFILE_DENIED", "device_kind": "LMS_RUNTIME_PROFILE_MISMATCH", "device_config_digest": "LMS_RUNTIME_PROFILE_MISMATCH", "software_digest": "LMS_RUNTIME_PROFILE_MISMATCH", } raise RuntimeContractError( codes.get(field, "LMS_IDENTITY_MISMATCH"), f"rank {field} values disagree with the expected identity", ) def validate_attestations( records: tuple[RankAttestation, ...], *, world_size: int, tensor_parallel_size: int, dimension_divisors: tuple[int, ...], expected_device_kind: str, expected_device_config_digest: str, expected_software_digest: str, expected_source_digest: str, expected_model_digest: str, expected_tokenizer_digest: str, expected_config_digest: str, expected_commit_sha: str, expected_code_digest: str, expected_placement_plan_digest: str, expected_storage_digest: str, expected_network_interface_digest: str, local_world_size: int, min_free_device_memory_bytes: int, min_free_host_memory_bytes: int, min_free_staging_bytes: int, ) -> None: """Require one complete, homogeneous, capacity-admitted fixed world.""" if len(records) != world_size: raise RuntimeContractError( "LMS_MEMBERSHIP_INVALID", "preflight requires exactly one attestation per rank", ) assert_rank_order(records, world_size=world_size) if tensor_parallel_size != world_size: raise RuntimeContractError( "LMS_TOPOLOGY_UNSUPPORTED", "tensor_parallel_size must equal the fixed world_size", ) if any(divisor % tensor_parallel_size for divisor in dimension_divisors): raise RuntimeContractError( "LMS_TOPOLOGY_UNSUPPORTED", "each configured dimension divisor must divide by the TP world", ) if len({record.device_digest for record in records}) != world_size: raise RuntimeContractError( "LMS_RANK_DEVICE_CONFLICT", "global device identities must be unique" ) if len({record.device_profile_digest for record in records}) != 1: raise RuntimeContractError("LMS_RUNTIME_PROFILE_MISMATCH", "rank device profiles disagree") if len({(record.host_digest, record.local_rank) for record in records}) != world_size: raise RuntimeContractError( "LMS_RANK_DEVICE_CONFLICT", "local ranks must be unique within each host", ) group_count = world_size // local_world_size if {record.group_rank for record in records} != set(range(group_count)): raise RuntimeContractError( "LMS_MEMBERSHIP_INVALID", "group ranks do not cover the fixed host groups" ) host_groups: set[str] = set() for group_rank in range(group_count): members = tuple(record for record in records if record.group_rank == group_rank) if ( len(members) != local_world_size or {record.local_rank for record in members} != set(range(local_world_size)) or {record.rank for record in members} != set( range( group_rank * local_world_size, (group_rank + 1) * local_world_size, ) ) or len({record.host_digest for record in members}) != 1 ): raise RuntimeContractError( "LMS_MEMBERSHIP_INVALID", "local ranks do not match the fixed host-group mapping", ) host_groups.add(members[0].host_digest) if len(host_groups) != group_count: raise RuntimeContractError( "LMS_MEMBERSHIP_INVALID", "one host identity spans multiple fixed groups" ) if len({record.software_digest for record in records}) != 1: raise RuntimeContractError( "LMS_RUNTIME_PROFILE_MISMATCH", "rank software identities disagree" ) if len({record.storage_digest for record in records}) != 1: raise RuntimeContractError( "LMS_STORAGE_PROFILE_MISMATCH", "rank storage identities disagree" ) expectations = { "world_size": world_size, "device_kind": expected_device_kind, "device_config_digest": expected_device_config_digest, "software_digest": expected_software_digest, "source_digest": expected_source_digest, "model_digest": expected_model_digest, "tokenizer_digest": expected_tokenizer_digest, "config_digest": expected_config_digest, "commit_sha": expected_commit_sha, "code_digest": expected_code_digest, "placement_plan_digest": expected_placement_plan_digest, "storage_digest": expected_storage_digest, "network_interface_digest": expected_network_interface_digest, "local_world_size": local_world_size, } for field, expected in expectations.items(): _require_expected(records, field, expected) if any(record.free_device_memory_bytes < min_free_device_memory_bytes for record in records): raise RuntimeContractError( "LMS_RESOURCE_ADMISSION_DENIED", "device memory headroom is below the configured floor", ) if any(record.free_host_memory_bytes < min_free_host_memory_bytes for record in records): raise RuntimeContractError( "LMS_RESOURCE_ADMISSION_DENIED", "host memory headroom is below the configured floor", ) if any(record.free_staging_bytes < min_free_staging_bytes for record in records): raise RuntimeContractError( "LMS_RESOURCE_ADMISSION_DENIED", "staging headroom is below the configured floor", )