tangle

Sleeping

tangle / huggingface_overlay /cloud_pipelines_backend /orchestrator_sql.py
Add debugging information to orchestrator_sql.py
8af11d2 verified about 1 month ago
44.8 kB
import copy
import json
import datetime
import logging
import time
import traceback
import typing
from typing import Any


import sqlalchemy as sql
from sqlalchemy import orm

from cloud_pipelines.orchestration.storage_providers import (
    interfaces as storage_provider_interfaces,
)
from cloud_pipelines.orchestration.launchers import naming_utils

from . import backend_types_sql as bts
from . import component_structures as structures
from .launchers import common_annotations
from .launchers import interfaces as launcher_interfaces


_logger = logging.getLogger(__name__)

_T = typing.TypeVar("_T")


class OrchestratorError(RuntimeError):
    pass


class OrchestratorService_Sql:
    def __init__(
        self,
        session_factory: typing.Callable[[], orm.Session],
        launcher: launcher_interfaces.ContainerTaskLauncher[
            launcher_interfaces.LaunchedContainer
        ],
        storage_provider: storage_provider_interfaces.StorageProvider,
        data_root_uri: str,
        logs_root_uri: str,
        default_task_annotations: dict[str, Any] | None = None,
        sleep_seconds_between_queue_sweeps: float = 1.0,
        output_data_purge_duration: datetime.timedelta = None,
    ):
        self._session_factory = session_factory
        self._launcher = launcher
        self._storage_provider = storage_provider
        self._data_root_uri = data_root_uri
        self._logs_root_uri = logs_root_uri
        self._default_task_annotations = default_task_annotations
        self._sleep_seconds_between_queue_sweeps = sleep_seconds_between_queue_sweeps
        self._queued_executions_queue_idle = False
        self._running_executions_queue_idle = False
        self._output_data_purge_duration = output_data_purge_duration

    def run_loop(self):
        while True:
            try:
                self.process_each_queue_once()
                time.sleep(self._sleep_seconds_between_queue_sweeps)
            except:
                _logger.exception("Error while calling `process_each_queue_once`")

    def process_each_queue_once(self):
        queue_handlers = [
            self.internal_process_queued_executions_queue,
            self.internal_process_running_executions_queue,
        ]
        for queue_handler in queue_handlers:
            try:
                with self._session_factory() as session:
                    queue_handler(session=session)
            except:
                _logger.exception(f"Error while executing {queue_handler=}")

    def internal_process_queued_executions_queue(self, session: orm.Session):
        query = (
            sql.select(bts.ExecutionNode).where(
                bts.ExecutionNode.container_execution_status.in_(
                    (
                        bts.ContainerExecutionStatus.UNINITIALIZED,
                        bts.ContainerExecutionStatus.QUEUED,
                    )
                )
            )
            # TODO: Maybe add last_processed_at
            # .order_by(bts.ExecutionNode.last_processed_at)
            # .order_by(bts.ExecutionNode.id)
            .limit(1)
        )
        queued_execution = session.scalar(query)
        if queued_execution:
            self._queued_executions_queue_idle = False
            _logger.info(f"Before processing {queued_execution.id=}")
            try:
                self.internal_process_one_queued_execution(
                    session=session, execution=queued_execution
                )
            except Exception as ex:
                _logger.exception(f"Error processing {queued_execution.id=}")
                session.rollback()
                queued_execution.container_execution_status = (
                    bts.ContainerExecutionStatus.SYSTEM_ERROR
                )
                record_system_error_exception(execution=queued_execution, exception=ex)
                session.commit()
            _logger.info(f"After processing {queued_execution.id=}")
            return True
        else:
            if not self._queued_executions_queue_idle:
                self._queued_executions_queue_idle = True
                _logger.debug(f"No queued executions found")
            return False

    def internal_process_running_executions_queue(self, session: orm.Session):
        query = (
            sql.select(bts.ContainerExecution)
            .where(
                bts.ContainerExecution.status.in_(
                    (
                        bts.ContainerExecutionStatus.PENDING,
                        bts.ContainerExecutionStatus.RUNNING,
                    )
                )
            )
            .order_by(bts.ContainerExecution.last_processed_at)
            .limit(1)
        )
        running_container_execution = session.scalar(query)
        if running_container_execution:
            self._running_executions_queue_idle = False
            try:
                _logger.info(f"Before processing {running_container_execution.id=}")
                self.internal_process_one_running_execution(
                    session=session, container_execution=running_container_execution
                )
                _logger.info(f"After processing {running_container_execution.id=}")
            except Exception as ex:
                _logger.exception(f"Error processing {running_container_execution.id=}")
                session.rollback()
                running_container_execution.status = (
                    bts.ContainerExecutionStatus.SYSTEM_ERROR
                )
                # Doing an intermediate commit here because it's most important to mark the problematic execution as SYSTEM_ERROR.
                session.commit()
                # Mark our ExecutionNode as SYSTEM_ERROR
                execution_nodes = running_container_execution.execution_nodes
                for execution_node in execution_nodes:
                    execution_node.container_execution_status = (
                        bts.ContainerExecutionStatus.SYSTEM_ERROR
                    )
                    record_system_error_exception(
                        execution=execution_node, exception=ex
                    )
                # Doing an intermediate commit here because it's most important to mark the problematic node as SYSTEM_ERROR.
                session.commit()
                # Skip downstream executions
                for execution_node in execution_nodes:
                    _mark_all_downstream_executions_as_skipped(
                        session=session, execution=execution_node
                    )
                session.commit()
            return True
        else:
            if not self._running_executions_queue_idle:
                _logger.debug(f"No running container executions found")
                self._running_executions_queue_idle = True
            return False

    def internal_process_one_queued_execution(
        self, session: orm.Session, execution: bts.ExecutionNode
    ):
        task_spec = structures.TaskSpec.from_json_dict(execution.task_spec)
        component_spec = task_spec.component_ref.spec
        if component_spec is None:
            raise OrchestratorError(
                f"Missing ComponentSpec. {task_spec.component_ref=}"
            )
        implementation = component_spec.implementation
        if not implementation:
            raise OrchestratorError(f"Missing implementation. {component_spec=}")
        if implementation is None or not isinstance(
            implementation, structures.ContainerImplementation
        ):
            raise OrchestratorError(
                f"Implementation must be container. {implementation=}"
            )
        container_spec = implementation.container
        input_artifact_data = dict(
            session.execute(
                sql.select(bts.InputArtifactLink.input_name, bts.ArtifactData)
                .where(bts.InputArtifactLink.execution == execution)
                .join(bts.InputArtifactLink.artifact)
                .join(bts.ArtifactNode.artifact_data, isouter=True)
            )
            .tuples()
            .all()
        )
        if None in input_artifact_data.values():
            _logger.info(
                f"Execution did not have all input artifact data present. Waiting for upstream. {execution.id=}"
            )
            execution.container_execution_status = (
                bts.ContainerExecutionStatus.WAITING_FOR_UPSTREAM
            )
            session.commit()
            return

        cache_key = _calculate_container_execution_cache_key(
            container_spec=container_spec, input_artifact_data=input_artifact_data
        )

        # Trying to reuse an older execution from cache.
        max_cache_staleness_str: str | None = None
        if task_spec.execution_options and task_spec.execution_options.caching_strategy:
            max_cache_staleness_str = (
                task_spec.execution_options.caching_strategy.max_cache_staleness
            )

        # TODO: Support other ways to express 0-length time period.
        if max_cache_staleness_str == "P0D":
            execution_candidates: typing.Sequence[bts.ExecutionNode] = []
        else:
            # TODO: Support cache reuse options (max cached data staleness etc).
            # TODO: It might be better to search for ContainerExecutions rather than ExecutionNodes.
            execution_candidates = session.scalars(
                sql.select(bts.ExecutionNode)
                .where(bts.ExecutionNode.container_execution_cache_key == cache_key)
                .where(
                    bts.ExecutionNode.container_execution_status.in_(
                        [
                            # We can reuse both succeeded executions and also non yet finished ones.
                            # Reusing still running executions is important since it allows cache reuse
                            # when multiple versions of a pipeline are submitted in parallel.
                            # bts.ContainerExecutionStatus.STARTING,  # Doesn't exist yet
                            bts.ContainerExecutionStatus.PENDING,
                            bts.ContainerExecutionStatus.RUNNING,
                            bts.ContainerExecutionStatus.SUCCEEDED,
                        ]
                    )
                )
                # TODO: Filter by purged==False. Or `expires_at`.
                .join(bts.ContainerExecution)
                .order_by(bts.ContainerExecution.created_at.desc())
            ).all()
        non_purged_candidates = [
            execution_candidate
            for execution_candidate in execution_candidates
            if not (execution_candidate.extra_data or {}).get("is_purged", False)
        ]
        if self._output_data_purge_duration:
            current_time = datetime.datetime.now(datetime.timezone.utc)
            data_purge_threshold_time = current_time - self._output_data_purge_duration
            non_purged_candidates = [
                execution_candidate
                for execution_candidate in execution_candidates
                # The `execution_candidate.container_execution` attribute access accesses the DB.
                # TODO: Move the time filtering to the DB side.
                # TODO: Filter by `ended_at`, not `created_at`.
                # Note: Need `.astimezone` to avoid `TypeError: can't compare offset-naive and offset-aware datetimes`
                if execution_candidate.container_execution.created_at.astimezone(
                    datetime.timezone.utc
                )
                > data_purge_threshold_time
            ]

        if non_purged_candidates:
            # Re-using the oldest candidate
            # TODO: Use better cached execution selection strategy:
            # 1. Try to reuse the execution with `max_cache_staleness in extra_data["used_for_max_cache_staleness"]`. There should be 0 or 1 of those.
            # 2. Else: Try to reuse the latest SUCCEEDED execution. Mark it `extra_data["used_for_max_cache_staleness"][max_cache_staleness] = True`.
            # 3. Else: Try to reuse the latest RUNNING/PENDING execution. Mark it `extra_data["used_for_max_cache_staleness"][max_cache_staleness] = True`.
            # There must be at least one SUCCEEDED/RUNNING/PENDING since non_purged_candidates is non-empty.
            old_execution = non_purged_candidates[-1]
            _logger.info(
                f"Execution {execution.id=} will reuse the {old_execution.id=} with "
                f"{old_execution.container_execution_id=}, {old_execution.container_execution_status=}"
            )
            # Reusing the execution:
            if not execution.extra_data:
                execution.extra_data = {}
            execution.extra_data["reused_from_execution_node_id"] = old_execution.id

            execution.container_execution_status = (
                old_execution.container_execution_status
            )

            # If the execution is still running, it's enough to just link execution node to it.
            # The container execution will set the outputs itself when it succeeds.
            execution.container_execution_id = old_execution.container_execution_id
            # However if the container execution has already ended, we need to copy the outputs ourselves.
            if (
                old_execution.container_execution_status
                == bts.ContainerExecutionStatus.SUCCEEDED
            ):
                # Copying the output artifact data (if the execution already succeeded).
                reused_execution_output_artifact_data_ids = {
                    output_name: artifact_data_id
                    for output_name, artifact_data_id in session.execute(
                        sql.select(
                            bts.OutputArtifactLink.output_name,
                            bts.ArtifactNode.artifact_data_id,
                        )
                        .join(bts.OutputArtifactLink.artifact)
                        .where(bts.OutputArtifactLink.execution_id == old_execution.id)
                    ).tuples()
                }
                # Setting artifact data
                for output_name, artifact_node in session.execute(
                    sql.select(bts.OutputArtifactLink.output_name, bts.ArtifactNode)
                    # TODO: Verify that this improves the performance
                    .with_for_update()
                    .join(bts.OutputArtifactLink.artifact)
                    .where(bts.OutputArtifactLink.execution_id == execution.id)
                ).tuples():
                    artifact_node.artifact_data_id = (
                        reused_execution_output_artifact_data_ids[output_name]
                    )
                # Waking up all direct downstream executions.
                # Many of them will get processed and go back to WAITING_FOR_UPSTREAM state.
                for downstream_execution in _get_direct_downstream_executions(
                    session=session, execution=execution
                ):
                    if (
                        downstream_execution.container_execution_status
                        == bts.ContainerExecutionStatus.WAITING_FOR_UPSTREAM
                    ):
                        downstream_execution.container_execution_status = (
                            bts.ContainerExecutionStatus.QUEUED
                        )
            session.commit()
            return

        # Launching the container execution

        pipeline_run = session.scalar(
            sql.select(bts.PipelineRun)
            .join(
                bts.ExecutionToAncestorExecutionLink,
                bts.ExecutionToAncestorExecutionLink.ancestor_execution_id
                == bts.PipelineRun.root_execution_id,
            )
            .where(bts.ExecutionToAncestorExecutionLink.execution_id == execution.id)
        )
        session.rollback()

        if pipeline_run is None:
            raise OrchestratorError(
                f"Execution {execution} is not associated with a PipelineRun."
            )

        # Handling cancellation.
        # In case of cancellation, we do not create container execution and we skip all downstream executions.
        should_terminate_execution = (execution.extra_data or {}).get(
            "desired_state"
        ) == "TERMINATED"
        should_terminate_run = (pipeline_run.extra_data or {}).get(
            "desired_state"
        ) == "TERMINATED"
        should_terminate = should_terminate_execution or should_terminate_run
        if should_terminate:
            _logger.info(
                f"Cancelling execution {execution.id} and skipping all downstream executions."
            )
            execution.container_execution_status = (
                bts.ContainerExecutionStatus.CANCELLED
            )
            _mark_all_downstream_executions_as_skipped(
                session=session, execution=execution
            )

        # Creating new container execution
        container_execution_uuid = _generate_random_id()

        _ARTIFACT_PATH_LAST_PART = "data"
        _LOG_FILE_NAME = "log.txt"
        # _TIMESTAMPED_LOG_FILE_NAME = "timestamped.log.txt"
        # _STDOUT_LOG_FILE_NAME = "stdout.log.txt"
        # _STDERR_LOG_FILE_NAME = "stderr.log.txt"

        # Input URIs might be needed to write constant input data.
        # Or to prevent the data from being purged during the execution.
        # TODO: Ensure unique artifact URIs (they can become non-unique due to sanitization)
        def generate_input_artifact_uri(
            root_dir: str,
            execution_id: str,
            input_name: str,
        ) -> str:
            root_dir = root_dir.rstrip("/")
            execution_id_str = naming_utils.sanitize_file_name(execution_id)
            input_name_str = naming_utils.sanitize_file_name(input_name)
            return f"{root_dir}/by_execution/{execution_id_str}/inputs/{input_name_str}/{_ARTIFACT_PATH_LAST_PART}"

        def generate_output_artifact_uri(
            root_dir: str,
            execution_id: str,
            output_name: str,
        ) -> str:
            root_dir = root_dir.rstrip("/")
            execution_id_str = naming_utils.sanitize_file_name(execution_id)
            output_name_str = naming_utils.sanitize_file_name(output_name)
            return f"{root_dir}/by_execution/{execution_id_str}/outputs/{output_name_str}/{_ARTIFACT_PATH_LAST_PART}"

        def generate_execution_log_uri(
            root_dir: str,
            execution_id: str,
        ) -> str:
            root_dir = root_dir.rstrip("/")
            execution_id_str = naming_utils.sanitize_file_name(execution_id)
            return f"{root_dir}/by_execution/{execution_id_str}/{_LOG_FILE_NAME}"

        log_uri = generate_execution_log_uri(
            root_dir=self._logs_root_uri, execution_id=container_execution_uuid
        )
        output_artifact_uris = {
            output_spec.name: generate_output_artifact_uri(
                root_dir=self._data_root_uri,
                execution_id=container_execution_uuid,
                output_name=output_spec.name,
            )
            for output_spec in component_spec.outputs or []
        }

        input_arguments = {
            input_name: launcher_interfaces.InputArgument(
                total_size=artifact_data.total_size,
                is_dir=artifact_data.is_dir,
                value=artifact_data.value,
                uri=artifact_data.uri,
                staging_uri=generate_input_artifact_uri(
                    root_dir=self._data_root_uri,
                    execution_id=container_execution_uuid,
                    input_name=input_name,
                ),
            )
            for input_name, artifact_data in input_artifact_data.items()
        }

        # Handling annotations

        full_annotations = {}
        _update_dict_recursive(
            full_annotations, copy.deepcopy(self._default_task_annotations or {})
        )
        _update_dict_recursive(
            full_annotations, copy.deepcopy(pipeline_run.annotations or {})
        )
        _update_dict_recursive(
            full_annotations, copy.deepcopy(task_spec.annotations or {})
        )

        full_annotations[common_annotations.PIPELINE_RUN_CREATED_BY_ANNOTATION_KEY] = (
            pipeline_run.created_by
        )
        full_annotations[common_annotations.PIPELINE_RUN_ID_ANNOTATION_KEY] = (
            pipeline_run.id
        )
        full_annotations[common_annotations.EXECUTION_NODE_ID_ANNOTATION_KEY] = (
            execution.id
        )
        full_annotations[common_annotations.CONTAINER_EXECUTION_ID_ANNOTATION_KEY] = (
            container_execution_uuid
        )

        try:
            launched_container: launcher_interfaces.LaunchedContainer = (
                self._launcher.launch_container_task(
                    component_spec=component_spec,
                    # The value and uri might be updated during launching.
                    input_arguments=input_arguments,
                    output_uris=output_artifact_uris,
                    log_uri=log_uri,
                    annotations=full_annotations,
                )
            )
            if launched_container.status not in (
                launcher_interfaces.ContainerStatus.PENDING,
                launcher_interfaces.ContainerStatus.RUNNING,
            ):
                raise OrchestratorError(
                    f"Unexpected status of just launched container: {launched_container.status=}, {launched_container=}"
                )
        except Exception as ex:
            session.rollback()
            with session.begin():
                # Logs whole exception
                _logger.exception(f"Error launching container for {execution.id=}")
                execution.container_execution_status = (
                    bts.ContainerExecutionStatus.SYSTEM_ERROR
                )
                record_system_error_exception(execution=execution, exception=ex)
                _mark_all_downstream_executions_as_skipped(
                    session=session, execution=execution
                )
            return

        current_time = _get_current_time()

        container_execution = bts.ContainerExecution(
            status=bts.ContainerExecutionStatus(launched_container.status),
            last_processed_at=current_time,
            created_at=current_time,
            launcher_data=launched_container.to_dict(),
            # Updating the map with uploaded value URIs and downloaded values.
            input_artifact_data_map={
                input_name: dict(
                    id=artifact_data.id,
                    is_dir=artifact_data.is_dir,
                    total_size=artifact_data.total_size,
                    hash=artifact_data.hash,
                    # The value and URI might have been updated during launching.
                    value=input_arguments[input_name].value,
                    uri=input_arguments[input_name].uri,
                )
                for input_name, artifact_data in input_artifact_data.items()
            },
            output_artifact_data_map={
                output_name: dict(
                    uri=uri,
                )
                for output_name, uri in output_artifact_uris.items()
            },
            log_uri=log_uri,
        )
        container_execution.id = container_execution_uuid

        # Need to commit/rollback before explicitly beginning non-nested transaction.
        session.rollback()
        with session.begin():
            session.add(container_execution)
            execution.container_execution = container_execution
            execution.container_execution_cache_key = cache_key
            execution.container_execution_status = container_execution.status
            # TODO: Maybe add artifact value and URI to input ArtifactData.

    def internal_process_one_running_execution(
        self, session: orm.Session, container_execution: bts.ContainerExecution
    ):
        # Should we update last_processed_at early? Should we do it always (even when there are errors).
        # ALways updating the last_processed_at such that a problematic container does not get queue stuck.
        session.expire_on_commit = False
        session.rollback()
        with session.begin():
            container_execution.last_processed_at = _get_current_time()
        launcher_data = _assert_not_none(container_execution.launcher_data)
        launched_container: launcher_interfaces.LaunchedContainer = (
            self._launcher.deserialize_launched_container_from_dict(launcher_data)
        )
        previous_status = launched_container.status
        if previous_status not in (
            bts.ContainerExecutionStatus.PENDING,
            bts.ContainerExecutionStatus.RUNNING,
        ):
            raise OrchestratorError(
                f"Unexpected running container status: {previous_status=}, {launched_container=}"
            )

        # Handling cancellation
        votes_to_terminate = []
        votes_to_not_terminate = []
        # TODO: Get the desired state from the pipeline runs, not execution nodes
        execution_nodes = container_execution.execution_nodes
        for execution_node in execution_nodes:
            should_terminate = (execution_node.extra_data or {}).get(
                "desired_state"
            ) == "TERMINATED"
            if should_terminate:
                votes_to_terminate.append(execution_node)
            else:
                votes_to_not_terminate.append(execution_node)

        if votes_to_terminate:
            # Terminate the container execution only when all execution nodes pointing to it are asked to terminate.
            terminated = False
            if votes_to_not_terminate:
                _logger.info(
                    f"Not terminating container execution {container_execution.id=} since some other executions ({[execution_node.id for execution_node in votes_to_not_terminate]}) are still using it."
                )
            else:
                _logger.info(
                    f"Terminating container execution {container_execution.id}."
                )
                # We should preserve the logs before terminating/deleting the container
                try:
                    _retry(lambda: launched_container.upload_log())
                except:
                    _logger.exception(
                        f"Error uploading logs for {container_execution.id=} before termination."
                    )
                # Requesting container termination.
                # Termination might not happen immediately (e.g. Kubernetes has grace period).
                launched_container.terminate()
                container_execution.ended_at = _get_current_time()
                # We need to mark the execution as CANCELLED otherwise orchestrator will continue polling it.
                container_execution.status = bts.ContainerExecutionStatus.CANCELLED
                terminated = True

            # Mark the execution nodes as cancelled only after the launched container is successfully terminated (if needed)
            for execution_node in votes_to_terminate:
                _logger.info(
                    f"Cancelling execution {execution_node.id} ({container_execution.id=}) and skipping all downstream executions."
                )
                execution_node.container_execution_status = (
                    bts.ContainerExecutionStatus.CANCELLED
                )
                _mark_all_downstream_executions_as_skipped(
                    session=session, execution=execution_node
                )
            session.commit()
            if terminated:
                return

        # Asking the launcher to refresh the container state.
        reloaded_launched_container: launcher_interfaces.LaunchedContainer = (
            self._launcher.get_refreshed_launched_container_from_dict(launcher_data)
        )
        current_time = _get_current_time()
        # Saving the updated launcher data
        reloaded_launcher_data = reloaded_launched_container.to_dict()
        if reloaded_launcher_data != launcher_data:
            session.rollback()
            with session.begin():
                import json

                try:
                    json.dumps(reloaded_launcher_data)
                except TypeError as e:
                    print("launcher_data JSON ERROR:", e)
                    print("launcher_data:", repr(reloaded_launcher_data))
                    raise
            
                container_execution.launcher_data = reloaded_launcher_data
                container_execution.updated_at = current_time
        new_status: launcher_interfaces.ContainerStatus = (
            reloaded_launched_container.status
        )
        if new_status == previous_status:
            _logger.info(
                f"Container execution {container_execution.id} remains in {new_status} state."
            )
            return
        _logger.info(
            f"Container execution {container_execution.id} is now in state {new_status} (was {previous_status})."
        )
        session.rollback()
        container_execution.updated_at = current_time
        execution_nodes = container_execution.execution_nodes
        if not execution_nodes:
            raise OrchestratorError(
                f"Could not find ExecutionNode associated with ContainerExecution. {container_execution.id=}"
            )
        if len(execution_nodes) > 1:
            execution_node_ids = [execution.id for execution in execution_nodes]
            _logger.warning(
                f"ContainerExecution is associated with multiple ExecutionNodes: {container_execution.id=}, {execution_node_ids=}"
            )

        if new_status == launcher_interfaces.ContainerStatus.RUNNING:
            container_execution.status = bts.ContainerExecutionStatus.RUNNING
            container_execution.started_at = reloaded_launched_container.started_at
            for execution_node in execution_nodes:
                execution_node.container_execution_status = (
                    bts.ContainerExecutionStatus.RUNNING
                )
        elif new_status == launcher_interfaces.ContainerStatus.SUCCEEDED:
            container_execution.status = bts.ContainerExecutionStatus.SUCCEEDED
            container_execution.exit_code = reloaded_launched_container.exit_code
            container_execution.started_at = reloaded_launched_container.started_at
            container_execution.ended_at = reloaded_launched_container.ended_at

            # Don't fail the execution if log upload fails.
            # Logs are important, but not so important that we should fail a successfully completed container execution.
            try:
                _retry(reloaded_launched_container.upload_log)
            except Exception as ex:
                _logger.exception(
                    f"! Error during `LaunchedContainer.upload_log` call: {ex}."
                )

            _MAX_PRELOAD_VALUE_SIZE = 255

            def _maybe_preload_value(
                uri_reader: storage_provider_interfaces.UriReader,
                data_info: storage_provider_interfaces.DataInfo,
            ) -> str | None:
                """Preloads artifact value is it's small enough (e.g. <=255 bytes)"""
                if (
                    not data_info.is_dir
                    and data_info.total_size < _MAX_PRELOAD_VALUE_SIZE
                ):
                    # Don't fail the execution if small value preloading fails.
                    # Those values may be useful for preservation, but not so important that we should fail a successfully completed container execution.
                    try:
                        data = uri_reader.download_as_bytes()
                    except Exception as ex:
                        _logger.exception(
                            f"Error during preloading small artifact values."
                        )
                        return None
                    try:
                        text = data.decode("utf-8")
                        return text
                    except:
                        pass

            output_artifact_uris: dict[str, str] = {
                output_name: output_artifact_info_dict["uri"]
                for output_name, output_artifact_info_dict in container_execution.output_artifact_data_map.items()
            }

            # We need to first check that the output data exists. Otherwise `uri_reader.get_info()`` throws `IndexError`
            # in cloud_pipelines/orchestration/storage_providers/interfaces.py", line 172, in _make_data_info_for_dir
            missing_output_names = [
                output_name
                for output_name, uri in output_artifact_uris.items()
                if not _retry(
                    lambda: self._storage_provider.make_uri(uri).get_reader().exists()
                )
            ]

            if missing_output_names:
                # Marking the container execution as FAILED (even though the program itself has completed successfully)
                container_execution.status = bts.ContainerExecutionStatus.FAILED
                orchestration_error_message = f"Container execution {container_execution.id} is marked as FAILED due to missing outputs: {missing_output_names}."
                _logger.error(orchestration_error_message)
                _record_orchestration_error_message(
                    container_execution=container_execution,
                    execution_nodes=execution_nodes,
                    message=orchestration_error_message,
                )
                # Skip downstream executions
                for execution_node in execution_nodes:
                    execution_node.container_execution_status = (
                        bts.ContainerExecutionStatus.FAILED
                    )
                    _mark_all_downstream_executions_as_skipped(
                        session=session, execution=execution_node
                    )
            else:
                output_artifact_data_info_map = {
                    output_name: _retry(
                        lambda: self._storage_provider.make_uri(uri)
                        .get_reader()
                        .get_info()
                    )
                    for output_name, uri in output_artifact_uris.items()
                }
                new_output_artifact_data_map = {
                    output_name: bts.ArtifactData(
                        total_size=data_info.total_size,
                        is_dir=data_info.is_dir,
                        # Using 1st hash only.
                        # TODO: Support multiple hashes.
                        hash=[
                            f"{hash_name}={hash_value}"
                            for hash_name, hash_value in data_info.hashes.items()
                        ][0],
                        uri=output_artifact_uris[output_name],
                        # Preloading artifact value is it's small enough (e.g. <=255 bytes)
                        value=_retry(
                            lambda: _maybe_get_small_artifact_value(
                                uri_reader=self._storage_provider.make_uri(
                                    output_artifact_uris[output_name]
                                ).get_reader(),
                                data_info=data_info,
                            )
                        ),
                        created_at=current_time,
                    )
                    for output_name, data_info in output_artifact_data_info_map.items()
                }

                session.add_all(new_output_artifact_data_map.values())
                for execution_node in execution_nodes:
                    execution_node.container_execution_status = (
                        bts.ContainerExecutionStatus.SUCCEEDED
                    )
                    # TODO: Optimize
                    for output_name, artifact_node in session.execute(
                        sql.select(bts.OutputArtifactLink.output_name, bts.ArtifactNode)
                        .join(bts.OutputArtifactLink.artifact)
                        .where(bts.OutputArtifactLink.execution_id == execution_node.id)
                    ).tuples():
                        artifact_node.artifact_data = new_output_artifact_data_map[
                            output_name
                        ]
                        artifact_node.had_data_in_past = True
                    # Waking up all direct downstream executions.
                    # Many of them will get processed and go back to WAITING_FOR_UPSTREAM state.
                    for downstream_execution in _get_direct_downstream_executions(
                        session=session, execution=execution_node
                    ):
                        if (
                            downstream_execution.container_execution_status
                            == bts.ContainerExecutionStatus.WAITING_FOR_UPSTREAM
                        ):
                            downstream_execution.container_execution_status = (
                                bts.ContainerExecutionStatus.QUEUED
                            )
        elif new_status == launcher_interfaces.ContainerStatus.FAILED:
            container_execution.status = bts.ContainerExecutionStatus.FAILED
            container_execution.exit_code = reloaded_launched_container.exit_code
            container_execution.started_at = reloaded_launched_container.started_at
            container_execution.ended_at = reloaded_launched_container.ended_at
            launcher_error = reloaded_launched_container.launcher_error_message
            if launcher_error:
                orchestration_error_message = f"Launcher error: {launcher_error}"
                _record_orchestration_error_message(
                    container_execution=container_execution,
                    execution_nodes=execution_nodes,
                    message=orchestration_error_message,
                )

            _retry(reloaded_launched_container.upload_log)
            # Skip downstream executions
            for execution_node in execution_nodes:
                execution_node.container_execution_status = (
                    bts.ContainerExecutionStatus.FAILED
                )
                _mark_all_downstream_executions_as_skipped(
                    session=session, execution=execution_node
                )
        else:
            _logger.error(
                f"Container execution {container_execution.id} is now in unexpected state {new_status}. System error. {container_execution=}"
            )
            # This SYSTEM_ERROR will be handled by the outer exception handler
            raise OrchestratorError(
                f"Unexpected running container status: {new_status=}, {launched_container=}"
            )
        session.commit()


