Merge pull request #2223 from codalab/fix-worker-logs

Compute worker - Improve status update and logs

Author: ObadaS

compute_worker/compute_worker.py

@@ -3,6 +3,7 @@
 import hashlib
 import json
 import os
+import traceback
 import shutil
 import signal
 import socket
@@ -18,7 +19,6 @@
 from rich.progress import Progress
 from rich.pretty import pprint
 import requests
-
 import websockets
 import yaml
 from billiard.exceptions import SoftTimeLimitExceeded
@@ -89,22 +89,31 @@
 
 def show_progress(line, progress):
     try:
-        if "Status: Image is up to date" in line["status"]:
-            logger.info(line["status"])
+        status = line.get("status") or ""
+        layer_id = line.get("id")
+        detail = line.get("progressDetail") or {}
+        current = detail.get("current")
+        total = detail.get("total")
+
+        if "Status: Image is up to date" in status:
+            logger.info(status)
+
+        if not layer_id:
+            return
 
         completed = False
-        if line["status"] == "Download complete":
+        if status == "Download complete":
             description = (
-                f"[blue][Download complete, waiting for extraction  {line['id']}]"
+                f"[blue][Download complete, waiting for extraction  {layer_id}]"
             )
             completed = True
-        elif line["status"] == "Downloading":
-            description = f"[bold][Downloading {line['id']}]"
-        elif line["status"] == "Pull complete":
-            description = f"[green][Extraction complete  {line['id']}]"
+        elif status == "Downloading":
+            description = f"[bold][Downloading {layer_id}]"
+        elif status == "Pull complete":
+            description = f"[green][Extraction complete  {layer_id}]"
             completed = True
-        elif line["status"] == "Extracting":
-            description = f"[blue][Extracting  {line['id']}]"
+        elif status == "Extracting":
+            description = f"[blue][Extracting  {layer_id}]"
 
         else:
             # skip other statuses, but show extraction progress
@@ -121,7 +130,7 @@ def show_progress(line, progress):
                 )
             else:
                 tasks[task_id] = progress.add_task(
-                    description, total=line["progressDetail"]["total"]
+                    description, total=total
                 )
         else:
             if completed:
@@ -134,12 +143,12 @@ def show_progress(line, progress):
             else:
                 progress.update(
                     tasks[task_id],
-                    completed=line["progressDetail"]["current"],
-                    total=line["progressDetail"]["total"],
+                    completed=current,
+                    total=total,
                 )
     except Exception as e:
-        logger.error("There was an error showing the progress bar")
-        logger.error(e)
+        if os.environ.get("LOG_LEVEL", "info").lower() == "debug":
+            logger.exception("There was an error showing the progress bar")
 
 
 # -----------------------------------------------
@@ -244,20 +253,44 @@ def run_wrapper(run_args):
     run_args.update(secret=str(run_args["secret"]))
     logger.info(f"Received run arguments: \n {colorize_run_args(json.dumps(run_args))}")
     run = Run(run_args)
-
     try:
         run.prepare()
         run.start()
         if run.is_scoring:
             run.push_scores()
         run.push_output()
     except DockerImagePullException as e:
-        run._update_status(STATUS_FAILED, str(e))
-    except SubmissionException as e:
-        run._update_status(STATUS_FAILED, str(e))
+        msg = str(e).strip()
+        if msg:
+            msg = f"Docker image pull failed: {msg}"
+        else:
+            msg = "Docker image pull failed."
+        run._update_status(STATUS_FAILED, extra_information=msg)
+        raise
     except SoftTimeLimitExceeded:
-        run._update_status(STATUS_FAILED, "Soft time limit exceeded!")
+        run._update_status(
+            STATUS_FAILED,
+            extra_information="Execution time limit exceeded.",
+        )
+        raise
+    except SubmissionException as e:
+        msg = str(e).strip()
+        if msg:
+            msg = f"Submission failed: {msg}. See logs for more details."
+        else:
+            msg = "Submission failed. See logs for more details."
+        run._update_status(STATUS_FAILED, extra_information=msg)
+        raise
+    except Exception as e:
+        # Catch any exception to avoid getting stuck in Running status
+        run._update_status(STATUS_FAILED, extra_information=traceback.format_exc())
+        raise
     finally:
+        try:
+            # Try to push logs before cleanup
+            run.push_logs()
+        except Exception:
+            logger.exception("push_logs failed")
         run.clean_up()
 
 
@@ -446,6 +479,22 @@ async def watch_detailed_results(self):
             if file_path:
                 await self.send_detailed_results(file_path)
 
