Remote Job Metadata - #708
Open
AdamBelfki3 wants to merge 3 commits into
Open
AdamBelfki3 wants to merge 3 commits into
AdamBelfki3 wants to merge 3 commits into
Conversation
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.
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.
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.
koriyoshi2041
left a comment
koriyoshi2041
left a comment
There was a problem hiding this comment.
There is a stale-report edge when a caller reuses an explicit blocking backend across traces. Remotable.trace(..., backend=...) accepts the backend as a public argument, but note() only assigns self.meta_data when the new response has a report.
The lifecycle reduces to:
backend = RemoteBackend(MODEL_KEY, host="http://ndif.test")
backend.note(ResponseModel(id="first", status=Status.COMPLETED, meta_data=META))
backend.note(ResponseModel(id="second", status=Status.COMPLETED)) # older server / no report
assert backend.meta_data is None # currently still META from "first"That can make the second job appear to have the first job's runtime/GPU footprint. The existing intermediate-update test correctly requires RUNNING without metadata not to clear the current job's eventual report, so clearing in every note() would be too broad. Could the start of a new blocking request reset meta_data, with a two-request regression that preserves the intermediate-frame behavior?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
meta_datatoResponseModel: what a remote job cost the server, reportedon the response that ends it.
The tracer holds the backend, so no new API is needed on the common path. Read
it after the block exits — the backend runs in
__exit__.What's here
MetaDatamodel inschema/response.py:runtime,max_memory_usage,max_mem_by_gpu,max_mem_pct_by_gpu, andalloc_shortfall_by_gpu.RemoteBackend.meta_data, populated innote()— the per-frame handling theblocking trace, the non-blocking poll and
await backendall share.stream()bypassesnote()by design, so it records on its own path.note()raises, so a failed job's report survives and canbe read from
except RemoteError:. On an out-of-memory failure that reportcarries
alloc_shortfall_by_gpu, which is the useful part.Testing
19 tests in
tests/test_remote_backend.pycovering both wire encodings (JSONand
torch.save), all four waiting modes, the failure path, unknown fields, andthe unreadable-report warning. Exercised end to end against a local NDIF stack
with a real GPU from the paired ndif branch.