[PyTorch] NCCL EP handle cache reaches its 4096 limit with only one live EpBuffer under allocator address churn

"""Observe EP tensor lifetime without retaining tensors or changing the cache."""

import atexit
import functools
import inspect
import json
import os
import time
import weakref
from pathlib import Path


class Audit:
    def __init__(self):
        self.created = 0
        self.dead = 0
        self.peak_live = 0
        self.attempts = 0
        self.succeeded = 0
        self.live = {}
        self.prepared = set()
        self.rank = int(os.environ.get("SLURM_PROCID", os.environ.get("RANK", "0")))
        self.output = None
        directory = os.environ.get("EP_AUDIT_DIR")
        if directory:
            Path(directory).mkdir(parents=True, exist_ok=True)
            self.output = (Path(directory) / f"rank-{self.rank}-pid-{os.getpid()}.jsonl").open("a", buffering=1)
        atexit.register(self.emit, "process_exit")

    def snapshot(self):
        return {
            "created": self.created,
            "dead": self.dead,
            "live_tensors": len(self.live),
            "live_addresses": len({item[1] for item in self.live.values()}),
            "peak_live": self.peak_live,
            "prepare_attempts": self.attempts,
            "prepare_successes": self.succeeded,
            "distinct_prepared_addresses": len(self.prepared),
        }

    def emit(self, event, **extra):
        row = {"event": event, "rank": self.rank, "pid": os.getpid(), "time": time.time(), **self.snapshot(), **extra}
        line = json.dumps(row, sort_keys=True)
        print("EP_AUDIT " + line, flush=True)
        if self.output:
            self.output.write(line + "\n")

    def observe(self, tensor):
        self.created += 1
        serial = self.created
        ptr = tensor.data_ptr()

        def released(_ref):
            self.live.pop(serial, None)
            self.dead += 1

        self.live[serial] = (weakref.ref(tensor, released), ptr)
        self.peak_live = max(self.peak_live, len(self.live))


def install(ep):
    if getattr(ep, "_cm_ep_audited", False):
        return ep._cm_ep_audit_state
    original_init = ep.EpBuffer.__init__
    original_dispatch = ep._ep_prepare_and_dispatch_fwd
    original_finalize = ep.ep_finalize
    expected = ("tokens", "topk_weights", "topk_idx", "buffer", "recv_tokens", "recv_topk_weights")
    if tuple(inspect.signature(original_dispatch).parameters) != expected:
        raise RuntimeError("EP observer requires the pinned TE prepare/dispatch signature")
    state = Audit()

    @functools.wraps(original_init)
    def initialize(self, *args, **kwargs):
        original_init(self, *args, **kwargs)
        state.observe(self.handle_mem)

    @functools.wraps(original_dispatch)
    def dispatch(tokens, topk_weights, topk_idx, buffer, recv_tokens, recv_topk_weights):
        ptr = buffer.handle_mem.data_ptr()
        state.attempts += 1
        sampled = state.attempts <= 4 or state.attempts % 128 == 0
        extra = {
            "ptr": ptr,
            "seen_before": ptr in state.prepared,
            "eager": buffer.eager,
            "handle_bytes": buffer.handle_mem.numel(),
        }
        if sampled:
            state.emit("before_prepare", **extra)
        try:
            result = original_dispatch(tokens, topk_weights, topk_idx, buffer, recv_tokens, recv_topk_weights)
        except BaseException as exc:
            state.emit("prepare_error", error=str(exc), **extra)
            raise
        state.succeeded += 1
        state.prepared.add(ptr)
        if sampled:
            state.emit("after_prepare", **extra)
        return result

    @functools.wraps(original_finalize)
    def finalize(*args, **kwargs):
        state.emit("before_finalize")
        result = original_finalize(*args, **kwargs)
        state.prepared.clear()
        state.emit("after_finalize")
        return result

    ep.EpBuffer.__init__ = initialize
    ep._ep_prepare_and_dispatch_fwd = dispatch
    ep.ep_finalize = finalize
    ep._cm_ep_audited = True
    ep._cm_ep_audit_state = state
    state.emit(
        "installed", te_ep_source=ep.__file__, cache_cap=os.environ.get("NVTE_EP_HANDLE_CACHE_SIZE", "4096(default)")
    )
    return state

"""TE-only sequential dispatch/combine/backward with measured tensor lifetime.

The churn mode allocates unrelated byte tensors after completed EP buffers die.
It never edits EpBuffer storage, private cache state, or the default cache cap.
This models allocator address churn while keeping only one EP call in flight.
"""

import argparse
import gc
import os
import sys
from datetime import timedelta
from importlib.metadata import version

