From 05a783964f41208bb3113ac59ad35b6a26a01fae Mon Sep 17 00:00:00 2001 From: Adam Belfki Date: Wed, 2 Sep 2026 14:26:04 -0400 Subject: [PATCH 1/3] remote: expose what a finished job cost on the server NDIF now reports a job's runtime and GPU footprint on its COMPLETED response. Nothing on the client read it, so the numbers stopped at the wire. Adds `meta_data` to ResponseModel and records it on the backend, which the tracer holds onto -- so the common path needs no new API: with model.trace(prompt, remote=True) as tracer: out = model.lm_head.output.save() tracer.backend.meta_data # {'runtime': 0.42, 'max_memory_usage': 2147483648, # 'max_mem_by_gpu': {'0': ...}, 'max_mem_pct_by_gpu': {'0': 20.0}} It is populated in note(), the per-frame handling every waiting mode shares, so the blocking trace, the non-blocking poll and `await backend` all get it from one place. stream() is the exception: it hands back raw updates without going through note(), so it records the report on its own path. Anything added to note() from here needs the same treatment or stream() silently won't have it. The field is a plain dict rather than a model, so the server can report a new measurement without a client release. It stays None on every non-COMPLETED status, on a failed job (note() raises before recording), and against a server too old to send it -- so read it defensively. --- docs/remote/ndif-overview.md | 21 +++++ src/nnsight/intervention/backends/remote.py | 10 +++ src/nnsight/schema/response.py | 8 +- tests/test_remote_backend.py | 93 ++++++++++++++++++++- 4 files changed, 127 insertions(+), 5 deletions(-) diff --git a/docs/remote/ndif-overview.md b/docs/remote/ndif-overview.md index 8be382da..e669fc81 100644 --- a/docs/remote/ndif-overview.md +++ b/docs/remote/ndif-overview.md @@ -77,6 +77,27 @@ Status updates are `ResponseModel` objects carrying one of these (`src/nnsight/s The client renders these as a single in-place status line (animated spinner in terminals, an in-place HTML element in Jupyter). See `StatusDisplay` (`src/nnsight/intervention/backends/display.py:58`). There is no `STREAM` status — the old `tracer.local()` hybrid-streaming path does not exist here. +## What the job cost + +The `COMPLETED` response also carries `meta_data` — what the run cost on the server. The backend keeps the last one it saw, and the tracer keeps the backend it ran on, so read it after the block exits (the backend runs in `__exit__`, so it is still `None` inside): + +```python +with model.trace("Hello", remote=True) as tracer: + out = model.lm_head.output.save() + +print(tracer.backend.meta_data) +# {'runtime': 0.42, # wall-clock seconds on the server +# 'max_memory_usage': 2147483648, # peak bytes on the worst-pressured card +# 'max_mem_by_gpu': {'0': 2147483648}, # ...per card +# 'max_mem_pct_by_gpu': {'0': 20.0}} # ...against the headroom the job had +``` + +The memory figures are what *your block* drove on top of the resident weights, not the card's total usage — the weights are the server's, and they are already there before your job starts. GPU keys are strings. + +It is a plain dict, not a model: a server can report more than the client knows about, and an older one that reports nothing leaves `meta_data` as `None`. Don't assume a key is present — `meta_data` is populated only on `COMPLETED`, so it stays `None` for a job that errored, and for a job that hasn't finished. + +A non-blocking job's `poll()` and an `AsyncRemoteBackend` record it the same way, on the backend you already hold. + ## What "meta device" means client-side When you instantiate `TransformersModel("meta-llama/Llama-3.1-70B")` without `dispatch=True`, the model is built on `torch.device("meta")` — the architecture is constructed (so `model.transformer.h[0].output` is a real envoy path) but no weights are allocated. This is what lets a machine with no GPU write intervention code against a 70B model. diff --git a/src/nnsight/intervention/backends/remote.py b/src/nnsight/intervention/backends/remote.py index b332c7f5..b16f3dfc 100644 --- a/src/nnsight/intervention/backends/remote.py +++ b/src/nnsight/intervention/backends/remote.py @@ -95,6 +95,7 @@ def __init__( self.blocking = blocking self.job_id = job_id self.status: Optional[Status] = None + self.meta_data: Optional[dict] = None self.host = host or CONFIG.API.HOST if not self.host.startswith(("http://", "https://")): raise ValueError( @@ -169,10 +170,17 @@ def note(self, response: ResponseModel) -> bool: The shared status handling for every update, however it arrived (websocket, poll): update the display, raise [`RemoteError`][nnsight.intervention.backends.remote.RemoteError] on ERROR, and return whether the status is COMPLETED (so the caller knows to fetch the result). + + Also where [`meta_data`][nnsight.intervention.backends.remote.RemoteBackend.meta_data] + is taken off the response, since every waiting mode but + [`stream`][nnsight.intervention.backends.remote.AsyncRemoteBackend.stream] + reaches its result through here. """ self.display.update(response) if response.status == Status.ERROR: raise RemoteError(response.description) + if response.meta_data is not None: + self.meta_data = response.meta_data return response.status == Status.COMPLETED def handle(self, response: ResponseModel) -> Optional[RESULT]: @@ -476,6 +484,8 @@ async def stream(self) -> AsyncIterator[Any]: try: while True: response = await self.receive() + if response.meta_data is not None: + self.meta_data = response.meta_data yield response if response.status == Status.COMPLETED: yield await self.download(response.data) diff --git a/src/nnsight/schema/response.py b/src/nnsight/schema/response.py index 2349a287..6c053b8a 100644 --- a/src/nnsight/schema/response.py +++ b/src/nnsight/schema/response.py @@ -3,7 +3,9 @@ Each update NDIF pushes for a job — whether streamed over a blocking websocket or saved to the object store and polled — arrives as a [`ResponseModel`][nnsight.schema.response.ResponseModel]. Its `Status` names where the job is in its lifecycle; on ``COMPLETED`` the -returned values ride in [`data`][nnsight.schema.response.ResponseModel.data]. +returned values ride in [`data`][nnsight.schema.response.ResponseModel.data] +and what the run cost the server rides in +[`meta_data`][nnsight.schema.response.ResponseModel.meta_data]. """ from __future__ import annotations @@ -37,7 +39,8 @@ class ResponseModel(BaseModel): Streamed over the websocket for a blocking job, or saved to the object store and fetched by the client for a non-blocking one. [`data`][nnsight.schema.response.ResponseModel.data] holds the - saved values only on a ``COMPLETED`` response. + saved values only on a ``COMPLETED`` response, and [`meta_data`][nnsight.schema.response.ResponseModel.meta_data] + what the run cost on the server -- also COMPLETED-only. """ model_config = ConfigDict(arbitrary_types_allowed=True, protected_namespaces=()) @@ -46,6 +49,7 @@ class ResponseModel(BaseModel): status: Status description: str = "" data: Optional[Any] = None + meta_data: Optional[Dict[str, Any]] = None def __str__(self) -> str: return f"[{self.id}] {self.status.value.ljust(12)} {self.description}" diff --git a/tests/test_remote_backend.py b/tests/test_remote_backend.py index 344ec3cf..ce8ec4cb 100644 --- a/tests/test_remote_backend.py +++ b/tests/test_remote_backend.py @@ -24,12 +24,13 @@ class _FakeConnection: that it was closed. recv/close are synchronous, like the real sync client the backend connects with (receive() runs recv off the event loop).""" - def __init__(self, statuses): + def __init__(self, statuses, meta_data=None): self._messages = [ ResponseModel( id="job", status=status, description="boom" if status == Status.ERROR else "", + meta_data=meta_data if status == Status.COMPLETED else None, ).model_dump_json() for status in statuses ] @@ -42,11 +43,11 @@ def close(self): self.closed = True -def _backend(statuses, result=None): +def _backend(statuses, result=None, meta_data=None): # Build a backend and drop a fake, already-subscribed connection onto it (so # __call__'s real submit is bypassed), stubbing the async download. backend = AsyncRemoteBackend(MODEL_KEY, host="http://ndif.test") - backend.connection = _FakeConnection(statuses) + backend.connection = _FakeConnection(statuses, meta_data=meta_data) async def _download(url): return result @@ -122,3 +123,89 @@ def test_is_a_remote_backend(self): from nnsight.intervention.backends.remote import RemoteBackend assert issubclass(AsyncRemoteBackend, RemoteBackend) + + +# What the server reports on COMPLETED: wall-clock seconds, and the peak GPU +# memory the request drove on top of the resident weights. Keys are the server's; +# the client stores the dict as-is and never interprets it. +META = { + "runtime": 1.25, + "max_memory_usage": 2048, + "max_mem_by_gpu": {"0": 2048, "1": 1024}, + "max_mem_pct_by_gpu": {"0": 12.5, "1": 6.25}, +} + + +class TestResponseMetaData: + """`meta_data` on the wire: it survives both encodings, and is optional.""" + + def test_survives_the_json_frame(self): + # Text frames — every status update, and a COMPLETED that carries a url. + response = ResponseModel(id="job", status=Status.COMPLETED, meta_data=META) + assert ResponseModel.model_validate_json( + response.model_dump_json() + ).meta_data == META + + def test_survives_the_pickled_frame(self): + # Binary frames — a COMPLETED whose data is the result blob itself. + response = ResponseModel(id="job", status=Status.COMPLETED, meta_data=META) + assert ResponseModel.unpickle(response.pickle()).meta_data == META + + def test_absent_from_an_older_server(self): + # A server that doesn't report cost sends no such key; parsing must not fail. + response = ResponseModel.model_validate_json( + '{"id": "job", "status": "COMPLETED"}' + ) + assert response.meta_data is None + + +class TestBackendMetaData: + """The backend keeps the finished job's cost report, whichever way it waited.""" + + def test_none_before_the_job_completes(self): + backend = AsyncRemoteBackend(MODEL_KEY, host="http://ndif.test") + assert backend.meta_data is None + + def test_recorded_when_awaited(self): + backend = _backend( + [Status.RUNNING, Status.COMPLETED], result={"out": 1}, meta_data=META + ) + asyncio.run(backend.resolve()) + assert backend.meta_data == META + + def test_recorded_when_streamed(self): + # stream() bypasses note(), so it records the report on its own path. + backend = _backend( + [Status.RUNNING, Status.COMPLETED], result={"out": 1}, meta_data=META + ) + + async def go(): + async for _ in backend: + pass + + asyncio.run(go()) + assert backend.meta_data == META + + def test_recorded_off_a_polled_response(self): + # The blocking and non-blocking paths both reach a response through note(). + from nnsight.intervention.backends.remote import RemoteBackend + + backend = RemoteBackend(MODEL_KEY, host="http://ndif.test") + assert backend.note( + ResponseModel(id="job", status=Status.COMPLETED, meta_data=META) + ) + assert backend.meta_data == META + + def test_intermediate_updates_leave_it_alone(self): + # RUNNING carries no report; it must not clear one already recorded. + from nnsight.intervention.backends.remote import RemoteBackend + + backend = RemoteBackend(MODEL_KEY, host="http://ndif.test") + backend.note(ResponseModel(id="job", status=Status.COMPLETED, meta_data=META)) + backend.note(ResponseModel(id="job", status=Status.RUNNING)) + assert backend.meta_data == META + + def test_stays_none_against_an_older_server(self): + backend = _backend([Status.RUNNING, Status.COMPLETED], result={"out": 1}) + asyncio.run(backend.resolve()) + assert backend.meta_data is None From 9f99cbcbeaa0e8fb16293fb553129d0f24df949c Mon Sep 17 00:00:00 2001 From: Adam Belfki Date: Wed, 2 Sep 2026 16:01:08 -0400 Subject: [PATCH 2/3] remote: keep a failed job's cost report instead of dropping it note() raised on ERROR before it recorded meta_data, so the one status that carries the most useful report was the one status that threw it away. A job that runs out of GPU memory now reports how far past its allowance it reached, and that was being discarded a line before it could be read. Swaps the two blocks. The tracer holds the backend and is bound at __enter__, so the report outlives the raise and is read where the failure is handled: try: with model.trace(prompt, remote=True) as tracer: acts = model.transformer.h[-1].output.save() except RemoteError: tracer.backend.meta_data["extra_memory_needed"] # {'0': 8025221248} extra_memory_needed is bytes per GPU and appears only on an out-of-memory failure, so its presence is the signal; every other outcome simply omits it. --- docs/remote/ndif-overview.md | 17 ++++++++++++++++- src/nnsight/intervention/backends/remote.py | 8 ++++++-- src/nnsight/schema/response.py | 4 +++- tests/test_remote_backend.py | 20 ++++++++++++++++++++ 4 files changed, 45 insertions(+), 4 deletions(-) diff --git a/docs/remote/ndif-overview.md b/docs/remote/ndif-overview.md index e669fc81..1b1eacd4 100644 --- a/docs/remote/ndif-overview.md +++ b/docs/remote/ndif-overview.md @@ -94,7 +94,22 @@ print(tracer.backend.meta_data) The memory figures are what *your block* drove on top of the resident weights, not the card's total usage — the weights are the server's, and they are already there before your job starts. GPU keys are strings. -It is a plain dict, not a model: a server can report more than the client knows about, and an older one that reports nothing leaves `meta_data` as `None`. Don't assume a key is present — `meta_data` is populated only on `COMPLETED`, so it stays `None` for a job that errored, and for a job that hasn't finished. +It is a plain dict, not a model: a server can report more than the client knows about, and an older one that reports nothing leaves `meta_data` as `None`. Don't assume a key is present. + +A **failed** job reports its cost too, which is when it matters most. The backend records it before raising, so catch the error and read it off the tracer: + +```python +from nnsight.intervention.backends.remote import RemoteError + +try: + with model.trace(prompt, remote=True) as tracer: + acts = model.transformer.h[-1].output.save() +except RemoteError: + print(tracer.backend.meta_data) + # {'runtime': 3.1, ..., 'extra_memory_needed': {'0': 1310000000}} +``` + +`extra_memory_needed` appears only when the server ran out of GPU memory. It maps **each card that ran out** to how many bytes past its allowance the block reached there — on a sharded model, *which* card is itself the finding. It is — the number that says how much you need to free, which the traceback cannot tell you: the allocation that failed is by definition the one that never counted. Treat it as approximate; the allocator drops cached blocks and retries before it gives up. A non-blocking job's `poll()` and an `AsyncRemoteBackend` record it the same way, on the backend you already hold. diff --git a/src/nnsight/intervention/backends/remote.py b/src/nnsight/intervention/backends/remote.py index b16f3dfc..8a768d9a 100644 --- a/src/nnsight/intervention/backends/remote.py +++ b/src/nnsight/intervention/backends/remote.py @@ -177,10 +177,14 @@ def note(self, response: ResponseModel) -> bool: reaches its result through here. """ self.display.update(response) - if response.status == Status.ERROR: - raise RemoteError(response.description) + # Recorded before the raise: a failed job reports its cost too, and on an + # OOM that report is the most useful thing it has to say. `tracer` is + # bound by the time the backend runs, so an `except RemoteError` can read + # `tracer.backend.meta_data`. if response.meta_data is not None: self.meta_data = response.meta_data + if response.status == Status.ERROR: + raise RemoteError(response.description) return response.status == Status.COMPLETED def handle(self, response: ResponseModel) -> Optional[RESULT]: diff --git a/src/nnsight/schema/response.py b/src/nnsight/schema/response.py index 6c053b8a..3ec2eb88 100644 --- a/src/nnsight/schema/response.py +++ b/src/nnsight/schema/response.py @@ -40,7 +40,9 @@ class ResponseModel(BaseModel): Streamed over the websocket for a blocking job, or saved to the object store and fetched by the client for a non-blocking one. [`data`][nnsight.schema.response.ResponseModel.data] holds the saved values only on a ``COMPLETED`` response, and [`meta_data`][nnsight.schema.response.ResponseModel.meta_data] - what the run cost on the server -- also COMPLETED-only. + what the run cost on the server -- on ``COMPLETED`` and on a failure alike, + since a job that ran out of memory or timed out is exactly when its cost is + worth reading. """ model_config = ConfigDict(arbitrary_types_allowed=True, protected_namespaces=()) diff --git a/tests/test_remote_backend.py b/tests/test_remote_backend.py index ce8ec4cb..bf333e8a 100644 --- a/tests/test_remote_backend.py +++ b/tests/test_remote_backend.py @@ -205,6 +205,26 @@ def test_intermediate_updates_leave_it_alone(self): backend.note(ResponseModel(id="job", status=Status.RUNNING)) assert backend.meta_data == META + def test_recorded_even_when_the_job_fails(self): + # The report is taken off the response *before* note() raises. A failed + # job is exactly when it earns its keep -- an OOM's meta_data carries + # extra_memory_needed, which the traceback cannot tell you. + from nnsight.intervention.backends.remote import RemoteBackend + + backend = RemoteBackend(MODEL_KEY, host="http://ndif.test") + failure = dict(META, extra_memory_needed={"0": 1_310_000_000}) + with pytest.raises(RemoteError, match="out of memory"): + backend.note( + ResponseModel( + id="job", + status=Status.ERROR, + description="CUDA out of memory", + meta_data=failure, + ) + ) + assert backend.meta_data == failure + assert backend.meta_data["extra_memory_needed"] == {"0": 1_310_000_000} + def test_stays_none_against_an_older_server(self): backend = _backend([Status.RUNNING, Status.COMPLETED], result={"out": 1}) asyncio.run(backend.resolve()) From 0086d7e4e0f42859afb03f97773d6b87c75c218b Mon Sep 17 00:00:00 2001 From: Adam Belfki Date: Wed, 2 Sep 2026 17:10:57 -0400 Subject: [PATCH 3/3] schema: give the cost report a type instead of a bare dict meta_data was Dict[str, Any], so its shape lived in prose -- a docstring here and a table in the ndif repo -- with nothing enforcing the two agreed. They had already drifted twice. MetaData makes the contract executable: nnsight declares it, the server constructs it, and a wrong key or a stray integer GPU id now fails in the actor that made the mistake rather than reaching a user as a differently shaped payload. tracer.backend.meta_data.runtime tracer.backend.meta_data.alloc_shortfall_by_gpu # {'0': 8025221248} Three settings carry the design and each is tested: extra="allow" -- a server may report a measurement this client has never heard of, and it stays readable as an attribute instead of being dropped. This was the whole reason to prefer a dict, and pydantic gives it for free. Every field optional -- an older server reporting nothing is ordinary, not an error. An unreadable report is discarded rather than raised. meta_data is diagnostic; data is the job. Failing a run that finished perfectly well in order to complain about the note attached to it is the wrong trade. It warns on the way out, though: with every field optional and extras allowed, the only way to land there is a known field with the wrong type -- a server bug or version mismatch, worth hearing about -- and silence would make it indistinguishable from an older server that reports nothing. The out-of-memory field is alloc_shortfall_by_gpu, keyed like the other per-device maps. It is the part of the refused allocation that would not fit, not the size of the allocation itself: asking for 2 GB with 1.9 GB free and asking for it with nothing free are the same request and completely different problems. --- docs/remote/ndif-overview.md | 17 +++-- src/nnsight/intervention/backends/remote.py | 4 +- src/nnsight/schema/response.py | 79 ++++++++++++++++++++- tests/test_remote_backend.py | 60 +++++++++++++--- 4 files changed, 136 insertions(+), 24 deletions(-) diff --git a/docs/remote/ndif-overview.md b/docs/remote/ndif-overview.md index 1b1eacd4..1c266702 100644 --- a/docs/remote/ndif-overview.md +++ b/docs/remote/ndif-overview.md @@ -85,16 +85,16 @@ The `COMPLETED` response also carries `meta_data` — what the run cost on the s with model.trace("Hello", remote=True) as tracer: out = model.lm_head.output.save() -print(tracer.backend.meta_data) -# {'runtime': 0.42, # wall-clock seconds on the server -# 'max_memory_usage': 2147483648, # peak bytes on the worst-pressured card -# 'max_mem_by_gpu': {'0': 2147483648}, # ...per card -# 'max_mem_pct_by_gpu': {'0': 20.0}} # ...against the headroom the job had +meta = tracer.backend.meta_data +meta.runtime # 0.42 wall-clock seconds on the server +meta.max_memory_usage # 2147483648 peak bytes on the worst-pressured card +meta.max_mem_by_gpu # {'0': 2147483648} ...per card +meta.max_mem_pct_by_gpu # {'0': 20.0} ...against the headroom the job had ``` The memory figures are what *your block* drove on top of the resident weights, not the card's total usage — the weights are the server's, and they are already there before your job starts. GPU keys are strings. -It is a plain dict, not a model: a server can report more than the client knows about, and an older one that reports nothing leaves `meta_data` as `None`. Don't assume a key is present. +It is a `MetaData` model (`src/nnsight/schema/response.py`), so the fields autocomplete and are documented on the type. Every field is optional — an older server may report none of them — so check before trusting one. Unknown fields from a newer server are kept and readable as attributes rather than dropped, and a report the client cannot parse is discarded on its own without failing the response that carried it. A **failed** job reports its cost too, which is when it matters most. The backend records it before raising, so catch the error and read it off the tracer: @@ -105,11 +105,10 @@ try: with model.trace(prompt, remote=True) as tracer: acts = model.transformer.h[-1].output.save() except RemoteError: - print(tracer.backend.meta_data) - # {'runtime': 3.1, ..., 'extra_memory_needed': {'0': 1310000000}} + print(tracer.backend.meta_data.alloc_shortfall_by_gpu) # {'0': 1310000000} ``` -`extra_memory_needed` appears only when the server ran out of GPU memory. It maps **each card that ran out** to how many bytes past its allowance the block reached there — on a sharded model, *which* card is itself the finding. It is — the number that says how much you need to free, which the traceback cannot tell you: the allocation that failed is by definition the one that never counted. Treat it as approximate; the allocator drops cached blocks and retries before it gives up. +`alloc_shortfall_by_gpu` appears only when the server ran out of GPU memory. For each card that ran out it gives the part of the refused allocation that would not fit — the number that says how much you need to free. That is not the size of the refused allocation: asking for 2 GB with 1.9 GB free and asking for it with nothing free are the same request and completely different problems. The traceback can't tell you either, since the allocation that failed is by definition the one that never counted. On a sharded model, *which* card is itself the finding. Treat it as approximate; the allocator drops cached blocks and retries before it gives up. A non-blocking job's `poll()` and an `AsyncRemoteBackend` record it the same way, on the backend you already hold. diff --git a/src/nnsight/intervention/backends/remote.py b/src/nnsight/intervention/backends/remote.py index 8a768d9a..c509da6b 100644 --- a/src/nnsight/intervention/backends/remote.py +++ b/src/nnsight/intervention/backends/remote.py @@ -27,7 +27,7 @@ from typing import Any, AsyncIterator, Optional, Union from ...schema.config import CONFIG -from ...schema.response import RESULT, ResponseModel, Status +from ...schema.response import RESULT, MetaData, ResponseModel, Status from ...schema.request import RequestModel from ...tracing.backend import Backend from ...tracing.tracer import Tracer, save @@ -95,7 +95,7 @@ def __init__( self.blocking = blocking self.job_id = job_id self.status: Optional[Status] = None - self.meta_data: Optional[dict] = None + self.meta_data: Optional[MetaData] = None self.host = host or CONFIG.API.HOST if not self.host.startswith(("http://", "https://")): raise ValueError( diff --git a/src/nnsight/schema/response.py b/src/nnsight/schema/response.py index 3ec2eb88..b9ab765d 100644 --- a/src/nnsight/schema/response.py +++ b/src/nnsight/schema/response.py @@ -11,15 +11,61 @@ from __future__ import annotations import io +import warnings from enum import Enum from typing import Any, Dict, Optional import torch -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator RESULT = Dict[str, Any] +class MetaData(BaseModel): + """What a remote job cost the server, reported on the response that ends it. + + The measurements NDIF takes around a run, shaped for the client. Present on + ``COMPLETED`` and on a failure alike -- a job that timed out or ran out of + memory is exactly when its cost is worth reading. + + Attributes: + runtime: Wall-clock seconds the block ran on the server. + max_memory_usage: Peak bytes on the worst-pressured single device. + max_mem_by_gpu: Bytes the request drove *above the resident weights*, + per device. Not the card's total usage -- the weights are the + server's, not the request's, and you cannot shrink them. + max_mem_pct_by_gpu: ``max_mem_by_gpu`` against the headroom the request + actually had (its share of the card, less what the weights already + hold), as a percentage. 100 means it filled everything left for it. + alloc_shortfall_by_gpu: Of the allocation the server refused, the bytes that would + not fit -- how much to free for the block to run -- per card that + ran out. Not the size of the refused allocation itself: asking for + 2 GB with 1.9 GB free and asking for it with nothing free are the + same request and completely different problems. Set **only** on an + out-of-memory failure, so ``None`` is itself the answer to "did this + run out of memory". Approximate: the allocator frees cached blocks + and retries before giving up. + + Every field is optional, and unknown ones are kept (``extra="allow"``): a + server may report a measurement this client has never heard of, and it stays + readable as an attribute rather than being dropped. So check a field before + trusting it -- an older server may send none of them. + + GPU maps are keyed by device id **as a string**. A response reaches the + client as JSON or as ``torch.save`` bytes, and only JSON stringifies dict + keys; declaring ``str`` here means a server that sends integers fails loudly + at the boundary instead of handing the two encodings different shapes. + """ + + model_config = ConfigDict(extra="allow") + + runtime: Optional[float] = None + max_memory_usage: Optional[int] = None + max_mem_by_gpu: Dict[str, int] = Field(default_factory=dict) + max_mem_pct_by_gpu: Dict[str, float] = Field(default_factory=dict) + alloc_shortfall_by_gpu: Optional[Dict[str, int]] = None + + class Status(str, Enum): """Where a remote job is in its lifecycle (or a transient log message).""" @@ -51,7 +97,36 @@ class ResponseModel(BaseModel): status: Status description: str = "" data: Optional[Any] = None - meta_data: Optional[Dict[str, Any]] = None + meta_data: Optional[MetaData] = None + + @field_validator("meta_data", mode="before") + @classmethod + def _drop_an_unreadable_report(cls, value: Any) -> Any: + """Never let the cost report fail the response carrying it. + + ``meta_data`` is diagnostic; ``data`` is the job. A server that sends a + malformed report would otherwise raise here and turn a run that finished + perfectly well into a client-side crash, which is a far worse outcome + than not knowing what it cost. So an unreadable report is dropped. + + Dropped *loudly*. Every field is optional and unknown ones are kept, so + the only way to land here is a known field with the wrong type -- a + server bug or a version mismatch, which is worth hearing about. Silence + would also make this indistinguishable from an older server that reports + nothing at all, and those want different responses from whoever hits it. + """ + if value is None or isinstance(value, MetaData): + return value + try: + return MetaData.model_validate(value) + except ValidationError as error: + warnings.warn( + f"Discarded an unreadable cost report from the server; the " + f"job itself is unaffected. {error}", + RuntimeWarning, + stacklevel=2, + ) + return None def __str__(self) -> str: return f"[{self.id}] {self.status.value.ljust(12)} {self.description}" diff --git a/tests/test_remote_backend.py b/tests/test_remote_backend.py index bf333e8a..da12a7da 100644 --- a/tests/test_remote_backend.py +++ b/tests/test_remote_backend.py @@ -14,7 +14,7 @@ import pytest from nnsight.intervention.backends.remote import AsyncRemoteBackend, RemoteError -from nnsight.schema.response import ResponseModel, Status +from nnsight.schema.response import MetaData, ResponseModel, Status MODEL_KEY = "nnsight.modeling.transformers.TransformersModel:{}" @@ -135,6 +135,10 @@ def test_is_a_remote_backend(self): "max_mem_pct_by_gpu": {"0": 12.5, "1": 6.25}, } +# The same report once parsed. The wire carries a dict; the client hands back a +# MetaData, so this is what every assertion below compares against. +EXPECTED = MetaData(**META) + class TestResponseMetaData: """`meta_data` on the wire: it survives both encodings, and is optional.""" @@ -144,12 +148,46 @@ def test_survives_the_json_frame(self): response = ResponseModel(id="job", status=Status.COMPLETED, meta_data=META) assert ResponseModel.model_validate_json( response.model_dump_json() - ).meta_data == META + ).meta_data == EXPECTED def test_survives_the_pickled_frame(self): # Binary frames — a COMPLETED whose data is the result blob itself. response = ResponseModel(id="job", status=Status.COMPLETED, meta_data=META) - assert ResponseModel.unpickle(response.pickle()).meta_data == META + assert ResponseModel.unpickle(response.pickle()).meta_data == EXPECTED + + def test_an_unknown_field_from_a_newer_server_is_kept(self): + # extra="allow": a server may report a measurement this client has never + # heard of, and it stays readable rather than being dropped on the floor. + response = ResponseModel.model_validate_json( + '{"id": "job", "status": "COMPLETED",' + ' "meta_data": {"runtime": 1.0, "future_metric": 42}}' + ) + assert response.meta_data.runtime == 1.0 + assert response.meta_data.future_metric == 42 + + def test_an_unreadable_report_does_not_fail_the_response(self): + # meta_data is diagnostic; data is the job. A malformed report must not + # turn a run that finished perfectly well into a client-side crash. + with pytest.warns(RuntimeWarning, match="unreadable cost report"): + response = ResponseModel.model_validate_json( + '{"id": "job", "status": "COMPLETED", "data": "s3://result",' + ' "meta_data": {"runtime": "not-a-number"}}' + ) + assert response.status is Status.COMPLETED + assert response.data == "s3://result" # the job survived + assert response.meta_data is None # only the report was dropped + + def test_a_server_that_reports_nothing_warns_about_nothing(self): + # None from an older server is normal and must stay quiet -- otherwise + # the warning above stops meaning anything. + import warnings as w + + with w.catch_warnings(): + w.simplefilter("error") + response = ResponseModel.model_validate_json( + '{"id": "job", "status": "COMPLETED"}' + ) + assert response.meta_data is None def test_absent_from_an_older_server(self): # A server that doesn't report cost sends no such key; parsing must not fail. @@ -171,7 +209,7 @@ def test_recorded_when_awaited(self): [Status.RUNNING, Status.COMPLETED], result={"out": 1}, meta_data=META ) asyncio.run(backend.resolve()) - assert backend.meta_data == META + assert backend.meta_data == EXPECTED def test_recorded_when_streamed(self): # stream() bypasses note(), so it records the report on its own path. @@ -184,7 +222,7 @@ async def go(): pass asyncio.run(go()) - assert backend.meta_data == META + assert backend.meta_data == EXPECTED def test_recorded_off_a_polled_response(self): # The blocking and non-blocking paths both reach a response through note(). @@ -194,7 +232,7 @@ def test_recorded_off_a_polled_response(self): assert backend.note( ResponseModel(id="job", status=Status.COMPLETED, meta_data=META) ) - assert backend.meta_data == META + assert backend.meta_data == EXPECTED def test_intermediate_updates_leave_it_alone(self): # RUNNING carries no report; it must not clear one already recorded. @@ -203,16 +241,16 @@ def test_intermediate_updates_leave_it_alone(self): backend = RemoteBackend(MODEL_KEY, host="http://ndif.test") backend.note(ResponseModel(id="job", status=Status.COMPLETED, meta_data=META)) backend.note(ResponseModel(id="job", status=Status.RUNNING)) - assert backend.meta_data == META + assert backend.meta_data == EXPECTED def test_recorded_even_when_the_job_fails(self): # The report is taken off the response *before* note() raises. A failed # job is exactly when it earns its keep -- an OOM's meta_data carries - # extra_memory_needed, which the traceback cannot tell you. + # alloc_shortfall_by_gpu, which the traceback cannot tell you. from nnsight.intervention.backends.remote import RemoteBackend backend = RemoteBackend(MODEL_KEY, host="http://ndif.test") - failure = dict(META, extra_memory_needed={"0": 1_310_000_000}) + failure = dict(META, alloc_shortfall_by_gpu={"0": 1_310_000_000}) with pytest.raises(RemoteError, match="out of memory"): backend.note( ResponseModel( @@ -222,8 +260,8 @@ def test_recorded_even_when_the_job_fails(self): meta_data=failure, ) ) - assert backend.meta_data == failure - assert backend.meta_data["extra_memory_needed"] == {"0": 1_310_000_000} + assert backend.meta_data == MetaData(**failure) + assert backend.meta_data.alloc_shortfall_by_gpu == {"0": 1_310_000_000} def test_stays_none_against_an_older_server(self): backend = _backend([Status.RUNNING, Status.COMPLETED], result={"out": 1})