← back to Exo
Download Progression in tiny chat
39a8a814035171dff95f53f9c94847b46eced33a · 2024-10-01 16:40:10 +0200 · DevEmilio96
Files touched
M exo/api/chatgpt_api.pyM main.pyM tinychat/examples/tinychat/index.cssM tinychat/examples/tinychat/index.htmlM tinychat/examples/tinychat/index.js
Diff
commit 39a8a814035171dff95f53f9c94847b46eced33a
Author: DevEmilio96 <minima998@gmail.com>
Date: Tue Oct 1 16:40:10 2024 +0200
Download Progression in tiny chat
---
exo/api/chatgpt_api.py | 17 ++++++++++
main.py | 19 +++++++----
tinychat/examples/tinychat/index.css | 4 ++-
tinychat/examples/tinychat/index.html | 15 +++++++++
tinychat/examples/tinychat/index.js | 62 +++++++++++++++++++++++++++++++++++
5 files changed, 110 insertions(+), 7 deletions(-)
diff --git a/exo/api/chatgpt_api.py b/exo/api/chatgpt_api.py
index ef97e5c1..3af2db72 100644
--- a/exo/api/chatgpt_api.py
+++ b/exo/api/chatgpt_api.py
@@ -9,6 +9,7 @@ from aiohttp import web
import aiohttp_cors
import traceback
from exo import DEBUG, VERSION
+from exo.download.download_progress import RepoProgressEvent
from exo.helpers import PrefixDict
from exo.inference.shard import Shard
from exo.inference.tokenizers import resolve_tokenizer
@@ -175,6 +176,8 @@ class ChatGPTAPI:
cors.add(self.app.router.add_post("/v1/chat/token/encode", self.handle_post_chat_token_encode), {"*": cors_options})
cors.add(self.app.router.add_post("/chat/completions", self.handle_post_chat_completions), {"*": cors_options})
cors.add(self.app.router.add_post("/v1/chat/completions", self.handle_post_chat_completions), {"*": cors_options})
+ # Endpoint for download progress tracking
+ cors.add(self.app.router.add_get("/v1/download/progress", self.handle_get_download_progress), {"*": cors_options})
self.static_dir = Path(__file__).parent.parent.parent/"tinychat/examples/tinychat"
self.app.router.add_get("/", self.handle_root)
@@ -203,6 +206,20 @@ class ChatGPTAPI:
tokenizer = await resolve_tokenizer(shard.model_id)
return web.json_response({"length": len(build_prompt(tokenizer, messages)[0])})
+ async def handle_get_download_progress(self, request):
+ progress_data = {}
+ for node_id, progress_event in self.node.node_download_progress.items():
+ if isinstance(progress_event, RepoProgressEvent):
+ # Convert to dict if not already
+ progress_data[node_id] = progress_event.to_dict()
+ elif isinstance(progress_event, dict):
+ progress_data[node_id] = progress_event
+ else:
+ # Handle unexpected types
+ progress_data[node_id] = str(progress_event)
+ return web.json_response(progress_data)
+
+
async def handle_post_chat_completions(self, request):
data = await request.json()
if DEBUG >= 2: print(f"Handling chat completions request from {request.remote}: {data}")
diff --git a/main.py b/main.py
index 6132bc67..17e25c22 100644
--- a/main.py
+++ b/main.py
@@ -87,6 +87,8 @@ node = StandardNode(
)
server = GRPCServer(node, args.node_host, args.node_port)
node.server = server
+node.download_progress = {} # Initialize download progress tracking
+node.node_download_progress = {} # For tracking per-node download progress
api = ChatGPTAPI(
node,
inference_engine.__class__.__name__,
@@ -116,12 +118,17 @@ if args.prometheus_client_port:
last_broadcast_time = 0
def throttled_broadcast(shard: Shard, event: RepoProgressEvent):
- global last_broadcast_time
- current_time = time.time()
- if event.status == "complete" or current_time - last_broadcast_time >= 0.1:
- last_broadcast_time = current_time
- asyncio.create_task(node.broadcast_opaque_status("", json.dumps({"type": "download_progress", "node_id": node.id, "progress": event.to_dict()})))
-
+ global last_broadcast_time
+ current_time = time.time()
+ if event.status == "complete" or current_time - last_broadcast_time >= 0.1:
+ last_broadcast_time = current_time
+ node.download_progress[event.repo_id] = event.to_dict()
+ node.node_download_progress[node.id] = event.to_dict()
+ asyncio.create_task(node.broadcast_opaque_status("", json.dumps({
+ "type": "download_progress",
+ "node_id": node.id,
+ "progress": event.to_dict()
+ })))
shard_downloader.on_progress.register("broadcast").on_next(throttled_broadcast)
diff --git a/tinychat/examples/tinychat/index.css b/tinychat/examples/tinychat/index.css
index c0f9d9d5..2129a72f 100644
--- a/tinychat/examples/tinychat/index.css
+++ b/tinychat/examples/tinychat/index.css
@@ -164,7 +164,9 @@ main {
border-right: 2px solid var(--secondary-color);
box-shadow: 10px 10px 20px 2px var(--secondary-color-transparent);
}
-
+.download-progress{
+ margin-bottom: 20em;
+}
.message > pre {
white-space: pre-wrap;
}
diff --git a/tinychat/examples/tinychat/index.html b/tinychat/examples/tinychat/index.html
index 9a5e07c2..0028b15b 100644
--- a/tinychat/examples/tinychat/index.html
+++ b/tinychat/examples/tinychat/index.html
@@ -149,6 +149,21 @@
$el.scrollTo({ top: $el.scrollHeight, behavior: 'smooth' });
" x-ref="messages" x-show="home === 2" x-transition="">
</div>
+
+<!-- Download Progress Section -->
+<template x-if="downloadProgress">
+<div class="download-progress message message-role-assistant">
+ <h2>Download Progress</h2>
+ <div class="download-progress-node">
+ <p><strong>Model:</strong> <span x-text="downloadProgress.repo_id + '@' + downloadProgress.repo_revision"></span></p>
+ <p><strong>Progress:</strong> <span x-text="`${downloadProgress.downloaded_bytes_display} / ${downloadProgress.total_bytes_display} (${downloadProgress.percentage}%)`"></span></p>
+ <p><strong>Speed:</strong> <span x-text="downloadProgress.overall_speed_display || 'N/A'"></span></p>
+ <p><strong>ETA:</strong> <span x-text="downloadProgress.overall_eta_display || 'N/A'"></span></p>
+ </div>
+</div>
+</template>
+
+
<div class="input-container">
<div class="input-performance">
<span class="input-performance-point">
diff --git a/tinychat/examples/tinychat/index.js b/tinychat/examples/tinychat/index.js
index 668e2012..bf3b831f 100644
--- a/tinychat/examples/tinychat/index.js
+++ b/tinychat/examples/tinychat/index.js
@@ -23,6 +23,9 @@ document.addEventListener("alpine:init", () => {
// image handling
imagePreview: null,
+ // download progress
+ downloadProgress: null,
+
removeHistory(cstate) {
const index = this.histories.findIndex((state) => {
return state.time === cstate.time;
@@ -32,6 +35,23 @@ document.addEventListener("alpine:init", () => {
localStorage.setItem("histories", JSON.stringify(this.histories));
}
},
+ // Utility functions
+ formatBytes(bytes) {
+ if (bytes === 0) return '0 B';
+ const k = 1024;
+ const sizes = ['B', 'KB', 'MB', 'GB', 'TB'];
+ const i = Math.floor(Math.log(bytes) / Math.log(k));
+ return parseFloat((bytes / Math.pow(k, i)).toFixed(2)) + ' ' + sizes[i];
+ },
+
+ formatDuration(seconds) {
+ const h = Math.floor(seconds / 3600);
+ const m = Math.floor((seconds % 3600) / 60);
+ const s = Math.floor(seconds % 60);
+ if (h > 0) return `${h}h ${m}m ${s}s`;
+ if (m > 0) return `${m}m ${s}s`;
+ return `${s}s`;
+ },
async handleImageUpload(event) {
const file = event.target.files[0];
@@ -254,6 +274,48 @@ document.addEventListener("alpine:init", () => {
}
}
},
+
+ async fetchDownloadProgress() {
+ try {
+ console.log("fetching download progress");
+ const response = await fetch(`${this.endpoint}/download/progress`);
+ if (response.ok) {
+ const data = await response.json();
+ // Since we're only interested in our own node's progress, extract it
+ // Assuming the server returns progress only for the local node
+ const progressArray = Object.values(data);
+ if (progressArray.length > 0) {
+ const progress = progressArray[0];
+ // Check if download is complete
+ if (progress.status === "complete" || progress.status === "failed") {
+ this.downloadProgress = null; // Hide the progress section
+ } else {
+ // Compute human-readable strings
+ progress.downloaded_bytes_display = this.formatBytes(progress.downloaded_bytes);
+ progress.total_bytes_display = this.formatBytes(progress.total_bytes);
+ progress.overall_speed_display = progress.overall_speed ? this.formatBytes(progress.overall_speed) + '/s' : '';
+ progress.overall_eta_display = progress.overall_eta ? this.formatDuration(progress.overall_eta) : '';
+ progress.percentage = ((progress.downloaded_bytes / progress.total_bytes) * 100).toFixed(2);
+
+ this.downloadProgress = progress;
+ }
+ } else {
+ // No ongoing download
+ this.downloadProgress = null;
+ }
+ }
+ } catch (error) {
+ console.error("Error fetching download progress:", error);
+ this.downloadProgress = null;
+ }
+ },
+ // Initialize the component
+ init() {
+ // Start polling for download progress every second
+ setInterval(() => {
+ this.fetchDownloadProgress();
+ }, 1000);
+ },
}));
});
← a7a9124d Hardware Requirements
·
back to Exo
·
download progress optimization in tinychat bef585d1 →