mrm-trace 0.1.0__tar.gz

mrm_trace-0.1.0/PKG-INFO

@@ -0,0 +1,234 @@
+Metadata-Version: 2.4
+Name: mrm-trace
+Version: 0.1.0
+Summary: LLM inference memory trace platform for MRM research
+License: MIT
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: pydantic>=2.0
+Requires-Dist: pyyaml>=6.0
+Requires-Dist: typer>=0.12
+Requires-Dist: rich>=13.0
+Requires-Dist: pandas>=2.0
+Requires-Dist: pyarrow>=15.0
+Requires-Dist: psutil>=5.9
+Requires-Dist: numpy>=1.26
+Provides-Extra: test
+Requires-Dist: pytest>=8.0; extra == "test"
+Requires-Dist: pytest-cov; extra == "test"
+Requires-Dist: pytest-xdist; extra == "test"
+Requires-Dist: pytest-benchmark; extra == "test"
+Requires-Dist: pytest-mock; extra == "test"
+Requires-Dist: hypothesis>=6.0; extra == "test"
+Requires-Dist: freezegun; extra == "test"
+Provides-Extra: plots
+Requires-Dist: matplotlib>=3.8; extra == "plots"
+Requires-Dist: seaborn>=0.13; extra == "plots"
+
+# mrm-trace
+
+A Python research package for collecting, parsing, labelling, and analysing LLM inference
+memory access traces. Designed as scientific instrumentation for
+**Managed-Retention Memory (MRM)** research - it characterises how model weights, KV cache,
+activations, and runtime allocations are actually accessed during inference.
+
+**Primary metrics:** retention duration · write-once ratio · read frequency · working set size
+
+---
+
+## Requirements
+
+| Requirement | Notes |
+|---|---|
+| Linux (WSL2 supported) | `perf mem` requires Linux PMU; WSL2 works |
+| Python ≥ 3.11 | Tested on 3.11 and 3.12 |
+| sudo / CAP_PERFMON | Required for `perf mem` collection |
+
+---
+
+## Install
+
+```bash
+# Clone and set up a virtual environment
+git clone <repo-url>
+cd mrm-trace
+python -m venv venv
+source venv/bin/activate      # Windows WSL: same command
+
+# Install package + test dependencies
+pip install -e ".[test]"
+
+# Optional: install matplotlib/seaborn for figures
+pip install -e ".[test,plots]"
+```
+
+---
+
+## Quick start
+
+```bash
+# Validate a config file
+mrm-trace validate --config config/default_experiment.yaml
+
+# Preview what a run would do (dry run)
+mrm-trace plan --config config/default_experiment.yaml
+
+# Run a full experiment (requires model files + sudo for perf)
+mrm-trace run --config config/default_experiment.yaml
+```
+
+---
+
+## Running tests
+
+```bash
+# Every commit - fast, no I/O
+pytest -m unit
+
+# Pre-merge - includes integration tests
+pytest -m "unit or integration"
+
+# Before dataset release - scientific correctness checks
+pytest -m validity
+
+# Property-based invariant tests (Hypothesis)
+pytest tests/property/
+
+# Performance benchmarks (excluded from default run)
+pytest -m performance --benchmark-only
+
+# Full suite (excludes slow + performance)
+pytest
+```
+
+The test suite has three tiers:
+
+| Tier | Marker | Purpose |
+|---|---|---|
+| 1 | `unit` | Individual functions behave correctly |
+| 2 | `integration` | Components work together |
+| 3 | `validity` | Measurements are scientifically sound |
+
+Tier-3 validity tests are the most important: they verify that known synthetic inputs produce
+known metric outputs (e.g. a 30s weight retention window must yield `retention_p99_s ≈ 30.0`).
+
+---
+
+## Output layout
+
+Each run writes to `results/<model_id>/<run_id>/`:
+
+```
+results/llama-7b/run_20240101_120000/
+├── trace.parquet                  ← labelled memory access trace
+├── region_map.parquet             ← one row per region (weight, kv_cache, …)
+├── kv_block_lifecycle.parquet     ← per-block write / read / eviction timestamps
+├── metrics.csv                    ← per-region-type summary (human-readable)
+├── metadata.json                  ← hardware, software, observer effect, run validity
+├── manifest.json                  ← SHA-256 checksums for all files
+└── raw/
+    ├── perf.data
+    ├── perf_script.txt
+    └── memray.bin
+```
+
+---
+
+## Run validity classification
+
+Every run is automatically classified based on observer overhead:
+
+| Class | Criteria |
+|---|---|
+| `clean` | observer CPU < 10 %, observer mem < 5 % of target RSS, no throttle, baseline CPU < 15 % |
+| `marginal` | observer CPU < 20 %, observer mem < 15 % of target RSS, ≤ 2 throttle events |
+| `contaminated` | anything worse than marginal |
+
+Contaminated runs are archived but excluded from aggregated metrics and paper figures.
+
+---
+
+## Architecture
+
+```
+mrm_trace/
+├── cli.py              CLI (typer)
+├── api.py              Python API (Experiment class)
+├── schema_version.py   Schema version registry and compatibility checking
+├── engines/            llama.cpp / vLLM wrappers
+├── collector/          perf mem / memray / process_monitor
+├── parser/             perf script + memray parsers → trace.parquet
+├── labeller/           symbol + address-range region classification
+├── analyser/           retention / write-once / read-freq / working-set / IAI / suitability
+├── telemetry/          baseline capture / thermal / observer effect / validity classifier
+├── reporter/           CSV + Parquet export / figures / manifest / RunExporter
+└── utils/              logging / IDs / file helpers
+```
+
+Key design decisions:
+- **Streaming parser** - generators throughout; never loads full trace into RAM (ADR-2)
+- **Phase-aware tracing** - `weight_load` / `generation` / `teardown` phases distinguish weight from KV (ADR-3)
+- **Observer effect as mandatory output** - every run records overhead and validity class (ADR-4)
+- **Parquet + zstd** - column-oriented, ~3× better compression than gzip (ADR-8)
+
+---
+
+## MRM suitability labels
+
+| Label | Criteria |
+|---|---|
+| `high_mrm` | write-once ratio ≥ 0.8 **and** retention p99 ≥ 10 s |
+| `medium_mrm` | write-once ratio ≥ 0.5 **and** retention p50 ≥ 1 s |
+| `low_mrm` | everything else |
+
+In practice: model weights → `high_mrm`, short-lived KV blocks → `low_mrm`.
+
+---
+
+## Schema versioning
+
+All output files carry a `mrm_trace_schema_version` in their Parquet metadata.
+The version registry is in `mrm_trace/schema_version.py`. Readers validate
+major-version compatibility on load; a major bump is a breaking change.
+
+```python
+from mrm_trace.schema_version import check_parquet_schema
+check_parquet_schema("results/.../trace.parquet", "trace")  # raises on incompatibility
+```
+
+---
+
+## Python API
+
+```python
+from mrm_trace.labeller import TraceLabeller
+from mrm_trace.analyser import compute_all
+from mrm_trace.reporter import RunExporter
+
+# Label a stream of raw trace rows
+labeller = TraceLabeller()
+labelled = list(labeller.label(raw_rows))
+region_map   = labeller.region_map()    # call after consuming label()
+kv_lifecycle = labeller.kv_lifecycle()
+
+# Analyse
+import pandas as pd
+trace = pd.DataFrame(labelled)
+results = compute_all(trace)
+# results keys: retention_per_region, retention_summary, write_once,
+#               read_freq, working_set_per_region, working_set_summary,
+#               locality_per_region, locality_summary, iai, suitability
+
+# Export a publication-ready run directory
+exporter = RunExporter("results/llama-7b/run_001")
+exporter.export(trace, region_map, kv_lifecycle, results,
+                metadata={"run_id": "run_001"}, run_id="run_001")
+```
+
+---
+
+## Collector hierarchy
+
+1. `perf mem` - primary; requires Linux PMU + root/sudo; WSL2 supported
+2. `memray` - fallback; Python-level allocations; no root needed
+3. `process_monitor` - always runs in parallel as coarse baseline (psutil)

