Upgrade pytest to 8.3.5 (#646)

* Put non-forked tests last

* Upgrade pytest dependency

* Remove timeouts from pytest

* Replace Python 3.8 with 3.12 in setup.py

* Remove the image from the app

* Add timeouts to process cleanup in test_cli_run_server_identity_path

* Clean up DHT and averager processes in test_adaptive_compression

Author: mryab

modal_ci.py

@@ -51,7 +51,7 @@
 )
 
 
-app = modal.App("hivemind-ci", image=image)
+app = modal.App("hivemind-ci")
 
 codecov_secret = modal.Secret.from_dict(
     {
@@ -99,9 +99,6 @@ def run_tests():
             "-v",
             "-n",
             "8",
-            "--dist",
-            "worksteal",
-            "--timeout=60",
             "tests",
         ],
         check=True,
@@ -120,7 +117,6 @@ def run_codecov():
             "hivemind",
             "--cov-config=pyproject.toml",
             "-v",
-            "--timeout=60",
             "tests",
         ],
         check=True,

requirements-dev.txt

@@ -1,4 +1,4 @@
-pytest==6.2.5  # see https://github.com/pytest-dev/pytest/issues/9621
+pytest==8.3.5
 pytest-forked
 pytest-asyncio==0.16.0
 pytest-cov

setup.py

@@ -183,10 +183,10 @@ def run(self):
         "Intended Audience :: Science/Research",
         "License :: OSI Approved :: MIT License",
         "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
         "Programming Language :: Python :: 3.10",
         "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
         "Topic :: Scientific/Engineering",
         "Topic :: Scientific/Engineering :: Mathematics",
         "Topic :: Scientific/Engineering :: Artificial Intelligence",

tests/test_compression.py

@@ -221,6 +221,8 @@ def test_adaptive_compression():
         greater_equal=STATE_FP16,
     )
 
+    dht_instances = launch_dht_instances(2)
+
     averager1 = hivemind.TrainingAverager(
         opt=torch.optim.Adam(make_params()),
         average_parameters=True,
@@ -232,7 +234,7 @@ def test_adaptive_compression():
         target_group_size=2,
         part_size_bytes=5_000,
         start=True,
-        dht=hivemind.DHT(start=True),
+        dht=dht_instances[0],
     )
 
     averager2 = hivemind.TrainingAverager(
@@ -246,7 +248,7 @@ def test_adaptive_compression():
         target_group_size=2,
         part_size_bytes=5_000,
         start=True,
-        dht=hivemind.DHT(initial_peers=averager1.dht.get_visible_maddrs(), start=True),
+        dht=dht_instances[1],
     )
 
     futures = [averager1.step(wait=False), averager2.step(wait=False)]
@@ -266,3 +268,6 @@ def test_adaptive_compression():
     assert STATE_FP16.mp_counter.value == len([tensor for tensor in state_tensors if tensor.numel() >= 500])
     assert STATE_FP32.mp_counter.value == len([tensor for tensor in state_tensors if tensor.numel() < 500])
     assert STATE_FP16.mp_part_size.value == STATE_FP32.mp_part_size.value == 0  # not partitioned
+
+    for instance in [averager1, averager2] + dht_instances:
+        instance.shutdown()

tests/test_dht.py

@@ -12,20 +12,6 @@
 from test_utils.networking import get_free_port
 
 
-@pytest.mark.asyncio
-async def test_startup_error():
-    with pytest.raises(hivemind.p2p.P2PDaemonError, match=r"(?i)Failed to connect to bootstrap peers"):
-        hivemind.DHT(
-            initial_peers=[f"/ip4/127.0.0.1/tcp/{get_free_port()}/p2p/QmdaK4LUeQaKhqSFPRu9N7MvXUEWDxWwtCvPrS444tCgd1"],
-            start=True,
-        )
-
-    dht = hivemind.DHT(start=True, await_ready=False)
-    with pytest.raises(concurrent.futures.TimeoutError):
-        dht.wait_until_ready(timeout=0.01)
-    dht.shutdown()
-
-
 @pytest.mark.forked
 def test_get_store(n_peers=10):
     peers = launch_dht_instances(n_peers)
@@ -122,3 +108,17 @@ async def test_dht_get_visible_maddrs():
 
     assert dht.get_visible_maddrs() == [dummy_endpoint.encapsulate(f"/p2p/{p2p.peer_id}")]
     dht.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_startup_error():
+    with pytest.raises(hivemind.p2p.P2PDaemonError, match=r"(?i)Failed to connect to bootstrap peers"):
+        hivemind.DHT(
+            initial_peers=[f"/ip4/127.0.0.1/tcp/{get_free_port()}/p2p/QmdaK4LUeQaKhqSFPRu9N7MvXUEWDxWwtCvPrS444tCgd1"],
+            start=True,
+        )
+
+    dht = hivemind.DHT(start=True, await_ready=False)
+    with pytest.raises(concurrent.futures.TimeoutError):
+        dht.wait_until_ready(timeout=0.01)
+    dht.shutdown()

tests/test_dht_crypto.py

