compute_worker.py

import asyncio
import glob
import hashlib
import json
import os
import traceback
import shutil
import signal
import socket
import tempfile
import time
import uuid
import requests
import websockets
import yaml
import docker
import logging
import sys  # This is only needed for the pytests to pass
from shutil import make_archive
from urllib.error import HTTPError
from urllib.parse import urlparse
from urllib.request import urlretrieve
from zipfile import ZipFile, BadZipFile
from urllib3 import Retry

from rich.pretty import pprint
from rich.progress import Progress
from kombu import Queue, Exchange
from celery import Celery, shared_task, utils, signals
from billiard.exceptions import SoftTimeLimitExceeded

from logs_loguru import configure_logging, colorize_run_args
from docker_image_update_checker import DockerImageStatus, DockerImageUpdateChecker

logger = logging.getLogger(__name__)

sys.path.append("/app/src/settings/")


# -----------------------------------------------
# Settings
# -----------------------------------------------
class Settings:

    @staticmethod
    def get(key, default=None):
        """
        Return the env var value if set, else default; returns None if not set and no default.
        """
        val = os.getenv(key)

        if val is not None:
            return val

        if default is not None:
            return default

        logger.warning(f"Environment variable '{key}' not found and no default provided.")
        return None

    @staticmethod
    def to_bool(val):
        try:
            if isinstance(val, bool):
                return val

            val_str = str(val).strip()

            if val_str in ("true", "True", "TRUE", "1"):
                return True
            if val_str in ("false", "False", "FALSE", "0"):
                return False

            logger.warning(f"Failed to parse boolean from '{val}'")
            return val

        except Exception as e:
            logger.warning(f"Failed to parse boolean from '{val}': {e}")
            return val

    # Directories
    # NOTE: we need to pass this directory to docker/podman so it knows where to store things!
    HOST_DIRECTORY = get("HOST_DIRECTORY", "/tmp/codabench/")
    MAX_CACHE_DIR_SIZE_GB = float(get("MAX_CACHE_DIR_SIZE_GB", 10))
    BASE_DIR = "/codabench/"  # base directory inside the container
    CACHE_DIR = os.path.join(BASE_DIR, "cache")

    # Constants
    DOCKER = "docker"
    PODMAN = "podman"
    LOG_LEVEL_DEBUG = "debug"

    # Defaults
    DEFAULT_SOCKETS = {
        DOCKER: "unix:///var/run/docker.sock",
        PODMAN: "unix:///run/user/1000/podman/podman.sock",
    }

    # env variables
    LOG_LEVEL = get("LOG_LEVEL", "INFO").lower()
    SERIALIZED = get("SERIALIZED", "false")

    USE_GPU = to_bool(get("USE_GPU", "false"))
    CONTAINER_ENGINE_EXECUTABLE = get("CONTAINER_ENGINE_EXECUTABLE", DOCKER).lower()
    GPU_DEVICE = get("GPU_DEVICE", "nvidia.com/gpu=all")

    CONTAINER_SOCKET = get("CONTAINER_SOCKET", DEFAULT_SOCKETS.get(CONTAINER_ENGINE_EXECUTABLE))

    COMPETITION_CONTAINER_NETWORK_DISABLED = to_bool(get("COMPETITION_CONTAINER_NETWORK_DISABLED", "False"))
    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"))

    WORKER_BUNDLE_URL_REWRITE = get("WORKER_BUNDLE_URL_REWRITE", "").strip()

    # Docker image config
    DOCKER_IMAGE_NAMESPACE = get("DOCKER_IMAGE_NAMESPACE", "codalab")
    DOCKER_IMAGE_REPOSITORY = get("DOCKER_IMAGE_REPOSITORY", "codabench-compute-worker")
    DOCKER_IMAGE_TAG = get("DOCKER_IMAGE_TAG", "latest")


# -----------------------------------------------
# Program Kind
# -----------------------------------------------
class ProgramKind:
    INGESTION_PROGRAM = "ingestion_program"
    SCORING_PROGRAM = "scoring_program"


# -----------------------------------------------
# Submission status
# -----------------------------------------------
class SubmissionStatus:
    NONE = "None"
    SUBMITTING = "Submitting"
    SUBMITTED = "Submitted"
    PREPARING = "Preparing"
    RUNNING = "Running"
    SCORING = "Scoring"
    FINISHED = "Finished"
    FAILED = "Failed"

    AVAILABLE_STATUSES = (
        NONE,
        SUBMITTING,
        SUBMITTED,
        PREPARING,
        RUNNING,
        SCORING,
        FINISHED,
        FAILED,
    )


# -----------------------------------------------
# Logging
# -----------------------------------------------
configure_logging(
    Settings.LOG_LEVEL, Settings.SERIALIZED
)

# -----------------------------------------------
# Initialize Docker or Podman depending on .env
# -----------------------------------------------
logger.info(
    f"Using {Settings.CONTAINER_ENGINE_EXECUTABLE} "
    f"{'with GPU capabilities: ' + Settings.GPU_DEVICE if Settings.USE_GPU else 'without GPU capabilities'}. "
    f"Network disabled for the competition container is set to {Settings.COMPETITION_CONTAINER_NETWORK_DISABLED}"
)

# Intializing client
# NOTE: CONTAINER_SOCKET is set in Settings based on CONTAINER_ENGINE_EXECUTABLE which must has either podman or docker
client = docker.APIClient(
    base_url=Settings.CONTAINER_SOCKET,
    version="auto",
)


# -----------------------------------------------
# Show Progress bar on downloading images
# -----------------------------------------------
tasks = {}


