code/models/tt_dit/utils/host_affinity.py

# SPDX-FileCopyrightText: (c) 2026 Tenstorrent USA, Inc.

# SPDX-License-Identifier: Apache-2.0

"""Pin the process to one hardware thread per physical core on SMT hosts.

Trace replay keeps the device busy, but the readback and dispatch threads that sit between the
stages (latent hand-offs, the VAE frame readback, the audio chain) are latency-critical host work.
On a host with SMT, two of those threads can land on sibling hardware threads of one core and run at
roughly half speed; on the 32-core/64-thread galaxy host restricting the affinity mask to one sibling
per core removes the sharing. A thread-count cap alone did nothing.

Two mechanisms are provided:

1. :func:`pin_one_thread_per_core` narrows the affinity mask of full-mask threads in the running
   process. This is a best-effort safety net: torch/OMP thread pools and tt-metal dispatch/reader
   threads are already spawned (at import and at device open) across both sibling sets before the LTX
   pipeline __init__ runs, so ``sched_setaffinity`` after the fact cannot migrate them. In-process
   pinning measured ~6.7-7.5 s on the galaxy ring traced replay. It is called explicitly from
   ``LTXPipeline.__init__`` (never as an import side effect) and, in an already-narrowed process, only
   caps torch's intra-op pool to the mask.

2. :func:`reexec_pinned_before_torch` sets the affinity mask and ``os.execv`` re-execs the python
   process BEFORE torch is imported. The re-execed process inherits the pinned mask from PID start, so
   torch/OMP pools and tt-metal device threads land inside the chosen cores -- equivalent to launching
   under ``taskset`` but fully in-process. This is the mechanism that reaches the taskset-class
   numbers: galaxy ring traced replay 6.2 s re-execed (measured, no taskset) vs 6.6-6.7 s without.
   The LTX test conftest calls it from the earliest pytest hook by default (``LTX_PIN_PREIMPORT=0``
   skips it); any other entry point (a server, a script) should call it as its first statement, before
   importing torch or ttnn and before any device is opened (re-exec after device open risks a wedge).

``LTX_PIN_CORES=0`` disables both; on hosts without SMT (or without the sysfs topology) both are no-ops.
"""

from __future__ import annotations

import glob
import os
import re
import sys
from pathlib import Path

from loguru import logger

_SYSFS_CPU = "/sys/devices/system/cpu"
_applied: set[int] | None = None


def _chosen_cores() -> set[int] | None:
    """The core set the pin targets: one hardware thread per physical core within the current mask.

    Pure host topology read (sysfs); does NOT import torch/ttnn and does NOT open the device.
    Returns ``None`` when there is nothing to narrow (no SMT / no topology / mask already narrow).
    """
    if not hasattr(os, "sched_getaffinity"):
        return None
    try:
        current = set(os.sched_getaffinity(0))
    except OSError:
        return None
    chosen = one_thread_per_core(allowed=current)
    if not chosen:
        return None
    chosen_set = set(chosen)
    if chosen_set == current:
        return None
    return chosen_set


def reexec_pinned_before_torch(reason: str = "LTX pipeline (pre-import re-exec)") -> None:
    """Pin to the chosen cores and re-exec this python process so torch/ttnn inherit the mask.

    Must be called as early as possible (before torch/ttnn import and before any device is opened).
    ``os.execv`` replaces the current process image; the re-execed process inherits the affinity mask
    from PID start, so torch/OMP thread pools and tt-metal dispatch/reader threads all land inside the
    chosen cores -- the taskset-equivalent placement, fully in-process.

    No-op (returns without re-exec) when:
      * ``LTX_PIN_CORES=0`` (pinning globally disabled), or
      * ``_LTX_REEXECED`` is already set (we are the re-execed child -- prevents an infinite loop), or
      * the current affinity mask already equals the chosen core set (nothing to do), or
      * the host has no SMT / no readable topology (``_chosen_cores`` returns ``None``).

    Importable and callable without torch/ttnn present; it never imports them and never opens a device.
    """
    if os.environ.get("LTX_PIN_CORES", "1") in ("0", "false", "False"):
        return
    if os.environ.get("_LTX_REEXECED"):
        return
    if not hasattr(os, "sched_getaffinity"):
        return
    chosen = _chosen_cores()
    if not chosen:
        return
    try:
        current = set(os.sched_getaffinity(0))
    except OSError:
        return
    if current == chosen:
        return
    try:
        os.sched_setaffinity(0, chosen)
    except OSError:
        return
    os.environ["_LTX_REEXECED"] = "1"
    logger.info(
        f"host affinity: {reason}: re-exec pinned to {len(chosen)} cores {sorted(chosen)} "
        f"(was {len(current)} CPUs); torch/ttnn will inherit this mask (LTX_PIN_CORES=0 disables)"
    )
    os.execv(sys.executable, [sys.executable] + sys.argv)


