Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Dev #37

Merged
merged 4 commits into from
Jan 2, 2024
Merged

Dev #37

Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/source/_templates/ndif_status.html
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
<script>
let statusIcon = document.querySelector('.ndif .fa-circle-check');
var statusIcon = document.querySelector('.ndif .fa-circle-check');
fetch("{{ndif_url}}")
.then((response) => {
if (response.status == 200) {
Expand Down
28 changes: 24 additions & 4 deletions src/nnsight/contexts/Runner.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -118,16 +120,34 @@ 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.
Expand Down
Loading