def show_progress(line, progress):
    try:
        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 status == "Download complete":
            description = (
                f"[blue][Download complete, waiting for extraction  {layer_id}]"
            )
            completed = True
        elif status == "Downloading":
            description = f"[bold][Downloading {layer_id}]"
        elif status == "Pull complete":
            description = f"[green][Extraction complete  {layer_id}]"
            completed = True
        elif status == "Extracting":
            description = f"[blue][Extracting  {layer_id}]"

        else:
            # skip other statuses, but show extraction progress
            return

        task_id = line["id"]
        if task_id not in tasks.keys():
            if completed:
                # some layers are really small that they download immediately without showing
                # anything as Downloading in the stream.
                # For that case, show a completed progress bar
                tasks[task_id] = progress.add_task(
                    description, total=100, completed=100
                )
            else:
                tasks[task_id] = progress.add_task(
                    description, total=total
                )
        else:
            if completed:
                # due to the stream, the Download complete output can happen before the Downloading
                # bar outputs the 100%. So when we detect that the download is in fact complete,
                # update the progress bar to show 100%
                progress.update(
                    tasks[task_id], description=description, total=100, completed=100
                )
            else:
                progress.update(
                    tasks[task_id],
                    completed=current,
                    total=total,
                )
    except Exception as e:
        if Settings.LOG_LEVEL == Settings.LOG_LEVEL_DEBUG:
            logger.exception(f"There was an error showing the progress bar: {e}")


# -----------------------------------------------
# Celery + Rabbit MQ
# -----------------------------------------------
@signals.setup_logging.connect
def setup_celery_logging(**kwargs):
    pass


# Init celery + rabbit queue definitions
app = Celery()
app.config_from_object("celery_config")  # grabs celery_config.py
app.conf.task_queues = [
    # Mostly defining queue here so we can set x-max-priority
    Queue(
        "compute-worker",
        Exchange("compute-worker"),
        routing_key="compute-worker",
        queue_arguments={"x-max-priority": 10},
    ),
]


# -----------------------------------------------
# Exceptions
# -----------------------------------------------
class SubmissionException(Exception):
    pass


class DockerImagePullException(Exception):
    pass


class ExecutionTimeLimitExceeded(Exception):
    pass


# -----------------------------------------------
# Local setups where public S3 URLs are not directly reachable from workers
# -----------------------------------------------
def rewrite_bundle_url_if_needed(url):
    """
    Optionally rewrite presigned bundle URLs for worker networking.

    Controlled by env: WORKER_BUNDLE_URL_REWRITE=FROM|TO

    Example: http://localhost:9000|http://minio:9000
    """
    rule = Settings.WORKER_BUNDLE_URL_REWRITE
    if not rule or "|" not in rule:
        return url
    src, dst = rule.split("|", 1)
    if url.startswith(src):
        new_url = dst + url[len(src):]
        logger.info(f"Rewriting bundle URL for worker: {url} -> {new_url}")
        return new_url
    return url


def check_docker_image_update():
    """
    Compare local and remote compute worker Docker images and log the
    synchronization status along with relevant image metadata.
    """
    checker = DockerImageUpdateChecker(
        namespace=Settings.DOCKER_IMAGE_NAMESPACE,
        repository=Settings.DOCKER_IMAGE_REPOSITORY,
        tag=Settings.DOCKER_IMAGE_TAG,
        docker_base_url=Settings.CONTAINER_SOCKET
    )
    result = checker.compare_local_vs_remote_images(container_id=socket.gethostname())
    status = result["status"]

    log_level = logging.INFO

    log_lines = [
        "",
        "=" * 60,
        "COMPUTE WORKER DOCKER IMAGE UPDATE CHECK",
        "=" * 60,
        f"Image: {result.get('image_name')}",
    ]

    remote = result.get("remote")
    local = result.get("local")

    if remote:
        log_lines.append(f"Remote: digest={remote.get('digest')}, date={remote.get('date')}")

    if local:
        log_lines.append(f"Local: id={local.get('id')}, date={local.get('date')}")

    log_lines.append("-" * 60)

    if status == DockerImageStatus.UP_TO_DATE:
        log_lines.append("Status: Local image is synchronized with remote")
        log_level = logging.INFO

    elif status == DockerImageStatus.BEHIND:
        log_lines.append("Status: Local image is behind remote version. For better submission processing and to avoid any submission errors, fetch the latest image!")
        log_level = logging.ERROR

    elif status == DockerImageStatus.LOCAL_MISSING:
        log_lines.append("Status: Local image not found. Pull required")
        log_level = logging.ERROR

    elif status == DockerImageStatus.REMOTE_UNAVAILABLE:
        log_lines.append("Status: Could not fetch remote image metadata")
        log_level = logging.ERROR

    elif status == "error":
        log_lines.append(f"Status: Image check failed: {result.get('error')}")
        log_level = logging.ERROR
    else:
        log_lines.append(f"Unknown image status: {status}")
        log_level = logging.ERROR

    log_lines.append("=" * 60)

    logger.log(log_level, "\n".join(log_lines))

    if status != DockerImageStatus.UP_TO_DATE:
        return "\n".join(log_lines) + "\n"
    return None


