Merge branch 'staging' into patch-2

Author: basfroman

bittensor/core/async_subtensor.py

@@ -1687,6 +1687,38 @@ async def get_neuron_for_pubkey_and_subnet(
                 reuse_block=reuse_block,
             )
 
+    async def get_next_epoch_start_block(
+        self,
+        netuid: int,
+        block: Optional[int] = None,
+        block_hash: Optional[str] = None,
+        reuse_block: bool = False,
+    ) -> Optional[int]:
+        """
+        Calculates the first block number of the next epoch for the given subnet.
+
+        If `block` is not provided, the current chain block will be used. Epochs are
+        determined based on the subnet's tempo (i.e., blocks per epoch). The result
+        is the block number at which the next epoch will begin.
+
+        Args:
+            netuid (int): The unique identifier of the subnet.
+            block (Optional[int], optional): The reference block to calculate from.
+                If None, uses the current chain block height.
+            block_hash (Optional[int]): The blockchain block number at which to perform the query.
+            reuse_block (bool): Whether to reuse the last-used blockchain block hash.
+
+
+        Returns:
+            int: The block number at which the next epoch will start.
+        """
+        block_hash = await self.determine_block_hash(block, block_hash, reuse_block)
+        if not block_hash and reuse_block:
+            block_hash = self.substrate.last_block_hash
+        block = await self.substrate.get_block_number(block_hash=block_hash)
+        tempo = await self.tempo(netuid=netuid, block_hash=block_hash)
+        return (((block // tempo) + 1) * tempo) + 1 if tempo else None
+
     async def get_owned_hotkeys(
         self,
         coldkey_ss58: str,

bittensor/core/metagraph.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextlib
 import copy
 import os
 import pickle
@@ -11,6 +12,7 @@
 import numpy as np
 from async_substrate_interface.errors import SubstrateRequestException
 from numpy.typing import NDArray
+from packaging import version
 
 from bittensor.core import settings
 from bittensor.core.chain_data import (
@@ -143,6 +145,27 @@ def latest_block_path(dir_path: str) -> str:
         return latest_file_full_path
 
 
+def safe_globals():
+    """
+    Context manager to load torch files for version 2.6+
+    """
+    if version.parse(torch.__version__).release < version.parse("2.6").release:
+        return contextlib.nullcontext()
+
+    np_core = (
+        np._core if version.parse(np.__version__) >= version.parse("2.0.0") else np.core
+    )
+    allow_list = [
+        np_core.multiarray._reconstruct,
+        np.ndarray,
+        np.dtype,
+        type(np.dtype(np.uint32)),
+        np.dtypes.Float32DType,
+        bytes,
+    ]
+    return torch.serialization.safe_globals(allow_list)
+
+
 class MetagraphMixin(ABC):
     """
     The metagraph class is a core component of the Bittensor network, representing the neural graph that forms the
@@ -1124,7 +1147,8 @@ def load_from_path(self, dir_path: str) -> "MetagraphMixin":
         """
 
         graph_file = latest_block_path(dir_path)
-        state_dict = torch.load(graph_file)
+        with safe_globals():
+            state_dict = torch.load(graph_file)
         self.n = torch.nn.Parameter(state_dict["n"], requires_grad=False)
         self.block = torch.nn.Parameter(state_dict["block"], requires_grad=False)
         self.uids = torch.nn.Parameter(state_dict["uids"], requires_grad=False)
@@ -1256,7 +1280,8 @@ def load_from_path(self, dir_path: str) -> "MetagraphMixin":
             try:
                 import torch as real_torch
 
-                state_dict = real_torch.load(graph_filename)
+                with safe_globals():
+                    state_dict = real_torch.load(graph_filename)
                 for key in METAGRAPH_STATE_DICT_NDARRAY_KEYS:
                     state_dict[key] = state_dict[key].detach().numpy()
                 del real_torch

bittensor/core/subtensor.py

@@ -1300,6 +1300,28 @@ def get_neuron_for_pubkey_and_subnet(
 
         return NeuronInfo.from_dict(result)
 
+    def get_next_epoch_start_block(
+        self, netuid: int, block: Optional[int] = None
+    ) -> Optional[int]:
+        """
+        Calculates the first block number of the next epoch for the given subnet.
+
+        If `block` is not provided, the current chain block will be used. Epochs are
+        determined based on the subnet's tempo (i.e., blocks per epoch). The result
+        is the block number at which the next epoch will begin.
+
+        Args:
+            netuid (int): The unique identifier of the subnet.
+            block (Optional[int], optional): The reference block to calculate from.
+                If None, uses the current chain block height.
+
+        Returns:
+            int: The block number at which the next epoch will start.
+        """
+        block = block or self.block
+        tempo = self.tempo(netuid=netuid, block=block)
+        return (((block // tempo) + 1) * tempo) + 1 if tempo else None
+
     def get_owned_hotkeys(
         self,
         coldkey_ss58: str,

pyproject.toml

@@ -59,10 +59,10 @@ dev = [
     "aioresponses==0.7.6",
     "factory-boy==3.3.0",
     "types-requests",
-    "torch>=1.13.1,<2.6.0"
+    "torch>=1.13.1,<3.0"
 ]
 torch = [
-    "torch>=1.13.1,<2.6.0"
+    "torch>=1.13.1,<3.0"
 ]
 cli = [
     "bittensor-cli>=9.0.2"

requirements/torch.txt

@@ -1 +1 @@
-torch>=1.13.1,<2.6.0
+torch>=1.13.1,<3.0