Implement validator source selection

app.py

@@ -1,3 +1,4 @@
+import json
 import os
 import time
 from dataclasses import dataclass
@@ -6,20 +7,20 @@ from zoneinfo import ZoneInfo
 
 import gradio as gr
 import plotly.graph_objects as go
-import schedule
 import wandb
 from substrateinterface import Keypair
-from wandb.apis.public import Run, Runs
+from wandb.apis.public import Run
 
-wandb_api = wandb.Api()
-demo = gr.Blocks(css=".typewriter {font-family: 'JMH Typewriter', sans-serif;}")
-
-SOURCE_VALIDATOR_UID = int(os.environ["SOURCE_VALIDATOR_UID"])
 WANDB_RUN_PATH = os.environ["WANDB_RUN_PATH"]
-
+SOURCE_VALIDATOR_UID = int(os.environ["SOURCE_VALIDATOR_UID"])
+REFRESH_RATE = 60 * 30  # 30 minutes
 BASELINE = 0.0
 GRAPH_HISTORY_DAYS = 30
-MAX_GRAPH_ENTRIES = 5
+MAX_GRAPH_ENTRIES = 10
+
+wandb_api = wandb.Api()
+demo = gr.Blocks(css=".typewriter {font-family: 'JMH Typewriter', sans-serif;}")
+runs: dict[int, list[Run]] = {}
 
 
 @dataclass
@@ -59,13 +60,10 @@ def is_valid_run(run: Run):
         return False
 
 
-def get_graph_entries(runs: Runs) -> dict[int, GraphEntry]:
-    graph_entries: dict[int, GraphEntry] = {}
+def get_graph_entries(runs: list[Run]) -> dict[int, GraphEntry]:
+    entries: dict[int, GraphEntry] = {}
 
     for run in reversed(runs[:GRAPH_HISTORY_DAYS]):
-        if not is_valid_run(run):
-            continue
-
         date = datetime.strptime(run.created_at, "%Y-%m-%dT%H:%M:%S")
 
         for key, value in run.summary.items():
@@ -76,24 +74,25 @@ def get_graph_entries(runs: Runs) -> dict[int, GraphEntry]:
             score = value["score"]
             model = value["model"]
 
-            if uid not in graph_entries:
-                graph_entries[uid] = GraphEntry([date], [score], [model], score)
+            if uid not in entries:
+                entries[uid] = GraphEntry([date], [score], [model], score)
             else:
-                if score > graph_entries[uid].max_score:
-                    graph_entries[uid].max_score = score
+                if score > entries[uid].max_score:
+                    entries[uid].max_score = score
 
-                data = graph_entries[uid]
+                data = entries[uid]
                 data.dates.append(date)
                 data.scores.append(data.max_score)
                 data.models.append(model)
 
-    return dict(sorted(graph_entries.items(), key=lambda entry: entry[1].max_score, reverse=True)[:MAX_GRAPH_ENTRIES])
+    return dict(sorted(entries.items(), key=lambda entry: entry[1].max_score, reverse=True)[:MAX_GRAPH_ENTRIES])
 
 
-def create_graph(graph_entries: dict[int, GraphEntry]):
+def create_graph(runs: list[Run]) -> go.Figure:
+    entries = get_graph_entries(runs)
     fig = go.Figure()
 
-    for uid, data in graph_entries.items():
+    for uid, data in entries.items():
         fig.add_trace(go.Scatter(
             x=data.dates,
             y=data.scores,
@@ -107,7 +106,7 @@ def create_graph(graph_entries: dict[int, GraphEntry]):
             ),
         ))
 
