feat(network): add timeouts for subtensor operations to prevent indefinite hangs and improve error handling

Author: erfanMhi

grail/infrastructure/network.py

@@ -12,6 +12,10 @@
 
 logger = logging.getLogger(__name__)
 
+# Timeouts for subtensor lifecycle operations to prevent indefinite hangs
+CLOSE_TIMEOUT = float(os.getenv("BT_CLOSE_TIMEOUT", "60.0"))
+INIT_TIMEOUT = float(os.getenv("BT_INIT_TIMEOUT", "120.0"))
+
 
 class ResilientSubtensor:
     """
@@ -104,30 +108,40 @@ def _should_open_circuit(self) -> bool:
         circuit_threshold = object.__getattribute__(self, "_circuit_threshold")
         return failure_count >= circuit_threshold
 
+    async def _close_subtensor_safe(self, subtensor: Any, timeout: float = CLOSE_TIMEOUT) -> None:
+        """Close subtensor with timeout, logging but not raising on failure."""
+        if not hasattr(subtensor, "close"):
+            return
+        try:
+            if asyncio.iscoroutinefunction(subtensor.close):
+                await asyncio.wait_for(subtensor.close(), timeout=timeout)
+            else:
+                subtensor.close()
+        except asyncio.TimeoutError:
+            logger.warning("⏱️ Timeout closing subtensor after %.0fs", timeout)
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            logger.warning("Failed to close subtensor: %s", exc)
+
     async def _restart_subtensor(self) -> None:
-        """Restart subtensor connection."""
+        """Restart subtensor connection with timeouts to prevent indefinite hangs."""
         logger.warning("🔄 Restarting subtensor connection...")
         subtensor = object.__getattribute__(self, "_subtensor")
-        network = (
-            subtensor.network
-            if hasattr(subtensor, "network")
-            else os.getenv("BT_NETWORK", "finney")
-        )
+        network = getattr(subtensor, "network", None) or os.getenv("BT_NETWORK", "finney")
 
-        # Close old subtensor to prevent resource leaks
-        try:
-            if hasattr(subtensor, "close"):
-                if asyncio.iscoroutinefunction(subtensor.close):
-                    await subtensor.close()
-                else:
-                    subtensor.close()
-                logger.debug("Closed old subtensor connection")
-        except Exception as exc:
-            logger.warning("Failed to close old subtensor: %s", exc)
+        await self._close_subtensor_safe(subtensor)
 
         new_subtensor = bt.async_subtensor(network=network)
-        await new_subtensor.initialize()
+        try:
+            await asyncio.wait_for(new_subtensor.initialize(), timeout=INIT_TIMEOUT)
+        except (asyncio.TimeoutError, asyncio.CancelledError):
+            logger.error("❌ Subtensor init failed, cleaning up")
+            await self._close_subtensor_safe(new_subtensor, timeout=5.0)
+            raise
+
         object.__setattr__(self, "_subtensor", new_subtensor)
+        object.__setattr__(self, "_last_call_timestamp", time.time())
         self._reset_circuit_breaker()
         logger.info("✅ Subtensor connection restarted")
 
@@ -136,13 +150,8 @@ async def restart(self) -> None:
         await self._restart_subtensor()
 
     async def close(self) -> None:
-        """Close the underlying subtensor connection."""
-        subtensor = object.__getattribute__(self, "_subtensor")
-        if hasattr(subtensor, "close"):
-            if asyncio.iscoroutinefunction(subtensor.close):
-                await subtensor.close()
-            else:
-                subtensor.close()
+        """Close the underlying subtensor connection with timeout."""
+        await self._close_subtensor_safe(object.__getattribute__(self, "_subtensor"))
 
     async def _handle_circuit_open(self, method_name: str, args: tuple) -> Any:
         """Handle method call when circuit breaker is open."""
@@ -199,7 +208,8 @@ def _handle_all_retries_failed(self, method_name: str, args: tuple, retries: int
 
         if self._should_open_circuit():
             self._open_circuit_breaker()
-            asyncio.create_task(self._restart_subtensor())
+            task = asyncio.create_task(self._restart_subtensor(), name="subtensor_restart")
+            task.add_done_callback(self._on_restart_complete)
 
         # Try to return cached metagraph as last resort
         if method_name == "metagraph" and args:
@@ -211,6 +221,14 @@ def _handle_all_retries_failed(self, method_name: str, args: tuple, retries: int
         logger.error("❌ %s failed after %d attempts", method_name, retries)
         raise TimeoutError(f"{method_name} failed after {retries} attempts")
 
+    def _on_restart_complete(self, task: asyncio.Task) -> None:
+        """Log errors from background restart task."""
+        try:
+            if (exc := task.exception()) is not None:
+                logger.error("⚠️ Background subtensor restart failed: %s", exc)
+        except (asyncio.CancelledError, asyncio.InvalidStateError):
+            pass
+
     async def _call_with_retry(
         self, method_name: str, method: Any, args: tuple, kwargs: dict
     ) -> Any:
@@ -227,21 +245,20 @@ async def _call_with_retry(
         if method_name == "metagraph":
             timeout = timeout * 2
 
-        # Double timeout if connection has been idle for 20+ seconds
-        # Research-based threshold:
-        # - Bittensor WebSocket auto-closes after 10s inactivity
-        # - Substrate layer closes after ~60s inactivity
-        # - 20s catches stale connections early without false positives
-        # - Critical for upload worker (40-300s idle during R2 uploads)
+        # Restart if connection idle for 60s+ (WebSocket likely stale)
         last_call_timestamp = object.__getattribute__(self, "_last_call_timestamp")
         idle_duration = time.time() - last_call_timestamp
         if idle_duration > 60.0:
             logger.warning(
-                "⏰ Connection idle for %.1fs, restarting subtensor and doubling timeout for %s",
-                idle_duration,
-                method_name,
+                "⏰ Connection idle for %.1fs, restarting for %s", idle_duration, method_name
             )
-            await self._restart_subtensor()
+            try:
+                await self._restart_subtensor()
+            except asyncio.CancelledError:
+                raise
+            except Exception as exc:
+                logger.warning("⚠️ Restart failed (%s), doubling timeout", exc)
+                timeout *= 2
 
         # Retry loop
         for retry in range(retries):
@@ -380,12 +397,27 @@ async def create_subtensor(*, resilient: bool = True) -> bt.subtensor | Resilien
         logger.info("Connecting to Bittensor %s (network=%s)", label, network)
         subtensor = bt.async_subtensor(network=network)
 
-    await await_with_stall_log(
-        subtensor.initialize(),
-        label="subtensor.initialize",
-        threshold_seconds=120.0,
-        log=logger,
-    )
+    try:
+        await asyncio.wait_for(
+            await_with_stall_log(
+                subtensor.initialize(),
+                label="subtensor.initialize",
+                threshold_seconds=60.0,
+                log=logger,
+            ),
+            timeout=INIT_TIMEOUT,
+        )
+    except asyncio.TimeoutError:
+        logger.error("❌ Subtensor init timed out after %.0fs", INIT_TIMEOUT)
+        if hasattr(subtensor, "close"):
+            try:
+                if asyncio.iscoroutinefunction(subtensor.close):
+                    await asyncio.wait_for(subtensor.close(), timeout=5.0)
+                else:
+                    subtensor.close()
+            except Exception:
+                pass
+        raise
 
     if resilient:
         # Wrap with resilience layer for production use