import torch
import torch.distributed as dist
from audit import install
from transformer_engine.pytorch import ep


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", choices=["reuse", "fresh", "churn"], required=True)
    parser.add_argument("--steps", type=int, default=4352)
    args = parser.parse_args()
    rank = int(os.environ.get("SLURM_PROCID", os.environ.get("RANK", "0")))
    world = int(os.environ.get("SLURM_NTASKS", os.environ.get("WORLD_SIZE", "1")))
    local = int(os.environ.get("SLURM_LOCALID", os.environ.get("LOCAL_RANK", "0")))
    if world != 4 or int(os.environ.get("SLURM_JOB_NUM_NODES", "1")) != 1:
        raise RuntimeError(f"This probe requires four ranks on one node; got world={world}")
    device = 0 if torch.cuda.device_count() == 1 else local
    torch.cuda.set_device(device)
    os.environ.update(RANK=str(rank), WORLD_SIZE=str(world))
    if os.environ.get("NVTE_EP_HANDLE_CACHE_SIZE") is not None:
        raise RuntimeError("This diagnostic requires the unmodified default handle cap")
    if not version("transformer_engine").startswith("2.21.0+d77ef88c"):
        raise RuntimeError("Unexpected TE revision")
    dist.init_process_group("nccl", timeout=timedelta(seconds=120), device_id=torch.device("cuda", device))
    state = install(ep)
    state.emit(
        "probe_start",
        mode=args.mode,
        steps=args.steps,
        torch=version("torch"),
        te=version("transformer_engine"),
        world=world,
        local_rank=local,
        device=device,
    )
    tokens_count, hidden, topk, local_experts = 64, 512, 2, 2
    ep.ep_bootstrap(
        dist.group.WORLD,
        num_experts=local_experts * world,
        max_tokens_per_rank=tokens_count,
        hidden_dim=hidden,
        num_topk=topk,
        recv_capacity_per_rank=None,
        zero_copy=False,
    )
    indices = (
        (
            torch.arange(tokens_count, device="cuda")[:, None]
            + rank * tokens_count
            + torch.arange(topk, device="cuda")[None, :]
        )
        % (local_experts * world)
    ).long()
    weights = torch.full((tokens_count, topk), 1 / topk, device="cuda", dtype=torch.float32)
    original = (torch.arange(tokens_count * hidden, device="cuda").reshape(tokens_count, hidden) % 13 / 16).bfloat16()
    fillers = []
    retained_bytes = 0
    reuse = None
    completed = 0
    hit_cap = False

    def make_buffer():
        return ep.EpBuffer(
            top_k=topk,
            max_tokens_per_rank=tokens_count,
            hidden_dim=hidden,
            num_local_experts=local_experts,
            recv_capacity_per_rank=None,
            alignment=0,
        )

    def one_step(buffer):
        x = original.detach().clone().requires_grad_(True)
        received, recv_weights, _counts = ep.ep_dispatch(buffer, x, indices, weights)
        # Dispatch weights are FP32; combine's public API requires BF16 tokens.
        weighted = (received * recv_weights[:, None]).to(received.dtype)
        out = ep.ep_combine(buffer, weighted, num_local_tokens=tokens_count)
        out.float().sum().backward()
        torch.cuda.synchronize()
        errors = torch.stack([(out - original).abs().max().float(), (x.grad - 1).abs().max().float()])
        dist.all_reduce(errors, op=dist.ReduceOp.MAX)
        if errors.max().item() > 0.0625:
            raise AssertionError(f"Round-trip/gradient mismatch: {errors.tolist()}")

    if args.mode == "reuse":
        reuse = make_buffer()
    for step in range(args.steps):
        buffer = reuse if reuse is not None else make_buffer()
        ptr, handle_bytes = buffer.handle_mem.data_ptr(), buffer.handle_mem.numel()
        # Coordinate before a known cache miss at capacity. Only full ranks invoke
        # the asserting operation, avoiding a peer hanging inside an EP collective.
        full = ptr not in state.prepared and len(state.prepared) == 4096
        any_full = torch.tensor(int(full), device="cuda")
        dist.all_reduce(any_full, op=dist.ReduceOp.MAX)
        if any_full.item():
            native_error = "peer reached capacity"
            if full:
                try:
                    ep.ep_prepare(buffer, indices)
                except RuntimeError as exc:
                    native_error = str(exc)
                    if "live handle limit (4096)" not in native_error:
                        raise
                else:
                    raise AssertionError("Observed pointers did not predict native cache occupancy")
            state.emit(
                "native_cap_confirmed",
                step=step,
                local_full=full,
                error=native_error,
                retained_unrelated_bytes=retained_bytes,
                handle_bytes=handle_bytes,
            )
            hit_cap = True
            del buffer
            break
        one_step(buffer)
        completed += 1
        del buffer
        if args.mode != "reuse" and state.live:
            gc.collect()
        if args.mode != "reuse" and state.live:
            raise AssertionError(f"EP handle tensor survived a completed step: {state.snapshot()}")
        if args.mode == "churn":
            # Normal PyTorch allocation, unrelated to TE: occupies allocator space
            # after the EP tensor is dead. Bound total storage to 2 GiB.
            if retained_bytes + handle_bytes > 2 * 1024**3:
                raise RuntimeError("Churn storage exceeded its diagnostic budget")
            filler = torch.empty(handle_bytes, dtype=torch.uint8, device="cuda")
            fillers.append(filler)
            retained_bytes += handle_bytes
        if step < 4 or (step + 1) % 128 == 0:
            state.emit("step_completed", step=step + 1, retained_unrelated_bytes=retained_bytes)
    state.emit(
        "probe_result",
        mode=args.mode,
        completed=completed,
        native_cap_reached=hit_cap,
        retained_unrelated_bytes=retained_bytes,
        mcore_imported="megatron.core" in sys.modules,
    )
    dist.barrier()
    ep.ep_finalize()
    dist.destroy_process_group()


if __name__ == "__main__":
    main()