-    date_range = max(graph_entries.values(), key=lambda entry: len(entry.dates)).dates
+    date_range = max(entries.values(), key=lambda entry: len(entry.dates)).dates
 
     fig.add_trace(go.Scatter(
         x=date_range,
@@ -127,62 +126,89 @@ def create_graph(graph_entries: dict[int, GraphEntry]):
         template="plotly_dark"
     )
 
-    gr.Plot(fig)
+    return fig
 
 
-def refresh_leaderboard():
-    now = datetime.now(tz=ZoneInfo("America/New_York"))
-    print(f"Refreshing Leaderboard at {now.strftime('%Y-%m-%d %H:%M:%S')}")
+def create_leaderboard(runs: list[Run]) -> list[tuple]:
+    entries: dict[int, LeaderboardEntry] = {}
 
-    demo.clear()
+    for run in runs:
+        has_data = False
+        for key, value in run.summary.items():
+            if key.startswith("_"):
+                continue
 
-    with demo:
-        with gr.Accordion("Contest #1 Submission Leader: New Dream SDXL on NVIDIA RTX 4090s"):
-            runs: Runs = wandb_api.runs(
-                WANDB_RUN_PATH,
-                filters={"config.type": "validator", "config.uid": SOURCE_VALIDATOR_UID},
-                order="-created_at",
-            )
+            has_data = True
+
+            try:
+                uid = int(key)
+
+                entries[uid] = LeaderboardEntry(
+                    uid=uid,
+                    rank=value["rank"],
+                    model=value["model"],
+                    score=value["score"],
+                    hotkey=value["hotkey"],
+                    previous_day_winner=value["multiday_winner"],
+                )
+            except Exception:
+                continue
 
-            entries: dict[int, LeaderboardEntry] = {}
+        if has_data:
+            break
 
-            for run in runs:
-                if not is_valid_run(run):
-                    continue
+    leaderboard: list[tuple] = [
+        (entry.rank + 1, entry.uid, entry.model, entry.score, entry.hotkey, entry.previous_day_winner)
+        for entry in sorted(entries.values(), key=lambda entry: (entry.score, entry.rank), reverse=True)
+    ]
 
-                has_data = False
-                for key, value in run.summary.items():
-                    if key.startswith("_"):
-                        continue
+    return leaderboard
 
-                    has_data = True
 
-                    try:
-                        uid = int(key)
+def get_run_validator_uid(run: Run) -> int:
+    json_config = json.loads(run.json_config)
+    uid = int(json_config["uid"]["value"])
+    return uid
 
-                        entries[uid] = LeaderboardEntry(
-                            uid=uid,
-                            rank=value["rank"],
-                            model=value["model"],
-                            score=value["score"],
-                            hotkey=value["hotkey"],
-                            previous_day_winner=value["multiday_winner"],
-                        )
-                    except Exception:
-                        continue
 
-                if has_data:
-                    break
+def fetch_wandb_data():
+    wandb_runs = wandb_api.runs(
+        WANDB_RUN_PATH,
+        filters={"config.type": "validator"},
+        order="-created_at",
+    )
+
+    global runs
+    runs.clear()
+    for run in wandb_runs:
+        if not is_valid_run(run):
+            continue
 
-            leaderboard: list[tuple] = [
-                (entry.rank + 1, entry.uid, entry.model, entry.score, entry.hotkey, entry.previous_day_winner)
-                for entry in sorted(entries.values(), key=lambda entry: (entry.score, entry.rank), reverse=True)
-            ]
+        uid = get_run_validator_uid(run)
+        if uid not in runs:
+            runs[uid] = []
+        runs[uid].append(run)
 
-            create_graph(get_graph_entries(runs))
+    runs = dict(sorted(runs.items(), key=lambda item: item[0]))
 
-            gr.components.Dataframe(
-                value=leaderboard,
+
+def refresh():
+    fetch_wandb_data()
+    demo.clear()
+    with demo:
+        with gr.Accordion("Contest #1 Submission Leader: New Dream SDXL on NVIDIA RTX 4090s"):
+            choices = list(runs.keys())
+            dropdown = gr.Dropdown(
+                choices,
+                value=SOURCE_VALIDATOR_UID,
+                interactive=True,
+                label="Source Validator"
+            )
+
+            graph = gr.Plot()
+
+            leaderboard = gr.components.Dataframe(
+                create_leaderboard(runs[dropdown.value]),
                 headers=["Rank", "Uid", "Model", "Score", "Hotkey", "Previous day winner"],
                 datatype=["number", "number", "markdown", "number", "markdown", "bool"],
                 elem_id="leaderboard-table",
@@ -190,16 +216,20 @@ def refresh_leaderboard():
                 visible=True,
             )
 
+            demo.load(lambda uid: create_graph(runs[uid]), [dropdown], [graph])
+
+            dropdown.change(lambda uid: create_graph(runs[uid]), [dropdown], [graph])
+            dropdown.change(lambda uid: create_leaderboard(runs[uid]), [dropdown], [leaderboard])
 
-def main():
-    refresh_leaderboard()
-    schedule.every(30).minutes.do(refresh_leaderboard)
 
+if __name__ == "__main__":
+    refresh()
     demo.launch(prevent_thread_lock=True)
 
     while True:
-        schedule.run_pending()
-        time.sleep(1)
+        time.sleep(REFRESH_RATE)
 
+        now = datetime.now(tz=ZoneInfo("America/New_York"))
+        print(f"Refreshing Leaderboard at {now.strftime('%Y-%m-%d %H:%M:%S')}")
 
-main()
+        refresh()