Hugging Face's logo

sunandbean
/
diffrhythm-1.2-p100a

Text-to-Audio
ONNX
blackhole
diffrhythm
music-generation
p100a
tenstorrent
tt-dit-server
tt-metal
tt-model-cache
tt-model-catalog
tt-model-container
tt-nn
ttnn
Model card Files
xet
 Community
diffrhythm-1.2-p100a / code /models /tt_dit /tests /unit /test_host_affinity.py
sunandbean's picture
sunandbean
Add files using upload-large-folder tool
33d23e1 verified
Raw History Blame Contribute Delete
8.74 kB
# SPDX-FileCopyrightText: (c) 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
"""Pure tests for models.tt_dit.utils.host_affinity (fake sysfs trees, no device)."""
import os
from models.tt_dit.utils import host_affinity as ha
def _fake_sysfs(tmp_path, siblings_by_cpu):
root = tmp_path / "cpu"
for cpu, siblings in siblings_by_cpu.items():
d = root / f"cpu{cpu}" / "topology"
d.mkdir(parents=True)
(d / "thread_siblings_list").write_text(siblings + "\n")
return str(root)
def test_parse_cpu_list_ranges_and_singles():
assert ha._parse_cpu_list("0,32") == [0, 32]
assert ha._parse_cpu_list("0-3,8-11") == [0, 1, 2, 3, 8, 9, 10, 11]
assert ha._parse_cpu_list("5\n") == [5]
assert ha._parse_cpu_list("") == []
def test_smt_host_keeps_one_sibling_per_core(tmp_path):
# 4 cores x 2 threads, siblings paired as (c, c+4) like the galaxy host's 0-31/32-63 layout.
sysfs = _fake_sysfs(tmp_path, {c: f"{c % 4},{c % 4 + 4}" for c in range(8)})
assert ha.one_thread_per_core(sysfs) == [0, 1, 2, 3]
def test_no_smt_returns_none(tmp_path):
sysfs = _fake_sysfs(tmp_path, {c: str(c) for c in range(4)})
assert ha.one_thread_per_core(sysfs) is None
def test_missing_topology_returns_none(tmp_path):
assert ha.one_thread_per_core(str(tmp_path / "nope")) is None
def test_allowed_mask_is_respected(tmp_path):
# Operator already restricted the process to CPUs 4-7 (the second siblings): pick those, never widen.
sysfs = _fake_sysfs(tmp_path, {c: f"{c % 4},{c % 4 + 4}" for c in range(8)})
assert ha.one_thread_per_core(sysfs, allowed={4, 5, 6, 7}) == [4, 5, 6, 7]
assert ha.one_thread_per_core(sysfs, allowed={9}) is None
def _reset(monkeypatch):
monkeypatch.setattr(ha, "_applied", None)
monkeypatch.setattr(ha, "_full", None)
def test_pin_env_disable(monkeypatch):
monkeypatch.setenv("LTX_PIN_CORES", "0")
_reset(monkeypatch)
assert ha.pin_one_thread_per_core("test") is None
def test_pin_narrows_only_full_mask_threads(monkeypatch, tmp_path):
if not hasattr(os, "sched_getaffinity"):
return # non-Linux host: nothing to pin
monkeypatch.delenv("LTX_PIN_CORES", raising=False)
_reset(monkeypatch)
sysfs = _fake_sysfs(tmp_path, {c: f"{c % 4},{c % 4 + 4}" for c in range(8)})
monkeypatch.setattr(ha, "_SYSFS_CPU", sysfs)
full = {0, 1, 2, 3, 4, 5, 6, 7}
masks = {
100: set(full),
101: {6},
102: set(full),
103: {0, 1, 2, 3},
} # 101: tt-metal's own pin; 103: already narrowed
monkeypatch.setattr(ha, "_thread_ids", lambda: sorted(masks))
monkeypatch.setattr(ha.os, "sched_getaffinity", lambda tid: set(masks[tid]) if tid else set(full))
calls = []
monkeypatch.setattr(ha.os, "sched_setaffinity", lambda tid, cpus: calls.append((tid, sorted(cpus))))
monkeypatch.setattr(ha, "_cap_torch_threads", lambda n: calls.append(("torch", n)))
assert ha.pin_one_thread_per_core("test") == [0, 1, 2, 3]
assert calls == [(100, [0, 1, 2, 3]), (102, [0, 1, 2, 3]), ("torch", 4)]
# a thread born later with the full mask gets narrowed on the next call; the pinned one is still left alone
masks[100] = {0, 1, 2, 3}
masks[102] = {0, 1, 2, 3}
masks[104] = set(full)
assert ha.pin_one_thread_per_core("test") == [0, 1, 2, 3]
assert calls[-1] == (104, [0, 1, 2, 3])
def test_pin_is_noop_when_mask_already_one_per_core(monkeypatch, tmp_path):
if not hasattr(os, "sched_getaffinity"):
return
monkeypatch.delenv("LTX_PIN_CORES", raising=False)
_reset(monkeypatch)
sysfs = _fake_sysfs(tmp_path, {c: f"{c % 4},{c % 4 + 4}" for c in range(8)})
monkeypatch.setattr(ha, "_SYSFS_CPU", sysfs)
monkeypatch.setattr(ha, "_thread_ids", lambda: [100])
monkeypatch.setattr(ha.os, "sched_getaffinity", lambda tid: {0, 1, 2, 3})
calls = []
monkeypatch.setattr(ha.os, "sched_setaffinity", lambda tid, cpus: calls.append((tid, sorted(cpus))))
assert ha.pin_one_thread_per_core("test") is None
assert calls == []
def test_pin_leaves_torch_alone_when_mask_already_narrow(monkeypatch, tmp_path):
"""The re-execed process already has the one-per-core mask: no narrowing and no torch cap (capping the pool
to the mask measured slower on the galaxy ring replay, 6.4-6.7 s vs 6.2 s)."""
if not hasattr(os, "sched_getaffinity"):
return
monkeypatch.delenv("LTX_PIN_CORES", raising=False)
_reset(monkeypatch)
sysfs = _fake_sysfs(tmp_path, {c: f"{c % 4},{c % 4 + 4}" for c in range(8)})
monkeypatch.setattr(ha, "_SYSFS_CPU", sysfs)
monkeypatch.setattr(ha, "_thread_ids", lambda: [100])
monkeypatch.setattr(ha.os, "sched_getaffinity", lambda tid: {0, 1, 2, 3})
monkeypatch.setattr(
ha.os, "sched_setaffinity", lambda tid, cpus: (_ for _ in ()).throw(AssertionError("no narrowing"))
)
caps = []
monkeypatch.setattr(ha, "_cap_torch_threads", lambda n: caps.append(n))
assert ha.pin_one_thread_per_core("test") is None
assert ha.pin_one_thread_per_core("test") is None
assert caps == []
def test_cap_torch_threads_narrows_only_the_pool(monkeypatch):
import torch
before = torch.get_num_threads()
try:
ha._cap_torch_threads(1)
assert torch.get_num_threads() == 1
ha._cap_torch_threads(64) # never raises the pool
assert torch.get_num_threads() == 1
finally:
try:
torch.set_num_threads(before)
except RuntimeError:
pass
# --- reexec_pinned_before_torch: pure no-op-condition tests (os.execv always mocked) ---
def _arm_reexec(monkeypatch, current, chosen, execv_calls, setaff_calls=None):
"""Wire host_affinity so reexec_pinned_before_torch computes ``chosen`` from ``current``.
os.execv is always mocked so the test process never actually re-execs.
"""
monkeypatch.setattr(ha.os, "sched_getaffinity", lambda pid: set(current))
monkeypatch.setattr(ha, "_chosen_cores", lambda: (set(chosen) if chosen is not None else None))
if setaff_calls is not None:
monkeypatch.setattr(ha.os, "sched_setaffinity", lambda pid, cpus: setaff_calls.append((pid, sorted(cpus))))
else:
monkeypatch.setattr(ha.os, "sched_setaffinity", lambda pid, cpus: None)
monkeypatch.setattr(ha.sys, "executable", "/usr/bin/python3")
monkeypatch.setattr(ha.sys, "argv", ["prog", "arg"])
monkeypatch.setattr(ha.os, "execv", lambda exe, argv: execv_calls.append((exe, list(argv))))
def test_reexec_noop_when_pin_cores_disabled(monkeypatch):
monkeypatch.setenv("LTX_PIN_CORES", "0")
monkeypatch.delenv("_LTX_REEXECED", raising=False)
calls = []
_arm_reexec(monkeypatch, current={0, 1, 2, 3, 4, 5, 6, 7}, chosen={0, 1, 2, 3}, execv_calls=calls)
ha.reexec_pinned_before_torch("test")
assert calls == []
def test_reexec_noop_when_already_reexeced(monkeypatch):
monkeypatch.delenv("LTX_PIN_CORES", raising=False)
monkeypatch.setenv("_LTX_REEXECED", "1")
calls = []
_arm_reexec(monkeypatch, current={0, 1, 2, 3, 4, 5, 6, 7}, chosen={0, 1, 2, 3}, execv_calls=calls)
ha.reexec_pinned_before_torch("test")
assert calls == []
def test_reexec_noop_when_mask_already_chosen(monkeypatch):
monkeypatch.delenv("LTX_PIN_CORES", raising=False)
monkeypatch.delenv("_LTX_REEXECED", raising=False)
calls = []
# _chosen_cores returns None when the mask already equals the chosen set.
_arm_reexec(monkeypatch, current={0, 1, 2, 3}, chosen=None, execv_calls=calls)
ha.reexec_pinned_before_torch("test")
assert calls == []
def test_reexec_noop_when_no_topology(monkeypatch):
monkeypatch.delenv("LTX_PIN_CORES", raising=False)
monkeypatch.delenv("_LTX_REEXECED", raising=False)
calls = []
_arm_reexec(monkeypatch, current={0, 1, 2, 3, 4, 5, 6, 7}, chosen=None, execv_calls=calls)
ha.reexec_pinned_before_torch("test")
assert calls == []
def test_reexec_pins_and_execs_when_armed(monkeypatch):
if not hasattr(os, "sched_getaffinity"):
return
monkeypatch.delenv("LTX_PIN_CORES", raising=False)
monkeypatch.delenv("_LTX_REEXECED", raising=False)
execv_calls = []
setaff_calls = []
_arm_reexec(
monkeypatch,
current={0, 1, 2, 3, 4, 5, 6, 7},
chosen={0, 1, 2, 3},
execv_calls=execv_calls,
setaff_calls=setaff_calls,
)
ha.reexec_pinned_before_torch("test")
# It set the mask to the chosen cores, marked the sentinel, and re-execed the same argv.
assert setaff_calls == [(0, [0, 1, 2, 3])]
assert os.environ.get("_LTX_REEXECED") == "1"
assert execv_calls == [("/usr/bin/python3", ["/usr/bin/python3", "prog", "arg"])]