Option to stop compute worker from sending logs to Codabench

compute_worker/compute_worker.py

@@ -109,10 +109,12 @@ def to_bool(val):
     COMPETITION_CONTAINER_HTTP_PROXY = get("COMPETITION_CONTAINER_HTTP_PROXY", "")
     COMPETITION_CONTAINER_HTTPS_PROXY = get("COMPETITION_CONTAINER_HTTPS_PROXY", "")
 
-    CODALAB_IGNORE_CLEANUP_STEP = to_bool(get("CODALAB_IGNORE_CLEANUP_STEP"))
+    COMPUTE_WORKER_NO_CLEANUP = to_bool(get("COMPUTE_WORKER_NO_CLEANUP", "False"))
 
     WORKER_BUNDLE_URL_REWRITE = get("WORKER_BUNDLE_URL_REWRITE", "").strip()
 
+    SILENT_COMPUTE_WORKER = to_bool(get("SILENT_COMPUTE_WORKER", "False"))
+
 
 # -----------------------------------------------
 # Program Kind
@@ -331,9 +333,15 @@ def run_wrapper(run_args):
     except SubmissionException as e:
         msg = str(e).strip()
         if msg:
-            msg = f"Submission failed: {msg}. See logs for more details."
+            if Settings.SILENT_COMPUTE_WORKER:
+                msg = f"Submission failed: {msg}. Contact the Organizer for more details."
+            else:
+                msg = f"Submission failed: {msg}. See logs for more details."
         else:
-            msg = "Submission failed. See logs for more details."
+            if Settings.SILENT_COMPUTE_WORKER:
+                msg = "Submission failed. Contact the Organizer for more details."
+            else:
+                msg = "Submission failed. See logs for more details."
         run._update_status(SubmissionStatus.FAILED, extra_information=msg)
         raise
     except Exception as e:
@@ -546,18 +554,36 @@ async def watch_detailed_results(self):
     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}")
+        if Settings.SILENT_COMPUTE_WORKER == False:
+            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}")
+
+        else:
+            try:
+                logs_path = os.path.join(self.root_dir, "logs")
+                with open(logs_path, "w") as f:
+                    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:
+                                f.write(str(data))
+            except Exception as e:
+                logger.exception(f"Failed best-effort log file creation: {e}")
+
 
     def get_detailed_results_file_path(self):
         default_detailed_results_path = os.path.join(
@@ -671,9 +697,14 @@ def _get_container_image(self, image_name):
                     self._update_submission(docker_pull_fail_data)
                     # Send error through web socket to the frontend
                     asyncio.run(self._send_data_through_socket(str(pull_error)))
-                    raise DockerImagePullException(
-                        f"Pull for {image_name} failed! Check the logs for more information"
-                    )
+                    if Settings.SILENT_COMPUTE_WORKER:
+                        raise DockerImagePullException(
+                            f"Pull for {image_name} failed! Contact the Organizer for more details."
+                        )
+                    else:
+                        raise DockerImagePullException(
+                            f"Pull for {image_name} failed! Check the logs for more information"
+                        )
                 else:
                     logger.warning("Failed. Retrying in 5 seconds...")
                     time.sleep(5)  # Wait 5 seconds before retrying
@@ -886,21 +917,24 @@ async def _run_container_engine_cmd(self, container, kind):
         # 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(f"Connecting to {websocket_url} for container {str(container.get('Id'))}")
-            websocket = await asyncio.wait_for(
-                websockets.connect(websocket_url), timeout=10.0
-            )
-            logger.debug(f"connected to {websocket_url} for container {str(container.get('Id'))}")
 
-        except Exception as e:
-            logger.error(
-                f"There was an error trying to connect to the websocket on the codabench instance: {e}"
-            )
+        # Do not create a websocket if the real time logs are not wanted (Silent Compute Worker)
+        if Settings.SILENT_COMPUTE_WORKER == False:
+            try:
+                websocket_url = f"{self.websocket_url}?kind={kind}"
+                logger.debug(f"Connecting to {websocket_url} for container {str(container.get('Id'))}")
+                websocket = await asyncio.wait_for(
+                    websockets.connect(websocket_url), timeout=10.0
+                )
+                logger.debug(f"connected to {websocket_url} for container {str(container.get('Id'))}")
 
-            if Settings.LOG_LEVEL == Settings.LOG_LEVEL_DEBUG:
-                logger.exception(e)
+            except Exception as e:
+                logger.error(
+                    f"There was an error trying to connect to the websocket on the codabench instance: {e}"
+                )
+
+                if Settings.LOG_LEVEL == Settings.LOG_LEVEL_DEBUG:
+                    logger.exception(e)
 
         start = time.time()
 
@@ -916,6 +950,7 @@ async def _run_container_engine_cmd(self, container, kind):
             )
 
             # If we enter the for loop after the container exited, the program will get stuck