def _parse_cpu_list(text: str) -> list[int]:
    """Parse a sysfs cpu list like ``0,32`` or ``0-3,8-11`` into sorted ints."""
    cpus: list[int] = []
    for part in text.strip().split(","):
        part = part.strip()
        if not part:
            continue
        m = re.fullmatch(r"(\d+)-(\d+)", part)
        if m:
            cpus.extend(range(int(m.group(1)), int(m.group(2)) + 1))
        else:
            cpus.append(int(part))
    return sorted(set(cpus))


def one_thread_per_core(sysfs_cpu: str | None = None, allowed: set[int] | None = None) -> list[int] | None:
    """Return one CPU per physical core (the lowest-numbered sibling), restricted to ``allowed``.

    Returns ``None`` when the topology is unavailable or the host has no SMT (every core has one
    sibling), so callers can leave the affinity mask alone.
    """
    sysfs_cpu = sysfs_cpu or _SYSFS_CPU
    sibling_files = glob.glob(os.path.join(sysfs_cpu, "cpu[0-9]*", "topology", "thread_siblings_list"))
    if not sibling_files:
        return None
    groups: dict[tuple[int, ...], None] = {}
    for f in sibling_files:
        try:
            siblings = tuple(_parse_cpu_list(Path(f).read_text()))
        except OSError:
            return None
        if siblings:
            groups[siblings] = None
    if not groups or all(len(g) == 1 for g in groups):
        return None
    chosen: list[int] = []
    for siblings in groups:
        candidates = [c for c in siblings if allowed is None or c in allowed]
        if candidates:
            chosen.append(min(candidates))
    return sorted(chosen) or None


def pin_one_thread_per_core(reason: str = "LTX pipeline") -> list[int] | None:
    """Narrow every thread that still carries the full CPU mask to one hardware thread per core.

    Only threads whose mask is the whole online set are touched: tt-metal pins its per-device dispatch
    and reader threads to single CPUs of its own choosing (spread over both sibling sets), and those
    placements must be left alone -- re-pinning them measured worse than not pinning at all. Threads
    created later inherit their creator's mask, and helpers born before the first call can still spawn
    full-mask threads, so the call is repeated after the device is open; it is cheap and idempotent per
    thread. torch's intra-op pool is capped to the chosen core count as well (it sizes itself from the
    CPU count it saw at import). Honours ``LTX_PIN_CORES=0``; no-op without SMT.
    """
    if os.environ.get("LTX_PIN_CORES", "1") in ("0", "false", "False"):
        return None
    if not hasattr(os, "sched_getaffinity"):
        return None
    try:
        full = set(os.sched_getaffinity(0)) if _applied is None else _full_set()
    except OSError:
        return None
    chosen = _applied if _applied is not None else set(one_thread_per_core(allowed=full) or [])
    if not chosen:
        return None
    if chosen == full:
        # Already one thread per core (the process re-execed by ``reexec_pinned_before_torch``, or a launch-time
        # taskset): nothing to narrow. torch's pool is deliberately left alone here even though the test conftest
        # sizes it from ``os.cpu_count()`` (64 threads on a 32-CPU mask): capping it to the mask was measured on
        # the galaxy ring traced replay at 6.4 s and 6.7 s vs 6.2 s with the 64-thread pool (jobs 917/918 vs
        # 919, 2026-09-18) -- the host-side stages got slower, not faster.
        if _applied is None:
            _set_applied(chosen, full)
        return None
    narrowed = 0
    for tid in _thread_ids():
        try:
            mask = set(os.sched_getaffinity(tid))
        except OSError:
            continue  # exited between listing and reading
        if mask != full:
            continue  # explicitly placed (tt-metal) or already narrowed: leave it
        try:
            os.sched_setaffinity(tid, chosen)
            narrowed += 1
        except OSError:
            continue
    first = _applied is None
    _set_applied(chosen, full)
    if first:
        _cap_torch_threads(len(chosen))
    if narrowed or first:
        logger.info(
            f"host affinity: {reason}: {len(chosen)} of {len(full)} CPUs (one hardware thread per core), "
            f"narrowed {narrowed} full-mask threads (LTX_PIN_CORES=0 disables)"
        )
    return sorted(chosen)


_full: set[int] | None = None


def _full_set() -> set[int]:
    return set(_full) if _full is not None else set(os.sched_getaffinity(0))


def _set_applied(chosen: set[int], full: set[int]) -> None:
    global _applied, _full
    _applied, _full = set(chosen), set(full)


def _cap_torch_threads(n: int) -> None:
    """torch sizes its intra-op pool from the CPU count it saw at import; cap it to the pinned cores."""
    try:
        import torch
    except ImportError:
        return  # torch absent: nothing to cap
    try:
        if torch.get_num_threads() > n:
            torch.set_num_threads(n)
    except RuntimeError:
        return  # pool already fixed: torch refuses once parallel work has started


def _thread_ids() -> list[int]:
    """All thread ids of this process (``/proc/self/task``); falls back to just the caller."""
    try:
        tids = sorted(int(t) for t in os.listdir("/proc/self/task"))
    except OSError:
        return [0]
    return tids or [0]