def _get_direct_downstream_executions(
    session: orm.Session, execution: bts.ExecutionNode
):
    return session.scalars(
        sql.select(bts.ExecutionNode)
        .select_from(bts.OutputArtifactLink)
        .where(bts.OutputArtifactLink.execution_id == execution.id)
        .join(bts.ArtifactNode)
        .join(bts.InputArtifactLink)
        .join(bts.ExecutionNode)
    ).all()


# TODO: Cover with tests
def _mark_all_downstream_executions_as_skipped(
    session: orm.Session,
    execution: bts.ExecutionNode,
    seen_execution_ids: set[bts.IdType] | None = None,
):
    if seen_execution_ids is None:
        seen_execution_ids = set()
    if execution.id in seen_execution_ids:
        return
    seen_execution_ids.add(execution.id)
    if execution.container_execution_status in {
        bts.ContainerExecutionStatus.WAITING_FOR_UPSTREAM
    }:
        execution.container_execution_status = bts.ContainerExecutionStatus.SKIPPED

    # for artifact_node in execution.output_artifact_nodes:
    #     for downstream_execution in artifact_node.downstream_executions:
    for downstream_execution in _get_direct_downstream_executions(session, execution):
        _mark_all_downstream_executions_as_skipped(
            session=session,
            execution=downstream_execution,
            seen_execution_ids=seen_execution_ids,
        )


