content: #  Copyright (c) ZenML GmbH 2025. All Rights Reserved.
  #
  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at:
  #
  #       https://www.apache.org/licenses/LICENSE-2.0
  #
  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
  #  or implied. See the License for the specific language governing
  #  permissions and limitations under the License.
  """Utility functions for runs."""
  
  from typing import TYPE_CHECKING, Dict, List, Optional, Set, cast
  
  from zenml.enums import ExecutionStatus, StackComponentType
  from zenml.exceptions import IllegalOperationError
  from zenml.logger import get_logger
  from zenml.models import PipelineRunResponse, PipelineRunUpdate, StepRunUpdate
  from zenml.orchestrators.base_orchestrator import BaseOrchestrator
  from zenml.stack.stack_component import StackComponent
  
  if TYPE_CHECKING:
      from zenml.zen_stores.base_zen_store import BaseZenStore
  
  logger = get_logger(__name__)
  
  
  def stop_run(run: PipelineRunResponse, graceful: bool = False) -> None:
      """Stop a pipeline run.
  
      Args:
          run: The pipeline run to stop.
          graceful: Whether to stop the run gracefully.
  
      Raises:
          IllegalOperationError: If the run is already stopped or being stopped.
          ValueError: If the stack is not accessible.
      """
      # Check if the stack is still accessible
      if run.stack is None:
          raise ValueError(
              "The stack that this pipeline run response was executed on "
              "is either not accessible or has been deleted."
          )
  
      # Check if pipeline can be stopped
      if run.status == ExecutionStatus.COMPLETED:
          raise IllegalOperationError(
              "Cannot stop a run that is already completed."
          )
  
      if run.status == ExecutionStatus.STOPPED:
          raise IllegalOperationError("Run is already stopped.")
  
      if run.status == ExecutionStatus.STOPPING:
          raise IllegalOperationError("Run is already being stopped.")
  
      # Check if the stack is still accessible
      orchestrator_list = run.stack.components.get(
          StackComponentType.ORCHESTRATOR, []
      )
      if len(orchestrator_list) == 0:
          raise ValueError(
              "The orchestrator that this pipeline run response was "
              "executed with is either not accessible or has been deleted."
          )
  
      orchestrator = cast(
          BaseOrchestrator,
          StackComponent.from_model(component_model=orchestrator_list[0]),
      )
  
      # Stop the run
      orchestrator.stop_run(run=run, graceful=graceful)
  
  
  def refresh_run_status(
      run: PipelineRunResponse,
      include_step_updates: bool = False,
      zen_store: Optional["BaseZenStore"] = None,
  ) -> PipelineRunResponse:
      """Refresh the status of a pipeline run if it is initializing/running.
  
      Args:
          run: The pipeline run to refresh.
          include_step_updates: Flag deciding whether we should also refresh
              the status of individual steps.
          zen_store: Optional ZenStore to use for updating the run. If not provided,
              the ZenStore will be fetched from the Client. This is mainly useful
              for running inside the ZenML server.
  
      Returns:
          The refreshed pipeline run.
  
      Raises:
          ValueError: If the stack or the orchestrator of the run is deleted.
      """
      # Check if the stack still accessible
      if run.stack is None:
          raise ValueError(
              "The stack that this pipeline run response was executed on"
              "is either not accessible or has been deleted."
          )
  
      # Check if the orchestrator is still accessible
      orchestrator_list = run.stack.components.get(
          StackComponentType.ORCHESTRATOR, []
      )
      if len(orchestrator_list) == 0:
          raise ValueError(
              "The orchestrator that this pipeline run response was "
              "executed with is either not accessible or has been deleted."
          )
  
      orchestrator = cast(
          BaseOrchestrator,
          StackComponent.from_model(component_model=orchestrator_list[0]),
      )
  
      # Fetch the status (and optionally step statuses)
      pipeline_status, step_statuses = orchestrator.fetch_status(
          run=run, include_steps=include_step_updates
      )
  
      if zen_store is None:
          from zenml.client import Client
  
          zen_store = Client().zen_store
      else:
          zen_store = zen_store
  
      # Update step statuses
      if include_step_updates:
          if step_statuses:
              current_steps = run.steps
              for step_name, step_status in step_statuses.items():
                  # If step is not in the run yet, skip
                  if step_name not in current_steps:
                      continue
  
                  # If the step is already finished, skip
                  if current_steps[step_name].status.is_finished:
                      continue
  
                  # Update only if the status has changed
                  if step_status != current_steps[step_name].status:
                      try:
                          zen_store.update_run_step(
                              step_run_id=current_steps[step_name].id,
                              step_run_update=StepRunUpdate(status=step_status),
                          )
                      except Exception as e:
                          logger.warning(
                              f"Failed to update status for step '{step_name}': {e}"
                          )
  
      # Update pipeline status
      if pipeline_status is not None and pipeline_status != run.status:
          return zen_store.update_run(
              run_id=run.id,
              run_update=PipelineRunUpdate(status=pipeline_status),
          )
  
      return run
  
  
  def build_dag(steps: Dict[str, List[str]]) -> Dict[str, Set[str]]:
      """Build DAG with downstream steps from a list of steps.
  
      Args:
          steps: The list of steps.
  
      Returns:
          The DAG with downstream steps.
      """
      dag: Dict[str, Set[str]] = {step: set() for step in steps}
  
      for step_name, upstream_steps in steps.items():
          for upstream_step in upstream_steps:
              dag[upstream_step].add(step_name)
  
      return dag
  
  
  def find_all_downstream_steps(
      step_name: str, dag: Dict[str, Set[str]]
  ) -> Set[str]:
      """Find all downstream steps of a given step.
  
      Args:
          step_name: The name of the step to find the downstream steps of.
          dag: The DAG.
  
      Returns:
          The set of downstream steps.
      """
      downstream_steps = set()
      for downstream_step in dag[step_name]:
          downstream_steps.add(downstream_step)
          downstream_steps.update(
              find_all_downstream_steps(downstream_step, dag)
          )
      return downstream_steps
  