Skip to content

Commit a77df26

Browse files
author
Roman
committed
Merge branch 'staging' into feat/thewhaleking/remove-requirements-dir
# Conflicts: # requirements/torch.txt
2 parents e1ad0c3 + a70bee8 commit a77df26

File tree

2 files changed

+29
-4
lines changed

2 files changed

+29
-4
lines changed

bittensor/core/metagraph.py

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import asyncio
2+
import contextlib
23
import copy
34
import os
45
import pickle
@@ -11,6 +12,7 @@
1112
import numpy as np
1213
from async_substrate_interface.errors import SubstrateRequestException
1314
from numpy.typing import NDArray
15+
from packaging import version
1416

1517
from bittensor.core import settings
1618
from bittensor.core.chain_data import (
@@ -143,6 +145,27 @@ def latest_block_path(dir_path: str) -> str:
143145
return latest_file_full_path
144146

145147

148+
def safe_globals():
149+
"""
150+
Context manager to load torch files for version 2.6+
151+
"""
152+
if version.parse(torch.__version__).release < version.parse("2.6").release:
153+
return contextlib.nullcontext()
154+
155+
np_core = (
156+
np._core if version.parse(np.__version__) >= version.parse("2.0.0") else np.core
157+
)
158+
allow_list = [
159+
np_core.multiarray._reconstruct,
160+
np.ndarray,
161+
np.dtype,
162+
type(np.dtype(np.uint32)),
163+
np.dtypes.Float32DType,
164+
bytes,
165+
]
166+
return torch.serialization.safe_globals(allow_list)
167+
168+
146169
class MetagraphMixin(ABC):
147170
"""
148171
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":
11241147
"""
11251148

11261149
graph_file = latest_block_path(dir_path)
1127-
state_dict = torch.load(graph_file)
1150+
with safe_globals():
1151+
state_dict = torch.load(graph_file)
11281152
self.n = torch.nn.Parameter(state_dict["n"], requires_grad=False)
11291153
self.block = torch.nn.Parameter(state_dict["block"], requires_grad=False)
11301154
self.uids = torch.nn.Parameter(state_dict["uids"], requires_grad=False)
@@ -1256,7 +1280,8 @@ def load_from_path(self, dir_path: str) -> "MetagraphMixin":
12561280
try:
12571281
import torch as real_torch
12581282

1259-
state_dict = real_torch.load(graph_filename)
1283+
with safe_globals():
1284+
state_dict = real_torch.load(graph_filename)
12601285
for key in METAGRAPH_STATE_DICT_NDARRAY_KEYS:
12611286
state_dict[key] = state_dict[key].detach().numpy()
12621287
del real_torch

pyproject.toml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,10 +59,10 @@ dev = [
5959
"aioresponses==0.7.6",
6060
"factory-boy==3.3.0",
6161
"types-requests",
62-
"torch>=1.13.1,<2.6.0"
62+
"torch>=1.13.1,<3.0"
6363
]
6464
torch = [
65-
"torch>=1.13.1,<2.6.0"
65+
"torch>=1.13.1,<3.0"
6666
]
6767
cli = [
6868
"bittensor-cli>=9.0.2"

0 commit comments

Comments
 (0)