+    def push_logs(self):
+        """Upload any collected logs, even in case of crash.
+        """
+        try:
+            for kind, logs in (self.logs or {}).items():
+                for stream_key in ("stdout", "stderr"):
+                    entry = logs.get(stream_key) if isinstance(logs, dict) else None
+                    if not entry:
+                        continue
+                    location = entry.get("location")
+                    data = entry.get("data") or b""
+                    if location:
+                        self._put_file(location, raw_data=data)
+        except Exception as e:
+            logger.exception(f"Failed best-effort log upload: {e}")
+
     def get_detailed_results_file_path(self):
         default_detailed_results_path = os.path.join(
             self.output_dir, "detailed_results.html"
@@ -467,7 +516,7 @@ async def send_detailed_results(self, file_path):
         )
         websocket_url = f"{self.websocket_url}?kind=detailed_results"
         logger.info(f"Connecting to {websocket_url} for detailed results")
-        # Wrap this with a Try ... Except otherwise a failure here will make the submission get stuck on Running
+        # Wrap this with a Try block to avoid getting stuck on Running
         try:
             websocket = await asyncio.wait_for(
                 websockets.connect(websocket_url), timeout=30.0
@@ -480,14 +529,8 @@ async def send_detailed_results(self, file_path):
                 )
             )
         except Exception as e:
-            logger.error(
-                f"This error might result in a Execution Time Exceeded error: {e}"
-            )
-            if os.environ.get("LOG_LEVEL", "info").lower() == "debug":
-                logger.exception(e)
-            raise SubmissionException(
-                "Could not connect to instance to update detailed result"
-            )
+            logger.exception(e)
+            return
 
     def _get_stdout_stderr_file_names(self, run_args):
         # run_args should be the run_args argument passed to __init__ from the run_wrapper.
@@ -513,7 +556,7 @@ def _update_submission(self, data):
 
         logger.info(f"Updating submission @ {url} with data = {data}")
 
-        resp = self.requests_session.patch(url, data, timeout=150)
+        resp = self.requests_session.patch(url, data=data, timeout=150)
         if resp.status_code == 200:
             logger.info("Submission updated successfully!")
         else:
@@ -523,23 +566,17 @@ def _update_submission(self, data):
             raise SubmissionException("Failure updating submission data.")
 
     def _update_status(self, status, extra_information=None):
+        # Update submission status
         if status not in AVAILABLE_STATUSES:
             raise SubmissionException(
                 f"Status '{status}' is not in available statuses: {AVAILABLE_STATUSES}"
             )
-
-        data = {
-            "status": status,
-            "status_details": extra_information,
-        }
-
-        # TODO: figure out if we should pull this task code later(submission.task should always be set)
-        # When we start
-        # if status == STATUS_SCORING:
-        #     data.update({
-        #         "task_pk": self.task_pk,
-        #     })
-        self._update_submission(data)
+        data = {"status": status, "status_details": extra_information}
+        try:
+            self._update_submission(data)
+        except Exception as e:
+            # Always catch exception and never raise error
+            logger.exception(f"Failed to update submission status to {status}: {e}")
 
     def _get_container_image(self, image_name):
         logger.info("Running pull for image: {}".format(image_name))
@@ -549,6 +586,8 @@ def _get_container_image(self, image_name):
                 with Progress() as progress:
                     resp = client.pull(image_name, stream=True, decode=True)
                     for line in resp:
+                        if isinstance(line, dict) and line.get("error"):
+                            raise DockerImagePullException(line["error"])
                         show_progress(line, progress)
                     break  # Break if the loop is successful to exit "with Progress() as progress"
 
@@ -684,8 +723,13 @@ async def _run_container_engine_cmd(self, container, kind):
         # Creating this and setting 2 values to None in case there is not enough time for the worker to get logs, otherwise we will have errors later on
         logs_Unified = [None, None]
 
+        # To store on-going logs and avoid empty logs returning to the platform
+        stdout_chunks = []
+        stderr_chunks = []
+
         # Create a websocket to send the logs in real time to the codabench instance
         # We need to set a timeout for the websocket connection otherwise the program will get stuck if he websocket does not connect.
+        websocket = None
         try:
             websocket_url = f"{self.websocket_url}?kind={kind}"
             logger.debug(
@@ -732,21 +776,27 @@ async def _run_container_engine_cmd(self, container, kind):
                     "Show the logs and stream them to codabench " + container.get("Id")
                 )
                 for log in container_LogsDemux:
-                    if str(log[0]) != "None":
+                    # Output
+                    if log[0] is not None:
+                        stdout_chunks.append(log[0])
                         logger.info(log[0].decode())
                         try:
-                            await websocket.send(
-                                json.dumps({"kind": kind, "message": log[0].decode()})
-                            )
+                            if websocket is not None:
+                                await websocket.send(
+                                    json.dumps({"kind": kind, "message": log[0].decode()})
+                                )
                         except Exception as e:
                             logger.error(e)
