From 99d8ba49e49794cb865ea5646620afa151105108 Mon Sep 17 00:00:00 2001
From: Jaden Fiotto-Kaufman <jadenfk@outlook.com>
Date: Tue, 2 Jan 2024 15:46:06 -0500
Subject: [PATCH] Progress bar for remote steaming result download.

---
 src/nnsight/contexts/Runner.py    | 27 +++++++++++++++++++++++----
 src/nnsight/pydantics/__init__.py |  4 ++--
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/src/nnsight/contexts/Runner.py b/src/nnsight/contexts/Runner.py
index cd53714e..57cc30c9 100644
--- a/src/nnsight/contexts/Runner.py
+++ b/src/nnsight/contexts/Runner.py
@@ -1,9 +1,11 @@
 from __future__ import annotations
 
+import io
 import pickle
 
 import requests
 import socketio
+from tqdm import tqdm
 
 from .. import CONFIG, pydantics
 from ..logger import logger
@@ -118,16 +120,33 @@ def blocking_response(data):
 
             # If the status of the response is completed, update the local nodes that the user specified to save.
             # Then disconnect and continue.
+
             if response.status == pydantics.ResponseModel.JobStatus.COMPLETED:
+                result_bytes = io.BytesIO()
+                result_bytes.seek(0)
+
                 with requests.get(
-                    url=f"https://{CONFIG.API.HOST}/retrieve/{response.id}", stream=True
+                    url=f"https://{CONFIG.API.HOST}/result/{response.id}", stream=True
                 ) as stream:
-                    result_response = pydantics.ResponseModel(**pickle.load(stream.raw))
+                    total_size = float(stream.headers["Content-length"])
+
+                    with tqdm(
+                        total=total_size, unit="B", unit_scale=True
+                    ) as progress_bar:
+                        for data in stream.iter_content(chunk_size=4000000):
+                            progress_bar.update(len(data))
+                            result_bytes.write(data)
+
+                result_bytes.seek(0)
+
+                result = pydantics.ResultModel(**pickle.load(result_bytes))
+
+                result_bytes.close()
 
-                for name, value in result_response.result.saves.items():
+                for name, value in result.saves.items():
                     self.graph.nodes[name].value = value
 
-                self.output = result_response.result.output
+                self.output = result.output
 
                 sio.disconnect()
             # Or if there was some error.
diff --git a/src/nnsight/pydantics/__init__.py b/src/nnsight/pydantics/__init__.py
index 425dee05..fb5cbacb 100644
--- a/src/nnsight/pydantics/__init__.py
+++ b/src/nnsight/pydantics/__init__.py
@@ -1,3 +1,3 @@
-from .Response import ResponseModel
+from .Config import ConfigModel
 from .Request import RequestModel
-from .Config import ConfigModel
\ No newline at end of file
+from .Response import ResponseModel, ResultModel