GPUWorker

GPUWorker

The one place nnsight gets into a vLLM worker process.

vLLM's worker builds its model runner in init_device; once it has, the runner's class is swapped for nnsight's subclass, which adds behaviour and no constructor state. _load names this class as vLLM's worker_cls, a supported engine argument, so no part of vLLM's own startup is patched — and a runner nnsight does not instrument is refused here, in the worker, rather than coming up silently uninstrumented.

NNsightGPUWorker

Bases: Worker

A vLLM GPU worker whose model runner interleaves interventions.

init_device

init_device() -> None

collect_nnsight

collect_nnsight(request_ids: list[str], finished_request_ids: Optional[list[str]] = None, outputs: Optional[Any] = None) -> Optional[bytes]

Return this worker's saved values, as collective_rpc reaches it here.

outputs arrives pickled (see NNsightLLMEngine.step): the RPC is msgpack-encoded on the way in, and bytes are what it carries natively.

nnsight_request_count

nnsight_request_count() -> int

How many requests this worker's runner still tracks, via collective_rpc.

nnsight_register

nnsight_register(registration_id: str, payload: bytes, name: str | None = None) -> None

Install a block this worker runs for every request (collective_rpc).

name is what requests may address it by (edits=[...]).

nnsight_clear_registered

nnsight_clear_registered(registration_id: str) -> None

Remove a registration from this worker (collective_rpc).