remote

remote

Running a captured trace on a remote NDIF server.

A trace captured locally can be shipped to NDIF and run against a model that never leaves the server: the block is serialized source-and-all, POSTed, and the model is named by a model_key rather than pickled. The backends here differ only in how the client waits for the job to finish and collect its saves:

• RemoteBackend blocking -- one /subscribe websocket, streamed status updates, result on COMPLETED (model.trace(..., remote=True)).
• RemoteBackend non-blocking (blocking=False) -- submit over HTTP and poll the server's stored status until the result lands.
• AsyncRemoteBackend -- the blocking stream, awaited on an asyncio event loop so a job runs without tying up a thread.

Status updates render through StatusDisplay; a failed job surfaces as RemoteError.

RemoteError

Bases: Exception

The remote job failed on the server.

RemoteBackend

RemoteBackend(model_key: str, host: Optional[str] = None, api_key: Optional[str] = None, env: Optional[dict] = None, blocking: bool = True, job_id: Optional[str] = None, verbose: bool = False)

Bases: Backend

Run a trace on a remote NDIF server, waiting for it in the calling thread.

Serializes the captured block (the model is named by model_key, never serialized), ships it, and returns the saved values once the job finishes. Two waiting modes, chosen by blocking:

• blocking=True (default): hold one /subscribe websocket open, stream status updates until COMPLETED, then download the saves and push them back into the caller's frame so its h = ...save() variables populate.
• blocking=False: submit the job over plain HTTP (no websocket) and store its job_id; the server records each status to the object store, and each later call poll\ s until the result is ready.

AsyncRemoteBackend runs the same blocking stream on an asyncio event loop instead.

model_key instance-attribute

model_key = model_key

env instance-attribute

env = env or {}

blocking instance-attribute

blocking = blocking

job_id instance-attribute

job_id = job_id

status instance-attribute

status: Optional[Status] = None

host instance-attribute

host = host or CONFIG.API.HOST

api_key instance-attribute

api_key = api_key or CONFIG.API.APIKEY or ''

compress instance-attribute

compress = CONFIG.API.COMPRESS

ws_host instance-attribute

ws_host = ('wss' if scheme == 'https' else 'ws') + '://' + rest

verbose instance-attribute

verbose = verbose or CONFIG.APP.DEBUG

display instance-attribute

display = StatusDisplay(enabled=CONFIG.APP.REMOTE_LOGGING or self.verbose, verbose=self.verbose)

__call__

__call__(tracer: Optional[Tracer] = None) -> Optional[RESULT]

parse staticmethod

parse(message: Union[str, bytes]) -> ResponseModel

Build a ResponseModel from one websocket frame.

A status update arrives as JSON in a text frame. A COMPLETED response may instead arrive as a binary frame — torch.save bytes whose data is the result blob itself, sent when the server hands the result back on the socket rather than staging it for download.

note

note(response: ResponseModel) -> bool

Render a status update, raise on ERROR, and report whether it's final.

The shared status handling for every update, however it arrived (websocket, poll): update the display, raise RemoteError on ERROR, and return whether the status is COMPLETED (so the caller knows to fetch the result).

handle

handle(response: ResponseModel) -> Optional[RESULT]

Process a single status update, returning the result on COMPLETED.

Returns None for intermediate statuses. On COMPLETED data is either the result blob itself — the server sent it back on the response — or a presigned url to download it from. Which one depends on the result's size against the server's limit, so both have to be handled.

download_result

download_result(url: Optional[str]) -> Optional[RESULT]

Deserialize the result, downloading it first if it is not already here.

url is whatever COMPLETED carried: a presigned url to fetch, or the result blob itself when the server sent it back on the response rather than staging it. A url is downloaded in chunks behind a tqdm progress bar (shown when remote logging is enabled); either way the bytes go to finalize, which decompresses under the same flag the request was compressed with and loads it with torch.load.

finalize

finalize(content: bytes) -> RESULT

Decompress (if the request was compressed) and load a result blob.

The shared tail of every route the bytes can arrive by — downloaded from a presigned url, or handed over on the response itself: everything after the bytes are in hand.

send

send(request: RequestModel, blob: bytes) -> Optional[RESULT]

Submit the request and handle the initial (blocking) response.

submit

submit(tracer: Tracer) -> None

Submit a non-blocking job: POST without a websocket and store the job id.

No /subscribe, so the request carries no session id; the server saves each status response to the object store for poll to read. Returns None — call poll (or the backend again) to fetch the result.

poll

poll(tracer: Optional[Tracer] = None) -> Optional[RESULT]

Fetch the latest status of a submitted non-blocking job.

GETs /response/{job_id}. Returns the saved-values dict on COMPLETED (the trace has long since exited, so the caller reads it from the return rather than from any frame), raises on ERROR, and returns None while the job is still running (or before its first status lands). status holds the latest.

request

request(request: RequestModel, tracer: Tracer) -> Optional[RESULT]

AsyncRemoteBackend

AsyncRemoteBackend(*args, **kwargs)

Bases: RemoteBackend

A RemoteBackend whose waiting is async, over the same websocket.

__call__ fires the request the way the blocking parent does — subscribe to the /subscribe websocket, take the session id, POST the payload — and returns without consuming any status updates. That initial connection is synchronous; only the waiting is async. Await the backend for the saves dict — which renders the status display and raises on a server error, like the blocking parent — or async-iterate it for the raw status updates to handle yourself, with the saves dict yielded as the final item:

backend = AsyncRemoteBackend(model.to_model_key())
with model.trace(prompt, backend=backend):
    out = model.output.save()

result = await backend                  # wait for COMPLETED, get the saves
# or
async for update in backend:            # raw ResponseModel status updates...
    if isinstance(update, dict):
        result = update                 # ...and the saves dict, yielded last
    else:
        print(update.status)            # you decide what to do with each

The websocket recv is blocking, so receive runs it through asyncio.to_thread to keep the event loop free while a job runs. Only that waiting differs from the parent — the connect, serialization, POST, and the decompress/load reuse the parent's synchronous methods.

The caller's frame is gone by the time the result lands (the trace has exited), so — like the parent's non-blocking path — the saved values aren't pushed back into it; they come out of the await / the iterator's final item.

connection instance-attribute

connection: Optional[object] = None

__call__

__call__(tracer: Optional[Tracer] = None) -> 'AsyncRemoteBackend'

__await__

__await__()

resolve async

resolve() -> Optional[RESULT]

Consume status updates until COMPLETED, returning the saves dict.

The trace has exited, so — like the non-blocking poll — nothing is pushed back into a frame; the caller reads the returned dict.

__aiter__

__aiter__() -> AsyncIterator[Any]

stream async

stream() -> AsyncIterator[Any]

Yield each raw status update as it lands, then the saves dict last.

Unlike resolve, this doesn't touch the display or raise on ERROR — it hands you each ResponseModel as-is to do with as you like. The final item (after COMPLETED) is the downloaded saves dict; an ERROR update just ends the stream (inspect it and raise yourself if you want).

receive async

receive() -> ResponseModel

Await the next status update off the websocket (blocking recv in a thread).

download async

download(url: Optional[str]) -> Optional[RESULT]

Async download_result: stream the blob with an async client, then hand it to the parent's shared decompress/load.

url is whatever COMPLETED carried: a presigned url to fetch, or the result blob itself when the server sent it back on the response. Both resolve and stream reach the result through here rather than through handle, so the bytes case is answered here too.

close

close() -> None

Close the subscription websocket (idempotent).