297 lines
13 KiB
Python
297 lines
13 KiB
Python
|
"""Init time and resident memory, which decide whether a model can be served
|
||
|
|
from an edge box at all.
|
||
|
|
|
||
|
|
Three quantities, because the plan's items are defined in terms of them and
|
||
|
|
the earlier version of this file measured none of them:
|
||
|
|
|
||
|
|
* **ready_s -- process start to serving.** The old ``init_s`` started its timer
|
||
|
|
*after* ``import vllm_omni``, so the 14.4 s of import cost the plan wants to
|
||
|
|
cut was invisible to the benchmark that was supposed to prove the cut. The
|
||
|
|
clock now starts at process creation, read from ``/proc/self/stat``, so
|
||
|
|
interpreter start-up and imports are inside it. ``init_s`` is still reported,
|
||
|
|
unchanged in meaning, so older rows stay comparable.
|
||
|
|
* **peak_rss_anon_mb -- the high-water mark, sampled.** What an edge device has
|
||
|
|
to survive is the *peak*, not the steady state: this model loads at 7338 MB
|
||
|
|
and settles at 3510 MB only after the allocator is asked for its arenas back,
|
||
|
|
and a device that cannot hold the peak never reaches the release. There is no
|
||
|
|
kernel high-water mark for RssAnon (``VmHWM`` and ``ru_maxrss`` both track
|
||
|
|
VmRSS, which counts the mmap'd checkpoint and moved 1-2 GB between identical
|
||
|
|
runs), so a sampler thread takes it.
|
||
|
|
* **rss_anon_at_context_mb -- memory once the cache is actually populated.**
|
||
|
|
The old default generated 8 tokens, so the KV cache was allocated and never
|
||
|
|
touched; "KV sizing saves only 188 MB" was measured there. ``--context-len``
|
||
|
|
fills it.
|
||
|
|
|
||
|
|
And one correction that changes what the older rows mean. Everything here used
|
||
|
|
to read ``/proc/self/status``, i.e. **this process only**. That is the right
|
||
|
|
number only when the model is in-process, which happens when
|
||
|
|
``VLLM_ENABLE_V1_MULTIPROCESSING=0`` keeps the executor at ``uni``
|
||
|
|
(``vllm/platforms/cpu.py:273`` promotes it to ``mp`` otherwise). Every earlier
|
||
|
|
run set that variable, so the published figures are in-process figures. Run the
|
||
|
|
same benchmark at vLLM's defaults and the weights land in a worker process the
|
||
|
|
benchmark never looked at: 796 MB "resident" for a model whose weights alone
|
||
|
|
are 1.4 GB. Resident memory is now summed over the **process tree**, which is
|
||
|
|
what an edge device actually has to hold, and both numbers are reported so the
|
||
|
|
older rows remain comparable.
|
||
|
|
|
||
|
|
Every row carries provenance, so it can be invalidated later instead of merely
|
||
|
|
doubted.
|
||
|
|
"""
|
||
|
|
import time
|
||
|
|
|
||
|
|
_T_IMPORT = time.time()
|
||
|
|
|
||
|
|
import argparse, json, os, resource, sys, threading # noqa: E402
|
||
|
|
|
||
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
|
|
from bench_harness import provenance # noqa: E402
|
||
|
|
|
||
|
|
|
||
|
|
def _rss(field: str, pid: str | int = "self") -> float | None:
|
||
|
|
try:
|
||
|
|
with open(f"/proc/{pid}/status") as f:
|
||
|
|
for line in f:
|
||
|
|
if line.startswith(field):
|
||
|
|
return round(int(line.split()[1]) / 1024, 1)
|
||
|
|
except OSError:
|
||
|
|
return None
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _descendants(root: int) -> list[int]:
|
||
|
|
"""Every process under ``root``. vLLM's default CPU path runs the model in
|
||
|
|
a worker process, so a self-only reading misses the weights entirely."""
|
||
|
|
children: dict[int, list[int]] = {}
|
||
|
|
for entry in os.listdir("/proc"):
|
||
|
|
if not entry.isdigit():
|
||
|
|
continue
|
||
|
|
try:
|
||
|
|
with open(f"/proc/{entry}/stat") as f:
|
||
|
|
fields = f.read().rsplit(")", 1)[1].split()
|
||
|
|
ppid = int(fields[1])
|
||
|
|
except (OSError, IndexError, ValueError):
|
||
|
|
continue
|
||
|
|
children.setdefault(ppid, []).append(int(entry))
|
||
|
|
out, stack = [], [root]
|
||
|
|
while stack:
|
||
|
|
pid = stack.pop()
|
||
|
|
for child in children.get(pid, []):
|
||
|
|
out.append(child)
|
||
|
|
stack.append(child)
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
def _tree_anon_mb(pids: list[int]) -> float:
|
||
|
|
return round(sum(_rss("RssAnon:", p) or 0.0 for p in pids), 1)
|
||
|
|
|
||
|
|
|
||
|
|
def _process_start_epoch() -> float | None:
|
||
|
|
"""When this process was created, so imports are inside the measurement."""
|
||
|
|
try:
|
||
|
|
with open("/proc/self/stat") as f:
|
||
|
|
fields = f.read().rsplit(")", 1)[1].split()
|
||
|
|
starttime_ticks = int(fields[19]) # /proc/self/stat field 22
|
||
|
|
hz = os.sysconf("SC_CLK_TCK")
|
||
|
|
with open("/proc/stat") as f:
|
||
|
|
for line in f:
|
||
|
|
if line.startswith("btime"):
|
||
|
|
return int(line.split()[1]) + starttime_ticks / hz
|
||
|
|
except Exception:
|
||
|
|
return None
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
class AnonSampler(threading.Thread):
|
||
|
|
"""Peak RssAnon has no kernel counter, so poll for it -- across the whole
|
||
|
|
process tree, because the model may not be in this process."""
|
||
|
|
|
||
|
|
def __init__(self, interval: float = 0.02, rescan_every: float = 0.5):
|
||
|
|
super().__init__(daemon=True)
|
||
|
|
self.interval, self.rescan_every = interval, rescan_every
|
||
|
|
self.peak = self.peak_self = 0.0
|
||
|
|
self.marks: list[dict] = []
|
||
|
|
self.pids = [os.getpid()]
|
||
|
|
self.max_procs = 1
|
||
|
|
# not `_stop`: threading.Thread already defines a private _stop().
|
||
|
|
self._done = threading.Event()
|
||
|
|
|
||
|
|
def _refresh(self) -> None:
|
||
|
|
me = os.getpid()
|
||
|
|
self.pids = [me] + _descendants(me)
|
||
|
|
self.max_procs = max(self.max_procs, len(self.pids))
|
||
|
|
|
||
|
|
def _sample(self) -> tuple[float, float]:
|
||
|
|
mine = _rss("RssAnon:") or 0.0
|
||
|
|
tree = _tree_anon_mb(self.pids)
|
||
|
|
self.peak_self = max(self.peak_self, mine)
|
||
|
|
self.peak = max(self.peak, tree)
|
||
|
|
return mine, tree
|
||
|
|
|
||
|
|
def run(self) -> None:
|
||
|
|
last_scan = 0.0
|
||
|
|
while not self._done.is_set():
|
||
|
|
now = time.monotonic()
|
||
|
|
if now - last_scan >= self.rescan_every:
|
||
|
|
self._refresh()
|
||
|
|
last_scan = now
|
||
|
|
self._sample()
|
||
|
|
self._done.wait(self.interval)
|
||
|
|
|
||
|
|
def mark(self, name: str) -> float | None:
|
||
|
|
"""Record a phase boundary with the peak reached so far."""
|
||
|
|
self._refresh()
|
||
|
|
mine, tree = self._sample()
|
||
|
|
self.marks.append({"phase": name, "anon_mb": mine,
|
||
|
|
"tree_anon_mb": tree, "n_procs": len(self.pids),
|
||
|
|
"peak_tree_so_far_mb": round(self.peak, 1)})
|
||
|
|
return mine
|
||
|
|
|
||
|
|
def stop(self) -> None:
|
||
|
|
self._done.set()
|
||
|
|
self.join(timeout=2.0)
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
ap = argparse.ArgumentParser()
|
||
|
|
ap.add_argument("--model", required=True)
|
||
|
|
ap.add_argument("--tag", required=True)
|
||
|
|
ap.add_argument("--max-model-len", type=int, default=2048)
|
||
|
|
ap.add_argument("--quantization", default=None)
|
||
|
|
ap.add_argument("--kv-cache-memory-bytes", type=int, default=None,
|
||
|
|
help="exact KV budget; VLLM_CPU_KVCACHE_SPACE is integer GiB "
|
||
|
|
"so it cannot express an edge-sized cache")
|
||
|
|
ap.add_argument("--max-num-batched-tokens", type=int, default=None,
|
||
|
|
help="vLLM sizes its warm-up forward from this; an edge "
|
||
|
|
"deployment serving one short stream does not need "
|
||
|
|
"the server default")
|
||
|
|
ap.add_argument("--trim", action="store_true",
|
||
|
|
help="return freed allocator arenas to the OS after init")
|
||
|
|
ap.add_argument("--generate", type=int, default=8,
|
||
|
|
help="tokens to generate, to prove the engine really works")
|
||
|
|
ap.add_argument("--context-len", type=int, default=0,
|
||
|
|
help="after init, run a request with this many prompt "
|
||
|
|
"tokens so the KV cache is populated, and report "
|
||
|
|
"resident memory there. 0 skips it -- and a memory "
|
||
|
|
"number taken without it says nothing about cache "
|
||
|
|
"sizing.")
|
||
|
|
ap.add_argument("--sample-interval", type=float, default=0.02)
|
||
|
|
ap.add_argument("--out", required=True)
|
||
|
|
a = ap.parse_args()
|
||
|
|
|
||
|
|
proc_start = _process_start_epoch()
|
||
|
|
sampler = AnonSampler(a.sample_interval)
|
||
|
|
sampler.start()
|
||
|
|
|
||
|
|
anon_start = sampler.mark("before_import")
|
||
|
|
import vllm_omni # noqa: F401
|
||
|
|
from vllm import LLM, SamplingParams
|
||
|
|
|
||
|
|
anon_after_import = sampler.mark("after_import")
|
||
|
|
import_s = time.time() - _T_IMPORT
|
||
|
|
|
||
|
|
t0 = time.perf_counter()
|
||
|
|
extra = {}
|
||
|
|
if a.kv_cache_memory_bytes:
|
||
|
|
extra["kv_cache_memory_bytes"] = a.kv_cache_memory_bytes
|
||
|
|
if a.max_num_batched_tokens:
|
||
|
|
extra["max_num_batched_tokens"] = a.max_num_batched_tokens
|
||
|
|
llm = LLM(model=a.model, dtype="bfloat16", max_model_len=a.max_model_len,
|
||
|
|
max_num_seqs=1, quantization=a.quantization,
|
||
|
|
enable_prefix_caching=False, trust_remote_code=False,
|
||
|
|
disable_log_stats=True, **extra)
|
||
|
|
init_s = time.perf_counter() - t0
|
||
|
|
ready_s = (time.time() - proc_start) if proc_start else None
|
||
|
|
|
||
|
|
anon_after_engine = sampler.mark("after_engine")
|
||
|
|
peak_during_load = round(sampler.peak, 1)
|
||
|
|
|
||
|
|
trimmed = None
|
||
|
|
if a.trim:
|
||
|
|
# The GPTQ repack unpacks 4-bit weights to int32 (4 bytes per weight)
|
||
|
|
# and throws the intermediates away, but a caching allocator keeps the
|
||
|
|
# arenas, so they stay resident. Ask for them back.
|
||
|
|
import ctypes
|
||
|
|
|
||
|
|
released = []
|
||
|
|
for lib, sym in ((None, "malloc_trim"),
|
||
|
|
("libtcmalloc_minimal.so.4",
|
||
|
|
"MallocExtension_ReleaseFreeMemory")):
|
||
|
|
try:
|
||
|
|
handle = ctypes.CDLL(lib) if lib else ctypes.CDLL("libc.so.6")
|
||
|
|
fn = getattr(handle, sym)
|
||
|
|
fn(0) if sym == "malloc_trim" else fn()
|
||
|
|
released.append(sym)
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
trimmed = released
|
||
|
|
anon_after_trim = sampler.mark("after_trim")
|
||
|
|
|
||
|
|
out = llm.generate([{"prompt_token_ids": [10]}],
|
||
|
|
SamplingParams(temperature=0.0, max_tokens=a.generate,
|
||
|
|
ignore_eos=True), use_tqdm=False)
|
||
|
|
n_out = len(out[0].outputs[0].token_ids)
|
||
|
|
sampler.mark("after_short_generate")
|
||
|
|
|
||
|
|
anon_at_context = context_tokens = None
|
||
|
|
if a.context_len:
|
||
|
|
# Fill the cache. Without this the cache is allocated and never
|
||
|
|
# written, so resident memory says nothing about how it was sized.
|
||
|
|
context_tokens = min(a.context_len, a.max_model_len - 16)
|
||
|
|
llm.generate([{"prompt_token_ids": [10] * context_tokens}],
|
||
|
|
SamplingParams(temperature=0.0, max_tokens=16,
|
||
|
|
ignore_eos=True), use_tqdm=False)
|
||
|
|
anon_at_context = sampler.mark("after_context_generate")
|
||
|
|
|
||
|
|
sampler.stop()
|
||
|
|
|
||
|
|
rec = {
|
||
|
|
"tag": a.tag, "model": a.model,
|
||
|
|
"ready_s": round(ready_s, 2) if ready_s else None,
|
||
|
|
"import_s": round(import_s, 2),
|
||
|
|
"init_s": round(init_s, 2),
|
||
|
|
# VmRSS counts the mmap'd checkpoint's file pages, which come and go
|
||
|
|
# with system page-cache pressure: identical configurations measured
|
||
|
|
# 1.2-2.0 GB apart. RssAnon is the allocated memory -- weights copied
|
||
|
|
# into tensors, the KV cache, activations -- and is what a memory
|
||
|
|
# budget is actually about.
|
||
|
|
# self-only, for comparability with rows taken before the tree fix
|
||
|
|
"rss_anon_mb": _rss("RssAnon:"), "rss_file_mb": _rss("RssFile:"),
|
||
|
|
# what the device actually has to hold
|
||
|
|
"rss_anon_tree_mb": _tree_anon_mb([os.getpid()] + _descendants(os.getpid())),
|
||
|
|
"peak_rss_anon_mb": round(sampler.peak, 1),
|
||
|
|
"peak_rss_anon_self_mb": round(sampler.peak_self, 1),
|
||
|
|
"peak_rss_anon_during_load_mb": peak_during_load,
|
||
|
|
"n_procs_max": sampler.max_procs,
|
||
|
|
"anon_before_import_mb": anon_start,
|
||
|
|
"anon_after_import_mb": anon_after_import,
|
||
|
|
"anon_after_engine_mb": anon_after_engine,
|
||
|
|
"anon_after_trim_mb": anon_after_trim,
|
||
|
|
"rss_anon_at_context_mb": anon_at_context,
|
||
|
|
"context_tokens": context_tokens,
|
||
|
|
"phase_marks": sampler.marks,
|
||
|
|
"trim_calls": trimmed,
|
||
|
|
"rss_mb": _rss("VmRSS:"), "rss_peak_mb": _rss("VmHWM:"),
|
||
|
|
"maxrss_mb": round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024, 1),
|
||
|
|
"kvcache_space_gib": os.environ.get("VLLM_CPU_KVCACHE_SPACE", ""),
|
||
|
|
"kv_cache_memory_bytes": a.kv_cache_memory_bytes,
|
||
|
|
"max_num_batched_tokens": a.max_num_batched_tokens,
|
||
|
|
"cache_root": os.environ.get("VLLM_CACHE_ROOT", "(default)"),
|
||
|
|
"aot_shim": os.environ.get("VLLM_OMNI_AOT_CACHE_SHIM", "1"),
|
||
|
|
"max_model_len": a.max_model_len,
|
||
|
|
"tokens_generated": n_out,
|
||
|
|
"provenance": provenance({
|
||
|
|
"quantization": a.quantization,
|
||
|
|
"sample_interval_s": a.sample_interval,
|
||
|
|
# uni vs mp decides whether the model is in this process at all
|
||
|
|
"v1_multiprocessing": os.environ.get(
|
||
|
|
"VLLM_ENABLE_V1_MULTIPROCESSING", "1"),
|
||
|
|
"memory_scope": "process tree",
|
||
|
|
}),
|
||
|
|
}
|
||
|
|
print(json.dumps(rec, indent=2))
|
||
|
|
with open(a.out, "a") as f:
|
||
|
|
f.write(json.dumps(rec) + "\n")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|