Catch MemoryError in model manager

Author: dbobrenko

neurons/validator.py

@@ -80,6 +80,8 @@ async def spawn_loops(task_queue: list, scoring_queue: list, reward_events: list
                 logger.debug(
                     f"Task Queue {len(task_queue)}. Scoring Queue {len(scoring_queue)}. Reward Events {len(reward_events)}"
                 )
+                if model_scheduler.memory_error is not None:
+                    raise model_scheduler.memory_error
         except asyncio.CancelledError:
             logger.info("spawn_loops received cancellation signal.")
             raise

prompting/llms/model_manager.py

@@ -255,6 +255,7 @@ class AsyncModelScheduler(AsyncLoopRunner):
     mp_lock: AcquirerProxy
     interval: int = 1200
     scoring_queue: list | None = None
+    memory_error: MemoryError | None = None
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
@@ -277,6 +278,9 @@ async def run_step(self):
             logger.info(f"Model {selected_model.llm_model_id} is already loaded.")
             return
 
-        await self.llm_model_manager.load_model(selected_model)
+        try:
+            await self.llm_model_manager.load_model(selected_model)
+        except MemoryError as e:
+            self.memory_error = e
         logger.debug(f"Active models: {self.llm_model_manager.active_models.keys()}")
         await asyncio.sleep(0.01)

prompting/tasks/task_sending.py

@@ -82,7 +82,6 @@ class TaskSender(AsyncLoopRunner):
 
     task_queue: list | None = None
     scoring_queue: list | None = None
-    # subtensor: bt.Subtensor | None = None
     miners_dict: dict | None = None
 
     class Config:
@@ -92,9 +91,6 @@ async def start(self, task_queue, scoring_queue, miners_dict, **kwargs):
         self.task_queue = task_queue
         self.scoring_queue = scoring_queue
         self.miners_dict = miners_dict
-
-        # # shared_settings is not initialised inside this process, meaning it cannot access any non-constants from here
-        # self.subtensor = bt.subtensor(network=shared_settings.SUBTENSOR_NETWORK)
         return await super().start(**kwargs)
 
     @property
@@ -133,7 +129,6 @@ async def run_step(self) -> ValidatorLoggingEvent | ErrorLoggingEvent | None:
                 task_id=task.task_id,
             )
             self.scoring_queue.append(scoring_config)
-            # logger.debug(f"Scoring queue length: {len(self.scoring_queue)}")
 
             # Log the step event.
             return ValidatorLoggingEvent(

shared/epistula.py

@@ -149,16 +149,16 @@ async def query_miners(
         responses_valid = 0
         responses_error = 0
         responses_exception = 0
-        exception_info: Exception | None = None
         results: list[SynapseStreamResult] = []
         for response, uid in zip(responses, uids):
             if isinstance(response, Exception):
                 responses_exception += 1
-                exception_info = response
                 results.append(SynapseStreamResult(exception=str(response)))
             elif isinstance(response, tuple) and isinstance(response[0], ChatCompletion):
                 if response and response[1]:
                     responses_valid += 1
+                else:
+                    responses_error += 1
                 results.append(
                     SynapseStreamResult(
                         uid=uid,
@@ -174,8 +174,8 @@ async def query_miners(
 
         logger.info(
             f"Responses success: {responses_valid}/{len(uids)}. "
-            f"Responses exception: {responses_exception}/{len(uids)} [{exception_info}]. "
-            f"Reponses error: {responses_error}/{len(uids)}"
+            f"Responses exception: {responses_exception}/{len(uids)}. "
+            f"Reponses invalid: {responses_error}/{len(uids)}"
         )
         return results
     except Exception as e: