Upload 8 files

app.py

@@ -1,187 +1,51 @@
-from threading import Thread
-import gradio as gr
-import inspect
-from gradio import routes
-from typing import List, Type
+from functools import partial
 
-import requests, os, re, asyncio, queue, sys, git
-import math
-import time
-import datetime
-import requests, json
-
-from pprint import pprint
 import hivemind
-from petals.constants import PUBLIC_INITIAL_PEERS
-from health import fetch_health_state
-
-dht = hivemind.DHT(initial_peers=PUBLIC_INITIAL_PEERS, client_mode=True, start=True)
-model_name = "quantumaikr/llama-2-70b-fb16-korean"
-
-loop = asyncio.get_event_loop()
-# Monkey patch
-def get_types(cls_set: List[Type], component: str):
-    docset = []
-    types = []
-    if component == "input":
-        for cls in cls_set:
-            doc = inspect.getdoc(cls)
-            doc_lines = doc.split("\n")
-            docset.append(doc_lines[1].split(":")[-1])
-            types.append(doc_lines[1].split(")")[0].split("(")[-1])
-    else:
-        for cls in cls_set:
-            doc = inspect.getdoc(cls)
-            doc_lines = doc.split("\n")
-            docset.append(doc_lines[-1].split(":")[-1])
-            types.append(doc_lines[-1].split(")")[0].split("(")[-1])
-    return docset, types
-routes.get_types = get_types
-
-# App code
-
-account_list = dict()
-
-account_list['id'] = "pass"
-
-name_list = dict()
-name_list['id'] = 'name'
-
-p2p_list = dict()
-p2p_list['id'] = '11111111'
-
-def chat(x):
-    
-    return "AI 응답입니다."
-
-
-def register(id, pw):
-    if id in account_list:
-        return "exist"
-    else:
-        account_list[id] = pw
-        return "ok"
+from flask import Flask, jsonify, request
+from flask_cors import CORS
 
-def login(id, pw):
-    if id in account_list:
-        if account_list[id] == pw:
-            return "ok"
-        else:
-            return "password error"
-    else:
-        return "no id"
+import config
+from p2p_utils import check_reachability
+from state_updater import StateUpdaterThread
 
-def add_name(id, name):
-    name_list[id] = name
-    return "ok"
+logger = hivemind.get_logger(__name__)
 
-def get_name(id):
-    if id in name_list:
-        return name_list[id]
-    else:
-        return "no id"
 
-def get_id(name):
-    reverse_dict= dict(map(reversed,name_list.items()))
-    if name in reverse_dict:
-        return reverse_dict[name]
-    else:
-        return "no name"
+logger.info("Connecting to DHT")
+dht = hivemind.DHT(initial_peers=config.INITIAL_PEERS, client_mode=True, num_workers=32, start=True)
 
-def add_p(id, p_id):
-    p2p_list[id] = p_id
-    return "ok"
+logger.info("Starting Flask app")
+app = Flask(__name__)
+CORS(app)
 
-def get_p(id):
-    if id in p2p_list:
-        return p2p_list[id]
-    else:
-        return "no id"
+logger.info("Starting updater")
+updater = StateUpdaterThread(dht, app, daemon=True)
+updater.start()
+updater.ready.wait()
 
-def get_id_from_p2p(i):
-    reverse_dict= dict(map(reversed,p2p_list.items()))
-    if i in reverse_dict:
-        return reverse_dict[i]
-    else:
-        return "no id"
 
-# Blockchain code
+@app.route("/")
+def main_page():
+    return updater.state_html
 
-def get_peers():
-    data = fetch_health_state(dht)
-    out = []
-    for d in data['model_reports']:
-        if d['name'] == model_name:
-            for r in d['server_rows']:
-                out.append(r['peer_id'])
 
-    return out
+@app.route("/api/v1/state")
+def api_v1_state():
+    return app.response_class(response=updater.state_json, status=200, mimetype="application/json")
 