def _assert_type(value: typing.Any, typ: typing.Type[_T]) -> _T:
    if not isinstance(value, typ):
        raise TypeError(f"Expected type {typ}, but got {type(value)}: {value}")
    return value


def _assert_not_none(value: _T | None) -> _T:
    if value is None:
        raise TypeError(f"Expected value to be not None, but got {value}.")
    return value


def _calculate_hash(s: str) -> str:
    import hashlib

    return "md5=" + hashlib.md5(s.encode("utf-8")).hexdigest()


def _calculate_container_execution_cache_key(
    container_spec: structures.ContainerSpec,
    input_artifact_data: dict[str, bts.ArtifactData],
):
    # Using ContainerSPec instead of the whole ComponentSpec.
    input_hashes = {
        input_name: _assert_not_none(artifact_data).hash
        for input_name, artifact_data in (input_artifact_data or {}).items()
    }
    cache_key_struct = {
        "container_spec": container_spec.to_json_dict(),
        "input_hashes": input_hashes,
    }
    # Need to sort keys to ensure consistency
    cache_key_str = json.dumps(cache_key_struct, separators=(",", ":"), sort_keys=True)
    # We could use a different hash function, but there is no reason to.
    cache_key = _calculate_hash(cache_key_str)
    return cache_key


def _get_current_time() -> datetime.datetime:
    return datetime.datetime.now(tz=datetime.timezone.utc)


