code/models/tt_transformers/host_mem_profiler.py

# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC

# SPDX-License-Identifier: Apache-2.0

"""
host_mem_profiler.py — Host-side memory profiler for model tests
===============================================================

What this script does
---------------------
Wraps a pytest run with `memory_profiler` to record the host process
memory (RSS) at 100 ms intervals while the model test executes.
After the run it produces a PNG plot showing memory over time with a
peak-memory reference line.  The script exits with pytest's own return
code, so CI pipelines see failures correctly and no plot is saved for a
failed run.

Measured memory includes all child processes spawned by the test
(e.g. the TT device runtime), giving a realistic view of total host
footprint during inference.

Dependencies
------------
`memory-profiler` and `matplotlib` are auto-installed at startup if
absent.  Everything else must already be present in the active
environment (loguru, pytest, and the tt-metal Python stack).

    pip install memory-profiler matplotlib   # manual one-time install

Arguments
---------
  --mesh-device   Value forwarded as the MESH_DEVICE environment variable
                  to the test process (e.g. N150, N300, T3K, TG).
                  Default: N150
  --hf-model      Hugging Face model ID forwarded as HF_MODEL.
                  Default: meta-llama/Llama-3.2-1B-Instruct
  --test          Pytest target — a file path, directory, or node id.
                  Default: models/tt_transformers/demo/simple_text_demo.py
  -k              Optional pytest -k filter expression to select a subset
                  of tests within the target.

Output
------
A timestamped directory is created under profiling_results/:

    profiling_results/<name>_<YYYYMMDD_HHMMSS>/
        memory_profile.png   ← time-series plot (peak annotated in red)

The script exits with pytest's return code; no plot is saved on failure.

Usage examples
--------------
# Minimal — profile the default Llama-3.2-1B demo on a single N150 card:
    python host_mem_profiler.py

# Llama-3.1-8B on a T3K (8-chip) host, full test file:
     python host_mem_profiler.py \\
        --mesh-device T3K \\
         --hf-model    meta-llama/Llama-3.1-8B-Instruct \\
         --test        models/tt_transformers/demo/simple_text_demo.py

# Same model but only the decode token accuracy test case:
     python host_mem_profiler.py \\
         --mesh-device T3K \\
        --hf-model    meta-llama/Llama-3.1-8B-Instruct \\
        --test        models/tt_transformers/demo/simple_text_demo.py \\
        -k            performance-ci-token-mathcing
"""

import argparse
import importlib
import os
import subprocess
import sys
from datetime import datetime
from pathlib import Path

from loguru import logger

_ALLOWED_PACKAGES = {"memory-profiler", "matplotlib"}


def _require(package, import_name=None):
    if package not in _ALLOWED_PACKAGES:
        raise ValueError(f"Package '{package}' is not in the allowed list: {_ALLOWED_PACKAGES}")
    import_name = import_name or package
    try:
        importlib.import_module(import_name)
    except ImportError:
        print(f"Installing missing dependency: {package}")
        subprocess.check_call([sys.executable, "-m", "pip", "install", package])
        importlib.invalidate_caches()  # flush finder cache so the new package is visible


_require("memory-profiler", "memory_profiler")
_require("matplotlib")

import matplotlib.pyplot as plt
from memory_profiler import memory_usage

PYTEST_EXTRA_ARGS = ["-v", "-s"]
OUTPUT_ROOT = Path("profiling_results")


# ── Runner ────────────────────────────────────────────────────────────────────
def make_runner(model_cfg, k_filter=None):
    """Return a zero-arg callable that runs a model's test suite."""

    def run():
        env = os.environ.copy()
        env.update(model_cfg["env"])
        cmd = [sys.executable, "-m", "pytest", model_cfg["test"], *PYTEST_EXTRA_ARGS]
        if k_filter:
            cmd += ["-k", k_filter]
        result = subprocess.run(cmd, capture_output=False, env=env)
        return result.returncode

    return run


# ── Plot ──────────────────────────────────────────────────────────────────────
def plot_results(name, elapsed, mib, out_dir):
    fig, ax = plt.subplots(figsize=(14, 6))
    ax.plot(elapsed, mib, "+-k", linewidth=1.0, markersize=4, label=name)
    ax.axhline(max(mib), color="red", linestyle="--", linewidth=0.8, label=f"Peak: {max(mib):.1f} MiB")

    ax.set_xlabel("Time (s)")
    ax.set_ylabel("Memory (MiB)")
    ax.set_title("Memory Usage Over Time")
    ax.legend()
    fig.tight_layout()
    fig.savefig(os.path.join(out_dir, "memory_profile.png"), dpi=150)
    plt.close(fig)
    logger.info(f"Saved at {os.path.join(out_dir, 'memory_profile.png')}")


# ── Main ──────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Profile host-side memory usage of a model test")
    parser.add_argument("--mesh-device", default="N150", help="MESH_DEVICE value (e.g. N150, T3K)")
    parser.add_argument("--hf-model", default="meta-llama/Llama-3.2-1B-Instruct", help="Hugging Face model ID")
    parser.add_argument("--test", default="models/tt_transformers/demo/simple_text_demo.py", help="Pytest target path")
    parser.add_argument("-k", default=None, help="pytest -k filter expression")
    args = parser.parse_args()

    name = args.hf_model.split("/")[-1] if "/" in args.hf_model else args.hf_model
    if args.k:
        name = f"{name}_{args.k}"

    model = {
        "name": name,
        "env": {
            "MESH_DEVICE": args.mesh_device,
            "HF_MODEL": args.hf_model,
        },
        "test": args.test,
    }

    logger.info(f"\n{'='*60}\nProfiling: {model['name']}\n{'='*60}")

    mem_ts, returncode = memory_usage(
        (make_runner(model, args.k), [], {}),
        interval=0.1,
        retval=True,
        timestamps=True,
        include_children=True,
    )
    if returncode:
        logger.error(f"pytest exited with code {returncode} — skipping plot")
        sys.exit(returncode)

    mib = [m for m, _ in mem_ts]
    ts = [t for _, t in mem_ts]
    t0 = ts[0]
    elapsed = [t - t0 for t in ts]

    logger.info(f"  Peak:     {max(mib):.1f} MiB")
    logger.info(f"  Baseline: {min(mib):.1f} MiB")

    safe_name = name.replace("/", "_").replace(" ", "_")
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_root = OUTPUT_ROOT.resolve()
    out_dir = (output_root / f"{safe_name}_{timestamp}").resolve()
    if not out_dir.is_relative_to(output_root):
        raise ValueError(f"Refusing to write outside output root: {out_dir}")
    out_dir.mkdir(parents=True, exist_ok=True)

    plot_results(model["name"], elapsed, mib, out_dir)