Migrate to pydantic >= 2.0

.github/workflows/run-tests.yml

@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [ '3.7', '3.8', '3.9', '3.10', '3.11' ]
+        python-version: [ '3.7', '3.8', '3.9', '3.10', '3.11', '3.12' ]
     timeout-minutes: 15
     steps:
       - uses: actions/checkout@v3

benchmarks/benchmark_optimizer.py

@@ -98,7 +98,7 @@ def run_trainer(batch_size: int, batch_time: float, client_mode: bool, verbose:
             grad_scaler = hivemind.GradScaler()
         else:
             # check that hivemind.Optimizer supports regular PyTorch grad scaler as well
-            grad_scaler = torch.cuda.amp.GradScaler(enabled=args.use_amp)
+            grad_scaler = torch.amp.GradScaler(enabled=args.use_amp)
 
         prev_time = time.perf_counter()
 
@@ -107,7 +107,7 @@ def run_trainer(batch_size: int, batch_time: float, client_mode: bool, verbose:
 
             batch = torch.randint(0, len(X_train), (batch_size,))
 
-            with torch.cuda.amp.autocast() if args.use_amp else nullcontext():
+            with torch.amp.autocast() if args.use_amp else nullcontext():
                 loss = F.cross_entropy(model(X_train[batch].to(args.device)), y_train[batch].to(args.device))
                 grad_scaler.scale(loss).backward()
 

examples/albert/requirements.txt

@@ -1,7 +1,7 @@
 transformers~=4.6
 datasets~=1.5
 torch_optimizer==0.1.0
-wandb==0.10.26
+wandb==0.17.1
 sentencepiece
 requests
 nltk==3.6.7

examples/albert/utils.py

@@ -1,6 +1,6 @@
 from typing import Dict, List, Tuple
 
-from pydantic import BaseModel, StrictFloat, confloat, conint
+from pydantic.v1 import BaseModel, StrictFloat, confloat, conint
 
 from hivemind.dht.crypto import RSASignatureValidator
 from hivemind.dht.schema import BytesWithPublicKey, SchemaValidator

hivemind/dht/schema.py

@@ -1,7 +1,7 @@
 import re
 from typing import Any, Dict, Optional, Type
 
-import pydantic
+import pydantic.v1 as pydantic
 
 from hivemind.dht.crypto import RSASignatureValidator
 from hivemind.dht.protocol import DHTProtocol

hivemind/optim/progress_tracker.py

@@ -6,7 +6,7 @@
 from typing import Dict, Optional
 
 import numpy as np
-from pydantic import BaseModel, StrictBool, StrictFloat, confloat, conint
+from pydantic.v1 import BaseModel, StrictBool, StrictFloat, confloat, conint
 
 from hivemind.dht import DHT
 from hivemind.dht.schema import BytesWithPublicKey, RSASignatureValidator, SchemaValidator

requirements.txt

@@ -12,5 +12,5 @@ configargparse>=1.2.3
 py-multihash>=0.2.3
 multiaddr @ git+https://github.com/multiformats/py-multiaddr.git@e01dbd38f2c0464c0f78b556691d655265018cce
 cryptography>=3.4.6
-pydantic>=1.8.1,<2.0
+pydantic>=2.5.3
 packaging>=20.9

tests/test_dht_schema.py

@@ -2,7 +2,7 @@
 from typing import Dict
 
 import pytest
-from pydantic import BaseModel, StrictInt, conint
+from pydantic.v1 import BaseModel, StrictInt, conint
 
 import hivemind
 from hivemind.dht.node import DHTNode
@@ -12,9 +12,9 @@
 
 
 class SampleSchema(BaseModel):
-    experiment_name: bytes
-    n_batches: Dict[bytes, conint(ge=0, strict=True)]
-    signed_data: Dict[BytesWithPublicKey, bytes]
+    experiment_name: bytes = None
+    n_batches: Dict[bytes, conint(ge=0, strict=True)] = None
+    signed_data: Dict[BytesWithPublicKey, bytes] = None
 
 
 @pytest.fixture
@@ -95,10 +95,10 @@ async def test_expecting_public_keys(dht_nodes_with_schema):
 @pytest.mark.asyncio
 async def test_keys_outside_schema(dht_nodes_with_schema):
     class Schema(BaseModel):
-        some_field: StrictInt
+        some_field: StrictInt = None
 
     class MergedSchema(BaseModel):
-        another_field: StrictInt
+        another_field: StrictInt = None
 
     for allow_extra_keys in [False, True]:
         validator = SchemaValidator(Schema, allow_extra_keys=allow_extra_keys)
@@ -122,7 +122,7 @@ class MergedSchema(BaseModel):
 @pytest.mark.asyncio
 async def test_prefix():
     class Schema(BaseModel):
-        field: StrictInt
+        field: StrictInt = None
 
     validator = SchemaValidator(Schema, allow_extra_keys=False, prefix="prefix")
 
@@ -154,11 +154,11 @@ def validate(self, record: DHTRecord) -> bool:
     assert not alice.protocol.record_validator.merge_with(second_validator)
 
     class SecondSchema(BaseModel):
-        some_field: StrictInt
-        another_field: str
+        some_field: StrictInt = 0
+        another_field: str = None
 
     class ThirdSchema(BaseModel):
-        another_field: StrictInt  # Allow it to be a StrictInt as well
+        another_field: StrictInt = 0  # Allow it to be a StrictInt as well
 
     for schema in [SecondSchema, ThirdSchema]:
         new_validator = SchemaValidator(schema, allow_extra_keys=False)

tests/test_dht_validation.py

@@ -2,7 +2,7 @@
 from typing import Dict
 
 import pytest
-from pydantic import BaseModel, StrictInt
+from pydantic.v1 import BaseModel, StrictInt
 
 import hivemind
 from hivemind.dht.crypto import RSASignatureValidator

tests/test_p2p_daemon.py

@@ -39,12 +39,12 @@ async def test_daemon_killed_on_del():
 
 @pytest.mark.asyncio
 async def test_startup_error_message():
-    with pytest.raises(P2PDaemonError, match=r"(?i)Failed to connect to bootstrap peers"):
+    with pytest.raises(P2PDaemonError, match=r"(?i)Failed to connect to bootstrap peers|Daemon failed to start"):
         await P2P.create(
             initial_peers=[f"/ip4/127.0.0.1/tcp/{get_free_port()}/p2p/QmdaK4LUeQaKhqSFPRu9N7MvXUEWDxWwtCvPrS444tCgd1"]
         )
 
-    with pytest.raises(P2PDaemonError, match=r"Daemon failed to start in .+ seconds"):
+    with pytest.raises(P2PDaemonError, match=r"Daemon failed to start|error accepting connection"):
         await P2P.create(startup_timeout=0.01)  # Test that startup_timeout works
 
 

tests/test_utils/p2p_daemon.py

@@ -5,16 +5,24 @@
 import time
 import uuid
 from contextlib import asynccontextmanager, suppress
+
+# from pathlib import Path
 from typing import NamedTuple
 
 from multiaddr import Multiaddr, protocols
-from pkg_resources import resource_filename
 
 from hivemind.p2p.p2p_daemon_bindings.p2pclient import Client
 
 from test_utils.networking import get_free_port
 
 TIMEOUT_DURATION = 30  # seconds
+
+# hivemind_dir = Path(__file__).parent / "hivemind" / "hivemind_cli"
+
+# P2PD_PATH = str(hivemind_dir / "p2pd")
+
+from pkg_resources import resource_filename
+
 P2PD_PATH = resource_filename("hivemind", "hivemind_cli/p2pd")