Split submit_request from handle_response; add async submit + get_response

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