[Misc] Add standalone safetensors retention diagnostic (#7145)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
1 parent
ce27ee6038
commit
f11bf1a05d
3 files changed
+182
No files matched your search
@@ -0,0 +1,57 @@
|
||||
# Host weight dependency memory diagnostic
|
||||
|
||||
`safetensors_retention.py` isolates repeated CPU `get_tensor()` calls from
|
||||
access to one cached tensor view. It imports neither vLLM nor HWR and uses a
|
||||
synthetic float32 payload controlled by `--tensor-elements` (default: 64).
|
||||
Tensor size can affect retention: a negative result for 16 elements does not
|
||||
rule out growth for 64 elements. The selected size is recorded in the JSON
|
||||
arguments and used for both file creation and the correctness check.
|
||||
HWR itself calls `get_tensor()` when
|
||||
acquiring a lease, then exposes cached views; this probe does not measure HWR
|
||||
requests or repeated lease acquisition.
|
||||
|
||||
Run one mode per fresh process in the same Python environment and CPU affinity:
|
||||
|
||||
```bash
|
||||
git rev-parse HEAD > revision.txt
|
||||
git status --short > working-tree.txt
|
||||
git diff > working-tree.patch
|
||||
# Select an allowed CPU from your process affinity; 0 is only an example.
|
||||
CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \
|
||||
benchmarks/host_weight_runtime/safetensors_retention.py \
|
||||
--mode get_tensor --iterations 10000 --sample-every 5000 > feasibility.json
|
||||
```
|
||||
|
||||
For a size comparison, run 16 and 64 elements twice each in fresh processes,
|
||||
with cached reuse as a control at both sizes:
|
||||
|
||||
```bash
|
||||
for repetition in 1 2; do
|
||||
for elements in 16 64; do
|
||||
for mode in get_tensor reuse; do
|
||||
CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \
|
||||
benchmarks/host_weight_runtime/safetensors_retention.py \
|
||||
--mode "$mode" --tensor-elements "$elements" \
|
||||
> "$mode-$elements-$repetition.json"
|
||||
done
|
||||
done
|
||||
done
|
||||
```
|
||||
|
||||
The diagnostic sets Torch to one CPU thread, warms up before the baseline,
|
||||
and reports preparation separately from timed loop work. Checkpoints run GC
|
||||
and report Linux private/anonymous memory, RSS, and open descriptors. It checks
|
||||
that transient tensor objects and the final cached view are released. Linux
|
||||
`/proc/self/smaps_rollup` must be readable; invalid arguments and missing proc
|
||||
support fail explicitly. Keep raw JSON and the repository snapshot with results.
|
||||
|
||||
Compare the final loop sample with the warmed baseline, then inspect the closed
|
||||
sample. Private-memory growth after Python tensor objects disappear is evidence
|
||||
of retention in the dependency/allocator path, not proof of a leak's root cause.
|
||||
RSS can include file-backed pages; it is not a private-memory metric. Samples
|
||||
also include small diagnostic bookkeeping allocations. Do not assert a portable
|
||||
memory threshold or turn this synthetic loop into a per-request service estimate.
|
||||
|
||||
Do not drop shared caches, trim allocators, or add serving-time GC to make the
|
||||
numbers look smaller. A dependency replacement or version change requires its
|
||||
own reproducible comparison and correctness/lifecycle validation.
|
||||
@@ -0,0 +1,120 @@
|
||||
# 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()
|
||||
Reference in new issue
Block a user