feat: Use Pylon client in the Facilitator

Facilitator now uses the pylon client to communicate with the chain
wherever possible.

Issue: COM-807
Impacts: facilitator

Author: kacper-wolkiewicz-reef

compute_horde/pyproject.toml

@@ -18,7 +18,7 @@ dependencies = [
     "compute-horde-sdk",
     "rich~=13.9.4",
     "web3==7.11.0",
-    "bittensor-pylon-client==1.3.0",
+    "bittensor-pylon-client==1.4.0",
 ]
 
 [build-system]

compute_horde/uv.lock

@@ -0,0 +1,13 @@
+from django.conf import settings
+from pylon_client.v1 import DEFAULT_RETRIES, Config, PylonClient
+from tenacity import Retrying
+
+
+def pylon_client(retries: Retrying = DEFAULT_RETRIES) -> PylonClient:
+    return PylonClient(
+        Config(
+            address=settings.PYLON_ADDRESS,
+            open_access_token=settings.PYLON_OPEN_ACCESS_TOKEN,
+            retry=retries,
+        )
+    )

executor/uv.lock

@@ -1,16 +1,17 @@
 from collections import defaultdict
 from datetime import timedelta
+from ipaddress import IPv4Address
 
 import requests
 import structlog
 from asgiref.sync import async_to_sync
 from celery.utils.log import get_task_logger
 from channels.layers import get_channel_layer
-from compute_horde.utils import get_validators
 from django.conf import settings
 from django.db import connection
 from django.utils.timezone import now
 from pydantic import BaseModel, ValidationError, parse_obj_as
+from pylon_client.v1 import NetUid
 from requests import RequestException
 
 from project.celery import app
@@ -26,6 +27,7 @@
     Validator,
 )
 from .models import MinerVersion as MinerVersionDTO
+from .pylon import pylon_client
 from .schemas import ForceDisconnect, HardwareSpec
 from .specs import normalize_gpu_name
 from .utils import fetch_compute_subnet_hardware
@@ -38,14 +40,18 @@
 @app.task
 def sync_metagraph() -> None:
     """Fetch current validators and miners from the network and store them in the database"""
-    import bittensor
-
-    with bittensor.subtensor(network=settings.BITTENSOR_NETWORK) as subtensor:
-        metagraph = subtensor.metagraph(netuid=settings.BITTENSOR_NETUID)
-        validators = get_validators(metagraph=metagraph)
-
-    sync_validators.delay([v.hotkey for v in validators])
-    sync_miners.delay([neuron.hotkey for neuron in metagraph.neurons if neuron.axon_info.is_serving])
+    netuid = NetUid(settings.BITTENSOR_NETUID)
+    with pylon_client() as client:
+        validators_response = client.open_access.get_latest_validators(netuid)
+        block_number = validators_response.block.number
+        neurons_response = client.open_access.get_neurons(netuid, block_number)
+
+    validator_hotkeys = [v.hotkey for v in validators_response.validators]
+    serving_miner_hotkeys = [
+        hotkey for hotkey, neuron in neurons_response.neurons.items() if neuron.axon_info.is_serving
+    ]
+    sync_validators.delay(validator_hotkeys)
+    sync_miners.delay(serving_miner_hotkeys)
 
 
 @app.task
@@ -191,14 +197,18 @@ def fetch_miner_version(hotkey: str, ip: str, port: int) -> None:
 def fetch_miner_versions() -> None:
     """
     Fetch miner & miner runner versions for every active miner.
-    The list of active miners is retrieved from the metagraph.
+    The list of active miners is retrieved from the metagraph via Pylon.
     """
-    import bittensor
-
-    metagraph = bittensor.metagraph(netuid=settings.BITTENSOR_NETUID, network=settings.BITTENSOR_NETWORK)
-    miners = [neuron for neuron in metagraph.neurons if neuron.axon_info.is_serving]
-    for miner in miners:
-        fetch_miner_version.delay(miner.hotkey, miner.axon_info.ip, miner.axon_info.port)
+    netuid = NetUid(settings.BITTENSOR_NETUID)
+    with pylon_client() as client:
+        response = client.open_access.get_latest_neurons(netuid)
+
+    for hotkey, neuron in response.neurons.items():
+        if neuron.axon_info.is_serving:
+            if not isinstance(neuron.axon_info.ip, IPv4Address):
+                log.warning("skipping non-IPv4 miner", miner_hotkey=hotkey, ip=neuron.axon_info.ip)
+                continue
+            fetch_miner_version.delay(hotkey, str(neuron.axon_info.ip), neuron.axon_info.port)
 
 
 @app.task