mrm_trace-0.1.0/README.md

@@ -0,0 +1,207 @@
+# mrm-trace
+
+A Python research package for collecting, parsing, labelling, and analysing LLM inference
+memory access traces. Designed as scientific instrumentation for
+**Managed-Retention Memory (MRM)** research - it characterises how model weights, KV cache,
+activations, and runtime allocations are actually accessed during inference.
+
+**Primary metrics:** retention duration · write-once ratio · read frequency · working set size
+
+---
+
+## Requirements
+
+| Requirement | Notes |
+|---|---|
+| Linux (WSL2 supported) | `perf mem` requires Linux PMU; WSL2 works |
+| Python ≥ 3.11 | Tested on 3.11 and 3.12 |
+| sudo / CAP_PERFMON | Required for `perf mem` collection |
+
+---
+
+## Install
+
+```bash
+# Clone and set up a virtual environment
+git clone <repo-url>
+cd mrm-trace
+python -m venv venv
+source venv/bin/activate      # Windows WSL: same command
+
+# Install package + test dependencies
+pip install -e ".[test]"
+
+# Optional: install matplotlib/seaborn for figures
+pip install -e ".[test,plots]"
+```
+
+---
+
+## Quick start
+
+```bash
+# Validate a config file
+mrm-trace validate --config config/default_experiment.yaml
+
+# Preview what a run would do (dry run)
+mrm-trace plan --config config/default_experiment.yaml
+
+# Run a full experiment (requires model files + sudo for perf)
+mrm-trace run --config config/default_experiment.yaml
+```
+
+---
+
+## Running tests
+
+```bash
+# Every commit - fast, no I/O
+pytest -m unit
+
+# Pre-merge - includes integration tests
+pytest -m "unit or integration"
+
+# Before dataset release - scientific correctness checks
+pytest -m validity
+
+# Property-based invariant tests (Hypothesis)
+pytest tests/property/
+
+# Performance benchmarks (excluded from default run)
+pytest -m performance --benchmark-only
+
+# Full suite (excludes slow + performance)
+pytest
+```
+
+The test suite has three tiers:
+
+| Tier | Marker | Purpose |
+|---|---|---|
+| 1 | `unit` | Individual functions behave correctly |
+| 2 | `integration` | Components work together |
+| 3 | `validity` | Measurements are scientifically sound |
+
+Tier-3 validity tests are the most important: they verify that known synthetic inputs produce
+known metric outputs (e.g. a 30s weight retention window must yield `retention_p99_s ≈ 30.0`).
+
+---
+
+## Output layout
+
+Each run writes to `results/<model_id>/<run_id>/`:
+
+```
+results/llama-7b/run_20240101_120000/
+├── trace.parquet                  ← labelled memory access trace
+├── region_map.parquet             ← one row per region (weight, kv_cache, …)
+├── kv_block_lifecycle.parquet     ← per-block write / read / eviction timestamps
+├── metrics.csv                    ← per-region-type summary (human-readable)
+├── metadata.json                  ← hardware, software, observer effect, run validity
+├── manifest.json                  ← SHA-256 checksums for all files
+└── raw/
+    ├── perf.data
+    ├── perf_script.txt
+    └── memray.bin
+```
+
+---
+
+## Run validity classification
+
+Every run is automatically classified based on observer overhead:
+
+| Class | Criteria |
+|---|---|
+| `clean` | observer CPU < 10 %, observer mem < 5 % of target RSS, no throttle, baseline CPU < 15 % |
+| `marginal` | observer CPU < 20 %, observer mem < 15 % of target RSS, ≤ 2 throttle events |
+| `contaminated` | anything worse than marginal |
+
+Contaminated runs are archived but excluded from aggregated metrics and paper figures.
+
+---
+
+## Architecture
+
+```
+mrm_trace/
+├── cli.py              CLI (typer)
+├── api.py              Python API (Experiment class)
+├── schema_version.py   Schema version registry and compatibility checking
+├── engines/            llama.cpp / vLLM wrappers
+├── collector/          perf mem / memray / process_monitor
+├── parser/             perf script + memray parsers → trace.parquet
+├── labeller/           symbol + address-range region classification
+├── analyser/           retention / write-once / read-freq / working-set / IAI / suitability
+├── telemetry/          baseline capture / thermal / observer effect / validity classifier
+├── reporter/           CSV + Parquet export / figures / manifest / RunExporter
+└── utils/              logging / IDs / file helpers
+```
+
+Key design decisions:
+- **Streaming parser** - generators throughout; never loads full trace into RAM (ADR-2)
+- **Phase-aware tracing** - `weight_load` / `generation` / `teardown` phases distinguish weight from KV (ADR-3)
+- **Observer effect as mandatory output** - every run records overhead and validity class (ADR-4)
+- **Parquet + zstd** - column-oriented, ~3× better compression than gzip (ADR-8)
+
+---
+
+## MRM suitability labels
+
+| Label | Criteria |
+|---|---|
+| `high_mrm` | write-once ratio ≥ 0.8 **and** retention p99 ≥ 10 s |
+| `medium_mrm` | write-once ratio ≥ 0.5 **and** retention p50 ≥ 1 s |
+| `low_mrm` | everything else |
+
+In practice: model weights → `high_mrm`, short-lived KV blocks → `low_mrm`.
+
+---
+
+## Schema versioning
+
+All output files carry a `mrm_trace_schema_version` in their Parquet metadata.
+The version registry is in `mrm_trace/schema_version.py`. Readers validate
+major-version compatibility on load; a major bump is a breaking change.
+
+```python
+from mrm_trace.schema_version import check_parquet_schema
+check_parquet_schema("results/.../trace.parquet", "trace")  # raises on incompatibility
+```
+
+---
+
+## Python API
+
+```python
+from mrm_trace.labeller import TraceLabeller
+from mrm_trace.analyser import compute_all
+from mrm_trace.reporter import RunExporter
+
+# Label a stream of raw trace rows
+labeller = TraceLabeller()
+labelled = list(labeller.label(raw_rows))
+region_map   = labeller.region_map()    # call after consuming label()
+kv_lifecycle = labeller.kv_lifecycle()
+
+# Analyse
+import pandas as pd
+trace = pd.DataFrame(labelled)
+results = compute_all(trace)
+# results keys: retention_per_region, retention_summary, write_once,
+#               read_freq, working_set_per_region, working_set_summary,
+#               locality_per_region, locality_summary, iai, suitability
+
+# Export a publication-ready run directory
+exporter = RunExporter("results/llama-7b/run_001")
+exporter.export(trace, region_map, kv_lifecycle, results,
+                metadata={"run_id": "run_001"}, run_id="run_001")
+```
+
+---
+
+## Collector hierarchy
+
+1. `perf mem` - primary; requires Linux PMU + root/sudo; WSL2 supported
+2. `memray` - fallback; Python-level allocations; no root needed
+3. `process_monitor` - always runs in parallel as coarse baseline (psutil)

