Overrides copy and deep copy for the metagraph (#2523)

Author: ibraheem-abe

bittensor/core/metagraph.py

@@ -1,3 +1,4 @@
+import copy
 import os
 import pickle
 import typing
@@ -911,6 +912,30 @@ def load_from_path(self, dir_path: str) -> "Metagraph":
             state files within it are accurate and consistent with the expected metagraph structure.
         """
 
+    def __deepcopy__(self, memo):
+        cls = self.__class__
+        new_instance = cls.__new__(cls)
+        memo[id(self)] = new_instance
+
+        for key, value in self.__dict__.items():
+            if key == "subtensor":
+                setattr(new_instance, key, None)
+            else:
+                setattr(new_instance, key, copy.deepcopy(value, memo))
+
+        return new_instance
+
+    def __copy__(self):
+        cls = self.__class__
+        new_instance = cls.__new__(cls)
+
+        for key, value in self.__dict__.items():
+            if key == "subtensor":
+                setattr(new_instance, key, None)
+            else:
+                setattr(new_instance, key, value)
+        return new_instance
+
 
 BaseClass: Union["torch.nn.Module", object] = torch.nn.Module if use_torch() else object
 """

tests/unit_tests/test_metagraph.py

@@ -20,6 +20,7 @@
 
 import numpy as np
 import pytest
+import copy
 
 from bittensor.core import settings
 from bittensor.core.metagraph import Metagraph
@@ -174,3 +175,97 @@ def test_sync_warning_cases(block, test_id, metagraph_instance, mock_subtensor,
     assert (
         expected_message in caplog.text
     ), f"Test ID: {test_id} - Expected warning message not found in Loguru sink."
+
+
+def test_deepcopy(mock_environment):
+    subtensor, neurons = mock_environment
+    metagraph = Metagraph(1, sync=False)
+    metagraph.neurons = neurons
+    metagraph.subtensor = subtensor
+
+    # Do a deep copy
+    copied_metagraph = copy.deepcopy(metagraph)
+
+    # Check that the subtensor attribute is None
+    assert copied_metagraph.subtensor is None
+
+    # Check that other attributes are copied correctly
+    assert copied_metagraph.n == metagraph.n
+    assert copied_metagraph.block == metagraph.block
+    assert np.array_equal(copied_metagraph.uids, metagraph.uids)
+    assert np.array_equal(copied_metagraph.stake, metagraph.stake)
+    assert np.array_equal(copied_metagraph.total_stake, metagraph.total_stake)
+    assert np.array_equal(copied_metagraph.ranks, metagraph.ranks)
+    assert np.array_equal(copied_metagraph.trust, metagraph.trust)
+    assert np.array_equal(copied_metagraph.consensus, metagraph.consensus)
+    assert np.array_equal(copied_metagraph.validator_trust, metagraph.validator_trust)
+    assert np.array_equal(copied_metagraph.incentive, metagraph.incentive)
+    assert np.array_equal(copied_metagraph.emission, metagraph.emission)
+    assert np.array_equal(copied_metagraph.dividends, metagraph.dividends)
+    assert np.array_equal(copied_metagraph.active, metagraph.active)
+    assert np.array_equal(copied_metagraph.last_update, metagraph.last_update)
+    assert np.array_equal(copied_metagraph.validator_permit, metagraph.validator_permit)
+    assert np.array_equal(copied_metagraph.weights, metagraph.weights)
+    assert np.array_equal(copied_metagraph.bonds, metagraph.bonds)
+
+    # Check that the neurons are different objects in the original and copied metagraphs
+    for original_neuron, copied_neuron in zip(
+        metagraph.neurons, copied_metagraph.neurons
+    ):
+        assert original_neuron is not copied_neuron
+        assert original_neuron.uid == copied_neuron.uid
+        assert original_neuron.trust == copied_neuron.trust
+        assert original_neuron.consensus == copied_neuron.consensus
+        assert original_neuron.incentive == copied_neuron.incentive
+        assert original_neuron.dividends == copied_neuron.dividends
+        assert original_neuron.rank == copied_neuron.rank
+        assert original_neuron.emission == copied_neuron.emission
+        assert original_neuron.active == copied_neuron.active
+        assert original_neuron.last_update == copied_neuron.last_update
+        assert original_neuron.validator_permit == copied_neuron.validator_permit
+        assert original_neuron.validator_trust == copied_neuron.validator_trust
+        assert original_neuron.total_stake.tao == copied_neuron.total_stake.tao
+        assert original_neuron.stake == copied_neuron.stake
+        assert original_neuron.axon_info == copied_neuron.axon_info
+        assert original_neuron.weights == copied_neuron.weights
+        assert original_neuron.bonds == copied_neuron.bonds
+
+
+def test_copy(mock_environment):
+    subtensor, neurons = mock_environment
+    metagraph = Metagraph(1, sync=False)
+    metagraph.neurons = neurons
+    metagraph.subtensor = subtensor
+
+    # Do a shallow copy
+    copied_metagraph = copy.copy(metagraph)
+
+    # Check that the subtensor attribute is None in the copied object
+    assert copied_metagraph.subtensor is None
+
+    # Check that other attributes are copied correctly
+    assert copied_metagraph.n == metagraph.n
+    assert copied_metagraph.block == metagraph.block
+    assert np.array_equal(copied_metagraph.uids, metagraph.uids)
+    assert np.array_equal(copied_metagraph.stake, metagraph.stake)
+    assert np.array_equal(copied_metagraph.total_stake, metagraph.total_stake)
+    assert np.array_equal(copied_metagraph.ranks, metagraph.ranks)
+    assert np.array_equal(copied_metagraph.trust, metagraph.trust)
+    assert np.array_equal(copied_metagraph.consensus, metagraph.consensus)
+    assert np.array_equal(copied_metagraph.validator_trust, metagraph.validator_trust)
+    assert np.array_equal(copied_metagraph.incentive, metagraph.incentive)
+    assert np.array_equal(copied_metagraph.emission, metagraph.emission)
+    assert np.array_equal(copied_metagraph.dividends, metagraph.dividends)
+    assert np.array_equal(copied_metagraph.active, metagraph.active)
+    assert np.array_equal(copied_metagraph.last_update, metagraph.last_update)
+    assert np.array_equal(copied_metagraph.validator_permit, metagraph.validator_permit)
+    assert copied_metagraph.axons == metagraph.axons
+    assert copied_metagraph.neurons == metagraph.neurons
+    assert np.array_equal(copied_metagraph.weights, metagraph.weights)
+    assert np.array_equal(copied_metagraph.bonds, metagraph.bonds)
+
+    # Check that the neurons are the same objects in the original and copied metagraphs
+    for original_neuron, copied_neuron in zip(
+        metagraph.neurons, copied_metagraph.neurons
+    ):
+        assert original_neuron is copied_neuron