176 lines
7.5 KiB
Python
176 lines
7.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Isolated real-model stream saturation, cancellation and owned-worker failure checks."""
|
|
|
|
import argparse
|
|
import asyncio
|
|
import copy
|
|
import os
|
|
import sys
|
|
import time
|
|
import traceback
|
|
from pathlib import Path
|
|
|
|
from profile_local_text import save
|
|
|
|
|
|
async def run(args):
|
|
import psutil
|
|
|
|
from vllm_omni.edge.local.engine import LocalTextEngine
|
|
from vllm_omni.edge.local.plan import plan_text_session
|
|
|
|
args.out.mkdir(parents=True, exist_ok=False)
|
|
report = {
|
|
"status": "running",
|
|
"start_unix": time.time(),
|
|
"argv": sys.argv,
|
|
"scope": (
|
|
"LocalTextEngine consumer queue and actual owned engine-process failure; "
|
|
"not proof of every internal queue being bounded"
|
|
),
|
|
}
|
|
engine = None
|
|
owned = []
|
|
save(args.out / "report.json", report)
|
|
try:
|
|
plan = plan_text_session(
|
|
args.model, max_model_len=4096, max_num_seqs=4, max_num_batched_tokens=512, enforce_eager=True
|
|
)
|
|
report["plan"] = plan.to_dict()
|
|
if not plan.admitted:
|
|
report["status"] = "rejected"
|
|
return
|
|
oversized = plan_text_session(
|
|
args.model, max_model_len=4096, max_num_seqs=100000, max_num_batched_tokens=512, enforce_eager=True
|
|
)
|
|
report["oversized_plan"] = oversized.to_dict()
|
|
report["oversized_plan_refused"] = not oversized.admitted
|
|
engine = LocalTextEngine(plan)
|
|
await engine.start()
|
|
prompt = "Describe a quiet garden with trees, flowers and a pond."
|
|
|
|
async def request(saturate=False):
|
|
session = engine.open_session()
|
|
rid, stream = await engine.submit(
|
|
session, prompt, max_tokens=128, temperature=0, ignore_eos=True, max_chunks=4
|
|
)
|
|
if saturate:
|
|
# Wait for actual capacity, rather than guessing a delay from device speed.
|
|
deadline = time.monotonic() + 30
|
|
while stream.stats()["high_water_chunks"] < 4 and time.monotonic() < deadline:
|
|
await asyncio.sleep(0.05)
|
|
report["saturated_queue_before_drain"] = stream.stats()
|
|
events = [e.to_dict() async for e in stream]
|
|
record = engine.records[rid].to_dict()
|
|
record["stream"] = stream.stats()
|
|
record["events"] = events
|
|
engine.close_session(session.session_id)
|
|
assert record["finished"] and not record["error"] and record["output_tokens"] == 128
|
|
return record
|
|
|
|
report["baseline"] = await request()
|
|
report["saturated"] = await request(saturate=True)
|
|
report["saturation_reached"] = report["saturated_queue_before_drain"]["high_water_chunks"] == 4
|
|
report["same_tokens_after_saturation"] = (
|
|
report["baseline"]["output_token_ids"] == report["saturated"]["output_token_ids"]
|
|
)
|
|
session = engine.open_session()
|
|
stale = copy.deepcopy(session)
|
|
rid, stream = await engine.submit(session, prompt, max_tokens=512, ignore_eos=True, max_chunks=4)
|
|
deadline = time.monotonic() + 30
|
|
while stream.stats()["high_water_chunks"] < 4 and time.monotonic() < deadline:
|
|
await asyncio.sleep(0.05)
|
|
report["cancel_queue_before"] = stream.stats()
|
|
report["cancel"] = await engine.cancel(rid)
|
|
report["cancel"]["late_events"] = [e.to_dict() async for e in stream]
|
|
try:
|
|
await engine.submit(stale, prompt, max_tokens=1)
|
|
except RuntimeError as error:
|
|
report["stale_handle_rejected"] = str(error)
|
|
else:
|
|
report["stale_handle_rejected"] = False
|
|
engine.close_session(session.session_id)
|
|
report["after_cancel"] = await request()
|
|
save(args.out / "report.json", report)
|
|
|
|
# Process handles come from this engine's actual manager, then must also
|
|
# belong to this probe's process tree. Never select arbitrary system PIDs.
|
|
client = engine._omni.engine.stage_clients[0]
|
|
manager = client.resources.engine_manager
|
|
processes = [p for p in manager.processes if p.is_alive()]
|
|
descendants = {p.pid: p for p in psutil.Process().children(recursive=True)}
|
|
assert processes and all(p.pid in descendants for p in processes)
|
|
manager_tree = {}
|
|
for handle in processes:
|
|
core = descendants[handle.pid]
|
|
manager_tree[core.pid] = core
|
|
manager_tree.update({child.pid: child for child in core.children(recursive=True)})
|
|
owned = [(p.pid, p.create_time()) for p in manager_tree.values()]
|
|
target = processes[0]
|
|
session = engine.open_session()
|
|
rid, stream = await engine.submit(session, prompt, max_tokens=512, ignore_eos=True)
|
|
first = await asyncio.wait_for(stream.get(), 30)
|
|
assert first is not None and first.kind == "token", "No token before worker-failure injection"
|
|
report["crash_first"] = first.to_dict()
|
|
report["crash_injection"] = {
|
|
"pid": target.pid,
|
|
"name": target.name,
|
|
"created": descendants[target.pid].create_time(),
|
|
"ownership": "engine manager handle and current process descendant",
|
|
}
|
|
start = time.perf_counter()
|
|
target.terminate()
|
|
|
|
async def drain():
|
|
return [e.to_dict() async for e in stream]
|
|
|
|
try:
|
|
events = await asyncio.wait_for(drain(), 30)
|
|
report["crash_injection"].update(
|
|
events=events,
|
|
observed_within_30s=True,
|
|
explicit_error_reported=bool(engine.records[rid].error) or any(e["kind"] == "error" for e in events),
|
|
request_record=engine.records[rid].to_dict(),
|
|
)
|
|
except asyncio.TimeoutError:
|
|
report["crash_injection"].update(observed_within_30s=False, error="No terminal stream event within 30s")
|
|
report["crash_injection"]["observation_s"] = time.perf_counter() - start
|
|
report["status"] = "completed"
|
|
except Exception as error:
|
|
report.update(status="failed", error=repr(error), traceback=traceback.format_exc())
|
|
traceback.print_exc()
|
|
finally:
|
|
if engine is not None:
|
|
try:
|
|
await engine.close()
|
|
except (Exception, asyncio.CancelledError) as error:
|
|
report.update(status="failed", cleanup_error=repr(error))
|
|
report.setdefault("error", f"Engine shutdown failed after probe: {error!r}")
|
|
await asyncio.sleep(2)
|
|
leftovers = []
|
|
for pid, created in owned:
|
|
try:
|
|
p = psutil.Process(pid)
|
|
if p.create_time() == created and p.is_running():
|
|
leftovers.append({"pid": pid, "created": created, "status": p.status()})
|
|
# Contain an isolated diagnostic's leak; report it, do not call cleanup a pass.
|
|
p.kill()
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
report["owned_processes_requiring_harness_cleanup"] = leftovers
|
|
report["end_unix"] = time.time()
|
|
save(args.out / "report.json", report)
|
|
return report["status"] == "completed"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
p = argparse.ArgumentParser(description=__doc__)
|
|
p.add_argument("--model", required=True)
|
|
p.add_argument("--out", type=Path, required=True)
|
|
args = p.parse_args()
|
|
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
|
from vllm_omni.windows.aio import install_selector_policy
|
|
|
|
install_selector_policy()
|
|
raise SystemExit(0 if asyncio.run(run(args)) else 1)
|