OBLITERATUS/obliteratus/distributed/launcher.py

"""Consume fixed torchrun membership without launching or provisioning workers."""

from __future__ import annotations

import ipaddress
import os
import socket
import threading
from collections.abc import Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import timedelta

import psutil  # type: ignore[import-untyped]
import torch.distributed as dist

from obliteratus.distributed.config import DistributedPreflightConfig
from obliteratus.distributed.contracts import (
    ContractError,
    RuntimeContractError,
    contract_digest,
)


_REQUIRED_ENV = frozenset(
    {
        "RANK",
        "LOCAL_RANK",
        "WORLD_SIZE",
        "LOCAL_WORLD_SIZE",
        "GROUP_RANK",
        "ROLE_RANK",
        "ROLE_WORLD_SIZE",
        "MASTER_ADDR",
        "MASTER_PORT",
        "TORCHELASTIC_RUN_ID",
        "TORCHELASTIC_RESTART_COUNT",
        "TORCHELASTIC_MAX_RESTARTS",
        "OBLITERATUS_RUN_ID",
        "GLOO_SOCKET_IFNAME",
        "NCCL_SOCKET_IFNAME",
    }
)
_SENSITIVE_ENV_PARTS = (
    "TOKEN",
    "SECRET",
    "PASSWORD",
    "PASSWD",
    "API_KEY",
    "ACCESS_KEY",
    "PRIVATE_KEY",
    "CREDENTIAL",
)
_SENSITIVE_ENV_PREFIXES = ("AWS_", "AZURE_", "GOOGLE_", "HF_", "OPENAI_", "ANTHROPIC_")
_PROXY_ENV = frozenset({"HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "NO_PROXY"})
_BACKEND_DIAGNOSTIC_LOCK = threading.RLock()
_CLEANUP_TIMEOUT_EXIT_CODE = 70


def _sensitive_environment_key(name: str) -> bool:
    upper = name.upper()
    return (
        upper in _PROXY_ENV
        or upper.startswith(_SENSITIVE_ENV_PREFIXES)
        or any(part in upper for part in _SENSITIVE_ENV_PARTS)
    )


def _env_integer(environ: Mapping[str, str], name: str, minimum: int, maximum: int) -> int:
    value = environ[name]
    if not value or len(value) > 10 or not value.isascii() or not value.isdecimal():
        raise ContractError(f"{name.lower()} must be an unsigned decimal integer")
    parsed = int(value)
    if not minimum <= parsed <= maximum:
        raise ContractError(f"{name.lower()} is outside its allowed range")
    return parsed


def _private_allowlisted_address(address: str, cidrs: tuple[str, ...]) -> bool:
    try:
        candidate = ipaddress.ip_address(address)
    except ValueError:
        return False
    private_ranges = (
        ipaddress.ip_network("10.0.0.0/8"),
        ipaddress.ip_network("172.16.0.0/12"),
        ipaddress.ip_network("192.168.0.0/16"),
        ipaddress.ip_network("fc00::/7"),
    )
    if not any(candidate in item for item in private_ranges):
        return False
    return any(candidate in ipaddress.ip_network(item, strict=True) for item in cidrs)


def validate_network_interface(
    interface: str,
    allowed_cidrs: tuple[str, ...],
    master_addr: str,
    *,
    coordinator: bool,
) -> None:
    """Require an approved private address before opening a rendezvous socket."""

    interfaces = psutil.net_if_addrs()
    if interface not in interfaces:
        raise RuntimeContractError(
            "LMS_NETWORK_PROFILE_DENIED", "configured network interface is unavailable"
        )
    networks = tuple(ipaddress.ip_network(item, strict=True) for item in allowed_cidrs)
    addresses: set[ipaddress.IPv4Address | ipaddress.IPv6Address] = set()
    for observed in interfaces[interface]:
        if observed.family not in {socket.AF_INET, socket.AF_INET6}:
            continue
        try:
            address = ipaddress.ip_address(observed.address.split("%", 1)[0])
        except ValueError:
            continue
        if any(address in network for network in networks if address.version == network.version):
            addresses.add(address)
    if not addresses:
        raise RuntimeContractError(
            "LMS_NETWORK_PROFILE_DENIED",
            "configured network interface has no private allowlisted address",
        )
    if coordinator and ipaddress.ip_address(master_addr) not in addresses:
        raise RuntimeContractError(
            "LMS_NETWORK_PROFILE_DENIED",
            "coordinator endpoint is not bound to the configured interface",
        )


@contextmanager
def _backend_diagnostic_guard() -> Iterator[None]:
    """Discard native fd-2 diagnostics and fail if the backend emits any bytes."""

    with _BACKEND_DIAGNOSTIC_LOCK:
        read_descriptor: int | None = None
        write_descriptor: int | None = None
        saved_stderr: int | None = None
        reader: threading.Thread | None = None
        observed = threading.Event()
        reader_failed = threading.Event()
        body_error: BaseException | None = None
        try:
            read_descriptor, write_descriptor = os.pipe()
            saved_stderr = os.dup(2)

            def discard() -> None:
                assert read_descriptor is not None
                try:
                    while chunk := os.read(read_descriptor, 8192):
                        observed.set()
                        del chunk
                except OSError:
                    reader_failed.set()
                finally:
                    os.close(read_descriptor)

            reader = threading.Thread(
                target=discard,
                daemon=True,
                name="obliteratus-backend-diagnostic-sink",
            )
            reader.start()
            os.dup2(write_descriptor, 2)
            os.close(write_descriptor)
            write_descriptor = None
            try:
                yield
            except BaseException as error:
                body_error = error
        except OSError as error:
            body_error = RuntimeContractError(
                "LMS_DIAGNOSTIC_REDACTION_FAILED",
                "backend diagnostic containment is unavailable",
            )
            del error
        finally:
            if saved_stderr is not None:
                try:
                    os.dup2(saved_stderr, 2)
                except OSError:
                    reader_failed.set()
                os.close(saved_stderr)
            if write_descriptor is not None:
                os.close(write_descriptor)
            if reader is not None:
                reader.join(1)
                if reader.is_alive():
                    reader_failed.set()
            elif read_descriptor is not None:
                os.close(read_descriptor)
        if (
            isinstance(body_error, RuntimeContractError)
            and body_error.code == "LMS_CLEANUP_INCOMPLETE"
        ):
            raise body_error.with_traceback(body_error.__traceback__)
        if observed.is_set() or reader_failed.is_set():
            raise RuntimeContractError(
                "LMS_DIAGNOSTIC_REDACTION_FAILED",
                "backend diagnostic output was suppressed",
            ) from None
        if body_error is not None:
            raise body_error.with_traceback(body_error.__traceback__)


@dataclass(frozen=True)
class TorchrunEnvironment:
    """One validated worker identity supplied by an external fixed scheduler."""

    rank: int
    local_rank: int
    world_size: int
    local_world_size: int
    group_rank: int
    role_rank: int
    role_world_size: int
    master_endpoint_digest: str
    run_id: str
    rendezvous_id: str
    network_interface: str

    @classmethod
    def from_environ(
        cls,
        environ: Mapping[str, str],
        config: DistributedPreflightConfig,
    ) -> "TorchrunEnvironment":
        if _REQUIRED_ENV - set(environ):
            raise RuntimeContractError(
                "LMS_LAUNCH_IDENTITY_INVALID",
                "required torchrun environment is incomplete",
            )
        if any(_sensitive_environment_key(name) for name in environ):
            raise RuntimeContractError(
                "LMS_SECRET_INPUT_REJECTED",
                "secret-bearing or proxy environment fields are forbidden",
            )
        allowed = _REQUIRED_ENV | set(config.allowed_environment_keys)
        if set(environ) - allowed:
            raise RuntimeContractError(
                "LMS_FORBIDDEN_RUNTIME_CAPABILITY",
                "environment contains fields outside the explicit allowlist",
            )
        rank = _env_integer(environ, "RANK", 0, 4095)
        local_rank = _env_integer(environ, "LOCAL_RANK", 0, 4095)
        world_size = _env_integer(environ, "WORLD_SIZE", 2, 4096)
        local_world_size = _env_integer(environ, "LOCAL_WORLD_SIZE", 1, 4096)
        group_rank = _env_integer(environ, "GROUP_RANK", 0, 4095)
        role_rank = _env_integer(environ, "ROLE_RANK", 0, 4095)
        role_world_size = _env_integer(environ, "ROLE_WORLD_SIZE", 1, 4096)
        master_port = _env_integer(environ, "MASTER_PORT", 1, 65535)
        restart_count = _env_integer(environ, "TORCHELASTIC_RESTART_COUNT", 0, 4096)
        max_restarts = _env_integer(environ, "TORCHELASTIC_MAX_RESTARTS", 0, 4096)

        if rank >= world_size:
            raise ContractError("rank must be smaller than world_size")
        if local_rank >= local_world_size:
            raise ContractError("local_rank must be smaller than local_world_size")
        if world_size != config.world_size:
            raise ContractError("world_size disagrees with the distributed profile")
        if local_world_size != config.local_world_size:
            raise ContractError("local_world_size disagrees with the distributed profile")
        expected_groups = world_size // local_world_size
        if group_rank >= expected_groups:
            raise ContractError("group_rank must be smaller than the fixed group count")
        if rank != group_rank * local_world_size + local_rank:
            raise ContractError("global, group, and local ranks disagree")
        if role_rank != rank or role_world_size != world_size:
            raise ContractError("role_world_size disagrees with fixed membership")
        if restart_count != 0 or max_restarts != 0:
            raise RuntimeContractError(
                "LMS_ELASTICITY_FORBIDDEN",
                "torchrun elasticity and restarts are forbidden",
            )
        if environ["OBLITERATUS_RUN_ID"] != config.run_id:
            raise ContractError("run_id disagrees with the distributed profile")
        if environ["TORCHELASTIC_RUN_ID"] != config.rendezvous_id:
            raise ContractError("rendezvous_id disagrees with the distributed profile")
        if environ["MASTER_ADDR"] != config.master_addr or master_port != config.master_port:
            raise ContractError("master endpoint disagrees with the distributed profile")
        if not _private_allowlisted_address(config.master_addr, config.allowed_master_cidrs):
            raise RuntimeContractError(
                "LMS_NETWORK_PROFILE_DENIED",
                "master endpoint must be a private allowlisted numeric address",
            )
        if (
            environ["GLOO_SOCKET_IFNAME"] != config.network_interface
            or environ["NCCL_SOCKET_IFNAME"] != config.network_interface
        ):
            raise ContractError("network interface disagrees with the distributed profile")
        return cls(
            rank=rank,
            local_rank=local_rank,
            world_size=world_size,
            local_world_size=local_world_size,
            group_rank=group_rank,
            role_rank=role_rank,
            role_world_size=role_world_size,
            master_endpoint_digest=contract_digest(
                {"address": config.master_addr, "port": config.master_port}
            ),
            run_id=config.run_id,
            rendezvous_id=config.rendezvous_id,
            network_interface=config.network_interface,
        )


@contextmanager
def control_group(
    config: DistributedPreflightConfig,
    launch: TorchrunEnvironment,
) -> Iterator[None]:
    """Own one bounded Gloo control group and always tear it down."""
    if not dist.is_available() or dist.is_initialized():
        raise ContractError("distributed control group is unavailable or already initialized")
    validate_network_interface(
        config.network_interface,
        config.allowed_master_cidrs,
        config.master_addr,
        coordinator=launch.rank == config.coordinator_rank,
    )

    def bounded_destroy() -> None:
        if not dist.is_initialized():
            return
        errors: list[BaseException] = []

        def destroy() -> None:
            try:
                dist.destroy_process_group()
            except BaseException as error:
                errors.append(error)

        thread = threading.Thread(target=destroy, daemon=True, name="obliteratus-gloo-teardown")
        thread.start()
        thread.join(config.teardown_timeout_seconds)
        if thread.is_alive():
            # A Python thread cannot be killed safely. Exiting this externally
            # launched worker while fd 2 is still guarded is the only bounded
            # way to prevent late native diagnostics or continued group use.
            os._exit(_CLEANUP_TIMEOUT_EXIT_CODE)
        if errors:
            raise RuntimeContractError(
                "LMS_CLEANUP_INCOMPLETE", "control group teardown failed"
            ) from None

    with _backend_diagnostic_guard():
        try:
            dist.init_process_group(
                backend="gloo",
                init_method="env://",
                rank=launch.rank,
                world_size=launch.world_size,
                timeout=timedelta(
                    seconds=min(
                        config.init_timeout_seconds,
                        config.collective_timeout_seconds,
                    )
                ),
            )
        except BaseException as error:
            bounded_destroy()
            if isinstance(error, TimeoutError) or type(error).__name__ == "DistStoreError":
                raise RuntimeContractError(
                    "LMS_MEMBERSHIP_TIMEOUT",
                    "fixed membership did not rendezvous within its bound",
                ) from None
            raise RuntimeContractError(
                "LMS_COLLECTIVE_FAILED", "control group initialization failed"
            ) from None
        try:
            yield
        finally:
            bounded_destroy()