facilitator/app/src/project/core/pylon.py

@@ -20,6 +20,8 @@
 from ...asgi import application
 from ..models import Channel, Job, Validator
 
+pytest_plugins = ["compute_horde.test_base.pylon"]
+
 
 @pytest_asyncio.fixture
 async def communicator():

facilitator/app/src/project/core/tasks.py

@@ -1,4 +1,6 @@
+from ipaddress import IPv4Address
 from typing import NamedTuple
+from unittest.mock import MagicMock, patch
 
 import pytest
 from asgiref.sync import sync_to_async
@@ -9,56 +11,38 @@
 
 class MockedAxonInfo(NamedTuple):
     is_serving: bool
-    ip: str = ""
-    port: int = 0
+    ip: IPv4Address = IPv4Address("0.0.0.0")
+    port: int = 8000
 
 
 class MockedNeuron(NamedTuple):
-    uid: int
     hotkey: str
     axon_info: MockedAxonInfo
-    stake: float
 
 
-class MockedMetagraph:
-    def __init__(self, neurons):
-        self.neurons = neurons
-        self.total_stake = [n.stake for n in neurons]
-
-
-class MockedSubtensor:
-    def __init__(self, metagraph: MockedMetagraph):
-        self._metagraph = metagraph
-
-    def __call__(self, *args, **kwargs):
-        return self
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, *args, **kwargs):
-        pass
+class MockedValidator(NamedTuple):
+    hotkey: str
 
-    def metagraph(self, *args, **kwargs):
-        return self._metagraph
 
+def configure_mock_pylon_client(
+    mock_client: MagicMock,
+    validators: list[str],
+    neurons: dict[str, MockedNeuron],
+) -> None:
+    validators_response = MagicMock()
+    validators_response.validators = [MockedValidator(hotkey=v) for v in validators]
+    validators_response.block.number = 12345
 
-validator_params = dict(
-    axon_info=MockedAxonInfo(is_serving=False),
-    stake=1000.0,
-)
+    neurons_response = MagicMock()
+    neurons_response.neurons = neurons
 
-miner_params = dict(
-    axon_info=MockedAxonInfo(is_serving=True),
-    stake=0.0,
-)
+    mock_client.open_access.get_latest_validators.return_value = validators_response
+    mock_client.open_access.get_neurons.return_value = neurons_response
 
 
 @pytest.mark.django_db(transaction=True)
