[FIX] Harden epoch RPC retry and snapshot resilience (#22)

Co-authored-by: Midori AI Agent <contact-us@midori-ai.xyz>

Author: lunamidori5

metahash/validator/epoch_validator.py

@@ -2,6 +2,8 @@
 from __future__ import annotations
 
 import asyncio
+import random
+import time
 from datetime import datetime
 from typing import Optional, Tuple, Any
 
@@ -41,6 +43,44 @@ def __init__(self, *args, log_interval_blocks: int = 2, **kwargs):
         self.current_strategy_out: Any = None  # can be dict (subnet bps) or list
 
     # ----------------------- helpers ---------------------------------- #
+    def _retry_rpc_call(self, fn, *, op_name: str):
+        """
+        Retry RPC calls forever (until shutdown) with exponential backoff.
+        This keeps the validator alive through transient RPC/node drift issues.
+        """
+        attempt = 0
+        delay_s = 1.0
+        cap_s = 30.0
+        last_error: Exception | None = None
+
+        while not self.should_exit:
+            try:
+                result = fn()
+                if attempt > 0:
+                    bt.logging.info(
+                        f"[epoch] RPC recovered for {op_name} after {attempt} retries."
+                    )
+                return result
+            except Exception as err:
+                last_error = err
+                attempt += 1
+
+                if attempt == 1 or attempt % 5 == 0:
+                    bt.logging.warning(
+                        f"[epoch] RPC error during {op_name} (attempt {attempt}): {err}. "
+                        f"Retrying in ~{int(min(cap_s, delay_s))}s."
+                    )
+
+                sleep_s = min(cap_s, delay_s)
+                sleep_s += random.uniform(0.0, min(1.0, sleep_s * 0.25))
+                time.sleep(sleep_s)
+                delay_s = min(cap_s, delay_s * 2.0)
+
+        raise RuntimeError(
+            f"[epoch] Stopping RPC retries for {op_name} due to shutdown "
+            f"(last_error={last_error!r})."
+        )
+
     def _discover_epoch_length(self) -> int:
         try:
             override = int(EPOCH_LENGTH_OVERRIDE or 0)
@@ -59,10 +99,19 @@ def _discover_epoch_length(self) -> int:
             return length
 
         self._override_active = False
-        tempo = self.subtensor.tempo(self.config.netuid) or 360
+        tempo = self._retry_rpc_call(
+            lambda: self.subtensor.tempo(self.config.netuid),
+            op_name="tempo",
+        ) or 360
         try:
-            head = self.subtensor.get_current_block()
-            next_head = self.subtensor.get_next_epoch_start_block(self.config.netuid)
+            head = self._retry_rpc_call(
+                self.subtensor.get_current_block,
+                op_name="get_current_block",
+            )
+            next_head = self._retry_rpc_call(
+                lambda: self.subtensor.get_next_epoch_start_block(self.config.netuid),
+                op_name="get_next_epoch_start_block",
+            )
             if next_head is None:
                 raise ValueError("RPC returned None")
             derived = next_head - (head - head % tempo)
@@ -77,7 +126,10 @@ def _discover_epoch_length(self) -> int:
         return length
 
     def _epoch_snapshot(self) -> Tuple[int, int, int, int, int]:
-        blk = self.subtensor.get_current_block()
+        blk = self._retry_rpc_call(
+            self.subtensor.get_current_block,
+            op_name="get_current_block",
+        )
         ep_l = self._epoch_len or self._discover_epoch_length()
         start = blk - (blk % ep_l)
         end = start + ep_l - 1
@@ -97,13 +149,19 @@ def _apply_epoch_state(self, blk: int, start: int, end: int, idx: int, ep_len: i
         )
 
     async def _wait_for_next_head(self):
-        head_block = self.subtensor.get_current_block()
+        head_block = self._retry_rpc_call(
+            self.subtensor.get_current_block,
+            op_name="get_current_block",
+        )
         ep_l = self._epoch_len or self._discover_epoch_length()
         target_head = head_block - (head_block % ep_l) + ep_l
         label = "override" if self._override_active else "chain"
 
         while not self.should_exit:
-            blk = self.subtensor.get_current_block()
+            blk = self._retry_rpc_call(
+                self.subtensor.get_current_block,
+                op_name="get_current_block",
+            )
             if blk >= target_head:
                 return
             remain = max(0, target_head - blk)
@@ -157,7 +215,13 @@ def run(self):  # noqa: D401
 
         async def _loop():
             while not self.should_exit:
-                blk, start, end, idx, ep_len = self._epoch_snapshot()
+                try:
+                    blk, start, end, idx, ep_len = self._epoch_snapshot()
+                except Exception as err:
+                    bt.logging.error(f"[epoch] snapshot failed before wait: {err}")
+                    bt.logging.debug("[epoch] snapshot traceback:", exc_info=True)
+                    await asyncio.sleep(max(1.0, float(BLOCKTIME)))
+                    continue
 
                 next_head = start + ep_len
                 into = blk - start
@@ -207,7 +271,13 @@ async def _loop():
 
                 # head snapshot
                 self._epoch_len = None
-                blk2, start2, end2, idx2, ep_len2 = self._epoch_snapshot()
+                try:
+                    blk2, start2, end2, idx2, ep_len2 = self._epoch_snapshot()
+                except Exception as err:
+                    bt.logging.error(f"[epoch] snapshot failed at epoch boundary: {err}")
+                    bt.logging.debug("[epoch] snapshot traceback:", exc_info=True)
+                    await asyncio.sleep(max(1.0, float(BLOCKTIME)))
+                    continue
                 self._apply_epoch_state(blk2, start2, end2, idx2, ep_len2)
 
                 try:

tests/test_epoch_validator.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+import importlib
+import sys
+import types
+
+import pytest
+
+
+def _install_epoch_validator_test_stubs():
+    if "bittensor" not in sys.modules:
+        bittensor_mod = types.ModuleType("bittensor")
+
+        class _Logging:
+            def info(self, *_args, **_kwargs):
+                return None
+
+            def warning(self, *_args, **_kwargs):
+                return None
+
+            def success(self, *_args, **_kwargs):
+                return None
+
+            def error(self, *_args, **_kwargs):
+                return None
+
+            def debug(self, *_args, **_kwargs):
+                return None
+
+        bittensor_mod.logging = _Logging()
+        bittensor_mod.BLOCKTIME = 12
+        sys.modules["bittensor"] = bittensor_mod
+
+    if "metahash.base.validator" not in sys.modules:
+        base_validator_mod = types.ModuleType("metahash.base.validator")
+
+        class _BaseValidatorNeuron:
+            pass
+
+        base_validator_mod.BaseValidatorNeuron = _BaseValidatorNeuron
+        sys.modules["metahash.base.validator"] = base_validator_mod
+
+    if "metahash.validator.strategy" not in sys.modules:
+        strategy_mod = types.ModuleType("metahash.validator.strategy")
+
+        class _Strategy:
+            def __init__(self, *args, **kwargs):
+                return None
+
+            def compute_weights_bps(self, **_kwargs):
+                return {}
+
+        strategy_mod.Strategy = _Strategy
+        sys.modules["metahash.validator.strategy"] = strategy_mod
+
+
+_install_epoch_validator_test_stubs()
+sys.modules.pop("metahash.validator.epoch_validator", None)
+epoch_validator_mod = importlib.import_module("metahash.validator.epoch_validator")
+
+
+def _make_neuron():
+    neuron = object.__new__(epoch_validator_mod.EpochValidatorNeuron)
+    neuron.should_exit = False
+    neuron.config = types.SimpleNamespace(netuid=73, no_epoch=True, force_epoch=False)
+    neuron._epoch_len = None
+    neuron._override_active = False
+    neuron._bootstrapped = False
+    neuron._running = False
+    neuron.block = 0
+    neuron.step = 0
+    neuron.sync = lambda: None
+    neuron.concurrent_forward = _noop_forward
+    neuron.axon = None
+    return neuron
+
+
+async def _noop_forward():
+    return None
+
+
+def test_discover_epoch_length_retries_tempo_until_recovery(monkeypatch):
+    neuron = _make_neuron()
+    monkeypatch.setattr(epoch_validator_mod, "EPOCH_LENGTH_OVERRIDE", 0)
+    monkeypatch.setattr(epoch_validator_mod.random, "uniform", lambda _a, _b: 0.0)
+    monkeypatch.setattr(epoch_validator_mod.time, "sleep", lambda _s: None)
+
+    class _Subtensor:
+        def __init__(self):
+            self.tempo_calls = 0
+
+        def tempo(self, _netuid):
+            self.tempo_calls += 1
+            if self.tempo_calls < 3:
+                raise RuntimeError("Header was not found in the database")
+            return 360
+
+        def get_current_block(self):
+            return 721
+
+        def get_next_epoch_start_block(self, _netuid):
+            return 1081
+
+    st = _Subtensor()
+    neuron.subtensor = st
+
+    epoch_len = neuron._discover_epoch_length()
+    assert epoch_len == 361
+    assert st.tempo_calls == 3
+
+
+def test_retry_rpc_call_uses_exponential_backoff_capped_at_30s(monkeypatch):
+    neuron = _make_neuron()
+    sleep_calls: list[float] = []
+
+    monkeypatch.setattr(epoch_validator_mod.random, "uniform", lambda _a, _b: 0.0)
+
+    def _fake_sleep(seconds: float):
+        sleep_calls.append(seconds)
+        if len(sleep_calls) >= 6:
+            neuron.should_exit = True
+
+    monkeypatch.setattr(epoch_validator_mod.time, "sleep", _fake_sleep)
+
+    def _always_fail():
+        raise RuntimeError("rpc down")
+
+    with pytest.raises(RuntimeError, match="Stopping RPC retries"):
+        neuron._retry_rpc_call(_always_fail, op_name="tempo")
+
+    assert sleep_calls == [1.0, 2.0, 4.0, 8.0, 16.0, 30.0]
+
+
+def test_discover_epoch_length_falls_back_when_next_head_is_none(monkeypatch):
+    neuron = _make_neuron()
+    monkeypatch.setattr(epoch_validator_mod, "EPOCH_LENGTH_OVERRIDE", 0)
+    monkeypatch.setattr(epoch_validator_mod.random, "uniform", lambda _a, _b: 0.0)
+    monkeypatch.setattr(epoch_validator_mod.time, "sleep", lambda _s: None)
+
+    class _Subtensor:
+        def tempo(self, _netuid):
+            return 360
+
+        def get_current_block(self):
+            return 720
+
+        def get_next_epoch_start_block(self, _netuid):
+            return None
+
+    neuron.subtensor = _Subtensor()
+    epoch_len = neuron._discover_epoch_length()
+    assert epoch_len == 361
+
+
+def test_discover_epoch_length_uses_override_without_rpc(monkeypatch):
+    neuron = _make_neuron()
+    monkeypatch.setattr(epoch_validator_mod, "EPOCH_LENGTH_OVERRIDE", 7)
+
+    class _Subtensor:
+        def tempo(self, _netuid):  # pragma: no cover - should not be called
+            raise AssertionError("tempo() should not be called when override is set")
+
+        def get_current_block(self):  # pragma: no cover - should not be called
+            raise AssertionError("get_current_block() should not be called when override is set")
+
+        def get_next_epoch_start_block(self, _netuid):  # pragma: no cover - should not be called
+            raise AssertionError("get_next_epoch_start_block() should not be called when override is set")
+
+    neuron.subtensor = _Subtensor()
+    assert neuron._discover_epoch_length() == 7
+
+
+def test_run_continues_after_snapshot_error(monkeypatch):
+    neuron = _make_neuron()
+    calls = {"snap": 0, "sleep": 0}
+
+    def _snapshot():
+        calls["snap"] += 1
+        neuron.should_exit = True
+        raise RuntimeError("boom")
+
+    async def _fake_async_sleep(_seconds: float):
+        calls["sleep"] += 1
+        return None
+
+    monkeypatch.setattr(epoch_validator_mod.asyncio, "sleep", _fake_async_sleep)
+    neuron._epoch_snapshot = _snapshot
+
+    neuron.run()
+
+    assert calls["snap"] == 1
+    assert calls["sleep"] == 1
+    assert neuron._running is False