feat: Use Pylon client for setting weights

Validator sets weights using Pylon client now. Failed connection to
Pylon service now results in set weights error system event instead of
generic chain error.

Impacts: validator
Issue: COM-800

Author: kacper-wolkiewicz-reef

validator/app/src/compute_horde_validator/settings.py

@@ -698,9 +698,9 @@ def wrapped(*args, **kwargs):
 
 BITTENSOR_NETUID = env.int("BITTENSOR_NETUID")
 BITTENSOR_NETWORK = env.str("BITTENSOR_NETWORK")
-PYLON_HOST = "pylon"
-PYLON_PORT = 8000
-PYLON_AUTH_TOKEN = "abc"
+PYLON_ADDRESS = env.str("PYLON_ADDRESS", "http://pylon:8000")
+PYLON_IDENTITY_NAME = env.str("PYLON_IDENTITY_NAME")
+PYLON_IDENTITY_TOKEN = env.str("PYLON_IDENTITY_TOKEN")
 # This can be explicitly set to None, which will cause some backfilling operations to never succeed. Useful when running
 # on a private staging net etc.
 BITTENSOR_ARCHIVE_NETWORK = env.str("BITTENSOR_ARCHIVE_NETWORK", "archive")

validator/app/src/compute_horde_validator/validator/pylon.py

@@ -0,0 +1,15 @@
+from django.conf import settings
+from pylon._internal.client.sync.client import PylonClient
+from pylon._internal.client.sync.config import DEFAULT_RETRIES, Config
+from tenacity import Retrying
+
+
+def pylon_client(retries: Retrying = DEFAULT_RETRIES) -> PylonClient:
+    return PylonClient(
+        Config(
+            address=settings.PYLON_ADDRESS,
+            identity_name=settings.PYLON_IDENTITY_NAME,
+            identity_token=settings.PYLON_IDENTITY_TOKEN,
+            retry=retries,
+        )
+    )

validator/app/src/compute_horde_validator/validator/scoring/pylon_client.py

@@ -1,7 +1,6 @@
 import contextlib
 import logging
 import random
-import time
 import traceback
 from functools import cached_property
 from typing import Union
@@ -16,6 +15,8 @@
 from django.db import transaction
 from numpy._typing import NDArray
 from pydantic import JsonValue
+from pylon.v1 import DEFAULT_RETRIES, BasePylonException, Hotkey, Weight
+from tenacity import stop_after_attempt, wait_fixed
 
 from compute_horde_validator.celery import app
 from compute_horde_validator.validator.allowance.utils.supertensor import supertensor
@@ -25,8 +26,8 @@
 from compute_horde_validator.validator.models.scoring.internal import (
     WeightSettingFinishedEvent,
 )
+from compute_horde_validator.validator.pylon import pylon_client
 from compute_horde_validator.validator.scoring import create_scoring_engine
-from compute_horde_validator.validator.scoring.pylon_client import PylonClientError, pylon_client
 
 if False:
     import torch
@@ -139,38 +140,35 @@ def set_scores() -> None:
                     max_weight_limit=hyperparameters["max_weights_limit"],
                 )
                 hk_by_uid = {n.uid: n.hotkey for n in neurons}