-def test__sync_metagraph__activation(monkeypatch):
-    import bittensor
-
-    validators = Validator.objects.bulk_create(
+def test__sync_metagraph__activation(mock_pylon_client):
+    Validator.objects.bulk_create(
         [
             Validator(ss58_address="remains_active", is_active=True),
             Validator(ss58_address="is_deactivated", is_active=True),
@@ -67,56 +51,55 @@ def test__sync_metagraph__activation(monkeypatch):
         ]
     )
 
-    metagraph = MockedMetagraph(
-        neurons=[
-            MockedNeuron(uid=0, hotkey="remains_active", **validator_params),
-            MockedNeuron(uid=1, hotkey="is_deactivated", **miner_params),
-            MockedNeuron(uid=2, hotkey="remains_inactive", **miner_params),
-            MockedNeuron(uid=3, hotkey="is_activated", **validator_params),
-            MockedNeuron(uid=4, hotkey="new_validator", **validator_params),
-            MockedNeuron(uid=5, hotkey="new_miner", **miner_params),
-        ]
-    )
-    subtensor = MockedSubtensor(metagraph=metagraph)
+    validator_hotkeys = ["remains_active", "is_activated", "new_validator"]
+    neurons = {
+        "remains_active": MockedNeuron(hotkey="remains_active", axon_info=MockedAxonInfo(is_serving=False)),
+        "is_deactivated": MockedNeuron(hotkey="is_deactivated", axon_info=MockedAxonInfo(is_serving=True)),
+        "remains_inactive": MockedNeuron(hotkey="remains_inactive", axon_info=MockedAxonInfo(is_serving=True)),
+        "is_activated": MockedNeuron(hotkey="is_activated", axon_info=MockedAxonInfo(is_serving=False)),
+        "new_validator": MockedNeuron(hotkey="new_validator", axon_info=MockedAxonInfo(is_serving=False)),
+        "new_miner": MockedNeuron(hotkey="new_miner", axon_info=MockedAxonInfo(is_serving=True)),
+    }
+
+    configure_mock_pylon_client(mock_pylon_client, validator_hotkeys, neurons)
 
-    with monkeypatch.context() as mp:
-        mp.setattr(bittensor, "subtensor", subtensor)
+    with patch("project.core.tasks.pylon_client", return_value=mock_pylon_client):
         sync_metagraph()
 
     validators = Validator.objects.order_by("id").values_list("ss58_address", "is_active")
     assert list(validators) == [
-        tuple(d.values())
-        for d in [
-            dict(ss58_address="remains_active", is_active=True),
-            dict(ss58_address="is_deactivated", is_active=False),
-            dict(ss58_address="remains_inactive", is_active=False),
-            dict(ss58_address="is_activated", is_active=True),
-            dict(ss58_address="new_validator", is_active=True),
-        ]
+        ("remains_active", True),
+        ("is_deactivated", False),
+        ("remains_inactive", False),
+        ("is_activated", True),
+        ("new_validator", True),
     ]
 
 
 @pytest.mark.asyncio
 @pytest.mark.django_db(transaction=True)
 async def test__websocket__disconnect_validator_if_become_inactive(
-    monkeypatch,
+    mock_pylon_client,
     communicator,
     authenticated,
     validator,
     job,
     dummy_job_params,
 ):
-    """Check that validator is disconnected if it becomes inactive"""
-    import bittensor
-
     await communicator.receive_json_from()
     assert await Channel.objects.filter(validator=validator).aexists()
 
-    metagraph = MockedMetagraph(neurons=[MockedNeuron(uid=0, hotkey=validator.ss58_address, **miner_params)])
-    subtensor = MockedSubtensor(metagraph=metagraph)
+    validator_hotkeys = []
+    neurons = {
+        validator.ss58_address: MockedNeuron(
+            hotkey=validator.ss58_address,
+            axon_info=MockedAxonInfo(is_serving=True),
+        ),
+    }
+
+    configure_mock_pylon_client(mock_pylon_client, validator_hotkeys, neurons)
 
-    with monkeypatch.context() as mp:
-        mp.setattr(bittensor, "subtensor", subtensor)
+    with patch("project.core.tasks.pylon_client", return_value=mock_pylon_client):
         await sync_to_async(sync_metagraph)()
 
     assert (await communicator.receive_output())["type"] == "websocket.close"

facilitator/app/src/project/core/tests/conftest.py

@@ -137,6 +137,8 @@ def wrapped(*args, **kwargs):
 
 BITTENSOR_NETUID = env.int("BITTENSOR_NETUID")
 BITTENSOR_NETWORK = env.str("BITTENSOR_NETWORK")
+PYLON_ADDRESS = env.str("PYLON_ADDRESS", "http://pylon:8000")
+PYLON_OPEN_ACCESS_TOKEN = env.str("PYLON_OPEN_ACCESS_TOKEN", "facilitator_token")
 SIGNATURE_EXPIRE_DURATION = env("SIGNATURE_EXPIRE_DURATION", default="300")
 
 WANDB_API_KEY = env("WANDB_API_KEY")

facilitator/app/src/pytest.ini …or/app/src/project/core/tests/pytest.inifacilitator/app/src/pytest.ini renamed to facilitator/app/src/project/core/tests/pytest.ini facilitator/app/src/pytest.ini renamed to facilitator/app/src/project/core/tests/pytest.ini

@@ -36,6 +36,8 @@ CHANNELS_BACKEND_URL=redis://localhost:7379/0
 BITTENSOR_NETUID=12
 BITTENSOR_NETWORK=test
 
+PYLON_OPEN_ACCESS_TOKEN=facilitator_token
+
 WANDB_API_KEY=
 
 EMAIL_BACKEND=django.core.mail.backends.filebased.EmailBackend