Add files using upload-large-folder tool

e89f16e verified 22 days ago

7.26 kB

# 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)

	# 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)