Merge branch 'staging' into fix/spelling

Author: thewhaleking

bittensor/core/axon.py

@@ -142,14 +142,6 @@ def run_in_thread(self):
             self.should_exit = True
             thread.join()
 
-    def _wrapper_run(self):
-        """
-        A wrapper method for the :func:`run_in_thread` context manager. This method is used internally by the ``start`` method to initiate the server's execution in a separate thread.
-        """
-        with self.run_in_thread():
-            while not self.should_exit:
-                time.sleep(1e-3)
-
     def start(self):
         """
         Starts the FastAPI server in a separate thread if it is not already running. This method sets up the server to handle HTTP requests concurrently, enabling the Axon server to efficiently manage incoming network requests.
@@ -158,7 +150,7 @@ def start(self):
         """
         if not self.is_running:
             self.should_exit = False
-            thread = threading.Thread(target=self._wrapper_run, daemon=True)
+            thread = threading.Thread(target=self.run, daemon=True)
             thread.start()
             self.is_running = True
 

bittensor/utils/btlogging/loggingmachine.py

@@ -582,6 +582,10 @@ def get_level(self) -> int:
         """Returns Logging level."""
         return self._logger.level
 
+    def setLevel(self, level):
+        """Set the specified level on the underlying logger."""
+        self._logger.setLevel(level)
+
     def check_config(self, config: "Config"):
         assert config.logging
 

tests/unit_tests/test_axon.py

@@ -17,20 +17,23 @@
 
 
 import re
+import threading
 import time
 from dataclasses import dataclass
 from typing import Any, Optional, Tuple
 from unittest import IsolatedAsyncioTestCase
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import aiohttp
 import fastapi
 import netaddr
 import pydantic
 import pytest
+import uvicorn
 from fastapi.testclient import TestClient
 from starlette.requests import Request
 
-from bittensor.core.axon import AxonMiddleware, Axon
+from bittensor.core.axon import Axon, AxonMiddleware, FastAPIThreadedServer
 from bittensor.core.errors import RunException
 from bittensor.core.settings import version_as_int
 from bittensor.core.stream import StreamingSynapse
@@ -785,3 +788,45 @@ async def forward_fn(synapse: streaming_synapse_cls):
                 "computed_body_hash": "a7ffc6f8bf1ed76651c14756a061d662f580ff4de43b49fa82d80a4b80f8434a",
             },
         )
+
+
+@pytest.mark.asyncio
+async def test_threaded_fastapi():
+    server_started = threading.Event()
+    server_stopped = threading.Event()
+
+    async def lifespan(app):
+        server_started.set()
+        yield
+        server_stopped.set()
+
+    app = fastapi.FastAPI(
+        lifespan=lifespan,
+    )
+    app.get("/")(lambda: "Hello World")
+
+    server = FastAPIThreadedServer(
+        uvicorn.Config(
+            app,
+        ),
+    )
+    server.start()
+
+    server_started.wait()
+
+    assert server.is_running is True
+
+    async with aiohttp.ClientSession(
+        base_url="http://127.0.0.1:8000",
+    ) as session:
+        async with session.get("/") as response:
+            assert await response.text() == '"Hello World"'
+
+        server.stop()
+
+        assert server.should_exit is True
+
+        server_stopped.wait()
+
+        with pytest.raises(aiohttp.ClientConnectorError):
+            await session.get("/")

tests/unit_tests/test_logging.py

@@ -212,3 +212,25 @@ def test_all_log_levels_output(logging_machine, caplog):
 def test_concat(msg, prefix, suffix, expected_result):
     """Test different options of message concatenation with prefix and suffix."""
     assert _concat_message(msg, prefix, suffix) == expected_result
+
+
+def test_logger_level(logging_machine, caplog):
+    """Test get/set Logger level."""
+
+    assert logging_machine.get_level() == stdlogging.WARN
+
+    logging_machine.info("info1")
+    logging_machine.warning("warn1")
+
+    assert "info1" not in caplog.text
+    assert "warn1" in caplog.text
+
+    logging_machine.setLevel(stdlogging.INFO)
+
+    assert logging_machine.get_level() == stdlogging.INFO
+
+    logging_machine.info("info2")
+    logging_machine.warning("warn2")
+
+    assert "info2" in caplog.text
+    assert "warn2" in caplog.text