mrm_trace-0.1.0/mrm_trace/__init__.py

@@ -0,0 +1,7 @@
+"""mrm-trace: LLM inference memory trace platform for MRM research."""
+
+__version__ = "0.1.0"
+
+from mrm_trace.api import Experiment
+
+__all__ = ["Experiment"]

mrm_trace-0.1.0/mrm_trace/analyser/__init__.py

@@ -0,0 +1,53 @@
+"""Analysis subsystem for mrm-trace — Phase 6."""
+
+from mrm_trace.analyser.iai import compute_iai
+from mrm_trace.analyser.locality import compute_locality
+from mrm_trace.analyser.read_freq import compute_read_freq
+from mrm_trace.analyser.retention import compute_retention
+from mrm_trace.analyser.suitability import classify_suitability
+from mrm_trace.analyser.working_set import compute_working_set
+from mrm_trace.analyser.write_once import compute_write_once
+
+__all__ = [
+    "compute_retention",
+    "compute_write_once",
+    "compute_read_freq",
+    "compute_working_set",
+    "compute_locality",
+    "compute_iai",
+    "classify_suitability",
+    "compute_all",
+]
+
+
+def compute_all(trace) -> dict:
+    """
+    Run all analysis modules on a trace DataFrame.
+
+    Returns a dict with keys:
+      retention_per_region, retention_summary,
+      write_once, read_freq,
+      working_set_per_region, working_set_summary,
+      locality_per_region, locality_summary,
+      iai, suitability
+    """
+    ret_per_region, ret_summary = compute_retention(trace)
+    wo = compute_write_once(trace)
+    rf = compute_read_freq(trace)
+    ws_per_region, ws_summary = compute_working_set(trace)
+    loc_per_region, loc_summary = compute_locality(trace)
+    iai = compute_iai(trace)
+    suit = classify_suitability(ret_summary, wo, rf)
+
+    return {
+        "retention_per_region": ret_per_region,
+        "retention_summary": ret_summary,
+        "write_once": wo,
+        "read_freq": rf,
+        "working_set_per_region": ws_per_region,
+        "working_set_summary": ws_summary,
+        "locality_per_region": loc_per_region,
+        "locality_summary": loc_summary,
+        "iai": iai,
+        "suitability": suit,
+    }

