121 lines
4.6 KiB
Python
121 lines
4.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Linux CPU diagnostic: repeated safetensors materialization versus view reuse."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import gc
|
|
import json
|
|
import os
|
|
import platform
|
|
import tempfile
|
|
import time
|
|
import weakref
|
|
from pathlib import Path
|
|
|
|
|
|
def positive_int(value: str) -> int:
|
|
parsed = int(value)
|
|
if parsed <= 0:
|
|
raise argparse.ArgumentTypeError("must be positive")
|
|
return parsed
|
|
|
|
|
|
def sample(phase: str, iterations: int, loop_seconds: float) -> dict[str, object]:
|
|
gc.collect()
|
|
counters = {}
|
|
for line in Path("/proc/self/smaps_rollup").read_text().splitlines():
|
|
fields = line.split()
|
|
if len(fields) == 3 and fields[2] == "kB":
|
|
counters[fields[0].removesuffix(":")] = int(fields[1]) * 1024
|
|
return {
|
|
"phase": phase,
|
|
"iterations": iterations,
|
|
"loop_seconds": loop_seconds,
|
|
"rss_bytes": counters["Rss"],
|
|
"private_bytes": counters["Private_Clean"] + counters["Private_Dirty"],
|
|
"anonymous_bytes": counters["Anonymous"],
|
|
"fd_count": len(tuple(Path("/proc/self/fd").iterdir())),
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--mode", choices=("get_tensor", "reuse"), required=True)
|
|
parser.add_argument(
|
|
"--tensor-elements", type=positive_int, default=64, help="Number of float32 elements in the synthetic tensor"
|
|
)
|
|
parser.add_argument("--iterations", type=positive_int, default=1_000_000)
|
|
parser.add_argument("--sample-every", type=positive_int, default=100_000)
|
|
parser.add_argument("--warmup", type=positive_int, default=1000)
|
|
args = parser.parse_args()
|
|
if not Path("/proc/self/smaps_rollup").is_file():
|
|
parser.error("requires Linux /proc/self/smaps_rollup")
|
|
|
|
preparation_start = time.monotonic()
|
|
import safetensors
|
|
import torch
|
|
from safetensors import safe_open
|
|
from safetensors.torch import save_file
|
|
|
|
torch.set_num_threads(1)
|
|
torch.set_num_interop_threads(1)
|
|
samples = [sample("before_file", 0, 0.0)]
|
|
with tempfile.TemporaryDirectory(prefix="safetensors-retention-") as directory:
|
|
path = Path(directory) / "synthetic.safetensors"
|
|
save_file({"weight": torch.arange(args.tensor_elements, dtype=torch.float32)}, str(path))
|
|
with safe_open(path, framework="pt", device="cpu") as reader:
|
|
cached = reader.get_tensor("weight")
|
|
assert torch.equal(cached, torch.arange(args.tensor_elements, dtype=torch.float32))
|
|
for _ in range(args.warmup):
|
|
tensor = reader.get_tensor("weight") if args.mode == "get_tensor" else cached
|
|
del tensor
|
|
samples.append(sample("baseline", 0, 0.0))
|
|
preparation_seconds = time.monotonic() - preparation_start
|
|
completed = 0
|
|
loop_seconds = 0.0
|
|
while completed < args.iterations:
|
|
count = min(args.sample_every, args.iterations - completed)
|
|
start = time.monotonic()
|
|
for _ in range(count):
|
|
tensor = reader.get_tensor("weight") if args.mode == "get_tensor" else cached
|
|
last_tensor = weakref.ref(tensor)
|
|
del tensor
|
|
loop_seconds += time.monotonic() - start
|
|
completed += count
|
|
snapshot = sample("loop", completed, loop_seconds)
|
|
snapshot["last_tensor_alive"] = last_tensor() is not None
|
|
if args.mode == "get_tensor":
|
|
assert last_tensor() is None, "diagnostic retained a transient tensor"
|
|
samples.append(snapshot)
|
|
del cached
|
|
del reader
|
|
samples.append(sample("closed", completed, loop_seconds))
|
|
assert last_tensor() is None, "last tensor survived reader/view teardown"
|
|
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"schema_version": 1,
|
|
"arguments": vars(args),
|
|
"python": platform.python_version(),
|
|
"torch": torch.__version__,
|
|
"safetensors": safetensors.__version__,
|
|
"kernel": platform.release(),
|
|
"hostname": platform.node(),
|
|
"pid": os.getpid(),
|
|
"cpu_affinity": sorted(os.sched_getaffinity(0)),
|
|
"torch_threads": torch.get_num_threads(),
|
|
"pythonmalloc": os.environ.get("PYTHONMALLOC", "default"),
|
|
"preparation_seconds": preparation_seconds,
|
|
"samples": samples,
|
|
},
|
|
indent=2,
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|