-
-                    elif str(log[1]) != "None":
+                    
+                    # Errors
+                    elif log[1] is not None:
+                        stderr_chunks.append(log[1])
                         logger.error(log[1].decode())
                         try:
-                            await websocket.send(
-                                json.dumps({"kind": kind, "message": log[1].decode()})
-                            )
+                            if websocket is not None:
+                                await websocket.send(
+                                    json.dumps({"kind": kind, "message": log[1].decode()})
+                                )
                         except Exception as e:
                             logger.error(e)
 
@@ -762,12 +812,17 @@ async def _run_container_engine_cmd(self, container, kind):
         # Get the return code of the competition container once done
         try:
             # Gets the logs of the container, sperating stdout and stderr (first and second position) thanks for demux=True
-            logs_Unified = client.attach(container, logs=True, demux=True)
             return_Code = client.wait(container)
+            logs_Unified = (b"".join(stdout_chunks), b"".join(stderr_chunks))
             logger.debug(
                 f"WORKER_MARKER: Disconnecting from {websocket_url}, program counter = {self.completed_program_counter}"
             )
-            await websocket.close()
+            if websocket is not None:
+                try:
+                    await websocket.close()
+                    await websocket.wait_closed()
+                except Exception as e:
+                    logger.error(e)
             client.remove_container(container, force=True)
 
             logger.debug(
@@ -785,6 +840,13 @@ async def _run_container_engine_cmd(self, container, kind):
             logger.error(e)
             return_Code = {"StatusCode": 1}
 
+        finally:
+            try:
+                # Last chance of removing container
+                client.remove_container(container_id, force=True)
+            except Exception:
+                pass
+
         self.logs[kind] = {
             "returncode": return_Code["StatusCode"],
             "start": start,
@@ -1055,9 +1117,8 @@ async def _run_program_directory(self, program_dir, kind):
         try:
             return await self._run_container_engine_cmd(container, kind=kind)
         except Exception as e:
-            logger.error(e)
-            if os.environ.get("LOG_LEVEL", "info").lower() == "debug":
-                logger.exception(e)
+            logger.exception("Program directory execution failed")
+            raise SubmissionException(str(e))
 
     def _put_dir(self, url, directory):
         """Zip the directory and send it to the given URL using _put_file."""
@@ -1099,7 +1160,7 @@ def _put_file(self, url, file=None, raw_data=None, content_type="application/zip
             logger.info("Putting file %s in %s" % (file, url))
             data = open(file, "rb")
             headers["Content-Length"] = str(os.path.getsize(file))
-        elif raw_data:
+        elif raw_data is not None:
             logger.info("Putting raw data %s in %s" % (raw_data, url))
             data = raw_data
         else:
@@ -1196,10 +1257,12 @@ def start(self):
         task_results = []  # will store results/exceptions from gather
         signal.signal(signal.SIGALRM, alarm_handler)
         signal.alarm(self.execution_time_limit)
+
         try:
             # run tasks
             # keep what gather returned so we can detect async errors later
             task_results = loop.run_until_complete(gathered_tasks) or []
+
         except ExecutionTimeLimitExceeded:
             error_message = f"Execution Time Limit exceeded. Limit was {self.execution_time_limit} seconds"
             logger.error(error_message)
@@ -1233,8 +1296,25 @@ def start(self):
             # Send error through web socket to the frontend
             asyncio.run(self._send_data_through_socket(error_message))
             raise SubmissionException(error_message)
+
         finally:
+            signal.alarm(0)
             self.watch = False
+
+            # Cancel any remaining pending tasks before closing the loop
+            pending = [t for t in asyncio.all_tasks(loop) if not t.done()]
+            for task in pending:
+                task.cancel()
+            if pending:
+                try:
+                    loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
+                except Exception:
+                    pass
+
+            # Close loop
+            asyncio.set_event_loop(None)
+            loop.close()
+
             for kind, logs in self.logs.items():
                 if logs["end"] is not None:
                     elapsed_time = logs["end"] - logs["start"]
@@ -1292,7 +1372,7 @@ def start(self):
                 program_results, BaseException
             ) and not isinstance(program_results, asyncio.CancelledError)
             program_rc = getattr(self, "program_exit_code", None)
-            failed_rc = program_rc not in (0, None)
+            failed_rc = (program_rc is None) or (program_rc != 0)
             if had_async_exc or failed_rc:
                 self._update_status(
                     STATUS_FAILED,
@@ -1301,6 +1381,7 @@ def start(self):
                 # Raise so upstream marks failed immediately
                 raise SubmissionException("Child task failed or non-zero return code")
             self._update_status(STATUS_FINISHED)
+
         else:
             self._update_status(STATUS_SCORING)