mrm_trace-0.1.0/mrm_trace/analyser/iai.py

@@ -0,0 +1,67 @@
+"""
+Inter-access interval (IAI) analysis — time between consecutive accesses.
+
+Groups by exact address (not address_page) to correctly handle sub-page KV
+blocks at 256-byte stride where multiple blocks share a 4K page.
+"""
+
+import numpy as np
+import pandas as pd
+
+
+def compute_iai(trace: pd.DataFrame) -> pd.DataFrame:
+    """
+    Compute inter-access interval distribution per region_type.
+
+    IAI = time gap between consecutive accesses to the *same exact address*,
+    sorted by timestamp_ns within each address group.
+
+    Returns
+    -------
+    DataFrame
+        Columns: region_type, n_intervals, iai_p50_ns, iai_p90_ns,
+                 iai_p99_ns, iai_mean_ns
+        One row per region_type (regions with < 2 accesses to any address
+        contribute 0 intervals and are excluded from the summary if they
+        have no intervals at all).
+    """
+    trace_sorted = trace.sort_values(["address", "timestamp_ns"])
+    all_intervals = []
+
+    for (address, region_type), grp in trace_sorted.groupby(
+        ["address", "region_type"], sort=False
+    ):
+        ts = grp["timestamp_ns"].to_numpy(dtype=np.int64)
+        if len(ts) < 2:
+            continue
+        intervals = np.diff(ts)
+        intervals = intervals[intervals > 0]  # ignore same-timestamp duplicates
+        if len(intervals) == 0:
+            continue
+        all_intervals.append(
+            pd.DataFrame({"region_type": region_type, "iai_ns": intervals})
+        )
+
+    if not all_intervals:
+        return pd.DataFrame(
+            columns=["region_type", "n_intervals", "iai_p50_ns",
+                     "iai_p90_ns", "iai_p99_ns", "iai_mean_ns"]
+        )
+
+    combined = pd.concat(all_intervals, ignore_index=True)
+
+    rows = []
+    for rtype, grp in combined.groupby("region_type"):
+        s = grp["iai_ns"]
+        rows.append(
+            {
+                "region_type": rtype,
+                "n_intervals": len(s),
+                "iai_p50_ns": int(s.quantile(0.50)),
+                "iai_p90_ns": int(s.quantile(0.90)),
+                "iai_p99_ns": int(s.quantile(0.99)),
+                "iai_mean_ns": float(s.mean()),
+            }
+        )
+
+    return pd.DataFrame(rows)