def _generate_random_id() -> str:
    import os
    import time

    random_bytes = os.urandom(4)
    nanoseconds = time.time_ns()
    milliseconds = nanoseconds // 1_000_000

    return ("%012x" % milliseconds) + random_bytes.hex()


def _update_dict_recursive(d1: dict, d2: dict):
    for k, v2 in d2.items():
        if k in d1:
            v1 = d1[k]
            if isinstance(v1, dict) and isinstance(v2, dict):
                _update_dict_recursive(v1, v2)
                continue
            # elif isinstance(v1, list) and isinstance(v2, list):
            # # Merging lists is not supported yet
        d1[k] = v2


def _retry(
    func: typing.Callable[[], _T], max_retries: int = 5, wait_seconds: float = 1.0
) -> _T:
    for i in range(max_retries):
        try:
            return func()
        except:
            _logger.exception(f"Exception calling {func}.")
            time.sleep(wait_seconds)
            if i == max_retries - 1:
                raise
    raise


def record_system_error_exception(execution: bts.ExecutionNode, exception: Exception):
    if execution.extra_data is None:
        execution.extra_data = {}
    execution.extra_data[
        bts.EXECUTION_NODE_EXTRA_DATA_SYSTEM_ERROR_EXCEPTION_MESSAGE_KEY
    ] = "".join(traceback.format_exception_only(type(exception), exception))
    execution.extra_data[
        bts.EXECUTION_NODE_EXTRA_DATA_SYSTEM_ERROR_EXCEPTION_FULL_KEY
    ] = traceback.format_exc()


