From 03afb473f05347a36fc77bbe4f8ff1011d1026cf Mon Sep 17 00:00:00 2001 From: Jaden Fiotto-Kaufman Date: Mon, 20 Apr 2026 17:04:33 -0400 Subject: [PATCH] Split submit_request from handle_response; add async variants submit_request now does only the POST and returns the initial ResponseModel (with self.job_id recorded). Dispatching the response to handle_response is left to the caller so that async / streaming consumers can yield the initial response before any side effects (display updates, RemoteException raises) run. Add async_submit_request and async_get_response using httpx.AsyncClient so callers inside an event loop don't need a threadpool or a duplicated POST path. The new parsing helper _parse_submit_response is shared between the sync and async submits. The three internal callers (blocking_request, async_request, non_blocking_request) now invoke handle_response explicitly on the initial response to preserve today's behaviour. --- src/nnsight/intervention/backends/remote.py | 91 ++++++++++++++++----- 1 file changed, 71 insertions(+), 20 deletions(-) diff --git a/src/nnsight/intervention/backends/remote.py b/src/nnsight/intervention/backends/remote.py index 630ff882..14033b3f 100755 --- a/src/nnsight/intervention/backends/remote.py +++ b/src/nnsight/intervention/backends/remote.py @@ -617,11 +617,37 @@ def handle_response( local_tracer.execute(fn) + def _parse_submit_response(self, response: httpx.Response) -> ResponseModel: + """Parse an HTTP POST /request response into a ResponseModel. + + Records ``self.job_id`` for subsequent status checks. Does not call + :meth:`handle_response` — callers decide when to dispatch status + updates so that async / streaming paths can yield the initial + response before any side effects run. + """ + from ...schema.response import ResponseModel + + if response.status_code != 200: + try: + msg = response.json()["detail"] + except Exception: + msg = response.reason_phrase + raise ConnectionError(msg) + + response_model = ResponseModel(**response.json()) + self.job_id = response_model.id + return response_model + def submit_request( self, data: bytes, headers: Dict[str, Any] - ) -> Optional[ResponseModel]: + ) -> ResponseModel: """Submit the serialized request to the remote server via HTTP POST. + Returns the initial :class:`ResponseModel` (with the assigned job id + stored on ``self.job_id``). The returned response has *not* been + passed through :meth:`handle_response`; the caller is responsible + for dispatching it when ready. + Args: data: Serialized request payload (potentially compressed). headers: HTTP headers including API key, version info, etc. @@ -633,10 +659,7 @@ def submit_request( ConnectionError: If the server returns a non-200 status code. httpx.TimeoutException: If the request times out. """ - from ...schema.response import ResponseModel - headers["Content-Type"] = "application/octet-stream" - timeout = httpx.Timeout(self.CONNECT_TIMEOUT, read=self.READ_TIMEOUT) with httpx.Client(timeout=timeout) as client: @@ -646,22 +669,29 @@ def submit_request( headers=headers, ) - if response.status_code == 200: - response_model = ResponseModel(**response.json()) + return self._parse_submit_response(response) - # Store job ID for subsequent status checks - self.job_id = response_model.id + async def async_submit_request( + self, data: bytes, headers: Dict[str, Any] + ) -> ResponseModel: + """Async version of :meth:`submit_request`. - self.handle_response(response_model) + Uses :class:`httpx.AsyncClient` so that callers inside an event loop + don't block the thread. See :meth:`submit_request` for the return + contract — the caller is responsible for invoking + :meth:`handle_response` on the returned model. + """ + headers["Content-Type"] = "application/octet-stream" + timeout = httpx.Timeout(self.CONNECT_TIMEOUT, read=self.READ_TIMEOUT) - return response_model - else: - # Extract error message from response - try: - msg = response.json()["detail"] - except Exception: - msg = response.reason_phrase - raise ConnectionError(msg) + async with httpx.AsyncClient(timeout=timeout) as client: + response = await client.post( + f"{self.address}/request", + content=data, + headers=headers, + ) + + return self._parse_submit_response(response) def get_response(self) -> Optional[RESULT]: """Poll the server for the current job status (non-blocking mode). @@ -692,6 +722,24 @@ def get_response(self) -> Optional[RESULT]: else: raise Exception(response.reason_phrase) + async def async_get_response(self) -> Optional[RESULT]: + """Async version of :meth:`get_response`.""" + from ...schema.response import ResponseModel + + timeout = httpx.Timeout(self.CONNECT_TIMEOUT, read=self.READ_TIMEOUT) + + async with httpx.AsyncClient(timeout=timeout) as client: + response = await client.get( + f"{self.address}/response/{self.job_id}", + headers={"ndif-api-key": self.api_key}, + ) + + if response.status_code == 200: + response_model = ResponseModel(**response.json()) + return self.handle_response(response_model) + else: + raise Exception(response.reason_phrase) + def _decompress_and_load(self, result_bytes: io.BytesIO) -> RESULT: """Decompress (if needed) and deserialize result bytes. @@ -846,7 +894,8 @@ def blocking_request(self, tracer: Tracer) -> Optional[RESULT]: # Prepare and submit the request data, headers = self.request(tracer) headers["ndif-session_id"] = sio.sid # Link WebSocket to this request - self.submit_request(data, headers) + initial = self.submit_request(data, headers) + self.handle_response(initial) try: # Register callback for streaming values back to server @@ -891,7 +940,8 @@ async def async_request(self, tracer: Tracer) -> Optional[RESULT]: data, headers = self.request(tracer) headers["ndif-session_id"] = sio.sid - self.submit_request(data, headers) + initial = await self.async_submit_request(data, headers) + self.handle_response(initial) try: LocalTracer.register(lambda data: self.stream_send(data, sio)) @@ -956,7 +1006,8 @@ def non_blocking_request(self, tracer: Tracer) -> Optional[RESULT]: if self.job_id is None: # First call: submit the job data, headers = self.request(tracer) - self.submit_request(data, headers) + initial = self.submit_request(data, headers) + self.handle_response(initial) # job_id is set by submit_request else: # Subsequent calls: poll for result