← back to Exo
topology endpoint to get current topology
67e05c87ce4ce1f4bc75741c8c0b9dbc8c4b65d0 · 2024-12-06 18:31:23 +0000 · Alex Cheema
Files touched
M exo/api/chatgpt_api.pyM exo/topology/topology.py
Diff
commit 67e05c87ce4ce1f4bc75741c8c0b9dbc8c4b65d0
Author: Alex Cheema <alexcheema123@gmail.com>
Date: Fri Dec 6 18:31:23 2024 +0000
topology endpoint to get current topology
---
exo/api/chatgpt_api.py | 47 +++++++++++++++++++++++++++++++----------------
exo/topology/topology.py | 20 ++++++++++++++++++++
2 files changed, 51 insertions(+), 16 deletions(-)
diff --git a/exo/api/chatgpt_api.py b/exo/api/chatgpt_api.py
index fbf328bc..4d4c59e3 100644
--- a/exo/api/chatgpt_api.py
+++ b/exo/api/chatgpt_api.py
@@ -182,6 +182,7 @@ class ChatGPTAPI:
cors.add(self.app.router.add_get("/initial_models", self.handle_get_initial_models), {"*": cors_options})
cors.add(self.app.router.add_post("/create_animation", self.handle_create_animation), {"*": cors_options})
cors.add(self.app.router.add_post("/download", self.handle_post_download), {"*": cors_options})
+ cors.add(self.app.router.add_get("/topology", self.handle_get_topology), {"*": cors_options})
if "__compiled__" not in globals():
self.static_dir = Path(__file__).parent.parent/"tinychat"
@@ -232,11 +233,11 @@ class ChatGPTAPI:
}
)
await response.prepare(request)
-
+
for model_name, pretty in pretty_name.items():
if model_name in model_cards:
model_info = model_cards[model_name]
-
+
if self.inference_engine_classname in model_info.get("repo", {}):
shard = build_base_shard(model_name, self.inference_engine_classname)
if shard:
@@ -244,11 +245,11 @@ class ChatGPTAPI:
downloader.current_shard = shard
downloader.current_repo_id = get_repo(shard.model_id, self.inference_engine_classname)
status = await downloader.get_shard_download_status()
-
+
download_percentage = status.get("overall") if status else None
total_size = status.get("total_size") if status else None
total_downloaded = status.get("total_downloaded") if status else False
-
+
model_data = {
model_name: {
"name": pretty,
@@ -258,17 +259,17 @@ class ChatGPTAPI:
"total_downloaded": total_downloaded
}
}
-
+
await response.write(f"data: {json.dumps(model_data)}\n\n".encode())
-
+
await response.write(b"data: [DONE]\n\n")
return response
-
+
except Exception as e:
print(f"Error in handle_model_support: {str(e)}")
traceback.print_exc()
return web.json_response(
- {"detail": f"Server error: {str(e)}"},
+ {"detail": f"Server error: {str(e)}"},
status=500
)
@@ -425,35 +426,35 @@ class ChatGPTAPI:
try:
model_name = request.match_info.get('model_name')
if DEBUG >= 2: print(f"Attempting to delete model: {model_name}")
-
+
if not model_name or model_name not in model_cards:
return web.json_response(
- {"detail": f"Invalid model name: {model_name}"},
+ {"detail": f"Invalid model name: {model_name}"},
status=400
)
shard = build_base_shard(model_name, self.inference_engine_classname)
if not shard:
return web.json_response(
- {"detail": "Could not build shard for model"},
+ {"detail": "Could not build shard for model"},
status=400
)
repo_id = get_repo(shard.model_id, self.inference_engine_classname)
if DEBUG >= 2: print(f"Repo ID for model: {repo_id}")
-
+
# Get the HF cache directory using the helper function
hf_home = get_hf_home()
cache_dir = get_repo_root(repo_id)
-
+
if DEBUG >= 2: print(f"Looking for model files in: {cache_dir}")
-
+
if os.path.exists(cache_dir):
if DEBUG >= 2: print(f"Found model files at {cache_dir}, deleting...")
try:
shutil.rmtree(cache_dir)
return web.json_response({
- "status": "success",
+ "status": "success",
"message": f"Model {model_name} deleted successfully",
"path": str(cache_dir)
})
@@ -465,7 +466,7 @@ class ChatGPTAPI:
return web.json_response({
"detail": f"Model files not found at {cache_dir}"
}, status=404)
-
+
except Exception as e:
print(f"Error in handle_delete_model: {str(e)}")
traceback.print_exc()
@@ -543,6 +544,20 @@ class ChatGPTAPI:
if DEBUG >= 2: traceback.print_exc()
return web.json_response({"error": str(e)}, status=500)
+ async def handle_get_topology(self, request):
+ try:
+ topology = self.node.current_topology
+ if topology:
+ return web.json_response(topology.to_json())
+ else:
+ return web.json_response({})
+ except Exception as e:
+ if DEBUG >= 2: traceback.print_exc()
+ return web.json_response(
+ {"detail": f"Error getting topology: {str(e)}"},
+ status=500
+ )
+
async def run(self, host: str = "0.0.0.0", port: int = 52415):
runner = web.AppRunner(self.app)
await runner.setup()
diff --git a/exo/topology/topology.py b/exo/topology/topology.py
index a002cdc1..e6eee05c 100644
--- a/exo/topology/topology.py
+++ b/exo/topology/topology.py
@@ -51,3 +51,23 @@ class Topology:
edges_str = ", ".join(f"{node}: {[f'{c.to_id}({c.description})' for c in conns]}"
for node, conns in self.peer_graph.items())
return f"Topology(Nodes: {{{nodes_str}}}, Edges: {{{edges_str}}})"
+
+ def to_json(self):
+ return {
+ "nodes": {
+ node_id: capabilities.to_dict()
+ for node_id, capabilities in self.nodes.items()
+ },
+ "peer_graph": {
+ node_id: [
+ {
+ "from_id": conn.from_id,
+ "to_id": conn.to_id,
+ "description": conn.description
+ }
+ for conn in connections
+ ]
+ for node_id, connections in self.peer_graph.items()
+ },
+ "active_node_id": self.active_node_id
+ }
← e8ece115 tweak sed, make compile_grpc.sh executable
·
back to Exo
·
add llama-3.3-70b 9d1f14a2 →