diff --git a/bittensor/core/axon.py b/bittensor/core/axon.py index 61573ae6fc..fa40a64449 100644 --- a/bittensor/core/axon.py +++ b/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, diff --git a/mypy.ini b/mypy.ini index d38bdc7172..054851282e 100644 --- a/mypy.ini +++ b/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 diff --git a/tests/unit_tests/test_axon.py b/tests/unit_tests/test_axon.py index 0712e50dc1..0cc2fd64a5 100644 --- a/tests/unit_tests/test_axon.py +++ b/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", [