Adds config to select only uids with unique coldkeys/ips

prompting/utils/config.py

@@ -98,7 +98,7 @@ def add_args(cls, parser):
         help="If set, we dont save events to a log file.",
         default=False,
     )
-    
+
     parser.add_argument(
         "--neuron.log_full",
         action="store_true",
@@ -291,6 +291,19 @@ def add_validator_args(cls, parser):
     )
 
 
+    parser.add_argument(
+        "--neuron.query_unique_coldkeys",
+        action="store_true",
+        help="Only query a single hotkey per coldkey.",
+        default=False,
+        )
+
+    parser.add_argument(
+        "--neuron.query_unique_ips",
+        action="store_true",
+        help="Only query a single hotkey per ip.",
+        default=False,
+        )
 
 def config(cls):
     """

prompting/utils/uids.py

@@ -5,24 +5,34 @@
 
 
 def check_uid_availability(
-    metagraph: "bt.metagraph.Metagraph", uid: int, vpermit_tao_limit: int
+    metagraph: "bt.metagraph.Metagraph", uid: int, vpermit_tao_limit: int, coldkeys: set = None, ips: set = None,
 ) -> bool:
     """Check if uid is available. The UID should be available if it is serving and has less than vpermit_tao_limit stake
     Args:
         metagraph (:obj: bt.metagraph.Metagraph): Metagraph object
         uid (int): uid to be checked
         vpermit_tao_limit (int): Validator permit tao limit
+        coldkeys (set): Set of coldkeys to exclude
+        ips (set): Set of ips to exclude
     Returns:
         bool: True if uid is available, False otherwise
     """
     # Filter non serving axons.
     if not metagraph.axons[uid].is_serving:
         bt.logging.debug(f"uid: {uid} is not serving")
         return False
+
     # Filter validator permit > 1024 stake.
     if metagraph.validator_permit[uid] and metagraph.S[uid] > vpermit_tao_limit:
-        bt.logging.debug(f"uid: {uid} has vpermit and stake ({metagraph.S[uid]}) > {vpermit_tao_limit}")
+        bt.logging.warning(f"uid: {uid} has vpermit and stake ({metagraph.S[uid]}) > {vpermit_tao_limit}")
         return False
+
+    if coldkeys and metagraph.axons[uid].coldkey in coldkeys:
+        return False
+
+    if ips and metagraph.axons[uid].ip in ips:
+        return False
+
     # Available otherwise.
     return True
 
@@ -41,19 +51,27 @@ def get_random_uids(
     """
     candidate_uids = []
     avail_uids = []
-
+    coldkeys = set()
+    ips = set()
     for uid in range(self.metagraph.n.item()):
         if uid == self.uid:
             continue
+
         uid_is_available = check_uid_availability(
-            self.metagraph, uid, self.config.neuron.vpermit_tao_limit
+            self.metagraph, uid, self.config.neuron.vpermit_tao_limit, coldkeys, ips,
         )
-        uid_is_not_excluded = exclude is None or uid not in exclude
+        if not uid_is_available:
+            continue
+
+        if self.config.neuron.query_unique_coldkeys:
+            coldkeys.add(self.metagraph.axons[uid].coldkey)
+
+        if self.config.neuron.query_unique_ips:
+            ips.add(self.metagraph.axons[uid].ip)
 
-        if uid_is_available:
-            avail_uids.append(uid)
-            if uid_is_not_excluded:
-                candidate_uids.append(uid)
+        avail_uids.append(uid)
+        if exclude is None or uid not in exclude:
+            candidate_uids.append(uid)
 
     # Check if candidate_uids contain enough for querying, if not grab all avaliable uids
     available_uids = candidate_uids

tests/test_uids.py

@@ -0,0 +1,48 @@
+
+import torch
+import pytest
+from types import SimpleNamespace
+from prompting.utils.uids import get_random_uids
+
+
+def make_mock_neuron(unique_coldkeys=False, unique_ips=False, vpermit_tao_limit=1000):
+
+    axons = [
+        SimpleNamespace(coldkey="a", ip="0.0.0.1", is_serving=True),
+        SimpleNamespace(coldkey="a", ip="0.0.0.0", is_serving=True),
+        SimpleNamespace(coldkey="b", ip="0.0.0.1", is_serving=True),
+        SimpleNamespace(coldkey="b", ip="0.0.0.0", is_serving=True),
+        SimpleNamespace(coldkey="c", ip="0.0.0.2", is_serving=True),
+    ]
+    metagraph = SimpleNamespace(
+        axons = axons,
+        validator_permit = torch.ones(len(axons), dtype=torch.bool),
+        S = torch.zeros(len(axons)),
+        n = torch.tensor(len(axons))
+    )
+
+    return SimpleNamespace(
+        uid = 4,
+        config = SimpleNamespace(
+            neuron = SimpleNamespace(
+                vpermit_tao_limit = vpermit_tao_limit,
+                query_unique_coldkeys = unique_coldkeys,
+                query_unique_ips = unique_ips,
+            )
+        ),
+        metagraph = metagraph
+    )
+
+@pytest.mark.parametrize(
+    "unique_coldkeys, unique_ips, k, expected_result", [
+        (False, False, 4, [0, 1, 2, 3]),
+        (True, False, 2, [0, 2]),
+        (False, True, 2, [0, 1]),
+        (True, True, 2, [0, 3])
+        ])
+def test_get_random_uids(unique_coldkeys, unique_ips, k, expected_result):
+
+    mock_neuron = make_mock_neuron(unique_coldkeys, unique_ips)
+
+    assert sorted(get_random_uids(mock_neuron, k).tolist()) == expected_result, "Incorrect uids returned."
+