-                hk_weight_mapping = {hk_by_uid[uid]: weight for uid, weight in zip(uids, weights)}
-
-                for try_number in range(WEIGHT_SETTING_ATTEMPTS):
-                    logger.debug(
-                        f"Setting weights (attempt #{try_number}):\nuids={uids}\nscores={weights}"
-                    )
-                    try:
-                        pylon_client().set_weights(hk_weight_mapping)
-                        logger.info("Successfully set weights!!!")
-                        break
-                    except PylonClientError:
-                        logger.warning(
-                            "Encountered when setting weights (attempt #{try_number}): ",
-                            exc_info=True,
-                        )
-                        save_weight_setting_failure(
-                            subtype=SystemEvent.EventSubType.WRITING_TO_CHAIN_GENERIC_ERROR,
+                hk_weight_mapping = {
+                    Hotkey(hk_by_uid[uid]): Weight(weight) for uid, weight in zip(uids, weights)
+                }
+                logger.debug(f"Setting weights:\nuids={uids}\nscores={weights}")
+                with pylon_client(
+                    retries=DEFAULT_RETRIES.copy(
+                        stop=stop_after_attempt(WEIGHT_SETTING_ATTEMPTS),
+                        wait=wait_fixed(WEIGHT_SETTING_FAILURE_BACKOFF),
+                        after=lambda retry_state: save_weight_setting_failure(
+                            subtype=SystemEvent.EventSubType.SET_WEIGHTS_ERROR,
                             long_description=traceback.format_exc(),
-                            data={"try_number": try_number, "operation": "setting/committing"},
-                        )
-                    time.sleep(WEIGHT_SETTING_FAILURE_BACKOFF)
-                else:
-                    raise MaximumNumberOfAttemptsExceeded()
+                            data={
+                                "try_number": retry_state.attempt_number,
+                                "operation": "setting/committing",
+                            },
+                        ),
+                    )
+                ) as p_client:
+                    p_client.identity.put_weights(hk_weight_mapping)
+                logger.info("Setting weights successfully scheduled in Pylon service!")
                 save_weight_setting_event(
                     type_=SystemEvent.EventType.WEIGHT_SETTING_SUCCESS,
                     subtype=SystemEvent.EventSubType.SET_WEIGHTS_SUCCESS,
                     long_description="",
                     data={},
                 )
-        except MaximumNumberOfAttemptsExceeded:
-            msg = f"Failed to set weights after {WEIGHT_SETTING_ATTEMPTS} attempts"
-            logger.warning(msg)
+        except BasePylonException as e:
+            msg = f"Failed to schedule setting weights in Pylon service after {WEIGHT_SETTING_ATTEMPTS} attempts!"
+            logger.warning(msg, exc_info=e)
             save_weight_setting_failure(
                 subtype=SystemEvent.EventSubType.GIVING_UP,
                 long_description=msg,

validator/app/src/compute_horde_validator/validator/scoring/tasks.py

@@ -1,12 +1,13 @@
 import logging
 import uuid
 from collections.abc import Generator
-from unittest.mock import patch
+from unittest.mock import create_autospec, patch
 
 import bittensor_wallet
 import pytest
 from compute_horde.executor_class import EXECUTOR_CLASS
 from compute_horde_core.executor_class import ExecutorClass
+from pylon._internal.client.sync.client import PylonClient
 
 from ..organic_jobs.miner_driver import execute_organic_job_request
 from .helpers import MockNeuron
@@ -94,3 +95,14 @@ def run_uuid():
 #         raise ValueError(
 #             "\n" + "\n".join(f"{task.get_name()}: {task.get_coro()}" for task in tasks)
 #         )
+
+
+@pytest.fixture
+def pylon_client_mock(mocker):
+    # This is a temporary solution until pylon client implements its own mocking utility.
+    mocked = create_autospec(PylonClient)
+    mocked.__enter__.return_value = mocked
+    mocked.open_access = create_autospec(PylonClient._open_access_api_cls, instance=True)
+    mocked.identity = create_autospec(PylonClient._identity_api_cls, instance=True)
+    mocker.patch("compute_horde_validator.validator.pylon.PylonClient", return_value=mocked)
+    return mocked

validator/app/src/compute_horde_validator/validator/tests/conftest.py

@@ -1,13 +1,10 @@
 import pytest
 
-from compute_horde_validator.validator.scoring.pylon_client import setup_mock_pylon_client
 from compute_horde_validator.validator.scoring.tasks import set_scores
 
 
 @pytest.mark.django_db
 @pytest.mark.override_config(SERVING=False)
-def test__migration__not_serving__should_not_set_scores():
-    with setup_mock_pylon_client() as mock_pylon_client:
-        set_scores()
-
-    assert mock_pylon_client.weights_submitted == []
+def test__migration__not_serving__should_not_set_scores(pylon_client_mock):
+    set_scores()
+    pylon_client_mock.identity.put_weights.assert_not_called()