diff --git a/subvortex/core/core_bittensor/subtensor/__init__.py b/subvortex/core/core_bittensor/subtensor/__init__.py index c9df6e62..454afcde 100644 --- a/subvortex/core/core_bittensor/subtensor/__init__.py +++ b/subvortex/core/core_bittensor/subtensor/__init__.py @@ -9,7 +9,6 @@ get_next_block, get_hyperparameter_value, get_number_of_uids, - RetryAsyncSubstrate, ) from .subtensor_settings import Settings @@ -26,5 +25,4 @@ "get_next_block", "get_hyperparameter_value", "get_number_of_uids", - "RetryAsyncSubstrate", ] \ No newline at end of file diff --git a/subvortex/core/core_bittensor/subtensor/subtensor.py b/subvortex/core/core_bittensor/subtensor/subtensor.py index 5e4b00f5..9f094067 100644 --- a/subvortex/core/core_bittensor/subtensor/subtensor.py +++ b/subvortex/core/core_bittensor/subtensor/subtensor.py @@ -31,132 +31,10 @@ import bittensor.utils.btlogging as btul import bittensor.utils.weight_utils as btuwu -from websockets.exceptions import ( - InvalidMessage, - WebSocketException, -) import bittensor.core.async_subtensor as btcas -from async_substrate_interface.errors import MaxRetriesExceeded -from async_substrate_interface.async_substrate import Websocket -from async_substrate_interface.substrate_addons import RETRY_METHODS U16_MAX = 65535 -# These are the exceptions that should trigger a retry or re-instantiation -RETRYABLE_EXCEPTIONS = ( - MaxRetriesExceeded, # Custom fallback exhaustion - ConnectionError, # Socket-level connection failure - WebSocketException, # Covers all websocket protocol issues (includes ConnectionClosed) - InvalidMessage, # Specific: bad WebSocket handshake - EOFError, # Stream ended early (common during abrupt disconnects) - socket.gaierror, # DNS resolution failure - asyncio.TimeoutError, # Node hang or await timeout - RuntimeError, # Catch asyncio loop mismatch (filtered by message content) -) - - -class RetryAsyncSubstrate(btcas.AsyncSubstrateInterface): - def __init__( - self, - url: str, - use_remote_preset: bool = False, - fallback_chains: Optional[list[str]] = None, - retry_forever: bool = False, - ss58_format: Optional[int] = None, - type_registry: Optional[dict] = None, - type_registry_preset: Optional[str] = None, - chain_name: str = "", - max_retries: int = 5, - retry_timeout: float = 60.0, - _mock: bool = False, - ): - fallback_chains = fallback_chains or [] - self.fallback_chains = ( - iter(fallback_chains) - if not retry_forever - else cycle(fallback_chains + [url]) - ) - self.use_remote_preset = use_remote_preset - self.chain_name = chain_name - self._mock = _mock - self.retry_timeout = retry_timeout - self.max_retries = max_retries - super().__init__( - url=url, - ss58_format=ss58_format, - type_registry=type_registry, - use_remote_preset=use_remote_preset, - type_registry_preset=type_registry_preset, - chain_name=chain_name, - _mock=_mock, - retry_timeout=retry_timeout, - max_retries=max_retries, - ) - self._original_methods = { - method: getattr(self, method) for method in RETRY_METHODS - } - for method in RETRY_METHODS: - setattr(self, method, partial(self._retry, method)) - - async def _reinstantiate_substrate(self, e: Optional[Exception] = None) -> None: - next_network = next(self.fallback_chains) - if e.__class__ == MaxRetriesExceeded: - btul.logging.error( - f"Max retries exceeded with {self.url}. Retrying with {next_network}." - ) - else: - btul.logging.error(f"Connection error. Trying again with {next_network}") - try: - await self.ws.shutdown() - except AttributeError: - pass - except Exception as shutdown_err: - btul.logging.debug( - f"Ignoring error during websocket shutdown: {shutdown_err}", - ) - - if self._forgettable_task is not None: - self._forgettable_task: asyncio.Task - self._forgettable_task.cancel() - try: - await self._forgettable_task - except asyncio.CancelledError: - pass - - self.chain_endpoint = next_network - self.url = next_network - self.ws = Websocket( - next_network, - options={ - "max_size": self.ws_max_size, - "write_limit": 2**16, - }, - ) - self._initialized = False - self._initializing = False - await self.initialize() - - async def _retry(self, method, *args, **kwargs): - method_ = self._original_methods[method] - try: - return await method_(*args, **kwargs) - except RETRYABLE_EXCEPTIONS as e: - if isinstance(e, RuntimeError) and not self._is_loop_mismatch_error(e): - # only retry RuntimeError if it's a loop mismatch - raise - - try: - await self._reinstantiate_substrate(e) - return await method_(*args, **kwargs) - except StopAsyncIteration: - btul.logging.error( - f"Max retries exceeded with {self.url}. No more fallback chains." - ) - raise MaxRetriesExceeded - - def _is_loop_mismatch_error(self, e: Exception) -> bool: - return isinstance(e, RuntimeError) and "attached to a different loop" in str(e) - def get_next_block(subtensor: btcs.Subtensor, block: int = 0): current_block = subtensor.get_current_block() @@ -294,6 +172,11 @@ async def watchdog(): f"No block received in the last {timeout} seconds" ) + btul.logging.trace( + "🚀 Starting block subscription and watchdog.", + prefix="ReliableSubtensor", + ) + # Run handler + watchdog concurrently await asyncio.gather( subtensor.substrate._get_block_handler( @@ -317,6 +200,10 @@ async def watchdog(): if attempt > 0: # Wait avg time of a block + btul.logging.debug( + "⏲️ Waiting before retrying subscription...", + prefix="ReliableSubtensor", + ) await asyncio.sleep(12) attempt += 1 diff --git a/subvortex/core/database/database.py b/subvortex/core/database/database.py index 2948dcf8..062775c8 100644 --- a/subvortex/core/database/database.py +++ b/subvortex/core/database/database.py @@ -1,8 +1,9 @@ +import asyncio from redis import asyncio as aioredis from packaging.version import parse as parse_version +from weakref import WeakKeyDictionary import bittensor.utils.btlogging as btul - from subvortex.core.database.database_utils import decode_value @@ -10,58 +11,78 @@ class Database: def __init__(self, settings): self.models = {} self.settings = settings - self.database = None + self._clients = WeakKeyDictionary() # Cache clients per event loop - async def connect(self): - self.database = aioredis.StrictRedis( + def _new_client(self): + return aioredis.StrictRedis( host=self.settings.database_host, port=self.settings.database_port, db=self.settings.database_index, password=self.settings.database_password, ) - btul.logging.info("Connected to Redis", prefix=self.settings.logging_name) + def _get_loop(self): + return asyncio.get_running_loop() + + async def get_client(self): + loop = self._get_loop() + + if loop in self._clients: + return self._clients[loop] + + client = self._new_client() + self._clients[loop] = client + + btul.logging.info( + "Created new Redis client for event loop", prefix=self.settings.logging_name + ) + return client async def is_connection_alive(self) -> bool: + client = await self.get_client() + try: - pong = await self.database.ping() + pong = await client.ping() return pong is True except Exception as e: - btul.logging.warning(f"Redis connection check failed: {e}") + btul.logging.warning( + f"Redis connection check failed: {e}", prefix=self.settings.logging_name + ) return False async def ensure_connection(self): - if self.database is None or not await self.is_connection_alive(): + client = await self.get_client() + + if not await self.is_connection_alive(): btul.logging.warning( - "Reconnecting to Redis...", + "Redis ping failed, but client will be reused", prefix=self.settings.logging_name, ) - await self.connect() + # You may optionally recreate here if needed async def wait_until_ready(self, name: str): - # Ensure the connection is ip and running await self.ensure_connection() + client = await self.get_client() + message_key = self._key(f"state:{name}") stream_key = self._key(f"state:{name}:stream") last_id = "$" try: - # Step 1: check the message key first - snapshot = await self.database.get(message_key) + snapshot = await client.get(message_key) if snapshot and snapshot.decode() == "ready": - btul.logging.info( + btul.logging.debug( f"{name} is already ready (via message key)", prefix=self.settings.logging_name, ) return - # Step 2: wait for stream messages btul.logging.debug( f"Waiting on stream: {stream_key}", prefix=self.settings.logging_name ) while True: - entries = await self.database.xread({stream_key: last_id}, block=0) + entries = await client.xread({stream_key: last_id}, block=0) if not entries: continue @@ -71,14 +92,14 @@ async def wait_until_ready(self, name: str): prefix=self.settings.logging_name, ) for msg_id, fields in messages: - state = fields.get("state".encode(), b"").decode() + state = fields.get(b"state", b"").decode() if state == "ready": - btul.logging.info( + btul.logging.debug( f"{name} is now ready (via stream)", prefix=self.settings.logging_name, ) return - last_id = msg_id # move forward + last_id = msg_id except Exception as err: btul.logging.warning( f"Failed to read the state of {name}: {err}", @@ -86,14 +107,9 @@ async def wait_until_ready(self, name: str): ) async def _get_migration_status(self, model_name: str): - """ - Returns: - - latest_version: the 'new' version - - active_versions: versions marked 'dual' or 'new', - or fallback to latest if none are active. - """ - # Ensure the connection is ip and running await self.ensure_connection() + + client = await self.get_client() latest = None active = [] @@ -101,7 +117,7 @@ async def _get_migration_status(self, model_name: str): all_versions = sorted(self.models[model_name].keys(), key=parse_version) for version in all_versions: - mode = await self.database.get(f"migration_mode:{version}") + mode = await client.get(f"migration_mode:{version}") mode = decode_value(mode) if mode == "new": diff --git a/subvortex/core/file/file_monitor.py b/subvortex/core/file/file_monitor.py index 7d9302fb..eb4c14fe 100644 --- a/subvortex/core/file/file_monitor.py +++ b/subvortex/core/file/file_monitor.py @@ -16,12 +16,11 @@ # DEALINGS IN THE SOFTWARE. import asyncio import threading -import bittensor.utils.btlogging as btul from enum import Enum +import bittensor.utils.btlogging as btul from subvortex.core.file.file_provider import FileProvider - LOGGER_NAME = "File Monitoring" @@ -32,21 +31,25 @@ class FileType(Enum): class FileMonitor(threading.Thread): def __init__(self): - super().__init__() + super().__init__(daemon=True) self.stop_flag = threading.Event() self.last_error_shown = None self.coroutines = [] - self.loop = asyncio.new_event_loop() + self.loop = None def add_file_provider(self, file_provider: FileProvider): - task = self.loop.create_task(self._check_file(file_provider)) - self.coroutines.append(task) + if self.loop: + task = asyncio.run_coroutine_threadsafe( + self._check_file(file_provider), self.loop + ) + self.coroutines.append(task) async def _check_file(self, file: FileProvider): - while not self.stop_flag.is_set(): - try: - # Wait a specific time before starting + try: + while not self.stop_flag.is_set(): await asyncio.sleep(file.check_interval) + if self.stop_flag.is_set(): + break btul.logging.debug( f"[{LOGGER_NAME}][{file.logger_name}] Checking file..." @@ -70,45 +73,38 @@ async def _check_file(self, file: FileProvider): # Reset the last error shown self.last_error_shown = None - except Exception as err: - error_message = f"[{LOGGER_NAME}][{file.logger_name}] Failed processing file: {err} {type(err)}" - if error_message != self.last_error_shown: - btul.logging.error(error_message) - self.last_error_shown = error_message - - async def _run_async(self): - while not self.stop_flag.is_set(): - try: - # Sleep for a second before gathering tasks - await asyncio.sleep(1) - if self.stop_flag.is_set(): - # Time to stop the file monitoring - # We wait until all the tasks are finished - await asyncio.gather(*self.coroutines) - except Exception as err: - error_message = ( - f"[{LOGGER_NAME}] Failed checking files: {err} {type(err)}" - ) - if error_message != self.last_error_shown: - btul.logging.error(error_message) - self.last_error_shown = error_message + except Exception as err: + error_message = f"[{LOGGER_NAME}][{file.logger_name}] Failed processing file: {err} {type(err)}" + if error_message != self.last_error_shown: + btul.logging.error(error_message) + self.last_error_shown = error_message def run(self): + self.loop = asyncio.new_event_loop() + asyncio.set_event_loop(self.loop) + + async def monitor(): + while not self.stop_flag.is_set(): + # Short sleep to yield to other tasks + await asyncio.sleep(0.1) + try: - self.loop.run_until_complete(self._run_async()) + self.loop.run_until_complete(monitor()) + except Exception as err: + btul.logging.error(f"[{LOGGER_NAME}] Loop error: {err}") finally: - self.loop.stop() - self.loop.run_until_complete(self.loop.shutdown_asyncgens()) + pending = asyncio.all_tasks(loop=self.loop) + for task in pending: + task.cancel() + self.loop.run_until_complete( + asyncio.gather(*pending, return_exceptions=True) + ) self.loop.close() - - btul.logging.debug(f"[{LOGGER_NAME}] run ended") - - def start(self): - super().start() - btul.logging.debug(f"[{LOGGER_NAME}] started") + btul.logging.debug(f"[{LOGGER_NAME}] Event loop closed") def stop(self): + btul.logging.info(f"[{LOGGER_NAME}] FileMonitor stopping") self.stop_flag.set() - super().join() - btul.logging.debug(f"[{LOGGER_NAME}] stopped") + self.join() + btul.logging.info(f"[{LOGGER_NAME}] FileMonitor stopped") diff --git a/subvortex/core/metagraph/database.py b/subvortex/core/metagraph/database.py index af462bd4..f925ea45 100644 --- a/subvortex/core/metagraph/database.py +++ b/subvortex/core/metagraph/database.py @@ -34,6 +34,9 @@ async def get_neuron(self, hotkey: str) -> scmm.Neuron: # Ensure the connection is up and running await self.ensure_connection() + # Get a client + client = await self.get_client() + # Get the active versions _, active = await self._get_migration_status("neuron") @@ -44,7 +47,7 @@ async def get_neuron(self, hotkey: str) -> scmm.Neuron: try: # Attempt to read the neuron using the model - neuron = await model.read(self.database, hotkey) + neuron = await model.read(client, hotkey) return neuron except Exception as ex: @@ -63,6 +66,9 @@ async def get_neurons(self) -> typing.Dict[str, scmm.Neuron]: # Ensure the connection is up and running await self.ensure_connection() + # Get a client + client = await self.get_client() + # Get the active versions _, active = await self._get_migration_status("neuron") @@ -73,7 +79,7 @@ async def get_neurons(self) -> typing.Dict[str, scmm.Neuron]: try: # Attempt to read all neurons using the model - neurons = await model.read_all(self.database) + neurons = await model.read_all(client) return neurons except Exception as ex: @@ -95,8 +101,11 @@ async def get_neuron_last_updated(self): # Ensure the connection is up and running await self.ensure_connection() + # Get a client + client = await self.get_client() + try: - raw = await self.database.get(self._key("state:neuron:last_updated")) + raw = await client.get(self._key("state:neuron:last_updated")) return int(decode_value(raw) or 0) except Exception as ex: @@ -133,6 +142,9 @@ async def update_neurons(self, neurons: typing.List[scmm.Neuron]): # Ensure the connection is up and running await self.ensure_connection() + # Get a client + client = await self.get_client() + # Get the active versions _, active = await self._get_migration_status("neuron") @@ -143,7 +155,7 @@ async def update_neurons(self, neurons: typing.List[scmm.Neuron]): try: # Attempt to write all neurons to the database - await model.write_all(self.database, neurons) + await model.write_all(client, neurons) except Exception as ex: btul.logging.warning( @@ -161,10 +173,13 @@ async def remove_neurons(self, neurons: list[Neuron]): # Ensure the connection is up and running await self.ensure_connection() + # Get a client + client = await self.get_client() + for version, model in self.models["neuron"].items(): try: # Attempt to delete all given neurons - await model.delete_all(self.database, neurons) + await model.delete_all(client, neurons) except Exception as ex: btul.logging.error( @@ -185,8 +200,11 @@ async def set_last_updated(self, block: int): # Ensure the connection is up and running await self.ensure_connection() + # Get a client + client = await self.get_client() + try: - await self.database.set(self._key("state:neuron:last_updated"), str(block)) + await client.set(self._key("state:neuron:last_updated"), str(block)) except Exception as ex: btul.logging.error( @@ -210,12 +228,15 @@ async def notify_state(self): # Ensure the connection is up and running await self.ensure_connection() + # Get a client + client = await self.get_client() + try: # Get the current state of the metagraph - state = await self.database.get(self._key("state:metagraph")) + state = await client.get(self._key("state:metagraph")) # Notify downstream via Redis stream - await self.database.xadd( + await client.xadd( self._key("state:metagraph:stream"), {"state": state} ) @@ -237,9 +258,12 @@ async def _set_state(self, state: str): # Ensure the connection is up and running await self.ensure_connection() + # Get a client + client = await self.get_client() + try: # Set the current metagraph state - await self.database.set(self._key("state:metagraph"), state) + await client.set(self._key("state:metagraph"), state) except Exception as ex: btul.logging.error( diff --git a/subvortex/core/metagraph/metagraph.py b/subvortex/core/metagraph/metagraph.py index 556e40e8..38d5940a 100644 --- a/subvortex/core/metagraph/metagraph.py +++ b/subvortex/core/metagraph/metagraph.py @@ -50,7 +50,9 @@ async def start(self): "🚀 MetagraphObserver service starting...", prefix=self.settings.logging_name, ) - btul.logging.debug(f"Settings: {self.settings}") + btul.logging.debug( + f"Settings: {self.settings}", prefix=self.settings.logging_name + ) # Load the neurons neurons = await self.database.get_neurons() @@ -61,13 +63,36 @@ async def start(self): try: while not self.should_exit.is_set(): try: - # Wait for next block to proceed - if not await scbs.wait_for_block(subtensor=self.subtensor): + + # Wait for either a new block OR a shutdown signal, whichever comes first. + done, _ = await asyncio.wait( + [ + self.subtensor.wait_for_block(), + self.should_exit.wait(), + ], + timeout=24, + return_when=asyncio.FIRST_COMPLETED, + ) + + # Timeout, no tasks completed + if not done: + btul.logging.warning( + "⏲️ No new block retrieved within 24 seconds. Retrying..." + ) + continue + + # If shutdown signal is received, break the loop immediately + if self.should_exit.is_set(): + break + + # If no new block was produced (e.g., shutdown happened or something failed), skip this round + # This guards against the case where wait_for_block() returned None or False + if not any(task.result() for task in done if not task.cancelled()): continue block = await self.subtensor.get_current_block() btul.logging.info( - f"📦 Block #{block} detected", prefix=self.settings.logging_name + f"📦 Block #{block}", prefix=self.settings.logging_name ) # Detect any new neuron registration @@ -125,6 +150,10 @@ async def start(self): # Store the new axons axons = new_axons + except ConnectionRefusedError as e: + btul.logging.error(f"Connection refused: {e}") + await asyncio.sleep(1) + except Exception as e: btul.logging.error( f"❌ Unhandled error in loop: {e}", @@ -135,7 +164,17 @@ async def start(self): ) finally: + # Ensure the metagraph is state as unready! + if not self.settings.dry_run: + btul.logging.debug( + f" Metagraph marked unready", prefix=self.settings.logging_name + ) + await self.database.mark_as_unready() + await self.database.notify_state() + + # Signal the run is completed self.run_complete.set() + btul.logging.info( "🛑 MetagraphObserver service exiting...", prefix=self.settings.logging_name, @@ -145,6 +184,10 @@ async def stop(self): """ Signals the observer to stop and waits for the loop to exit cleanly. """ + btul.logging.info( + f"MetagraphObserver stopping...", prefix=self.settings.logging_name + ) + # Signal the service to exit self.should_exit.set() diff --git a/subvortex/core/tests/src/test_metagraph_database.py b/subvortex/core/tests/src/test_metagraph_database.py index 1a9fb37d..da6d2a59 100644 --- a/subvortex/core/tests/src/test_metagraph_database.py +++ b/subvortex/core/tests/src/test_metagraph_database.py @@ -43,6 +43,7 @@ def database(): # Patch the database and connection s.database = mock_redis s.ensure_connection = AsyncMock() + s.get_client = AsyncMock(return_value=mock_redis) # Patch _get_migration_status to return mock version s._get_migration_status = AsyncMock(return_value=(None, ["2.1.0"])) @@ -57,6 +58,7 @@ def database(): mock_model.delete_all = AsyncMock() s.models["neuron"]["2.1.0"] = mock_model s._mock_model = mock_model # optional for assertions + s._mock_redis = mock_redis return s @@ -77,7 +79,7 @@ async def test_update_neurons(database): neurons = [scmm.Neuron(hotkey="hk1"), scmm.Neuron(hotkey="hk2")] await database.update_neurons(neurons) - database._mock_model.write_all.assert_awaited_once_with(database.database, neurons) + database._mock_model.write_all.assert_awaited_once_with(database._mock_redis, neurons) @pytest.mark.asyncio @@ -85,7 +87,7 @@ async def test_get_neuron_found(database): neuron = await database.get_neuron("hk1") assert neuron is not None assert neuron.hotkey == "hk1" - database._mock_model.read.assert_awaited_once_with(database.database, "hk1") + database._mock_model.read.assert_awaited_once_with(database._mock_redis, "hk1") @pytest.mark.asyncio @@ -101,7 +103,7 @@ async def test_get_neurons(database): async def test_remove_neurons(database): neurons = [scmm.Neuron(hotkey="hk1")] await database.remove_neurons(neurons) - database._mock_model.delete_all.assert_awaited_once_with(database.database, neurons) + database._mock_model.delete_all.assert_awaited_once_with(database._mock_redis, neurons) @pytest.mark.asyncio diff --git a/subvortex/core/tests/src/test_metagraph_observer.py b/subvortex/core/tests/src/test_metagraph_observer.py index 6ac8c91a..9186b359 100644 --- a/subvortex/core/tests/src/test_metagraph_observer.py +++ b/subvortex/core/tests/src/test_metagraph_observer.py @@ -237,30 +237,30 @@ async def test_start_and_stop(observer): observer._has_neuron_ip_changed = AsyncMock(return_value=(False, {})) observer.subtensor.get_current_block = AsyncMock(return_value=123) + # ✅ Patch self.subtensor.wait_for_block directly instead of global patch + observer.subtensor.wait_for_block = AsyncMock() + # Force exit after one loop using `should_exit.set()` inside the loop # Patch wait_for_block so it returns immediately - with patch( - "subvortex.core.core_bittensor.subtensor.wait_for_block", new_callable=AsyncMock - ) as mock_wait: - - async def wait_for_block_mock(*args, **kwargs): - await asyncio.sleep(0.1) - return True - - mock_wait.side_effect = wait_for_block_mock - - task = asyncio.create_task(observer.start()) + async def wait_for_block_mock(*args, **kwargs): await asyncio.sleep(0.1) - await observer.stop() + return True + + observer.subtensor.wait_for_block.side_effect = wait_for_block_mock + + task = asyncio.create_task(observer.start()) + await asyncio.sleep(0.2) + await observer.stop() - # Await the start task to finish gracefully - await asyncio.wait_for(task, timeout=1.0) + # Await the start task to finish gracefully + await asyncio.wait_for(task, timeout=1.0) # Assertions assert observer.run_complete.is_set() observer._resync.assert_called_once() observer._notify_if_needed.assert_called_once() observer._has_new_neuron_registered.assert_called_once() + observer._has_neuron_ip_changed.assert_called_once() @pytest.mark.asyncio diff --git a/subvortex/miner/metagraph/requirements.txt b/subvortex/miner/metagraph/requirements.txt index e7e28654..5938fe66 100644 --- a/subvortex/miner/metagraph/requirements.txt +++ b/subvortex/miner/metagraph/requirements.txt @@ -1,4 +1,4 @@ -bittensor==9.6.0 +bittensor==9.7.0 bittensor-wallet==3.0.10 loguru==0.7.0 numpy==2.0.1 diff --git a/subvortex/miner/metagraph/src/main.py b/subvortex/miner/metagraph/src/main.py index 26ba4b85..3c3cb3ff 100644 --- a/subvortex/miner/metagraph/src/main.py +++ b/subvortex/miner/metagraph/src/main.py @@ -9,11 +9,8 @@ import bittensor.core.config as btcc import bittensor.core.async_subtensor as btcas import bittensor.core.metagraph as btcm -import bittensor.core.settings as btcs -import bittensor_wallet.utils as btwu import subvortex.core.core_bittensor.config.config_utils as scccu -import subvortex.core.core_bittensor.subtensor as sccs import subvortex.core.metagraph.metagraph as scmm import subvortex.core.metagraph.database as scmms import subvortex.core.version as scv @@ -21,6 +18,9 @@ load_dotenv(override=True) +# An asyncio event to signal when shutdown is complete +shutdown_complete = asyncio.Event() + class Runner: def __init__(self): @@ -74,17 +74,6 @@ async def start(self): # Initialize the subtensor self.subtensor = btcas.AsyncSubtensor(config=config, retry_forever=True) - - # TODO: remove once OTF patched it - self.subtensor.substrate = sccs.RetryAsyncSubstrate( - url=self.subtensor.chain_endpoint, - ss58_format=btwu.SS58_FORMAT, - type_registry=btcs.TYPE_REGISTRY, - retry_forever=True, - use_remote_preset=True, - chain_name="Bittensor", - _mock=False, - ) await self.subtensor.initialize() btul.logging.info(str(self.subtensor)) @@ -123,28 +112,43 @@ async def shutdown(self): btul.logging.info("Shutting down completed") -if __name__ == "__main__": +async def main(): + # Initialize runner runner = Runner() - def _handle_signal(sig, frame): - loop.call_soon_threadsafe(lambda: asyncio.create_task(runner.shutdown())) + # Get the current asyncio event loop + loop = asyncio.get_running_loop() + + # Define a signal handler that schedules the shutdown coroutine + def _signal_handler(): + # Schedule graceful shutdown without blocking the signal handler + loop.create_task(_shutdown(runner)) + + # Register signal handlers for SIGINT (Ctrl+C) and SIGTERM (kill command) + for sig in (signal.SIGINT, signal.SIGTERM): + loop.add_signal_handler(sig, _signal_handler) - # Create and set a new event loop - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) + # Start the main service logic + await runner.start() - # Register signal handlers (in main thread) - signal.signal(signal.SIGINT, _handle_signal) - signal.signal(signal.SIGTERM, _handle_signal) + # Block here until shutdown is signaled and completed + await shutdown_complete.wait() + +async def _shutdown(runner: Runner): + # Gracefully shut down the service + await runner.shutdown() + + # Notify the main function that shutdown is complete + shutdown_complete.set() + + +if __name__ == "__main__": try: - # Run the metagraph - loop.run_until_complete(runner.start()) + # Start the main asyncio loop + asyncio.run(main()) except Exception as e: + # Log any unexpected exceptions that bubble up btul.logging.error(f"Unhandled exception: {e}") btul.logging.debug(traceback.format_exc()) - - finally: - loop.run_until_complete(loop.shutdown_asyncgens()) - loop.close() diff --git a/subvortex/miner/neuron/requirements.txt b/subvortex/miner/neuron/requirements.txt index 3296f323..9e5fce64 100644 --- a/subvortex/miner/neuron/requirements.txt +++ b/subvortex/miner/neuron/requirements.txt @@ -1,4 +1,4 @@ -bittensor==9.6.0 +bittensor==9.7.0 bittensor-wallet==3.0.10 loguru==0.7.0 netfilterqueue==1.1.0; sys_platform == 'linux' diff --git a/subvortex/miner/neuron/src/database.py b/subvortex/miner/neuron/src/database.py index 0920ec2c..0b48cd50 100644 --- a/subvortex/miner/neuron/src/database.py +++ b/subvortex/miner/neuron/src/database.py @@ -1,3 +1,8 @@ +import typing +import traceback +import bittensor.utils.btlogging as btul + +from subvortex.miner.neuron.src.models.score import MinerScore100, Score from subvortex.core.metagraph.database import NeuronReadOnlyDatabase @@ -20,3 +25,95 @@ def __init__(self, settings): super().__init__(settings=settings) self.setup_neuron_models() + self.models["score"] = {x.version: x for x in [MinerScore100()]} + + async def get_scores(self) -> typing.List[Score]: + # Ensure the connection is up and running + await self.ensure_connection() + + # Get a client + client = await self.get_client() + + # Get the active versions + _, active = await self._get_migration_status("score") + + for version in reversed(active): + model = self.models["score"][version] + if not model: + continue + + try: + # Attempt to read all neurons using the model + neurons = await model.read_all(client) + return neurons + + except Exception as ex: + btul.logging.warning( + f"[get_scores] Failed to read all scores using version={version}: {ex}", + prefix=self.settings.logging_name, + ) + btul.logging.debug( + f"[get_scores] Exception type: {type(ex).__name__}, Traceback:\n{traceback.format_exc()}", + prefix=self.settings.logging_name, + ) + + return [] + + async def save_scores(self, score: Score, max_entries: int = 100): + """ + Bulk update for a list of miners using active model versions. + """ + await self.ensure_connection() + + # Get a client + client = await self.get_client() + + _, active = await self._get_migration_status("score") + + for version in reversed(active): + model = self.models["score"][version] + if not model: + continue + + try: + await model.write(client, score) + + except Exception as ex: + btul.logging.warning( + f"[{version}] Update score failed: {ex}", + prefix=self.settings.logging_name, + ) + btul.logging.debug( + f"[update_score] Exception type: {type(ex).__name__}, Traceback:\n{traceback.format_exc()}", + prefix=self.settings.logging_name, + ) + + return None + + async def prune_scores(self, max_entries: int): + await self.ensure_connection() + + # Get a client + client = await self.get_client() + + _, active = await self._get_migration_status("score") + + for version in reversed(active): + model = self.models["score"][version] + if not model: + continue + + try: + await model.prune(client, max_entries) + + except Exception as ex: + btul.logging.warning( + f"[{version}] Update score failed: {ex}", + prefix=self.settings.logging_name, + ) + btul.logging.debug( + f"[update_score] Exception type: {type(ex).__name__}, Traceback:\n{traceback.format_exc()}", + prefix=self.settings.logging_name, + ) + + return None \ No newline at end of file diff --git a/subvortex/miner/neuron/src/main.py b/subvortex/miner/neuron/src/main.py index e932200a..1ee9f20f 100644 --- a/subvortex/miner/neuron/src/main.py +++ b/subvortex/miner/neuron/src/main.py @@ -26,9 +26,6 @@ import bittensor_wallet.mock as btwm from dotenv import load_dotenv -import bittensor.core.settings as btcs -import bittensor_wallet.utils as btwu - from subvortex.core.protocol import Score from subvortex.core.shared.neuron import wait_until_registered from subvortex.core.shared.substrate import ( @@ -42,7 +39,6 @@ from subvortex.core.core_bittensor.metagraph import SubVortexMetagraph from subvortex.core.core_bittensor.axon import SubVortexAxon from subvortex.core.core_bittensor.synapse import Synapse -from subvortex.core.core_bittensor.subtensor import wait_for_block, RetryAsyncSubstrate from subvortex.core.model.neuron.neuron import Neuron from subvortex.core.sse.sse_thread import SSEThread @@ -68,6 +64,9 @@ # Load the environment variables for the whole process load_dotenv(override=True) +# An asyncio event to signal when shutdown is complete +shutdown_complete = asyncio.Event() + class Miner: @classmethod @@ -145,6 +144,8 @@ async def run(self): self.version = get_version() btul.logging.debug(f"Version: {self.version}") + self.loop = asyncio.get_running_loop() + await self._initialize() await self._serve() await self._main_loop() @@ -152,6 +153,34 @@ async def run(self): # Signal the neuron has finished self.run_complete.set() + async def shutdown(self): + btul.logging.info("Shutting down miner...") + + # Notify the miner to stop + self.should_exit.set() + + # Wait the neuron to stop + await self.run_complete.wait() + + if getattr(self, "axon", None): + self.axon.stop() + btul.logging.debug("Axon stopped") + + if getattr(self, "subtensor", None): + await self.subtensor.close() + btul.logging.debug("Subtensor stopped") + + if getattr(self, "sse", None): + self.sse.stop() + + if getattr(self, "firewall", None): + self.firewall.stop() + + if getattr(self, "file_monitor", None): + self.file_monitor.stop() + + btul.logging.info("Shutting down miner completed") + async def _initialize(self): self.wallet = ( btwm.MockWallet(config=self.config) @@ -169,16 +198,6 @@ async def _initialize(self): config=self.config, network=network, retry_forever=True ) ) - # TODO: remove once OTF patched it - self.subtensor.substrate = RetryAsyncSubstrate( - url=self.subtensor.chain_endpoint, - ss58_format=btwu.SS58_FORMAT, - type_registry=btcs.TYPE_REGISTRY, - retry_forever=True, - use_remote_preset=True, - chain_name="Bittensor", - _mock=False, - ) await self.subtensor.initialize() # Initialize database @@ -232,17 +251,20 @@ async def _serve(self): wallet=self.wallet, config=self.config, external_ip=btun.get_external_ip(), - blacklist_fn=self._blacklist, + blacklist_fn=self._sync_blacklist_handler, ) if self.config.mock else SubVortexAxon( wallet=self.wallet, config=self.config, external_ip=btun.get_external_ip(), - blacklist_fn=self._blacklist, + blacklist_fn=self._sync_blacklist_handler, ) ) - self.axon.attach(forward_fn=self._score, blacklist_fn=self._blacklist_score) + self.axon.attach( + forward_fn=self._score, + blacklist_fn=self._sync_blacklist_score_handler, + ) if not self.settings.dry_run: # Start the axon @@ -256,13 +278,39 @@ async def _serve(self): async def _main_loop(self): while not self.should_exit.is_set(): try: - # Wait for the next block - if not await wait_for_block(subtensor=self.subtensor): + # Wait for either a new block OR a shutdown signal, whichever comes first. + done, _ = await asyncio.wait( + [ + self.subtensor.wait_for_block(), + self.should_exit.wait(), + ], + timeout=24, + return_when=asyncio.FIRST_COMPLETED, + ) + + # Timeout, no tasks completed + if not done: + btul.logging.warning( + "⏲️ No new block retrieved within 24 seconds. Retrying..." + ) + continue + + # If shutdown signal is received, break the loop immediately + if self.should_exit.is_set(): + break + + # If no new block was produced (e.g., shutdown happened or something failed), skip this round + # This guards against the case where wait_for_block() returned None or False + if not any(task.result() for task in done if not task.cancelled()): continue # Get the current block current_block = await self.subtensor.get_current_block() - btul.logging.debug(f"Block #{current_block}") + btul.logging.debug(f"📦 Block #{current_block}") + + # Ensure the metagraph is ready + btul.logging.debug("Ensure metagraph readiness") + await self.database.wait_until_ready("metagraph") # Get the last time the neurons have been updated last_updated = await self.database.get_neuron_last_updated() @@ -320,6 +368,10 @@ async def _main_loop(self): # We already display a log, so need to do more here pass + except ConnectionRefusedError as e: + btul.logging.error(f"Connection refused: {e}") + await asyncio.sleep(1) + except Exception as ex: btul.logging.error(f"Unhandled exception in main loop: {ex}") btul.logging.debug(traceback.format_exc()) @@ -354,34 +406,6 @@ async def _update_firewall(self): ) btul.logging.debug("Firewall updated") - async def _shutdown(self): - btul.logging.info("Shutting down miner...") - - # Notify the miner to stop - self.should_exit.set() - - # Wait the neuron to stop - await self.run_complete.wait() - - if getattr(self, "axon", None): - self.axon.stop() - btul.logging.debug("Axon stopped") - - if getattr(self, "subtensor", None): - await self.subtensor.close() - btul.logging.debug("Subtensor stopped") - - if getattr(self, "sse", None): - self.sse.stop() - - if getattr(self, "firewall", None): - self.firewall.stop() - - if getattr(self, "file_monitor", None): - self.file_monitor.stop() - - btul.logging.info("Shutting down miner completed") - async def _blacklist(self, synapse: Synapse) -> typing.Tuple[bool, str]: caller = synapse.dendrite.hotkey caller_version = synapse.dendrite.neuron_version or 0 @@ -471,32 +495,65 @@ def _score(self, synapse: Score) -> Score: return synapse - async def _blacklist_score(self, synapse: Score) -> typing.Tuple[bool, str]: - return await self._blacklist(synapse) + def _sync_blacklist_handler(self, synapse: Synapse) -> typing.Tuple[bool, str]: + result = self._run_safe(self._blacklist(synapse)) + return result or (True, "Error during blacklist check") + def _sync_blacklist_score_handler(self, synapse: Score) -> typing.Tuple[bool, str]: + result = self._run_safe(self._blacklist(synapse)) + return result or (True, "Error during blacklist check") -if __name__ == "__main__": + def _run_safe(self, coro: typing.Coroutine, timeout: float = 10): + try: + future = asyncio.run_coroutine_threadsafe(coro, self.loop) + return future.result(timeout=timeout) + + except ConnectionRefusedError as e: + pass + + except Exception as e: + btul.logging.error(f"Error in threaded Axon handler: {e}") + btul.logging.debug(traceback.format_exc()) + return None + + +async def main(): + # Initialize miner miner = Miner() - def _handle_signal(sig, frame): - loop.call_soon_threadsafe(lambda: asyncio.create_task(miner._shutdown())) + # Get the current asyncio event loop + loop = asyncio.get_running_loop() + + # Define a signal handler that schedules the shutdown coroutine + def _signal_handler(): + # Schedule graceful shutdown without blocking the signal handler + loop.create_task(_shutdown(miner)) + + # Register signal handlers for SIGINT (Ctrl+C) and SIGTERM (kill command) + for sig in (signal.SIGINT, signal.SIGTERM): + loop.add_signal_handler(sig, _signal_handler) + + # Start the main service logic + await miner.run() - # Create and set a new event loop - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) + # Block here until shutdown is signaled and completed + await shutdown_complete.wait() - # Register signal handlers (in main thread) - signal.signal(signal.SIGINT, _handle_signal) - signal.signal(signal.SIGTERM, _handle_signal) +async def _shutdown(miner: Miner): + # Gracefully shut down the service + await miner.shutdown() + + # Notify the main function that shutdown is complete + shutdown_complete.set() + + +if __name__ == "__main__": try: - # Run the miner - loop.run_until_complete(miner.run()) + # Start the main asyncio loop + asyncio.run(main()) except Exception as e: + # Log any unexpected exceptions that bubble up btul.logging.error(f"Unhandled exception: {e}") btul.logging.debug(traceback.format_exc()) - finally: - # Cleanup async generators and close loop - loop.run_until_complete(loop.shutdown_asyncgens()) - loop.close() diff --git a/subvortex/validator/metagraph/src/main.py b/subvortex/validator/metagraph/src/main.py index 8d67efed..36d0689e 100644 --- a/subvortex/validator/metagraph/src/main.py +++ b/subvortex/validator/metagraph/src/main.py @@ -9,11 +9,8 @@ import bittensor.core.config as btcc import bittensor.core.async_subtensor as btcas import bittensor.core.metagraph as btcm -import bittensor.core.settings as btcs -import bittensor_wallet.utils as btwu import subvortex.core.core_bittensor.config.config_utils as scccu -import subvortex.core.core_bittensor.subtensor as sccs import subvortex.core.metagraph.metagraph as scmm import subvortex.core.metagraph.database as scmms import subvortex.core.version as scv @@ -21,6 +18,9 @@ load_dotenv(override=True) +# An asyncio event to signal when shutdown is complete +shutdown_complete = asyncio.Event() + class Runner: def __init__(self): @@ -74,17 +74,6 @@ async def start(self): # Initialize the subtensor self.subtensor = btcas.AsyncSubtensor(config=config, retry_forever=True) - - # TODO: remove once OTF patched it - self.subtensor.substrate = sccs.RetryAsyncSubstrate( - url=self.subtensor.chain_endpoint, - ss58_format=btwu.SS58_FORMAT, - type_registry=btcs.TYPE_REGISTRY, - retry_forever=True, - use_remote_preset=True, - chain_name="Bittensor", - _mock=False, - ) await self.subtensor.initialize() btul.logging.info(str(self.subtensor)) @@ -123,28 +112,43 @@ async def shutdown(self): btul.logging.info("Shutting down completed") -if __name__ == "__main__": +async def main(): + # Initialize runner runner = Runner() - def _handle_signal(sig, frame): - loop.call_soon_threadsafe(lambda: asyncio.create_task(runner.shutdown())) + # Get the current asyncio event loop + loop = asyncio.get_running_loop() + + # Define a signal handler that schedules the shutdown coroutine + def _signal_handler(): + # Schedule graceful shutdown without blocking the signal handler + loop.create_task(_shutdown(runner)) + + # Register signal handlers for SIGINT (Ctrl+C) and SIGTERM (kill command) + for sig in (signal.SIGINT, signal.SIGTERM): + loop.add_signal_handler(sig, _signal_handler) - # Create and set a new event loop - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) + # Start the main service logic + await runner.start() - # Register signal handlers (in main thread) - signal.signal(signal.SIGINT, _handle_signal) - signal.signal(signal.SIGTERM, _handle_signal) + # Block here until shutdown is signaled and completed + await shutdown_complete.wait() + +async def _shutdown(runner: Runner): + # Gracefully shut down the service + await runner.shutdown() + + # Notify the main function that shutdown is complete + shutdown_complete.set() + + +if __name__ == "__main__": try: - # Run the metagraph - loop.run_until_complete(runner.start()) + # Start the main asyncio loop + asyncio.run(main()) except Exception as e: + # Log any unexpected exceptions that bubble up btul.logging.error(f"Unhandled exception: {e}") btul.logging.debug(traceback.format_exc()) - - finally: - loop.run_until_complete(loop.shutdown_asyncgens()) - loop.close() diff --git a/subvortex/validator/neuron/src/database.py b/subvortex/validator/neuron/src/database.py index d788cf85..74cb8391 100644 --- a/subvortex/validator/neuron/src/database.py +++ b/subvortex/validator/neuron/src/database.py @@ -48,6 +48,10 @@ async def get_selected_miners(self, ss58_address: str): Return selected uids for a hotkey using versioned selection models. """ await self.ensure_connection() + + # Get a client + client = await self.get_client() + _, active = await self._get_migration_status("selection") for version in reversed(active): # Prefer latest version @@ -56,7 +60,7 @@ async def get_selected_miners(self, ss58_address: str): continue try: - uids = await model.read(self.database, ss58_address) + uids = await model.read(client, ss58_address) return uids except Exception as err: @@ -76,6 +80,10 @@ async def set_selection_miners(self, ss58_address: str, uids: list[int]): Store selected miner UIDs in all active selection model versions. """ await self.ensure_connection() + + # Get a client + client = await self.get_client() + _, active = await self._get_migration_status("selection") for version in active: @@ -84,7 +92,7 @@ async def set_selection_miners(self, ss58_address: str, uids: list[int]): continue try: - await model.write(self.database, ss58_address, uids) + await model.write(client, ss58_address, uids) except Exception as err: btul.logging.error( @@ -100,6 +108,10 @@ async def set_selection_miners(self, ss58_address: str, uids: list[int]): async def get_miner(self, hotkey: str) -> Miner: # Ensure the connection is up and running await self.ensure_connection() + + # Get a client + client = await self.get_client() + _, active = await self._get_migration_status("miner") for version in reversed(active): @@ -108,7 +120,7 @@ async def get_miner(self, hotkey: str) -> Miner: continue try: - miner = await model.read(self.database, hotkey) + miner = await model.read(client, hotkey) return miner except Exception as ex: @@ -126,6 +138,10 @@ async def get_miner(self, hotkey: str) -> Miner: async def get_miners(self) -> dict[str, Miner]: # Ensure the connection is up and running await self.ensure_connection() + + # Get a client + client = await self.get_client() + _, active = await self._get_migration_status("miner") for version in reversed(active): @@ -134,7 +150,7 @@ async def get_miners(self) -> dict[str, Miner]: continue try: - miners = await model.read_all(self.database) + miners = await model.read_all(client) return miners except Exception as ex: @@ -154,6 +170,10 @@ async def add_miner(self, miner: Miner): Add a new miner record to all active versions of the miner schema. """ await self.ensure_connection() + + # Get a client + client = await self.get_client() + _, active = await self._get_migration_status("miner") for version in active: @@ -162,7 +182,7 @@ async def add_miner(self, miner: Miner): continue try: - await model.write(self.database, miner) + await model.write(client, miner) except Exception as ex: btul.logging.error( @@ -181,6 +201,10 @@ async def update_miner(self, miner: Miner): Update an existing miner record. """ await self.ensure_connection() + + # Get a client + client = await self.get_client() + _, active = await self._get_migration_status("miner") for version in reversed(active): @@ -189,7 +213,7 @@ async def update_miner(self, miner: Miner): continue try: - await model.write(self.database, miner) + await model.write(client, miner) except Exception as ex: btul.logging.warning( @@ -208,6 +232,10 @@ async def update_miners(self, miners: List[Miner]): Bulk update for a list of miners using active model versions. """ await self.ensure_connection() + + # Get a client + client = await self.get_client() + _, active = await self._get_migration_status("miner") for version in reversed(active): @@ -216,7 +244,7 @@ async def update_miners(self, miners: List[Miner]): continue try: - await model.write_all(self.database, miners) + await model.write_all(client, miners) except Exception as ex: btul.logging.warning( @@ -236,9 +264,12 @@ async def remove_miner(self, miner: Miner): """ await self.ensure_connection() + # Get a client + client = await self.get_client() + for version, model in self.models["miner"].items(): try: - await model.delete(self.database, miner) + await model.delete(client, miner) except Exception as ex: btul.logging.error( @@ -258,9 +289,12 @@ async def remove_miners(self, miners: List[Miner]): """ await self.ensure_connection() + # Get a client + client = await self.get_client() + for version, model in self.models["miner"].items(): try: - await model.delete_all(self.database, miners) + await model.delete_all(client, miners) except Exception as ex: btul.logging.error( diff --git a/subvortex/validator/neuron/src/main.py b/subvortex/validator/neuron/src/main.py index a6a16522..9ad0700a 100644 --- a/subvortex/validator/neuron/src/main.py +++ b/subvortex/validator/neuron/src/main.py @@ -67,6 +67,8 @@ # Load the environment variables for the whole process load_dotenv(override=True) +# An asyncio event to signal when shutdown is complete +shutdown_complete = asyncio.Event() class Validator: """ @@ -218,6 +220,10 @@ async def run(self): current_block = 0 while not self.should_exit.is_set(): try: + # Ensure the metagraph is ready + btul.logging.debug("Ensure metagraph readiness") + await self.database.wait_until_ready("metagraph") + # Get the last time neurons have been updated last_updated = await self.database.get_neuron_last_updated() if last_updated == 0: @@ -394,28 +400,43 @@ async def _shutdown(self): btul.logging.info("✅ Shutting down validator completed") -if __name__ == "__main__": +async def main(): + # Initialize miner validator = Validator() - def _handle_signal(sig, frame): - loop.call_soon_threadsafe(lambda: asyncio.create_task(validator._shutdown())) + # Get the current asyncio event loop + loop = asyncio.get_running_loop() + + # Define a signal handler that schedules the shutdown coroutine + def _signal_handler(): + # Schedule graceful shutdown without blocking the signal handler + loop.create_task(_shutdown(validator)) + + # Register signal handlers for SIGINT (Ctrl+C) and SIGTERM (kill command) + for sig in (signal.SIGINT, signal.SIGTERM): + loop.add_signal_handler(sig, _signal_handler) + + # Start the main service logic + await validator.run() - # Create and set a new event loop - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) + # Block here until shutdown is signaled and completed + await shutdown_complete.wait() - # Register signal handlers (in main thread) - signal.signal(signal.SIGINT, _handle_signal) - signal.signal(signal.SIGTERM, _handle_signal) +async def _shutdown(validator: Validator): + # Gracefully shut down the service + await validator.shutdown() + + # Notify the main function that shutdown is complete + shutdown_complete.set() + + +if __name__ == "__main__": try: - # Run the miner - loop.run_until_complete(validator.run()) + # Start the main asyncio loop + asyncio.run(main()) except Exception as e: + # Log any unexpected exceptions that bubble up btul.logging.error(f"Unhandled exception: {e}") btul.logging.debug(traceback.format_exc()) - finally: - # Cleanup async generators and close loop - loop.run_until_complete(loop.shutdown_asyncgens()) - loop.close() diff --git a/subvortex/validator/neuron/tests/src/test_database.py b/subvortex/validator/neuron/tests/src/test_database.py index 7cf140aa..a8292b0b 100644 --- a/subvortex/validator/neuron/tests/src/test_database.py +++ b/subvortex/validator/neuron/tests/src/test_database.py @@ -12,15 +12,19 @@ class DummySettings: key_prefix = "sv" logging_name = "test" - redis_url = "redis://localhost:6379/0" + database_host="localhost" + database_port=6379 + database_index=0 + database_password=None @pytest_asyncio.fixture async def db(): db = Database(DummySettings()) db.ensure_connection = AsyncMock() - db.database = AsyncMock() - db.client = AsyncMock() + db.get_client = AsyncMock() + client = AsyncMock() + db.get_client.return_value = client return db @@ -49,7 +53,7 @@ async def test_set_selection_miners_calls_write(db): await db.set_selection_miners("hotkey2", uids) db.models["selection"][version].write.assert_called_once_with( - db.database, "hotkey2", uids + db.get_client.return_value, "hotkey2", uids ) @@ -114,7 +118,7 @@ async def test_add_miner_calls_write(db): db.models["miner"][version].write = AsyncMock() await db.add_miner(miner) - db.models["miner"][version].write.assert_called_once_with(db.database, miner) + db.models["miner"][version].write.assert_called_once_with(db.get_client.return_value, miner) @pytest.mark.asyncio @@ -126,7 +130,7 @@ async def test_update_miners_batch_success(db): db.models["miner"][version].write_all = AsyncMock() await db.update_miners(miners) - db.models["miner"][version].write_all.assert_called_once_with(db.database, miners) + db.models["miner"][version].write_all.assert_called_once_with(db.get_client.return_value, miners) @pytest.mark.asyncio @@ -138,12 +142,12 @@ async def test_remove_miner_calls_delete(db): db.models["miner"][version].delete = AsyncMock() await db.remove_miner(miner) - db.models["miner"][version].delete.assert_called_once_with(db.database, miner) + db.models["miner"][version].delete.assert_called_once_with(db.get_client.return_value, miner) @pytest.mark.asyncio async def test_get_last_update_success(db): - db.database.get = AsyncMock(return_value=b"1000") + db.get_client.return_value.get = AsyncMock(return_value=b"1000") result = await db.get_neuron_last_updated() assert result == 1000 @@ -152,12 +156,12 @@ async def test_get_last_update_success(db): @pytest.mark.asyncio async def test_get_migration_status_returns_active_versions(db): db.models["selection"] = {"2.0.0": SelectionModel200()} - db.database.get = AsyncMock(return_value=b"new") + db.get_client.return_value.get = AsyncMock(return_value=b"new") latest, active = await db._get_migration_status("selection") assert latest == "2.0.0" assert active == ["2.0.0"] - db.database.get.assert_called_once_with("migration_mode:2.0.0") + db.get_client.return_value.get.assert_called_once_with("migration_mode:2.0.0") @pytest.mark.asyncio @@ -166,7 +170,7 @@ async def test_get_migration_status_fallback(db): "2.0.0": SelectionModel200(), "2.1.0": SelectionModel200(), } - db.database.get = AsyncMock(return_value=None) + db.get_client.return_value.get = AsyncMock(return_value=None) latest, active = await db._get_migration_status("selection") assert latest == "2.1.0"