# -----------------------------------------------------------------------------
# The main compute worker entrypoint, this is how a job is ran at the highest
# level.
# -----------------------------------------------------------------------------
@shared_task(name="compute_worker_run")
def run_wrapper(run_args):
    # Check for docker image update
    docker_image_warning = check_docker_image_update()

    # We need to convert the UUID given by celery into a byte like object otherwise things will break
    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)
    if docker_image_warning:
        run.docker_image_warning = docker_image_warning.encode()
    try:
        run.prepare()
        run.start()
        if run.is_scoring:
            run.push_scores()
        run.push_output()
    except DockerImagePullException as e:
        msg = str(e).strip()
        if msg:
            msg = f"Docker image pull failed: {msg}"
        else:
            msg = "Docker image pull failed."
        run._update_status(SubmissionStatus.FAILED, extra_information=msg)
        raise
    except SoftTimeLimitExceeded:
        run._update_status(
            SubmissionStatus.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(SubmissionStatus.FAILED, extra_information=msg)
        raise
    except Exception:
        # Catch any exception to avoid getting stuck in Running status
        run._update_status(SubmissionStatus.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()


def replace_legacy_metadata_command(
    command, kind, is_scoring, ingestion_only_during_scoring=False
):
    vars_to_replace = [
        ("$input", "/app/input_data" if kind == ProgramKind.INGESTION_PROGRAM else "/app/input"),
        ("$output", "/app/output"),
        (
            "$program",
            "/app/ingestion_program"
            if ingestion_only_during_scoring and is_scoring
            else "/app/program",
        ),
        ("$ingestion_program", "/app/program"),
        ("$hidden", "/app/input/ref"),
        ("$shared", "/app/shared"),
        ("$submission_program", "/app/ingested_program"),
        # for v1.8 compatibility
        ("$tmp", "/app/output"),
        ("$predictions", "/app/input/res" if is_scoring else "/app/output"),
    ]
    for var_string, var_replacement in vars_to_replace:
        command = command.replace(var_string, var_replacement)
    return command


def md5(filename):
    """Given some file return its md5, works well on large files"""
    hash_md5 = hashlib.md5()
    with open(filename, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()


def get_folder_size_in_gb(folder):
    # Check if the folder exists; if not, return 0 GB
    if not os.path.exists(folder):
        return 0

    total_size = 0  # Initialize total size accumulator (in bytes)

    # Walk through the folder and all its subdirectories
    for root, dirs, files in os.walk(folder):
        for f in files:
            # Construct full path to the file
            fp = os.path.join(root, f)
            # Add the file size to total_size
            total_size += os.path.getsize(fp)

    # Convert bytes to gigabytes using decimal system (1 GB = 1000^3 bytes)
    return total_size / (1000 ** 3)


def delete_files_in_folder(folder):
    if not os.path.isdir(folder):
        return
    for filename in os.listdir(folder):
        file_path = os.path.join(folder, filename)
        if os.path.isfile(file_path) or os.path.islink(file_path):
            os.unlink(file_path)
        elif os.path.isdir(file_path):
            shutil.rmtree(file_path)


def is_valid_zip(zip_path):
    # Check zip integrity
    try:
        with ZipFile(zip_path, "r") as zf:
            return zf.testzip() is None
    except BadZipFile:
        return False


def alarm_handler(signum, frame):
    raise ExecutionTimeLimitExceeded


# -----------------------------------------------
# Class Run
# Responsible for running a submission inside a docker/podman container
# -----------------------------------------------
class Run:
    """A "Run" in Codabench is composed of some program, some data to work with, and some signed URLs to upload results
    to. There is also a secret key to do special commands for just this submission.

    Some example API's you can hit using this secret key are:

        push_scores

        (maybe later:
            get previous submission
            get sibling submission
            get top submission
            get some different dataset
            post results to twitter)
    """

    def __init__(self, run_args):
        self.run_related_name = (
            f"uPK-{run_args['user_pk']}_sID-{run_args['id']}"
        )
        # Directories for the run
        self.watch = True
        self.completed_program_counter = 0
        self.root_dir = tempfile.mkdtemp(prefix=f'{self.run_related_name}__', dir=Settings.BASE_DIR)
        self.bundle_dir = os.path.join(self.root_dir, "bundles")
        self.input_dir = os.path.join(self.root_dir, "input")
        self.output_dir = os.path.join(self.root_dir, "output")
        self.data_dir = os.path.join(Settings.HOST_DIRECTORY, "data")  # absolute path to data in the host
        self.logs = {}
        self.docker_image_warning = None

        # Details for submission
        self.is_scoring = run_args["is_scoring"]
        self.user_pk = run_args["user_pk"]
        self.submission_id = run_args["id"]
        self.submissions_api_url = run_args["submissions_api_url"]
        self.container_image = run_args["docker_image"]
        self.secret = run_args["secret"]
        self.prediction_result = run_args["prediction_result"]
        self.scoring_result = run_args.get("scoring_result")
        self.execution_time_limit = run_args["execution_time_limit"]
        # stdout and stderr
        self.stdout, self.stderr, self.ingestion_stdout, self.ingestion_stderr = (
            self._get_stdout_stderr_file_names(run_args)
        )
        # Setting up container names for ingestion, scoring and submission
        self.ingestion_program_container_name = f"ingestion_{self.run_related_name}"
        self.scoring_program_container_name = f"scoring_{self.run_related_name}"

        # Setting up ingestion, scoring and submission data
        self.ingestion_program_data = run_args.get("ingestion_program_data")
        self.scoring_program_data = run_args.get("scoring_program_data")
        self.submission_data = run_args.get("submission_data")

        self.input_data = run_args.get("input_data")
        self.reference_data = run_args.get("reference_data")
        self.ingestion_only_during_scoring = run_args.get("ingestion_only_during_scoring")
        self.detailed_results_url = run_args.get("detailed_results_url")

        self.ingestion_program_exit_code = None
        self.ingestion_program_elapsed_time = None
        self.scoring_program_exit_code = None
        self.scoring_program_elapsed_time = None

        # Socket connection to stream output of submission
        submission_api_url_parsed = urlparse(self.submissions_api_url)
        websocket_host = submission_api_url_parsed.netloc
        websocket_scheme = "ws" if submission_api_url_parsed.scheme == "http" else "wss"
        self.websocket_url = f"{websocket_scheme}://{websocket_host}/submission_input/{self.user_pk}/{self.submission_id}/{self.secret}/"

        # Nice requests adapter with generous retries/etc.
        self.requests_session = requests.Session()
        adapter = requests.adapters.HTTPAdapter(
            max_retries=Retry(
                total=3,
                backoff_factor=1,
            )
        )
        self.requests_session.mount("http://", adapter)
        self.requests_session.mount("https://", adapter)

    async def watch_detailed_results(self):
        """Watches files alongside scoring + program containers, currently only used
        for detailed_results.html"""
        if not self.detailed_results_url:
            return
        file_path = self.get_detailed_results_file_path()
        last_modified_time = None
        start = time.time()
        expiration_seconds = 60

        # When running scoring program, we have at least one program to run i.e. scoring_program
        # Sometimes when ingestion_only_during_scoring is True, we have two programs to run
        expected_completed_program_counters = 1 + int(bool(self.ingestion_only_during_scoring))
        while self.watch and self.completed_program_counter < expected_completed_program_counters:
            if file_path:
                new_time = os.path.getmtime(file_path)
                if new_time != last_modified_time:
                    last_modified_time = new_time
                    await self.send_detailed_results(file_path)
            else:
                logger.info(time.time() - start)
                if time.time() - start > expiration_seconds:
                    timeout_error_message = (
                        "WARNING: Detailed results not written before the execution."
                    )
                    logger.warning(timeout_error_message)
            await asyncio.sleep(5)
            file_path = self.get_detailed_results_file_path()
        else:
            # make sure we always send the final version of the file
            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 self.docker_image_warning and stream_key == "stderr":
                        data = self.docker_image_warning + data
                    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"
        )
        if os.path.exists(default_detailed_results_path):
            return default_detailed_results_path
        else:
            # v1.5 compatibility - get the first html file if detailed_results.html doesn't exists
            html_files = glob.glob(os.path.join(self.output_dir, "*.html"))
            if html_files:
                return html_files[0]

    async def send_detailed_results(self, file_path):
        logger.info(
            f"Updating detailed results {file_path} - {self.detailed_results_url}"
        )
        self._put_file(
            self.detailed_results_url, file=file_path, content_type="text/html"
        )
        websocket_url = f"{self.websocket_url}?kind=detailed_results"
        logger.info(f"Connecting to {websocket_url} for detailed results")
        # 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
            )
            await websocket.send(
                json.dumps(
                    {
                        "kind": "detailed_result_update",
                    }
                )
            )
        except Exception as e:
            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.
        if not self.is_scoring:
            DETAILED_OUTPUT_NAMES = [
                "prediction_stdout",
                "prediction_stderr",
                "prediction_ingestion_stdout",
                "prediction_ingestion_stderr",
            ]
        else:
            DETAILED_OUTPUT_NAMES = [
                "scoring_stdout",
                "scoring_stderr",
                "scoring_ingestion_stdout",
                "scoring_ingestion_stderr",
            ]
        return [run_args[name] for name in DETAILED_OUTPUT_NAMES]

    def _update_submission(self, data):
        url = f"{self.submissions_api_url}/submissions/{self.submission_id}/"
        data["secret"] = self.secret

        logger.info(f"Updating submission @ {url} with data = {data}")

        resp = self.requests_session.patch(url, data=data, timeout=150)
        if resp.status_code == 200:
            logger.info("Submission updated successfully!")
        else:
            logger.error(
                f"Submission patch failed with status = {resp.status_code}, and response = \n{resp.content}"
            )
            raise SubmissionException("Failure updating submission data.")

    def _update_status(self, status, extra_information=None):
        # Update submission status
        if status not in SubmissionStatus.AVAILABLE_STATUSES:
            raise SubmissionException(
                f"Status '{status}' is not in available statuses: {SubmissionStatus.AVAILABLE_STATUSES}"
            )
        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))
        retries, max_retries = (0, 3)
        while retries < max_retries:
            try:
                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"

            except (docker.errors.APIError, Exception) as pull_error:
                retries += 1
                if retries >= max_retries:
                    logger.error(
                        "There was a problem pulling the image : " + str(pull_error)
                    )
                    # Prepare data to be sent to submissions api
                    docker_pull_fail_data = {
                        "type": "Docker_Image_Pull_Fail",
                        "error_message": pull_error,
                        "is_scoring": self.is_scoring,
                    }
                    # Send data to be written to ingestion logs
                    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"
                    )
                else:
                    logger.warning("Failed. Retrying in 5 seconds...")
                    time.sleep(5)  # Wait 5 seconds before retrying

    async def _send_data_through_socket(self, error_message):
        """
        This function gets an error messages and sends it through a web socket. This function is used for sending
        - Docker image pull failure logs
        - Execution time limit exceeded logs
        """
        # Create a unique websocket URL for error messages
        websocket_url = f"{self.websocket_url}?kind=error_logs"
        logger.info(f"Connecting to {websocket_url} to send error message")

        logger.info(f"Connecting to {websocket_url} to send docker image pull error")

        # connect to web socket
        websocket = await asyncio.wait_for(
            websockets.connect(websocket_url), timeout=10.0
        )

        # define websocket errors
        websocket_errors = (
            socket.gaierror,
            websockets.WebSocketException,
            websockets.ConnectionClosedError,
            ConnectionRefusedError,
        )

        try:
            # send message
            await websocket.send(
                json.dumps({"kind": "stderr", "message": error_message})
            )

        except websocket_errors:
            # handle websocket errors
            logger.error("Error sending failed through websocket")
            try:
                await websocket.close()
            except Exception as e:
                logger.error(e)
        else:
            # no error in websocket message sending
            logger.info("Error sent successfully through websocket")

        logger.info(f"Disconnecting from websocket {websocket_url}")

        # close websocket
        await websocket.close()

    def _get_bundle(self, url, destination, cache=True):
        """Downloads zip from url and unzips into destination. If cache=True then url is hashed and checked
        against existence in CACHE_DIR/<hashed_url> and only downloaded if needed. Cache size is checked
        during the prepare step and cleared if it's over MAX_CACHE_DIR_SIZE_GB.

        :returns zip file path"""
        logger.info(f"Getting bundle {url} to unpack @ {destination}")
        download_needed = True

        # Try to find the bundle in the cache of the worker
        if cache:
            # Hash url and download it if it doesn't exist
            url_without_params = url.split("?")[0]
            url_hash = hashlib.sha256(url_without_params.encode("utf8")).hexdigest()
            bundle_file = os.path.join(Settings.CACHE_DIR, url_hash)
            download_needed = not os.path.exists(bundle_file)
        else:
            if not os.path.exists(self.bundle_dir):
                os.mkdir(self.bundle_dir)
            bundle_file = tempfile.NamedTemporaryFile(
                dir=self.bundle_dir, delete=False
            ).name

        # Fetch and extract
        retries, max_retries = (0, 10)
        while retries < max_retries:
            if download_needed:
                try:
                    # Download the bundle
                    url = rewrite_bundle_url_if_needed(url)
                    urlretrieve(url, bundle_file)
                except HTTPError:
                    raise SubmissionException(
                        f"Problem fetching {url} to put in {destination}"
                    )
            try:
                # Extract the contents to destination directory
                with ZipFile(bundle_file, "r") as z:
                    z.extractall(os.path.join(self.root_dir, destination))
                break  # Break if the loop is successful
            except BadZipFile:
                retries += 1
                if retries >= max_retries:
                    raise SubmissionException("Bad or empty zip file")
                else:
                    logger.warning("Failed. Retrying in 20 seconds...")
                    time.sleep(20)  # Wait 20 seconds before retrying
        # Return the zip file path for other uses, e.g. for creating a MD5 hash to identify it
        return bundle_file

    def _create_container(
        self,
        container_name: str,
        command: str,
        volumes_host: list,
        volumes_config: dict
    ):
        """
        Helper to create and configure a container for ingestion, scoring, or submission.
        Returns the container object.
        """

        logger.info(
            "Creating Container with: "
            f"Container Name: {container_name} \n"
            f"Command: {command} \n"
            "Volumes config:"
        )
        pprint(volumes_config)

        cap_drop_list = [
            "AUDIT_WRITE",
            "CHOWN",
            "DAC_OVERRIDE",
            "FOWNER",
            "FSETID",
            "KILL",
            "MKNOD",
            "NET_BIND_SERVICE",
            "NET_RAW",
            "SETFCAP",
            "SETGID",
            "SETPCAP",
            "SETUID",
            "SYS_CHROOT",
        ]

        # Configure whether or not we use the GPU. Also setting auto_remove to False because
        if Settings.CONTAINER_ENGINE_EXECUTABLE == Settings.DOCKER:
            security_options = ["no-new-privileges"]
        else:
            security_options = ["label=disable"]

        # Setting the device ID like this allows users to specify which gpu to use in the .env file, with all being the default if no value is given
        device_id = [Settings.GPU_DEVICE]
        if Settings.USE_GPU:
            logger.info("Running the container with GPU capabilities")
            host_config = client.create_host_config(
                auto_remove=False,
                cap_drop=cap_drop_list,
                binds=volumes_config,
                userns_mode="host",
                security_opt=security_options,
                device_requests=[
                    {
                        "Driver": "cdi",
                        "DeviceIDs": device_id,
                    },
                ],
            )
        else:
            host_config = client.create_host_config(
                auto_remove=False,
                cap_drop=cap_drop_list,
                binds=volumes_config,
                userns_mode="host",
                security_opt=security_options,
            )

        # Creating container
        # COMPETITION_CONTAINER_NETWORK_DISABLED: Disable or not the competition container access to Internet (False by default)
        # HTTP and HTTPS proxy for the competition container if needed
        container = client.create_container(
            self.container_image,
            name=container_name,
            host_config=host_config,
            detach=False,
            volumes=volumes_host,
            command=command,
            working_dir="/app/program",
            environment=[
                "PYTHONUNBUFFERED=1",
                "http_proxy=" + Settings.COMPETITION_CONTAINER_HTTP_PROXY,
                "https_proxy=" + Settings.COMPETITION_CONTAINER_HTTPS_PROXY,
            ],
            network_disabled=Settings.COMPETITION_CONTAINER_NETWORK_DISABLED,
        )

        logger.debug("Created container: " + str(container))

        return container

    async def _run_container_engine_cmd(self, container, kind):
        """This runs a command and asynchronously writes the data to both a storage file
        and a socket

        :param engine_cmd: the list of container engine command arguments
        :param kind: either 'ingestion' or 'program'
        :return:
        """

        # 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(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}"
            )

            if Settings.LOG_LEVEL == Settings.LOG_LEVEL_DEBUG:
                logger.exception(e)

        start = time.time()

        # Stream the logs of competition container while also sending them to the codabench instance
        try:
            logger.debug("Starting container " + container.get("Id"))
            client.start(container=container.get("Id"))
            logger.debug(
                "Attaching to started container to get the logs :" + container.get("Id")
            )
            container_LogsDemux = client.attach(
                container, demux=True, stream=True, logs=True
            )

            # If we enter the for loop after the container exited, the program will get stuck
            if client.inspect_container(container)["State"]["Status"].lower() == "running":
                logger.debug(
                    "Show the logs and stream them to codabench " + container.get("Id")
                )
                for log in container_LogsDemux:
                    # Output
                    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)

                    # 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)

        except (docker.errors.NotFound, docker.errors.APIError) as e:
            logger.error(e)
        except Exception as e:
            logger.error(
                f"There was an error while starting the container and getting the logs: {e}"
            )
            if Settings.LOG_LEVEL == Settings.LOG_LEVEL_DEBUG:
                logger.exception(e)

        # 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
            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)
            client.remove_container(container, force=True)

            logger.debug(f"Container {container.get('Id')} exited with status code : {str(return_Code['StatusCode'])}")

        except (
            requests.exceptions.ReadTimeout,
            docker.errors.APIError,
            Exception,
        ) as e:
            logger.error(e)
            return_Code = {"StatusCode": 1}

        finally:
            try:
                # Last chance of removing container
                client.remove_container(container.get("Id"), force=True)
            except Exception:
                pass

        self.logs[kind] = {
            "returncode": return_Code["StatusCode"],
            "start": start,
            "end": None,
            "stdout": {
                "data": logs_Unified[0],
                "stream": logs_Unified[0],
                "continue": True,
                "location": self.stdout if kind == ProgramKind.SCORING_PROGRAM else self.ingestion_stdout,
            },
            "stderr": {
                "data": logs_Unified[1],
                "stream": logs_Unified[1],
                "continue": True,
                "location": self.stderr if kind == ProgramKind.SCORING_PROGRAM else self.ingestion_stderr,
            },
        }

        self.logs[kind]["end"] = time.time()

        # Communicate that the program is closing
        self.completed_program_counter += 1

    def _get_host_path(self, *paths):
        """Turns an absolute path inside our container, into what the path
        would be on the host machine. We also ensure that the directory exists,
        docker will create if necessary, but other container engines such as
        podman may not."""
        # Take our list of paths and smash 'em together
        path = os.path.join(*paths)

        # pull front of path, which points to the location inside the container
        path = path[len(Settings.BASE_DIR):]

        # add host to front, so when we run commands in the container on the host they
        # can be seen properly
        path = os.path.join(Settings.HOST_DIRECTORY, path)

        # Create if necessary
        os.makedirs(path, exist_ok=True)

        return path

    async def _run_program_directory(self, kind, program_dir):
        """
        Function responsible for running
            - ingestion program
            - scoring program

        Args:
            kind: `ingestion_program` or `scoring_program`
            program_dir: path to the program to run
        """
        # Return if directory does not exist
        if not os.path.exists(program_dir):
            logger.warning(f"{program_dir} for {kind} not found, no program to execute")

            # Communicate that the program is closing
            self.completed_program_counter += 1
            return

        # Find metadata file.
        # Raise error if metadata is not found
        if os.path.exists(os.path.join(program_dir, "metadata.yaml")):
            metadata_path = "metadata.yaml"
        elif os.path.exists(os.path.join(program_dir, "metadata")):
            metadata_path = "metadata"
        else:
            error_message = f"{program_dir} for {kind} missing 'metadata.yaml/metadata' file."
            logger.error(error_message)
            raise SubmissionException(error_message)

        # Metadata file is found
        logger.info(f"Metadata path is {os.path.join(program_dir, metadata_path)}")

        # Reading metadata file to find command.
        # Raise error if command is not found for ingestion or scoring
        with open(os.path.join(program_dir, metadata_path), "r") as metadata_file:
            command = None
            try:
                metadata = yaml.safe_load(metadata_file.read())
                logger.info(f"Metadata contains:\n {metadata}")
                if isinstance(metadata, dict):
                    command = metadata.get("command")

            except yaml.YAMLError as e:

                logger.error(f"Error parsing YAML file: {e}")

            if not command:
                raise SubmissionException(f"Missing 'command' for {kind} in metadata or metadata format is not correct!")

        # Prepare volumes_host and volumes_config
        volumes_host = [
            self._get_host_path(program_dir),
            self._get_host_path(self.output_dir),
            self.data_dir,
            self._get_host_path(self.root_dir, "submission")
        ]
        volumes_config = {
            volumes_host[0]: {"bind": "/app/program", "mode": "z"},
            volumes_host[1]: {"bind": "/app/output", "mode": "z"},
            volumes_host[2]: {"bind": "/app/data", "mode": "ro"},
            volumes_host[3]: {"bind": "/app/ingested_program", "mode": "ro"},
        }

        # During Scoring: Add `shared` and `/app/input` to volumes_host and update volumes_config
        if kind == ProgramKind.SCORING_PROGRAM:
            # For scoring program, we want to have a shared directory just in case we have an ingestion program.
            volumes_host.extend([self._get_host_path(self.root_dir, "shared")])
            volumes_config.update({volumes_host[-1]: {"bind": "/app/shared"}})

            # Input dir for scoring program
            volumes_host.extend([self._get_host_path(self.input_dir)])
            volumes_config.update({volumes_host[-1]: {"bind": "/app/input"}})

        # During Ingestion: Add `/app/input_data` to volumes_host and update volumes_config
        if kind == ProgramKind.INGESTION_PROGRAM:
            # NOTE: self.input_data is valid when running an ingestion program and competition task has input data
            if self.input_data:
                volumes_host.extend([self._get_host_path(self.root_dir, "input_data")])
                volumes_config.update({volumes_host[-1]: {"bind": "/app/input_data"}})

        # Handle Legacy competitions by replacing anything in the run command
        command = replace_legacy_metadata_command(
            command=command,
            kind=kind,
            is_scoring=self.is_scoring,
            ingestion_only_during_scoring=self.ingestion_only_during_scoring,
        )

        # Create container
        container_name = self.ingestion_program_container_name if kind == ProgramKind.INGESTION_PROGRAM else self.scoring_program_container_name
        container = self._create_container(
            container_name=container_name,
            command=command,
            volumes_host=volumes_host,
            volumes_config=volumes_config
        )

        # This runs the container engine command and asynchronously passes data back via websocket
        try:
            return await self._run_container_engine_cmd(container, kind=kind)
        except Exception as 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."""
        logger.info("Putting dir %s in %s" % (directory, url))
        retries, max_retries = (0, 3)
        while retries < max_retries:
            # Zip the directory
            start_time = time.time()
            zip_path = make_archive(
                os.path.join(self.root_dir, str(uuid.uuid4())), "zip", directory
            )
            duration = time.time() - start_time
            logger.info(f"Time needed to zip archive: {duration} seconds.")
            if is_valid_zip(zip_path):  # Check zip integrity
                self._put_file(url, file=zip_path)  # Send the file
                break  # Leave the loop in case of success
            else:
                retries += 1
                if retries >= max_retries:
                    raise Exception("ZIP file is corrupted or incomplete.")
                else:
                    logger.info("Failed. Retrying in 30 seconds...")
                    time.sleep(30)  # Wait 30 seconds before retrying

    def _put_file(self, url, file=None, raw_data=None, content_type="application/zip"):
        """Send the file in the storage."""
        url = rewrite_bundle_url_if_needed(url)
        if file and raw_data:
            raise Exception("Cannot put both a file and raw_data")

        headers = {
            # For Azure only, other systems ignore these headers
            "x-ms-blob-type": "BlockBlob",
            "x-ms-version": "2018-03-28",
        }
        if content_type:
            headers["Content-Type"] = content_type
        if file:
            logger.info("Putting file %s in %s" % (file, url))
            data = open(file, "rb")
            headers["Content-Length"] = str(os.path.getsize(file))
        elif raw_data is not None:
            logger.info("Putting raw data %s in %s" % (raw_data, url))
            data = raw_data
        else:
            raise SubmissionException(
                "Must provide data, both file and raw_data cannot be empty"
            )

        resp = self.requests_session.put(
            url,
            data=data,
            headers=headers,
        )
        logger.info("*** PUT RESPONSE: ***")
        logger.info(f"response: {resp}")
        logger.info(f"content: {resp.content}")

    def _prep_cache_dir(self, max_size=Settings.MAX_CACHE_DIR_SIZE_GB):
        if not os.path.exists(Settings.CACHE_DIR):
            os.mkdir(Settings.CACHE_DIR)
        logger.info("Checking if cache directory needs to be pruned...")
        if get_folder_size_in_gb(Settings.CACHE_DIR) > max_size:
            logger.info("Pruning cache directory")
            delete_files_in_folder(Settings.CACHE_DIR)
        else:
            logger.info("Cache directory does not need to be pruned!")

    def _copy_submission_to_input_res(self):
        """
        Temporary backward-compatibility function.

        Earlier, scoring programs expected submission files in ingestion output:
        /app/input/res/

        Newer changes expose submission under:
        /app/ingested_program/

        To avoid breaking older scoring programs, we copy the submission
        directory into input/res
        """

        submission_directory = os.path.join(self.root_dir, "submission")
        ingestion_res_directory = os.path.join(self.root_dir, "input/res")

        # copy from submission_directory ingestion_res_directory
        try:
            shutil.copytree(submission_directory, ingestion_res_directory, dirs_exist_ok=True)
            logger.info("Copied submission files to input/res successfully")
        except Exception as e:
            logger.error(f"Failed to copy submission to input/res: {e}")

    def prepare(self):
        hostname = utils.nodenames.gethostname()
        if self.is_scoring:
            self._update_status(SubmissionStatus.RUNNING, extra_information=f"scoring_hostname-{hostname}")
        else:
            # Only during prediction step do we want to announce "preparing"
            self._update_status(SubmissionStatus.PREPARING, extra_information=f"ingestion_hostname-{hostname}")

        # Ensure output_dir exists on the host
        self._get_host_path(self.output_dir)

        # Setup cache and prune if it's out of control
        self._prep_cache_dir()

        # A run *may* contain the following bundles, let's grab them and dump them in the appropriate
        # sub folder.
        bundles = [
            # (url to file, relative folder destination)
            (self.ingestion_program_data, "ingestion_program"),
            (self.scoring_program_data, "scoring_program"),
            (self.submission_data, "submission"),
            (self.input_data, "input_data"),
            (self.reference_data, "input/ref"),
        ]
        if self.is_scoring:
            # Send along submission result so scoring_program can get access
            bundles += [(self.prediction_result, "input/res")]

        for url, path in bundles:
            if url is not None:
                # At the moment let's just cache input & reference data
                cache_this_bundle = path in ("input_data", "input/ref")
                zip_file = self._get_bundle(url, path, cache=cache_this_bundle)

                # Computing checksum of the submission file during ingestion run
                if url == self.submission_data and not self.is_scoring:
                    # We want to get a checksum of submissions so we can check if they are
                    # a solution, or maybe match them against other submissions later
                    logger.info(f"Beginning MD5 checksum of submission: {zip_file}")
                    checksum = md5(zip_file)
                    logger.info(f"Checksum result: {checksum}")
                    self._update_submission({"md5": checksum})

        # During scoring: copy submission files into "input/res"
        if self.is_scoring:
            # NOTE: Temporary compatibility hook (To be removed in the future)
            self._copy_submission_to_input_res()

        # For logging purposes let's dump file names
        for filename in glob.iglob(self.root_dir + "**/*.*", recursive=True):
            logger.info(filename)

        # Before the run starts we want to download images, they may take a while to download
        # and to do this during the run would subtract from the participants time.
        self._get_container_image(self.container_image)
        self._update_status(SubmissionStatus.RUNNING)

    def start(self):

        logger.info(f"Preparing to run: {ProgramKind.SCORING_PROGRAM if self.is_scoring else ProgramKind.INGESTION_PROGRAM}")

        # Define directories for ingestion, scoring and submission
        ingestion_program_dir = os.path.join(self.root_dir, "ingestion_program")
        scoring_program_dir = os.path.join(self.root_dir, "scoring_program")

        loop = asyncio.new_event_loop()
        # Set the event loop for the gather
        asyncio.set_event_loop(loop)

        tasks = []
        if self.is_scoring:
            # During scoring, run scoring program directory
            tasks.append(
                self._run_program_directory(kind=ProgramKind.SCORING_PROGRAM, program_dir=scoring_program_dir)
            )

            # If ingestion_only_during_scoring is true, we also run ingestion program directory in parallel to scoring program
            if self.ingestion_only_during_scoring:
                tasks.append(
                    self._run_program_directory(kind=ProgramKind.INGESTION_PROGRAM, program_dir=ingestion_program_dir)
                )

            # During scoring we watch for detailed results
            tasks.append(
                self.watch_detailed_results()
            )
        else:
            # During ingestion we run ingestion program directory and submission directory
            tasks.extend([
                self._run_program_directory(kind=ProgramKind.INGESTION_PROGRAM, program_dir=ingestion_program_dir),
            ])

        gathered_tasks = asyncio.gather(*tasks, return_exceptions=True)

        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)
            # Prepare data to be sent to submissions api
            execution_time_limit_exceeded_data = {
                "type": "Execution_Time_Limit_Exceeded",
                "error_message": error_message,
                "is_scoring": self.is_scoring,
            }

            # Cleanup containers
            containers_to_kill = [
                self.ingestion_program_container_name,
                self.scoring_program_container_name
            ]
            logger.debug("Trying to kill and remove container " + str(containers_to_kill))

            for container in containers_to_kill:
                try:
                    client.remove_container(str(container), force=True)
                except docker.errors.APIError as e:
                    logger.error(e)
                except Exception as e:
                    logger.error(f"There was a problem killing {containers_to_kill}: {e}")
                    if Settings.LOG_LEVEL == Settings.LOG_LEVEL_DEBUG:
                        logger.exception(e)

            # Send data to be written to ingestion/scoring std_err
            self._update_submission(execution_time_limit_exceeded_data)
            # 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"]
                else:
                    elapsed_time = self.execution_time_limit
                # Normalize the return_code
                return_code = (
                    logs["returncode"]
                    if logs["returncode"] is None or isinstance(logs["returncode"], int)
                    else 1
                )
                if return_code is None:
                    logger.warning("No return code from Process. Killing it")
                    if kind == ProgramKind.INGESTION_PROGRAM:
                        containers_to_kill = self.ingestion_program_container_name
                    else:
                        containers_to_kill = self.scoring_program_container_name
                    try:
                        client.kill(containers_to_kill)
                        client.remove_container(containers_to_kill, force=True)
                    except docker.errors.APIError as e:
                        logger.error(e)
                    except Exception as e:
                        logger.error(
                            f"There was a problem killing {containers_to_kill}: {e}"
                        )
                        if Settings.LOG_LEVEL == Settings.LOG_LEVEL_DEBUG:
                            logger.exception(e)
                if kind == ProgramKind.SCORING_PROGRAM:
                    self.scoring_program_exit_code = return_code
                    self.scoring_program_elapsed_time = elapsed_time
                elif kind == ProgramKind.INGESTION_PROGRAM:
                    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"])

                # set logs of this kind to None, since we handled them already
                logger.info("Program finished")
        signal.alarm(0)

        if self.is_scoring:
            # Check if scoring program failed
            # We have can have 2 or 3 gathered tasks: 3 gathered tasks in case when `ingestion_only_during_scoring` is True, 2 otherwise
            if self.ingestion_only_during_scoring:
                program_results, _, _ = task_results
            else:
                program_results, _ = task_results
            # Gather returns either normal values or exception instances when return_exceptions=True
            had_async_exc = isinstance(
                program_results, BaseException
            ) and not isinstance(program_results, asyncio.CancelledError)
            program_rc = getattr(self, "scoring_program_exit_code", None)
            failed_rc = (program_rc is None) or (program_rc != 0)
            if had_async_exc or failed_rc:
                self._update_status(
                    SubmissionStatus.FAILED,
                    extra_information=f"program_rc={program_rc}, async={task_results}",
                )
                # Raise so upstream marks failed immediately
                raise SubmissionException("Child task failed or non-zero return code")
            self._update_status(SubmissionStatus.FINISHED)

        else:
            self._update_status(SubmissionStatus.SCORING)

    def push_scores(self):
        """This is only ran at the end of the scoring step"""
        # POST to some endpoint:
        # {
        #     "correct": 1.0
        # }
        if os.path.exists(os.path.join(self.output_dir, "scores.json")):
            scores_file = os.path.join(self.output_dir, "scores.json")
            with open(scores_file) as f:
                try:
                    scores = json.load(f)
                except json.decoder.JSONDecodeError as e:
                    raise SubmissionException(
                        f"Could not decode scores json properly, it contains an error.\n{e.msg}"
                    )

        elif os.path.exists(os.path.join(self.output_dir, "scores.txt")):
            scores_file = os.path.join(self.output_dir, "scores.txt")
            with open(scores_file) as f:
                scores = yaml.safe_load(f)
        else:
            raise SubmissionException(
                "Could not find scores file, did the scoring program output it?"
            )

        url = (
            f"{self.submissions_api_url}/upload_submission_scores/{self.submission_id}/"
        )
        data = {
            "secret": self.secret,
            "scores": scores,
        }
        logger.info(f"Submitting these scores to {url}: {scores} with data = {data}")
        resp = self.requests_session.post(url, json=data)
        logger.info(resp)
        logger.info(str(resp.content))

    def push_output(self):
        """Output is pushed at the end of both prediction and scoring steps."""
        # V1.5 compatibility, write program statuses to metadata file
        prog_status = {
            "exitCode": self.scoring_program_exit_code,
            # for v1.5 compat, send `ingestion_elapsed_time` if no `program_elapsed_time`
            "elapsedTime": self.scoring_program_elapsed_time or self.ingestion_program_elapsed_time,
            "ingestionExitCode": self.ingestion_program_exit_code,
            "ingestionElapsedTime": self.ingestion_program_elapsed_time,
        }

        logger.info(f"Metadata output: {prog_status}")

        # Create output_dir if does not exist
        os.makedirs(self.output_dir, exist_ok=True)
        metadata_path = os.path.join(self.output_dir, "metadata")

        if os.path.exists(metadata_path):
            raise SubmissionException(
                "Error, the output directory already contains a metadata file. This file is used "
                "to store exitCode and other data, do not write to this file manually."
            )
        try:
            with open(metadata_path, "w") as f:
                f.write(yaml.dump(prog_status, default_flow_style=False))
        except Exception as e:
            logger.error(e)
            raise SubmissionException("Failed to write metadata file.")

        if not self.is_scoring:
            self._put_dir(self.prediction_result, self.output_dir)
        else:
            self._put_dir(self.scoring_result, self.output_dir)

    def clean_up(self):
        if Settings.CODALAB_IGNORE_CLEANUP_STEP:
            logger.warning(
                f"CODALAB_IGNORE_CLEANUP_STEP mode enabled, ignoring clean up of: {self.root_dir}"
            )
            return

        logger.info(f"Destroying submission temp dir: {self.root_dir}")
        shutil.rmtree(self.root_dir)