Pass the computed priority_fn value into the axon executor queue

bittensor/core/axon.py

@@ -12,7 +12,7 @@
 import uuid
 import warnings
 from inspect import signature, Signature, Parameter
-from typing import Any, Awaitable, Callable, Optional, Tuple
+from typing import Awaitable, Callable, Optional, Tuple
 
 from async_substrate_interface.utils import json
 from pydantic import ValidationError
@@ -1422,22 +1422,17 @@ async def priority(self, synapse: "Synapse"):
 
         async def submit_task(
             executor: "PriorityThreadPoolExecutor", priority: float
-        ) -> tuple[float, Any]:
-            """
-            Submits the given priority function to the specified executor for asynchronous execution.
-            The function will run in the provided executor and return the priority value along with the result.
-
-            Parameters:
-                executor: The executor in which the priority function will be run.
-                priority: The priority function to be executed.
+        ) -> float:
+            """Submit the computed ``priority`` to the executor's priority queue.
 
-            Returns:
-                A tuple containing the priority value and the result of the priority function execution.
+            ``loop.run_in_executor(executor, fn)`` calls ``executor.submit(fn)`` without
+            the ``priority`` keyword, so the queue fell back to a random priority and the
+            value returned by ``priority_fn`` was effectively discarded. Submit directly
+            with the computed ``priority`` so it actually reaches the priority queue.
             """
-            loop = asyncio.get_running_loop()
-            future = loop.run_in_executor(executor, lambda: priority)
-            result_ = await future
-            return priority, result_
+            future = executor.submit(lambda: priority, priority=priority)
+            await asyncio.wrap_future(future)
+            return priority
 
         # If a priority function exists for the request's name
         if priority_fn:
@@ -1451,7 +1446,7 @@ async def submit_task(
 
                 # Submit the task to the thread pool for execution with the given priority.
                 # The submit_task function will handle the execution and return the result.
-                _, result = await submit_task(self.axon.thread_pool, priority)
+                await submit_task(self.axon.thread_pool, priority)
 
             except TimeoutError as e:
                 # If the execution of the priority function exceeds the timeout,

mypy.ini

@@ -1,6 +1,13 @@
 [mypy]
 ignore_missing_imports = True
 ignore_errors = True
+follow_imports_for_stubs = True
+
+[mypy-numpy.*]
+follow_imports = skip
+
+[mypy-numpy]
+follow_imports = skip
 
 [mypy-*.axon.*]
 ignore_errors = False

tests/unit_tests/test_axon.py

@@ -261,6 +261,30 @@ async def test_priority_pass(middleware):
     assert synapse.axon.status_code != 408
 
 
+@pytest.mark.asyncio
+async def test_priority_passes_computed_priority_to_executor(middleware, mocker):
+    """The value returned by priority_fn must reach the executor's submit call as
+    ``priority=...``. Previously it was discarded: a no-op was submitted via
+    ``loop.run_in_executor`` (no ``priority`` kwarg), so the queue used a random priority."""
+
+    def my_priority_fn(synapse) -> float:
+        return 42.0
+
+    synapse = SynapseMock()
+    middleware.axon.priority_fns = {"SynapseMock": my_priority_fn}
+
+    fake_future = asyncio.Future()
+    fake_future.set_result(None)
+    mock_pool = mocker.MagicMock()
+    mock_pool.submit.return_value = fake_future
+    middleware.axon.thread_pool = mock_pool
+
+    await middleware.priority(synapse)
+
+    assert mock_pool.submit.call_count == 1
+    assert mock_pool.submit.call_args.kwargs.get("priority") == 42.0
+
+
 @pytest.mark.parametrize(
     "body, expected",
     [