Harden KOTH weight inputs and fallback

gas/koth_weights.py

@@ -1,8 +1,10 @@
 """Build on-chain weights for King-of-the-Hill discriminator lanes."""
 
+from numbers import Real
 from typing import Callable, Dict, Iterable, List, Optional
 
 import numpy as np
+from bittensor.utils import is_valid_ss58_address
 
 KOTH_SPLIT = {
     "image": 0.40,
@@ -14,6 +16,103 @@
 # Unused slots (no prior distinct king) roll up to the current king.
 KOTH_LANE_RESIDUAL = (0.85, 0.10, 0.05)
 KOTH_CHAIN_ROLES = ("current", "previous", "two_back")
+KOTH_MODALITIES = ("image", "video", "audio")
+
+
+def resolve_koth_split(split: Optional[Dict[str, float]]) -> Dict[str, float]:
+    """Overlay an API split on protocol defaults and validate the result."""
+    if split is not None and not isinstance(split, dict):
+        raise ValueError("Invalid KOTH split")
+
+    resolved = dict(KOTH_SPLIT)
+    for key in KOTH_SPLIT:
+        if split is None or key not in split:
+            continue
+        value = split[key]
+        if isinstance(value, bool) or not isinstance(value, Real):
+            raise ValueError("Invalid KOTH split")
+        try:
+            resolved[key] = float(value)
+        except (OverflowError, TypeError, ValueError) as exc:
+            raise ValueError("Invalid KOTH split") from exc
+
+    if any(
+        not np.isfinite(value) or value < 0 for value in resolved.values()
+    ) or not np.isclose(sum(resolved.values()), 1.0, rtol=0.0, atol=1e-9):
+        raise ValueError("Invalid KOTH split")
+    return resolved
+
+
+def validate_koth_payload(payload: object) -> dict:
+    """Return the trusted subset of a current-kings API response."""
+    if not isinstance(payload, dict):
+        raise ValueError("Invalid current-kings payload")
+
+    raw_kings = payload.get("kings")
+    if raw_kings is None:
+        raw_kings = []
+    if not isinstance(raw_kings, list):
+        raise ValueError("Invalid current-kings payload")
+
+    kings = []
+    king_addresses = {}
+    for king in raw_kings:
+        if not isinstance(king, dict):
+            raise ValueError("Invalid current-kings payload")
+        modality = king.get("modality")
+        address = king.get("ss58_address")
+        if modality not in KOTH_MODALITIES or modality in king_addresses:
+            raise ValueError("Invalid current-kings payload")
+        if not isinstance(address, str) or not is_valid_ss58_address(address):
+            raise ValueError("Invalid current-kings payload")
+        king_addresses[modality] = address
+        kings.append({"modality": modality, "ss58_address": address})
+
+    raw_chain = payload.get("chain")
+    if raw_chain is None:
+        raw_chain = {}
+    if not isinstance(raw_chain, dict):
+        raise ValueError("Invalid current-kings payload")
+
+    chain = {}
+    for modality in KOTH_MODALITIES:
+        raw_members = raw_chain.get(modality)
+        if raw_members is None:
+            continue
+        if not isinstance(raw_members, list):
+            raise ValueError("Invalid current-kings payload")
+
+        addresses = []
+        for member in raw_members:
+            if isinstance(member, str):
+                address = member
+            elif isinstance(member, dict):
+                address = member.get("ss58_address")
+            else:
+                raise ValueError("Invalid current-kings payload")
+            if not isinstance(address, str) or not is_valid_ss58_address(address):
+                raise ValueError("Invalid current-kings payload")
+            addresses.append(address)
+
+        distinct_addresses = list(dict.fromkeys(addresses))
+        if len(distinct_addresses) > len(KOTH_LANE_RESIDUAL):
+            raise ValueError("Invalid current-kings payload")
+        if (
+            distinct_addresses
+            and modality in king_addresses
+            and distinct_addresses[0] != king_addresses[modality]
+        ):
+            raise ValueError("Invalid current-kings payload")
+        if distinct_addresses:
+            chain[modality] = [
+                {"ss58_address": address} for address in distinct_addresses
+            ]
+
+    return {
+        "kings": kings,
+        "chain": chain,
+        "split": resolve_koth_split(payload.get("split")),
+    }
 
 
 def kings_by_modality(payload: Optional[dict]) -> Dict[str, str]:
@@ -24,7 +123,7 @@ def kings_by_modality(payload: Optional[dict]) -> Dict[str, str]:
     for king in payload.get("kings") or []:
         modality = king.get("modality")
         hotkey = king.get("ss58_address")
-        if modality in ("image", "video", "audio") and hotkey:
+        if modality in KOTH_MODALITIES and hotkey:
             out[modality] = hotkey
     return out
 
@@ -61,7 +160,7 @@ def chains_by_modality(payload: Optional[dict]) -> Dict[str, List[Dict[str, obje
     raw = payload.get("chain") or {}
     kings = kings_by_modality(payload)
     out: Dict[str, List[Dict[str, object]]] = {}
-    for modality in ("image", "video", "audio"):
+    for modality in KOTH_MODALITIES:
         members = raw.get(modality) or []
         hotkeys = []
         for member in members:
@@ -95,24 +194,18 @@ def build_koth_weights(
     king. An unresolvable current king burns its share; unresolvable previous
     kings roll to the current king when that UID resolved.
     """
-    split = dict(KOTH_SPLIT if split is None else split)
-    try:
-        split = {key: float(split[key]) for key in KOTH_SPLIT}
-    except (KeyError, TypeError, ValueError) as exc:
-        raise ValueError("Invalid KOTH split") from exc
-    if any(
-        not np.isfinite(value) or value < 0 for value in split.values()
-    ) or not np.isclose(sum(split.values()), 1.0):
-        raise ValueError("Invalid KOTH split")
+    split = resolve_koth_split(split)
     weights = np.zeros(n, dtype=np.float64)
+    if burn_uid is None or not 0 <= burn_uid < n:
+        raise ValueError("Owner/burn UID is unavailable")
     if len(scores) < n:
         scores = np.append(scores, np.zeros(n - len(scores)))
     elif len(scores) > n:
         scores = scores[:n]
 
     king_uids = set()
     burned = 0.0
-    for modality in ("image", "video", "audio"):
+    for modality in KOTH_MODALITIES:
         pct = split[modality]
         members = list((chains or {}).get(modality) or [])
         if not members:
@@ -150,7 +243,13 @@ def build_koth_weights(
                 burned += pct * leftover
 
     generator_pct = split["generator"]
-    active = [uid for uid in generator_uids if 0 <= uid < n and uid not in king_uids]
+    active = list(
+        dict.fromkeys(
+            uid
+            for uid in generator_uids
+            if 0 <= uid < n and uid not in king_uids
+        )
+    )
     if active and generator_pct > 0:
         gen_scores = np.array([max(float(scores[uid]), 0.0) for uid in active])
         total = float(np.sum(gen_scores))
@@ -163,10 +262,14 @@ def build_koth_weights(
         burned += generator_pct
 
     if burned > 0:
-        if burn_uid is None or not 0 <= burn_uid < n:
-            raise ValueError(
-                f"Cannot allocate {burned:.4f} burn weight: burn UID is unavailable"
-            )
         weights[burn_uid] += burned
 
+    if (
+        len(weights) != n
+        or not np.all(np.isfinite(weights))
+        or np.any(weights < 0)
+        or not np.isclose(float(np.sum(weights)), 1.0, rtol=0.0, atol=1e-9)
+    ):
+        raise ValueError("Invalid final KOTH weight vector")
+
     return weights

neurons/validator/validator.py

@@ -11,7 +11,13 @@
 
 from gas import __spec_version__ as spec_version
 from gas.protocol.validator_requests import get_benchmark_results, get_current_kings
-from gas.koth_weights import build_koth_weights, chains_by_modality, kings_by_modality
+from gas.koth_weights import (
+    KOTH_SPLIT,
+    build_koth_weights,
+    chains_by_modality,
+    kings_by_modality,
+    validate_koth_payload,
+)
 from gas.utils.autoupdater import autoupdate
 from gas.cache import ContentManager
 from gas.utils.metagraph import create_set_weights
@@ -41,21 +47,48 @@
 
 MAINNET_UID = 34
 BURN_SS58 = "5HjBSeeoz52CLfvDWDkzupqrYLHz1oToDPHjdmJjc4TF68LQ"
+KINGS_CACHE_TTL_SECONDS = 24 * 60 * 60
 
 
 class _KingsState:
     """Persist last-known KOTH kings across validator restarts."""
 
     def __init__(self):
         self.payload = None
+        self.fetched_at = None
+
+    def update(self, payload: dict, now: float = None) -> None:
+        self.payload = payload
+        self.fetched_at = time.time() if now is None else now
+
+    def fresh_payload(self, now: float = None):
+        if self.payload is None or isinstance(self.fetched_at, bool):
+            return None
+        try:
+            age = (time.time() if now is None else now) - float(self.fetched_at)
+        except (TypeError, ValueError):
+            return None
+        if 0 <= age <= KINGS_CACHE_TTL_SECONDS:
+            return self.payload
+        return None
+
+    def age_seconds(self, now: float = None):
+        if self.payload is None or isinstance(self.fetched_at, bool):
+            return None
+        try:
+            return (time.time() if now is None else now) - float(self.fetched_at)
+        except (TypeError, ValueError):
+            return None
 
     def save_state(self, save_dir: str, filename: str) -> None:
         import json
         import os
 
         path = os.path.join(save_dir, filename)
         with open(path, "w") as f:
-            json.dump(self.payload, f)
+            json.dump(
+                {"payload": self.payload, "fetched_at": self.fetched_at}, f
+            )
 
     def load_state(self, save_dir: str, filename: str) -> bool:
         import json
@@ -65,7 +98,14 @@ def load_state(self, save_dir: str, filename: str) -> bool:
         if not os.path.exists(path):
             return False
         with open(path) as f:
-            self.payload = json.load(f)
+            state = json.load(f)
+        if isinstance(state, dict) and "payload" in state:
+            self.payload = state.get("payload")
+            self.fetched_at = state.get("fetched_at")
+        else:
+            # Legacy caches have no trustworthy age and are therefore stale.
+            self.payload = state
+            self.fetched_at = None
         return True
 
 
@@ -201,16 +241,45 @@ async def set_weights(self, block):
             self.wallet.hotkey, base_url=self.config.benchmark_api_url
         )
         if kings_payload is not None:
-            self.kings_state.payload = kings_payload
-        elif self.kings_state.payload is not None:
-            bt.logging.warning("current-kings API unavailable; using last known kings")
-            kings_payload = self.kings_state.payload
-        else:
+            try:
+                kings_payload = validate_koth_payload(kings_payload)
+                self.kings_state.update(kings_payload)
+            except ValueError as e:
+                bt.logging.error(f"Rejected current-kings API response: {e}")
+                kings_payload = None
+
+        if kings_payload is None:
+            cached_payload = self.kings_state.fresh_payload()
+            if cached_payload is not None:
+                try:
+                    kings_payload = validate_koth_payload(cached_payload)
+                    age = self.kings_state.age_seconds()
+                    bt.logging.warning(
+                        "current-kings API unavailable or invalid; using "
+                        f"last known good response ({age / 3600:.1f}h old)"
+                    )
+                except ValueError as e:
+                    bt.logging.error(f"Rejected cached current-kings response: {e}")
+                    kings_payload = None
+
+        if kings_payload is None:
+            cache_age = self.kings_state.age_seconds()
+            if cache_age is not None:
+                bt.logging.warning(
+                    "current-kings cache expired; discriminator shares will burn"
+                )
+            else:
+                bt.logging.warning(
+                    "No valid current-kings cache; discriminator shares will burn"
+                )
             bt.logging.warning(
-                "current-kings API unavailable and no cached kings; "
-                "discriminator shares will burn"
+                "Using the local default KOTH split with no current kings"
             )
-            kings_payload = {"kings": []}
+            kings_payload = {
+                "kings": [],
+                "chain": {},
+                "split": dict(KOTH_SPLIT),
+            }
 
         kings = kings_by_modality(kings_payload)
         chains = chains_by_modality(kings_payload)

tests/test_koth_cache.py

@@ -0,0 +1,36 @@
+"""Tests for the validator's bounded last-known-good KOTH cache."""
+
+import json
+
+from neurons.validator.validator import KINGS_CACHE_TTL_SECONDS, _KingsState
+
+
+def test_kings_cache_expires_after_24_hours():
+    state = _KingsState()
+    payload = {"kings": [], "chain": {}, "split": {}}
+    state.update(payload, now=100.0)
+
+    assert state.fresh_payload(now=100.0 + KINGS_CACHE_TTL_SECONDS) == payload
+    assert state.fresh_payload(now=100.0 + KINGS_CACHE_TTL_SECONDS + 1) is None
+
+
+def test_kings_cache_persists_timestamp(tmp_path):
+    state = _KingsState()
+    payload = {"kings": [], "chain": {}, "split": {}}
+    state.update(payload, now=123.0)
+    state.save_state(str(tmp_path), "kings.json")
+
+    restored = _KingsState()
+    assert restored.load_state(str(tmp_path), "kings.json") is True
+    assert restored.payload == payload
+    assert restored.fetched_at == 123.0
+
+
+def test_legacy_cache_without_timestamp_is_stale(tmp_path):
+    payload = {"kings": []}
+    (tmp_path / "kings.json").write_text(json.dumps(payload))
+
+    state = _KingsState()
+    assert state.load_state(str(tmp_path), "kings.json") is True
+    assert state.payload == payload
+    assert state.fresh_payload(now=100.0) is None