mrm_trace-0.1.0/mrm_trace/analyser/locality.py

@@ -0,0 +1,83 @@
+"""
+Access locality analysis — stride distribution and same-page access fraction.
+
+Locality measures spatial regularity: small strides and high same-page fraction
+indicate cache-friendly, predictable access patterns.
+"""
+
+from typing import Tuple
+
+import numpy as np
+import pandas as pd
+
+
+def compute_locality(trace: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]:
+    """
+    Compute locality statistics per region.
+
+    Stride = |address[i+1] - address[i]| within consecutive accesses to same region.
+    Same-page fraction = fraction of consecutive accesses that stay on the same page.
+
+    Returns
+    -------
+    per_region : DataFrame
+        Columns: region_id, region_type, n_accesses, mean_stride_bytes,
+                 median_stride_bytes, same_page_fraction
+    type_summary : DataFrame
+        Columns: region_type, mean_stride_bytes, median_stride_bytes,
+                 same_page_fraction
+    """
+    rows = []
+    trace_sorted = trace.sort_values(["region_id", "timestamp_ns"])
+
+    for region_id, grp in trace_sorted.groupby("region_id"):
+        region_type = grp["region_type"].iloc[0]
+        n = len(grp)
+
+        if n < 2:
+            rows.append(
+                {
+                    "region_id": region_id,
+                    "region_type": region_type,
+                    "n_accesses": n,
+                    "mean_stride_bytes": 0.0,
+                    "median_stride_bytes": 0.0,
+                    "same_page_fraction": 0.0,
+                }
+            )
+            continue
+
+        addrs = grp["address"].to_numpy()
+        pages = grp["address_page"].to_numpy()
+        strides = np.abs(np.diff(addrs.astype(np.int64)))
+        same_page = np.sum(pages[1:] == pages[:-1])
+
+        rows.append(
+            {
+                "region_id": region_id,
+                "region_type": region_type,
+                "n_accesses": n,
+                "mean_stride_bytes": float(strides.mean()),
+                "median_stride_bytes": float(np.median(strides)),
+                "same_page_fraction": float(same_page / len(strides)),
+            }
+        )
+
+    per_region = pd.DataFrame(rows)
+    if per_region.empty:
+        return per_region, pd.DataFrame(
+            columns=["region_type", "mean_stride_bytes", "median_stride_bytes",
+                     "same_page_fraction"]
+        )
+
+    type_summary = (
+        per_region.groupby("region_type")
+        .agg(
+            mean_stride_bytes=("mean_stride_bytes", "mean"),
+            median_stride_bytes=("median_stride_bytes", "mean"),
+            same_page_fraction=("same_page_fraction", "mean"),
+        )
+        .reset_index()
+    )
+
+    return per_region, type_summary

