Remote Job Metadata

docs/remote/ndif-overview.md

@@ -77,6 +77,41 @@ 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()
+
+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 `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:
+
+```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.alloc_shortfall_by_gpu)   # {'0': 1310000000}
+```
+
+`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.
+
 ## 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.

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,6 +95,7 @@ def __init__(
         self.blocking = blocking
         self.job_id = job_id
         self.status: Optional[Status] = None
+        self.meta_data: Optional[MetaData] = None
         self.host = host or CONFIG.API.HOST
         if not self.host.startswith(("http://", "https://")):
             raise ValueError(
@@ -169,8 +170,19 @@ 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)
+        # 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
@@ -476,6 +488,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)

src/nnsight/schema/response.py

@@ -3,21 +3,69 @@
 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
 
 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)."""
 
@@ -37,7 +85,10 @@ 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 -- 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=())
@@ -46,6 +97,36 @@ class ResponseModel(BaseModel):
     status: Status
     description: str = ""
     data: Optional[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}"

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:{}"
 
@@ -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,147 @@ 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},
+}
+
+# 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."""
+
+    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 == 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 == 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.
+        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 == EXPECTED
+
+    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 == EXPECTED
+
+    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 == EXPECTED
+
+    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 == 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
+        # 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, alloc_shortfall_by_gpu={"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 == 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})
+        asyncio.run(backend.resolve())
+        assert backend.meta_data is None