Add debugging information to orchestrator_sql.py

huggingface_overlay/cloud_pipelines_backend/orchestrator_sql.py

@@ -0,0 +1,1005 @@
+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