+            # Do not send the real time logs if they are not wanted (Silent Compute Worker)
             if client.inspect_container(container)["State"]["Status"].lower() == "running":
                 logger.debug(
                     "Show the logs and stream them to codabench " + container.get("Id")
@@ -925,25 +960,27 @@ async def _run_container_engine_cmd(self, container, kind):
                     if log[0] is not None:
                         stdout_chunks.append(log[0])
                         logger.info(log[0].decode())
-                        try:
-                            if websocket is not None:
-                                await websocket.send(
-                                    json.dumps({"kind": kind, "message": log[0].decode()})
-                                )
-                        except Exception as e:
-                            logger.error(e)
+                        if Settings.SILENT_COMPUTE_WORKER == False:
+                            try:
+                                if websocket is not None:
+                                    await websocket.send(
+                                        json.dumps({"kind": kind, "message": log[0].decode()})
+                                    )
+                            except Exception as e:
+                                logger.error(e)
 
                     # Errors
                     elif log[1] is not None:
                         stderr_chunks.append(log[1])
                         logger.error(log[1].decode())
-                        try:
-                            if websocket is not None:
-                                await websocket.send(
-                                    json.dumps({"kind": kind, "message": log[1].decode()})
-                                )
-                        except Exception as e:
-                            logger.error(e)
+                        if Settings.SILENT_COMPUTE_WORKER == False:
+                            try:
+                                if websocket is not None:
+                                    await websocket.send(
+                                        json.dumps({"kind": kind, "message": log[1].decode()})
+                                    )
+                            except Exception as e:
+                                logger.error(e)
 
         except (docker.errors.NotFound, docker.errors.APIError) as e:
             logger.error(e)
@@ -959,15 +996,17 @@ async def _run_container_engine_cmd(self, container, kind):
             # Gets the logs of the container, sperating stdout and stderr (first and second position) thanks for 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}"
-            )
-            if websocket is not None:
-                try:
-                    await websocket.close()
-                    await websocket.wait_closed()
-                except Exception as e:
-                    logger.error(e)
+
+            if Settings.SILENT_COMPUTE_WORKER == False:
+                logger.debug(
+                    f"WORKER_MARKER: Disconnecting from {websocket_url}, program counter = {self.completed_program_counter}"
+                )
+                if websocket is not None:
+                    try:
+                        await websocket.close()
+                        await websocket.wait_closed()
+                    except Exception as e:
+                        logger.error(e)
             client.remove_container(container, v=True, force=True)
 
             logger.debug(f"Container {container.get('Id')} exited with status code : {str(return_Code['StatusCode'])}")
@@ -1317,6 +1356,7 @@ def start(self):
                 self._run_program_directory(kind=ProgramKind.INGESTION_PROGRAM, program_dir=ingestion_program_dir),
             ])
 
+        logger.info(tasks)
         gathered_tasks = asyncio.gather(*tasks, return_exceptions=True)
 
         task_results = []  # will store results/exceptions from gather
@@ -1414,12 +1454,14 @@ def start(self):
                     self.ingestion_program_exit_code = return_code
                     self.ingestion_program_elapsed_time = elapsed_time
                 logger.info(f"[exited with {logs['returncode']}]")
-                for key, value in logs.items():
-                    if key not in ["stdout", "stderr"]:
-                        continue
-                    if value["data"]:
-                        logger.info(f"[{key}]\n{value['data']}")
-                        self._put_file(value["location"], raw_data=value["data"])
+                if Settings.SILENT_COMPUTE_WORKER == False:
+                    for key, value in logs.items():
+                        if key not in ["stdout", "stderr"]:
+                            continue
+                        if value["data"]:
+                            logger.info(f"[{key}]\n{value['data']}")
+                            self._put_file(value["location"], raw_data=value["data"])
+
 
                 # set logs of this kind to None, since we handled them already
                 logger.info("Program finished")
@@ -1522,9 +1564,9 @@ def push_output(self):
             self._put_dir(self.scoring_result, self.output_dir)
 
     def clean_up(self):
-        if Settings.CODALAB_IGNORE_CLEANUP_STEP:
+        if Settings.COMPUTE_WORKER_NO_CLEANUP:
             logger.warning(
-                f"CODALAB_IGNORE_CLEANUP_STEP mode enabled, ignoring clean up of: {self.root_dir}"
+                f"COMPUTE_WORKER_NO_CLEANUP mode enabled, ignoring clean up of: {self.root_dir}"
             )
             return
 

docker-compose.yml

@@ -244,7 +244,7 @@ services:
     environment:
       - BROKER_URL=pyamqp://${RABBITMQ_DEFAULT_USER}:${RABBITMQ_DEFAULT_PASS}@${RABBITMQ_HOST}:${RABBITMQ_PORT}//
       # Make the worker leave behind the submission so we can examine it
-      - CODALAB_IGNORE_CLEANUP_STEP=1
+      - COMPUTE_WORKER_NO_CLEANUP=True
     tty: true
     logging:
       options:

documentation/docs/Organizers/Running_a_benchmark/Compute-Worker-Management---Setup.md

@@ -59,6 +59,12 @@ HOST_DIRECTORY=/codabench
 CONTAINER_ENGINE_EXECUTABLE=docker
 #USE_GPU=True
 #GPU_DEVICE=nvidia.com/gpu=all
+# This option removes the ability of the compute worker to send logs to
+# codabench, instead writing them locally on disk. Combine with
+# COMPUTE_WORKER_NO_CLEANUP=true to stop the worker's cleanup to keep
+# all the logs locally only
+#SILENT_COMPUTE_WORKER=False
+#COMPUTE_WORKER_NO_CLEANUP=False
 
 #######################################################################
 #                       Network                                       #

src/settings/logs_loguru.py

@@ -130,7 +130,17 @@ def colorize_run_args(json_str):
         json_str,
     )
     json_str = re.sub(
-        r'("ingestion_program": ")(.*?)(",)',
+        r'("ingestion_program_data": ")(.*?)(",)',
+        rf"\1{yellow}\2{reset}\3{lineskip}",
+        json_str,
+    )
+    json_str = re.sub(
+        r'("submission_data": ")(.*?)(",)',
+        rf"\1{yellow}\2{reset}\3{lineskip}",
+        json_str,
+    )
+    json_str = re.sub(
+        r'("scoring_program_data": ")(.*?)(",)',
         rf"\1{yellow}\2{reset}\3{lineskip}",
         json_str,
     )