Track killed worker for memory usage (#108)

* track killed worker log event from nanny

* add configurability to requeue

* request not found

* rename requeue env var

* change env variable

* check all qos status

* qa

Author: francesconazzaro

cads_broker/dispatcher.py

@@ -268,14 +268,30 @@ def sync_database(self, session: sa.orm.Session) -> None:
             # if it doesn't find the request: re-queue it
             else:
                 # FIXME: check if request status has changed
-                logger.info(
-                    "request not found: re-queueing", job_id={request.request_uid}
-                )
-                db.requeue_request(request_uid=request.request_uid, session=session)
-                self.queue.add(request.request_uid, request)
-                self.qos.notify_end_of_request(
-                    request, session, scheduler=self.internal_scheduler
-                )
+                if os.getenv(
+                    "BROKER_REQUEUE_ON_LOST_REQUESTS", False
+                ) and request.request_metadata.get("resubmit", 0) < os.getenv(
+                    "BROKER_REQUEUE_LIMIT", 3
+                ):
+                    logger.info(
+                        "request not found: re-queueing", job_id={request.request_uid}
+                    )
+                    db.requeue_request(request_uid=request.request_uid, session=session)
+                    self.queue.add(request.request_uid, request)
+                    self.qos.notify_end_of_request(
+                        request, session, scheduler=self.internal_scheduler
+                    )
+                else:
+                    db.set_request_status(
+                        request_uid=request.request_uid,
+                        status="failed",
+                        error_message="Request not found in dask scheduler",
+                        error_reason="not_found",
+                        session=session,
+                    )
+                    self.qos.notify_end_of_request(
+                        request, session, scheduler=self.internal_scheduler
+                    )
 
     @perf_logger
     def sync_qos_rules(self, session_write) -> None:
@@ -298,10 +314,6 @@ def sync_qos_rules(self, session_write) -> None:
     def on_future_done(self, future: distributed.Future) -> None:
         job_status = DASK_STATUS_TO_STATUS.get(future.status, "accepted")
         logger_kwargs: dict[str, Any] = {}
-        log = list(self.client.get_events(f"{future.key}/log"))
-        user_visible_log = list(
-            self.client.get_events(f"{future.key}/user_visible_log")
-        )
         with self.session_maker_write() as session:
             if future.status == "finished":
                 result = future.result()
@@ -310,27 +322,53 @@ def on_future_done(self, future: distributed.Future) -> None:
                     job_status,
                     cache_id=result,
                     session=session,
-                    log=log,
-                    user_visible_log=user_visible_log,
                 )
             elif future.status == "error":
                 exception = future.exception()
                 error_message = "".join(traceback.format_exception(exception))
                 error_reason = exception.__class__.__name__
-                if error_reason == "distributed.scheduler.KilledWorker" and os.getenv(
-                    "BROKER_REQUEUE_ON_KILLED_WORKER", False
-                ):
-                    logger.info("worker killed: re-queueing", job_id=future.key)
-                    db.requeue_request(request_uid=future.key, session=session)
-                    self.queue.add(request.request_uid, request)
+                request = db.get_request(future.key, session=session)
+                requeue = os.getenv("BROKER_REQUEUE_ON_KILLED_WORKER_REQUESTS", False)
+                if error_reason == "KilledWorker":
+                    worker_restart_events = self.client.get_events(
+                        "worker-restart-memory"
+                    )
+                    # get info on worker and pid of the killed request
+                    _, worker_pid_event = self.client.get_events(future.key)[0]
+                    if worker_restart_events:
+                        for event in worker_restart_events:
+                            _, job = event
+                            if (
+                                job["worker"] == worker_pid_event["worker"]
+                                and job["pid"] == worker_pid_event["pid"]
+                            ):
+                                db.add_event(
+                                    event_type="killed_worker",
+                                    request_uid=future.key,
+                                    message="Worker has been killed by the Nanny due to memory usage. "
+                                    f"{job['worker']=}, {job['pid']=}, {job['rss']=}",
+                                    session=session,
+                                )
+                                request = db.set_request_status(
+                                    future.key,
+                                    "failed",
+                                    error_message=error_message,
+                                    error_reason=error_reason,
+                                    session=session,
+                                )
+                                requeue = False
+                    if requeue and request.request_metadata.get(
+                        "resubmit", 0
+                    ) < os.getenv("BROKER_REQUEUE_LIMIT", 3):
+                        logger.info("worker killed: re-queueing", job_id=future.key)
+                        db.requeue_request(request_uid=future.key, session=session)
+                        self.queue.add(future.key, request)
                 else:
                     request = db.set_request_status(
                         future.key,
                         job_status,
                         error_message=error_message,
                         error_reason=error_reason,
-                        log=log,
-                        user_visible_log=user_visible_log,
                         session=session,
                     )
             else:
@@ -340,9 +378,8 @@ def on_future_done(self, future: distributed.Future) -> None:
                     job_status,
                     session=session,
                     resubmit=True,
-                    log=log,
-                    user_visible_log=user_visible_log,
                 )
+                self.queue.add(future.key, request)
                 logger.warning(
                     "unknown dask status, re-queing",
                     job_status={future.status},
@@ -367,20 +404,17 @@ def submit_requests(
     ) -> None:
         queue = sorted(
             candidates,
-            key=lambda candidate: self.qos.priority(
-                candidate, session_write
-            ),
+            key=lambda candidate: self.qos.priority(candidate, session_write),
             reverse=True,
         )
         requests_counter = 0
         for request in queue:
             if self.qos.can_run(
                 request, session=session_write, scheduler=self.internal_scheduler
             ):
-                self.submit_request(request, session=session_write)
+                if requests_counter <= int(number_of_requests * WORKERS_MULTIPLIER):
+                    self.submit_request(request, session=session_write)
                 requests_counter += 1
-                if requests_counter == int(number_of_requests * WORKERS_MULTIPLIER):
-                    break
 
     def submit_request(
         self, request: db.SystemRequest, session: sa.orm.Session
@@ -421,9 +455,7 @@ def run(self) -> None:
             with self.session_maker_read() as session_read:
                 if (rules_hash := get_rules_hash(self.qos.path)) != self.qos.rules_hash:
                     logger.info("reloading qos rules")
-                    self.qos.reload_rules(
-                        session=session_read
-                    )
+                    self.qos.reload_rules(session=session_read)
                     self.qos.rules_hash = rules_hash
                 self.qos.environment.set_session(session_read)
                 # expire_on_commit=False is used to detach the accepted requests without an error

cads_broker/entry_points.py

@@ -126,7 +126,7 @@ def delete_requests(
     days: float = 0,
     skip_confirmation: Annotated[bool, typer.Option("--yes", "-y")] = False,
 ) -> None:
-    """Set the status of records in the system_requests table to 'dismissed' if they are in the specified status.
+    """Set the status of records in the system_requests table to 'dismissed'.
 
     Parameters
     ----------
@@ -148,7 +148,8 @@ def delete_requests(
         number_of_requests = session.execute(statement).rowcount
         if not skip_confirmation:
             if not typer.confirm(
-                f"Setting status to 'dismissed' for {number_of_requests} {status} requests. Do you want to continue?",
+                f"Setting status to 'dismissed' for {number_of_requests} {status} requests. "
+                "Do you want to continue?",
                 abort=True,
                 default=True,
             ):