schd: Requeue task if training server is unavailable

Restructure the logic around server response handling to allow for more
fine-grained control. If the training server is reported unavailable due
to a training job already running, the task is requeued instead of
dropped. A shortcoming of this approach, given that RabbitMQ attempts to
place requeued messages in their original positions or near the head of
the queue, is that long-running tasks (e.g. training/evaluation) can
block the execution of others that could be processed in parallel on the
same servers (CMS only allows running a single training task at a time
to work around resource limitations, but places no restrictions on the
execution of short-running tasks). This could be addressed by
introducing a separate queue for long-running tasks, or by increasing
the priority of short-running tasks in the queue. In any case, we expect
to have a more robust solution once we move to a more sophisticated,
resource and task-aware scheduler.

Signed-off-by: Phoevos Kalemkeris <[email protected]>

Author: phoevos

cogstack_model_gateway/scheduler/scheduler.py

@@ -1,7 +1,7 @@
 import logging
 import time
 
-import requests
+from requests import Response, request
 
 from cogstack_model_gateway.common.object_store import ObjectStoreManager
 from cogstack_model_gateway.common.queue import QueueManager
@@ -30,7 +30,6 @@ def run(self):
         self.queue_manager.consume(self.process_task)
 
     def process_task(self, task: dict, ack: callable, nack: callable) -> None:
-        # FIXME: Handle ACK and NACK appropriately
         task_uuid = task["uuid"]
         log.info(f"Processing task '{task_uuid}'")
 
@@ -41,85 +40,41 @@ def process_task(self, task: dict, ack: callable, nack: callable) -> None:
         task_obj = self.handle_server_response(task_uuid, res, err_msg, ack, nack)
         self.send_notification(task_obj)
 
-    def route_task(self, task: dict) -> tuple[requests.Response, str]:
+    def route_task(self, task: dict) -> tuple[Response, str]:
         log.info(f"Routing task '{task['uuid']}' to model server at {task['url']}")
-        request = self._prepare_request(task)
-
+        req = self._prepare_request(task)
+        response = None
         try:
-            log.debug(f"Request: {request}")
-            response = requests.request(
-                method=request["method"],
-                url=request["url"],
-                headers=request["headers"],
-                params=request["params"],
-                data=request["data"],
-                files=request["files"],
+            log.debug(f"Request: {req}")
+            response = request(
+                method=req["method"],
+                url=req["url"],
+                headers=req["headers"],
+                params=req["params"],
+                data=req["data"],
+                files=req["files"],
             )
-        except Exception as e:
-            err_msg = f"Failed to forward task '{task['uuid']}': {e}"
-            log.error(err_msg)
-            return None, err_msg
-
-        try:
             log.debug(f"Response: {response.text}")
             response.raise_for_status()
             log.info(f"Task '{task['uuid']}' forwarded successfully to {task['url']}")
             return response, None
-        except requests.HTTPError:
-            err_msg = f"Failed to process task '{task['uuid']}']: {response.json()}"
+        except Exception as e:
+            err_msg = f"Failed to forward task '{task['uuid']}': {e}"
             log.error(err_msg)
-            return None, err_msg
+            return response, err_msg
 
     def handle_server_response(
         self,
         task_uuid: str,
-        response: requests.Response,
+        response: Response,
         err_msg: str,
         ack: callable,
         nack: callable,
     ) -> Task:
-        if response is None:
-            # FIXME: Perhaps set task to a different status?
-            # Pending and requeued? Or failed and done with?
-            # Should we reprocess failed tasks? How can we tell transient failures?
-            ack()
-            return self.task_manager.update_task(
-                task_uuid, status=Status.FAILED, error_message=err_msg or "Failed to process task"
-            )
-        ack()
-
-        if response.status_code == 202:
-            log.info(f"Task '{task_uuid}' accepted for processing, waiting for results")
-            tracking_id = response.json().get("run_id") if response.json() else None
-            self.task_manager.update_task(
-                task_uuid,
-                status=Status.RUNNING,
-                expected_status=Status.SCHEDULED,
-                tracking_id=tracking_id,
-            )
-
-            results = self.poll_task_status(task_uuid, tracking_id)
-            if results["status"] == Status.FAILED:
-                log.error(f"Task '{task_uuid}' failed: {results['error']}")
-                return self.task_manager.update_task(
-                    task_uuid, status=Status.FAILED, error_message=str(results["error"])
-                )
-            else:
-                log.info(f"Task '{task_uuid}' completed, writing results to object store")
-                object_key = self.results_object_store_manager.upload_object(
-                    results["url"].encode(), "results.url", prefix=task_uuid
-                )
-                return self.task_manager.update_task(
-                    task_uuid, status=Status.SUCCEEDED, result=object_key
-                )
+        if response is None or response.status_code >= 400:
+            return self._handle_task_failure(task_uuid, response, err_msg, nack)
         else:
-            log.info(f"Task '{task_uuid}' completed, writing results to object store")
-            object_key = self.results_object_store_manager.upload_object(
-                response.content, "results.json", prefix=task_uuid
-            )
-            return self.task_manager.update_task(
-                task_uuid, status=Status.SUCCEEDED, result=object_key
-            )
+            return self._handle_task_success(task_uuid, response, ack)
 
     def poll_task_status(self, task_uuid: str, tracking_id: str = None) -> dict:
         while True:
@@ -137,8 +92,9 @@ def poll_task_status(self, task_uuid: str, tracking_id: str = None) -> dict:
                 time.sleep(5)
 
     def send_notification(self, task: Task):
-        # FIXME: notify user
-        log.info(f"Task '{task.uuid}' {task.status.value}: {task.result or task.error_message}")
+        # FIXME: notify user if task is completed
+        if task.status.is_final():
+            log.info(f"Task '{task.uuid}' {task.status.value}: {task.result or task.error_message}")
 
     def _get_payload_from_refs(self, refs: list) -> str:
         if len(refs) > 1:
@@ -180,3 +136,68 @@ def _prepare_request(self, task: dict) -> dict:
             "files": files,
             "headers": headers,
         }