@@ -88,24 +88,6 @@ def get_signed_record(conn: mp.connection.Connection) -> DHTRecord:
     return record
 
 
-def test_signing_in_different_process():
-    parent_conn, child_conn = mp.Pipe()
-    process = mp.Process(target=get_signed_record, args=[child_conn])
-    process.start()
-
-    validator = RSASignatureValidator()
-    parent_conn.send(validator)
-
-    record = DHTRecord(
-        key=b"key", subkey=b"subkey" + validator.local_public_key, value=b"value", expiration_time=get_dht_time() + 10
-    )
-    parent_conn.send(record)
-
-    signed_record = parent_conn.recv()
-    assert b"[signature:" in signed_record.value
-    assert validator.validate(signed_record)
-
-
 @pytest.mark.forked
 @pytest.mark.asyncio
 async def test_dhtnode_signatures():
@@ -134,3 +116,21 @@ async def test_dhtnode_signatures():
     store_ok = await mallory.store(key, b"updated_fake_value", hivemind.get_dht_time() + 10, subkey=subkey)
     assert not store_ok
     assert (await alice.get(key, latest=True)).value[subkey].value == b"updated_true_value"
+
+
+def test_signing_in_different_process():
+    parent_conn, child_conn = mp.Pipe()
+    process = mp.Process(target=get_signed_record, args=[child_conn])
+    process.start()
+
+    validator = RSASignatureValidator()
+    parent_conn.send(validator)
+
+    record = DHTRecord(
+        key=b"key", subkey=b"subkey" + validator.local_public_key, value=b"value", expiration_time=get_dht_time() + 10
+    )
+    parent_conn.send(record)
+
+    signed_record = parent_conn.recv()
+    assert b"[signature:" in signed_record.value
+    assert validator.validate(signed_record)

tests/test_dht_experts.py

@@ -126,6 +126,37 @@ def test_dht_single_node():
     node.shutdown()
 
 
+@pytest.mark.forked
+@pytest.mark.asyncio
+async def test_negative_caching(n_peers=10):
+    dht_kwargs = {"cache_locally": False}
+
+    peers = [hivemind.DHT(start=True, **dht_kwargs)]
+    initial_peers = peers[0].get_visible_maddrs()
+    peers += [hivemind.DHT(initial_peers=initial_peers, start=True, **dht_kwargs) for _ in range(n_peers - 1)]
+
+    writer_peer = random.choice(peers)
+    assert all(declare_experts(writer_peer, ["ffn.1.2.3", "ffn.3.4.5"], get_dht_time() + 30).values())
+
+    neighbors = sum([peer.get_visible_maddrs() for peer in random.sample(peers, min(3, len(peers)))], [])
+    neg_caching_peer = hivemind.DHT(initial_peers=neighbors, start=True, **dht_kwargs)
+    beam_search = MoEBeamSearcher(neg_caching_peer, uid_prefix="ffn.", grid_size=(10, 10, 10), negative_caching=True)
+    # get prefixes by the peer with negative caching. Cache "no data" entries for ffn.0.*, ffn.2.*, ffn.4.*, ffn.5.*
+    assert len(beam_search.get_initial_beam(scores=[0.1, 0.2, 0.3, 0.4, 0.5, 0.6], beam_size=3)) == 2
+
+    node = await DHTNode.create(initial_peers=neighbors)
+    fetched = await asyncio.gather(*(node.get(f"ffn.{i}.") for i in range(10)))
+    for i in range(6):
+        assert fetched[i] is not None, f"node should have cached ffn.{i}."
+    for i in range(6, len(fetched)):
+        assert fetched[i] is None, f"node shouldn't have cached ffn.{i}."
+
+    await node.shutdown()
+    neg_caching_peer.shutdown()
+    for peer in peers:
+        peer.shutdown()
+
+
 def test_uid_patterns():
     valid_experts = [
         "expert.1",
@@ -188,34 +219,3 @@ def test_uid_patterns():
         assert not is_valid_uid(uid), f"UID {uid} is not valid, but was perceived as valid"
     for pfx in invalid_prefixes:
         assert not is_valid_prefix(pfx), f"Prefix {pfx} is not valid, but was perceived as valid"
-
-
-@pytest.mark.forked
-@pytest.mark.asyncio
-async def test_negative_caching(n_peers=10):
-    dht_kwargs = {"cache_locally": False}
-
-    peers = [hivemind.DHT(start=True, **dht_kwargs)]
-    initial_peers = peers[0].get_visible_maddrs()
-    peers += [hivemind.DHT(initial_peers=initial_peers, start=True, **dht_kwargs) for _ in range(n_peers - 1)]
-
-    writer_peer = random.choice(peers)
-    assert all(declare_experts(writer_peer, ["ffn.1.2.3", "ffn.3.4.5"], get_dht_time() + 30).values())
-
-    neighbors = sum([peer.get_visible_maddrs() for peer in random.sample(peers, min(3, len(peers)))], [])
-    neg_caching_peer = hivemind.DHT(initial_peers=neighbors, start=True, **dht_kwargs)
-    beam_search = MoEBeamSearcher(neg_caching_peer, uid_prefix="ffn.", grid_size=(10, 10, 10), negative_caching=True)
-    # get prefixes by the peer with negative caching. Cache "no data" entries for ffn.0.*, ffn.2.*, ffn.4.*, ffn.5.*
-    assert len(beam_search.get_initial_beam(scores=[0.1, 0.2, 0.3, 0.4, 0.5, 0.6], beam_size=3)) == 2
-
-    node = await DHTNode.create(initial_peers=neighbors)
-    fetched = await asyncio.gather(*(node.get(f"ffn.{i}.") for i in range(10)))
-    for i in range(6):
-        assert fetched[i] is not None, f"node should have cached ffn.{i}."
-    for i in range(6, len(fetched)):
-        assert fetched[i] is None, f"node shouldn't have cached ffn.{i}."
-
-    await node.shutdown()
-    neg_caching_peer.shutdown()
-    for peer in peers:
-        peer.shutdown()

tests/test_dht_validation.py

@@ -29,6 +29,34 @@ def validators_for_app():
     }
 
 
+@pytest.mark.forked
+def test_dht_add_validators(validators_for_app):
+    # One app may create a DHT with its validators
+    dht = hivemind.DHT(start=False, record_validators=validators_for_app["A"])
+
+    # While the DHT process is not started, you can't send a command to append new validators
+    with pytest.raises(RuntimeError):
+        dht.add_validators(validators_for_app["B"])
+    dht.run_in_background(await_ready=True)
+
+    # After starting the process, other apps may add new validators to the existing DHT
+    dht.add_validators(validators_for_app["B"])
+
+    assert dht.store("field_a", b"bytes_value", hivemind.get_dht_time() + 10)
+    assert dht.get("field_a", latest=True).value == b"bytes_value"
+
+    assert not dht.store("field_a", 666, hivemind.get_dht_time() + 10)
+    assert dht.get("field_a", latest=True).value == b"bytes_value"
+
+    local_public_key = validators_for_app["A"][0].local_public_key
+    assert dht.store("field_b", 777, hivemind.get_dht_time() + 10, subkey=local_public_key)
+    dictionary = dht.get("field_b", latest=True).value
+    assert len(dictionary) == 1 and dictionary[local_public_key].value == 777
+
+    assert not dht.store("unknown_key", 666, hivemind.get_dht_time() + 10)
+    assert dht.get("unknown_key", latest=True) is None
+
+
 def test_composite_validator(validators_for_app):
     validator = CompositeValidator(validators_for_app["A"])
     assert [type(item) for item in validator._validators] == [SchemaValidator, RSASignatureValidator]
@@ -63,31 +91,3 @@ def test_composite_validator(validators_for_app):
     assert signed_record.value.count(b"[signature:") == 0
     # Expect failed validation since `unknown_key` is not a part of any schema
     assert not validator.validate(signed_record)
-
-
-@pytest.mark.forked
-def test_dht_add_validators(validators_for_app):
-    # One app may create a DHT with its validators
-    dht = hivemind.DHT(start=False, record_validators=validators_for_app["A"])
-
-    # While the DHT process is not started, you can't send a command to append new validators
-    with pytest.raises(RuntimeError):
-        dht.add_validators(validators_for_app["B"])
-    dht.run_in_background(await_ready=True)
-
-    # After starting the process, other apps may add new validators to the existing DHT
-    dht.add_validators(validators_for_app["B"])
-
-    assert dht.store("field_a", b"bytes_value", hivemind.get_dht_time() + 10)
-    assert dht.get("field_a", latest=True).value == b"bytes_value"
-
-    assert not dht.store("field_a", 666, hivemind.get_dht_time() + 10)
-    assert dht.get("field_a", latest=True).value == b"bytes_value"
-
-    local_public_key = validators_for_app["A"][0].local_public_key
-    assert dht.store("field_b", 777, hivemind.get_dht_time() + 10, subkey=local_public_key)
-    dictionary = dht.get("field_b", latest=True).value
-    assert len(dictionary) == 1 and dictionary[local_public_key].value == 777
-
-    assert not dht.store("unknown_key", 666, hivemind.get_dht_time() + 10)
-    assert dht.get("unknown_key", latest=True) is None

tests/test_start_server.py

@@ -9,6 +9,15 @@
 from hivemind.moe.server import background_server
 
 
+def cleanup_process(process, timeout=5):
+    try:
+        process.terminate()
+        process.wait(timeout=timeout)  # Add timeout to wait
+    except:
+        process.kill()
+        process.wait(timeout=timeout)
+
+
 @pytest.mark.xfail(reason="Flaky test", strict=False)
 def test_background_server_identity_path():
     with TemporaryDirectory() as tempdir:
@@ -96,10 +105,5 @@ def test_cli_run_server_identity_path():
         assert addrs_1 != addrs_3
         assert addrs_2 != addrs_3
 
-        server_1_proc.terminate()
-        server_2_proc.terminate()
-        server_3_proc.terminate()
-
-        server_1_proc.wait()
-        server_2_proc.wait()
-        server_3_proc.wait()
+        for p in [server_1_proc, server_2_proc, server_3_proc]:
+            cleanup_process(p)