remote¶
Remote backend for executing nnsight interventions on NDIF servers.
This module provides the infrastructure for submitting intervention requests to a remote NDIF (Neural Network Distributed Inference Framework) server and receiving results back. It handles the complete lifecycle of remote job execution:
1. Serializing and submitting intervention requests via HTTP
2. Maintaining WebSocket connections for real-time status updates
3. Downloading and decompressing results
4. Displaying job progress to users (terminal and notebook)
Architecture
The remote execution flow uses a hybrid HTTP/WebSocket approach: - HTTP POST to submit the initial request - WebSocket for real-time status updates (QUEUED → RUNNING → COMPLETED) - HTTP streaming to download large results
This design allows for efficient long-polling without blocking HTTP connections while providing immediate feedback on job status changes.
Key components
- RemoteBackend: Main backend class that orchestrates remote execution
- JobStatusDisplay: Handles user-facing status output (terminal/notebook)
- LocalTracer: Tracer subclass for handling streamed remote values locally
Modes of operation
- Blocking: Wait for job completion via WebSocket (default)
- Non-blocking: Submit and poll for results separately
-
Async: Asyncio-compatible versions of blocking operations
Examples:
from nnsight import NNsight model = NNsight(model_key="openai-community/gpt2") with model.trace("Hello", remote=True): ... hidden = model.transformer.h[0].output.save() print(hidden.shape)
Manages single-line status display for remote job execution.
Provides a consistent user interface for displaying job status updates in both terminal and Jupyter notebook environments. Features include: - Animated spinners for active states (QUEUED, RUNNING) - Color-coded status indicators - Elapsed time tracking per status and total job time - In-place line updates (no scrolling spam) - Notebook-compatible HTML rendering
| PARAMETER | DESCRIPTION |
|---|---|
enabled
|
Whether to display status updates. Controlled by CONFIG.APP.REMOTE_LOGGING.
TYPE:
|
verbose
|
If True, preserve each status on its own line instead of overwriting.
TYPE:
|
Update the status display with new job information.
| PARAMETER | DESCRIPTION |
|---|---|
job_id
|
The remote job identifier.
TYPE:
|
status_name
|
Current status (RECEIVED, QUEUED, RUNNING, COMPLETED, ERROR, etc.)
TYPE:
|
description
|
Optional additional context to display.
TYPE:
|
If called with no arguments, refreshes the display with the last known status (useful for updating spinner animation and elapsed time).
RemoteBackend(model_key: str, host: Optional[str] = None, blocking: bool = True, job_id: Optional[str] = None, api_key: str = '', callback: str = '', verbose: bool = False)
Bases: Backend
Backend for executing nnsight interventions on a remote NDIF server.
This backend serializes intervention graphs and submits them to a remote service for execution on cloud-hosted models. It supports both synchronous (blocking) and asynchronous execution modes.
The execution flow
- Serialize the intervention graph via RequestModel
- Submit via HTTP POST, receive job ID
- Listen for status updates via WebSocket
- Download results when job completes
| PARAMETER | DESCRIPTION |
|---|---|
model_key
|
Identifier for the remote model (e.g., "openai-community/gpt2").
TYPE:
|
host
|
Remote server URL. Defaults to CONFIG.API.HOST.
TYPE:
|
blocking
|
If True (default), wait for job completion. If False, submit and return immediately (use job_id to poll for results later).
TYPE:
|
job_id
|
Existing job ID to retrieve results for (non-blocking mode).
TYPE:
|
api_key
|
NDIF API key. Falls back to NDIF_API_KEY env var or CONFIG.
TYPE:
|
callback
|
Optional webhook URL to receive job completion notification.
TYPE:
|
verbose
|
If True, preserve each status update on its own line.
TYPE:
|
| ATTRIBUTE | DESCRIPTION |
|---|---|
address |
HTTP address of the remote server.
TYPE:
|
ws_address |
WebSocket address (derived from HTTP address).
TYPE:
|
job_id |
Current job ID (set after submission).
TYPE:
|
job_status |
Last known job status.
TYPE:
|
compress |
Whether to use zstd compression for requests/responses.
TYPE:
|
instance-attribute
¶
status_display: JobStatusDisplay = JobStatusDisplay(enabled=REMOTE_LOGGING, verbose=verbose)
request(tracer: Tracer) -> Tuple[bytes, Dict[str, str]]
Prepare a request payload and headers for submission to the remote server.
Extracts interventions from the tracer, serializes them into a RequestModel, and builds the HTTP headers required by the NDIF API.
| PARAMETER | DESCRIPTION |
|---|---|
tracer
|
The tracer containing the intervention graph to execute.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
Tuple[bytes, Dict[str, str]]
|
Tuple of (serialized_data, headers_dict) ready for HTTP POST. |
Execute the backend, dispatching to the appropriate request mode.
Routes to async, blocking, or non-blocking execution based on the tracer's configuration and the backend's blocking setting.
| PARAMETER | DESCRIPTION |
|---|---|
tracer
|
The tracer to execute. May be None for non-blocking result retrieval.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
Optional[RESULT]
|
The execution result, or None if non-blocking and job not yet complete. |
handle_response(response: ResponseModel, tracer: Optional[Tracer] = None) -> Optional[RESULT]
Process an incoming response from the remote server.
Handles all response types: status updates, errors, completion, and streaming. Updates the status display and takes appropriate action based on job status.
| PARAMETER | DESCRIPTION |
|---|---|
response
|
The response model from the server.
TYPE:
|
tracer
|
The original tracer, needed for STREAM responses to access model info.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
Optional[RESULT]
|
For COMPLETED status: the result data (URL string or actual data). |
Optional[RESULT]
|
For other statuses: None (job still in progress). |
| RAISES | DESCRIPTION |
|---|---|
RemoteException
|
If the job status is ERROR. |
submit_request(data: bytes, headers: Dict[str, Any]) -> Optional[ResponseModel]
Submit the serialized request to the remote server via HTTP POST.
| PARAMETER | DESCRIPTION |
|---|---|
data
|
Serialized request payload (potentially compressed).
TYPE:
|
headers
|
HTTP headers including API key, version info, etc.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
Optional[ResponseModel]
|
The initial ResponseModel containing the assigned job ID. |
| RAISES | DESCRIPTION |
|---|---|
ConnectionError
|
If the server returns a non-200 status code. |
TimeoutException
|
If the request times out. |
get_response() -> Optional[RESULT]
Poll the server for the current job status (non-blocking mode).
Used when not connected via WebSocket to check if a previously submitted job has completed.
| RETURNS | DESCRIPTION |
|---|---|
Optional[RESULT]
|
The result if job is complete, None otherwise. |
| RAISES | DESCRIPTION |
|---|---|
Exception
|
If the server returns a non-200 status code. |
TimeoutException
|
If the request times out. |
get_result(url: str, content_length: Optional[float] = None) -> RESULT
Download and deserialize the result from the server.
For large results, the server returns a URL instead of inline data. This method streams the result with a progress bar, decompresses if needed, and deserializes via torch.load.
| PARAMETER | DESCRIPTION |
|---|---|
url
|
URL to download the result from (typically a presigned S3 URL).
TYPE:
|
content_length
|
Optional content length hint for progress bar.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
RESULT
|
The deserialized result object. |
async
¶
async_get_result(url: str, content_length: Optional[float] = None) -> RESULT
Async version of get_result(). See get_result() for full documentation.
Execute the intervention request and wait for completion via WebSocket.
This is the primary execution path for remote jobs. It establishes a WebSocket connection for real-time status updates while the job executes on the server.
| PARAMETER | DESCRIPTION |
|---|---|
tracer
|
The tracer containing the intervention graph.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
Optional[RESULT]
|
The execution result after the job completes. |
| RAISES | DESCRIPTION |
|---|---|
RemoteException
|
If the job fails on the server. |
async
¶
Async version of blocking_request(). See blocking_request() for full documentation.
Send computed values back to the server during hybrid execution.
When the server streams a function to execute locally (via STREAM status), local results may need to be uploaded back. This method serializes and sends those values over the WebSocket connection.
| PARAMETER | DESCRIPTION |
|---|---|
values
|
Dictionary of values to send (keyed by intervention ID).
TYPE:
|
sio
|
The active WebSocket client connection.
TYPE:
|
Submit a job or poll for results without blocking.
This mode allows submitting a job and retrieving results in separate calls, useful for long-running jobs or when you want to do other work while waiting.
First call (job_id is None): Submits the request and stores the job_id. Subsequent calls (job_id is set): Polls for completion and returns result.
| PARAMETER | DESCRIPTION |
|---|---|
tracer
|
The tracer to execute (used only on first call).
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
Optional[RESULT]
|
None on first call (job submitted). |
Optional[RESULT]
|
The result on subsequent calls if job is complete, None if still running. |
Bases: Tracer
Tracer subclass for executing streamed functions locally during hybrid execution.
When the server streams a function to execute on the client (via STREAM response), LocalTracer handles the local execution and sends results back to the server.
This enables hybrid execution patterns where some computations happen locally (e.g., on user's GPU) while others happen remotely.
Class Attributes
_send: Callback function to send values back to the server (set via register()).
classmethod
¶
Register the send callback for uploading values to the server.
Execute a streamed function with remote value tracking.
| PARAMETER | DESCRIPTION |
|---|---|
fn
|
The function streamed from the server to execute locally.
TYPE:
|
Push local state and send remote-marked values back to server.
Inspects the caller's local variables, pushes them to the parent tracer, then filters for remote-marked objects and sends them to the server.
Pull the NDIF environment information from the remote server, and register any locally-available modules not present remotely.