mrm_trace-0.1.0/mrm_trace/analyser/read_freq.py

@@ -0,0 +1,40 @@
+"""
+Read frequency analysis — total accesses, read fraction, reads-per-write.
+"""
+
+import pandas as pd
+
+
+def compute_read_freq(trace: pd.DataFrame) -> pd.DataFrame:
+    """
+    Compute read frequency statistics per region.
+
+    Returns
+    -------
+    DataFrame
+        Columns: region_id, region_type, total_reads, total_writes,
+                 total_accesses, read_fraction, reads_per_write
+    """
+    region_types = trace.groupby("region_id")["region_type"].first()
+
+    counts = trace.groupby(["region_id", "op_type"]).size().unstack(fill_value=0)
+
+    # Ensure both columns exist even if all ops are one type
+    for col in ("load", "store"):
+        if col not in counts.columns:
+            counts[col] = 0
+
+    counts = counts.rename(columns={"load": "total_reads", "store": "total_writes"})
+    counts["total_accesses"] = counts["total_reads"] + counts["total_writes"]
+    counts["read_fraction"] = (counts["total_reads"] / counts["total_accesses"]).fillna(0.0)
+    counts["reads_per_write"] = (
+        counts["total_reads"] / counts["total_writes"].replace(0, float("nan"))
+    ).fillna(0.0)
+
+    result = counts.reset_index()
+    result["region_type"] = result["region_id"].map(region_types)
+
+    return result[
+        ["region_id", "region_type", "total_reads", "total_writes",
+         "total_accesses", "read_fraction", "reads_per_write"]
+    ]