def _record_orchestration_error_message(
    container_execution: bts.ContainerExecution,
    execution_nodes: list[bts.ExecutionNode],
    message: str,
):
    if container_execution.extra_data is None:
        container_execution.extra_data = {}
    container_execution.extra_data[
        bts.CONTAINER_EXECUTION_EXTRA_DATA_ORCHESTRATION_ERROR_MESSAGE_KEY
    ] = message

    for execution_node in execution_nodes:
        if execution_node.extra_data is None:
            execution_node.extra_data = {}
        execution_node.extra_data[
            bts.EXECUTION_NODE_EXTRA_DATA_ORCHESTRATION_ERROR_MESSAGE_KEY
        ] = message


_MAX_PRELOAD_VALUE_SIZE = 255


def _maybe_get_small_artifact_value(
    uri_reader: storage_provider_interfaces.UriReader,
    data_info: storage_provider_interfaces.DataInfo,
) -> str | None:
    """Preloads artifact value is it's small enough (e.g. <=255 bytes)"""
    if not data_info.is_dir and data_info.total_size < _MAX_PRELOAD_VALUE_SIZE:
        # Don't fail the execution if small value preloading fails.
        # Those values may be useful for preservation, but not so important that we should fail a successfully completed container execution.
        try:
            data = uri_reader.download_as_bytes()
        except Exception as ex:
            _logger.exception(f"Error during preloading small artifact values.")
            return None
        try:
            text = data.decode("utf-8")
            return text
        except:
            pass