-get_peers()
 
-with gr.Blocks() as demo:
-    count = 0
-    aa = gr.Interface(
-      fn=chat,
-      inputs=["text"],
-      outputs="text",
-      description="chat, ai 응답을 반환합니다.\n /run/predict",
+@app.route("/api/v1/is_reachable/<peer_id>")
+def api_v1_is_reachable(peer_id):
+    peer_id = hivemind.PeerID.from_base58(peer_id)
+    rpc_info = dht.run_coroutine(partial(check_reachability, peer_id, use_cache=False))
+    return jsonify(
+        success=rpc_info["ok"],
+        message=rpc_info.get("error"),
+        your_ip=request.remote_addr,
     )
 
-    rr = gr.Interface(
-      fn=register,
-      inputs=["text", "text"],
-      outputs="text",
-      description="register, 회원가입(성공시:ok, 중복시:exist 반환)\n /run/predict_1",
-    )
-
-    ll = gr.Interface(
-      fn=login,
-      inputs=["text", "text"],
-      outputs="text",
-      description="login, 로그인(성공시: ok, 실패시: password error, 아이디가 없으면: no id) \n /run/predict_2",
-    )
-
-    ad = gr.Interface(
-      fn=add_name,
-      inputs=["text", "text"],
-      outputs="text",
-      description="add_name, id로 닉네임 추가. ok 반환.\n /run/predict_3",
-    )
-    
-    nn = gr.Interface(
-      fn=get_name,
-      inputs=["text"],
-      outputs="text",
-      description="get_name, id로 닉네임 반환(없으면 no id)\n /run/predict_4",
-    )
 
-    nnn = gr.Interface(
-      fn=get_id,
-      inputs=["text"],
-      outputs="text",
-      description="get_name, 닉네임으로 id 반환(없으면 no name)\n /run/predict_5",
-    )
-    
-    adp = gr.Interface(
-      fn=add_p,
-      inputs=["text", "text"],
-      outputs="text",
-      description="add_p, id로 p2p id 추가. ok 반환. \n /run/predict_6",
-    )
-    
-    nnp = gr.Interface(
-      fn=get_p,
-      inputs=["text"],
-      outputs="text",
-      description="get_p, id로 p2p id 반환. 없으면 no id. \n /run/predict_7",
-    )
-
-    nnp = gr.Interface(
-      fn=get_id_from_p2p,
-      inputs=["text"],
-      outputs="text",
-      description="get_p, p2p id로 일반 id 반환. 없으면 no id. \n /run/predict_8",
-    )
-    
-    demo.queue(max_size=32).launch(enable_queue=True)
+@app.route("/metrics")
+@app.route("/api/prometheus")
+def metrics():
+    return app.response_class(response=updater.prometheus_metrics, status=200, mimetype="text/plain")

config.py

@@ -0,0 +1,31 @@
+from petals.constants import PUBLIC_INITIAL_PEERS
+
+from data_structures import ModelInfo
+
+INITIAL_PEERS = PUBLIC_INITIAL_PEERS
+
+MODELS = [
+    ModelInfo(
+        dht_prefix="StableBeluga2-hf",
+        repository="https://huggingface.co/petals-team/StableBeluga2",
+        num_blocks=80,
+    ),
+    ModelInfo(
+        dht_prefix="falcon-180B-chat",
+        repository="https://huggingface.co/tiiuae/falcon-180B-chat",
+        num_blocks=80,
+        limited=True,
+    ),
+    ModelInfo(
+        dht_prefix="Llama-2-70b-chat-hf",
+        repository="https://huggingface.co/meta-llama/Llama-2-70b-chat-hf",
+        num_blocks=80,
+    ),
+    ModelInfo(
+        dht_prefix="Llama-2-70b-hf",
+        repository="https://huggingface.co/meta-llama/Llama-2-70b-hf",
+        num_blocks=80,
+    ),
+]
+
+UPDATE_PERIOD = 60

data_structures.py

@@ -0,0 +1,20 @@
+from typing import Optional
+from urllib.parse import urlparse
+
+import petals
+import pydantic
+
+
+@pydantic.dataclasses.dataclass
+class ModelInfo(petals.data_structures.ModelInfo):
+    dht_prefix: Optional[str] = None
+    official: bool = True
+    limited: bool = False
+
+    @property
+    def name(self) -> str:
+        return urlparse(self.repository).path.lstrip("/")
+
+    @property
+    def short_name(self) -> str:
+        return self.name.split("/")[-1]

health.py

@@ -0,0 +1,124 @@
+import datetime
+import time
+from collections import Counter
+from contextlib import suppress
+from dataclasses import asdict
+from functools import partial
+
+import hivemind
+import numpy as np
+from multiaddr import Multiaddr
+from petals.data_structures import UID_DELIMITER, ServerState
+from petals.utils.dht import compute_spans, get_remote_module_infos
+
+import config
+from data_structures import ModelInfo
+from p2p_utils import check_reachability_parallel, get_peers_ips, extract_peer_ip_info
+
+logger = hivemind.get_logger(__name__)
+
+
+def fetch_health_state(dht: hivemind.DHT) -> dict:
+    start_time = time.perf_counter()
+    bootstrap_peer_ids = []
+    for addr in config.INITIAL_PEERS:
+        peer_id = hivemind.PeerID.from_base58(Multiaddr(addr)["p2p"])
+        if peer_id not in bootstrap_peer_ids:
+            bootstrap_peer_ids.append(peer_id)
+
+    reach_infos = dht.run_coroutine(partial(check_reachability_parallel, bootstrap_peer_ids))
+    bootstrap_states = ["online" if reach_infos[peer_id]["ok"] else "unreachable" for peer_id in bootstrap_peer_ids]
+
+    models = config.MODELS[:]
+    model_index = dht.get("_petals.models", latest=True)
+    if model_index is not None and isinstance(model_index.value, dict):
+        official_dht_prefixes = {model.dht_prefix for model in models}
+        custom_models = []
+        for dht_prefix, model in model_index.value.items():
+            if dht_prefix in official_dht_prefixes:
+                continue
+            with suppress(TypeError, ValueError):
+                model_info = ModelInfo.from_dict(model.value)
+                if model_info.repository is None or not model_info.repository.startswith("https://huggingface.co/"):
+                    continue
+                model_info.dht_prefix = dht_prefix
+                model_info.official = False
+                custom_models.append(model_info)
+        models.extend(sorted(custom_models, key=lambda info: (-info.num_blocks, info.dht_prefix)))
+    logger.info(f"Fetching info for models {[info.name for info in models]}")
+
+    block_uids = [f"{model.dht_prefix}{UID_DELIMITER}{i}" for model in models for i in range(model.num_blocks)]
+    module_infos = get_remote_module_infos(dht, block_uids, latest=True)
+
+    model_servers = {}
+    all_servers = {}
+    offset = 0
+    for model in models:
+        model_servers[model.dht_prefix] = compute_spans(
+            module_infos[offset : offset + model.num_blocks], min_state=ServerState.OFFLINE
+        )
+        all_servers.update(model_servers[model.dht_prefix])
+        offset += model.num_blocks
+
+    online_servers = [peer_id for peer_id, span in all_servers.items() if span.state == ServerState.ONLINE]
+
+    reach_infos.update(dht.run_coroutine(partial(check_reachability_parallel, online_servers, fetch_info=True)))
+    peers_info = {str(peer.peer_id): {"location": extract_peer_ip_info(str(peer.addrs[0])), "multiaddrs": [str(multiaddr) for multiaddr in peer.addrs]} for peer in dht.run_coroutine(get_peers_ips)}
+
+    top_contributors = Counter()
+    model_reports = []
+    for model in models:
+        block_healthy = np.zeros(model.num_blocks, dtype=bool)
+        server_rows = []
+        for peer_id, span in sorted(model_servers[model.dht_prefix].items()):
+            reachable = reach_infos[peer_id]["ok"] if peer_id in reach_infos else True
+            state = span.state.name.lower() if reachable else "unreachable"
+            if state == "online":
+                block_healthy[span.start : span.end] = True
+
+            show_public_name = state == "online" and span.length >= 10
+            if model.official and span.server_info.public_name and show_public_name:
+                top_contributors[span.server_info.public_name] += span.length
+
+            row = {
+                "short_peer_id": "..." + str(peer_id)[-6:],
+                "peer_id": peer_id,
+                "peer_ip_info": peers_info.get(str(peer_id), "unknown"),
+                "show_public_name": show_public_name,
+                "state": state,
+                "span": span,
+                "adapters": [dict(name=name, short_name=name.split("/")[-1]) for name in span.server_info.adapters],
+                "pings_to_me": {
+                    str(origin_id): origin.server_info.next_pings[str(peer_id)]
+                    for origin_id, origin in model_servers[model.dht_prefix].items()
+                    if origin.server_info.next_pings is not None and str(peer_id) in origin.server_info.next_pings
+                },
+            }
+            if span.server_info.cache_tokens_left is not None:
+                # We use num_blocks * 2 to account for both keys and values
+                row["cache_tokens_left_per_block"] = span.server_info.cache_tokens_left // (span.length * 2)
+            server_rows.append(row)
+
+        model_reports.append(
+            dict(
+                name=model.name,
+                short_name=model.short_name,
+                state="healthy" if block_healthy.all() else "broken",
+                server_rows=server_rows,
+                **asdict(model),
+            )
+        )
+
+    reachability_issues = [
+        dict(peer_id=peer_id, err=info["error"]) for peer_id, info in sorted(reach_infos.items()) if not info["ok"]
+    ]
+
+    return dict(
+        bootstrap_states=bootstrap_states,
+        top_contributors=top_contributors,
+        model_reports=model_reports,
+        reachability_issues=reachability_issues,
+        last_updated=datetime.datetime.now(datetime.timezone.utc),
+        update_period=config.UPDATE_PERIOD,
+        update_duration=time.perf_counter() - start_time
+    )

metrics.py

@@ -0,0 +1,118 @@
+from collections import Counter, defaultdict
+from typing import List
+
+import numpy as np
+
+
+def get_servers_metrics(model_reports) -> List[str]:
+    servers_num_total = 0
+    servers_num_relay = 0
+    num_peers = 0
+    pings = []
+    num_ping_infs = 0
+    version_counts = Counter()
+    result = ["# SERVER LEVEL METRICS"]
+
+    for model_reports in model_reports:
+        for server in model_reports["server_rows"]:
+            if server["span"].server_info is not None:
+                next_pings = server["span"].server_info.next_pings
+                if next_pings is not None:
+                    servers_num_total += 1
+                    num_peers += len(next_pings)
+                    pings_not_inf = [v for k, v in next_pings.items() if v != float("inf")]
+                    pings.extend(pings_not_inf)
+                    num_ping_infs += len([v for v in next_pings.values() if v == float("inf")])
+
+                if server["span"].server_info.using_relay:
+                    servers_num_relay += 1
+
+                version = server["span"].server_info.version
+                if version:
+                    version_counts[version] += 1
+
+    if servers_num_total > 0 and pings:
+        peers_per_srv = (len(pings) + num_ping_infs) / servers_num_total
+        pings_inf_share = num_ping_infs / (num_ping_infs + len(pings))
+
+        result.extend(
+            [
+                f"peers_per_srv {peers_per_srv:.1f}",
+                f"pings_inf_share {pings_inf_share:.3f}",
+            ]
+        )
+
+    result.append(f"servers_num_total {servers_num_total}")
+    result.append(f"servers_num_relay {servers_num_relay}")
+
+    if pings:
+        result.append("# PINGS")
+        pings = np.sort(pings).tolist()
+        for pct in (25, 50, 75, 90, 95):
+            result.append(f'ping_pct{{pct="{pct}"}} {np.percentile(pings, pct):.4f}')
+
+    result.append("# VERSIONS")
+    for version_number, version_count in version_counts.items():
+        result.append(f'server_version{{version_number="{version_number}"}} {version_count}')
+
+    return result
+
+
+def get_models_metrics(model_reports) -> List[str]:
+    result = [
+        "# MODEL LEVEL METRICS",
+    ]
+
+    for model_reports in model_reports:
+        model_name = model_reports["dht_prefix"]
+
+        result.append(f"# MODEL: {model_name} {'-' * 50}")
+
+        blocks = defaultdict(lambda: np.zeros(model_reports["num_blocks"]))
+
+        for server in model_reports["server_rows"]:
+            for block_idx in range(server["span"].start, server["span"].end):
+                blocks["total"][block_idx] += 1
+                blocks[server["state"]][block_idx] += 1
+
+                if server["span"].server_info is not None:
+                    for rps in ("network_rps", "inference_rps", "forward_rps"):
+                        rps_value = getattr(server["span"].server_info, rps, 0)
+                        if rps_value is not None:
+                            blocks[rps][block_idx] += rps_value
+
+        result.extend(
+            [
+                f'n_blocks{{model="{model_name}"}} {model_reports["num_blocks"]}',
+                f'servers_num{{model="{model_name}"}} {len(model_reports["server_rows"])}',
+                f'blocks_total{{model="{model_name}"}} {blocks["total"].sum()}',
+                f'blocks_online_min{{model="{model_name}"}} {blocks["online"].min()}',
+            ]
+        )
+
+        for block_state in ("online", "joining", "offline", "unreachable"):
+            result.append(f'blocks{{model="{model_name}",state="{block_state}"}} {blocks[block_state].sum():.0f}')
+
+        for rps in ("network_rps", "inference_rps", "forward_rps"):
+            rps_type = rps.split("_")[0]
+            result.append(f'rps_avg{{model="{model_name}",rps="{rps_type}"}} {blocks[rps].mean():.1f}')
+            result.append(f'rps_min{{model="{model_name}",rps="{rps_type}"}} {blocks[rps].min():.1f}')
+
+    return result
+
+
+def get_prometheus_metrics(state_dict) -> str:
+    """prepares metrics in Prometeus format
+    description: https://prometheus.io/docs/instrumenting/exposition_formats/
+    returns multline string with single metric per line
+    """
+    result = []
+
+    result.append("# GENERAL METRICS")
+    result.append(f"update_duration {state_dict.get('update_duration', None):.1f}")
+
+    result.extend(get_servers_metrics(state_dict["model_reports"]))
+
+    result.extend(get_models_metrics(state_dict["model_reports"]))
+
+    return "\n".join(result)

p2p_utils.py

@@ -0,0 +1,67 @@
+import re
+import asyncio
+import requests
+import hivemind
+import functools
+from async_timeout import timeout
+from petals.server.handler import TransformerConnectionHandler
+
+info_cache = hivemind.TimedStorage()
+
+
+async def check_reachability(peer_id, _, node, *, fetch_info=False, connect_timeout=5, expiration=300, use_cache=True):
+    if use_cache:
+        entry = info_cache.get(peer_id)
+        if entry is not None:
+            return entry.value
+
+    try:
+        with timeout(connect_timeout):
+            if fetch_info:  # For Petals servers
+                stub = TransformerConnectionHandler.get_stub(node.p2p, peer_id)
+                response = await stub.rpc_info(hivemind.proto.runtime_pb2.ExpertUID())
+                rpc_info = hivemind.MSGPackSerializer.loads(response.serialized_info)
+                rpc_info["ok"] = True
+            else:  # For DHT-only bootstrap peers
+                await node.p2p._client.connect(peer_id, [])
+                await node.p2p._client.disconnect(peer_id)
+                rpc_info = {"ok": True}
+    except Exception as e:
+        # Actual connection error
+        if not isinstance(e, asyncio.TimeoutError):
+            message = str(e) if str(e) else repr(e)
+            if message == "protocol not supported":
+                # This may be returned when a server is joining, see https://github.com/petals-infra/health.petals.dev/issues/1
+                return {"ok": True}
+        else:
+            message = f"Failed to connect in {connect_timeout:.0f} sec. Firewall may be blocking connections"
+        rpc_info = {"ok": False, "error": message}
+
+    info_cache.store(peer_id, rpc_info, hivemind.get_dht_time() + expiration)
+    return rpc_info
+
+
+async def check_reachability_parallel(peer_ids, dht, node, *, fetch_info=False):
+    rpc_infos = await asyncio.gather(
+        *[check_reachability(peer_id, dht, node, fetch_info=fetch_info) for peer_id in peer_ids]
+    )
+    return dict(zip(peer_ids, rpc_infos))
+
+
+async def get_peers_ips(dht, dht_node):
+    return await dht_node.p2p.list_peers()
+
+@functools.cache
+def get_location(ip_address):
+    try:
+        response = requests.get(f"http://ip-api.com/json/{ip_address}")
+        if response.status_code == 200:
+            return response.json()
+    except Exception:
+        pass
+    return {}
+
+def extract_peer_ip_info(multiaddr_str):
+    if ip_match := re.search(r"/ip4/(\d+\.\d+\.\d+\.\d+)", multiaddr_str):
+        return get_location(ip_match[1])
+    return {}

pyproject.toml

@@ -0,0 +1,10 @@
+[tool.black]
+line-length = 120
+required-version = "22.3.0"
+
+[tool.isort]
+profile = "black"
+line_length = 120
+combine_as_imports = true
+combine_star = true
+known_local_folder = ["tests", "cli"]

state_updater.py

@@ -0,0 +1,57 @@
+import datetime
+import threading
+import time
+from dataclasses import asdict, is_dataclass
+from enum import Enum
+
+import hivemind
+import simplejson
+from flask import Flask, render_template
+
+import config
+from health import fetch_health_state
+from metrics import get_prometheus_metrics
+
+logger = hivemind.get_logger(__name__)
+
+
+class StateUpdaterThread(threading.Thread):
+    def __init__(self, dht: hivemind.DHT, app: Flask, **kwargs):
+        super().__init__(**kwargs)
+        self.dht = dht
+        self.app = app
+
+        self.state_json = self.state_html = None
+        self.ready = threading.Event()
+
+    def run(self):
+        while True:
+            start_time = time.perf_counter()
+            try:
+                state_dict = fetch_health_state(self.dht)
+                with self.app.app_context():
+                    self.state_html = render_template("index.html", **state_dict)
+                    self.prometheus_metrics = get_prometheus_metrics(state_dict)
+                self.state_json = simplejson.dumps(state_dict, indent=2, ignore_nan=True, default=json_default)
+
+                self.ready.set()
+                logger.info(f"Fetched new state in {time.perf_counter() - start_time:.1f} sec")
+            except Exception:
+                logger.error("Failed to update state:", exc_info=True)
+
+            delay = config.UPDATE_PERIOD - (time.perf_counter() - start_time)
+            if delay < 0:
+                logger.warning("Update took more than update_period, consider increasing it")
+            time.sleep(max(delay, 0))
+
+
+def json_default(value):
+    if is_dataclass(value):
+        return asdict(value)
+    if isinstance(value, Enum):
+        return value.name.lower()
+    if isinstance(value, hivemind.PeerID):
+        return value.to_base58()
+    if isinstance(value, datetime.datetime):
+        return value.timestamp()
+    raise TypeError(f"Can't serialize {repr(value)}")