+
+    def _handle_task_failure(
+        self, task_uuid: str, response: Response, err_msg: str, nack: callable
+    ) -> Task:
+        # FIXME: Add fine-grained error handling for different status codes
+        if not response:
+            nack(requeue=False)
+            return self.task_manager.update_task(
+                task_uuid, status=Status.FAILED, error_message=err_msg or "Failed to process task"
+            )
+        elif (
+            response.status_code == 503
+            and (experiment_id := response.json().get("experiment_id"))
+            and (run_id := response.json().get("run_id"))
+        ):
+            warn_msg = (
+                f"Task '{task_uuid}' wasn't accepted for processing: a training run is already in"
+                f" progress (experiment_id={experiment_id}, run_id={run_id}). Requeuing task..."
+            )
+            log.warning(warn_msg)
+            nack()
+            return self.task_manager.update_task(
+                task_uuid, status=Status.PENDING, error_message=warn_msg
+            )
+        else:
+            log.error(f"Task '{task_uuid}' failed with unexpected error: {response.text}")
+            nack(requeue=False)
+            return self.task_manager.update_task(
+                task_uuid, status=Status.FAILED, error_message=response.text
+            )
+
+    def _handle_task_success(self, task_uuid: str, response: Response, ack: callable) -> Task:
+        ack()
+        if response.status_code == 202:
+            log.info(f"Task '{task_uuid}' accepted for processing, waiting for results")
+            tracking_id = response.json().get("run_id") if response.json() else None
+            self.task_manager.update_task(
+                task_uuid,
+                status=Status.RUNNING,
+                expected_status=Status.SCHEDULED,
+                tracking_id=tracking_id,
+            )
+
+            results = self.poll_task_status(task_uuid, tracking_id)
+            if results["status"] == Status.FAILED:
+                log.error(f"Task '{task_uuid}' failed: {results['error']}")
+                return self.task_manager.update_task(
+                    task_uuid, status=Status.FAILED, error_message=str(results["error"])
+                )
+            else:
+                log.info(f"Task '{task_uuid}' completed, writing results to object store")
+                object_key = self.results_object_store_manager.upload_object(
+                    results["url"].encode(), "results.url", prefix=task_uuid
+                )
+                return self.task_manager.update_task(
+                    task_uuid, status=Status.SUCCEEDED, result=object_key
+                )
+        else:
+            log.info(f"Task '{task_uuid}' completed, writing results to object store")
+            object_key = self.results_object_store_manager.upload_object(
+                response.content, "results.json", prefix=task_uuid
+            )
+            return self.task_manager.update_task(
+                task_uuid, status=Status.SUCCEEDED, result=object_key
+            )